Sync public snapshot from freebuff-private

Author: github-actions[bot]

bun.lock

@@ -0,0 +1,74 @@
+import { describe, expect, test } from 'bun:test'
+
+import { appendTextToRootStream } from '../block-operations'
+
+import type { ContentBlock, TextContentBlock } from '../../types/chat'
+
+const stream = (text: string, chunkSize: number): ContentBlock[] => {
+  let blocks: ContentBlock[] = []
+  for (let i = 0; i < text.length; i += chunkSize) {
+    blocks = appendTextToRootStream(blocks, {
+      type: 'text',
+      text: text.slice(i, i + chunkSize),
+    })
+  }
+  return blocks
+}
+
+// thinkingId is a global counter, so compare everything else.
+const shape = (blocks: ContentBlock[]) =>
+  blocks.map((block) => {
+    const { content, textType, thinkingOpen } = block as TextContentBlock
+    return { type: block.type, content, textType, thinkingOpen }
+  })
+
+describe('appendTextToRootStream seam scan', () => {
+  const corpus = [
+    'plain text with no tags at all',
+    'before<think>thought</think>after',
+    '<think>a</think><think>b</think>c',
+    'code: if (a < b && c > d) { return <Item value={a} /> }',
+    'stray </think> close then more',
+    'a <thinker> is not a tag, nor is </thin k>',
+    `${'x'.repeat(5_000)}<think>late</think>tail`,
+    '<think></think>',
+    'open <think>never closed',
+  ]
+
+  test('chunked deltas produce the blocks of a single append', () => {
+    for (const text of corpus) {
+      const whole = shape(stream(text, text.length))
+      for (const chunkSize of [7, 8, 9, 64]) {
+        expect({ text, chunkSize, blocks: shape(stream(text, chunkSize)) })
+          .toEqual({ text, chunkSize, blocks: whole })
+      }
+    }
+  })
+
+  test('a tag split at every position after a long block opens thinking', () => {
+    const prefix = 'p'.repeat(50_000)
+    for (let split = 1; split < '<think>'.length; split++) {
+      let blocks = appendTextToRootStream([], {
+        type: 'text',
+        text: prefix + '<think>'.slice(0, split),
+      })
+      blocks = appendTextToRootStream(blocks, {
+        type: 'text',
+        text: '<think>'.slice(split) + 'body',
+      })
+      expect(shape(blocks)).toEqual([
+        { type: 'text', content: prefix, textType: 'text', thinkingOpen: undefined },
+        { type: 'text', content: 'body', textType: 'reasoning', thinkingOpen: true },
+      ])
+    }
+  })
+
+  test('a long reply full of less-than signs stays one text block', () => {
+    const line = 'if (a < b) { return <Row key={i} /> }\n'
+    const text = line.repeat(500)
+    const blocks = stream(text, 13)
+    expect(shape(blocks)).toEqual([
+      { type: 'text', content: text, textType: 'text', thinkingOpen: undefined },
+    ])
+  })
+})

cli/src/utils/__tests__/block-operations-streaming.test.ts

@@ -643,3 +643,66 @@ describe('createBatchedMessageUpdater timer behavior', () => {
     expect(intervalCallbacks.size).toBe(0)
   })
 })
+
+describe('appendRootText', () => {
+  const shape = (blocks: ContentBlock[] | undefined) =>
+    (blocks ?? []).map((block) => {
+      const { content, textType, thinkingOpen } = block as TextContentBlock
+      return { type: block.type, content, textType, thinkingOpen }
+    })
+
+  const deltas = [
+    { type: 'text' as const, text: 'answer ' },
+    { type: 'text' as const, text: 'with <think>a tagged' },
+    { type: 'text' as const, text: ' thought</think> and ' },
+    { type: 'reasoning' as const, text: 'native reasoning' },
+    { type: 'reasoning' as const, text: ' continues' },
+    { type: 'text' as const, text: 'then text' },
+  ]
+
+  test('a batched flush matches applying each delta immediately', () => {
+    let immediate: ChatMessage[] = [{ ...baseMessages[0] }]
+    const plain = createMessageUpdater('ai-1', (fn) => {
+      immediate = fn(immediate)
+    })
+    for (const delta of deltas) plain.appendRootText(delta)
+
+    let batchedState: ChatMessage[] = [{ ...baseMessages[0] }]
+    let setCalls = 0
+    const batched = createBatchedMessageUpdater(
+      'ai-1',
+      (fn) => {
+        setCalls++
+        batchedState = fn(batchedState)
+      },
+      60_000,
+    )
+    for (const delta of deltas) batched.appendRootText(delta)
+    batched.flush()
+    batched.dispose()
+
+    expect(setCalls).toBe(1)
+    expect(shape(batchedState[0].blocks)).toEqual(shape(immediate[0].blocks))
+  })
+
+  test('other updates keep their place between text deltas', () => {
+    let state: ChatMessage[] = [{ ...baseMessages[0] }]
+    const batched = createBatchedMessageUpdater(
+      'ai-1',
+      (fn) => {
+        state = fn(state)
+      },
+      60_000,
+    )
+    batched.appendRootText({ type: 'text', text: 'before ' })
+    batched.appendRootText({ type: 'text', text: 'tool' })
+    batched.addBlock({ type: 'text', content: 'TOOL', textType: 'reasoning' })
+    batched.appendRootText({ type: 'text', text: 'after' })
+    batched.flush()
+    batched.dispose()
+
+    expect(
+      (state[0].blocks ?? []).map((b) => (b as TextContentBlock).content),
+    ).toEqual(['before tool', 'TOOL', 'after'])
+  })
+})

cli/src/utils/__tests__/message-updater.test.ts

@@ -170,13 +170,14 @@ const appendTextWithThinkParsingToBlocks = (
     lastBlock?.type === 'text' &&
     lastBlock.textType === 'text'
   ) {
-    const combinedText = lastBlock.content + text
-
-    if (
-      combinedText.includes(THINK_OPEN_TAG) ||
-      combinedText.includes(THINK_CLOSE_TAG)
-    ) {
-      textToParse = combinedText
+    // A plain text block never holds a whole tag (parsing consumed it), so a
+    // tag can only be completed across the seam: the block's last
+    // tag-length-minus-one characters plus this delta. Scanning just that keeps
+    // each delta O(delta) instead of O(block), which was quadratic over a long reply.
+    const seam = lastBlock.content.slice(-(THINK_CLOSE_TAG.length - 1)) + text
+
+    if (seam.includes(THINK_OPEN_TAG) || seam.includes(THINK_CLOSE_TAG)) {
+      textToParse = lastBlock.content + text
       nextBlocks.pop()
     }
   }

cli/src/utils/block-operations.ts

@@ -1,5 +1,9 @@
+import { appendTextToRootStream } from './block-operations'
+
 import type { ChatMessage, ContentBlock, TextContentBlock } from '../types/chat'
 
+export type RootTextDelta = { type: 'text' | 'reasoning'; text: string }
+
 // Small wrapper to avoid repeating the ai-message map/update pattern.
 export type SetMessagesFn = (
   updater: (messages: ChatMessage[]) => ChatMessage[],
@@ -10,6 +14,9 @@ export type MessageUpdater = {
   updateAiMessageBlocks: (
     blockUpdater: (blocks: ContentBlock[]) => ContentBlock[],
   ) => void
+  /** Append a streamed delta to the root message's blocks. The batched
+   *  updater merges consecutive deltas of one type into a single append. */
+  appendRootText: (delta: RootTextDelta) => void
   markComplete: (metadata?: Partial<ChatMessage>) => void
   setError: (message: string) => void
   /** Clears the userError field (e.g., when a new message is sent successfully) */
@@ -47,6 +54,10 @@ export const createMessageUpdater = (
     }))
   }
 
+  const appendRootText = (delta: RootTextDelta) => {
+    updateAiMessageBlocks((blocks) => appendTextToRootStream(blocks, delta))
+  }
+
   const addBlock = (block: ContentBlock) => {
     updateAiMessage((msg) => ({
       ...msg,
@@ -113,6 +124,7 @@ export const createMessageUpdater = (
   return {
     updateAiMessage,
     updateAiMessageBlocks,
+    appendRootText,
     markComplete,
     setError,
     clearUserError,
@@ -126,8 +138,13 @@ export const createBatchedMessageUpdater = (
   setMessages: SetMessagesFn,
   flushIntervalMs: number = DEFAULT_FLUSH_INTERVAL_MS,
 ): BatchedMessageUpdater => {
-  // Queue of message updater functions to be applied on next flush
-  const pendingUpdaters: Array<(msg: ChatMessage) => ChatMessage> = []
+  // Queue of message updater functions to be applied on next flush. A root
+  // text append stays a delta, not a closure, until the flush so the next
+  // delta of the same type can extend it: one block-array copy per flush
+  // instead of one per token.
+  const pendingUpdaters: Array<
+    ((msg: ChatMessage) => ChatMessage) | RootTextDelta
+  > = []
   let intervalId: ReturnType<typeof setInterval> | null = null
   let isDisposed = false
 
@@ -139,7 +156,13 @@ export const createBatchedMessageUpdater = (
 
     // Compose all pending updaters into a single transform
     const composedUpdater = (msg: ChatMessage): ChatMessage => {
-      return updaters.reduce((m, fn) => fn(m), msg)
+      return updaters.reduce(
+        (m, entry) =>
+          typeof entry === 'function'
+            ? entry(m)
+            : { ...m, blocks: appendTextToRootStream(m.blocks ?? [], entry) },
+        msg,
+      )
     }
 
     // Apply composed update to the target message
@@ -182,6 +205,23 @@ export const createBatchedMessageUpdater = (
     }))
   }
 
+  const appendRootText = (delta: RootTextDelta) => {
+    if (!delta.text) return
+    if (isDisposed) {
+      updateAiMessageBlocks((blocks) => appendTextToRootStream(blocks, delta))
+      return
+    }
+    const last = pendingUpdaters[pendingUpdaters.length - 1]
+    if (last && typeof last !== 'function' && last.type === delta.type) {
+      pendingUpdaters[pendingUpdaters.length - 1] = {
+        type: last.type,
+        text: last.text + delta.text,
+      }
+      return
+    }
+    pendingUpdaters.push({ type: delta.type, text: delta.text })
+  }
+
   const addBlock = (block: ContentBlock) => {
     updateAiMessage((msg) => ({
       ...msg,
@@ -269,6 +309,7 @@ export const createBatchedMessageUpdater = (
   return {
     updateAiMessage,
     updateAiMessageBlocks,
+    appendRootText,
     markComplete,
     setError,
     clearUserError,

cli/src/utils/message-updater.ts

@@ -7,7 +7,6 @@ import {
 } from '../ads/ad-signals'
 
 import {
-  appendTextToRootStream,
   appendToolToAgentBlock,
   closeNativeReasoningBlock,
   closeNativeReasoningInAgent,
@@ -131,9 +130,7 @@ const appendRootChunk = (state: EventHandlerState, delta: TextDelta) => {
     return
   }
 
-  state.message.updater.updateAiMessageBlocks((blocks) =>
-    appendTextToRootStream(blocks, delta),
-  )
+  state.message.updater.appendRootText(delta)
 
   if (
     state.mode.agentMode === 'PLAN' &&