From 2b2aacecef2684e202f394c8396ee72743188aa8 Mon Sep 17 00:00:00 2001 From: Vladimir Krivosheev Date: Mon, 29 Jun 2026 07:18:18 +0200 Subject: [PATCH] IJPL-248623 Avoid eager document text copies in Agent Workbench Use immutable document character sequences on prompt and context paths so completion checks and context extraction do not materialize whole document strings before they need payload text. Console test-runner context now streams and truncates normalized output from the immutable sequence, preserving existing trimming and truncation semantics while avoiding full console text copies. GitOrigin-RevId: 78ab1df8dc97b2436a60aa6f778de74da5f539ac --- .../src/AgentPromptEditorContextSupport.kt | 15 +- ...ntPromptTestSelectionContextContributor.kt | 140 ++++++++++++++---- ...omptTestSelectionContextContributorTest.kt | 34 +++++ ...gentPromptClaudeSlashCompletionProvider.kt | 31 ++-- .../AgentPromptPaletteSessionController.kt | 4 +- ...AgentWorkbenchOpenScratchAsPromptAction.kt | 7 +- 6 files changed, 181 insertions(+), 50 deletions(-) diff --git a/plugins/agent-workbench/prompt/context/src/AgentPromptEditorContextSupport.kt b/plugins/agent-workbench/prompt/context/src/AgentPromptEditorContextSupport.kt index 35321de281a6..e86ef2925487 100644 --- a/plugins/agent-workbench/prompt/context/src/AgentPromptEditorContextSupport.kt +++ b/plugins/agent-workbench/prompt/context/src/AgentPromptEditorContextSupport.kt @@ -19,7 +19,6 @@ import com.intellij.openapi.editor.Editor import com.intellij.openapi.fileEditor.FileDocumentManager import com.intellij.openapi.fileEditor.FileEditorManager import com.intellij.openapi.project.Project -import com.intellij.openapi.util.TextRange import com.intellij.openapi.vfs.VirtualFile import com.intellij.psi.PsiDocumentManager import com.intellij.psi.PsiElement @@ -216,10 +215,11 @@ object AgentPromptEditorContextSupport { .map { range -> range.shiftRight(elementStartOffset) } .firstOrNull { range -> range.containsOffset(offset) } ?: return null - if (absoluteReferenceRange.startOffset < 0 || absoluteReferenceRange.endOffset > document.textLength) { + val documentChars = document.immutableCharSequence + if (absoluteReferenceRange.startOffset < 0 || absoluteReferenceRange.endOffset > documentChars.length) { return null } - val referenceText = document.getText(absoluteReferenceRange) + val referenceText = documentChars.subSequence(absoluteReferenceRange.startOffset, absoluteReferenceRange.endOffset).toString() return extractReferenceSymbolName(referenceText, offset - absoluteReferenceRange.startOffset) } @@ -313,11 +313,12 @@ object AgentPromptEditorContextSupport { } private fun truncate(document: Document, startOffset: Int, endOffsetExclusive: Int): TruncateOutcome { - val safeStartOffset = startOffset.coerceIn(0, document.textLength) - val safeEndOffset = endOffsetExclusive.coerceIn(safeStartOffset, document.textLength) + val documentChars = document.immutableCharSequence + val safeStartOffset = startOffset.coerceIn(0, documentChars.length) + val safeEndOffset = endOffsetExclusive.coerceIn(safeStartOffset, documentChars.length) val originalChars = safeEndOffset - safeStartOffset if (originalChars <= AGENT_PROMPT_MAX_SNIPPET_CHARS) { - val text = document.getText(TextRange(safeStartOffset, safeEndOffset)) + val text = documentChars.subSequence(safeStartOffset, safeEndOffset).toString() return TruncateOutcome( text = text, originalChars = originalChars, @@ -327,7 +328,7 @@ object AgentPromptEditorContextSupport { ) } val includedEndOffset = safeStartOffset + AGENT_PROMPT_MAX_SNIPPET_CHARS - val prefix = document.getText(TextRange(safeStartOffset, includedEndOffset)) + val prefix = documentChars.subSequence(safeStartOffset, includedEndOffset).toString() val truncatedText = buildString(AGENT_PROMPT_MAX_SNIPPET_CHARS + 32) { append(prefix) append("\n...[truncated]") diff --git a/plugins/agent-workbench/prompt/testrunner/src/context/AgentPromptTestSelectionContextContributor.kt b/plugins/agent-workbench/prompt/testrunner/src/context/AgentPromptTestSelectionContextContributor.kt index 2b38733dedab..0df3fa012c6a 100644 --- a/plugins/agent-workbench/prompt/testrunner/src/context/AgentPromptTestSelectionContextContributor.kt +++ b/plugins/agent-workbench/prompt/testrunner/src/context/AgentPromptTestSelectionContextContributor.kt @@ -146,18 +146,25 @@ internal class AgentPromptTestSelectionContextContributor : AgentPromptContextCo private fun extractConsoleOutput(dataContext: DataContext): ConsoleOutputExcerpt? { val editor = consoleEditor(dataContext) ?: return null - val selectedText = editor.selectionModel.selectedText - ?.let(::normalizeConsoleOutput) - ?.takeIf { it.isNotEmpty() } - if (selectedText != null) { - return truncateConsoleOutput(selectedText, fromSelection = true) + val documentChars = editor.document.immutableCharSequence + val selectionModel = editor.selectionModel + if (selectionModel.hasSelection()) { + buildConsoleOutputExcerpt( + rawText = documentChars, + startOffset = selectionModel.selectionStart, + endOffsetExclusive = selectionModel.selectionEnd, + fromSelection = true, + )?.let { excerpt -> + return excerpt + } } - val documentText = normalizeConsoleOutput(editor.document.text) - if (documentText.isEmpty()) { - return null - } - return truncateConsoleOutput(documentText, fromSelection = false) + return buildConsoleOutputExcerpt( + rawText = documentChars, + startOffset = 0, + endOffsetExclusive = documentChars.length, + fromSelection = false, + ) } private fun consoleEditor(dataContext: DataContext): Editor? { @@ -263,34 +270,109 @@ internal class AgentPromptTestSelectionContextContributor : AgentPromptContextCo } } -private fun normalizeConsoleOutput(rawText: String): String { - val normalizedNewlines = rawText - .replace("\r\n", "\n") - .replace('\r', '\n') - val lines = normalizedNewlines.lines() - val firstNonBlank = lines.indexOfFirst { line -> line.isNotBlank() } - if (firstNonBlank < 0) { - return "" - } - val lastNonBlank = lines.indexOfLast { line -> line.isNotBlank() } - return lines.subList(firstNonBlank, lastNonBlank + 1).joinToString(separator = "\n") -} +private fun buildConsoleOutputExcerpt( + rawText: CharSequence, + startOffset: Int, + endOffsetExclusive: Int, + fromSelection: Boolean, +): ConsoleOutputExcerpt? { + val safeStartOffset = startOffset.coerceIn(0, rawText.length) + val safeEndOffset = endOffsetExclusive.coerceIn(safeStartOffset, rawText.length) + val includedText = StringBuilder(MAX_CONSOLE_OUTPUT_CHARS.coerceAtMost(safeEndOffset - safeStartOffset)) + var originalChars = 0 + var hasRetainedLine = false + var pendingBlankLines = 0 -private fun truncateConsoleOutput(text: String, fromSelection: Boolean): ConsoleOutputExcerpt { - val includedText = if (text.length <= MAX_CONSOLE_OUTPUT_CHARS) { - text + forEachNormalizedLine(rawText, safeStartOffset, safeEndOffset) { lineStart, lineEnd -> + if (isBlankLine(rawText, lineStart, lineEnd)) { + if (hasRetainedLine) { + pendingBlankLines++ + } + return@forEachNormalizedLine + } + + if (hasRetainedLine) { + repeat(pendingBlankLines + 1) { + appendIncludedNewline(includedText) + originalChars++ + } + } + + appendIncludedRange(includedText, rawText, lineStart, lineEnd) + originalChars += lineEnd - lineStart + hasRetainedLine = true + pendingBlankLines = 0 } - else { - text.take(MAX_CONSOLE_OUTPUT_CHARS) + + if (!hasRetainedLine) { + return null } + return ConsoleOutputExcerpt( - text = includedText, + text = includedText.toString(), fromSelection = fromSelection, - originalChars = text.length, + originalChars = originalChars, includedChars = includedText.length, ) } +private inline fun forEachNormalizedLine( + text: CharSequence, + startOffset: Int, + endOffsetExclusive: Int, + action: (lineStart: Int, lineEnd: Int) -> Unit, +) { + var lineStart = startOffset + var index = startOffset + while (index < endOffsetExclusive) { + when (text[index]) { + '\n' -> { + action(lineStart, index) + index++ + lineStart = index + } + '\r' -> { + action(lineStart, index) + index++ + if (index < endOffsetExclusive && text[index] == '\n') { + index++ + } + lineStart = index + } + else -> index++ + } + } + action(lineStart, endOffsetExclusive) +} + +private fun isBlankLine(text: CharSequence, startOffset: Int, endOffsetExclusive: Int): Boolean { + for (index in startOffset until endOffsetExclusive) { + if (!text[index].isWhitespace()) { + return false + } + } + return true +} + +private fun appendIncludedRange( + target: StringBuilder, + text: CharSequence, + startOffset: Int, + endOffsetExclusive: Int, +) { + var index = startOffset + while (index < endOffsetExclusive && target.length < MAX_CONSOLE_OUTPUT_CHARS) { + target.append(text[index]) + index++ + } +} + +private fun appendIncludedNewline(target: StringBuilder) { + if (target.length < MAX_CONSOLE_OUTPUT_CHARS) { + target.append('\n') + } +} + private fun renderLine(entry: SelectedTestContext): String { val anchor = entry.reference val assertionMessage = entry.assertionMessage ?: return "${entry.status}: $anchor" diff --git a/plugins/agent-workbench/prompt/testrunner/testSrc/context/AgentPromptTestSelectionContextContributorTest.kt b/plugins/agent-workbench/prompt/testrunner/testSrc/context/AgentPromptTestSelectionContextContributorTest.kt index fbcc313e777d..058673422571 100644 --- a/plugins/agent-workbench/prompt/testrunner/testSrc/context/AgentPromptTestSelectionContextContributorTest.kt +++ b/plugins/agent-workbench/prompt/testrunner/testSrc/context/AgentPromptTestSelectionContextContributorTest.kt @@ -101,6 +101,38 @@ class AgentPromptTestSelectionContextContributorTest { } } + @Test + fun trimsFocusedConsoleOutputBlankEdges() { + val project = ProjectManager.getInstance().defaultProject + val documentText = "\n \nAssertionError: boom\n\n at MainTest.test(MainTest.kt:42)\n\n" + val expectedOutput = "AssertionError: boom\n\n at MainTest.test(MainTest.kt:42)" + + runInEdtAndWait { + val editorFactory = EditorFactory.getInstance() + val editor = editorFactory.createViewer(editorFactory.createDocument(documentText), project, EditorKind.CONSOLE) + try { + val dataContext = SimpleDataContext.builder() + .add(CommonDataKeys.EDITOR, editor) + .add(AbstractTestProxy.DATA_KEY, testProxy(name = "testSingle", isDefect = true, errorMessage = "single failure")) + .build() + + val result = contributor.collect(invocationData(dataContext = dataContext)) + + assertThat(result).hasSize(1) + val item = result.single() + val payload = item.payload.objOrNull()!! + assertThat(payload.string("consoleOutput")).isEqualTo(expectedOutput) + assertThat(payload.bool("consoleOutputFromSelection")).isFalse() + assertThat(item.truncation.originalChars).isEqualTo(item.body.length + expectedOutput.length) + assertThat(item.truncation.includedChars).isEqualTo(item.body.length + expectedOutput.length) + assertThat(item.truncation.reason).isEqualTo(AgentPromptContextTruncationReason.NONE) + } + finally { + editorFactory.releaseEditor(editor) + } + } + } + @Test fun prefersSelectedConsoleTextOverFullDocument() { val project = ProjectManager.getInstance().defaultProject @@ -278,6 +310,8 @@ class AgentPromptTestSelectionContextContributorTest { val payload = item.payload.objOrNull()!! assertThat(payload.string("consoleOutput")).hasSize(4_000) assertThat(payload.bool("consoleOutputFromSelection")).isFalse() + assertThat(item.truncation.originalChars).isEqualTo(item.body.length + documentText.length) + assertThat(item.truncation.includedChars).isEqualTo(item.body.length + 4_000) assertThat(item.truncation.reason).isEqualTo(AgentPromptContextTruncationReason.SOURCE_LIMIT) } finally { diff --git a/plugins/agent-workbench/prompt/ui/src/AgentPromptClaudeSlashCompletionProvider.kt b/plugins/agent-workbench/prompt/ui/src/AgentPromptClaudeSlashCompletionProvider.kt index 6c2264db3725..6d67c7408a6f 100644 --- a/plugins/agent-workbench/prompt/ui/src/AgentPromptClaudeSlashCompletionProvider.kt +++ b/plugins/agent-workbench/prompt/ui/src/AgentPromptClaudeSlashCompletionProvider.kt @@ -60,32 +60,43 @@ internal class AgentPromptClaudeSlashCompletionProvider( } } -internal fun findClaudeSlashCompletionPrefix(text: String, offset: Int): String? { +internal fun findClaudeSlashCompletionPrefix(text: CharSequence, offset: Int): String? { val safeOffset = offset.coerceIn(0, text.length) - val tokenStart = text.lastIndexOfAny(charArrayOf(' ', '\n', '\t', '\r'), startIndex = safeOffset - 1) + 1 + val tokenStart = findCompletionTokenStart(text, safeOffset) if (tokenStart >= safeOffset) { return null } - val prefix = text.substring(tokenStart, safeOffset) + val prefix = text.subSequence(tokenStart, safeOffset).toString() return prefix.takeIf { candidate -> candidate.startsWith('/') } } -internal fun findCodexSkillCompletionPrefix(text: String, offset: Int): String? { +internal fun findCodexSkillCompletionPrefix(text: CharSequence, offset: Int): String? { val safeOffset = offset.coerceIn(0, text.length) - val tokenStart = text.lastIndexOfAny(charArrayOf(' ', '\n', '\t', '\r'), startIndex = safeOffset - 1) + 1 + val tokenStart = findCompletionTokenStart(text, safeOffset) if (tokenStart >= safeOffset) { return null } - val prefix = text.substring(tokenStart, safeOffset) + val prefix = text.subSequence(tokenStart, safeOffset).toString() return prefix.takeIf { candidate -> candidate.startsWith('$') } } +private fun findCompletionTokenStart(text: CharSequence, offset: Int): Int { + var index = offset - 1 + while (index >= 0) { + when (text[index]) { + ' ', '\n', '\t', '\r' -> return index + 1 + } + index-- + } + return 0 +} + internal fun shouldAutoPopupClaudeSlashCompletion( selectedProvider: AgentSessionProvider?, workingProjectPaths: Iterable, - text: String, + text: CharSequence, offsetAfterChange: Int, insertedFragment: CharSequence, ): Boolean { @@ -95,7 +106,7 @@ internal fun shouldAutoPopupClaudeSlashCompletion( if (insertedFragment.length != 1 || insertedFragment[0] != '/') { return false } - if (offsetAfterChange != 1 || !text.startsWith('/')) { + if (offsetAfterChange != 1 || text.isEmpty() || text[0] != '/') { return false } if (findClaudeSlashCompletionPrefix(text, offsetAfterChange) != "/") { @@ -106,7 +117,7 @@ internal fun shouldAutoPopupClaudeSlashCompletion( internal fun shouldAutoPopupCodexSkillCompletion( selectedProvider: AgentSessionProvider?, - text: String, + text: CharSequence, offsetAfterChange: Int, insertedFragment: CharSequence, ): Boolean { @@ -116,7 +127,7 @@ internal fun shouldAutoPopupCodexSkillCompletion( if (insertedFragment.length != 1 || insertedFragment[0] != '$') { return false } - if (offsetAfterChange != 1 || !text.startsWith('$')) { + if (offsetAfterChange != 1 || text.isEmpty() || text[0] != '$') { return false } return findCodexSkillCompletionPrefix(text, offsetAfterChange) == CODEX_SKILL_PREFIX.toString() diff --git a/plugins/agent-workbench/prompt/ui/src/AgentPromptPaletteSessionController.kt b/plugins/agent-workbench/prompt/ui/src/AgentPromptPaletteSessionController.kt index 7f760e33dce6..e03485431ee7 100644 --- a/plugins/agent-workbench/prompt/ui/src/AgentPromptPaletteSessionController.kt +++ b/plugins/agent-workbench/prompt/ui/src/AgentPromptPaletteSessionController.kt @@ -520,7 +520,7 @@ internal class AgentPromptPaletteSessionController( } val selectedProvider = providerSelector.selectedProvider?.bridge?.provider - val documentText = event.document.immutableCharSequence.toString() + val documentText = event.document.immutableCharSequence val sourceProjectBasePath = launcherProvider() ?.resolveSourceProject(invocationData) ?.basePath @@ -582,7 +582,7 @@ internal class AgentPromptPaletteSessionController( if (LookupManager.getActiveLookup(editor) != null) { return@withContext } - val text = editor.document.immutableCharSequence.toString() + val text = editor.document.immutableCharSequence val caretOffset = editor.caretModel.offset val currentPrefix = when (expectedPrefix) { '/' -> findClaudeSlashCompletionPrefix(text, caretOffset) diff --git a/plugins/agent-workbench/prompt/ui/src/actions/AgentWorkbenchOpenScratchAsPromptAction.kt b/plugins/agent-workbench/prompt/ui/src/actions/AgentWorkbenchOpenScratchAsPromptAction.kt index 773c81da9a08..67e9ea1d55d3 100644 --- a/plugins/agent-workbench/prompt/ui/src/actions/AgentWorkbenchOpenScratchAsPromptAction.kt +++ b/plugins/agent-workbench/prompt/ui/src/actions/AgentWorkbenchOpenScratchAsPromptAction.kt @@ -84,8 +84,11 @@ internal fun resolveScratchMarkdownPromptText(dataContext: DataContext): String? return null } - val promptText = document.immutableCharSequence.toString() - return promptText.takeIf(String::isNotBlank) + val promptText = document.immutableCharSequence + if (promptText.isBlank()) { + return null + } + return promptText.toString() } private fun openGlobalPromptWithInitialText(e: AnActionEvent, promptText: String) {