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;