jb/paseo
1
0
Fork 0
Code Issues 1 Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
packages/app/src/types/stream.test.ts
+189 -1
View File
@@ -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 = [
packages/app/src/types/stream.ts
+57 -9
View File
@@ -336,7 +336,36 @@ function mergeUnknownValue(existing: unknown | null, incoming: unknown | null):
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") {
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,
packages/server/src/server/agent/agent-manager-stream-coalescing.test.ts
+166 -19
View File
@@ -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 {
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);
packages/server/src/server/agent/agent-stream-coalescer.test.ts
+170 -34
View File
@@ -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<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", () => {
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" },
]);
});
});
packages/server/src/server/agent/agent-stream-coalescer.ts
+100 -39
View File
@@ -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<AgentTimelineItem, { type: ChunkableTimelineKind }>;
type ChunkableTimelineEvent = Extract<AgentStreamEvent, { type: "timeline" }> & {
item: ChunkableTimelineItem;
type CoalescableTextKind = "assistant_message" | "reasoning";
type CoalescableTimelineKind = CoalescableTextKind | "tool_call";
type CoalescableTextItem = Extract<AgentTimelineItem, { type: CoalescableTextKind }>;
type CoalescableTimelineItem = Extract<AgentTimelineItem, { type: CoalescableTimelineKind }>;
type CoalescableTimelineEvent = Extract<AgentStreamEvent, { type: "timeline" }> & {
item: CoalescableTimelineItem;
};
export type AgentStreamCoalescerTimers = {
@@ -15,7 +17,7 @@ export type AgentStreamCoalescerTimers = {
export type AgentStreamCoalescerFlush = {
agentId: string;
item: Extract<AgentTimelineItem, { type: "assistant_message" | "reasoning" }>;
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<AgentTimelineItem, { type: "tool_call" }>;
provider: AgentProvider;
turnId?: string;
};
type PendingAgentStreamEntry = PendingTextEntry | PendingToolCallEntry;
type PendingAgentStreamBuffer = {
agentId: string;
chunks: PendingAgentStreamChunk[];
entries: PendingAgentStreamEntry[];
toolCallEntryIndexes: Map<string, number>;
timer: ReturnType<typeof setTimeout> | 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;