Preserve Codex assistant message ids (#989)

packages/app/src/types/stream-event.test.ts

@@ -5,12 +5,13 @@ import { applyStreamEvent } from "@/types/stream";
 
 const baseTimestamp = new Date(0);
 
-const assistantChunk = (text: string): AgentStreamEventPayload => ({
+const assistantChunk = (text: string, messageId?: string): AgentStreamEventPayload => ({
   type: "timeline",
   provider: "codex",
   item: {
     type: "assistant_message",
     text,
+    ...(messageId ? { messageId } : {}),
   },
 });
 
@@ -131,6 +132,36 @@ describe("applyStreamEvent", () => {
     expect(result.tail[0].kind).toBe("assistant_message");
   });
 
+  it("does not continue a tail assistant message when the incoming message id differs", () => {
+    const result = applyStreamEvent({
+      tail: [
+        {
+          kind: "assistant_message",
+          id: "msg-first",
+          messageId: "msg-first",
+          text: "First answer.",
+          timestamp: baseTimestamp,
+        },
+      ],
+      head: [],
+      event: assistantChunk("Second answer.", "msg-second"),
+      timestamp: baseTimestamp,
+    });
+
+    expect(result.tail).toHaveLength(1);
+    expect(result.head).toHaveLength(1);
+    expect(result.tail[0].kind).toBe("assistant_message");
+    expect(result.head[0].kind).toBe("assistant_message");
+    if (
+      result.tail[0].kind === "assistant_message" &&
+      result.head[0].kind === "assistant_message"
+    ) {
+      expect(result.tail[0].text).toBe("First answer.");
+      expect(result.head[0].text).toBe("Second answer.");
+      expect(result.head[0].messageId).toBe("msg-second");
+    }
+  });
+
   it("flushes reasoning when assistant message starts", () => {
     let result = applyStreamEvent({
       tail: [],

packages/app/src/types/stream.test.ts

@@ -18,11 +18,12 @@ type CanonicalToolStatus = "running" | "completed" | "failed" | "canceled";
 function assistantTimeline(
   text: string,
   provider: AgentProvider = "claude",
+  messageId?: string,
 ): AgentStreamEventPayload {
   return {
     type: "timeline",
     provider,
-    item: { type: "assistant_message", text },
+    item: { type: "assistant_message", text, ...(messageId ? { messageId } : {}) },
   };
 }
 
@@ -312,6 +313,65 @@ describe("stream reducer canonical tool calls", () => {
     assert.strictEqual(assistantMessage?.text, "Hello world");
   });
 
+  it("keeps adjacent assistant timeline items separate when message ids differ", () => {
+    const state = hydrateStreamState([
+      {
+        event: assistantTimeline("First answer.", "codex", "msg-first"),
+        timestamp: new Date("2025-01-01T10:01:00Z"),
+      },
+      {
+        event: assistantTimeline("Second answer.", "codex", "msg-second"),
+        timestamp: new Date("2025-01-01T10:01:01Z"),
+      },
+    ]);
+
+    assert.deepStrictEqual(
+      state.map((item) => (item.kind === "assistant_message" ? item.text : item.kind)),
+      ["First answer.", "Second answer."],
+    );
+    assert.deepStrictEqual(
+      state.map((item) => (item.kind === "assistant_message" ? item.messageId : null)),
+      ["msg-first", "msg-second"],
+    );
+  });
+
+  it("merges adjacent assistant deltas when message ids match", () => {
+    const state = hydrateStreamState([
+      {
+        event: assistantTimeline("Hel", "codex", "msg-same"),
+        timestamp: new Date("2025-01-01T10:02:00Z"),
+      },
+      {
+        event: assistantTimeline("lo", "codex", "msg-same"),
+        timestamp: new Date("2025-01-01T10:02:01Z"),
+      },
+    ]);
+
+    assert.strictEqual(state.length, 1);
+    assert.strictEqual(state[0]?.kind, "assistant_message");
+    if (state[0]?.kind === "assistant_message") {
+      assert.strictEqual(state[0].text, "Hello");
+      assert.strictEqual(state[0].id, "msg-same");
+      assert.strictEqual(state[0].messageId, "msg-same");
+    }
+  });
+
+  it("preserves old assistant merge behavior when message ids are absent", () => {
+    const state = hydrateStreamState([
+      {
+        event: assistantTimeline("Hel", "codex"),
+        timestamp: new Date("2025-01-01T10:03:00Z"),
+      },
+      {
+        event: assistantTimeline("lo", "codex"),
+        timestamp: new Date("2025-01-01T10:03:01Z"),
+      },
+    ]);
+
+    assert.strictEqual(state.length, 1);
+    assert.strictEqual(state[0]?.kind === "assistant_message" ? state[0].text : null, "Hello");
+  });
+
   it("merges running and completed events by callId", () => {
     const callId = "tool-merge-1";
     const updates = [

packages/app/src/types/stream.ts

@@ -66,6 +66,7 @@ export interface UserMessageItem {
 export interface AssistantMessageItem {
   kind: "assistant_message";
   id: string;
+  messageId?: string;
   text: string;
   timestamp: Date;
   blockGroupId?: string;
@@ -227,6 +228,7 @@ function appendAssistantMessage(
   text: string,
   timestamp: Date,
   source: StreamUpdateSource,
+  messageId?: string,
 ): StreamItem[] {
   const { chunk, hasContent } = normalizeChunk(text);
   if (!chunk) {
@@ -234,7 +236,11 @@ function appendAssistantMessage(
   }
 
   const last = state[state.length - 1];
-  if (last && last.kind === "assistant_message") {
+  const shouldAppendToLast =
+    last &&
+    last.kind === "assistant_message" &&
+    (messageId === undefined || last.messageId === messageId);
+  if (shouldAppendToLast) {
     const updated: AssistantMessageItem = {
       ...last,
       text: `${last.text}${chunk}`,
@@ -249,7 +255,8 @@ function appendAssistantMessage(
   if (
     source === "live" &&
     last?.kind === "user_message" &&
-    secondLast?.kind === "assistant_message"
+    secondLast?.kind === "assistant_message" &&
+    (messageId === undefined || secondLast.messageId === messageId)
   ) {
     const updated: AssistantMessageItem = {
       ...secondLast,
@@ -264,9 +271,11 @@ function appendAssistantMessage(
   }
 
   const idSeed = chunk.trim() || chunk;
+  const entryId = messageId ?? createUniqueTimelineId(state, "assistant", idSeed, timestamp);
   const item: AssistantMessageItem = {
     kind: "assistant_message",
-    id: createUniqueTimelineId(state, "assistant", idSeed, timestamp),
+    id: entryId,
+    ...(messageId ? { messageId } : {}),
     text: chunk,
     timestamp,
   };
@@ -647,7 +656,9 @@ function reduceTimelineEvent(
     case "user_message":
       return finalizeActiveThoughts(appendUserMessage(state, item.text, timestamp, item.messageId));
     case "assistant_message":
-      return finalizeActiveThoughts(appendAssistantMessage(state, item.text, timestamp, source));
+      return finalizeActiveThoughts(
+        appendAssistantMessage(state, item.text, timestamp, source, item.messageId),
+      );
     case "reasoning":
       return appendThought(state, item.text, timestamp);
     case "tool_call":
@@ -993,7 +1004,14 @@ export function applyStreamEvent(params: {
 
   if (incomingKind === "assistant_message" && nextHead.length === 0) {
     const tailAssistant = nextTail.at(-1);
-    if (tailAssistant?.kind === "assistant_message") {
+    const incomingMessageId =
+      event.type === "timeline" && event.item.type === "assistant_message"
+        ? event.item.messageId
+        : undefined;
+    const shouldContinueTailAssistant =
+      tailAssistant?.kind === "assistant_message" &&
+      (incomingMessageId === undefined || tailAssistant.messageId === incomingMessageId);
+    if (shouldContinueTailAssistant) {
       nextTail = nextTail.slice(0, -1);
       nextHead = [tailAssistant];
       changedTail = true;