diff --git a/packages/app/src/types/stream.test.ts b/packages/app/src/types/stream.test.ts index 3b0ee8730..5695a6454 100644 --- a/packages/app/src/types/stream.test.ts +++ b/packages/app/src/types/stream.test.ts @@ -3,6 +3,7 @@ import { describe, it } from "vitest"; import { hydrateStreamState, + mergeToolCallDetail, reduceStreamUpdate, type AgentToolCallItem, type StreamItem, @@ -10,7 +11,7 @@ import { } from "./stream"; import type { AgentProvider, ToolCallDetail } from "@server/server/agent/agent-sdk-types"; import type { AgentStreamEventPayload } from "@server/shared/messages"; -import { buildToolCallDisplayModel } from "@/utils/tool-call-display"; +import { buildToolCallDisplayModel } from "../../../server/src/shared/tool-call-display"; type CanonicalToolStatus = "running" | "completed" | "failed" | "canceled"; @@ -99,6 +100,193 @@ function findToolByCallId(state: StreamItem[], callId: string): AgentToolCallIte ); } +describe("stream reducer tool call idempotency", () => { + it("returns the same detail reference when tool call detail is identical", () => { + const existing: ToolCallDetail = { + type: "shell", + command: "npm test", + cwd: "/tmp/repo", + }; + const incoming: ToolCallDetail = { + type: "shell", + command: "npm test", + cwd: "/tmp/repo", + }; + + const merged = mergeToolCallDetail(existing, incoming); + + assert.strictEqual(merged, existing); + }); + + it("returns a new detail reference when tool call detail changes", () => { + const existing: ToolCallDetail = { + type: "shell", + command: "npm test", + cwd: "/tmp/repo", + }; + const incoming: ToolCallDetail = { + type: "shell", + command: "npm run typecheck", + cwd: "/tmp/repo", + }; + + const merged = mergeToolCallDetail(existing, incoming); + + assert.notStrictEqual(merged, existing); + assert.deepStrictEqual(merged, incoming); + }); + + it("returns the same state array when status, error, detail, and metadata are identical", () => { + const callId = "idempotent-tool-call"; + const initialState = reduceStreamUpdate( + [], + canonicalToolTimeline({ + provider: "codex", + callId, + name: "shell", + status: "running", + detail: { + type: "shell", + command: "npm test", + cwd: "/tmp/repo", + }, + metadata: { + paneId: "%1", + }, + }), + new Date("2025-01-01T12:00:00Z"), + ); + + const nextState = reduceStreamUpdate( + initialState, + canonicalToolTimeline({ + provider: "codex", + callId, + name: "shell", + status: "running", + detail: { + type: "shell", + command: "npm test", + cwd: "/tmp/repo", + }, + metadata: { + paneId: "%1", + }, + }), + new Date("2025-01-01T12:00:01Z"), + ); + + assert.strictEqual(nextState, initialState); + }); + + it("returns a new state array when tool call status changes", () => { + const callId = "status-change-tool-call"; + const initialState = reduceStreamUpdate( + [], + canonicalToolTimeline({ + provider: "codex", + callId, + name: "shell", + status: "running", + detail: { + type: "shell", + command: "npm test", + }, + }), + new Date("2025-01-01T12:10:00Z"), + ); + + const nextState = reduceStreamUpdate( + initialState, + canonicalToolTimeline({ + provider: "codex", + callId, + name: "shell", + status: "completed", + detail: { + type: "shell", + command: "npm test", + }, + }), + new Date("2025-01-01T12:10:01Z"), + ); + + assert.notStrictEqual(nextState, initialState); + }); + + it("returns a new state array when tool call detail changes", () => { + const callId = "detail-change-tool-call"; + const initialState = reduceStreamUpdate( + [], + canonicalToolTimeline({ + provider: "codex", + callId, + name: "shell", + status: "running", + detail: { + type: "shell", + command: "npm test", + }, + }), + new Date("2025-01-01T12:20:00Z"), + ); + + const nextState = reduceStreamUpdate( + initialState, + canonicalToolTimeline({ + provider: "codex", + callId, + name: "shell", + status: "running", + detail: { + type: "shell", + command: "npm run typecheck", + }, + }), + new Date("2025-01-01T12:20:01Z"), + ); + + assert.notStrictEqual(nextState, initialState); + }); + + it("returns a new state array when tool call error changes", () => { + const callId = "error-change-tool-call"; + const initialState = reduceStreamUpdate( + [], + canonicalToolTimeline({ + provider: "codex", + callId, + name: "shell", + status: "failed", + error: { message: "first failure" }, + detail: { + type: "shell", + command: "npm test", + }, + }), + new Date("2025-01-01T12:30:00Z"), + ); + + const nextState = reduceStreamUpdate( + initialState, + canonicalToolTimeline({ + provider: "codex", + callId, + name: "shell", + status: "failed", + error: { message: "second failure" }, + detail: { + type: "shell", + command: "npm test", + }, + }), + new Date("2025-01-01T12:30:01Z"), + ); + + assert.notStrictEqual(nextState, initialState); + }); +}); + describe("stream reducer canonical tool calls", () => { it("is deterministic for equivalent hydration sequences", () => { const updates = [ diff --git a/packages/app/src/types/stream.ts b/packages/app/src/types/stream.ts index 86f979870..7d74e6998 100644 --- a/packages/app/src/types/stream.ts +++ b/packages/app/src/types/stream.ts @@ -336,7 +336,36 @@ function mergeUnknownValue(existing: unknown | null, incoming: unknown | null): return incoming; } -function mergeToolCallDetail(existing: ToolCallDetail, incoming: ToolCallDetail): ToolCallDetail { +function hasSameIncomingFields>( + existing: T, + incoming: T, +): boolean { + return Object.entries(incoming).every(([key, value]) => existing[key] === value); +} + +function mergeToolCallMetadata( + existing: Record | undefined, + incoming: Record | undefined, +): Record | undefined { + if (!incoming) { + return existing; + } + + if (!existing) { + return incoming; + } + + if (hasSameIncomingFields(existing, incoming)) { + return existing; + } + + return { ...existing, ...incoming }; +} + +export function mergeToolCallDetail( + existing: ToolCallDetail, + incoming: ToolCallDetail, +): ToolCallDetail { if (existing.type === "unknown" && incoming.type !== "unknown") { return incoming; } @@ -346,14 +375,24 @@ function mergeToolCallDetail(existing: ToolCallDetail, incoming: ToolCallDetail) } if (existing.type === "unknown" && incoming.type === "unknown") { + const input = mergeUnknownValue(existing.input, incoming.input); + const output = mergeUnknownValue(existing.output, incoming.output); + if (input === existing.input && output === existing.output) { + return existing; + } + return { type: "unknown", - input: mergeUnknownValue(existing.input, incoming.input), - output: mergeUnknownValue(existing.output, incoming.output), + input, + output, }; } if (existing.type === incoming.type) { + if (hasSameIncomingFields(existing, incoming)) { + return existing; + } + return { ...existing, ...incoming } as ToolCallDetail; } @@ -391,8 +430,7 @@ function appendAgentToolCall( const existingIndex = findExistingAgentToolCallIndex(state, data.callId); if (existingIndex >= 0) { - const next = [...state]; - const existing = next[existingIndex]; + const existing = state[existingIndex]; if (!existing || !isAgentToolCallItem(existing)) { return state; } @@ -401,12 +439,22 @@ function appendAgentToolCall( mergedStatus === "failed" ? (data.error ?? existing.payload.data.error ?? { message: "Tool call failed" }) : null; - const mergedMetadata = - data.metadata || existing.payload.data.metadata - ? { ...existing.payload.data.metadata, ...data.metadata } - : undefined; + const mergedMetadata = mergeToolCallMetadata(existing.payload.data.metadata, data.metadata); const mergedDetail = mergeToolCallDetail(existing.payload.data.detail, data.detail); + if ( + data.provider === existing.payload.data.provider && + data.callId === existing.payload.data.callId && + data.name === existing.payload.data.name && + mergedStatus === existing.payload.data.status && + mergedError === existing.payload.data.error && + mergedDetail === existing.payload.data.detail && + mergedMetadata === existing.payload.data.metadata + ) { + return state; + } + + const next = [...state]; next[existingIndex] = { ...existing, timestamp, diff --git a/packages/server/src/server/agent/agent-manager-stream-coalescing.test.ts b/packages/server/src/server/agent/agent-manager-stream-coalescing.test.ts index dc340f32e..b5aeed002 100644 --- a/packages/server/src/server/agent/agent-manager-stream-coalescing.test.ts +++ b/packages/server/src/server/agent/agent-manager-stream-coalescing.test.ts @@ -57,6 +57,28 @@ const TOOL_CALL: AgentTimelineItem = { }, }; +function toolCall(options?: { + callId?: string; + status?: "running" | "completed" | "failed" | "canceled"; + output?: string; + error?: unknown; +}): Extract { + const status = options?.status ?? "running"; + return { + type: "tool_call", + callId: options?.callId ?? "tool-1", + name: "shell", + status, + error: status === "failed" ? (options?.error ?? "failed") : null, + detail: { + type: "shell", + command: "printf ok", + output: options?.output ?? "", + exitCode: status === "completed" ? 0 : null, + }, + }; +} + class TestAgentSession implements AgentSession { readonly capabilities = TEST_CAPABILITIES; readonly id: string; @@ -347,7 +369,7 @@ afterEach(() => { }); describe("target coalesced behavior", () => { - test("coalesces a same-tick assistant burst after the 33ms window", async () => { + test("coalesces a same-tick assistant burst after the 200ms window", async () => { vi.useFakeTimers(); const harness = createHarness(); try { @@ -358,7 +380,7 @@ describe("target coalesced behavior", () => { } await waitForSessionEventQueue(); - await vi.advanceTimersByTimeAsync(32); + await vi.advanceTimersByTimeAsync(199); expect(await harness.manager.getTimelineRows(agentId)).toHaveLength(0); expect(getTimelineStreamEvents(harness.events, agentId)).toHaveLength(0); @@ -386,7 +408,7 @@ describe("target coalesced behavior", () => { } }); - test("coalesces same-tick reasoning chunks after the 33ms window", async () => { + test("coalesces same-tick reasoning chunks after the 200ms window", async () => { vi.useFakeTimers(); const harness = createHarness(); try { @@ -397,7 +419,7 @@ describe("target coalesced behavior", () => { } await waitForSessionEventQueue(); - await vi.advanceTimersByTimeAsync(33); + await vi.advanceTimersByTimeAsync(200); const rows = await harness.manager.getTimelineRows(agentId); const events = getTimelineStreamEvents(harness.events, agentId); @@ -410,21 +432,26 @@ describe("target coalesced behavior", () => { } }); - test("preserves order when an atomic timeline event lands between two large bursts", async () => { + test("coalesces running tool calls without flushing buffered text", async () => { vi.useFakeTimers(); const harness = createHarness(); try { const { agentId, session } = await createManagedSession(harness); + const runningToolCall = toolCall({ output: "running" }); for (let i = 0; i < 500; i++) { session.pushEvent(assistant("a")); } - session.pushEvent(timelineEvent(TOOL_CALL)); + session.pushEvent(timelineEvent(runningToolCall)); await waitForSessionEventQueue(); + expect(await harness.manager.getTimelineRows(agentId)).toHaveLength(0); + expect(getTimelineStreamEvents(harness.events, agentId)).toHaveLength(0); + + await vi.advanceTimersByTimeAsync(200); expect(getTimelineItems(await harness.manager.getTimelineRows(agentId))).toEqual([ { type: "assistant_message", text: "a".repeat(500) }, - TOOL_CALL, + runningToolCall, ]); expectContiguousLiveSeqs(getTimelineStreamEvents(harness.events, agentId), [1, 2]); @@ -432,7 +459,7 @@ describe("target coalesced behavior", () => { session.pushEvent(assistant("b")); } await waitForSessionEventQueue(); - await vi.advanceTimersByTimeAsync(32); + await vi.advanceTimersByTimeAsync(199); expect(await harness.manager.getTimelineRows(agentId)).toHaveLength(2); await vi.advanceTimersByTimeAsync(1); @@ -441,7 +468,7 @@ describe("target coalesced behavior", () => { expect(getTimelineItems(rows)).toEqual([ { type: "assistant_message", text: "a".repeat(500) }, - TOOL_CALL, + runningToolCall, { type: "assistant_message", text: "b".repeat(500) }, ]); expectContiguousRowSeqs(rows, [1, 2, 3]); @@ -451,6 +478,126 @@ describe("target coalesced behavior", () => { } }); + test("coalesces same-window tool call updates to the latest snapshot", async () => { + vi.useFakeTimers(); + const harness = createHarness(); + try { + const { agentId, session } = await createManagedSession(harness); + + for (let i = 0; i < 200; i++) { + session.pushEvent(timelineEvent(toolCall({ output: `chunk-${i}` }))); + } + await waitForSessionEventQueue(); + + await vi.advanceTimersByTimeAsync(199); + expect(await harness.manager.getTimelineRows(agentId)).toHaveLength(0); + expect(getTimelineStreamEvents(harness.events, agentId)).toHaveLength(0); + + await vi.advanceTimersByTimeAsync(1); + const rows = await harness.manager.getTimelineRows(agentId); + const events = getTimelineStreamEvents(harness.events, agentId); + + expect(getTimelineItems(rows)).toEqual([toolCall({ output: "chunk-199" })]); + expect(events).toHaveLength(1); + expectContiguousRowSeqs(rows, [1]); + expectContiguousLiveSeqs(events, [1]); + } finally { + harness.cleanup(); + } + }); + + test("coalesces interleaved tool calls independently and preserves first arrival order", async () => { + vi.useFakeTimers(); + const harness = createHarness(); + try { + const { agentId, session } = await createManagedSession(harness); + + session.pushEvent(timelineEvent(toolCall({ callId: "tool-1", output: "one-a" }))); + session.pushEvent(timelineEvent(toolCall({ callId: "tool-2", output: "two-a" }))); + session.pushEvent(timelineEvent(toolCall({ callId: "tool-1", output: "one-b" }))); + session.pushEvent(timelineEvent(toolCall({ callId: "tool-2", output: "two-b" }))); + await waitForSessionEventQueue(); + await vi.advanceTimersByTimeAsync(200); + + const rows = await harness.manager.getTimelineRows(agentId); + const events = getTimelineStreamEvents(harness.events, agentId); + + expect(getTimelineItems(rows)).toEqual([ + toolCall({ callId: "tool-1", output: "one-b" }), + toolCall({ callId: "tool-2", output: "two-b" }), + ]); + expect(events).toHaveLength(2); + expectContiguousRowSeqs(rows, [1, 2]); + expectContiguousLiveSeqs(events, [1, 2]); + } finally { + harness.cleanup(); + } + }); + + test("terminal tool call statuses flush immediately with pending text", async () => { + vi.useFakeTimers(); + const harness = createHarness(); + try { + const { agentId, session } = await createManagedSession(harness); + + session.pushEvent(assistant("before")); + session.pushEvent(timelineEvent(toolCall({ output: "running" }))); + await waitForSessionEventQueue(); + expect(await harness.manager.getTimelineRows(agentId)).toHaveLength(0); + + session.pushEvent(timelineEvent(toolCall({ status: "completed", output: "done" }))); + await waitForSessionEventQueue(); + + const rows = await harness.manager.getTimelineRows(agentId); + const events = getTimelineStreamEvents(harness.events, agentId); + + expect(getTimelineItems(rows)).toEqual([ + { type: "assistant_message", text: "before" }, + toolCall({ status: "completed", output: "done" }), + ]); + expect(events).toHaveLength(2); + expectContiguousRowSeqs(rows, [1, 2]); + expectContiguousLiveSeqs(events, [1, 2]); + + await vi.advanceTimersByTimeAsync(200); + expect(await harness.manager.getTimelineRows(agentId)).toHaveLength(2); + expect(getTimelineStreamEvents(harness.events, agentId)).toHaveLength(2); + } finally { + harness.cleanup(); + } + }); + + test("preserves mixed text and tool call order after coalescing", async () => { + vi.useFakeTimers(); + const harness = createHarness(); + try { + const { agentId, session } = await createManagedSession(harness); + + session.pushEvent(assistant("a")); + session.pushEvent(timelineEvent(toolCall({ output: "first" }))); + session.pushEvent(reasoning("r")); + session.pushEvent(timelineEvent(toolCall({ output: "latest" }))); + session.pushEvent(assistant("b")); + await waitForSessionEventQueue(); + await vi.advanceTimersByTimeAsync(200); + + const rows = await harness.manager.getTimelineRows(agentId); + const events = getTimelineStreamEvents(harness.events, agentId); + + expect(getTimelineItems(rows)).toEqual([ + { type: "assistant_message", text: "a" }, + toolCall({ output: "latest" }), + { type: "reasoning", text: "r" }, + { type: "assistant_message", text: "b" }, + ]); + expect(events).toHaveLength(4); + expectContiguousRowSeqs(rows, [1, 2, 3, 4]); + expectContiguousLiveSeqs(events, [1, 2, 3, 4]); + } finally { + harness.cleanup(); + } + }); + test("preserves assistant and reasoning interleave", async () => { vi.useFakeTimers(); const harness = createHarness(); @@ -462,7 +609,7 @@ describe("target coalesced behavior", () => { session.pushEvent(reasoning("r2")); session.pushEvent(assistant("a2")); await waitForSessionEventQueue(); - await vi.advanceTimersByTimeAsync(33); + await vi.advanceTimersByTimeAsync(200); const rows = await harness.manager.getTimelineRows(agentId); const events = getTimelineStreamEvents(harness.events, agentId); @@ -491,7 +638,7 @@ describe("target coalesced behavior", () => { session.pushEvent(reasoning(`r${i}`)); } await waitForSessionEventQueue(); - await vi.advanceTimersByTimeAsync(33); + await vi.advanceTimersByTimeAsync(200); const rows = await harness.manager.getTimelineRows(agentId); const events = getTimelineStreamEvents(harness.events, agentId); @@ -542,7 +689,7 @@ describe("target coalesced behavior", () => { first.session.pushEvent(assistant("a")); second.session.pushEvent(assistant("b")); await waitForSessionEventQueue(); - await vi.advanceTimersByTimeAsync(33); + await vi.advanceTimersByTimeAsync(200); const firstRows = await harness.manager.getTimelineRows(first.agentId); const secondRows = await harness.manager.getTimelineRows(second.agentId); @@ -572,7 +719,7 @@ describe("target coalesced behavior", () => { session.pushEvent(assistant("codex", "codex")); session.pushEvent(assistant("claude", "claude")); await waitForSessionEventQueue(); - await vi.advanceTimersByTimeAsync(33); + await vi.advanceTimersByTimeAsync(200); const rows = await harness.manager.getTimelineRows(agentId); const events = getTimelineStreamEvents(harness.events, agentId); @@ -605,7 +752,7 @@ describe("target coalesced behavior", () => { session.pushEvent(assistant("one", "codex", "turn-1")); session.pushEvent(assistant("two", "codex", "turn-2")); await waitForSessionEventQueue(); - await vi.advanceTimersByTimeAsync(33); + await vi.advanceTimersByTimeAsync(200); const rows = await harness.manager.getTimelineRows(agentId); const events = getTimelineStreamEvents(harness.events, agentId); @@ -646,7 +793,7 @@ describe("target coalesced behavior", () => { session.pushEvent(assistant("\n\t")); session.pushEvent(assistant("done")); await waitForSessionEventQueue(); - await vi.advanceTimersByTimeAsync(33); + await vi.advanceTimersByTimeAsync(200); const rows = await harness.manager.getTimelineRows(agentId); const events = getTimelineStreamEvents(harness.events, agentId); @@ -704,7 +851,7 @@ describe("target coalesced behavior", () => { : undefined, ).toBe(undefined); - await vi.advanceTimersByTimeAsync(33); + await vi.advanceTimersByTimeAsync(200); expect(await harness.manager.getTimelineRows(agentId)).toHaveLength( expectedTimelineEventCount, ); @@ -804,7 +951,7 @@ describe("target coalesced behavior", () => { }); second.session.pushEvent(assistant("new")); await waitForSessionEventQueue(); - await vi.advanceTimersByTimeAsync(100); + await vi.advanceTimersByTimeAsync(200); expect(getTimelineItems(await reuseHarness.manager.getTimelineRows(second.agentId))).toEqual([ { type: "assistant_message", text: "new" }, @@ -878,7 +1025,7 @@ describe("target coalesced behavior", () => { session.pushEvent(assistant("A")); await waitForSessionEventQueue(); - await vi.advanceTimersByTimeAsync(33); + await vi.advanceTimersByTimeAsync(200); await waitForSessionEventQueue(); const rowsAfterFirstFlush = await reentryHarness.manager.getTimelineRows(agentId); @@ -887,7 +1034,7 @@ describe("target coalesced behavior", () => { ]); expect(getTimelineStreamEvents(reentryHarness.events, agentId)).toHaveLength(1); - await vi.advanceTimersByTimeAsync(33); + await vi.advanceTimersByTimeAsync(200); await waitForSessionEventQueue(); const rowsAfterSecondFlush = await reentryHarness.manager.getTimelineRows(agentId); diff --git a/packages/server/src/server/agent/agent-stream-coalescer.test.ts b/packages/server/src/server/agent/agent-stream-coalescer.test.ts index a276c373d..270ba8253 100644 --- a/packages/server/src/server/agent/agent-stream-coalescer.test.ts +++ b/packages/server/src/server/agent/agent-stream-coalescer.test.ts @@ -60,6 +60,33 @@ function reasoning( return timeline({ type: "reasoning", text }, options); } +function toolCall(options?: { + callId?: string; + status?: "running" | "completed" | "failed" | "canceled"; + output?: string; + provider?: AgentProvider; + turnId?: string; + error?: unknown; +}): Extract { + const status = options?.status ?? "running"; + return timeline( + { + type: "tool_call", + callId: options?.callId ?? "tool-1", + name: "shell", + status, + error: status === "failed" ? (options?.error ?? "failed") : null, + detail: { + type: "shell", + command: "printf ok", + output: options?.output ?? "", + exitCode: status === "completed" ? 0 : null, + }, + }, + options, + ); +} + describe("AgentStreamCoalescer", () => { beforeEach(() => { vi.useFakeTimers(); @@ -76,7 +103,7 @@ describe("AgentStreamCoalescer", () => { expect(coalescer.handle("agent-1", assistant("lo"))).toBe(true); expect(flushes).toEqual([]); - await vi.advanceTimersByTimeAsync(32); + await vi.advanceTimersByTimeAsync(199); expect(flushes).toEqual([]); await vi.advanceTimersByTimeAsync(1); @@ -92,7 +119,7 @@ describe("AgentStreamCoalescer", () => { test("uses constructor windowMs instead of a hard-coded value", async () => { const { coalescer, flushes } = createHarness(10); - expect(AGENT_STREAM_COALESCE_DEFAULT_WINDOW_MS).toBe(33); + expect(AGENT_STREAM_COALESCE_DEFAULT_WINDOW_MS).toBe(200); expect(coalescer.handle("agent-1", assistant("fast"))).toBe(true); await vi.advanceTimersByTimeAsync(9); @@ -108,7 +135,7 @@ describe("AgentStreamCoalescer", () => { ]); }); - test("flushes before non-chunkable boundaries when caller invokes flushFor", () => { + test("flushFor drains pending coalesced events", () => { const { coalescer, flushes } = createHarness(); expect(coalescer.handle("agent-1", assistant("before"))).toBe(true); @@ -121,30 +148,12 @@ describe("AgentStreamCoalescer", () => { provider: "codex", }, ]); - expect( - coalescer.handle( - "agent-1", - timeline({ - type: "tool_call", - callId: "tool-1", - name: "shell", - status: "running", - error: null, - }), - ), - ).toBe(false); + expect(coalescer.handle("agent-1", toolCall())).toBe(true); }); test("does not consume non-chunkable events", () => { const { coalescer, flushes } = createHarness(); const nonChunkableEvents: AgentStreamEvent[] = [ - timeline({ - type: "tool_call", - callId: "tool-1", - name: "shell", - status: "completed", - error: null, - }), timeline({ type: "todo", items: [{ text: "ship it", completed: false }] }), timeline({ type: "user_message", text: "hi", messageId: "message-1" }), timeline({ type: "error", message: "boom" }), @@ -194,7 +203,7 @@ describe("AgentStreamCoalescer", () => { coalescer.handle("agent-1", reasoning("r1")); coalescer.handle("agent-1", assistant("a2")); - await vi.advanceTimersByTimeAsync(33); + await vi.advanceTimersByTimeAsync(200); expect(flushes).toEqual([ { agentId: "agent-1", @@ -222,7 +231,7 @@ describe("AgentStreamCoalescer", () => { coalescer.handle("agent-1", assistant("a2")); coalescer.handle("agent-1", reasoning("r2")); - await vi.advanceTimersByTimeAsync(33); + await vi.advanceTimersByTimeAsync(200); expect(flushes.map((flush) => flush.item)).toEqual([ { type: "assistant_message", text: "a1" }, { type: "reasoning", text: "r1" }, @@ -237,7 +246,7 @@ describe("AgentStreamCoalescer", () => { coalescer.handle("agent-1", assistant("c", { provider: "codex" })); coalescer.handle("agent-1", assistant("o", { provider: "opencode" })); - await vi.advanceTimersByTimeAsync(33); + await vi.advanceTimersByTimeAsync(200); expect(flushes).toEqual([ { agentId: "agent-1", @@ -258,7 +267,7 @@ describe("AgentStreamCoalescer", () => { coalescer.handle("agent-1", assistant("one", { turnId: "turn-1" })); coalescer.handle("agent-1", assistant("two", { turnId: "turn-2" })); - await vi.advanceTimersByTimeAsync(33); + await vi.advanceTimersByTimeAsync(200); expect(flushes).toEqual([ { agentId: "agent-1", @@ -281,7 +290,7 @@ describe("AgentStreamCoalescer", () => { expect(coalescer.handle("agent-1", assistant(""))).toBe(true); expect(coalescer.handle("agent-2", reasoning(""))).toBe(true); - await vi.advanceTimersByTimeAsync(33); + await vi.advanceTimersByTimeAsync(200); expect(flushes).toEqual([]); }); @@ -293,7 +302,7 @@ describe("AgentStreamCoalescer", () => { coalescer.handle("agent-1", assistant("\t")); coalescer.handle("agent-1", assistant(" mixed \n\t whitespace ")); - await vi.advanceTimersByTimeAsync(33); + await vi.advanceTimersByTimeAsync(200); expect(flushes).toEqual([ { agentId: "agent-1", @@ -323,7 +332,7 @@ describe("AgentStreamCoalescer", () => { coalescer.handle("agent-1", assistant(chunk)); } - await vi.advanceTimersByTimeAsync(33); + await vi.advanceTimersByTimeAsync(200); expect(flushes).toEqual([ { agentId: "agent-1", @@ -341,7 +350,7 @@ describe("AgentStreamCoalescer", () => { coalescer.handle("agent-1", assistant("b")); coalescer.handle("agent-2", assistant("y")); - await vi.advanceTimersByTimeAsync(33); + await vi.advanceTimersByTimeAsync(200); expect(flushes).toEqual([ { agentId: "agent-1", @@ -422,7 +431,7 @@ describe("AgentStreamCoalescer", () => { }, ]); - await vi.advanceTimersByTimeAsync(33); + await vi.advanceTimersByTimeAsync(200); expect(flushes).toEqual([ { agentId: "agent-1", @@ -443,7 +452,7 @@ describe("AgentStreamCoalescer", () => { coalescer.handle("agent-1", assistant("manual")); coalescer.flushFor("agent-1"); - await vi.advanceTimersByTimeAsync(33); + await vi.advanceTimersByTimeAsync(200); expect(flushes).toEqual([ { agentId: "agent-1", @@ -459,7 +468,7 @@ describe("AgentStreamCoalescer", () => { coalescer.handle("agent-1", assistant("durable")); coalescer.flushAndDiscard("agent-1"); - await vi.advanceTimersByTimeAsync(33); + await vi.advanceTimersByTimeAsync(200); expect(flushes).toEqual([ { agentId: "agent-1", @@ -523,7 +532,7 @@ describe("AgentStreamCoalescer", () => { coalescer.handle("agent-1", timeline(firstItem)); coalescer.handle("agent-1", assistant("llo")); - await vi.advanceTimersByTimeAsync(33); + await vi.advanceTimersByTimeAsync(200); expect(flushes).toEqual([ { agentId: "agent-1", @@ -536,4 +545,131 @@ describe("AgentStreamCoalescer", () => { }, ]); }); + + test("coalesces tool call updates by callId with latest snapshot winning", async () => { + const { coalescer, flushes } = createHarness(); + + expect(coalescer.handle("agent-1", toolCall({ output: "first" }))).toBe(true); + expect(coalescer.handle("agent-1", toolCall({ output: "second" }))).toBe(true); + + await vi.advanceTimersByTimeAsync(200); + + expect(flushes).toEqual([ + { + agentId: "agent-1", + item: { + type: "tool_call", + callId: "tool-1", + name: "shell", + status: "running", + error: null, + detail: { + type: "shell", + command: "printf ok", + output: "second", + exitCode: null, + }, + }, + provider: "codex", + }, + ]); + }); + + test("coalesces interleaved tool call updates independently by callId", async () => { + const { coalescer, flushes } = createHarness(); + + coalescer.handle("agent-1", toolCall({ callId: "tool-1", output: "one-a" })); + coalescer.handle("agent-1", toolCall({ callId: "tool-2", output: "two-a" })); + coalescer.handle("agent-1", toolCall({ callId: "tool-1", output: "one-b" })); + coalescer.handle("agent-1", toolCall({ callId: "tool-2", output: "two-b" })); + + await vi.advanceTimersByTimeAsync(200); + + expect(flushes.map((flush) => flush.item)).toEqual([ + { + type: "tool_call", + callId: "tool-1", + name: "shell", + status: "running", + error: null, + detail: { + type: "shell", + command: "printf ok", + output: "one-b", + exitCode: null, + }, + }, + { + type: "tool_call", + callId: "tool-2", + name: "shell", + status: "running", + error: null, + detail: { + type: "shell", + command: "printf ok", + output: "two-b", + exitCode: null, + }, + }, + ]); + }); + + test("terminal tool call statuses flush immediately", () => { + const { coalescer, flushes } = createHarness(); + + coalescer.handle("agent-1", toolCall({ output: "running" })); + expect(flushes).toEqual([]); + + expect(coalescer.handle("agent-1", toolCall({ status: "completed", output: "done" }))).toBe( + true, + ); + + expect(flushes.map((flush) => flush.item)).toEqual([ + { + type: "tool_call", + callId: "tool-1", + name: "shell", + status: "completed", + error: null, + detail: { + type: "shell", + command: "printf ok", + output: "done", + exitCode: 0, + }, + }, + ]); + }); + + test("preserves mixed text and tool call arrival order within a flush", async () => { + const { coalescer, flushes } = createHarness(); + + coalescer.handle("agent-1", assistant("a")); + coalescer.handle("agent-1", toolCall({ output: "running" })); + coalescer.handle("agent-1", reasoning("r")); + coalescer.handle("agent-1", toolCall({ output: "latest" })); + coalescer.handle("agent-1", assistant("b")); + + await vi.advanceTimersByTimeAsync(200); + + expect(flushes.map((flush) => flush.item)).toEqual([ + { type: "assistant_message", text: "a" }, + { + type: "tool_call", + callId: "tool-1", + name: "shell", + status: "running", + error: null, + detail: { + type: "shell", + command: "printf ok", + output: "latest", + exitCode: null, + }, + }, + { type: "reasoning", text: "r" }, + { type: "assistant_message", text: "b" }, + ]); + }); }); diff --git a/packages/server/src/server/agent/agent-stream-coalescer.ts b/packages/server/src/server/agent/agent-stream-coalescer.ts index 878794a72..60ca89bcf 100644 --- a/packages/server/src/server/agent/agent-stream-coalescer.ts +++ b/packages/server/src/server/agent/agent-stream-coalescer.ts @@ -1,11 +1,13 @@ import type { AgentProvider, AgentStreamEvent, AgentTimelineItem } from "./agent-sdk-types.js"; -export const AGENT_STREAM_COALESCE_DEFAULT_WINDOW_MS = 33; +export const AGENT_STREAM_COALESCE_DEFAULT_WINDOW_MS = 200; -type ChunkableTimelineKind = "assistant_message" | "reasoning"; -type ChunkableTimelineItem = Extract; -type ChunkableTimelineEvent = Extract & { - item: ChunkableTimelineItem; +type CoalescableTextKind = "assistant_message" | "reasoning"; +type CoalescableTimelineKind = CoalescableTextKind | "tool_call"; +type CoalescableTextItem = Extract; +type CoalescableTimelineItem = Extract; +type CoalescableTimelineEvent = Extract & { + item: CoalescableTimelineItem; }; export type AgentStreamCoalescerTimers = { @@ -15,7 +17,7 @@ export type AgentStreamCoalescerTimers = { export type AgentStreamCoalescerFlush = { agentId: string; - item: Extract; + item: CoalescableTimelineItem; provider: AgentProvider; turnId?: string; }; @@ -26,25 +28,48 @@ export type AgentStreamCoalescerOptions = { onFlush: (payload: AgentStreamCoalescerFlush) => void; }; -type PendingAgentStreamChunk = { - item: ChunkableTimelineItem; +type PendingTextEntry = { + kind: "text"; + item: CoalescableTextItem; text: string; provider: AgentProvider; turnId?: string; }; +type PendingToolCallEntry = { + kind: "tool_call"; + item: Extract; + provider: AgentProvider; + turnId?: string; +}; + +type PendingAgentStreamEntry = PendingTextEntry | PendingToolCallEntry; + type PendingAgentStreamBuffer = { agentId: string; - chunks: PendingAgentStreamChunk[]; + entries: PendingAgentStreamEntry[]; + toolCallEntryIndexes: Map; timer: ReturnType | null; flushing: boolean; }; -function isChunkableTimelineEvent(event: AgentStreamEvent): event is ChunkableTimelineEvent { +function isCoalescableTimelineEvent(event: AgentStreamEvent): event is CoalescableTimelineEvent { return ( event.type === "timeline" && - (event.item.type === "assistant_message" || event.item.type === "reasoning") && - typeof event.item.text === "string" + (event.item.type === "assistant_message" || + event.item.type === "reasoning" || + event.item.type === "tool_call") + ); +} + +function isTextTimelineItem(item: CoalescableTimelineItem): item is CoalescableTextItem { + return item.type === "assistant_message" || item.type === "reasoning"; +} + +function isTerminalToolCall(item: CoalescableTimelineItem): boolean { + return ( + item.type === "tool_call" && + (item.status === "completed" || item.status === "failed" || item.status === "canceled") ); } @@ -61,21 +86,21 @@ export class AgentStreamCoalescer { } handle(agentId: string, event: AgentStreamEvent): boolean { - if (!isChunkableTimelineEvent(event)) { + if (!isCoalescableTimelineEvent(event)) { return false; } - if (event.item.text === "") { + if (isTextTimelineItem(event.item) && event.item.text === "") { return true; } const buffer = this.getOrCreateBuffer(agentId); - buffer.chunks.push({ - item: event.item, - text: event.item.text, - provider: event.provider, - ...(event.turnId !== undefined ? { turnId: event.turnId } : {}), - }); + this.appendToBuffer(buffer, event); + + if (isTerminalToolCall(event.item)) { + this.flushBuffer(agentId); + return true; + } if (!buffer.timer) { this.scheduleFlush(buffer); @@ -111,7 +136,8 @@ export class AgentStreamCoalescer { const buffer: PendingAgentStreamBuffer = { agentId, - chunks: [], + entries: [], + toolCallEntryIndexes: new Map(), timer: null, flushing: false, }; @@ -119,6 +145,35 @@ export class AgentStreamCoalescer { return buffer; } + private appendToBuffer(buffer: PendingAgentStreamBuffer, event: CoalescableTimelineEvent): void { + if (isTextTimelineItem(event.item)) { + buffer.entries.push({ + kind: "text", + item: event.item, + text: event.item.text, + provider: event.provider, + ...(event.turnId !== undefined ? { turnId: event.turnId } : {}), + }); + return; + } + + const existingIndex = buffer.toolCallEntryIndexes.get(event.item.callId); + const entry: PendingToolCallEntry = { + kind: "tool_call", + item: event.item, + provider: event.provider, + ...(event.turnId !== undefined ? { turnId: event.turnId } : {}), + }; + + if (existingIndex !== undefined) { + buffer.entries[existingIndex] = entry; + return; + } + + buffer.toolCallEntryIndexes.set(event.item.callId, buffer.entries.length); + buffer.entries.push(entry); + } + private scheduleFlush(buffer: PendingAgentStreamBuffer): void { const timer = this.timers.setTimeout(() => { this.flushBuffer(buffer.agentId, buffer); @@ -148,24 +203,28 @@ export class AgentStreamCoalescer { } this.clearTimer(buffer); - if (buffer.chunks.length === 0) { + if (buffer.entries.length === 0) { return; } - const chunks = buffer.chunks; - buffer.chunks = []; + const entries = buffer.entries; + buffer.entries = []; + buffer.toolCallEntryIndexes.clear(); buffer.flushing = true; try { - for (const chunk of this.collapseChunks(chunks)) { + for (const entry of this.collapseEntries(entries)) { this.onFlush({ agentId, - item: { - ...chunk.item, - text: chunk.text, - }, - provider: chunk.provider, - ...(chunk.turnId !== undefined ? { turnId: chunk.turnId } : {}), + item: + entry.kind === "text" + ? { + ...entry.item, + text: entry.text, + } + : entry.item, + provider: entry.provider, + ...(entry.turnId !== undefined ? { turnId: entry.turnId } : {}), }); } } finally { @@ -173,22 +232,24 @@ export class AgentStreamCoalescer { } } - private collapseChunks(chunks: PendingAgentStreamChunk[]): PendingAgentStreamChunk[] { - const collapsed: PendingAgentStreamChunk[] = []; + private collapseEntries(entries: PendingAgentStreamEntry[]): PendingAgentStreamEntry[] { + const collapsed: PendingAgentStreamEntry[] = []; - for (const chunk of chunks) { + for (const entry of entries) { const previous = collapsed.at(-1); if ( previous && - previous.item.type === chunk.item.type && - previous.provider === chunk.provider && - previous.turnId === chunk.turnId + previous.kind === "text" && + entry.kind === "text" && + previous.item.type === entry.item.type && + previous.provider === entry.provider && + previous.turnId === entry.turnId ) { - previous.text += chunk.text; + previous.text += entry.text; continue; } - collapsed.push({ ...chunk }); + collapsed.push({ ...entry }); } return collapsed;