diff --git a/packages/protocol/src/timeline-identity.test.ts b/packages/protocol/src/timeline-identity.test.ts index 19f77fd7bf..96994689cb 100644 --- a/packages/protocol/src/timeline-identity.test.ts +++ b/packages/protocol/src/timeline-identity.test.ts @@ -25,4 +25,11 @@ describe("timelineItemIdentity", () => { ).toBe("review/row-1"); expect(timelineItemIdentity({ type: "reasoning", text: "hi" })).toBeNull(); }); + + it("identifies an assistant message by its provider message id", () => { + expect( + timelineItemIdentity({ type: "assistant_message", text: "hi", messageId: "msg-1" }), + ).toBe("message:msg-1"); + expect(timelineItemIdentity({ type: "assistant_message", text: "hi" })).toBeNull(); + }); }); diff --git a/packages/protocol/src/timeline-identity.ts b/packages/protocol/src/timeline-identity.ts index e41737f60b..8b593b2afa 100644 --- a/packages/protocol/src/timeline-identity.ts +++ b/packages/protocol/src/timeline-identity.ts @@ -3,5 +3,8 @@ import type { AgentTimelineItem } from "./agent-types.js"; export function timelineItemIdentity(item: AgentTimelineItem): string | null { if (item.type === "tool_call") return item.callId; if (item.type === "plugin") return `${item.pluginId}/${item.id}`; + // Streamed chunks of one provider assistant message share its id, so they + // stay one item even when another row is recorded between two chunks. + if (item.type === "assistant_message" && item.messageId) return `message:${item.messageId}`; return null; } diff --git a/packages/server/src/server/agent/agent-manager.test.ts b/packages/server/src/server/agent/agent-manager.test.ts index c107b3da0b..89ae20a769 100644 --- a/packages/server/src/server/agent/agent-manager.test.ts +++ b/packages/server/src/server/agent/agent-manager.test.ts @@ -897,6 +897,67 @@ test("orders buffered pre-steer output before an immediately accepted steer", as } }); +test("keeps one assistant message whole when a steer lands between its chunks", async () => { + const session = new SteeringTestSession({ provider: "codex", cwd: process.cwd() }); + const workdir = mkdtempSync(join(tmpdir(), "agent-manager-steer-split-")); + const manager = new AgentManager({ + clients: { + codex: new (class extends TestAgentClient { + override async createSession(): Promise { + return session; + } + })(), + }, + agentStreamCoalesceWindowMs: 1, + logger, + }); + let agentId: string | null = null; + try { + const agent = await manager.createAgent({ provider: "codex", cwd: workdir }, undefined, { + workspaceId: undefined, + }); + agentId = agent.id; + const run = manager.streamAgent(agent.id, "initial"); + void (async () => { + for await (const _event of run) { + } + })(); + await manager.waitForAgentRunStart(agent.id); + + session.pushEvent({ + type: "timeline", + provider: "codex", + turnId: "active-turn-1", + item: { type: "assistant_message", text: "with the uncommit", messageId: "msg-1" }, + }); + await new Promise((resolve) => setTimeout(resolve, 20)); + await expect( + manager.steerAgentRun(agent.id, "keep going", { clientMessageId: "keep-going-client" }), + ).resolves.toEqual({ status: "accepted" }); + session.pushEvent({ + type: "timeline", + provider: "codex", + turnId: "active-turn-1", + item: { type: "assistant_message", text: "ted changes.", messageId: "msg-1" }, + }); + await new Promise((resolve) => setTimeout(resolve, 20)); + + const rows = manager + .fetchTimeline(agent.id, { limit: 0 }) + .rows.filter( + (row) => row.item.type === "assistant_message" || row.item.type === "user_message", + ) + .map((row) => row.item); + expect(rows).toMatchObject([ + { type: "assistant_message", text: "with the uncommitted changes.", messageId: "msg-1" }, + { type: "user_message", text: "keep going" }, + ]); + } finally { + if (agentId) await manager.closeAgent(agentId).catch(() => undefined); + rmSync(workdir, { recursive: true, force: true }); + } +}); + test("orders a concurrent replacement after a pending accepted steer", async () => { const entered = deferred(); const release = deferred(); diff --git a/packages/server/src/server/agent/agent-timeline-store.test.ts b/packages/server/src/server/agent/agent-timeline-store.test.ts index 895aec2684..cf4a368c9a 100644 --- a/packages/server/src/server/agent/agent-timeline-store.test.ts +++ b/packages/server/src/server/agent/agent-timeline-store.test.ts @@ -2,6 +2,27 @@ import { describe, expect, it } from "vitest"; import { InMemoryAgentTimelineStore } from "./agent-timeline-store.js"; describe("InMemoryAgentTimelineStore", () => { + it("keeps a streaming assistant message whole when a steering prompt is recorded mid-stream", () => { + const store = new InMemoryAgentTimelineStore(); + store.initialize("agent-1"); + store.append("agent-1", { + type: "assistant_message", + text: "with the uncommit", + messageId: "msg-1", + }); + store.append("agent-1", { type: "user_message", text: "keep going" }); + store.append("agent-1", { + type: "assistant_message", + text: "ted changes.", + messageId: "msg-1", + }); + + expect(store.getItems("agent-1")).toEqual([ + { type: "assistant_message", text: "with the uncommitted changes.", messageId: "msg-1" }, + { type: "user_message", text: "keep going" }, + ]); + }); + it("clamps an overshooting before cursor into the bounded tail window", () => { const store = new InMemoryAgentTimelineStore(); store.initialize("agent-1", { diff --git a/packages/server/src/server/agent/timeline-projection.test.ts b/packages/server/src/server/agent/timeline-projection.test.ts index 2d8cf9d4a1..4487d3c7b1 100644 --- a/packages/server/src/server/agent/timeline-projection.test.ts +++ b/packages/server/src/server/agent/timeline-projection.test.ts @@ -64,6 +64,33 @@ describe("projectTimelineRows", () => { }); }); + test("keeps one assistant message whole when a steering user message lands mid-stream", () => { + const rows: AgentTimelineRow[] = [ + { + seq: 1, + timestamp: "2026-02-13T00:00:00.000Z", + item: { type: "assistant_message", text: "with the uncommit", messageId: "msg-1" }, + }, + { + seq: 2, + timestamp: "2026-02-13T00:00:00.100Z", + item: { type: "user_message", text: "keep going" }, + }, + { + seq: 3, + timestamp: "2026-02-13T00:00:00.200Z", + item: { type: "assistant_message", text: "ted changes.", messageId: "msg-1" }, + }, + ]; + + const projected = projectTimelineRows({ rows, mode: "projected" }); + + expect(projected.map((entry) => entry.item)).toEqual([ + { type: "assistant_message", text: "with the uncommitted changes.", messageId: "msg-1" }, + { type: "user_message", text: "keep going" }, + ]); + }); + test("keeps adjacent assistant chunks with different message ids separate 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 d8128a6a2b..17c1760480 100644 --- a/packages/server/src/server/agent/timeline-projection.ts +++ b/packages/server/src/server/agent/timeline-projection.ts @@ -155,6 +155,17 @@ function mergeIdentityEntries(existing: WorkingEntry, entry: WorkingEntry): Work seqEnd: Math.max(existing.seqEnd, entry.seqEnd), ...mergeIdentityMetadata(existing, entry, "tool_lifecycle"), }; + case "assistant_message": + if (existing.item.type !== "assistant_message" || existing.turnId !== entry.turnId) { + return null; + } + return { + ...existing, + item: { ...existing.item, text: `${existing.item.text}${entry.item.text}` }, + timestamp: entry.timestamp, + seqEnd: Math.max(existing.seqEnd, entry.seqEnd), + ...mergeIdentityMetadata(existing, entry, "assistant_merge"), + }; case "plugin": if (existing.item.type !== "plugin") return null; return {