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

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]")

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"

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 {

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<String?>,
-  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()

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)

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