From 8f83876a3cbd64b148078f0b7d2f03acb6ec7615 Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Thu, 2 Apr 2026 16:56:47 +0700 Subject: [PATCH] fix(stream): restore live timeline emission and preserve timeline events --- .../src/server/agent/agent-manager.test.ts | 43 +++++++--- .../server/src/server/agent/agent-manager.ts | 84 +------------------ 2 files changed, 34 insertions(+), 93 deletions(-) diff --git a/packages/server/src/server/agent/agent-manager.test.ts b/packages/server/src/server/agent/agent-manager.test.ts index ef82def56..ae2e81788 100644 --- a/packages/server/src/server/agent/agent-manager.test.ts +++ b/packages/server/src/server/agent/agent-manager.test.ts @@ -1390,7 +1390,7 @@ describe("AgentManager", () => { expect(fetched.rows.map((row) => row.seq)).toEqual([1, 2]); }); - test("buffers assistant chunks provisionally and streams one finalized assistant row", async () => { + test("streams assistant chunks incrementally and persists canonical chunk rows", async () => { const workdir = mkdtempSync(join(tmpdir(), "agent-manager-provisional-timeline-")); const storagePath = join(workdir, "agents"); const storage = new AgentStorage(storagePath, logger); @@ -1440,29 +1440,45 @@ describe("AgentManager", () => { } } - const assistantTimelineEvents = streamEvents.filter((event) => event.itemType === "assistant_message"); - expect(assistantTimelineEvents).toHaveLength(1); + const assistantTimelineEvents = streamEvents.filter( + (event) => event.itemType === "assistant_message", + ); + expect(assistantTimelineEvents).toHaveLength(2); expect(assistantTimelineEvents[0]).toMatchObject({ eventType: "timeline", itemType: "assistant_message", - text: "final reply", + text: "final ", seq: 1, }); + expect(assistantTimelineEvents[1]).toMatchObject({ + eventType: "timeline", + itemType: "assistant_message", + text: "reply", + seq: 2, + }); expect(manager.getTimeline(snapshot.id)).toEqual([ { type: "assistant_message", - text: "final reply", + text: "final ", + }, + { + type: "assistant_message", + text: "reply", }, ]); const fetched = await manager.fetchTimeline(snapshot.id, { direction: "tail", limit: 0, }); - expect(fetched.rows).toHaveLength(1); + expect(fetched.rows).toHaveLength(2); expect(fetched.rows[0]?.item).toEqual({ type: "assistant_message", - text: "final reply", + text: "final ", + }); + expect(fetched.rows[1]?.item).toEqual({ + type: "assistant_message", + text: "reply", }); }); @@ -1560,7 +1576,7 @@ describe("AgentManager", () => { expect(fetched.window.maxSeq).toBe(3); }); - test("hydrateTimeline canonicalizes tool-interleaved assistant replay into the committed turn shape", async () => { + test("hydrateTimeline preserves assistant chunk, reasoning, and tool timeline history", async () => { const workdir = mkdtempSync(join(tmpdir(), "agent-manager-history-canonical-assistant-")); const storagePath = join(workdir, "agents"); const storage = new AgentStorage(storagePath, logger); @@ -1645,6 +1661,9 @@ describe("AgentManager", () => { await manager.hydrateTimelineFromProvider(snapshot.id); expect(manager.getTimeline(snapshot.id)).toEqual([ + { type: "assistant_message", text: "chunk one " }, + { type: "assistant_message", text: "chunk two" }, + { type: "reasoning", text: "internal" }, { type: "tool_call", callId: "call-history-1", @@ -1658,11 +1677,11 @@ describe("AgentManager", () => { }, error: null, }, - { type: "assistant_message", text: "chunk one chunk twofinal answer" }, + { type: "assistant_message", text: "final answer" }, ]); }); - test("hydrateTimeline canonicalizes reasoning-interleaved assistant replay into one committed assistant row", async () => { + test("hydrateTimeline preserves reasoning between assistant chunks", async () => { const workdir = mkdtempSync(join(tmpdir(), "agent-manager-history-reasoning-interleave-")); const storagePath = join(workdir, "agents"); const storage = new AgentStorage(storagePath, logger); @@ -1727,8 +1746,10 @@ describe("AgentManager", () => { expect(manager.getTimeline(snapshot.id)).toEqual([ { type: "assistant_message", - text: "before reasoning after reasoning", + text: "before reasoning ", }, + { type: "reasoning", text: "internal step" }, + { type: "assistant_message", text: "after reasoning" }, ]); }); diff --git a/packages/server/src/server/agent/agent-manager.ts b/packages/server/src/server/agent/agent-manager.ts index d25f1cf0d..07a3ba0a4 100644 --- a/packages/server/src/server/agent/agent-manager.ts +++ b/packages/server/src/server/agent/agent-manager.ts @@ -151,7 +151,6 @@ type ManagedAgentBase = { currentModeId: string | null; pendingPermissions: Map; pendingReplacement: boolean; - provisionalAssistantText: string | null; persistence: AgentPersistenceHandle | null; historyPrimed: boolean; lastUserMessageAt: Date | null; @@ -657,7 +656,6 @@ export class AgentManager { await this.cancelAgentRun(agentId); existing = this.requireSessionAgent(agentId); } - const preservedProvisionalAssistantText = existing.provisionalAssistantText; const preservedHistoryPrimed = existing.historyPrimed; const preservedLastUsage = existing.lastUsage; const preservedLastError = existing.lastError; @@ -700,7 +698,6 @@ export class AgentManager { createdAt: existing.createdAt, updatedAt: existing.updatedAt, lastUserMessageAt: existing.lastUserMessageAt, - provisionalAssistantText: preservedProvisionalAssistantText, historyPrimed: preservedHistoryPrimed, lastUsage: preservedLastUsage, lastError: preservedLastError, @@ -1789,7 +1786,6 @@ export class AgentManager { timeline?: AgentTimelineItem[]; timelineRows?: AgentTimelineRow[]; timelineNextSeq?: number; - provisionalAssistantText?: string | null; historyPrimed?: boolean; lastUsage?: AgentUsage; lastError?: string; @@ -1848,7 +1844,6 @@ export class AgentManager { activeForegroundTurnId: null, foregroundTurnWaiters: new Set(), unsubscribeSession: null, - provisionalAssistantText: options?.provisionalAssistantText ?? null, persistence: attachPersistenceCwd(session.describePersistence(), config.cwd), historyPrimed: options?.historyPrimed ?? durableTimelineHasRows, lastUserMessageAt: options?.lastUserMessageAt ?? null, @@ -2137,36 +2132,13 @@ export class AgentManager { } agent.historyPrimed = true; const canonicalUserMessagesById = this.timelineStore.getCanonicalUserMessagesById(agent.id); - const pendingTurnItems: AgentTimelineItem[] = []; - let bufferedAssistantText = ""; - const flushPendingTurn = () => { - for (const item of pendingTurnItems) { - this.recordTimeline(agent.id, item); - } - pendingTurnItems.length = 0; - if (bufferedAssistantText) { - this.recordTimeline(agent.id, { - type: "assistant_message", - text: bufferedAssistantText, - }); - bufferedAssistantText = ""; - } - }; try { for await (const event of agent.session.streamHistory()) { if (event.type !== "timeline") { - if ( - event.type === "turn_completed" || - event.type === "turn_failed" || - event.type === "turn_canceled" - ) { - flushPendingTurn(); - } continue; } if (event.item.type === "user_message") { - flushPendingTurn(); const eventMessageId = normalizeMessageId(event.item.messageId); if (eventMessageId) { const canonicalText = canonicalUserMessagesById.get(eventMessageId); @@ -2174,26 +2146,10 @@ export class AgentManager { continue; } } - this.recordTimeline(agent.id, event.item); - continue; } - if (event.item.type === "assistant_message") { - bufferedAssistantText += event.item.text; - continue; - } - - if (event.item.type === "reasoning") { - continue; - } - - if (event.item.type === "tool_call" && event.item.status === "running") { - continue; - } - - pendingTurnItems.push(event.item); + this.recordTimeline(agent.id, event.item); } - flushPendingTurn(); } catch { // ignore history failures } @@ -2211,7 +2167,6 @@ export class AgentManager { const isForegroundEvent = Boolean( eventTurnId && agent.activeForegroundTurnId === eventTurnId, ); - let suppressLiveDispatch = false; // Only update timestamp for live events, not history replay if (!options?.fromHistory) { @@ -2265,17 +2220,6 @@ export class AgentManager { } } } - if (event.item.type === "assistant_message") { - agent.provisionalAssistantText = `${agent.provisionalAssistantText ?? ""}${event.item.text}`; - suppressLiveDispatch = true; - break; - } - if (event.item.type === "reasoning") { - break; - } - if (event.item.type === "tool_call" && event.item.status === "running") { - break; - } timelineRow = this.recordTimeline(agent.id, event.item); if (!options?.fromHistory && event.item.type === "user_message") { agent.lastUserMessageAt = new Date(); @@ -2292,27 +2236,6 @@ export class AgentManager { }, "handleStreamEvent: turn_completed", ); - if (agent.provisionalAssistantText) { - const item: AgentTimelineItem = { - type: "assistant_message", - text: agent.provisionalAssistantText, - }; - timelineRow = this.recordTimeline(agent.id, item); - if (!options?.fromHistory) { - this.dispatchStream( - agent.id, - { - type: "timeline", - item, - provider: event.provider, - }, - { - seq: timelineRow.seq, - }, - ); - } - agent.provisionalAssistantText = null; - } agent.lastUsage = event.usage; agent.lastError = undefined; // For autonomous turns (not foreground), transition to idle @@ -2336,7 +2259,6 @@ export class AgentManager { }, "handleStreamEvent: turn_failed", ); - agent.provisionalAssistantText = null; // For autonomous turns, set error state directly if (!isForegroundEvent) { agent.lifecycle = "error"; @@ -2373,7 +2295,6 @@ export class AgentManager { }, "handleStreamEvent: turn_canceled", ); - agent.provisionalAssistantText = null; // For autonomous turns, transition to idle // unless a replacement is pending (avoid idle flash during replace) if (!isForegroundEvent && !agent.pendingReplacement) { @@ -2405,7 +2326,6 @@ export class AgentManager { }, "handleStreamEvent: turn_started", ); - agent.provisionalAssistantText = null; // For autonomous turn_started (no foreground match), set running if (!isForegroundEvent) { (agent as ActiveManagedAgent).lifecycle = "running"; @@ -2435,7 +2355,7 @@ export class AgentManager { } // Skip dispatching individual stream events during history replay. - if (!options?.fromHistory && !suppressLiveDispatch) { + if (!options?.fromHistory) { this.dispatchStream( agent.id, event,