fix(server,app): coalesce tool call stream events to prevent JS thread freeze

Pi's ACP binary streams 200+ redundant tool_call_update events per tool call,
bypassing the server coalescer and flooding the app with synchronous state
updates. Two-layer fix:

Server: extend AgentStreamCoalescer to handle tool_call items via last-write-wins
per callId (200ms window, up from 33ms). Terminal statuses flush immediately.

App: mergeToolCallDetail and appendAgentToolCall now return existing references
when fields are unchanged, preserving React.memo equality checks.
This commit is contained in:
Mohamed Boudra
2026-04-19 14:26:50 +07:00
parent 2238171955
commit 5a232d8cf1
5 changed files with 682 additions and 102 deletions
+189 -1
View File
@@ -3,6 +3,7 @@ import { describe, it } from "vitest";
import { import {
hydrateStreamState, hydrateStreamState,
mergeToolCallDetail,
reduceStreamUpdate, reduceStreamUpdate,
type AgentToolCallItem, type AgentToolCallItem,
type StreamItem, type StreamItem,
@@ -10,7 +11,7 @@ import {
} from "./stream"; } from "./stream";
import type { AgentProvider, ToolCallDetail } from "@server/server/agent/agent-sdk-types"; import type { AgentProvider, ToolCallDetail } from "@server/server/agent/agent-sdk-types";
import type { AgentStreamEventPayload } from "@server/shared/messages"; 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"; 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", () => { describe("stream reducer canonical tool calls", () => {
it("is deterministic for equivalent hydration sequences", () => { it("is deterministic for equivalent hydration sequences", () => {
const updates = [ const updates = [
+57 -9
View File
@@ -336,7 +336,36 @@ function mergeUnknownValue(existing: unknown | null, incoming: unknown | null):
return incoming; return incoming;
} }
function mergeToolCallDetail(existing: ToolCallDetail, incoming: ToolCallDetail): ToolCallDetail { function hasSameIncomingFields<T extends Record<string, unknown>>(
existing: T,
incoming: T,
): boolean {
return Object.entries(incoming).every(([key, value]) => existing[key] === value);
}
function mergeToolCallMetadata(
existing: Record<string, unknown> | undefined,
incoming: Record<string, unknown> | undefined,
): Record<string, unknown> | 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") { if (existing.type === "unknown" && incoming.type !== "unknown") {
return incoming; return incoming;
} }
@@ -346,14 +375,24 @@ function mergeToolCallDetail(existing: ToolCallDetail, incoming: ToolCallDetail)
} }
if (existing.type === "unknown" && incoming.type === "unknown") { 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 { return {
type: "unknown", type: "unknown",
input: mergeUnknownValue(existing.input, incoming.input), input,
output: mergeUnknownValue(existing.output, incoming.output), output,
}; };
} }
if (existing.type === incoming.type) { if (existing.type === incoming.type) {
if (hasSameIncomingFields(existing, incoming)) {
return existing;
}
return { ...existing, ...incoming } as ToolCallDetail; return { ...existing, ...incoming } as ToolCallDetail;
} }
@@ -391,8 +430,7 @@ function appendAgentToolCall(
const existingIndex = findExistingAgentToolCallIndex(state, data.callId); const existingIndex = findExistingAgentToolCallIndex(state, data.callId);
if (existingIndex >= 0) { if (existingIndex >= 0) {
const next = [...state]; const existing = state[existingIndex];
const existing = next[existingIndex];
if (!existing || !isAgentToolCallItem(existing)) { if (!existing || !isAgentToolCallItem(existing)) {
return state; return state;
} }
@@ -401,12 +439,22 @@ function appendAgentToolCall(
mergedStatus === "failed" mergedStatus === "failed"
? (data.error ?? existing.payload.data.error ?? { message: "Tool call failed" }) ? (data.error ?? existing.payload.data.error ?? { message: "Tool call failed" })
: null; : null;
const mergedMetadata = const mergedMetadata = mergeToolCallMetadata(existing.payload.data.metadata, data.metadata);
data.metadata || existing.payload.data.metadata
? { ...existing.payload.data.metadata, ...data.metadata }
: undefined;
const mergedDetail = mergeToolCallDetail(existing.payload.data.detail, data.detail); 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] = { next[existingIndex] = {
...existing, ...existing,
timestamp, timestamp,
@@ -57,6 +57,28 @@ const TOOL_CALL: AgentTimelineItem = {
}, },
}; };
function toolCall(options?: {
callId?: string;
status?: "running" | "completed" | "failed" | "canceled";
output?: string;
error?: unknown;
}): Extract<AgentTimelineItem, { type: "tool_call" }> {
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 { class TestAgentSession implements AgentSession {
readonly capabilities = TEST_CAPABILITIES; readonly capabilities = TEST_CAPABILITIES;
readonly id: string; readonly id: string;
@@ -347,7 +369,7 @@ afterEach(() => {
}); });
describe("target coalesced behavior", () => { 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(); vi.useFakeTimers();
const harness = createHarness(); const harness = createHarness();
try { try {
@@ -358,7 +380,7 @@ describe("target coalesced behavior", () => {
} }
await waitForSessionEventQueue(); await waitForSessionEventQueue();
await vi.advanceTimersByTimeAsync(32); await vi.advanceTimersByTimeAsync(199);
expect(await harness.manager.getTimelineRows(agentId)).toHaveLength(0); expect(await harness.manager.getTimelineRows(agentId)).toHaveLength(0);
expect(getTimelineStreamEvents(harness.events, 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(); vi.useFakeTimers();
const harness = createHarness(); const harness = createHarness();
try { try {
@@ -397,7 +419,7 @@ describe("target coalesced behavior", () => {
} }
await waitForSessionEventQueue(); await waitForSessionEventQueue();
await vi.advanceTimersByTimeAsync(33); await vi.advanceTimersByTimeAsync(200);
const rows = await harness.manager.getTimelineRows(agentId); const rows = await harness.manager.getTimelineRows(agentId);
const events = getTimelineStreamEvents(harness.events, 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(); vi.useFakeTimers();
const harness = createHarness(); const harness = createHarness();
try { try {
const { agentId, session } = await createManagedSession(harness); const { agentId, session } = await createManagedSession(harness);
const runningToolCall = toolCall({ output: "running" });
for (let i = 0; i < 500; i++) { for (let i = 0; i < 500; i++) {
session.pushEvent(assistant("a")); session.pushEvent(assistant("a"));
} }
session.pushEvent(timelineEvent(TOOL_CALL)); session.pushEvent(timelineEvent(runningToolCall));
await waitForSessionEventQueue(); 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([ expect(getTimelineItems(await harness.manager.getTimelineRows(agentId))).toEqual([
{ type: "assistant_message", text: "a".repeat(500) }, { type: "assistant_message", text: "a".repeat(500) },
TOOL_CALL, runningToolCall,
]); ]);
expectContiguousLiveSeqs(getTimelineStreamEvents(harness.events, agentId), [1, 2]); expectContiguousLiveSeqs(getTimelineStreamEvents(harness.events, agentId), [1, 2]);
@@ -432,7 +459,7 @@ describe("target coalesced behavior", () => {
session.pushEvent(assistant("b")); session.pushEvent(assistant("b"));
} }
await waitForSessionEventQueue(); await waitForSessionEventQueue();
await vi.advanceTimersByTimeAsync(32); await vi.advanceTimersByTimeAsync(199);
expect(await harness.manager.getTimelineRows(agentId)).toHaveLength(2); expect(await harness.manager.getTimelineRows(agentId)).toHaveLength(2);
await vi.advanceTimersByTimeAsync(1); await vi.advanceTimersByTimeAsync(1);
@@ -441,7 +468,7 @@ describe("target coalesced behavior", () => {
expect(getTimelineItems(rows)).toEqual([ expect(getTimelineItems(rows)).toEqual([
{ type: "assistant_message", text: "a".repeat(500) }, { type: "assistant_message", text: "a".repeat(500) },
TOOL_CALL, runningToolCall,
{ type: "assistant_message", text: "b".repeat(500) }, { type: "assistant_message", text: "b".repeat(500) },
]); ]);
expectContiguousRowSeqs(rows, [1, 2, 3]); 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 () => { test("preserves assistant and reasoning interleave", async () => {
vi.useFakeTimers(); vi.useFakeTimers();
const harness = createHarness(); const harness = createHarness();
@@ -462,7 +609,7 @@ describe("target coalesced behavior", () => {
session.pushEvent(reasoning("r2")); session.pushEvent(reasoning("r2"));
session.pushEvent(assistant("a2")); session.pushEvent(assistant("a2"));
await waitForSessionEventQueue(); await waitForSessionEventQueue();
await vi.advanceTimersByTimeAsync(33); await vi.advanceTimersByTimeAsync(200);
const rows = await harness.manager.getTimelineRows(agentId); const rows = await harness.manager.getTimelineRows(agentId);
const events = getTimelineStreamEvents(harness.events, agentId); const events = getTimelineStreamEvents(harness.events, agentId);
@@ -491,7 +638,7 @@ describe("target coalesced behavior", () => {
session.pushEvent(reasoning(`r${i}`)); session.pushEvent(reasoning(`r${i}`));
} }
await waitForSessionEventQueue(); await waitForSessionEventQueue();
await vi.advanceTimersByTimeAsync(33); await vi.advanceTimersByTimeAsync(200);
const rows = await harness.manager.getTimelineRows(agentId); const rows = await harness.manager.getTimelineRows(agentId);
const events = getTimelineStreamEvents(harness.events, agentId); const events = getTimelineStreamEvents(harness.events, agentId);
@@ -542,7 +689,7 @@ describe("target coalesced behavior", () => {
first.session.pushEvent(assistant("a")); first.session.pushEvent(assistant("a"));
second.session.pushEvent(assistant("b")); second.session.pushEvent(assistant("b"));
await waitForSessionEventQueue(); await waitForSessionEventQueue();
await vi.advanceTimersByTimeAsync(33); await vi.advanceTimersByTimeAsync(200);
const firstRows = await harness.manager.getTimelineRows(first.agentId); const firstRows = await harness.manager.getTimelineRows(first.agentId);
const secondRows = await harness.manager.getTimelineRows(second.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("codex", "codex"));
session.pushEvent(assistant("claude", "claude")); session.pushEvent(assistant("claude", "claude"));
await waitForSessionEventQueue(); await waitForSessionEventQueue();
await vi.advanceTimersByTimeAsync(33); await vi.advanceTimersByTimeAsync(200);
const rows = await harness.manager.getTimelineRows(agentId); const rows = await harness.manager.getTimelineRows(agentId);
const events = getTimelineStreamEvents(harness.events, 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("one", "codex", "turn-1"));
session.pushEvent(assistant("two", "codex", "turn-2")); session.pushEvent(assistant("two", "codex", "turn-2"));
await waitForSessionEventQueue(); await waitForSessionEventQueue();
await vi.advanceTimersByTimeAsync(33); await vi.advanceTimersByTimeAsync(200);
const rows = await harness.manager.getTimelineRows(agentId); const rows = await harness.manager.getTimelineRows(agentId);
const events = getTimelineStreamEvents(harness.events, agentId); const events = getTimelineStreamEvents(harness.events, agentId);
@@ -646,7 +793,7 @@ describe("target coalesced behavior", () => {
session.pushEvent(assistant("\n\t")); session.pushEvent(assistant("\n\t"));
session.pushEvent(assistant("done")); session.pushEvent(assistant("done"));
await waitForSessionEventQueue(); await waitForSessionEventQueue();
await vi.advanceTimersByTimeAsync(33); await vi.advanceTimersByTimeAsync(200);
const rows = await harness.manager.getTimelineRows(agentId); const rows = await harness.manager.getTimelineRows(agentId);
const events = getTimelineStreamEvents(harness.events, agentId); const events = getTimelineStreamEvents(harness.events, agentId);
@@ -704,7 +851,7 @@ describe("target coalesced behavior", () => {
: undefined, : undefined,
).toBe(undefined); ).toBe(undefined);
await vi.advanceTimersByTimeAsync(33); await vi.advanceTimersByTimeAsync(200);
expect(await harness.manager.getTimelineRows(agentId)).toHaveLength( expect(await harness.manager.getTimelineRows(agentId)).toHaveLength(
expectedTimelineEventCount, expectedTimelineEventCount,
); );
@@ -804,7 +951,7 @@ describe("target coalesced behavior", () => {
}); });
second.session.pushEvent(assistant("new")); second.session.pushEvent(assistant("new"));
await waitForSessionEventQueue(); await waitForSessionEventQueue();
await vi.advanceTimersByTimeAsync(100); await vi.advanceTimersByTimeAsync(200);
expect(getTimelineItems(await reuseHarness.manager.getTimelineRows(second.agentId))).toEqual([ expect(getTimelineItems(await reuseHarness.manager.getTimelineRows(second.agentId))).toEqual([
{ type: "assistant_message", text: "new" }, { type: "assistant_message", text: "new" },
@@ -878,7 +1025,7 @@ describe("target coalesced behavior", () => {
session.pushEvent(assistant("A")); session.pushEvent(assistant("A"));
await waitForSessionEventQueue(); await waitForSessionEventQueue();
await vi.advanceTimersByTimeAsync(33); await vi.advanceTimersByTimeAsync(200);
await waitForSessionEventQueue(); await waitForSessionEventQueue();
const rowsAfterFirstFlush = await reentryHarness.manager.getTimelineRows(agentId); const rowsAfterFirstFlush = await reentryHarness.manager.getTimelineRows(agentId);
@@ -887,7 +1034,7 @@ describe("target coalesced behavior", () => {
]); ]);
expect(getTimelineStreamEvents(reentryHarness.events, agentId)).toHaveLength(1); expect(getTimelineStreamEvents(reentryHarness.events, agentId)).toHaveLength(1);
await vi.advanceTimersByTimeAsync(33); await vi.advanceTimersByTimeAsync(200);
await waitForSessionEventQueue(); await waitForSessionEventQueue();
const rowsAfterSecondFlush = await reentryHarness.manager.getTimelineRows(agentId); const rowsAfterSecondFlush = await reentryHarness.manager.getTimelineRows(agentId);
@@ -60,6 +60,33 @@ function reasoning(
return timeline({ type: "reasoning", text }, options); 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<AgentStreamEvent, { type: "timeline" }> {
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", () => { describe("AgentStreamCoalescer", () => {
beforeEach(() => { beforeEach(() => {
vi.useFakeTimers(); vi.useFakeTimers();
@@ -76,7 +103,7 @@ describe("AgentStreamCoalescer", () => {
expect(coalescer.handle("agent-1", assistant("lo"))).toBe(true); expect(coalescer.handle("agent-1", assistant("lo"))).toBe(true);
expect(flushes).toEqual([]); expect(flushes).toEqual([]);
await vi.advanceTimersByTimeAsync(32); await vi.advanceTimersByTimeAsync(199);
expect(flushes).toEqual([]); expect(flushes).toEqual([]);
await vi.advanceTimersByTimeAsync(1); await vi.advanceTimersByTimeAsync(1);
@@ -92,7 +119,7 @@ describe("AgentStreamCoalescer", () => {
test("uses constructor windowMs instead of a hard-coded value", async () => { test("uses constructor windowMs instead of a hard-coded value", async () => {
const { coalescer, flushes } = createHarness(10); 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); expect(coalescer.handle("agent-1", assistant("fast"))).toBe(true);
await vi.advanceTimersByTimeAsync(9); 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(); const { coalescer, flushes } = createHarness();
expect(coalescer.handle("agent-1", assistant("before"))).toBe(true); expect(coalescer.handle("agent-1", assistant("before"))).toBe(true);
@@ -121,30 +148,12 @@ describe("AgentStreamCoalescer", () => {
provider: "codex", provider: "codex",
}, },
]); ]);
expect( expect(coalescer.handle("agent-1", toolCall())).toBe(true);
coalescer.handle(
"agent-1",
timeline({
type: "tool_call",
callId: "tool-1",
name: "shell",
status: "running",
error: null,
}),
),
).toBe(false);
}); });
test("does not consume non-chunkable events", () => { test("does not consume non-chunkable events", () => {
const { coalescer, flushes } = createHarness(); const { coalescer, flushes } = createHarness();
const nonChunkableEvents: AgentStreamEvent[] = [ 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: "todo", items: [{ text: "ship it", completed: false }] }),
timeline({ type: "user_message", text: "hi", messageId: "message-1" }), timeline({ type: "user_message", text: "hi", messageId: "message-1" }),
timeline({ type: "error", message: "boom" }), timeline({ type: "error", message: "boom" }),
@@ -194,7 +203,7 @@ describe("AgentStreamCoalescer", () => {
coalescer.handle("agent-1", reasoning("r1")); coalescer.handle("agent-1", reasoning("r1"));
coalescer.handle("agent-1", assistant("a2")); coalescer.handle("agent-1", assistant("a2"));
await vi.advanceTimersByTimeAsync(33); await vi.advanceTimersByTimeAsync(200);
expect(flushes).toEqual([ expect(flushes).toEqual([
{ {
agentId: "agent-1", agentId: "agent-1",
@@ -222,7 +231,7 @@ describe("AgentStreamCoalescer", () => {
coalescer.handle("agent-1", assistant("a2")); coalescer.handle("agent-1", assistant("a2"));
coalescer.handle("agent-1", reasoning("r2")); coalescer.handle("agent-1", reasoning("r2"));
await vi.advanceTimersByTimeAsync(33); await vi.advanceTimersByTimeAsync(200);
expect(flushes.map((flush) => flush.item)).toEqual([ expect(flushes.map((flush) => flush.item)).toEqual([
{ type: "assistant_message", text: "a1" }, { type: "assistant_message", text: "a1" },
{ type: "reasoning", text: "r1" }, { type: "reasoning", text: "r1" },
@@ -237,7 +246,7 @@ describe("AgentStreamCoalescer", () => {
coalescer.handle("agent-1", assistant("c", { provider: "codex" })); coalescer.handle("agent-1", assistant("c", { provider: "codex" }));
coalescer.handle("agent-1", assistant("o", { provider: "opencode" })); coalescer.handle("agent-1", assistant("o", { provider: "opencode" }));
await vi.advanceTimersByTimeAsync(33); await vi.advanceTimersByTimeAsync(200);
expect(flushes).toEqual([ expect(flushes).toEqual([
{ {
agentId: "agent-1", agentId: "agent-1",
@@ -258,7 +267,7 @@ describe("AgentStreamCoalescer", () => {
coalescer.handle("agent-1", assistant("one", { turnId: "turn-1" })); coalescer.handle("agent-1", assistant("one", { turnId: "turn-1" }));
coalescer.handle("agent-1", assistant("two", { turnId: "turn-2" })); coalescer.handle("agent-1", assistant("two", { turnId: "turn-2" }));
await vi.advanceTimersByTimeAsync(33); await vi.advanceTimersByTimeAsync(200);
expect(flushes).toEqual([ expect(flushes).toEqual([
{ {
agentId: "agent-1", agentId: "agent-1",
@@ -281,7 +290,7 @@ describe("AgentStreamCoalescer", () => {
expect(coalescer.handle("agent-1", assistant(""))).toBe(true); expect(coalescer.handle("agent-1", assistant(""))).toBe(true);
expect(coalescer.handle("agent-2", reasoning(""))).toBe(true); expect(coalescer.handle("agent-2", reasoning(""))).toBe(true);
await vi.advanceTimersByTimeAsync(33); await vi.advanceTimersByTimeAsync(200);
expect(flushes).toEqual([]); expect(flushes).toEqual([]);
}); });
@@ -293,7 +302,7 @@ describe("AgentStreamCoalescer", () => {
coalescer.handle("agent-1", assistant("\t")); coalescer.handle("agent-1", assistant("\t"));
coalescer.handle("agent-1", assistant(" mixed \n\t whitespace ")); coalescer.handle("agent-1", assistant(" mixed \n\t whitespace "));
await vi.advanceTimersByTimeAsync(33); await vi.advanceTimersByTimeAsync(200);
expect(flushes).toEqual([ expect(flushes).toEqual([
{ {
agentId: "agent-1", agentId: "agent-1",
@@ -323,7 +332,7 @@ describe("AgentStreamCoalescer", () => {
coalescer.handle("agent-1", assistant(chunk)); coalescer.handle("agent-1", assistant(chunk));
} }
await vi.advanceTimersByTimeAsync(33); await vi.advanceTimersByTimeAsync(200);
expect(flushes).toEqual([ expect(flushes).toEqual([
{ {
agentId: "agent-1", agentId: "agent-1",
@@ -341,7 +350,7 @@ describe("AgentStreamCoalescer", () => {
coalescer.handle("agent-1", assistant("b")); coalescer.handle("agent-1", assistant("b"));
coalescer.handle("agent-2", assistant("y")); coalescer.handle("agent-2", assistant("y"));
await vi.advanceTimersByTimeAsync(33); await vi.advanceTimersByTimeAsync(200);
expect(flushes).toEqual([ expect(flushes).toEqual([
{ {
agentId: "agent-1", agentId: "agent-1",
@@ -422,7 +431,7 @@ describe("AgentStreamCoalescer", () => {
}, },
]); ]);
await vi.advanceTimersByTimeAsync(33); await vi.advanceTimersByTimeAsync(200);
expect(flushes).toEqual([ expect(flushes).toEqual([
{ {
agentId: "agent-1", agentId: "agent-1",
@@ -443,7 +452,7 @@ describe("AgentStreamCoalescer", () => {
coalescer.handle("agent-1", assistant("manual")); coalescer.handle("agent-1", assistant("manual"));
coalescer.flushFor("agent-1"); coalescer.flushFor("agent-1");
await vi.advanceTimersByTimeAsync(33); await vi.advanceTimersByTimeAsync(200);
expect(flushes).toEqual([ expect(flushes).toEqual([
{ {
agentId: "agent-1", agentId: "agent-1",
@@ -459,7 +468,7 @@ describe("AgentStreamCoalescer", () => {
coalescer.handle("agent-1", assistant("durable")); coalescer.handle("agent-1", assistant("durable"));
coalescer.flushAndDiscard("agent-1"); coalescer.flushAndDiscard("agent-1");
await vi.advanceTimersByTimeAsync(33); await vi.advanceTimersByTimeAsync(200);
expect(flushes).toEqual([ expect(flushes).toEqual([
{ {
agentId: "agent-1", agentId: "agent-1",
@@ -523,7 +532,7 @@ describe("AgentStreamCoalescer", () => {
coalescer.handle("agent-1", timeline(firstItem)); coalescer.handle("agent-1", timeline(firstItem));
coalescer.handle("agent-1", assistant("llo")); coalescer.handle("agent-1", assistant("llo"));
await vi.advanceTimersByTimeAsync(33); await vi.advanceTimersByTimeAsync(200);
expect(flushes).toEqual([ expect(flushes).toEqual([
{ {
agentId: "agent-1", 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" },
]);
});
}); });
@@ -1,11 +1,13 @@
import type { AgentProvider, AgentStreamEvent, AgentTimelineItem } from "./agent-sdk-types.js"; 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 CoalescableTextKind = "assistant_message" | "reasoning";
type ChunkableTimelineItem = Extract<AgentTimelineItem, { type: ChunkableTimelineKind }>; type CoalescableTimelineKind = CoalescableTextKind | "tool_call";
type ChunkableTimelineEvent = Extract<AgentStreamEvent, { type: "timeline" }> & { type CoalescableTextItem = Extract<AgentTimelineItem, { type: CoalescableTextKind }>;
item: ChunkableTimelineItem; type CoalescableTimelineItem = Extract<AgentTimelineItem, { type: CoalescableTimelineKind }>;
type CoalescableTimelineEvent = Extract<AgentStreamEvent, { type: "timeline" }> & {
item: CoalescableTimelineItem;
}; };
export type AgentStreamCoalescerTimers = { export type AgentStreamCoalescerTimers = {
@@ -15,7 +17,7 @@ export type AgentStreamCoalescerTimers = {
export type AgentStreamCoalescerFlush = { export type AgentStreamCoalescerFlush = {
agentId: string; agentId: string;
item: Extract<AgentTimelineItem, { type: "assistant_message" | "reasoning" }>; item: CoalescableTimelineItem;
provider: AgentProvider; provider: AgentProvider;
turnId?: string; turnId?: string;
}; };
@@ -26,25 +28,48 @@ export type AgentStreamCoalescerOptions = {
onFlush: (payload: AgentStreamCoalescerFlush) => void; onFlush: (payload: AgentStreamCoalescerFlush) => void;
}; };
type PendingAgentStreamChunk = { type PendingTextEntry = {
item: ChunkableTimelineItem; kind: "text";
item: CoalescableTextItem;
text: string; text: string;
provider: AgentProvider; provider: AgentProvider;
turnId?: string; turnId?: string;
}; };
type PendingToolCallEntry = {
kind: "tool_call";
item: Extract<AgentTimelineItem, { type: "tool_call" }>;
provider: AgentProvider;
turnId?: string;
};
type PendingAgentStreamEntry = PendingTextEntry | PendingToolCallEntry;
type PendingAgentStreamBuffer = { type PendingAgentStreamBuffer = {
agentId: string; agentId: string;
chunks: PendingAgentStreamChunk[]; entries: PendingAgentStreamEntry[];
toolCallEntryIndexes: Map<string, number>;
timer: ReturnType<typeof setTimeout> | null; timer: ReturnType<typeof setTimeout> | null;
flushing: boolean; flushing: boolean;
}; };
function isChunkableTimelineEvent(event: AgentStreamEvent): event is ChunkableTimelineEvent { function isCoalescableTimelineEvent(event: AgentStreamEvent): event is CoalescableTimelineEvent {
return ( return (
event.type === "timeline" && event.type === "timeline" &&
(event.item.type === "assistant_message" || event.item.type === "reasoning") && (event.item.type === "assistant_message" ||
typeof event.item.text === "string" 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 { handle(agentId: string, event: AgentStreamEvent): boolean {
if (!isChunkableTimelineEvent(event)) { if (!isCoalescableTimelineEvent(event)) {
return false; return false;
} }
if (event.item.text === "") { if (isTextTimelineItem(event.item) && event.item.text === "") {
return true; return true;
} }
const buffer = this.getOrCreateBuffer(agentId); const buffer = this.getOrCreateBuffer(agentId);
buffer.chunks.push({ this.appendToBuffer(buffer, event);
item: event.item,
text: event.item.text, if (isTerminalToolCall(event.item)) {
provider: event.provider, this.flushBuffer(agentId);
...(event.turnId !== undefined ? { turnId: event.turnId } : {}), return true;
}); }
if (!buffer.timer) { if (!buffer.timer) {
this.scheduleFlush(buffer); this.scheduleFlush(buffer);
@@ -111,7 +136,8 @@ export class AgentStreamCoalescer {
const buffer: PendingAgentStreamBuffer = { const buffer: PendingAgentStreamBuffer = {
agentId, agentId,
chunks: [], entries: [],
toolCallEntryIndexes: new Map(),
timer: null, timer: null,
flushing: false, flushing: false,
}; };
@@ -119,6 +145,35 @@ export class AgentStreamCoalescer {
return buffer; 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 { private scheduleFlush(buffer: PendingAgentStreamBuffer): void {
const timer = this.timers.setTimeout(() => { const timer = this.timers.setTimeout(() => {
this.flushBuffer(buffer.agentId, buffer); this.flushBuffer(buffer.agentId, buffer);
@@ -148,24 +203,28 @@ export class AgentStreamCoalescer {
} }
this.clearTimer(buffer); this.clearTimer(buffer);
if (buffer.chunks.length === 0) { if (buffer.entries.length === 0) {
return; return;
} }
const chunks = buffer.chunks; const entries = buffer.entries;
buffer.chunks = []; buffer.entries = [];
buffer.toolCallEntryIndexes.clear();
buffer.flushing = true; buffer.flushing = true;
try { try {
for (const chunk of this.collapseChunks(chunks)) { for (const entry of this.collapseEntries(entries)) {
this.onFlush({ this.onFlush({
agentId, agentId,
item: { item:
...chunk.item, entry.kind === "text"
text: chunk.text, ? {
}, ...entry.item,
provider: chunk.provider, text: entry.text,
...(chunk.turnId !== undefined ? { turnId: chunk.turnId } : {}), }
: entry.item,
provider: entry.provider,
...(entry.turnId !== undefined ? { turnId: entry.turnId } : {}),
}); });
} }
} finally { } finally {
@@ -173,22 +232,24 @@ export class AgentStreamCoalescer {
} }
} }
private collapseChunks(chunks: PendingAgentStreamChunk[]): PendingAgentStreamChunk[] { private collapseEntries(entries: PendingAgentStreamEntry[]): PendingAgentStreamEntry[] {
const collapsed: PendingAgentStreamChunk[] = []; const collapsed: PendingAgentStreamEntry[] = [];
for (const chunk of chunks) { for (const entry of entries) {
const previous = collapsed.at(-1); const previous = collapsed.at(-1);
if ( if (
previous && previous &&
previous.item.type === chunk.item.type && previous.kind === "text" &&
previous.provider === chunk.provider && entry.kind === "text" &&
previous.turnId === chunk.turnId previous.item.type === entry.item.type &&
previous.provider === entry.provider &&
previous.turnId === entry.turnId
) { ) {
previous.text += chunk.text; previous.text += entry.text;
continue; continue;
} }
collapsed.push({ ...chunk }); collapsed.push({ ...entry });
} }
return collapsed; return collapsed;