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) {