From 896cd76d674fedaca3b1412726c5c7a6c449bdb5 Mon Sep 17 00:00:00 2001 From: paseo-bot Date: Wed, 23 Sep 2026 12:58:13 +0200 Subject: [PATCH] fix(server): keep a streamed assistant message whole across a steer Steering an agent mid-stream records the user message at the next timeline seq, between two chunks of the assistant message that is still streaming. The projection merged assistant chunks only when they were adjacent, so the interposed user row split one provider message into two timeline rows at whatever character had arrived, often mid-word. An assistant message that carries a provider message id now has a timeline identity, like a tool call does, so its chunks collapse into one item wherever they land. Chunks without a message id keep the existing adjacency merge. --- .../protocol/src/timeline-identity.test.ts | 7 +++ packages/protocol/src/timeline-identity.ts | 3 + .../src/server/agent/agent-manager.test.ts | 61 +++++++++++++++++++ .../server/agent/agent-timeline-store.test.ts | 21 +++++++ .../server/agent/timeline-projection.test.ts | 27 ++++++++ .../src/server/agent/timeline-projection.ts | 11 ++++ 6 files changed, 130 insertions(+) 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 {