diff --git a/packages/app/src/types/stream-event.test.ts b/packages/app/src/types/stream-event.test.ts index 2185b9339..adf171da8 100644 --- a/packages/app/src/types/stream-event.test.ts +++ b/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: [], diff --git a/packages/app/src/types/stream.test.ts b/packages/app/src/types/stream.test.ts index e9d3bcf91..395c743f4 100644 --- a/packages/app/src/types/stream.test.ts +++ b/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 = [ diff --git a/packages/app/src/types/stream.ts b/packages/app/src/types/stream.ts index f1f9ce83a..a6829eeb1 100644 --- a/packages/app/src/types/stream.ts +++ b/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; diff --git a/packages/server/src/server/agent/agent-sdk-types.ts b/packages/server/src/server/agent/agent-sdk-types.ts index 5639bb8b4..5f3caad67 100644 --- a/packages/server/src/server/agent/agent-sdk-types.ts +++ b/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 }[] } diff --git a/packages/server/src/server/agent/providers/codex-app-server-agent.test.ts b/packages/server/src/server/agent/providers/codex-app-server-agent.test.ts index 1ea8a4fe5..7729df953 100644 --- a/packages/server/src/server/agent/providers/codex-app-server-agent.test.ts +++ b/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.", }, }, diff --git a/packages/server/src/server/agent/providers/codex-app-server-agent.ts b/packages/server/src/server/agent/providers/codex-app-server-agent.ts index 7aee4a885..2fe97fa86 100644 --- a/packages/server/src/server/agent/providers/codex-app-server-agent.ts +++ b/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 }, }); } diff --git a/packages/server/src/server/agent/timeline-projection.test.ts b/packages/server/src/server/agent/timeline-projection.test.ts index aa98a7c16..c7f53abc6 100644 --- a/packages/server/src/server/agent/timeline-projection.test.ts +++ b/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[] = [ { diff --git a/packages/server/src/server/agent/timeline-projection.ts b/packages/server/src/server/agent/timeline-projection.ts index 3d8f55ca0..62477ac5b 100644 --- a/packages/server/src/server/agent/timeline-projection.ts +++ b/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; + if ( + entryAssistant.messageId !== undefined && + previousAssistant.messageId !== entryAssistant.messageId + ) { + output.push(entry); + continue; + } const collapsedKinds = new Set([ ...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), diff --git a/packages/server/src/shared/messages.ts b/packages/server/src/shared/messages.ts index da7893eee..309c2b6c7 100644 --- a/packages/server/src/shared/messages.ts +++ b/packages/server/src/shared/messages.ts @@ -505,6 +505,7 @@ export const AgentTimelineItemPayloadSchema: z.ZodType { + 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();