From fc81355fe1469e6f97ff8f2d3ed0a3caccda60c2 Mon Sep 17 00:00:00 2001
From: OrcaWin <alpha-eng@stably.ai>
Date: Sat, 12 Sep 2026 21:05:50 -0700
Subject: [PATCH] perf: accelerate cancellable remote transcript line scanning
 (#20351)

* perf: search remote transcript newlines directly

* fix: bound newline search by the yield window so cancellation stays observable

A newline-free segment jumped straight to the next line break, skipping the
character-count yield and its abort checks. Cap each jump at the yield window
and yield there so a large single-line transcript still stops promptly.

---------

Co-authored-by: Orca Worker <orca-worker@localhost>
Co-authored-by: Neil <neil@stably.ai>
---
 .../remote-session-content-lines.test.ts      | 59 +++++++++++++++++++
 .../ai-vault/remote-session-content-lines.ts  | 24 ++++++--
 2 files changed, 79 insertions(+), 4 deletions(-)
 create mode 100644 src/main/ai-vault/remote-session-content-lines.test.ts

diff --git a/src/main/ai-vault/remote-session-content-lines.test.ts b/src/main/ai-vault/remote-session-content-lines.test.ts
new file mode 100644
index 00000000000..348087dc728
--- /dev/null
+++ b/src/main/ai-vault/remote-session-content-lines.test.ts
@@ -0,0 +1,59 @@
+import { describe, expect, it } from 'vitest'
+import { remoteSessionContentLines } from './remote-session-content-lines'
+
+// Mirrors REMOTE_CONTENT_YIELD_CHAR_COUNT in the implementation.
+const YIELD_CHAR_COUNT = 256 * 1024
+
+async function collect(content: string, signal: AbortSignal): Promise<string[]> {
+  const lines: string[] = []
+  for await (const line of remoteSessionContentLines(content, signal)) {
+    lines.push(line)
+  }
+  return lines
+}
+
+describe('remote session content lines', () => {
+  it.each([
+    ['', ['']],
+    ['\n', ['', '']],
+    ['one\r\ntwo\n', ['one', 'two', '']],
+    ['one\rtwo\r', ['one\rtwo']],
+    ['x'.repeat(300_000), ['x'.repeat(300_000)]],
+    [`${'x'.repeat(YIELD_CHAR_COUNT + 1)}\ny`, ['x'.repeat(YIELD_CHAR_COUNT + 1), 'y']],
+    [`${'x'.repeat(YIELD_CHAR_COUNT)}\r\ny`, ['x'.repeat(YIELD_CHAR_COUNT), 'y']],
+    [`${'x'.repeat(YIELD_CHAR_COUNT * 2)}\ny`, ['x'.repeat(YIELD_CHAR_COUNT * 2), 'y']]
+  ])('preserves line boundaries for input %#', async (content, expected) => {
+    expect(await collect(content as string, new AbortController().signal)).toEqual(expected)
+  })
+
+  it('rejects an already cancelled scan even for empty content', async () => {
+    const controller = new AbortController()
+    controller.abort()
+    await expect(collect('', controller.signal)).rejects.toThrow()
+  })
+
+  it.each(['\n'.repeat(400), `${'x'.repeat(300_000)}\nlast`])(
+    'observes cancellation at an event-loop yield for input %#',
+    async (content) => {
+      const controller = new AbortController()
+      setImmediate(() => controller.abort())
+      await expect(collect(content, controller.signal)).rejects.toThrow()
+    }
+  )
+
+  it('observes cancellation inside a newline-free segment before emitting its line', async () => {
+    const controller = new AbortController()
+    const seen: string[] = []
+    setImmediate(() => controller.abort())
+    const scan = (async () => {
+      for await (const line of remoteSessionContentLines(
+        'x'.repeat(YIELD_CHAR_COUNT * 3),
+        controller.signal
+      )) {
+        seen.push(line)
+      }
+    })()
+    await expect(scan).rejects.toThrow()
+    expect(seen).toEqual([])
+  })
+})
diff --git a/src/main/ai-vault/remote-session-content-lines.ts b/src/main/ai-vault/remote-session-content-lines.ts
index 9a107fd1ef3..c20c6ec10ac 100644
--- a/src/main/ai-vault/remote-session-content-lines.ts
+++ b/src/main/ai-vault/remote-session-content-lines.ts
@@ -27,7 +27,19 @@ async function* cancellableContentLines(
 
   for (let index = 0; index <= content.length; index++) {
     if (index < content.length && content.charCodeAt(index) !== 10) {
-      continue
+      const newline = content.indexOf('\n', index)
+      const lineBreak = newline === -1 ? content.length : newline
+      // Bound the jump so a newline-free segment still observes cancellation.
+      const windowEnd = yieldStart + REMOTE_CONTENT_YIELD_CHAR_COUNT
+      if (lineBreak > windowEnd) {
+        await yieldUnlessCancelled(signal)
+        linesSinceYield = 0
+        yieldStart = windowEnd
+        // Resume the search at windowEnd itself; the loop increment lands there.
+        index = windowEnd - 1
+        continue
+      }
+      index = lineBreak
     }
     const lineEnd = index > lineStart && content.charCodeAt(index - 1) === 13 ? index - 1 : index
     yield content.slice(lineStart, lineEnd)
@@ -37,11 +49,15 @@ async function* cancellableContentLines(
       linesSinceYield >= REMOTE_CONTENT_YIELD_LINE_COUNT ||
       index - yieldStart >= REMOTE_CONTENT_YIELD_CHAR_COUNT
     ) {
-      throwIfAiVaultScanCancelled(signal)
-      await yieldToEventLoop()
-      throwIfAiVaultScanCancelled(signal)
+      await yieldUnlessCancelled(signal)
       linesSinceYield = 0
       yieldStart = index
     }
   }
 }
+
+async function yieldUnlessCancelled(signal: AbortSignal): Promise<void> {
+  throwIfAiVaultScanCancelled(signal)
+  await yieldToEventLoop()
+  throwIfAiVaultScanCancelled(signal)
+}