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;

packages/server/src/server/agent/agent-sdk-types.ts

@@ -304,7 +304,7 @@ export interface CompactionTimelineItem {
 
 export type AgentTimelineItem =
   | { type: "user_message"; text: string; messageId?: string }
-  | { type: "assistant_message"; text: string }
+  | { type: "assistant_message"; text: string; messageId?: string }
   | { type: "reasoning"; text: string }
   | ToolCallTimelineItem
   | { type: "todo"; items: { text: string; completed: boolean }[] }

packages/server/src/server/agent/providers/codex-app-server-agent.test.ts

@@ -1315,6 +1315,66 @@ describe("Codex app-server provider", () => {
         item: {
           type: "assistant_message",
           text: "History loaded.",
+          messageId: "message-history",
+        },
+      },
+    ]);
+  });
+
+  test("preserves Codex app-server assistant item ids in persisted history", async () => {
+    const session = createSession();
+    session.client = {
+      request: vi.fn(async (method: string) => {
+        if (method !== "thread/read") {
+          return {};
+        }
+        return {
+          thread: {
+            turns: [
+              {
+                items: [
+                  {
+                    type: "agentMessage",
+                    id: "before-tool-message",
+                    text: "I checked the workspace.",
+                  },
+                  {
+                    type: "agentMessage",
+                    id: "after-tool-message",
+                    text: "The tests are green.",
+                  },
+                ],
+              },
+            ],
+          },
+        };
+      }),
+    };
+
+    await asInternals(session).loadPersistedHistory();
+
+    const history: AgentStreamEvent[] = [];
+    for await (const event of session.streamHistory()) {
+      history.push(event);
+    }
+
+    expect(history).toEqual([
+      {
+        type: "timeline",
+        provider: "codex",
+        item: {
+          type: "assistant_message",
+          text: "I checked the workspace.",
+          messageId: "before-tool-message",
+        },
+      },
+      {
+        type: "timeline",
+        provider: "codex",
+        item: {
+          type: "assistant_message",
+          text: "The tests are green.",
+          messageId: "after-tool-message",
         },
       },
     ]);
@@ -1768,13 +1828,13 @@ describe("Codex app-server provider", () => {
         type: "timeline",
         provider: "codex",
         turnId: "test-turn",
-        item: { type: "assistant_message", text: "Hel" },
+        item: { type: "assistant_message", text: "Hel", messageId: "assistant-item-1" },
       },
       {
         type: "timeline",
         provider: "codex",
         turnId: "test-turn",
-        item: { type: "assistant_message", text: "lo" },
+        item: { type: "assistant_message", text: "lo", messageId: "assistant-item-1" },
       },
     ]);
   });
@@ -1805,19 +1865,19 @@ describe("Codex app-server provider", () => {
         type: "timeline",
         provider: "codex",
         turnId: "test-turn",
-        item: { type: "assistant_message", text: "Hel" },
+        item: { type: "assistant_message", text: "Hel", messageId: "assistant-item-2" },
       },
       {
         type: "timeline",
         provider: "codex",
         turnId: "test-turn",
-        item: { type: "assistant_message", text: "lo" },
+        item: { type: "assistant_message", text: "lo", messageId: "assistant-item-2" },
       },
       {
         type: "timeline",
         provider: "codex",
         turnId: "test-turn",
-        item: { type: "assistant_message", text: "!" },
+        item: { type: "assistant_message", text: "!", messageId: "assistant-item-2" },
       },
     ]);
   });
@@ -1852,6 +1912,7 @@ describe("Codex app-server provider", () => {
         turnId: "test-turn",
         item: {
           type: "assistant_message",
+          messageId: "assistant-item-3",
           text: "I’m in the waiting phase now. The next read is intentionally delayed so we get meaningful CI state instead of churn.",
         },
       },
@@ -1861,6 +1922,7 @@ describe("Codex app-server provider", () => {
         turnId: "test-turn",
         item: {
           type: "assistant_message",
+          messageId: "assistant-item-4",
           text: "\n\n---\n\nCI is still cooking. I’m staying on the current run rather than jumping around, because the first red job will tell us exactly whether anything else needs work.",
         },
       },

packages/server/src/server/agent/providers/codex-app-server-agent.ts

@@ -1534,11 +1534,14 @@ function threadItemToTimeline(
   switch (normalizedType) {
     case "userMessage":
       return mapCodexThreadUserMessageItem(normalizedItem, includeUserMessage);
-    case "agentMessage":
+    case "agentMessage": {
+      const messageId = nonEmptyString(normalizedItem.id);
       return {
         type: "assistant_message",
         text: typeof normalizedItem.text === "string" ? normalizedItem.text : "",
+        ...(messageId ? { messageId } : {}),
       };
+    }
     case "plan":
       return mapCodexThreadPlanItem(normalizedItem);
     case "reasoning":
@@ -4073,6 +4076,7 @@ class CodexAppServerAgentSession implements AgentSession {
       if (subAgentCallId) {
         this.upsertSubAgentChildItem(subAgentCallId, parsed.itemId, {
           type: "assistant_message",
+          messageId: parsed.itemId,
           text,
         });
         this.emitSubAgentActivityUpdate(subAgentCallId, "running");
@@ -4084,6 +4088,7 @@ class CodexAppServerAgentSession implements AgentSession {
         provider: CODEX_PROVIDER,
         item: {
           type: "assistant_message",
+          messageId: parsed.itemId,
           text:
             isFirstDeltaForItem && this.pendingAssistantMessageBoundary
               ? `${ASSISTANT_MESSAGE_BOUNDARY_MARKDOWN}${parsed.delta}`
@@ -4443,7 +4448,14 @@ class CodexAppServerAgentSession implements AgentSession {
     this.emitEvent({
       type: "timeline",
       provider: CODEX_PROVIDER,
-      item: { type: timelineItem.type, text: suffix },
+      item:
+        timelineItem.type === "assistant_message"
+          ? {
+              type: timelineItem.type,
+              text: suffix,
+              ...(timelineItem.messageId ? { messageId: timelineItem.messageId } : {}),
+            }
+          : { type: timelineItem.type, text: suffix },
     });
   }
 

packages/server/src/server/agent/timeline-projection.test.ts

@@ -39,6 +39,59 @@ describe("projectTimelineRows", () => {
     expect(projected[0]?.collapsed).toContain("assistant_merge");
   });
 
+  test("merges adjacent assistant chunks with the same message id in projected mode", () => {
+    const rows: AgentTimelineRow[] = [
+      {
+        seq: 1,
+        timestamp: "2026-02-13T00:00:00.000Z",
+        item: { type: "assistant_message", text: "Hel", messageId: "msg-1" },
+      },
+      {
+        seq: 2,
+        timestamp: "2026-02-13T00:00:00.100Z",
+        item: { type: "assistant_message", text: "lo", messageId: "msg-1" },
+      },
+    ];
+
+    const projected = projectTimelineRows({ rows, mode: "projected" });
+
+    expect(projected).toHaveLength(1);
+    expect(projected[0]?.item).toEqual({
+      type: "assistant_message",
+      text: "Hello",
+      messageId: "msg-1",
+    });
+  });
+
+  test("keeps adjacent assistant chunks with different message ids separate in projected mode", () => {
+    const rows: AgentTimelineRow[] = [
+      {
+        seq: 1,
+        timestamp: "2026-02-13T00:00:00.000Z",
+        item: { type: "assistant_message", text: "First answer.", messageId: "msg-1" },
+      },
+      {
+        seq: 2,
+        timestamp: "2026-02-13T00:00:00.100Z",
+        item: { type: "assistant_message", text: "Second answer.", messageId: "msg-2" },
+      },
+    ];
+
+    const projected = projectTimelineRows({ rows, mode: "projected" });
+
+    expect(projected).toHaveLength(2);
+    expect(projected[0]?.item).toEqual({
+      type: "assistant_message",
+      text: "First answer.",
+      messageId: "msg-1",
+    });
+    expect(projected[1]?.item).toEqual({
+      type: "assistant_message",
+      text: "Second answer.",
+      messageId: "msg-2",
+    });
+  });
+
   test("merges adjacent reasoning chunks in projected mode", () => {
     const rows: AgentTimelineRow[] = [
       {

packages/server/src/server/agent/timeline-projection.ts

@@ -210,6 +210,13 @@ function mergeAssistantChunks(entries: readonly WorkingEntry[]): WorkingEntry[]
       { type: "assistant_message" }
     >;
     const entryAssistant = entry.item as Extract<AgentTimelineItem, { type: "assistant_message" }>;
+    if (
+      entryAssistant.messageId !== undefined &&
+      previousAssistant.messageId !== entryAssistant.messageId
+    ) {
+      output.push(entry);
+      continue;
+    }
 
     const collapsedKinds = new Set<TimelineProjectionKind>([
       ...previous.collapsed,
@@ -222,6 +229,7 @@ function mergeAssistantChunks(entries: readonly WorkingEntry[]): WorkingEntry[]
       item: {
         type: "assistant_message",
         text: `${previousAssistant.text}${entryAssistant.text}`,
+        ...(previousAssistant.messageId ? { messageId: previousAssistant.messageId } : {}),
       },
       seqEnd: entry.seqEnd,
       sourceSeqRanges: mergeSeqRanges(previous.sourceSeqRanges, entry.sourceSeqRanges),

packages/server/src/shared/messages.ts

@@ -505,6 +505,7 @@ export const AgentTimelineItemPayloadSchema: z.ZodType<AgentTimelineItem, z.ZodT
     z.object({
       type: z.literal("assistant_message"),
       text: z.string(),
+      messageId: z.string().optional(),
     }),
     z.object({
       type: z.literal("reasoning"),

packages/server/src/shared/wire-compat.test.ts

@@ -194,6 +194,29 @@ async function emitTimelineResponse(
 }
 
 describe("wire compatibility", () => {
+  test("assistant timeline message ids are optional on the wire", () => {
+    expect(
+      AgentTimelineItemPayloadSchema.parse({
+        type: "assistant_message",
+        text: "old daemon shape",
+      }),
+    ).toEqual({
+      type: "assistant_message",
+      text: "old daemon shape",
+    });
+    expect(
+      AgentTimelineItemPayloadSchema.parse({
+        type: "assistant_message",
+        text: "new daemon shape",
+        messageId: "msg-1",
+      }),
+    ).toEqual({
+      type: "assistant_message",
+      text: "new daemon shape",
+      messageId: "msg-1",
+    });
+  });
+
   test("downgrades reasoning_merge for clients that do not declare the capability", async () => {
     const response = await emitTimelineResponse();