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.
This commit is contained in:
paseo-bot
2026-09-23 12:58:13 +02:00
parent d161b5987d
commit 896cd76d67
6 changed files with 130 additions and 0 deletions
@@ -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();
});
});
@@ -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;
}
@@ -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<AgentSession> {
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<void>();
const release = deferred<void>();
@@ -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", {
@@ -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[] = [
{
@@ -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 {