From 70876c9790ec2e3733b7843a95e123ead762f412 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Sun, 2 Aug 2026 14:07:24 -0700 Subject: [PATCH] perf(gateway): remove streaming hot-path rescans (#118192) --- .../llm-core/src/utils/event-stream.test.ts | 52 +++++++++++ packages/llm-core/src/utils/event-stream.ts | 13 ++- .../embedded-agent-runner/run/payloads.ts | 4 +- ...-agent-subscribe.handlers.messages.test.ts | 23 ++++- ...edded-agent-subscribe.handlers.messages.ts | 40 ++++++--- ...embedded-agent-subscribe.handlers.types.ts | 3 + src/agents/embedded-agent-subscribe.ts | 8 +- src/agents/embedded-agent-utils.test.ts | 72 ++++++++++++++++ src/agents/embedded-agent-utils.ts | 86 ++++++++++++------- src/agents/session-activity-notes.test.ts | 63 ++++++++++++++ src/agents/session-activity-notes.ts | 45 +++++++++- src/gateway/chat-abort.ts | 14 +-- .../chat-display-projection.sanitize.ts | 6 +- src/gateway/server-chat-state.ts | 33 +++++++ src/gateway/server-chat.agent-events.test.ts | 50 +++++++++++ src/gateway/server-chat.ts | 20 ++--- .../server-methods/chat-abort-handler.ts | 2 +- .../server-methods/chat-abort-runtime.ts | 2 +- src/shared/chat-message-content.test.ts | 23 ++++- src/shared/chat-message-content.ts | 62 ++++++++----- 20 files changed, 518 insertions(+), 103 deletions(-) create mode 100644 packages/llm-core/src/utils/event-stream.test.ts create mode 100644 src/agents/session-activity-notes.test.ts diff --git a/packages/llm-core/src/utils/event-stream.test.ts b/packages/llm-core/src/utils/event-stream.test.ts new file mode 100644 index 000000000000..9808dc9f8125 --- /dev/null +++ b/packages/llm-core/src/utils/event-stream.test.ts @@ -0,0 +1,52 @@ +import { describe, expect, it } from "vitest"; +import { EventStream } from "./event-stream.js"; + +function createNumberStream(): EventStream { + return new EventStream( + (event) => event === -1, + (event) => event, + ); +} + +describe("EventStream", () => { + it("preserves interleaved queued and waiting push/pull order", async () => { + const stream = createNumberStream(); + const iterator = stream[Symbol.asyncIterator](); + + stream.push(1); + expect(await iterator.next()).toEqual({ value: 1, done: false }); + + const waiting = iterator.next(); + stream.push(2); + expect(await waiting).toEqual({ value: 2, done: false }); + + stream.push(3); + stream.end(); + expect(await iterator.next()).toEqual({ value: 3, done: false }); + expect(await iterator.next()).toEqual({ value: undefined, done: true }); + }); + + it("compacts a consumed queue prefix at the cursor boundary without losing events", async () => { + const stream = createNumberStream(); + const iterator = stream[Symbol.asyncIterator](); + for (let value = 0; value < 2048; value += 1) { + stream.push(value); + } + + for (let value = 0; value < 1024; value += 1) { + expect(await iterator.next()).toEqual({ value, done: false }); + } + const queueState = stream as unknown as { queue: number[]; queueHead: number }; + expect(queueState.queueHead).toBe(0); + expect(queueState.queue).toHaveLength(1024); + + for (let value = 2048; value < 2052; value += 1) { + stream.push(value); + } + stream.end(); + for (let value = 1024; value < 2052; value += 1) { + expect(await iterator.next()).toEqual({ value, done: false }); + } + expect(await iterator.next()).toEqual({ value: undefined, done: true }); + }); +}); diff --git a/packages/llm-core/src/utils/event-stream.ts b/packages/llm-core/src/utils/event-stream.ts index c4c91a7f873e..2678bcf8922c 100644 --- a/packages/llm-core/src/utils/event-stream.ts +++ b/packages/llm-core/src/utils/event-stream.ts @@ -8,6 +8,7 @@ import type { /** Generic async-iterable event stream with a separately awaited final result. */ export class EventStream implements AsyncIterable { private queue: T[] = []; + private queueHead = 0; private waiting: ((value: IteratorResult) => void)[] = []; private done = false; private finalResultPromise: Promise; @@ -63,10 +64,16 @@ export class EventStream implements AsyncIterable { async *[Symbol.asyncIterator](): AsyncIterator { while (true) { - if (this.queue.length > 0) { - for (const event of this.queue.splice(0, 1)) { - yield event; + if (this.queueHead < this.queue.length) { + const event = this.queue[this.queueHead] as T; + this.queueHead += 1; + // Compact only after a substantial consumed prefix reaches half the + // backing array, keeping dequeue amortized O(1) when consumers lag. + if (this.queueHead >= 1024 && this.queueHead * 2 >= this.queue.length) { + this.queue = this.queue.slice(this.queueHead); + this.queueHead = 0; } + yield event; } else if (this.done) { return; } else { diff --git a/src/agents/embedded-agent-runner/run/payloads.ts b/src/agents/embedded-agent-runner/run/payloads.ts index fefd1ac41e80..20c91fbe9d79 100644 --- a/src/agents/embedded-agent-runner/run/payloads.ts +++ b/src/agents/embedded-agent-runner/run/payloads.ts @@ -111,7 +111,7 @@ function resolveRawAssistantAnswerText(lastAssistant: AssistantMessage | undefin const record = block as { type?: unknown; textSignature?: unknown }; return ( isAssistantTextContentBlockType(record.type) && - Boolean(parseAssistantTextSignature(record.textSignature)?.phase) + Boolean(parseAssistantTextSignature(record)?.phase) ); }); if (!hasExplicitPhasedTextBlock) { @@ -121,7 +121,7 @@ function resolveRawAssistantAnswerText(lastAssistant: AssistantMessage | undefin return null; } const record = block as { type?: unknown; text?: unknown; textSignature?: unknown }; - const signature = parseAssistantTextSignature(record.textSignature); + const signature = parseAssistantTextSignature(record); if ( !isAssistantTextContentBlockType(record.type) || typeof record.text !== "string" || diff --git a/src/agents/embedded-agent-subscribe.handlers.messages.test.ts b/src/agents/embedded-agent-subscribe.handlers.messages.test.ts index 9b899b162d6c..bba3cf0ec884 100644 --- a/src/agents/embedded-agent-subscribe.handlers.messages.test.ts +++ b/src/agents/embedded-agent-subscribe.handlers.messages.test.ts @@ -24,6 +24,7 @@ import { createOpenAiResponsesTextBlock, createOpenAiResponsesTextEvent as createTextUpdateEvent, } from "./embedded-agent-subscribe.openai-responses.test-helpers.js"; +import { createThinkingTagStreamState } from "./embedded-agent-utils.js"; function updateMessage( context: EmbeddedAgentSubscribeContext, @@ -55,6 +56,7 @@ function createMessageUpdateContext( sourceReplyDeliveryMode?: "automatic" | "message_tool_only"; consumePartialReplyDirectives?: ReturnType; stripBlockTags?: ReturnType; + emitReasoningStream?: ReturnType; state?: Record; } = {}, ) { @@ -81,6 +83,7 @@ function createMessageUpdateContext( reasoningStreamOpen: false, streamReasoning: false, deltaBuffer: "", + thinkingTagStream: createThinkingTagStreamState(), blockBuffer: "", partialBlockState: { thinking: false, @@ -107,7 +110,7 @@ function createMessageUpdateContext( vi.fn((text: string, options?: { final?: boolean }) => partialReplyDirectiveAccumulator.consume(text, options), ), - emitReasoningStream: vi.fn(), + emitReasoningStream: params.emitReasoningStream ?? vi.fn(), flushBlockReplyBuffer: params.flushBlockReplyBuffer ?? vi.fn(), resetAssistantMessageState: params.resetAssistantMessageState ?? vi.fn(), recordAssistantUsage: vi.fn(), @@ -483,6 +486,24 @@ describe("handleMessageUpdate current-source message-tool previews", () => { }); describe("handleMessageUpdate text signatures", () => { + it("emits the full incrementally extracted reasoning value on every delta", () => { + const emitReasoningStream = vi.fn(); + const context = createMessageUpdateContext({ emitReasoningStream }); + + for (const chunk of ["reason", "ing"]) { + updateMessage( + context, + createTextUpdateEvent({ type: "text_delta", text: chunk, delta: chunk }), + ); + } + + expect(emitReasoningStream.mock.calls.map(([text]) => text)).toEqual([ + "", + "reason", + "reasoning", + ]); + }); + it("uses incremental text deltas for unphased OpenAI Responses streams", () => { const onAgentEvent = vi.fn(); const stripBlockTags = vi.fn((text: string) => text); diff --git a/src/agents/embedded-agent-subscribe.handlers.messages.ts b/src/agents/embedded-agent-subscribe.handlers.messages.ts index f08c84610e3e..faa5fa61f77d 100644 --- a/src/agents/embedded-agent-subscribe.handlers.messages.ts +++ b/src/agents/embedded-agent-subscribe.handlers.messages.ts @@ -40,6 +40,7 @@ import { extractAssistantThinking, extractAssistantCommentaryText, extractAssistantVisibleText, + createThinkingTagStreamState, extractThinkingFromTaggedStream, extractThinkingFromTaggedText, promoteThinkingTagsToBlocks, @@ -209,8 +210,11 @@ function resolveAssistantStreamItemId(params: { ? (indexedBlock as { type?: unknown }) : undefined; const hasIndexedTextBlock = indexedRecord?.type === "text"; - const candidateBlocks = hasIndexedTextBlock ? [indexedBlock] : content.toReversed(); - for (const block of candidateBlocks) { + const candidateStart = + hasIndexedTextBlock && contentIndex !== undefined ? contentIndex : content.length - 1; + const candidateEnd = hasIndexedTextBlock ? candidateStart : 0; + for (let index = candidateStart; index >= candidateEnd; index -= 1) { + const block = content[index]; if (!block || typeof block !== "object") { continue; } @@ -218,7 +222,7 @@ function resolveAssistantStreamItemId(params: { if (record.type !== "text") { continue; } - const signature = parseAssistantTextSignature(record.textSignature); + const signature = parseAssistantTextSignature(record); if (signature?.id) { return signature.id; } @@ -239,17 +243,24 @@ function scopeAssistantMessageToStreamBlock( return message; } const indexedBlock = contentIndex === undefined ? undefined : message.content[contentIndex]; - const block = + let block = indexedBlock && typeof indexedBlock === "object" && indexedBlock.type === "text" ? indexedBlock - : itemId - ? message.content.toReversed().find((candidate) => { - if (!candidate || typeof candidate !== "object" || candidate.type !== "text") { - return false; - } - return parseAssistantTextSignature(candidate.textSignature)?.id === itemId; - }) - : undefined; + : undefined; + if (!block && itemId) { + for (let index = message.content.length - 1; index >= 0; index -= 1) { + const candidate = message.content[index]; + if ( + candidate && + typeof candidate === "object" && + candidate.type === "text" && + parseAssistantTextSignature(candidate)?.id === itemId + ) { + block = candidate; + break; + } + } + } if (!block) { return message; } @@ -1035,7 +1046,9 @@ export function handleMessageUpdate( // Handle partial tags: stream whatever reasoning is visible so far. // Emit-always: emitReasoningStream reaches the bus/archive; rendering + // message_tool_only suppression are gated downstream (#92738). - ctx.emitReasoningStream(extractThinkingFromTaggedStream(ctx.state.deltaBuffer)); + ctx.emitReasoningStream( + extractThinkingFromTaggedStream(ctx.state.deltaBuffer, ctx.state.thinkingTagStream), + ); const wasThinking = ctx.state.partialBlockState.thinking; let visibleDelta = ""; // A text_start partial may already contain text that the following text_delta replays. @@ -1360,6 +1373,7 @@ export function handleMessageEnd( const finalizeMessageEnd = () => { ctx.state.deltaBuffer = ""; + ctx.state.thinkingTagStream = createThinkingTagStreamState(); ctx.state.blockBuffer = ""; ctx.blockChunker?.reset(); ctx.state.blockState.thinking = false; diff --git a/src/agents/embedded-agent-subscribe.handlers.types.ts b/src/agents/embedded-agent-subscribe.handlers.types.ts index 5eda9f1e813b..91762416a650 100644 --- a/src/agents/embedded-agent-subscribe.handlers.types.ts +++ b/src/agents/embedded-agent-subscribe.handlers.types.ts @@ -24,6 +24,7 @@ import type { BlockReplyChunking, SubscribeEmbeddedAgentSessionParams, } from "./embedded-agent-subscribe.types.js"; +import type { ThinkingTagStreamState } from "./embedded-agent-utils.js"; import type { McpAppChannelView } from "./mcp-ui-resource.js"; import type { AgentRunTimeoutPhase } from "./run-timeout-attribution.js"; import type { AgentMessage } from "./runtime/index.js"; @@ -101,6 +102,8 @@ export type EmbeddedAgentSubscribeState = { streamReasoning: boolean; deltaBuffer: string; + /** Scanner state shares deltaBuffer's lifecycle so each provider byte is parsed once. */ + thinkingTagStream: ThinkingTagStreamState; blockBuffer: string; blockState: { thinking: boolean; diff --git a/src/agents/embedded-agent-subscribe.ts b/src/agents/embedded-agent-subscribe.ts index e1891f35ff44..545f205826c0 100644 --- a/src/agents/embedded-agent-subscribe.ts +++ b/src/agents/embedded-agent-subscribe.ts @@ -58,7 +58,11 @@ import { filterToolResultMediaUrls, } from "./embedded-agent-subscribe.tools.js"; import type { SubscribeEmbeddedAgentSessionParams } from "./embedded-agent-subscribe.types.js"; -import { stripDowngradedToolCallText, THINKING_TAG_SCAN_RE } from "./embedded-agent-utils.js"; +import { + createThinkingTagStreamState, + stripDowngradedToolCallText, + THINKING_TAG_SCAN_RE, +} from "./embedded-agent-utils.js"; import { mediaUrlsFromGeneratedAttachments } from "./generated-attachments.js"; import { hasGeneratedMediaCompletionEvent } from "./internal-event-contract.js"; import type { AgentInternalEvent } from "./internal-events.js"; @@ -219,6 +223,7 @@ export function subscribeEmbeddedAgentSession(params: SubscribeEmbeddedAgentSess canShowReasoning && typeof params.onReasoningStream === "function", deltaBuffer: "", + thinkingTagStream: createThinkingTagStreamState(), blockBuffer: "", // Track if a streamed chunk opened a block (stateful across chunks). blockState: { thinking: false, final: false, inlineCode: createInlineCodeState() }, @@ -451,6 +456,7 @@ export function subscribeEmbeddedAgentSession(params: SubscribeEmbeddedAgentSess const resetAssistantMessageState = (nextAssistantTextBaseline: number) => { state.deltaBuffer = ""; + state.thinkingTagStream = createThinkingTagStreamState(); state.blockBuffer = ""; blockChunker?.reset(); replyDirectiveAccumulator.reset(); diff --git a/src/agents/embedded-agent-utils.test.ts b/src/agents/embedded-agent-utils.test.ts index 6f64cebc7d80..49fd35802566 100644 --- a/src/agents/embedded-agent-utils.test.ts +++ b/src/agents/embedded-agent-utils.test.ts @@ -8,11 +8,57 @@ import { extractAssistantText, extractAssistantThinking, extractAssistantVisibleText, + createThinkingTagStreamState, + extractThinkingFromTaggedStream, + extractThinkingFromTaggedText, formatReasoningMessage, promoteThinkingTagsToBlocks, stripDowngradedToolCallText, } from "./embedded-agent-utils.js"; +const REFERENCE_THINKING_TAG_NAME_PATTERN = String.raw`(?:(?:antml:|mm:)?(?:think(?:ing)?|thought)|antthinking)`; +const REFERENCE_THINKING_TAG_OPEN_RE = new RegExp( + String.raw`<\s*${REFERENCE_THINKING_TAG_NAME_PATTERN}\s*>`, + "gi", +); +const REFERENCE_THINKING_TAG_CLOSE_RE = new RegExp( + String.raw`<\s*\/\s*${REFERENCE_THINKING_TAG_NAME_PATTERN}\s*>`, + "gi", +); + +function extractThinkingFromTaggedStreamReference(text: string): string { + if (!text) { + return ""; + } + const closed = extractThinkingFromTaggedText(text); + if (closed) { + return closed; + } + const openMatches = [...text.matchAll(REFERENCE_THINKING_TAG_OPEN_RE)]; + const lastOpen = openMatches.at(-1); + if (!lastOpen) { + return ""; + } + const lastClose = [...text.matchAll(REFERENCE_THINKING_TAG_CLOSE_RE)].at(-1); + if (lastClose && (lastClose.index ?? -1) > (lastOpen.index ?? -1)) { + return closed; + } + return text.slice((lastOpen.index ?? 0) + lastOpen[0].length).trim(); +} + +function randomChunks(text: string, seed: number): string[] { + const chunks: string[] = []; + let offset = 0; + let value = seed; + while (offset < text.length) { + value = (value * 1664525 + 1013904223) >>> 0; + const length = Math.min(1 + (value % 7), text.length - offset); + chunks.push(text.slice(offset, offset + length)); + offset += length; + } + return chunks; +} + function makeAssistantMessage( message: Omit< AssistantMessage, @@ -40,6 +86,32 @@ function makeAssistantMessage( } as unknown as AssistantMessage; } +describe("extractThinkingFromTaggedStream", () => { + it("matches full-buffer extraction at every randomized chunk boundary", () => { + const cases = [ + "plain text with and no reasoning", + "Before first line\nsecond line after", + "prefixonemiddletwosuffix", + "surrounding text unfinished reasoning tail", + "< \nAnTThinking \t>spaced tag content< / antthinking > visible", + "closedunclosed trailing reasoning", + ]; + + for (const text of cases) { + for (let seed = 1; seed <= 24; seed += 1) { + const state = createThinkingTagStreamState(); + let prefix = ""; + for (const chunk of randomChunks(text, seed)) { + prefix += chunk; + expect(extractThinkingFromTaggedStream(prefix, state), `${text} (seed ${seed})`).toBe( + extractThinkingFromTaggedStreamReference(prefix), + ); + } + } + } + }); +}); + describe("extractAssistantText", () => { it("strips tool-only Minimax invocation XML from text", () => { const cases = [ diff --git a/src/agents/embedded-agent-utils.ts b/src/agents/embedded-agent-utils.ts index fbcd5af72035..7e69b6bf9a7e 100644 --- a/src/agents/embedded-agent-utils.ts +++ b/src/agents/embedded-agent-utils.ts @@ -84,7 +84,7 @@ function extractAssistantTextForPhase( if (!isAssistantTextContentBlockType(record.type)) { return false; } - return Boolean(parseAssistantTextSignature(record.textSignature)?.phase); + return Boolean(parseAssistantTextSignature(record)?.phase); }); let hadRequestedPhase = false; @@ -97,7 +97,7 @@ function extractAssistantTextForPhase( if (!isAssistantTextContentBlockType(record.type) || typeof record.text !== "string") { return null; } - const signature = parseAssistantTextSignature(record.textSignature); + const signature = parseAssistantTextSignature(record); const resolvedPhase = signature?.phase ?? (hasExplicitPhasedTextBlocks ? undefined : messagePhase); if (!shouldIncludeContent(resolvedPhase)) { @@ -208,19 +208,33 @@ const THINKING_TAG_CLOSE_RE = new RegExp( String.raw`<\s*\/\s*${THINKING_TAG_NAME_PATTERN}\s*>`, "i", ); -const THINKING_TAG_OPEN_GLOBAL_RE = new RegExp( - String.raw`<\s*${THINKING_TAG_NAME_PATTERN}\s*>`, - "gi", -); -const THINKING_TAG_CLOSE_GLOBAL_RE = new RegExp( - String.raw`<\s*\/\s*${THINKING_TAG_NAME_PATTERN}\s*>`, - "gi", -); /** Global regex used to scan provider-emitted thinking tags. */ export const THINKING_TAG_SCAN_RE = new RegExp( String.raw`<\s*(\/?)\s*${THINKING_TAG_NAME_PATTERN}\s*>`, "gi", ); +const THINKING_TAG_EXACT_RE = new RegExp( + String.raw`^<\s*(\/?)\s*${THINKING_TAG_NAME_PATTERN}\s*>$`, + "i", +); + +export type ThinkingTagStreamState = { + scannedOffset: number; + pendingTagStart?: number; + inThinking: boolean; + extracted: string; + lastMatchEnd: number; + lastTag?: { type: "open" | "close"; end: number }; +}; + +export function createThinkingTagStreamState(): ThinkingTagStreamState { + return { + scannedOffset: 0, + inThinking: false, + extracted: "", + lastMatchEnd: 0, + }; +} /** Split text that starts with thinking tags into structured thinking/text blocks. */ function splitThinkingTaggedText(text: string): ThinkTaggedSplitBlock[] | null { @@ -355,31 +369,41 @@ export function extractThinkingFromTaggedText(text: string): string { return result.trim(); } -/** Extract thinking-tag content from a possibly incomplete streaming payload. */ -export function extractThinkingFromTaggedStream(text: string): string { - if (!text) { - return ""; - } - const closed = extractThinkingFromTaggedText(text); - if (closed) { - return closed; +/** Incrementally extract thinking-tag content from a growing streaming payload. */ +export function extractThinkingFromTaggedStream( + text: string, + state: ThinkingTagStreamState, +): string { + for (let index = state.scannedOffset; index < text.length; index += 1) { + const char = text[index]; + if (char === "<") { + state.pendingTagStart = index; + continue; + } + if (char !== ">" || state.pendingTagStart === undefined) { + continue; + } + const start = state.pendingTagStart; + state.pendingTagStart = undefined; + const match = THINKING_TAG_EXACT_RE.exec(text.slice(start, index + 1)); + if (!match) { + continue; + } + if (state.inThinking) { + state.extracted += text.slice(state.lastMatchEnd, start); + } + const isClose = match[1] === "/"; + state.inThinking = !isClose; + state.lastMatchEnd = index + 1; + state.lastTag = { type: isClose ? "close" : "open", end: index + 1 }; } + state.scannedOffset = text.length; - const openMatches = [...text.matchAll(THINKING_TAG_OPEN_GLOBAL_RE)]; - if (openMatches.length === 0) { - return ""; - } - const closeMatches = [...text.matchAll(THINKING_TAG_CLOSE_GLOBAL_RE)]; - const lastOpen = openMatches.at(-1); - const lastClose = closeMatches.at(-1); - if (!lastOpen) { - return ""; - } - if (lastClose && (lastClose.index ?? -1) > (lastOpen.index ?? -1)) { + const closed = state.extracted.trim(); + if (closed || state.lastTag?.type !== "open") { return closed; } - const start = (lastOpen.index ?? 0) + lastOpen[0].length; - return text.slice(start).trim(); + return text.slice(state.lastTag.end).trim(); } /** Infer compact display metadata for a tool call from its args. */ diff --git a/src/agents/session-activity-notes.test.ts b/src/agents/session-activity-notes.test.ts new file mode 100644 index 000000000000..472cd70964ea --- /dev/null +++ b/src/agents/session-activity-notes.test.ts @@ -0,0 +1,63 @@ +import { describe, expect, it } from "vitest"; +import type { AgentEventPayload } from "../infra/agent-events.js"; +import { + createSessionActivityNoteState, + flushSessionActivityAssistantNote, + noteSessionActivityEvent, +} from "./session-activity-notes.js"; + +function assistantEvent(ts: number, text: string, delta: string): AgentEventPayload { + return { + runId: "run-1", + seq: ts, + stream: "assistant", + ts, + data: { text, delta }, + }; +} + +describe("session activity assistant buffering", () => { + it("defers cumulative rescans inside the throttle and flushes the exact latest text", () => { + const state = createSessionActivityNoteState(); + noteSessionActivityEvent(state, assistantEvent(1_000, "first", "first")); + noteSessionActivityEvent(state, assistantEvent(1_050, "first second", " second")); + + expect(state.assistantBuffer).toBe("first"); + expect(state.assistantBufferDirty).toBe(true); + + flushSessionActivityAssistantNote(state); + + expect(state.assistantBuffer).toBe("first second"); + expect(state.notes.at(-1)?.text).toBe("Assistant: first second"); + }); + + it("preserves split internal-context filtering when the forced flush follows a burst", () => { + const state = createSessionActivityNoteState(); + noteSessionActivityEvent( + state, + assistantEvent(1_000, "visible\n<<>>\n", "visible"), + ); + noteSessionActivityEvent( + state, + assistantEvent( + 1_050, + "visible\n<<>>\nsecret\n<<>>\nafter", + "secret\n<<>>\nafter", + ), + ); + + flushSessionActivityAssistantNote(state); + + expect(state.notes.at(-1)?.text).toBe("Assistant: visible after"); + }); + + it("keeps delta-only producers on the bounded incremental path", () => { + const state = createSessionActivityNoteState(); + noteSessionActivityEvent(state, assistantEvent(1_000, "", "a".repeat(5_000))); + noteSessionActivityEvent(state, assistantEvent(1_001, "", "tail")); + + expect(state.assistantBuffer).toHaveLength(4_096); + expect(state.assistantBuffer.endsWith("tail")).toBe(true); + expect(state.assistantBufferDirty).toBe(false); + }); +}); diff --git a/src/agents/session-activity-notes.ts b/src/agents/session-activity-notes.ts index be18b6fb1890..ec63ba3c29f2 100644 --- a/src/agents/session-activity-notes.ts +++ b/src/agents/session-activity-notes.ts @@ -17,6 +17,9 @@ export type SessionActivityNoteState = { noteBytes: number; itemStatuses: Map; assistantBuffer: string; + assistantRawBuffer: string; + assistantBufferDirty: boolean; + lastAssistantBufferAt: number; lastAssistantNote?: string; planProgress?: { completed: number; total: number }; }; @@ -26,10 +29,20 @@ const MAX_NOTE_BYTES = 8 * 1024; const DEFAULT_NOTE_MAX_CHARS = 360; const ASSISTANT_NOTE_MAX_CHARS = 240; const ASSISTANT_BUFFER_MAX_CHARS = 4096; +const ASSISTANT_BUFFER_THROTTLE_MS = 150; const MAX_ITEM_STATUSES = 160; export function createSessionActivityNoteState(): SessionActivityNoteState { - return { noteSequence: 0, notes: [], noteBytes: 0, itemStatuses: new Map(), assistantBuffer: "" }; + return { + noteSequence: 0, + notes: [], + noteBytes: 0, + itemStatuses: new Map(), + assistantBuffer: "", + assistantRawBuffer: "", + assistantBufferDirty: false, + lastAssistantBufferAt: 0, + }; } // Preserve an unmatched BEGIN while truncating so a later END can still strip the private block. @@ -49,6 +62,18 @@ function assembleAssistantBuffer(value: string, maxChars: number): string { return `${head}${INTERNAL_RUNTIME_CONTEXT_BEGIN}${body}`; } +function syncAssistantBuffer(state: SessionActivityNoteState, at = Date.now()): void { + if (!state.assistantBufferDirty) { + return; + } + state.assistantBuffer = assembleAssistantBuffer( + state.assistantRawBuffer, + ASSISTANT_BUFFER_MAX_CHARS, + ); + state.assistantBufferDirty = false; + state.lastAssistantBufferAt = at; +} + function keepUtf16SafeTail(value: string, maxChars: number): string { if (value.length <= maxChars) { return value; @@ -153,6 +178,8 @@ export function flushSessionActivityAssistantNote( state: SessionActivityNoteState, noteMaxChars: number = DEFAULT_NOTE_MAX_CHARS, ): void { + // Consumers force the latest cumulative snapshot; periodic assembly only keeps live state warm. + syncAssistantBuffer(state); // Redact assembled prose so split secrets match and raw fragments do not count as notes. if (!state.assistantBuffer || state.assistantBuffer.includes(INTERNAL_RUNTIME_CONTEXT_BEGIN)) { return; @@ -253,12 +280,26 @@ export function noteSessionActivityEvent( const full = readString(data.text); const delta = readString(data.delta); if (full) { - state.assistantBuffer = assembleAssistantBuffer(full, ASSISTANT_BUFFER_MAX_CHARS); + state.assistantRawBuffer = full; } else if (delta) { + // Delta-only producers never expose an unbounded cumulative string. Keep their historical + // bounded assembly path; its scan cost is capped independently of turn length. + syncAssistantBuffer(state, event.ts); state.assistantBuffer = assembleAssistantBuffer( state.assistantBuffer + delta, ASSISTANT_BUFFER_MAX_CHARS, ); + state.assistantRawBuffer = state.assistantBuffer; + state.lastAssistantBufferAt = event.ts; + return; + } else { + return; + } + // Runtime-context markers can straddle the retained tail, so only the full raw snapshot can + // be normalized safely. Bound that scan while preserving exact output at every consumer. + state.assistantBufferDirty = true; + if (event.ts - state.lastAssistantBufferAt >= ASSISTANT_BUFFER_THROTTLE_MS) { + syncAssistantBuffer(state, event.ts); } return; } diff --git a/src/gateway/chat-abort.ts b/src/gateway/chat-abort.ts index bf0579b267a8..dcec1c067177 100644 --- a/src/gateway/chat-abort.ts +++ b/src/gateway/chat-abort.ts @@ -264,7 +264,7 @@ function normalizeActiveAgentId(agentId: string | undefined): string | undefined */ export function resolveInFlightRunSnapshot(params: { chatAbortControllers: Map; - chatRunState: Pick; + chatRunState: Pick; requestedSessionKey: string; canonicalSessionKey: string; agentId?: string; @@ -330,10 +330,10 @@ export function resolveInFlightRunSnapshot(params: { // should still adopt the run and show a `streaming` status (not idle) and // render the result cleanly when it lands. const run = params.chatRunState.runs.get(best.runId); - const bufferedText = run?.buffer ?? ""; - const projected = projectLiveAssistantBufferedText(bufferedText, { - suppressLeadFragments: true, - }); + const projected = projectLiveAssistantBufferedText( + params.chatRunState.resolveBuffer(best.runId).text, + { suppressLeadFragments: true }, + ); const plan = run?.planSnapshot; return { runId: best.runId, @@ -383,7 +383,7 @@ export function boundInFlightRunSnapshotForChatHistory(params: { export type ChatAbortOps = { chatAbortControllers: Map; - chatRunState: Pick; + chatRunState: Pick; removeChatRun: ( sessionId: string, clientRunId: string, @@ -533,7 +533,7 @@ export function abortChatRunById( return { aborted: false }; } - const bufferedText = ops.chatRunState.runs.get(runId)?.buffer; + const bufferedText = ops.chatRunState.resolveBuffer(runId).text; const partialText = bufferedText && bufferedText.trim() ? bufferedText : undefined; ops.chatRunState.getOrCreate(runId).abortMarker = createChatAbortMarker(); if (stopReason) { diff --git a/src/gateway/chat-display-projection.sanitize.ts b/src/gateway/chat-display-projection.sanitize.ts index 931184358c57..adec0af4958a 100644 --- a/src/gateway/chat-display-projection.sanitize.ts +++ b/src/gateway/chat-display-projection.sanitize.ts @@ -136,9 +136,7 @@ function sanitizeAssistantPhasedContentBlocks(content: unknown[]): { return false; } const entry = block as { type?: unknown; textSignature?: unknown }; - return ( - entry.type === "text" && Boolean(parseAssistantTextSignature(entry.textSignature)?.phase) - ); + return entry.type === "text" && Boolean(parseAssistantTextSignature(entry)?.phase); }); if (!hasExplicitPhasedText) { return { content, changed: false }; @@ -151,7 +149,7 @@ function sanitizeAssistantPhasedContentBlocks(content: unknown[]): { if (entry.type !== "text") { return true; } - return parseAssistantTextSignature(entry.textSignature)?.phase === "final_answer"; + return parseAssistantTextSignature(entry)?.phase === "final_answer"; }); return { content: filtered, diff --git a/src/gateway/server-chat-state.ts b/src/gateway/server-chat-state.ts index 896affc2818e..60ac6ced5cdf 100644 --- a/src/gateway/server-chat-state.ts +++ b/src/gateway/server-chat-state.ts @@ -2,6 +2,10 @@ import type { AgentPlanStep } from "../channels/streaming.js"; // Gateway chat run state registries. // Tracks active runs, delta buffers, tool recipients, and session subscribers. import type { AgentEventPayload } from "../infra/agent-events.js"; +import { + normalizeLiveAssistantBufferedText, + projectLiveAssistantBufferedText, +} from "./live-chat-projector.js"; export type ChatRunTiming = { ackedAtMs: number; @@ -102,6 +106,8 @@ type ChatRunRecord = { registrations?: ChatRunEntry[]; rawBuffer?: string; buffer?: string; + /** Projection stays valid only while source matches rawBuffer; readers refresh it lazily. */ + bufferProjection?: { source: string; suppress: boolean }; planSnapshot?: ChatRunPlanSnapshot; /** Last time any buffered assistant text changed, including suppressed raw buffers. */ bufferUpdatedAt?: number; @@ -227,6 +233,7 @@ export type ChatRunState = { registry: ChatRunRegistry; toolEventRecipients: ToolEventRecipientRegistry; getOrCreate: (runId: string) => ChatRunRecord; + resolveBuffer: (runId: string) => { text: string; suppress: boolean }; hasAbortMarker: (runId: string) => boolean; deleteAbortMarker: (runId: string) => void; clearRun: (runId: string) => void; @@ -246,6 +253,7 @@ export function createChatRunState(): ChatRunState { } delete record.rawBuffer; delete record.buffer; + delete record.bufferProjection; delete record.planSnapshot; delete record.bufferUpdatedAt; delete record.deltaSentAt; @@ -259,11 +267,36 @@ export function createChatRunState(): ChatRunState { store.runs.clear(); }; + const resolveBuffer = (runId: string) => { + const record = store.runs.get(runId); + if (!record) { + return projectLiveAssistantBufferedText(""); + } + const rawText = record.rawBuffer; + if (rawText === undefined) { + return projectLiveAssistantBufferedText(record.buffer ?? ""); + } + if (record.bufferProjection?.source === rawText && record.buffer !== undefined) { + return { + text: record.buffer, + suppress: record.bufferProjection.suppress, + }; + } + // Protected blocks and directive tags can span delta frames, so the + // projection cache belongs to the complete merged raw buffer. + const normalizedText = normalizeLiveAssistantBufferedText(rawText); + const projected = projectLiveAssistantBufferedText(normalizedText); + record.buffer = projected.text; + record.bufferProjection = { source: rawText, suppress: projected.suppress }; + return projected; + }; + return { runs: store.runs, registry, toolEventRecipients, getOrCreate: store.getOrCreate, + resolveBuffer, hasAbortMarker: (runId) => store.runs.get(runId)?.abortMarker !== undefined, deleteAbortMarker: (runId) => { const record = store.runs.get(runId); diff --git a/src/gateway/server-chat.agent-events.test.ts b/src/gateway/server-chat.agent-events.test.ts index e617b4b32010..b73cb449c619 100644 --- a/src/gateway/server-chat.agent-events.test.ts +++ b/src/gateway/server-chat.agent-events.test.ts @@ -21,6 +21,7 @@ import { subscribePluginSessionsChanged } from "../plugins/gateway-events.js"; const persistGatewaySessionLifecycleEventMock = vi.fn(); const logErrorMock = vi.fn(); +const normalizeLiveAssistantBufferedTextMock = vi.hoisted(() => vi.fn()); vi.mock("./server-chat.persist-session-lifecycle.runtime.js", () => ({ persistGatewaySessionLifecycleEvent: (...args: unknown[]) => @@ -31,6 +32,17 @@ vi.mock("../logger.js", () => ({ logError: (...args: unknown[]) => logErrorMock(...args), })); +vi.mock("./live-chat-projector.js", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + normalizeLiveAssistantBufferedText: (text: string) => { + normalizeLiveAssistantBufferedTextMock(text); + return actual.normalizeLiveAssistantBufferedText(text); + }, + }; +}); + vi.mock("../config/io.js", () => ({ getRuntimeConfig: vi.fn(() => ({})), })); @@ -115,6 +127,7 @@ describe("agent event handler", () => { vi.mocked(loadGatewaySessionRow).mockReset().mockReturnValue(null); persistGatewaySessionLifecycleEventMock.mockReset().mockResolvedValue(undefined); logErrorMock.mockReset(); + normalizeLiveAssistantBufferedTextMock.mockReset(); }); afterEach(() => { @@ -568,6 +581,43 @@ describe("agent event handler", () => { nowSpy?.mockRestore(); }); + it("sanitizes only broadcasted assistant buffers while preserving cross-frame tags", () => { + let now = 10_000; + const nowSpy = vi.spyOn(Date, "now").mockImplementation(() => now); + const { broadcast, chatRunState, handler } = createHarness(); + registerNamedChatRun(chatRunState, "lazy-sanitize"); + + const deltas = [ + "Visible", + `\n${INTERNAL_RUNTIME_CONTEXT_BEGIN.slice(0, 20)}`, + `${INTERNAL_RUNTIME_CONTEXT_BEGIN.slice(20)}\nprivate runtime detail\n`, + ...Array.from({ length: 16 }, (_, index) => `private fragment ${index}\n`), + INTERNAL_RUNTIME_CONTEXT_END.slice(0, 18), + `${INTERNAL_RUNTIME_CONTEXT_END.slice(18)}\nAfter [[reply_`, + "to_current]] done", + ]; + deltas.forEach((delta, index) => { + now = 10_000 + index; + emitAgentEvent(handler, "run-lazy-sanitize", "assistant", { delta }, { seq: index + 1 }); + }); + + expect(normalizeLiveAssistantBufferedTextMock).toHaveBeenCalledTimes(1); + emitLifecycleEnd(handler, "run-lazy-sanitize", deltas.length + 1); + expect(normalizeLiveAssistantBufferedTextMock).toHaveBeenCalledTimes(2); + + const payloads = chatBroadcastCalls(broadcast).map(([, payload]) => payload) as Array<{ + state?: string; + message?: { content?: Array<{ text?: string }> }; + }>; + expect(payloads.map((payload) => payload.message?.content?.[0]?.text)).toEqual([ + "Visible", + "Visible\n\nAfter done", + "Visible\n\nAfter done", + ]); + expect(JSON.stringify(payloads)).not.toContain("private runtime detail"); + nowSpy.mockRestore(); + }); + it("emits the first assistant chat.send timing event to the originating Control UI", () => { const { broadcastToConnIds, chatRunState, handler, nowSpy } = createHarness({ now: 1_000 }); registerChatRun(chatRunState, "run-1", "session-1", "client-1", { diff --git a/src/gateway/server-chat.ts b/src/gateway/server-chat.ts index 80aa9623888d..6feef24f4583 100644 --- a/src/gateway/server-chat.ts +++ b/src/gateway/server-chat.ts @@ -35,7 +35,6 @@ import { import { resolveAssistantEventPhase } from "../shared/chat-message-content.js"; import { setSafeTimeout } from "../utils/timer-delay.js"; import { - normalizeLiveAssistantBufferedText, projectLiveAssistantBufferedText, resolveAssistantLiveChatInput, resolveMergedAssistantText, @@ -909,22 +908,15 @@ export function createAgentEventHandler({ const now = Date.now(); run.rawBuffer = mergedRawText; run.bufferUpdatedAt = now; - // Sanitize only after merging. Protected blocks and directive tags can span - // delta frames; cleaning each frame independently can expose their contents. - const normalizedText = normalizeLiveAssistantBufferedText(mergedRawText); - const projected = projectLiveAssistantBufferedText(normalizedText); - const mergedText = projected.text; - run.buffer = mergedText; - if (projected.suppress) { - return; - } - if (shouldHideHeartbeatChatOutput(clientRunId, sourceRunId)) { - return; - } const last = run.deltaSentAt ?? 0; if (now - last < 150) { return; } + const projected = chatRunState.resolveBuffer(clientRunId); + const mergedText = projected.text; + if (projected.suppress || shouldHideHeartbeatChatOutput(clientRunId, sourceRunId)) { + return; + } const broadcastDelta = resolveBroadcastDelta({ text: mergedText, previousBroadcastText: run.deltaLastBroadcastText, @@ -964,7 +956,7 @@ export function createAgentEventHandler({ sourceRunId: string, options?: { suppressLeadFragments?: boolean }, ) => { - const bufferedText = (chatRunState.runs.get(clientRunId)?.buffer ?? "").trim(); + const bufferedText = chatRunState.resolveBuffer(clientRunId).text.trim(); const normalizedHeartbeatText = normalizeHeartbeatChatFinalText({ runId: clientRunId, sourceRunId, diff --git a/src/gateway/server-methods/chat-abort-handler.ts b/src/gateway/server-methods/chat-abort-handler.ts index 75246705b284..1f35ff260586 100644 --- a/src/gateway/server-methods/chat-abort-handler.ts +++ b/src/gateway/server-methods/chat-abort-handler.ts @@ -294,7 +294,7 @@ export async function handleChatAbortRequestWithLifecycle( return; } - const partialText = context.chatRunState.runs.get(runId)?.buffer; + const partialText = context.chatRunState.resolveBuffer(runId).text; const res = abortChatRunById(ops, { runId, sessionKey: active.sessionKey, diff --git a/src/gateway/server-methods/chat-abort-runtime.ts b/src/gateway/server-methods/chat-abort-runtime.ts index 58e436c2e281..cf9521bdbfa2 100644 --- a/src/gateway/server-methods/chat-abort-runtime.ts +++ b/src/gateway/server-methods/chat-abort-runtime.ts @@ -48,7 +48,7 @@ function collectSessionAbortPartials(params: { if (!params.runIds.has(runId)) { continue; } - const text = params.chatRunState.runs.get(runId)?.buffer; + const text = params.chatRunState.resolveBuffer(runId).text; if (!text || !text.trim()) { continue; } diff --git a/src/shared/chat-message-content.test.ts b/src/shared/chat-message-content.test.ts index 479cecca9aac..35a28d3b4bd7 100644 --- a/src/shared/chat-message-content.test.ts +++ b/src/shared/chat-message-content.test.ts @@ -1,9 +1,10 @@ // Chat message content tests cover visible text extraction from message parts. -import { describe, expect, it } from "vitest"; +import { describe, expect, it, vi } from "vitest"; import { extractAssistantTextForPhase, extractAssistantVisibleText, extractFirstTextBlock, + parseAssistantTextSignature, resolveAssistantMessagePhase, } from "./chat-message-content.js"; @@ -181,6 +182,26 @@ describe("resolveAssistantMessagePhase", () => { ); }); + it("reuses a block signature parse until the live block signature changes", () => { + const block = { + type: "text", + text: "streaming text", + textSignature: JSON.stringify({ v: 1, id: "msg_1", phase: "commentary" }), + }; + const parseSpy = vi.spyOn(JSON, "parse"); + + expect(parseAssistantTextSignature(block)).toEqual({ id: "msg_1", phase: "commentary" }); + block.text += " delta"; + expect(parseAssistantTextSignature(block)).toEqual({ id: "msg_1", phase: "commentary" }); + expect(parseSpy).toHaveBeenCalledTimes(1); + + block.textSignature = JSON.stringify({ v: 1, id: "msg_2", phase: "final_answer" }); + expect(parseAssistantTextSignature(block)).toEqual({ id: "msg_2", phase: "final_answer" }); + expect(parseSpy).toHaveBeenCalledTimes(2); + + parseSpy.mockRestore(); + }); + it("resolves a single explicit phase from textSignature metadata", () => { expect( resolveAssistantMessagePhase({ diff --git a/src/shared/chat-message-content.ts b/src/shared/chat-message-content.ts index 1d610e8e8fb9..af4ed953dacf 100644 --- a/src/shared/chat-message-content.ts +++ b/src/shared/chat-message-content.ts @@ -23,6 +23,16 @@ export function extractFirstTextBlock(message: unknown): string | undefined { export type AssistantPhase = "commentary" | "final_answer"; +type AssistantTextSignature = { id?: string; phase?: AssistantPhase } | null; +type AssistantTextSignatureBlock = { textSignature?: unknown }; + +// Provider partials mutate blocks in place. Pair the stable block identity with +// its current signature text so a replacement can never reuse a stale parse. +const assistantTextSignatureCache = new WeakMap< + object, + { text: unknown; result: AssistantTextSignature } +>(); + function isAssistantTextContentBlockType(value: unknown): boolean { return value === "text" || value === "input_text" || value === "output_text"; } @@ -34,28 +44,36 @@ export function normalizeAssistantPhase(value: unknown): AssistantPhase | undefi /** Parses assistant text block signatures, preserving legacy raw ids when not JSON encoded. */ export function parseAssistantTextSignature( - value: unknown, -): { id?: string; phase?: AssistantPhase } | null { + block: AssistantTextSignatureBlock, +): AssistantTextSignature { + const value = block.textSignature; + const cached = assistantTextSignatureCache.get(block); + if (cached && cached.text === value) { + return cached.result; + } + let result: AssistantTextSignature; if (typeof value !== "string" || value.trim().length === 0) { - return null; - } - if (!value.startsWith("{")) { - return { id: value }; - } - try { - const parsed = JSON.parse(value) as { id?: unknown; phase?: unknown; v?: unknown }; - if (parsed.v !== 1) { - return null; + result = null; + } else if (!value.startsWith("{")) { + result = { id: value }; + } else { + try { + const parsed = JSON.parse(value) as { id?: unknown; phase?: unknown; v?: unknown }; + result = + parsed.v === 1 + ? { + ...(typeof parsed.id === "string" ? { id: parsed.id } : {}), + ...(normalizeAssistantPhase(parsed.phase) + ? { phase: normalizeAssistantPhase(parsed.phase) } + : {}), + } + : null; + } catch { + result = null; } - return { - ...(typeof parsed.id === "string" ? { id: parsed.id } : {}), - ...(normalizeAssistantPhase(parsed.phase) - ? { phase: normalizeAssistantPhase(parsed.phase) } - : {}), - }; - } catch { - return null; } + assistantTextSignatureCache.set(block, { text: value, result }); + return result; } /** Resolves a message phase only when the top-level phase or all explicit blocks agree. */ @@ -80,7 +98,7 @@ export function resolveAssistantMessagePhase(message: unknown): AssistantPhase | if (!isAssistantTextContentBlockType(record.type)) { continue; } - const phase = parseAssistantTextSignature(record.textSignature)?.phase; + const phase = parseAssistantTextSignature(record)?.phase; if (phase) { explicitPhases.add(phase); } @@ -163,7 +181,7 @@ export function extractAssistantTextForPhase( if (!isAssistantTextContentBlockType(record.type)) { return false; } - return Boolean(parseAssistantTextSignature(record.textSignature)?.phase); + return Boolean(parseAssistantTextSignature(record)?.phase); }); // Once explicit phased blocks exist, unphased extraction should not revive legacy text. @@ -180,7 +198,7 @@ export function extractAssistantTextForPhase( if (!isAssistantTextContentBlockType(record.type) || typeof record.text !== "string") { return null; } - const signature = parseAssistantTextSignature(record.textSignature); + const signature = parseAssistantTextSignature(record); const resolvedPhase = signature?.phase ?? (hasExplicitPhasedTextBlocks ? undefined : messagePhase); if (!shouldIncludeContent(resolvedPhase)) {