From 75b7cab2703e99352a88b06ea5fd8fe652b18d54 Mon Sep 17 00:00:00 2001 From: ThePharmer Date: Fri, 9 Oct 2026 14:17:44 -0400 Subject: [PATCH 1/2] fix(claude): show each thinking and text block once in multi-block messages While streaming, the Claude Agent SDK emits one assistant message per completed non-empty content block, each carrying only that block and sharing the message id. The timeline assembler treated each of those as a snapshot of the whole message, so when a message had two or more thinking (or text) blocks, the last block's snapshot did not start with the accumulated text, the emitted cursor reset, and the block was sent again. The Thinking sheet then showed the last paragraph twice. Track text per content block, keyed by the stream block index, and match each snapshot to the earliest block of its kind that has not had one. A snapshot only appends text the deltas missed. Snapshot-only blocks become their own blocks, and a redelivered snapshot uuid is ignored. --- .../providers/claude/agent.redesign.test.ts | 190 ++++++++++++++++++ .../server/agent/providers/claude/agent.ts | 189 +++++++++++------ 2 files changed, 318 insertions(+), 61 deletions(-) diff --git a/packages/server/src/server/agent/providers/claude/agent.redesign.test.ts b/packages/server/src/server/agent/providers/claude/agent.redesign.test.ts index 5393926520..ca49d34ed3 100644 --- a/packages/server/src/server/agent/providers/claude/agent.redesign.test.ts +++ b/packages/server/src/server/agent/providers/claude/agent.redesign.test.ts @@ -2060,3 +2060,193 @@ test("does not use stream_event uuid as assistant message identity when message_ await session.close(); }); + +const MULTI_BLOCK_SESSION_ID = "multi-block-session"; + +function streamFrame(event: Record) { + return { type: "stream_event", parent_tool_use_id: null, event }; +} + +function streamedBlock(index: number, kind: "thinking" | "text", deltas: string[]) { + const deltaType = kind === "thinking" ? "thinking_delta" : "text_delta"; + return [ + streamFrame({ type: "content_block_start", index, content_block: { type: kind, [kind]: "" } }), + ...deltas.map((text) => + streamFrame({ type: "content_block_delta", index, delta: { type: deltaType, [kind]: text } }), + ), + streamFrame({ type: "content_block_stop", index }), + ]; +} + +// The SDK's per-block assistant message: same message id as its siblings, one block of content. +function blockSnapshot(uuid: string, kind: "thinking" | "text", text: string) { + return { + type: "assistant", + uuid, + parent_tool_use_id: null, + session_id: MULTI_BLOCK_SESSION_ID, + message: { id: "msg-multi", role: "assistant", content: [{ type: kind, [kind]: text }] }, + }; +} + +async function runScriptedTurn(frames: Record[]): Promise { + sdkQueryFactory.mockImplementation(({ prompt }: { prompt: AsyncIterable }) => { + const readPromptUuid = createPromptUuidReader(prompt); + let step = 0; + return createBaseQueryMock( + vi.fn(async () => { + const index = step++; + if (index === 0) { + return { + done: false, + value: { ...claudeTurnInit(), session_id: MULTI_BLOCK_SESSION_ID }, + }; + } + if (index === 1) { + return { + done: false, + value: { + type: "user", + message: { role: "user", content: "multi-block prompt" }, + parent_tool_use_id: null, + uuid: (await readPromptUuid()) ?? "missing-prompt-uuid", + session_id: MULTI_BLOCK_SESSION_ID, + }, + }; + } + const frame = frames[index - 2]; + if (frame) { + return { done: false, value: frame }; + } + if (index === frames.length + 2) { + return { done: false, value: successResult() }; + } + return { done: true, value: undefined }; + }), + ); + }); + + const session = await createSession(); + const events = await collectUntilTerminal(streamSession(session, "multi-block prompt")); + await session.close(); + return events.flatMap((event) => (event.type === "timeline" ? [event.item] : [])); +} + +// Each chunk as emitted, so a test sees duplicates and lost chunks, not just the final text. +function emittedChunks( + items: AgentTimelineItem[], + type: "reasoning" | "assistant_message", +): string[] { + return items.flatMap((item) => (item.type === type ? [item.text] : [])); +} + +test("shows each thinking block of a multi-block message once", async () => { + const items = await runScriptedTurn([ + streamFrame({ type: "message_start", message: { id: "msg-multi", role: "assistant" } }), + ...streamedBlock(0, "thinking", ["First paragraph.\n\n", "Second paragraph.\n\n"]), + blockSnapshot("uuid-block-0", "thinking", "First paragraph.\n\nSecond paragraph.\n\n"), + ...streamedBlock(1, "thinking", ["Last paragraph.\n\n"]), + blockSnapshot("uuid-block-1", "thinking", "Last paragraph.\n\n"), + streamFrame({ type: "message_stop" }), + ]); + + expect(emittedChunks(items, "reasoning")).toEqual([ + "First paragraph.\n\n", + "Second paragraph.\n\n", + "Last paragraph.\n\n", + ]); +}); + +test("shows each text block of a multi-block message once", async () => { + const items = await runScriptedTurn([ + streamFrame({ type: "message_start", message: { id: "msg-multi", role: "assistant" } }), + ...streamedBlock(0, "text", ["Checking ", "the tests."]), + blockSnapshot("uuid-block-0", "text", "Checking the tests."), + ...streamedBlock(1, "text", [" All green."]), + blockSnapshot("uuid-block-1", "text", " All green."), + streamFrame({ type: "message_stop" }), + ]); + + expect(emittedChunks(items, "assistant_message")).toEqual([ + "Checking ", + "the tests.", + " All green.", + ]); +}); + +test("fills in text a block snapshot has beyond what streamed", async () => { + const items = await runScriptedTurn([ + streamFrame({ type: "message_start", message: { id: "msg-multi", role: "assistant" } }), + ...streamedBlock(0, "thinking", ["Partial"]), + blockSnapshot("uuid-block-0", "thinking", "Partial thought."), + ...streamedBlock(1, "thinking", ["Next."]), + blockSnapshot("uuid-block-1", "thinking", "Next."), + streamFrame({ type: "message_stop" }), + ]); + + expect(emittedChunks(items, "reasoning")).toEqual(["Partial", " thought.", "Next."]); +}); + +test("keeps identical blocks that arrive only as snapshots", async () => { + const items = await runScriptedTurn([ + blockSnapshot("uuid-block-0", "text", "Same line."), + blockSnapshot("uuid-block-1", "text", "Same line."), + ]); + + expect(emittedChunks(items, "assistant_message")).toEqual(["Same line.", "Same line."]); +}); + +test("ignores a redelivered block snapshot", async () => { + const items = await runScriptedTurn([ + blockSnapshot("uuid-block-0", "text", "Once."), + blockSnapshot("uuid-block-0", "text", "Once."), + ]); + + expect(emittedChunks(items, "assistant_message")).toEqual(["Once."]); +}); + +test("continues the open block with deltas that carry no index", async () => { + const items = await runScriptedTurn([ + streamFrame({ type: "message_start", message: { id: "msg-multi", role: "assistant" } }), + streamFrame({ + type: "content_block_start", + index: 0, + content_block: { type: "text", text: "" }, + }), + streamFrame({ + type: "content_block_delta", + index: 0, + delta: { type: "text_delta", text: "A" }, + }), + streamFrame({ type: "content_block_delta", delta: { type: "text_delta", text: "B" } }), + streamFrame({ type: "content_block_stop", index: 0 }), + blockSnapshot("uuid-block-0", "text", "AB"), + streamFrame({ type: "content_block_delta", delta: { type: "text_delta", text: "C" } }), + streamFrame({ type: "content_block_stop", index: 1 }), + blockSnapshot("uuid-block-1", "text", "C"), + streamFrame({ type: "message_stop" }), + ]); + + expect(emittedChunks(items, "assistant_message")).toEqual(["A", "B", "C"]); +}); + +// With thinking display "omitted", a thinking block streams an empty delta and a signature, and +// the SDK sends no snapshot for it because it only snapshots non-empty blocks. +test("does not hold a snapshot slot for an empty thinking block", async () => { + const items = await runScriptedTurn([ + streamFrame({ type: "message_start", message: { id: "msg-multi", role: "assistant" } }), + ...streamedBlock(0, "thinking", [""]).slice(0, -1), + streamFrame({ + type: "content_block_delta", + index: 0, + delta: { type: "signature_delta", signature: "sig" }, + }), + streamFrame({ type: "content_block_stop", index: 0 }), + ...streamedBlock(1, "thinking", ["Visible."]).slice(0, -1), + blockSnapshot("uuid-block-1", "thinking", "Visible."), + streamFrame({ type: "content_block_stop", index: 1 }), + streamFrame({ type: "message_stop" }), + ]); + + expect(emittedChunks(items, "reasoning")).toEqual(["Visible."]); +}); diff --git a/packages/server/src/server/agent/providers/claude/agent.ts b/packages/server/src/server/agent/providers/claude/agent.ts index b797ce2c94..cd374c8d83 100644 --- a/packages/server/src/server/agent/providers/claude/agent.ts +++ b/packages/server/src/server/agent/providers/claude/agent.ts @@ -1144,12 +1144,26 @@ interface TimelineFragment { text: string; } +// While streaming, the SDK emits one assistant message per completed non-empty content block, +// carrying only that block and sent before the block's content_block_stop (see +// SDKAssistantMessage). Snapshots therefore reconcile against their own block, never against the +// whole message. Empty blocks get no snapshot, so they are never registered as blocks. +interface TimelineBlockState { + kind: TimelineFragment["kind"]; + text: string; + emittedLength: number; + snapshotApplied: boolean; +} + interface TimelineMessageState { id: string; - assistantText: string; - reasoningText: string; - emittedAssistantLength: number; - emittedReasoningLength: number; + // Insertion order is content order. Streamed blocks are keyed by stream block index. + blocks: Map; + // The block that deltas without an index continue, until content_block_stop or its snapshot. + // The Messages API always sends an index; index-less deltas only come from older fixtures. + openStreamBlockKey: string | null; + appliedSnapshotUuids: Set; + syntheticBlockCount: number; stopped: boolean; } @@ -1189,8 +1203,15 @@ class TimelineAssembler { return []; } const state = this.ensureMessageState(messageId, runId); + const snapshotUuid = readTrimmedString(message.uuid); + if (snapshotUuid) { + if (state.appliedSnapshotUuids.has(snapshotUuid)) { + return []; + } + state.appliedSnapshotUuids.add(snapshotUuid); + } const fragments = this.extractFragments(message.message?.content); - return this.applyAbsoluteFragments(state, fragments); + return this.applyBlockSnapshots(state, fragments); } private consumeStreamEvent( @@ -1228,21 +1249,46 @@ class TimelineAssembler { } if (eventType === "content_block_start") { - return this.consumeDeltaContent(event.content_block, runId, streamEventMessageId); + return this.consumeDeltaContent({ + content: event.content_block, + blockIndex: event.index, + runId, + messageIdHint: streamEventMessageId, + }); } if (eventType === "content_block_delta") { - return this.consumeDeltaContent(event.delta, runId, streamEventMessageId); + return this.consumeDeltaContent({ + content: event.delta, + blockIndex: event.index, + runId, + messageIdHint: streamEventMessageId, + }); + } + + if (eventType === "content_block_stop") { + const messageId = this.resolveMessageId({ + runId, + createIfMissing: false, + messageId: streamEventMessageId, + }); + const state = messageId ? this.messages.get(messageId) : undefined; + if (state) { + state.openStreamBlockKey = null; + } + return []; } return []; } - private consumeDeltaContent( - content: unknown, - runId: string | null, - messageIdHint: string | null, - ): AgentTimelineItem[] { + private consumeDeltaContent(input: { + content: unknown; + blockIndex: unknown; + runId: string | null; + messageIdHint: string | null; + }): AgentTimelineItem[] { + const { content, blockIndex, runId, messageIdHint } = input; const fragments = this.extractFragments(content); if (fragments.length === 0) { return []; @@ -1256,49 +1302,69 @@ class TimelineAssembler { return []; } const state = this.ensureMessageState(messageId, runId); - return this.appendFragments(state, fragments); + for (const fragment of fragments) { + const key = this.resolveStreamBlockKey(state, blockIndex, fragment.kind); + state.openStreamBlockKey = key; + const block = state.blocks.get(key) ?? this.addBlock(state, key, fragment.kind); + block.text += fragment.text; + } + return this.emitNewContent(state); } - private appendFragments( + private resolveStreamBlockKey( + state: TimelineMessageState, + blockIndex: unknown, + kind: TimelineFragment["kind"], + ): string { + if (typeof blockIndex === "number") { + return `stream:${blockIndex}`; + } + const openKey = state.openStreamBlockKey; + if (openKey && state.blocks.get(openKey)?.kind === kind) { + return openKey; + } + return `stream:unindexed:${++state.syntheticBlockCount}`; + } + + private applyBlockSnapshots( state: TimelineMessageState, fragments: TimelineFragment[], ): AgentTimelineItem[] { for (const fragment of fragments) { - if (fragment.kind === "assistant") { - state.assistantText += fragment.text; - } else { - state.reasoningText += fragment.text; + // Snapshots arrive in content order, so each belongs to the earliest block of its kind + // that has not had one yet. + const match = [...state.blocks.entries()].find( + ([, candidate]) => candidate.kind === fragment.kind && !candidate.snapshotApplied, + ); + if (!match) { + const key = `snapshot:${++state.syntheticBlockCount}`; + const snapshotOnly = this.addBlock(state, key, fragment.kind); + snapshotOnly.text = fragment.text; + snapshotOnly.snapshotApplied = true; + continue; + } + const [key, block] = match; + block.snapshotApplied = true; + // A snapshot completes its block, so later deltas without an index start a new one. + if (state.openStreamBlockKey === key) { + state.openStreamBlockKey = null; + } + // Emitted text cannot be retracted, so a snapshot that rewrites streamed text is ignored. + if (fragment.text.startsWith(block.text)) { + block.text = fragment.text; } } return this.emitNewContent(state); } - private applyAbsoluteFragments( + private addBlock( state: TimelineMessageState, - fragments: TimelineFragment[], - ): AgentTimelineItem[] { - const assistantText = fragments - .filter((fragment) => fragment.kind === "assistant") - .map((fragment) => fragment.text) - .join(""); - const reasoningText = fragments - .filter((fragment) => fragment.kind === "reasoning") - .map((fragment) => fragment.text) - .join(""); - - if (assistantText.length > 0) { - if (!assistantText.startsWith(state.assistantText)) { - state.emittedAssistantLength = 0; - } - state.assistantText = assistantText; - } - if (reasoningText.length > 0) { - if (!reasoningText.startsWith(state.reasoningText)) { - state.emittedReasoningLength = 0; - } - state.reasoningText = reasoningText; - } - return this.emitNewContent(state); + key: string, + kind: TimelineFragment["kind"], + ): TimelineBlockState { + const block: TimelineBlockState = { kind, text: "", emittedLength: 0, snapshotApplied: false }; + state.blocks.set(key, block); + return block; } private finalizeMessage(messageId: string, runId: string | null): AgentTimelineItem[] { @@ -1318,20 +1384,21 @@ class TimelineAssembler { private emitNewContent(state: TimelineMessageState): AgentTimelineItem[] { const items: AgentTimelineItem[] = []; - const nextAssistantText = state.assistantText.slice(state.emittedAssistantLength); - if ( - nextAssistantText.length > 0 && - nextAssistantText !== INTERRUPT_TOOL_USE_PLACEHOLDER && - !isClaudeTranscriptNoiseText(nextAssistantText) - ) { - state.emittedAssistantLength = state.assistantText.length; - items.push({ type: "assistant_message", text: nextAssistantText, messageId: state.id }); - } - - const nextReasoningText = state.reasoningText.slice(state.emittedReasoningLength); - if (nextReasoningText.length > 0) { - state.emittedReasoningLength = state.reasoningText.length; - items.push({ type: "reasoning", text: nextReasoningText }); + for (const block of state.blocks.values()) { + const nextText = block.text.slice(block.emittedLength); + if (nextText.length === 0) { + continue; + } + if (block.kind === "reasoning") { + block.emittedLength = block.text.length; + items.push({ type: "reasoning", text: nextText }); + continue; + } + if (nextText === INTERRUPT_TOOL_USE_PLACEHOLDER || isClaudeTranscriptNoiseText(nextText)) { + continue; + } + block.emittedLength = block.text.length; + items.push({ type: "assistant_message", text: nextText, messageId: state.id }); } return items; } @@ -1347,10 +1414,10 @@ class TimelineAssembler { } const created: TimelineMessageState = { id: messageId, - assistantText: "", - reasoningText: "", - emittedAssistantLength: 0, - emittedReasoningLength: 0, + blocks: new Map(), + appliedSnapshotUuids: new Set(), + openStreamBlockKey: null, + syntheticBlockCount: 0, stopped: false, }; this.messages.set(messageId, created); From ff5b7b8022e857f74f96208942edecff7a605be4 Mon Sep 17 00:00:00 2001 From: ThePharmer Date: Fri, 9 Oct 2026 15:27:41 -0400 Subject: [PATCH 2/2] refactor(claude): object params for new assembler helpers; snapshot before stop in tests Address review on #6466: - resolveStreamBlockKey, applyBlockSnapshots, addBlock and the new test helpers take a single object parameter, per packages/server/CLAUDE.md. - The multi-block tests send each block snapshot before its content_block_stop, the order the Agent SDK documents. The index-less test keeps the snapshot-after-stop order as separate coverage. --- .../providers/claude/agent.redesign.test.ts | 139 ++++++++++-------- .../server/agent/providers/claude/agent.ts | 39 ++--- 2 files changed, 101 insertions(+), 77 deletions(-) diff --git a/packages/server/src/server/agent/providers/claude/agent.redesign.test.ts b/packages/server/src/server/agent/providers/claude/agent.redesign.test.ts index ca49d34ed3..3966de92da 100644 --- a/packages/server/src/server/agent/providers/claude/agent.redesign.test.ts +++ b/packages/server/src/server/agent/providers/claude/agent.redesign.test.ts @@ -2067,19 +2067,26 @@ function streamFrame(event: Record) { return { type: "stream_event", parent_tool_use_id: null, event }; } -function streamedBlock(index: number, kind: "thinking" | "text", deltas: string[]) { +// A block's start and deltas. Its stop is separate because the SDK sends the block's snapshot +// before content_block_stop. +function streamedBlock(input: { index: number; kind: "thinking" | "text"; deltas: string[] }) { + const { index, kind, deltas } = input; const deltaType = kind === "thinking" ? "thinking_delta" : "text_delta"; return [ streamFrame({ type: "content_block_start", index, content_block: { type: kind, [kind]: "" } }), ...deltas.map((text) => streamFrame({ type: "content_block_delta", index, delta: { type: deltaType, [kind]: text } }), ), - streamFrame({ type: "content_block_stop", index }), ]; } +function blockStop(index: number) { + return streamFrame({ type: "content_block_stop", index }); +} + // The SDK's per-block assistant message: same message id as its siblings, one block of content. -function blockSnapshot(uuid: string, kind: "thinking" | "text", text: string) { +function blockSnapshot(input: { uuid: string; kind: "thinking" | "text"; text: string }) { + const { uuid, kind, text } = input; return { type: "assistant", uuid, @@ -2133,24 +2140,41 @@ async function runScriptedTurn(frames: Record[]): Promise (item.type === type ? [item.text] : [])); } +const MESSAGE_START = streamFrame({ + type: "message_start", + message: { id: "msg-multi", role: "assistant" }, +}); +const MESSAGE_STOP = streamFrame({ type: "message_stop" }); + test("shows each thinking block of a multi-block message once", async () => { const items = await runScriptedTurn([ - streamFrame({ type: "message_start", message: { id: "msg-multi", role: "assistant" } }), - ...streamedBlock(0, "thinking", ["First paragraph.\n\n", "Second paragraph.\n\n"]), - blockSnapshot("uuid-block-0", "thinking", "First paragraph.\n\nSecond paragraph.\n\n"), - ...streamedBlock(1, "thinking", ["Last paragraph.\n\n"]), - blockSnapshot("uuid-block-1", "thinking", "Last paragraph.\n\n"), - streamFrame({ type: "message_stop" }), + MESSAGE_START, + ...streamedBlock({ + index: 0, + kind: "thinking", + deltas: ["First paragraph.\n\n", "Second paragraph.\n\n"], + }), + blockSnapshot({ + uuid: "uuid-block-0", + kind: "thinking", + text: "First paragraph.\n\nSecond paragraph.\n\n", + }), + blockStop(0), + ...streamedBlock({ index: 1, kind: "thinking", deltas: ["Last paragraph.\n\n"] }), + blockSnapshot({ uuid: "uuid-block-1", kind: "thinking", text: "Last paragraph.\n\n" }), + blockStop(1), + MESSAGE_STOP, ]); - expect(emittedChunks(items, "reasoning")).toEqual([ + expect(emittedChunks({ items, type: "reasoning" })).toEqual([ "First paragraph.\n\n", "Second paragraph.\n\n", "Last paragraph.\n\n", @@ -2159,15 +2183,17 @@ test("shows each thinking block of a multi-block message once", async () => { test("shows each text block of a multi-block message once", async () => { const items = await runScriptedTurn([ - streamFrame({ type: "message_start", message: { id: "msg-multi", role: "assistant" } }), - ...streamedBlock(0, "text", ["Checking ", "the tests."]), - blockSnapshot("uuid-block-0", "text", "Checking the tests."), - ...streamedBlock(1, "text", [" All green."]), - blockSnapshot("uuid-block-1", "text", " All green."), - streamFrame({ type: "message_stop" }), + MESSAGE_START, + ...streamedBlock({ index: 0, kind: "text", deltas: ["Checking ", "the tests."] }), + blockSnapshot({ uuid: "uuid-block-0", kind: "text", text: "Checking the tests." }), + blockStop(0), + ...streamedBlock({ index: 1, kind: "text", deltas: [" All green."] }), + blockSnapshot({ uuid: "uuid-block-1", kind: "text", text: " All green." }), + blockStop(1), + MESSAGE_STOP, ]); - expect(emittedChunks(items, "assistant_message")).toEqual([ + expect(emittedChunks({ items, type: "assistant_message" })).toEqual([ "Checking ", "the tests.", " All green.", @@ -2176,77 +2202,72 @@ test("shows each text block of a multi-block message once", async () => { test("fills in text a block snapshot has beyond what streamed", async () => { const items = await runScriptedTurn([ - streamFrame({ type: "message_start", message: { id: "msg-multi", role: "assistant" } }), - ...streamedBlock(0, "thinking", ["Partial"]), - blockSnapshot("uuid-block-0", "thinking", "Partial thought."), - ...streamedBlock(1, "thinking", ["Next."]), - blockSnapshot("uuid-block-1", "thinking", "Next."), - streamFrame({ type: "message_stop" }), + MESSAGE_START, + ...streamedBlock({ index: 0, kind: "thinking", deltas: ["Partial"] }), + blockSnapshot({ uuid: "uuid-block-0", kind: "thinking", text: "Partial thought." }), + blockStop(0), + ...streamedBlock({ index: 1, kind: "thinking", deltas: ["Next."] }), + blockSnapshot({ uuid: "uuid-block-1", kind: "thinking", text: "Next." }), + blockStop(1), + MESSAGE_STOP, ]); - expect(emittedChunks(items, "reasoning")).toEqual(["Partial", " thought.", "Next."]); + expect(emittedChunks({ items, type: "reasoning" })).toEqual(["Partial", " thought.", "Next."]); }); test("keeps identical blocks that arrive only as snapshots", async () => { const items = await runScriptedTurn([ - blockSnapshot("uuid-block-0", "text", "Same line."), - blockSnapshot("uuid-block-1", "text", "Same line."), + blockSnapshot({ uuid: "uuid-block-0", kind: "text", text: "Same line." }), + blockSnapshot({ uuid: "uuid-block-1", kind: "text", text: "Same line." }), ]); - expect(emittedChunks(items, "assistant_message")).toEqual(["Same line.", "Same line."]); + expect(emittedChunks({ items, type: "assistant_message" })).toEqual(["Same line.", "Same line."]); }); test("ignores a redelivered block snapshot", async () => { const items = await runScriptedTurn([ - blockSnapshot("uuid-block-0", "text", "Once."), - blockSnapshot("uuid-block-0", "text", "Once."), + blockSnapshot({ uuid: "uuid-block-0", kind: "text", text: "Once." }), + blockSnapshot({ uuid: "uuid-block-0", kind: "text", text: "Once." }), ]); - expect(emittedChunks(items, "assistant_message")).toEqual(["Once."]); + expect(emittedChunks({ items, type: "assistant_message" })).toEqual(["Once."]); }); +// Older fixtures send deltas without an index. This also covers a snapshot arriving after its +// block's stop, the reverse of the documented order. test("continues the open block with deltas that carry no index", async () => { const items = await runScriptedTurn([ - streamFrame({ type: "message_start", message: { id: "msg-multi", role: "assistant" } }), - streamFrame({ - type: "content_block_start", - index: 0, - content_block: { type: "text", text: "" }, - }), - streamFrame({ - type: "content_block_delta", - index: 0, - delta: { type: "text_delta", text: "A" }, - }), + MESSAGE_START, + ...streamedBlock({ index: 0, kind: "text", deltas: ["A"] }), streamFrame({ type: "content_block_delta", delta: { type: "text_delta", text: "B" } }), - streamFrame({ type: "content_block_stop", index: 0 }), - blockSnapshot("uuid-block-0", "text", "AB"), + blockStop(0), + blockSnapshot({ uuid: "uuid-block-0", kind: "text", text: "AB" }), streamFrame({ type: "content_block_delta", delta: { type: "text_delta", text: "C" } }), - streamFrame({ type: "content_block_stop", index: 1 }), - blockSnapshot("uuid-block-1", "text", "C"), - streamFrame({ type: "message_stop" }), + blockStop(1), + blockSnapshot({ uuid: "uuid-block-1", kind: "text", text: "C" }), + MESSAGE_STOP, ]); - expect(emittedChunks(items, "assistant_message")).toEqual(["A", "B", "C"]); + expect(emittedChunks({ items, type: "assistant_message" })).toEqual(["A", "B", "C"]); }); // With thinking display "omitted", a thinking block streams an empty delta and a signature, and // the SDK sends no snapshot for it because it only snapshots non-empty blocks. test("does not hold a snapshot slot for an empty thinking block", async () => { const items = await runScriptedTurn([ - streamFrame({ type: "message_start", message: { id: "msg-multi", role: "assistant" } }), - ...streamedBlock(0, "thinking", [""]).slice(0, -1), + MESSAGE_START, + ...streamedBlock({ index: 0, kind: "thinking", deltas: [""] }), streamFrame({ type: "content_block_delta", index: 0, delta: { type: "signature_delta", signature: "sig" }, }), - streamFrame({ type: "content_block_stop", index: 0 }), - ...streamedBlock(1, "thinking", ["Visible."]).slice(0, -1), - blockSnapshot("uuid-block-1", "thinking", "Visible."), - streamFrame({ type: "content_block_stop", index: 1 }), - streamFrame({ type: "message_stop" }), + blockStop(0), + ...streamedBlock({ index: 1, kind: "thinking", deltas: ["Visible."] }), + blockSnapshot({ uuid: "uuid-block-1", kind: "thinking", text: "Visible." }), + blockStop(1), + MESSAGE_STOP, ]); - expect(emittedChunks(items, "reasoning")).toEqual(["Visible."]); + expect(emittedChunks({ items, type: "reasoning" })).toEqual(["Visible."]); }); diff --git a/packages/server/src/server/agent/providers/claude/agent.ts b/packages/server/src/server/agent/providers/claude/agent.ts index cd374c8d83..3c10eaadc3 100644 --- a/packages/server/src/server/agent/providers/claude/agent.ts +++ b/packages/server/src/server/agent/providers/claude/agent.ts @@ -1211,7 +1211,7 @@ class TimelineAssembler { state.appliedSnapshotUuids.add(snapshotUuid); } const fragments = this.extractFragments(message.message?.content); - return this.applyBlockSnapshots(state, fragments); + return this.applyBlockSnapshots({ state, fragments }); } private consumeStreamEvent( @@ -1303,19 +1303,20 @@ class TimelineAssembler { } const state = this.ensureMessageState(messageId, runId); for (const fragment of fragments) { - const key = this.resolveStreamBlockKey(state, blockIndex, fragment.kind); + const key = this.resolveStreamBlockKey({ state, blockIndex, kind: fragment.kind }); state.openStreamBlockKey = key; - const block = state.blocks.get(key) ?? this.addBlock(state, key, fragment.kind); + const block = state.blocks.get(key) ?? this.addBlock({ state, key, kind: fragment.kind }); block.text += fragment.text; } return this.emitNewContent(state); } - private resolveStreamBlockKey( - state: TimelineMessageState, - blockIndex: unknown, - kind: TimelineFragment["kind"], - ): string { + private resolveStreamBlockKey(input: { + state: TimelineMessageState; + blockIndex: unknown; + kind: TimelineFragment["kind"]; + }): string { + const { state, blockIndex, kind } = input; if (typeof blockIndex === "number") { return `stream:${blockIndex}`; } @@ -1326,10 +1327,11 @@ class TimelineAssembler { return `stream:unindexed:${++state.syntheticBlockCount}`; } - private applyBlockSnapshots( - state: TimelineMessageState, - fragments: TimelineFragment[], - ): AgentTimelineItem[] { + private applyBlockSnapshots(input: { + state: TimelineMessageState; + fragments: TimelineFragment[]; + }): AgentTimelineItem[] { + const { state, fragments } = input; for (const fragment of fragments) { // Snapshots arrive in content order, so each belongs to the earliest block of its kind // that has not had one yet. @@ -1338,7 +1340,7 @@ class TimelineAssembler { ); if (!match) { const key = `snapshot:${++state.syntheticBlockCount}`; - const snapshotOnly = this.addBlock(state, key, fragment.kind); + const snapshotOnly = this.addBlock({ state, key, kind: fragment.kind }); snapshotOnly.text = fragment.text; snapshotOnly.snapshotApplied = true; continue; @@ -1357,11 +1359,12 @@ class TimelineAssembler { return this.emitNewContent(state); } - private addBlock( - state: TimelineMessageState, - key: string, - kind: TimelineFragment["kind"], - ): TimelineBlockState { + private addBlock(input: { + state: TimelineMessageState; + key: string; + kind: TimelineFragment["kind"]; + }): TimelineBlockState { + const { state, key, kind } = input; const block: TimelineBlockState = { kind, text: "", emittedLength: 0, snapshotApplied: false }; state.blocks.set(key, block); return block;