Merge remote-tracking branch 'origin/pr-6466' into fix/6199-live-narration-764ea71e

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<string, unknown>) {
+  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<string, unknown>[]): Promise<AgentTimelineItem[]> {
+  sdkQueryFactory.mockImplementation(({ prompt }: { prompt: AsyncIterable<unknown> }) => {
+    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."]);
+});

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<string, TimelineBlockState>;
+  // 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<string>;
+  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);