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..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 @@ -2060,3 +2060,214 @@ 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 }; +} + +// 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 } }), + ), + ]; +} + +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(input: { uuid: string; kind: "thinking" | "text"; text: string }) { + const { uuid, kind, text } = input; + 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(input: { + items: AgentTimelineItem[]; + type: "reasoning" | "assistant_message"; +}): string[] { + const { items, type } = input; + return items.flatMap((item) => (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([ + 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, type: "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([ + 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, type: "assistant_message" })).toEqual([ + "Checking ", + "the tests.", + " All green.", + ]); +}); + +test("fills in text a block snapshot has beyond what streamed", async () => { + const items = await runScriptedTurn([ + 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, type: "reasoning" })).toEqual(["Partial", " thought.", "Next."]); +}); + +test("keeps identical blocks that arrive only as snapshots", async () => { + const items = await runScriptedTurn([ + blockSnapshot({ uuid: "uuid-block-0", kind: "text", text: "Same line." }), + blockSnapshot({ uuid: "uuid-block-1", kind: "text", text: "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: "uuid-block-0", kind: "text", text: "Once." }), + blockSnapshot({ uuid: "uuid-block-0", kind: "text", text: "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([ + MESSAGE_START, + ...streamedBlock({ index: 0, kind: "text", deltas: ["A"] }), + streamFrame({ type: "content_block_delta", delta: { type: "text_delta", text: "B" } }), + blockStop(0), + blockSnapshot({ uuid: "uuid-block-0", kind: "text", text: "AB" }), + streamFrame({ type: "content_block_delta", delta: { type: "text_delta", text: "C" } }), + blockStop(1), + blockSnapshot({ uuid: "uuid-block-1", kind: "text", text: "C" }), + MESSAGE_STOP, + ]); + + 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([ + MESSAGE_START, + ...streamedBlock({ index: 0, kind: "thinking", deltas: [""] }), + streamFrame({ + type: "content_block_delta", + index: 0, + delta: { type: "signature_delta", signature: "sig" }, + }), + 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, 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 8ea12610ea..e529b3fd55 100644 --- a/packages/server/src/server/agent/providers/claude/agent.ts +++ b/packages/server/src/server/agent/providers/claude/agent.ts @@ -1146,12 +1146,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; } @@ -1191,8 +1205,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( @@ -1230,21 +1251,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 []; @@ -1258,49 +1304,72 @@ class TimelineAssembler { return []; } const state = this.ensureMessageState(messageId, runId); - return this.appendFragments(state, fragments); - } - - private appendFragments( - state: TimelineMessageState, - fragments: TimelineFragment[], - ): AgentTimelineItem[] { for (const fragment of fragments) { - if (fragment.kind === "assistant") { - state.assistantText += fragment.text; - } else { - state.reasoningText += fragment.text; + const key = this.resolveStreamBlockKey({ state, blockIndex, kind: fragment.kind }); + state.openStreamBlockKey = key; + const block = state.blocks.get(key) ?? this.addBlock({ state, key, kind: fragment.kind }); + block.text += fragment.text; + } + return this.emitNewContent(state); + } + + private resolveStreamBlockKey(input: { + state: TimelineMessageState; + blockIndex: unknown; + kind: TimelineFragment["kind"]; + }): string { + const { state, blockIndex, kind } = input; + 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(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. + 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, kind: 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( - 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); + 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; } private finalizeMessage(messageId: string, runId: string | null): AgentTimelineItem[] { @@ -1320,20 +1389,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; } @@ -1349,10 +1419,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);