From ccb956f98f825a159a8f72b041eb6bacb2d52c16 Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Sat, 18 Apr 2026 16:08:54 +0700 Subject: [PATCH] wip(server): coalesce agent stream events and add runtime metrics Introduce AgentStreamCoalescer to batch rapid-fire agent events before broadcasting, reducing WebSocket chatter during tool-heavy turns. Wire websocket-server runtime metrics so we can observe coalescing behavior in tests. Work in progress. --- .../agent-manager-stream-coalescing.test.ts | 1049 +++++++++++++++++ .../src/server/agent/agent-manager.test.ts | 28 +- .../server/src/server/agent/agent-manager.ts | 138 ++- .../agent/agent-stream-coalescer.test.ts | 539 +++++++++ .../server/agent/agent-stream-coalescer.ts | 196 +++ .../websocket-server.runtime-metrics.test.ts | 350 ++++++ .../server/src/server/websocket-server.ts | 77 ++ 7 files changed, 2320 insertions(+), 57 deletions(-) create mode 100644 packages/server/src/server/agent/agent-manager-stream-coalescing.test.ts create mode 100644 packages/server/src/server/agent/agent-stream-coalescer.test.ts create mode 100644 packages/server/src/server/agent/agent-stream-coalescer.ts create mode 100644 packages/server/src/server/websocket-server.runtime-metrics.test.ts 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 new file mode 100644 index 000000000..dc340f32e --- /dev/null +++ b/packages/server/src/server/agent/agent-manager-stream-coalescing.test.ts @@ -0,0 +1,1049 @@ +import { mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, describe, expect, test, vi } from "vitest"; + +import { createTestLogger } from "../../test-utils/test-logger.js"; +import { AgentManager, type AgentManagerEvent } from "./agent-manager.js"; +import type { AgentTimelineRow } from "./agent-timeline-store-types.js"; +import type { + AgentCapabilityFlags, + AgentClient, + AgentLaunchContext, + AgentModelDefinition, + AgentPersistenceHandle, + AgentPromptInput, + AgentProvider, + AgentRunOptions, + AgentRunResult, + AgentSession, + AgentSessionConfig, + AgentStreamEvent, + AgentTimelineItem, +} from "./agent-sdk-types.js"; + +/** + * Contract for AgentManager pre-record stream coalescing. + * Assistant/reasoning chunks coalesce before recordTimeline() assigns canonical seqs. + */ + +const TEST_CAPABILITIES: AgentCapabilityFlags = { + supportsStreaming: false, + supportsSessionPersistence: false, + supportsDynamicModes: false, + supportsMcpServers: false, + supportsReasoningStream: false, + supportsToolInvocations: false, +}; + +const AGENT_IDS = [ + "00000000-0000-4000-8000-000000000001", + "00000000-0000-4000-8000-000000000002", + "00000000-0000-4000-8000-000000000003", + "00000000-0000-4000-8000-000000000004", +] as const; + +const TOOL_CALL: AgentTimelineItem = { + type: "tool_call", + callId: "tool-1", + name: "shell", + status: "completed", + error: null, + detail: { + type: "shell", + command: "printf ok", + output: "ok", + exitCode: 0, + }, +}; + +class TestAgentSession implements AgentSession { + readonly capabilities = TEST_CAPABILITIES; + readonly id: string; + private historyEvents: AgentStreamEvent[] = []; + private readonly subscribers = new Set<(event: AgentStreamEvent) => void>(); + private turnIdCounter = 0; + + constructor( + readonly provider: AgentProvider, + private readonly config: AgentSessionConfig, + sessionId: string, + ) { + this.id = sessionId; + } + + async run(): Promise { + return { + sessionId: this.id, + finalText: "", + timeline: [], + }; + } + + async startTurn( + _prompt: AgentPromptInput, + _options?: AgentRunOptions, + ): Promise<{ turnId: string }> { + return { turnId: `turn-${++this.turnIdCounter}` }; + } + + subscribe(callback: (event: AgentStreamEvent) => void): () => void { + this.subscribers.add(callback); + return () => { + this.subscribers.delete(callback); + }; + } + + pushEvent(event: AgentStreamEvent): void { + for (const callback of this.subscribers) { + callback(event); + } + } + + setHistory(events: AgentStreamEvent[]): void { + this.historyEvents = [...events]; + } + + async *streamHistory(): AsyncGenerator { + for (const event of this.historyEvents) { + yield event; + } + } + + async getRuntimeInfo() { + return { + provider: this.provider, + sessionId: this.id, + model: this.config.model ?? null, + modeId: this.config.modeId ?? null, + }; + } + + async getAvailableModes() { + return []; + } + + async getCurrentMode() { + return null; + } + + async setMode(): Promise {} + + getPendingPermissions() { + return []; + } + + async respondToPermission(): Promise {} + + describePersistence(): AgentPersistenceHandle { + return { + provider: this.provider, + sessionId: this.id, + }; + } + + async interrupt(): Promise {} + + async close(): Promise {} +} + +class TestAgentClient implements AgentClient { + readonly capabilities = TEST_CAPABILITIES; + private sessionCounter = 0; + readonly sessions = new Map(); + + constructor(readonly provider: AgentProvider = "codex") {} + + async createSession( + config: AgentSessionConfig, + _launchContext?: AgentLaunchContext, + ): Promise { + const session = new TestAgentSession( + config.provider, + config, + `${config.provider}-session-${++this.sessionCounter}`, + ); + this.sessions.set(config.cwd, session); + return session; + } + + async resumeSession( + _handle: AgentPersistenceHandle, + config?: Partial, + _launchContext?: AgentLaunchContext, + ): Promise { + const resolvedConfig: AgentSessionConfig = { + provider: this.provider, + cwd: config?.cwd ?? process.cwd(), + ...config, + }; + return this.createSession(resolvedConfig); + } + + async listModels(): Promise { + return [ + { + provider: this.provider, + id: "test-model", + label: "Test Model", + isDefault: true, + }, + ]; + } + + async isAvailable(): Promise { + return true; + } + + getSession(cwd: string): TestAgentSession { + const session = this.sessions.get(cwd); + if (!session) { + throw new Error(`No test session for cwd ${cwd}`); + } + return session; + } +} + +type Harness = { + manager: AgentManager; + client: TestAgentClient; + events: AgentManagerEvent[]; + workdir: string; + cleanup: () => void; +}; + +function createHarness(options?: { provider?: AgentProvider }): Harness { + const workdir = mkdtempSync(join(tmpdir(), "agent-manager-stream-coalescing-")); + const client = new TestAgentClient(options?.provider ?? "codex"); + const manager = new AgentManager({ + clients: { [client.provider]: client }, + idFactory: createIdFactory(), + logger: createTestLogger(), + }); + const events: AgentManagerEvent[] = []; + manager.subscribe((event) => events.push(event), { replayState: false }); + + return { + manager, + client, + events, + workdir, + cleanup: () => rmSync(workdir, { recursive: true, force: true }), + }; +} + +function createIdFactory(): () => string { + let index = 0; + return () => AGENT_IDS[index++] ?? `00000000-0000-4000-8000-${String(index).padStart(12, "0")}`; +} + +async function createManagedSession( + harness: Harness, + options?: { agentId?: string; provider?: AgentProvider; workdir?: string }, +): Promise<{ agentId: string; session: TestAgentSession; workdir: string }> { + const agentId = options?.agentId ?? AGENT_IDS[0]; + const workdir = options?.workdir ?? harness.workdir; + await harness.manager.createAgent( + { + provider: options?.provider ?? harness.client.provider, + cwd: workdir, + }, + agentId, + ); + return { + agentId, + session: harness.client.getSession(workdir), + workdir, + }; +} + +function assistant( + text: string, + provider: AgentProvider = "codex", + turnId?: string, +): AgentStreamEvent { + return { + type: "timeline", + provider, + turnId, + item: { type: "assistant_message", text }, + }; +} + +function reasoning( + text: string, + provider: AgentProvider = "codex", + turnId?: string, +): AgentStreamEvent { + return { + type: "timeline", + provider, + turnId, + item: { type: "reasoning", text }, + }; +} + +function timelineEvent( + item: AgentTimelineItem, + provider: AgentProvider = "codex", + turnId?: string, +): AgentStreamEvent { + return { + type: "timeline", + provider, + turnId, + item, + }; +} + +function terminalEvent( + type: "turn_completed" | "turn_failed" | "turn_canceled", + turnId: string, +): AgentStreamEvent { + if (type === "turn_completed") { + return { type, provider: "codex", turnId }; + } + if (type === "turn_failed") { + return { type, provider: "codex", turnId, error: "failed" }; + } + return { type, provider: "codex", turnId, reason: "canceled" }; +} + +function getStreamEvents(events: AgentManagerEvent[], agentId: string): AgentManagerEvent[] { + return events.filter((event) => event.type === "agent_stream" && event.agentId === agentId); +} + +function getTimelineStreamEvents( + events: AgentManagerEvent[], + agentId: string, +): AgentManagerEvent[] { + return getStreamEvents(events, agentId).filter( + (event) => event.type === "agent_stream" && event.event.type === "timeline", + ); +} + +function getTimelineItems(rows: AgentTimelineRow[]): AgentTimelineItem[] { + return rows.map((row) => row.item); +} + +function expectContiguousRowSeqs(rows: AgentTimelineRow[], expected: number[]): void { + expect(rows.map((row) => row.seq)).toEqual(expected); +} + +function expectContiguousLiveSeqs(events: AgentManagerEvent[], expected: number[]): void { + expect(events.map((event) => (event.type === "agent_stream" ? event.seq : undefined))).toEqual( + expected, + ); +} + +async function waitForSessionEventQueue(): Promise { + for (let i = 0; i < 10_000; i++) { + await Promise.resolve(); + } +} + +afterEach(() => { + vi.useRealTimers(); +}); + +describe("target coalesced behavior", () => { + test("coalesces a same-tick assistant burst after the 33ms window", async () => { + vi.useFakeTimers(); + const harness = createHarness(); + try { + const { agentId, session } = await createManagedSession(harness); + + for (let i = 0; i < 1000; i++) { + session.pushEvent(assistant("x")); + } + await waitForSessionEventQueue(); + + await vi.advanceTimersByTimeAsync(32); + 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(rows).toHaveLength(1); + expect(events).toHaveLength(1); + expect(rows[0]?.item).toEqual({ + type: "assistant_message", + text: "x".repeat(1000), + }); + expect(events.map((event) => (event.type === "agent_stream" ? event.event : null))).toEqual([ + { + type: "timeline", + provider: "codex", + item: { type: "assistant_message", text: "x".repeat(1000) }, + }, + ]); + expectContiguousRowSeqs(rows, [1]); + expectContiguousLiveSeqs(events, [1]); + } finally { + harness.cleanup(); + } + }); + + test("coalesces same-tick reasoning chunks after the 33ms window", async () => { + vi.useFakeTimers(); + const harness = createHarness(); + try { + const { agentId, session } = await createManagedSession(harness); + + for (let i = 0; i < 100; i++) { + session.pushEvent(reasoning("r")); + } + await waitForSessionEventQueue(); + + await vi.advanceTimersByTimeAsync(33); + const rows = await harness.manager.getTimelineRows(agentId); + const events = getTimelineStreamEvents(harness.events, agentId); + + expect(getTimelineItems(rows)).toEqual([{ type: "reasoning", text: "r".repeat(100) }]); + expect(events).toHaveLength(1); + expectContiguousRowSeqs(rows, [1]); + expectContiguousLiveSeqs(events, [1]); + } finally { + harness.cleanup(); + } + }); + + test("preserves order when an atomic timeline event lands between two large bursts", async () => { + vi.useFakeTimers(); + const harness = createHarness(); + try { + const { agentId, session } = await createManagedSession(harness); + + for (let i = 0; i < 500; i++) { + session.pushEvent(assistant("a")); + } + session.pushEvent(timelineEvent(TOOL_CALL)); + await waitForSessionEventQueue(); + + expect(getTimelineItems(await harness.manager.getTimelineRows(agentId))).toEqual([ + { type: "assistant_message", text: "a".repeat(500) }, + TOOL_CALL, + ]); + expectContiguousLiveSeqs(getTimelineStreamEvents(harness.events, agentId), [1, 2]); + + for (let i = 0; i < 500; i++) { + session.pushEvent(assistant("b")); + } + await waitForSessionEventQueue(); + await vi.advanceTimersByTimeAsync(32); + expect(await harness.manager.getTimelineRows(agentId)).toHaveLength(2); + + await vi.advanceTimersByTimeAsync(1); + const rows = await harness.manager.getTimelineRows(agentId); + const events = getTimelineStreamEvents(harness.events, agentId); + + expect(getTimelineItems(rows)).toEqual([ + { type: "assistant_message", text: "a".repeat(500) }, + TOOL_CALL, + { type: "assistant_message", text: "b".repeat(500) }, + ]); + expectContiguousRowSeqs(rows, [1, 2, 3]); + expectContiguousLiveSeqs(events, [1, 2, 3]); + } finally { + harness.cleanup(); + } + }); + + test("preserves assistant and reasoning interleave", async () => { + vi.useFakeTimers(); + const harness = createHarness(); + try { + const { agentId, session } = await createManagedSession(harness); + + session.pushEvent(assistant("a1")); + session.pushEvent(reasoning("r1")); + session.pushEvent(reasoning("r2")); + session.pushEvent(assistant("a2")); + await waitForSessionEventQueue(); + await vi.advanceTimersByTimeAsync(33); + + const rows = await harness.manager.getTimelineRows(agentId); + const events = getTimelineStreamEvents(harness.events, agentId); + + expect(getTimelineItems(rows)).toEqual([ + { type: "assistant_message", text: "a1" }, + { type: "reasoning", text: "r1r2" }, + { type: "assistant_message", text: "a2" }, + ]); + expect(events).toHaveLength(3); + expectContiguousRowSeqs(rows, [1, 2, 3]); + expectContiguousLiveSeqs(events, [1, 2, 3]); + } finally { + harness.cleanup(); + } + }); + + test("preserves strict alternating assistant/reasoning order", async () => { + vi.useFakeTimers(); + const harness = createHarness(); + try { + const { agentId, session } = await createManagedSession(harness); + + for (let i = 1; i <= 5; i++) { + session.pushEvent(assistant(`a${i}`)); + session.pushEvent(reasoning(`r${i}`)); + } + await waitForSessionEventQueue(); + await vi.advanceTimersByTimeAsync(33); + + const rows = await harness.manager.getTimelineRows(agentId); + const events = getTimelineStreamEvents(harness.events, agentId); + + expect(getTimelineItems(rows)).toEqual([ + { type: "assistant_message", text: "a1" }, + { type: "reasoning", text: "r1" }, + { type: "assistant_message", text: "a2" }, + { type: "reasoning", text: "r2" }, + { type: "assistant_message", text: "a3" }, + { type: "reasoning", text: "r3" }, + { type: "assistant_message", text: "a4" }, + { type: "reasoning", text: "r4" }, + { type: "assistant_message", text: "a5" }, + { type: "reasoning", text: "r5" }, + ]); + expect(rows).toHaveLength(10); + expect(events).toHaveLength(10); + expectContiguousRowSeqs( + rows, + Array.from({ length: 10 }, (_, index) => index + 1), + ); + expectContiguousLiveSeqs( + events, + Array.from({ length: 10 }, (_, index) => index + 1), + ); + } finally { + harness.cleanup(); + } + }); + + test("isolates buffers per agent", async () => { + vi.useFakeTimers(); + const harness = createHarness(); + try { + const first = await createManagedSession(harness, { + agentId: AGENT_IDS[0], + workdir: harness.workdir, + }); + const secondWorkdir = mkdtempSync(join(tmpdir(), "agent-manager-stream-coalescing-")); + const second = await createManagedSession(harness, { + agentId: AGENT_IDS[1], + workdir: secondWorkdir, + }); + + first.session.pushEvent(assistant("a")); + second.session.pushEvent(assistant("b")); + first.session.pushEvent(assistant("a")); + second.session.pushEvent(assistant("b")); + await waitForSessionEventQueue(); + await vi.advanceTimersByTimeAsync(33); + + const firstRows = await harness.manager.getTimelineRows(first.agentId); + const secondRows = await harness.manager.getTimelineRows(second.agentId); + const firstEvents = getTimelineStreamEvents(harness.events, first.agentId); + const secondEvents = getTimelineStreamEvents(harness.events, second.agentId); + + expect(getTimelineItems(firstRows)).toEqual([{ type: "assistant_message", text: "aa" }]); + expect(getTimelineItems(secondRows)).toEqual([{ type: "assistant_message", text: "bb" }]); + expect(firstEvents).toHaveLength(1); + expect(secondEvents).toHaveLength(1); + expectContiguousRowSeqs(firstRows, [1]); + expectContiguousRowSeqs(secondRows, [1]); + expectContiguousLiveSeqs(firstEvents, [1]); + expectContiguousLiveSeqs(secondEvents, [1]); + rmSync(secondWorkdir, { recursive: true, force: true }); + } finally { + harness.cleanup(); + } + }); + + test("provider boundaries break collapse", async () => { + vi.useFakeTimers(); + const harness = createHarness(); + try { + const { agentId, session } = await createManagedSession(harness); + + session.pushEvent(assistant("codex", "codex")); + session.pushEvent(assistant("claude", "claude")); + await waitForSessionEventQueue(); + await vi.advanceTimersByTimeAsync(33); + + const rows = await harness.manager.getTimelineRows(agentId); + const events = getTimelineStreamEvents(harness.events, agentId); + + expect(getTimelineItems(rows)).toEqual([ + { type: "assistant_message", text: "codex" }, + { type: "assistant_message", text: "claude" }, + ]); + expect(events.map((event) => (event.type === "agent_stream" ? event.event : null))).toEqual([ + { type: "timeline", provider: "codex", item: { type: "assistant_message", text: "codex" } }, + { + type: "timeline", + provider: "claude", + item: { type: "assistant_message", text: "claude" }, + }, + ]); + expectContiguousRowSeqs(rows, [1, 2]); + expectContiguousLiveSeqs(events, [1, 2]); + } finally { + harness.cleanup(); + } + }); + + test("turnId boundaries break collapse", async () => { + vi.useFakeTimers(); + const harness = createHarness(); + try { + const { agentId, session } = await createManagedSession(harness); + + session.pushEvent(assistant("one", "codex", "turn-1")); + session.pushEvent(assistant("two", "codex", "turn-2")); + await waitForSessionEventQueue(); + await vi.advanceTimersByTimeAsync(33); + + const rows = await harness.manager.getTimelineRows(agentId); + const events = getTimelineStreamEvents(harness.events, agentId); + + expect(getTimelineItems(rows)).toEqual([ + { type: "assistant_message", text: "one" }, + { type: "assistant_message", text: "two" }, + ]); + expect(events.map((event) => (event.type === "agent_stream" ? event.event : null))).toEqual([ + { + type: "timeline", + provider: "codex", + turnId: "turn-1", + item: { type: "assistant_message", text: "one" }, + }, + { + type: "timeline", + provider: "codex", + turnId: "turn-2", + item: { type: "assistant_message", text: "two" }, + }, + ]); + expectContiguousRowSeqs(rows, [1, 2]); + expectContiguousLiveSeqs(events, [1, 2]); + } finally { + harness.cleanup(); + } + }); + + test("drops empty text but preserves whitespace byte-exactly", async () => { + vi.useFakeTimers(); + const harness = createHarness(); + try { + const { agentId, session } = await createManagedSession(harness); + + session.pushEvent(assistant("")); + session.pushEvent(assistant(" ")); + session.pushEvent(assistant("\n\t")); + session.pushEvent(assistant("done")); + await waitForSessionEventQueue(); + await vi.advanceTimersByTimeAsync(33); + + const rows = await harness.manager.getTimelineRows(agentId); + const events = getTimelineStreamEvents(harness.events, agentId); + + expect(getTimelineItems(rows)).toEqual([{ type: "assistant_message", text: " \n\tdone" }]); + expect(events).toHaveLength(1); + expectContiguousRowSeqs(rows, [1]); + expectContiguousLiveSeqs(events, [1]); + } finally { + harness.cleanup(); + } + }); + + test("flushes pending chunks before terminal lifecycle events by dispatch order", async () => { + vi.useFakeTimers(); + + for (const terminalType of ["turn_completed", "turn_failed", "turn_canceled"] as const) { + const harness = createHarness(); + try { + const { agentId, session } = await createManagedSession(harness); + + session.pushEvent(assistant(`${terminalType}-text`, "codex", terminalType)); + session.pushEvent(terminalEvent(terminalType, terminalType)); + await waitForSessionEventQueue(); + + const streamEvents = getStreamEvents(harness.events, agentId); + const timelineEvents = getTimelineStreamEvents(harness.events, agentId); + const rows = await harness.manager.getTimelineRows(agentId); + + expect(getTimelineItems(rows)[0]).toEqual({ + type: "assistant_message", + text: `${terminalType}-text`, + }); + const expectedTimelineEventCount = terminalType === "turn_failed" ? 2 : 1; + expect(timelineEvents).toHaveLength(expectedTimelineEventCount); + expect(streamEvents[0]).toMatchObject({ + type: "agent_stream", + event: { + type: "timeline", + item: { type: "assistant_message", text: `${terminalType}-text` }, + }, + seq: 1, + }); + const terminalEventIndex = streamEvents.findIndex( + (event) => event.type === "agent_stream" && event.event.type === terminalType, + ); + expect(terminalEventIndex).toBeGreaterThan(0); + expect(streamEvents[terminalEventIndex]).toMatchObject({ + type: "agent_stream", + event: { type: terminalType }, + }); + expect( + streamEvents[terminalEventIndex]?.type === "agent_stream" + ? streamEvents[terminalEventIndex].seq + : undefined, + ).toBe(undefined); + + await vi.advanceTimersByTimeAsync(33); + expect(await harness.manager.getTimelineRows(agentId)).toHaveLength( + expectedTimelineEventCount, + ); + expect(getTimelineStreamEvents(harness.events, agentId)).toHaveLength( + expectedTimelineEventCount, + ); + } finally { + harness.cleanup(); + } + } + }); + + test("lifecycle paths are timer-safe", async () => { + vi.useFakeTimers(); + + const closeHarness = createHarness(); + try { + const { agentId, session } = await createManagedSession(closeHarness); + session.pushEvent(assistant("close")); + await waitForSessionEventQueue(); + await closeHarness.manager.closeAgent(agentId); + await vi.advanceTimersByTimeAsync(100); + + expect(getTimelineStreamEvents(closeHarness.events, agentId)).toHaveLength(1); + expect(getTimelineStreamEvents(closeHarness.events, agentId)[0]).toMatchObject({ + type: "agent_stream", + event: { type: "timeline", item: { type: "assistant_message", text: "close" } }, + seq: 1, + }); + } finally { + closeHarness.cleanup(); + } + + const reloadHarness = createHarness(); + try { + const { agentId, session } = await createManagedSession(reloadHarness); + session.pushEvent(assistant("reload")); + await waitForSessionEventQueue(); + await reloadHarness.manager.reloadAgentSession(agentId); + await vi.advanceTimersByTimeAsync(100); + + const rows = await reloadHarness.manager.getTimelineRows(agentId); + expect(getTimelineItems(rows)).toEqual([{ type: "assistant_message", text: "reload" }]); + expect(getTimelineStreamEvents(reloadHarness.events, agentId)).toHaveLength(1); + } finally { + reloadHarness.cleanup(); + } + + const deletionHarness = createHarness(); + try { + const { agentId, session } = await createManagedSession(deletionHarness); + session.pushEvent(assistant("delete")); + await waitForSessionEventQueue(); + await deletionHarness.manager.closeAgent(agentId); + await vi.advanceTimersByTimeAsync(100); + + expect(deletionHarness.manager.getAgent(agentId)).toBeNull(); + expect(getTimelineStreamEvents(deletionHarness.events, agentId)).toHaveLength(1); + } finally { + deletionHarness.cleanup(); + } + }); + + test("old timers cannot write into reused agent ids", async () => { + vi.useFakeTimers(); + + const reloadHarness = createHarness(); + try { + const { agentId, session } = await createManagedSession(reloadHarness); + session.pushEvent(assistant("before-reload")); + await waitForSessionEventQueue(); + await reloadHarness.manager.reloadAgentSession(agentId); + await vi.advanceTimersByTimeAsync(100); + + expect(getTimelineItems(await reloadHarness.manager.getTimelineRows(agentId))).toEqual([ + { type: "assistant_message", text: "before-reload" }, + ]); + expect(getTimelineStreamEvents(reloadHarness.events, agentId)).toHaveLength(1); + } finally { + reloadHarness.cleanup(); + } + + const reuseHarness = createHarness(); + try { + const first = await createManagedSession(reuseHarness, { + agentId: AGENT_IDS[0], + workdir: reuseHarness.workdir, + }); + first.session.pushEvent(assistant("old")); + await waitForSessionEventQueue(); + await reuseHarness.manager.closeAgent(first.agentId); + + const nextWorkdir = mkdtempSync(join(tmpdir(), "agent-manager-stream-coalescing-")); + const second = await createManagedSession(reuseHarness, { + agentId: AGENT_IDS[0], + workdir: nextWorkdir, + }); + second.session.pushEvent(assistant("new")); + await waitForSessionEventQueue(); + await vi.advanceTimersByTimeAsync(100); + + expect(getTimelineItems(await reuseHarness.manager.getTimelineRows(second.agentId))).toEqual([ + { type: "assistant_message", text: "new" }, + ]); + expect(getTimelineStreamEvents(reuseHarness.events, first.agentId)).toHaveLength(2); + expect( + getTimelineStreamEvents(reuseHarness.events, first.agentId).map((event) => + event.type === "agent_stream" && event.event.type === "timeline" + ? event.event.item + : null, + ), + ).toEqual([ + { type: "assistant_message", text: "old" }, + { type: "assistant_message", text: "new" }, + ]); + rmSync(nextWorkdir, { recursive: true, force: true }); + } finally { + reuseHarness.cleanup(); + } + }); + + test("flush helper is idempotent and re-entry safe", async () => { + vi.useFakeTimers(); + const idempotencyHarness = createHarness(); + try { + const { agentId, session } = await createManagedSession(idempotencyHarness); + + session.pushEvent(assistant("a")); + session.pushEvent(assistant("b")); + session.pushEvent(assistant("c")); + session.pushEvent(timelineEvent(TOOL_CALL)); + session.pushEvent(timelineEvent(TOOL_CALL)); + await waitForSessionEventQueue(); + + const rows = await idempotencyHarness.manager.getTimelineRows(agentId); + const timelineEvents = getTimelineStreamEvents(idempotencyHarness.events, agentId); + + expect(getTimelineItems(rows)).toEqual([ + { type: "assistant_message", text: "abc" }, + TOOL_CALL, + TOOL_CALL, + ]); + expect(timelineEvents).toHaveLength(3); + expectContiguousRowSeqs(rows, [1, 2, 3]); + expectContiguousLiveSeqs(timelineEvents, [1, 2, 3]); + } finally { + idempotencyHarness.cleanup(); + } + + const reentryHarness = createHarness(); + try { + const { agentId, session } = await createManagedSession(reentryHarness); + let didPushDuringFlush = false; + reentryHarness.manager.subscribe( + (event) => { + if ( + didPushDuringFlush || + event.type !== "agent_stream" || + event.agentId !== agentId || + event.event.type !== "timeline" || + event.event.item.type !== "assistant_message" + ) { + return; + } + didPushDuringFlush = true; + session.pushEvent(assistant("B")); + }, + { replayState: false }, + ); + + session.pushEvent(assistant("A")); + await waitForSessionEventQueue(); + + await vi.advanceTimersByTimeAsync(33); + await waitForSessionEventQueue(); + + const rowsAfterFirstFlush = await reentryHarness.manager.getTimelineRows(agentId); + expect(getTimelineItems(rowsAfterFirstFlush)).toEqual([ + { type: "assistant_message", text: "A" }, + ]); + expect(getTimelineStreamEvents(reentryHarness.events, agentId)).toHaveLength(1); + + await vi.advanceTimersByTimeAsync(33); + await waitForSessionEventQueue(); + + const rowsAfterSecondFlush = await reentryHarness.manager.getTimelineRows(agentId); + const timelineEvents = getTimelineStreamEvents(reentryHarness.events, agentId); + expect(getTimelineItems(rowsAfterSecondFlush)).toEqual([ + { type: "assistant_message", text: "A" }, + { type: "assistant_message", text: "B" }, + ]); + expect(timelineEvents).toHaveLength(2); + expectContiguousRowSeqs(rowsAfterSecondFlush, [1, 2]); + expectContiguousLiveSeqs(timelineEvents, [1, 2]); + expect(didPushDuringFlush).toBe(true); + } finally { + reentryHarness.cleanup(); + } + }); + + test("manager flush drains pending chunks before background tasks", async () => { + vi.useFakeTimers(); + const harness = createHarness(); + try { + const { agentId, session } = await createManagedSession(harness); + + session.pushEvent(assistant("flush")); + await waitForSessionEventQueue(); + + expect(await harness.manager.getTimelineRows(agentId)).toHaveLength(0); + await harness.manager.flush(); + + const rows = await harness.manager.getTimelineRows(agentId); + const events = getTimelineStreamEvents(harness.events, agentId); + expect(getTimelineItems(rows)).toEqual([{ type: "assistant_message", text: "flush" }]); + expect(events).toHaveLength(1); + + await vi.advanceTimersByTimeAsync(100); + expect(await harness.manager.getTimelineRows(agentId)).toHaveLength(1); + expect(getTimelineStreamEvents(harness.events, agentId)).toHaveLength(1); + } finally { + harness.cleanup(); + } + }); + + test("history hydration is uncoalesced and immediate", async () => { + vi.useFakeTimers(); + const harness = createHarness(); + try { + const { agentId, session } = await createManagedSession(harness); + session.setHistory([assistant("a1"), assistant("a2"), reasoning("r1"), reasoning("r2")]); + + await harness.manager.hydrateTimelineFromProvider(agentId); + + const rows = await harness.manager.getTimelineRows(agentId); + expect(getTimelineItems(rows)).toEqual([ + { type: "assistant_message", text: "a1" }, + { type: "assistant_message", text: "a2" }, + { type: "reasoning", text: "r1" }, + { type: "reasoning", text: "r2" }, + ]); + expectContiguousRowSeqs(rows, [1, 2, 3, 4]); + expect(getTimelineStreamEvents(harness.events, agentId)).toHaveLength(0); + + await vi.advanceTimersByTimeAsync(100); + expect(getTimelineItems(await harness.manager.getTimelineRows(agentId))).toEqual([ + { type: "assistant_message", text: "a1" }, + { type: "assistant_message", text: "a2" }, + { type: "reasoning", text: "r1" }, + { type: "reasoning", text: "r2" }, + ]); + } finally { + harness.cleanup(); + } + }); + + test("foreground stream yields coalesced timeline events before terminal events", async () => { + vi.useFakeTimers(); + const harness = createHarness(); + try { + const { agentId, session } = await createManagedSession(harness); + const stream = harness.manager.streamAgent(agentId, "prompt"); + const firstEvent = stream.next(); + await harness.manager.waitForAgentRunStart(agentId); + + session.pushEvent(assistant("hello ", "codex", "turn-1")); + session.pushEvent(assistant("world", "codex", "turn-1")); + session.pushEvent(terminalEvent("turn_completed", "turn-1")); + await waitForSessionEventQueue(); + + await expect(firstEvent).resolves.toEqual({ + done: false, + value: { + type: "timeline", + provider: "codex", + turnId: "turn-1", + item: { type: "assistant_message", text: "hello world" }, + }, + }); + await expect(stream.next()).resolves.toEqual({ + done: false, + value: { type: "turn_completed", provider: "codex", turnId: "turn-1" }, + }); + await expect(stream.next()).resolves.toEqual({ done: true, value: undefined }); + + const rows = await harness.manager.getTimelineRows(agentId); + const streamEvents = getStreamEvents(harness.events, agentId); + expect(getTimelineItems(rows)).toEqual([{ type: "assistant_message", text: "hello world" }]); + expect(streamEvents[0]).toMatchObject({ + type: "agent_stream", + event: { + type: "timeline", + item: { type: "assistant_message", text: "hello world" }, + }, + }); + expect(streamEvents[1]).toMatchObject({ + type: "agent_stream", + event: { type: "turn_completed" }, + }); + } finally { + harness.cleanup(); + } + }); + + test("every collapsed item is one timeline row with monotonic seqs", async () => { + vi.useFakeTimers(); + const harness = createHarness(); + try { + const { agentId, session } = await createManagedSession(harness); + + session.pushEvent(assistant("a1", "codex", "turn-1")); + session.pushEvent(assistant("a2", "codex", "turn-1")); + session.pushEvent(reasoning("r1", "codex", "turn-1")); + session.pushEvent(reasoning("r2", "codex", "turn-1")); + session.pushEvent(assistant("b1", "codex", "turn-1")); + session.pushEvent(assistant("b2", "codex", "turn-1")); + session.pushEvent(terminalEvent("turn_completed", "turn-1")); + await waitForSessionEventQueue(); + + const rows = await harness.manager.getTimelineRows(agentId); + expect(getTimelineItems(rows)).toEqual([ + { type: "assistant_message", text: "a1a2" }, + { type: "reasoning", text: "r1r2" }, + { type: "assistant_message", text: "b1b2" }, + ]); + expectContiguousRowSeqs(rows, [1, 2, 3]); + const timelineEvents = getTimelineStreamEvents(harness.events, agentId); + expect(timelineEvents).toHaveLength(3); + expectContiguousLiveSeqs(timelineEvents, [1, 2, 3]); + + await vi.advanceTimersByTimeAsync(100); + expect(getTimelineItems(await harness.manager.getTimelineRows(agentId))).toEqual([ + { type: "assistant_message", text: "a1a2" }, + { type: "reasoning", text: "r1r2" }, + { type: "assistant_message", text: "b1b2" }, + ]); + expect(getTimelineStreamEvents(harness.events, agentId)).toHaveLength(3); + } finally { + harness.cleanup(); + } + }); +}); diff --git a/packages/server/src/server/agent/agent-manager.test.ts b/packages/server/src/server/agent/agent-manager.test.ts index 6904d529f..643f8920e 100644 --- a/packages/server/src/server/agent/agent-manager.test.ts +++ b/packages/server/src/server/agent/agent-manager.test.ts @@ -1426,7 +1426,7 @@ describe("AgentManager", () => { expect(fetched.rows.map((row) => row.seq)).toEqual([1, 2]); }); - test("streams assistant chunks incrementally and persists canonical chunk rows", async () => { + test("coalesces assistant chunks and persists the canonical row", async () => { const workdir = mkdtempSync(join(tmpdir(), "agent-manager-provisional-timeline-")); const storagePath = join(workdir, "agents"); const storage = new AgentStorage(storagePath, logger); @@ -1481,46 +1481,30 @@ describe("AgentManager", () => { const assistantTimelineEvents = streamEvents.filter( (event) => event.itemType === "assistant_message", ); - expect(assistantTimelineEvents).toHaveLength(2); + expect(assistantTimelineEvents).toHaveLength(1); expect(assistantTimelineEvents[0]).toMatchObject({ eventType: "timeline", itemType: "assistant_message", - text: "final ", + text: "final reply", seq: 1, epoch: expect.any(String), }); - expect(assistantTimelineEvents[1]).toMatchObject({ - eventType: "timeline", - itemType: "assistant_message", - text: "reply", - seq: 2, - epoch: expect.any(String), - }); expect(manager.getTimeline(snapshot.id)).toEqual([ { type: "assistant_message", - text: "final ", - }, - { - type: "assistant_message", - text: "reply", + text: "final reply", }, ]); const fetched = await manager.fetchTimeline(snapshot.id, { direction: "tail", limit: 0, }); - expect(fetched.rows).toHaveLength(2); + expect(fetched.rows).toHaveLength(1); expect(assistantTimelineEvents[0]?.epoch).toBe(fetched.epoch); - expect(assistantTimelineEvents[1]?.epoch).toBe(fetched.epoch); expect(fetched.rows[0]?.item).toEqual({ type: "assistant_message", - text: "final ", - }); - expect(fetched.rows[1]?.item).toEqual({ - type: "assistant_message", - text: "reply", + text: "final reply", }); }); diff --git a/packages/server/src/server/agent/agent-manager.ts b/packages/server/src/server/agent/agent-manager.ts index d0c722fb9..12cf9813c 100644 --- a/packages/server/src/server/agent/agent-manager.ts +++ b/packages/server/src/server/agent/agent-manager.ts @@ -44,6 +44,10 @@ import type { AgentTimelineRow, AgentTimelineStore, } from "./agent-timeline-store-types.js"; +import { + AGENT_STREAM_COALESCE_DEFAULT_WINDOW_MS, + AgentStreamCoalescer, +} from "./agent-stream-coalescer.js"; import { getAgentProviderDefinition } from "./provider-manifest.js"; export { AGENT_LIFECYCLE_STATUSES, type AgentLifecycleStatus }; @@ -97,6 +101,7 @@ export type AgentManagerOptions = { durableTimelineStore?: AgentTimelineStore; terminalManager?: TerminalManager | null; mcpBaseUrl?: string; + agentStreamCoalesceWindowMs?: number; logger: Logger; }; @@ -310,6 +315,7 @@ export class AgentManager { private readonly durableTimelineStore?: AgentTimelineStore; private readonly previousStatuses = new Map(); private readonly backgroundTasks = new Set>(); + private readonly agentStreamCoalescer: AgentStreamCoalescer; private mcpBaseUrl: string | null; private onAgentAttention?: AgentAttentionCallback; private logger: Logger; @@ -321,6 +327,14 @@ export class AgentManager { this.onAgentAttention = options?.onAgentAttention; this.mcpBaseUrl = options?.mcpBaseUrl ?? null; this.logger = options.logger.child({ module: "agent", component: "agent-manager" }); + this.agentStreamCoalescer = new AgentStreamCoalescer({ + windowMs: options.agentStreamCoalesceWindowMs ?? AGENT_STREAM_COALESCE_DEFAULT_WINDOW_MS, + timers: { setTimeout, clearTimeout }, + onFlush: ({ agentId, item, provider, turnId }) => { + const event = this.recordAndDispatchTimelineItem(agentId, item, provider, turnId); + this.notifyForegroundTurnWaiters(agentId, event); + }, + }); if (options?.clients) { for (const [provider, client] of Object.entries(options.clients)) { if (client) { @@ -699,6 +713,7 @@ export class AgentManager { ? await client.resumeSession(handle, normalizedConfig, launchContext) : await client.createSession(normalizedConfig, launchContext); + this.agentStreamCoalescer.flushAndDiscard(agentId); // Remove the existing agent entry before swapping sessions this.agents.delete(agentId); if (existing.unsubscribeSession) { @@ -1937,6 +1952,7 @@ export class AgentManager { agent: LiveManagedAgent, cancelReason: string, ): ManagedAgentClosed { + this.agentStreamCoalescer.flushAndDiscard(agent.id); this.agents.delete(agent.id); this.previousStatuses.delete(agent.id); if (agent.unsubscribeSession) { @@ -2018,7 +2034,11 @@ export class AgentManager { (waiter) => waiter.turnId === turnId && !waiter.settled, ); - await this.handleStreamEvent(agent, event); + const shouldNotifyWaiters = await this.handleStreamEvent(agent, event); + + if (!shouldNotifyWaiters) { + return; + } for (const waiter of matchingWaiters) { waiter.callback(event); @@ -2192,6 +2212,24 @@ export class AgentManager { } } + private notifyForegroundTurnWaiters(agentId: string, event: AgentStreamEvent): void { + const turnId = (event as { turnId?: string }).turnId; + if (turnId == null) { + return; + } + + const agent = this.agents.get(agentId); + if (!agent) { + return; + } + + for (const waiter of agent.foregroundTurnWaiters) { + if (waiter.turnId === turnId && !waiter.settled) { + waiter.callback(event); + } + } + } + private async handleStreamEvent( agent: ActiveManagedAgent, event: AgentStreamEvent, @@ -2199,17 +2237,21 @@ export class AgentManager { fromHistory?: boolean; canonicalUserMessagesById?: ReadonlyMap; }, - ): Promise { + ): Promise { const eventTurnId = (event as { turnId?: string }).turnId; const isForegroundEvent = Boolean(eventTurnId && agent.activeForegroundTurnId === eventTurnId); // Only update timestamp for live events, not history replay if (!options?.fromHistory) { this.touchUpdatedAt(agent); + if (this.agentStreamCoalescer.handle(agent.id, event)) { + return false; + } + this.agentStreamCoalescer.flushFor(agent.id); } - let timelineRow: AgentTimelineRow | null = null; let shouldDispatchEvent = true; + let shouldNotifyWaiters = true; switch (event.type) { case "thread_started": @@ -2230,36 +2272,48 @@ export class AgentManager { this.emitState(agent); break; case "timeline": - // Skip provider-replayed user_message items during history hydration. - if (options?.fromHistory && event.item.type === "user_message") { - const eventMessageId = normalizeMessageId(event.item.messageId); - if (eventMessageId) { - const canonicalText = options?.canonicalUserMessagesById?.get(eventMessageId); - if (canonicalText === event.item.text) { - break; + { + // Skip provider-replayed user_message items during history hydration. + if (options?.fromHistory && event.item.type === "user_message") { + const eventMessageId = normalizeMessageId(event.item.messageId); + if (eventMessageId) { + const canonicalText = options?.canonicalUserMessagesById?.get(eventMessageId); + if (canonicalText === event.item.text) { + shouldDispatchEvent = false; + shouldNotifyWaiters = false; + break; + } } } - } - // Suppress user_message echoes for the active foreground turn — - // these are already recorded by recordUserMessage(). - if (!options?.fromHistory && event.item.type === "user_message" && isForegroundEvent) { - const eventMessageId = normalizeMessageId(event.item.messageId); - const eventText = event.item.text; - if (eventMessageId) { + + // Suppress user_message echoes for the active foreground turn. + if (!options?.fromHistory && event.item.type === "user_message" && isForegroundEvent) { + const eventMessageId = normalizeMessageId(event.item.messageId); if ( - await this.hasCommittedUserMessageFromStores(agent.id, { + eventMessageId && + (await this.hasCommittedUserMessageFromStores(agent.id, { messageId: eventMessageId, - text: eventText, - }) + text: event.item.text, + })) ) { break; } } - } - timelineRow = this.recordTimeline(agent.id, event.item); - if (!options?.fromHistory && event.item.type === "user_message") { - agent.lastUserMessageAt = new Date(); - this.emitState(agent); + + if (options?.fromHistory) { + this.recordTimeline(agent.id, event.item); + shouldDispatchEvent = false; + shouldNotifyWaiters = false; + break; + } + + this.recordAndDispatchTimelineItem(agent.id, event.item, event.provider, event.turnId); + if (event.item.type === "user_message") { + agent.lastUserMessageAt = new Date(); + this.emitState(agent); + } + shouldDispatchEvent = false; + shouldNotifyWaiters = true; } break; case "turn_completed": @@ -2397,17 +2451,30 @@ export class AgentManager { // Skip dispatching individual stream events during history replay. if (!options?.fromHistory && shouldDispatchEvent) { - this.dispatchStream( - agent.id, - event, - timelineRow - ? { - seq: timelineRow.seq, - epoch: this.timelineStore.getEpoch(agent.id), - } - : undefined, - ); + this.dispatchStream(agent.id, event); } + + return shouldNotifyWaiters; + } + + private recordAndDispatchTimelineItem( + agentId: string, + item: AgentTimelineItem, + provider: AgentProvider, + turnId?: string, + ): AgentStreamEvent { + const row = this.recordTimeline(agentId, item); + const event: AgentStreamEvent = { + type: "timeline", + item, + provider, + ...(turnId !== undefined ? { turnId } : {}), + }; + this.dispatchStream(agentId, event, { + seq: row.seq, + epoch: this.timelineStore.getEpoch(agentId), + }); + return event; } private async appendSystemErrorTimelineMessage( @@ -2584,6 +2651,7 @@ export class AgentManager { * Used by daemon shutdown paths to avoid unhandled rejections after cleanup. */ async flush(): Promise { + this.agentStreamCoalescer.flushAll(); // Drain tasks, including tasks spawned while awaiting. while (this.backgroundTasks.size > 0) { const pending = Array.from(this.backgroundTasks); diff --git a/packages/server/src/server/agent/agent-stream-coalescer.test.ts b/packages/server/src/server/agent/agent-stream-coalescer.test.ts new file mode 100644 index 000000000..a276c373d --- /dev/null +++ b/packages/server/src/server/agent/agent-stream-coalescer.test.ts @@ -0,0 +1,539 @@ +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; + +import type { AgentProvider, AgentStreamEvent } from "./agent-sdk-types.js"; +import { + AGENT_STREAM_COALESCE_DEFAULT_WINDOW_MS, + AgentStreamCoalescer, + type AgentStreamCoalescerFlush, + type AgentStreamCoalescerTimers, +} from "./agent-stream-coalescer.js"; + +function createHarness(windowMs?: number) { + const flushes: AgentStreamCoalescerFlush[] = []; + const timers: AgentStreamCoalescerTimers = { + setTimeout, + clearTimeout, + }; + const coalescer = new AgentStreamCoalescer({ + ...(windowMs !== undefined ? { windowMs } : {}), + timers, + onFlush: (payload) => { + flushes.push(payload); + }, + }); + + return { coalescer, flushes }; +} + +function timeline( + item: Extract["item"], + options?: { + provider?: AgentProvider; + turnId?: string; + }, +): Extract { + return { + type: "timeline", + item, + provider: options?.provider ?? "codex", + ...(options?.turnId !== undefined ? { turnId: options.turnId } : {}), + }; +} + +function assistant( + text: string, + options?: { + provider?: AgentProvider; + turnId?: string; + }, +): Extract { + return timeline({ type: "assistant_message", text }, options); +} + +function reasoning( + text: string, + options?: { + provider?: AgentProvider; + turnId?: string; + }, +): Extract { + return timeline({ type: "reasoning", text }, options); +} + +describe("AgentStreamCoalescer", () => { + beforeEach(() => { + vi.useFakeTimers(); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + test("coalesces same-tick assistant chunks after the configured window", async () => { + const { coalescer, flushes } = createHarness(); + + expect(coalescer.handle("agent-1", assistant("hel"))).toBe(true); + expect(coalescer.handle("agent-1", assistant("lo"))).toBe(true); + + expect(flushes).toEqual([]); + await vi.advanceTimersByTimeAsync(32); + expect(flushes).toEqual([]); + + await vi.advanceTimersByTimeAsync(1); + expect(flushes).toEqual([ + { + agentId: "agent-1", + item: { type: "assistant_message", text: "hello" }, + provider: "codex", + }, + ]); + }); + + 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(coalescer.handle("agent-1", assistant("fast"))).toBe(true); + + await vi.advanceTimersByTimeAsync(9); + expect(flushes).toEqual([]); + + await vi.advanceTimersByTimeAsync(1); + expect(flushes).toEqual([ + { + agentId: "agent-1", + item: { type: "assistant_message", text: "fast" }, + provider: "codex", + }, + ]); + }); + + test("flushes before non-chunkable boundaries when caller invokes flushFor", () => { + const { coalescer, flushes } = createHarness(); + + expect(coalescer.handle("agent-1", assistant("before"))).toBe(true); + coalescer.flushFor("agent-1"); + + expect(flushes).toEqual([ + { + agentId: "agent-1", + item: { type: "assistant_message", text: "before" }, + provider: "codex", + }, + ]); + expect( + 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", () => { + 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" }), + timeline({ type: "compaction", status: "loading", trigger: "auto" }), + { type: "thread_started", provider: "codex", sessionId: "session-1" }, + { type: "turn_started", provider: "codex", turnId: "turn-1" }, + { type: "turn_completed", provider: "codex", turnId: "turn-1" }, + { type: "turn_failed", provider: "codex", error: "failed", turnId: "turn-1" }, + { type: "turn_canceled", provider: "codex", reason: "canceled", turnId: "turn-1" }, + { type: "usage_updated", provider: "codex", usage: { inputTokens: 1 }, turnId: "turn-1" }, + { + type: "permission_requested", + provider: "codex", + turnId: "turn-1", + request: { + id: "permission-1", + provider: "codex", + name: "shell", + kind: "tool", + }, + }, + { + type: "permission_resolved", + provider: "codex", + turnId: "turn-1", + requestId: "permission-1", + resolution: { behavior: "allow" }, + }, + { + type: "attention_required", + provider: "codex", + reason: "permission", + timestamp: "2026-04-18T00:00:00.000Z", + }, + ]; + + for (const event of nonChunkableEvents) { + expect(coalescer.handle("agent-1", event)).toBe(false); + } + expect(flushes).toEqual([]); + }); + + test("preserves kind boundaries", async () => { + const { coalescer, flushes } = createHarness(); + + coalescer.handle("agent-1", assistant("a1")); + coalescer.handle("agent-1", reasoning("r1")); + coalescer.handle("agent-1", assistant("a2")); + + await vi.advanceTimersByTimeAsync(33); + expect(flushes).toEqual([ + { + agentId: "agent-1", + item: { type: "assistant_message", text: "a1" }, + provider: "codex", + }, + { + agentId: "agent-1", + item: { type: "reasoning", text: "r1" }, + provider: "codex", + }, + { + agentId: "agent-1", + item: { type: "assistant_message", text: "a2" }, + provider: "codex", + }, + ]); + }); + + test("preserves strict alternating assistant/reasoning order", async () => { + const { coalescer, flushes } = createHarness(); + + coalescer.handle("agent-1", assistant("a1")); + coalescer.handle("agent-1", reasoning("r1")); + coalescer.handle("agent-1", assistant("a2")); + coalescer.handle("agent-1", reasoning("r2")); + + await vi.advanceTimersByTimeAsync(33); + expect(flushes.map((flush) => flush.item)).toEqual([ + { type: "assistant_message", text: "a1" }, + { type: "reasoning", text: "r1" }, + { type: "assistant_message", text: "a2" }, + { type: "reasoning", text: "r2" }, + ]); + }); + + test("does not collapse across provider boundaries", async () => { + const { coalescer, flushes } = createHarness(); + + coalescer.handle("agent-1", assistant("c", { provider: "codex" })); + coalescer.handle("agent-1", assistant("o", { provider: "opencode" })); + + await vi.advanceTimersByTimeAsync(33); + expect(flushes).toEqual([ + { + agentId: "agent-1", + item: { type: "assistant_message", text: "c" }, + provider: "codex", + }, + { + agentId: "agent-1", + item: { type: "assistant_message", text: "o" }, + provider: "opencode", + }, + ]); + }); + + test("does not collapse across turnId boundaries", async () => { + const { coalescer, flushes } = createHarness(); + + coalescer.handle("agent-1", assistant("one", { turnId: "turn-1" })); + coalescer.handle("agent-1", assistant("two", { turnId: "turn-2" })); + + await vi.advanceTimersByTimeAsync(33); + expect(flushes).toEqual([ + { + agentId: "agent-1", + item: { type: "assistant_message", text: "one" }, + provider: "codex", + turnId: "turn-1", + }, + { + agentId: "agent-1", + item: { type: "assistant_message", text: "two" }, + provider: "codex", + turnId: "turn-2", + }, + ]); + }); + + test("drops empty text chunks", async () => { + const { coalescer, flushes } = createHarness(); + + expect(coalescer.handle("agent-1", assistant(""))).toBe(true); + expect(coalescer.handle("agent-2", reasoning(""))).toBe(true); + + await vi.advanceTimersByTimeAsync(33); + expect(flushes).toEqual([]); + }); + + test("preserves whitespace byte-exactly", async () => { + const { coalescer, flushes } = createHarness(); + + coalescer.handle("agent-1", assistant(" ")); + coalescer.handle("agent-1", assistant("\n")); + coalescer.handle("agent-1", assistant("\t")); + coalescer.handle("agent-1", assistant(" mixed \n\t whitespace ")); + + await vi.advanceTimersByTimeAsync(33); + expect(flushes).toEqual([ + { + agentId: "agent-1", + item: { + type: "assistant_message", + text: " \n\t mixed \n\t whitespace ", + }, + provider: "codex", + }, + ]); + }); + + test("reconstructs bytes exactly for fragmented text", async () => { + const { coalescer, flushes } = createHarness(); + const text = "Hello, 世界.\nPunctuation: ,.!?;: — but JS strings stay intact.\n"; + const chunks = [ + "Hello", + ", ", + "世界", + ".\n", + "Punctuation: ,.!?;: ", + "—", + " but JS strings stay intact.\n", + ]; + + for (const chunk of chunks) { + coalescer.handle("agent-1", assistant(chunk)); + } + + await vi.advanceTimersByTimeAsync(33); + expect(flushes).toEqual([ + { + agentId: "agent-1", + item: { type: "assistant_message", text }, + provider: "codex", + }, + ]); + }); + + test("isolates buffers per agent", async () => { + const { coalescer, flushes } = createHarness(); + + coalescer.handle("agent-1", assistant("a")); + coalescer.handle("agent-2", assistant("x")); + coalescer.handle("agent-1", assistant("b")); + coalescer.handle("agent-2", assistant("y")); + + await vi.advanceTimersByTimeAsync(33); + expect(flushes).toEqual([ + { + agentId: "agent-1", + item: { type: "assistant_message", text: "ab" }, + provider: "codex", + }, + { + agentId: "agent-2", + item: { type: "assistant_message", text: "xy" }, + provider: "codex", + }, + ]); + }); + + test("flushAll flushes every pending agent once", () => { + const { coalescer, flushes } = createHarness(); + + coalescer.handle("agent-1", assistant("a")); + coalescer.handle("agent-2", assistant("b")); + + coalescer.flushAll(); + coalescer.flushAll(); + + expect(flushes).toEqual([ + { + agentId: "agent-1", + item: { type: "assistant_message", text: "a" }, + provider: "codex", + }, + { + agentId: "agent-2", + item: { type: "assistant_message", text: "b" }, + provider: "codex", + }, + ]); + }); + + test("flushFor is idempotent", () => { + const { coalescer, flushes } = createHarness(); + + coalescer.handle("agent-1", assistant("once")); + coalescer.flushFor("agent-1"); + coalescer.flushFor("agent-1"); + + expect(flushes).toEqual([ + { + agentId: "agent-1", + item: { type: "assistant_message", text: "once" }, + provider: "codex", + }, + ]); + }); + + test("chunks appended during onFlush use a later buffer", async () => { + const flushes: AgentStreamCoalescerFlush[] = []; + let coalescer!: AgentStreamCoalescer; + coalescer = new AgentStreamCoalescer({ + timers: { + setTimeout, + clearTimeout, + }, + onFlush: (payload) => { + flushes.push(payload); + if (payload.agentId === "agent-1" && payload.item.text === "first") { + coalescer.handle("agent-1", assistant("second")); + } + }, + }); + + coalescer.handle("agent-1", assistant("first")); + coalescer.flushFor("agent-1"); + + expect(flushes).toEqual([ + { + agentId: "agent-1", + item: { type: "assistant_message", text: "first" }, + provider: "codex", + }, + ]); + + await vi.advanceTimersByTimeAsync(33); + expect(flushes).toEqual([ + { + agentId: "agent-1", + item: { type: "assistant_message", text: "first" }, + provider: "codex", + }, + { + agentId: "agent-1", + item: { type: "assistant_message", text: "second" }, + provider: "codex", + }, + ]); + }); + + test("manual flush prevents late duplicate output after timers advance", async () => { + const { coalescer, flushes } = createHarness(); + + coalescer.handle("agent-1", assistant("manual")); + coalescer.flushFor("agent-1"); + + await vi.advanceTimersByTimeAsync(33); + expect(flushes).toEqual([ + { + agentId: "agent-1", + item: { type: "assistant_message", text: "manual" }, + provider: "codex", + }, + ]); + }); + + test("flushAndDiscard flushes, clears timer, and invalidates stale callbacks", async () => { + const { coalescer, flushes } = createHarness(); + + coalescer.handle("agent-1", assistant("durable")); + coalescer.flushAndDiscard("agent-1"); + + await vi.advanceTimersByTimeAsync(33); + expect(flushes).toEqual([ + { + agentId: "agent-1", + item: { type: "assistant_message", text: "durable" }, + provider: "codex", + }, + ]); + }); + + test("stale timers cannot flush reused agent ids", () => { + const scheduled: Array<() => void> = []; + const flushes: AgentStreamCoalescerFlush[] = []; + const setTimeoutShim: AgentStreamCoalescerTimers["setTimeout"] = (callback, delay, ...args) => { + if (typeof callback === "function") { + scheduled.push(() => { + callback(...args); + }); + } + return setTimeout(callback, delay, ...args); + }; + const coalescer = new AgentStreamCoalescer({ + timers: { + setTimeout: setTimeoutShim, + clearTimeout, + }, + onFlush: (payload) => { + flushes.push(payload); + }, + }); + + coalescer.handle("agent-1", assistant("old")); + const oldTimerCallback = scheduled[0]; + coalescer.flushAndDiscard("agent-1"); + coalescer.handle("agent-1", assistant("new")); + + oldTimerCallback(); + coalescer.flushFor("agent-1"); + + expect(flushes).toEqual([ + { + agentId: "agent-1", + item: { type: "assistant_message", text: "old" }, + provider: "codex", + }, + { + agentId: "agent-1", + item: { type: "assistant_message", text: "new" }, + provider: "codex", + }, + ]); + }); + + test("preserves first chunk item shape and replaces only text on collapse", async () => { + const { coalescer, flushes } = createHarness(); + const firstItem = { + type: "assistant_message" as const, + text: "he", + futureOptionalField: { preserved: true }, + }; + + coalescer.handle("agent-1", timeline(firstItem)); + coalescer.handle("agent-1", assistant("llo")); + + await vi.advanceTimersByTimeAsync(33); + expect(flushes).toEqual([ + { + agentId: "agent-1", + item: { + type: "assistant_message", + text: "hello", + futureOptionalField: { preserved: true }, + }, + provider: "codex", + }, + ]); + }); +}); diff --git a/packages/server/src/server/agent/agent-stream-coalescer.ts b/packages/server/src/server/agent/agent-stream-coalescer.ts new file mode 100644 index 000000000..878794a72 --- /dev/null +++ b/packages/server/src/server/agent/agent-stream-coalescer.ts @@ -0,0 +1,196 @@ +import type { AgentProvider, AgentStreamEvent, AgentTimelineItem } from "./agent-sdk-types.js"; + +export const AGENT_STREAM_COALESCE_DEFAULT_WINDOW_MS = 33; + +type ChunkableTimelineKind = "assistant_message" | "reasoning"; +type ChunkableTimelineItem = Extract; +type ChunkableTimelineEvent = Extract & { + item: ChunkableTimelineItem; +}; + +export type AgentStreamCoalescerTimers = { + setTimeout: typeof setTimeout; + clearTimeout: typeof clearTimeout; +}; + +export type AgentStreamCoalescerFlush = { + agentId: string; + item: Extract; + provider: AgentProvider; + turnId?: string; +}; + +export type AgentStreamCoalescerOptions = { + windowMs?: number; + timers: AgentStreamCoalescerTimers; + onFlush: (payload: AgentStreamCoalescerFlush) => void; +}; + +type PendingAgentStreamChunk = { + item: ChunkableTimelineItem; + text: string; + provider: AgentProvider; + turnId?: string; +}; + +type PendingAgentStreamBuffer = { + agentId: string; + chunks: PendingAgentStreamChunk[]; + timer: ReturnType | null; + flushing: boolean; +}; + +function isChunkableTimelineEvent(event: AgentStreamEvent): event is ChunkableTimelineEvent { + return ( + event.type === "timeline" && + (event.item.type === "assistant_message" || event.item.type === "reasoning") && + typeof event.item.text === "string" + ); +} + +export class AgentStreamCoalescer { + private readonly buffers = new Map(); + private readonly onFlush: (payload: AgentStreamCoalescerFlush) => void; + private readonly timers: AgentStreamCoalescerTimers; + private readonly windowMs: number; + + constructor(options: AgentStreamCoalescerOptions) { + this.windowMs = options.windowMs ?? AGENT_STREAM_COALESCE_DEFAULT_WINDOW_MS; + this.timers = options.timers; + this.onFlush = options.onFlush; + } + + handle(agentId: string, event: AgentStreamEvent): boolean { + if (!isChunkableTimelineEvent(event)) { + return false; + } + + if (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 } : {}), + }); + + if (!buffer.timer) { + this.scheduleFlush(buffer); + } + + return true; + } + + flushFor(agentId: string): void { + this.flushBuffer(agentId); + } + + flushAll(): void { + for (const agentId of Array.from(this.buffers.keys())) { + this.flushBuffer(agentId); + } + } + + flushAndDiscard(agentId: string): void { + this.flushBuffer(agentId); + const buffer = this.buffers.get(agentId); + if (buffer) { + this.clearTimer(buffer); + this.buffers.delete(agentId); + } + } + + private getOrCreateBuffer(agentId: string): PendingAgentStreamBuffer { + const existing = this.buffers.get(agentId); + if (existing) { + return existing; + } + + const buffer: PendingAgentStreamBuffer = { + agentId, + chunks: [], + timer: null, + flushing: false, + }; + this.buffers.set(agentId, buffer); + return buffer; + } + + private scheduleFlush(buffer: PendingAgentStreamBuffer): void { + const timer = this.timers.setTimeout(() => { + this.flushBuffer(buffer.agentId, buffer); + }, this.windowMs); + timer.unref?.(); + buffer.timer = timer; + } + + private clearTimer(buffer: PendingAgentStreamBuffer): void { + if (!buffer.timer) { + return; + } + this.timers.clearTimeout(buffer.timer); + buffer.timer = null; + } + + private flushBuffer(agentId: string, expectedBuffer?: PendingAgentStreamBuffer): void { + const buffer = this.buffers.get(agentId); + if (!buffer) { + return; + } + if (expectedBuffer && buffer !== expectedBuffer) { + return; + } + if (buffer.flushing) { + return; + } + + this.clearTimer(buffer); + if (buffer.chunks.length === 0) { + return; + } + + const chunks = buffer.chunks; + buffer.chunks = []; + buffer.flushing = true; + + try { + for (const chunk of this.collapseChunks(chunks)) { + this.onFlush({ + agentId, + item: { + ...chunk.item, + text: chunk.text, + }, + provider: chunk.provider, + ...(chunk.turnId !== undefined ? { turnId: chunk.turnId } : {}), + }); + } + } finally { + buffer.flushing = false; + } + } + + private collapseChunks(chunks: PendingAgentStreamChunk[]): PendingAgentStreamChunk[] { + const collapsed: PendingAgentStreamChunk[] = []; + + for (const chunk of chunks) { + const previous = collapsed.at(-1); + if ( + previous && + previous.item.type === chunk.item.type && + previous.provider === chunk.provider && + previous.turnId === chunk.turnId + ) { + previous.text += chunk.text; + continue; + } + + collapsed.push({ ...chunk }); + } + + return collapsed; + } +} diff --git a/packages/server/src/server/websocket-server.runtime-metrics.test.ts b/packages/server/src/server/websocket-server.runtime-metrics.test.ts new file mode 100644 index 000000000..40251d9df --- /dev/null +++ b/packages/server/src/server/websocket-server.runtime-metrics.test.ts @@ -0,0 +1,350 @@ +import { describe, expect, it, vi } from "vitest"; +import type pino from "pino"; +import type { SessionOutboundMessage, WSOutboundMessage } from "./messages.js"; + +const wsModuleMock = vi.hoisted(() => { + class MockWebSocketServer { + readonly handlers = new Map void>(); + + constructor(_options: unknown) {} + + on(event: string, handler: (...args: any[]) => void) { + this.handlers.set(event, handler); + return this; + } + + close() { + // no-op + } + } + + return { MockWebSocketServer }; +}); + +vi.mock("ws", () => ({ + WebSocketServer: wsModuleMock.MockWebSocketServer, +})); + +import { VoiceAssistantWebSocketServer } from "./websocket-server.js"; +import { wrapSessionMessage } from "./messages.js"; + +type RuntimeMetricsLog = { + outboundMessageTypesTop: Array<[string, number]>; + outboundSessionMessageTypesTop: Array<[string, number]>; + outboundAgentStreamTypesTop: Array<[string, number]>; + outboundAgentStreamAgentsTop: Array<[string, number]>; + outboundBinaryFrameTypesTop: Array<[string, number]>; + bufferedAmount: { + p95: number; + max: number; + }; +}; + +type TestSocket = { + readyState: number; + bufferedAmount: number; + sent: Array; + afterSendBufferedAmounts: number[]; + send: (data: string | Uint8Array | ArrayBuffer) => void; + close: () => void; + on: () => void; + once: () => void; +}; + +function createLogger() { + const logger = { + child: vi.fn(() => logger), + trace: vi.fn(), + debug: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }; + return logger; +} + +function createServer(logger: ReturnType) { + const daemonConfigStore = { + onChange: vi.fn(() => () => {}), + }; + + return new VoiceAssistantWebSocketServer( + {} as any, + logger as unknown as pino.Logger, + "srv-test", + { + setAgentAttentionCallback: vi.fn(), + getAgent: vi.fn(() => null), + getMetricsSnapshot: vi.fn(() => ({ + totalAgents: 0, + idleAgents: 0, + runningAgents: 0, + pendingPermissionAgents: 0, + erroredAgents: 0, + })), + } as any, + {} as any, + {} as any, + "/tmp/paseo-test", + daemonConfigStore as any, + null, + { allowedOrigins: new Set() }, + undefined, + undefined, + undefined, + undefined, + undefined, + "1.2.3-test", + undefined, + undefined, + undefined, + {} as any, + {} as any, + {} as any, + { + subscribe: vi.fn(), + scheduleRefreshForCwd: vi.fn(), + getMetrics: vi.fn(() => ({ + checkoutDiffTargetCount: 0, + checkoutDiffSubscriptionCount: 0, + checkoutDiffWatcherCount: 0, + checkoutDiffFallbackRefreshTargetCount: 0, + })), + dispose: vi.fn(), + } as any, + ); +} + +function createSocket(afterSendBufferedAmounts: number[]): TestSocket { + const socket: TestSocket = { + readyState: 1, + bufferedAmount: 0, + sent: [], + afterSendBufferedAmounts, + send: vi.fn((data) => { + socket.sent.push(data); + socket.bufferedAmount = afterSendBufferedAmounts.shift() ?? socket.bufferedAmount; + }), + close: vi.fn(), + on: vi.fn(), + once: vi.fn(), + }; + return socket; +} + +function flushRuntimeMetrics(server: VoiceAssistantWebSocketServer): void { + (server as any).flushRuntimeMetrics({ final: true }); +} + +function getRuntimeMetricsLog(logger: ReturnType): RuntimeMetricsLog { + const metricsCall = logger.info.mock.calls.find((call) => call[1] === "ws_runtime_metrics"); + expect(metricsCall).toBeDefined(); + return metricsCall![0] as RuntimeMetricsLog; +} + +function sendToClient( + server: VoiceAssistantWebSocketServer, + socket: TestSocket, + message: WSOutboundMessage, +) { + (server as any).sendToClient(socket, message); +} + +function sendBinaryToClient( + server: VoiceAssistantWebSocketServer, + socket: TestSocket, + frame: Uint8Array, +) { + (server as any).sendBinaryToClient(socket, frame); +} + +function attachSessionSocket(server: VoiceAssistantWebSocketServer, socket: TestSocket): void { + (server as any).sessions.set(socket, { + session: {}, + clientId: "client-1", + appVersion: null, + connectionLogger: createLogger(), + sockets: new Set([socket]), + externalDisconnectCleanupTimeout: null, + }); +} + +function agentStreamMessage(params: { + agentId: string; + event: Extract["payload"]["event"]; +}): WSOutboundMessage { + return wrapSessionMessage({ + type: "agent_stream", + payload: { + agentId: params.agentId, + event: params.event, + timestamp: "2026-04-17T00:00:00.000Z", + }, + }); +} + +describe("VoiceAssistantWebSocketServer runtime metrics", () => { + it("records outbound message type counts in the ws runtime metrics window", () => { + const logger = createLogger(); + const server = createServer(logger); + const broadcastSocket = createSocket([0]); + const directSocket = createSocket([0, 0]); + attachSessionSocket(server, broadcastSocket); + + server.broadcast( + agentStreamMessage({ + agentId: "agent-1", + event: { + type: "turn_completed", + provider: "codex", + }, + }), + ); + sendToClient( + server, + directSocket, + wrapSessionMessage({ + type: "status", + payload: { + status: "ok", + message: "ready", + }, + }), + ); + sendToClient(server, directSocket, { type: "pong" }); + + flushRuntimeMetrics(server); + + const metrics = getRuntimeMetricsLog(logger); + expect(metrics.outboundMessageTypesTop).toEqual([ + ["session_message", 2], + ["pong", 1], + ]); + expect(metrics.outboundSessionMessageTypesTop).toEqual([ + ["agent_stream", 1], + ["status", 1], + ]); + }); + + it("records agent_stream subtypes and top agents", () => { + const logger = createLogger(); + const server = createServer(logger); + const socket = createSocket([0, 0, 0, 0]); + + sendToClient( + server, + socket, + agentStreamMessage({ + agentId: "agent-1", + event: { + type: "timeline", + provider: "codex", + item: { type: "assistant_message", text: "hello" }, + }, + }), + ); + sendToClient( + server, + socket, + agentStreamMessage({ + agentId: "agent-1", + event: { + type: "timeline", + provider: "codex", + item: { type: "reasoning", text: "thinking" }, + }, + }), + ); + sendToClient( + server, + socket, + agentStreamMessage({ + agentId: "agent-1", + event: { + type: "turn_completed", + provider: "codex", + }, + }), + ); + sendToClient( + server, + socket, + agentStreamMessage({ + agentId: "agent-2", + event: { + type: "timeline", + provider: "codex", + item: { type: "assistant_message", text: "there" }, + }, + }), + ); + + flushRuntimeMetrics(server); + + const metrics = getRuntimeMetricsLog(logger); + expect(metrics.outboundAgentStreamTypesTop).toEqual([ + ["timeline:assistant_message", 2], + ["timeline:reasoning", 1], + ["turn_completed", 1], + ]); + expect(metrics.outboundAgentStreamAgentsTop).toEqual([ + ["agent-1", 3], + ["agent-2", 1], + ]); + }); + + it("records bufferedAmount p95 and max from samples taken after send", () => { + const logger = createLogger(); + const server = createServer(logger); + const broadcastSocket = createSocket([0]); + const directSocket = createSocket([10, 50]); + const binarySocket = createSocket([100]); + attachSessionSocket(server, broadcastSocket); + + server.broadcast( + agentStreamMessage({ + agentId: "agent-1", + event: { + type: "turn_completed", + provider: "codex", + }, + }), + ); + sendToClient( + server, + directSocket, + wrapSessionMessage({ + type: "status", + payload: { status: "ok" }, + }), + ); + sendToClient(server, directSocket, { type: "pong" }); + sendBinaryToClient(server, binarySocket, new Uint8Array([1, 2, 3])); + + flushRuntimeMetrics(server); + + const metrics = getRuntimeMetricsLog(logger); + expect(metrics.bufferedAmount).toEqual({ + p95: 100, + max: 100, + }); + }); + + it("counts binary frames without decoding", () => { + const logger = createLogger(); + const server = createServer(logger); + const socket = createSocket([12, 24]); + const frame = new Proxy(Buffer.from([0xff, 0xfe, 0xfd, 0x00, 0xc0]), { + get(_target, property) { + throw new Error(`Binary frame payload was unexpectedly accessed via ${String(property)}`); + }, + }) as unknown as Uint8Array; + + expect(() => sendBinaryToClient(server, socket, frame)).not.toThrow(); + + flushRuntimeMetrics(server); + + const metrics = getRuntimeMetricsLog(logger); + expect(metrics.outboundBinaryFrameTypesTop).toEqual([["binary", 1]]); + }); +}); diff --git a/packages/server/src/server/websocket-server.ts b/packages/server/src/server/websocket-server.ts index 01551a61a..a33eb8b23 100644 --- a/packages/server/src/server/websocket-server.ts +++ b/packages/server/src/server/websocket-server.ts @@ -15,6 +15,7 @@ import type { CheckoutDiffManager, CheckoutDiffMetrics } from "./checkout-diff-m import type { DaemonConfigStore, MutableDaemonConfig } from "./daemon-config-store.js"; import { type ServerInfoStatusPayload, + type SessionOutboundMessage, type WSHelloMessage, WSInboundMessageSchema, type ServerCapabilityState, @@ -346,6 +347,12 @@ export class VoiceAssistantWebSocketServer { }; private readonly inboundMessageCounts = new Map(); private readonly inboundSessionRequestCounts = new Map(); + private readonly outboundMessageCounts = new Map(); + private readonly outboundSessionMessageCounts = new Map(); + private readonly outboundAgentStreamCounts = new Map(); + private readonly outboundAgentStreamByAgentCounts = new Map(); + private readonly outboundBinaryFrameCounts = new Map(); + private readonly bufferedAmountSamples: number[] = []; private readonly requestLatencies = new Map(); private runtimeMetricsInterval: ReturnType | null = null; private unsubscribeSpeechReadiness: (() => void) | null = null; @@ -512,6 +519,7 @@ export class VoiceAssistantWebSocketServer { // WebSocket.OPEN = 1 if (ws.readyState === 1) { ws.send(payload); + this.recordOutboundMessage(message, ws); } } } @@ -623,6 +631,7 @@ export class VoiceAssistantWebSocketServer { // WebSocket.OPEN = 1 if (ws.readyState === 1) { ws.send(JSON.stringify(message)); + this.recordOutboundMessage(message, ws); } } @@ -631,6 +640,7 @@ export class VoiceAssistantWebSocketServer { return; } ws.send(frame); + this.recordOutboundBinaryFrame(ws); } private sendToConnection(connection: SessionConnection, message: WSOutboundMessage): void { @@ -1322,6 +1332,44 @@ export class VoiceAssistantWebSocketServer { this.incrementCount(this.inboundSessionRequestCounts, type); } + private recordOutboundMessage(message: WSOutboundMessage, ws: WebSocketLike): void { + if (message.type !== "session") { + this.incrementCount(this.outboundMessageCounts, message.type); + this.recordBufferedAmount(ws); + return; + } + + this.incrementCount(this.outboundMessageCounts, "session_message"); + this.incrementCount(this.outboundSessionMessageCounts, message.message.type); + + if (message.message.type === "agent_stream") { + this.recordOutboundAgentStreamMessage(message.message.payload); + } + + this.recordBufferedAmount(ws); + } + + private recordOutboundAgentStreamMessage( + payload: Extract["payload"], + ): void { + const { agentId, event } = payload; + const eventType = event.type === "timeline" ? `timeline:${event.item.type}` : event.type; + this.incrementCount(this.outboundAgentStreamCounts, eventType); + this.incrementCount(this.outboundAgentStreamByAgentCounts, agentId); + } + + private recordOutboundBinaryFrame(ws: WebSocketLike): void { + this.incrementCount(this.outboundBinaryFrameCounts, "binary"); + this.recordBufferedAmount(ws); + } + + private recordBufferedAmount(ws: WebSocketLike): void { + if (typeof ws.bufferedAmount !== "number") { + return; + } + this.bufferedAmountSamples.push(ws.bufferedAmount); + } + private recordRequestLatency(type: string, durationMs: number): void { let latencies = this.requestLatencies.get(type); if (!latencies) { @@ -1365,6 +1413,22 @@ export class VoiceAssistantWebSocketServer { return stats.slice(0, 15); } + private computeBufferedAmountStats(): { + p95: number; + max: number; + } { + if (this.bufferedAmountSamples.length === 0) { + return { p95: 0, max: 0 }; + } + + const samples = [...this.bufferedAmountSamples].sort((a, b) => a - b); + const p95Index = Math.ceil(samples.length * 0.95) - 1; + return { + p95: samples[p95Index] ?? 0, + max: samples[samples.length - 1] ?? 0, + }; + } + private collectSessionRuntimeMetrics(): WebSocketRuntimeMetrics { const uniqueConnections = new Set(this.externalSessionsByKey.values()); let terminalDirectorySubscriptionCount = 0; @@ -1402,6 +1466,7 @@ export class VoiceAssistantWebSocketServer { ).length; const sessionMetrics = this.collectSessionRuntimeMetrics(); const latencyStats = this.computeLatencyStats(); + const bufferedAmountStats = this.computeBufferedAmountStats(); const agentSnapshot = this.agentManager.getMetricsSnapshot(); this.logger.info( @@ -1420,6 +1485,12 @@ export class VoiceAssistantWebSocketServer { counters: { ...this.runtimeCounters }, inboundMessageTypesTop: this.getTopCounts(this.inboundMessageCounts, 12), inboundSessionRequestTypesTop: this.getTopCounts(this.inboundSessionRequestCounts, 20), + outboundMessageTypesTop: this.getTopCounts(this.outboundMessageCounts, 12), + outboundSessionMessageTypesTop: this.getTopCounts(this.outboundSessionMessageCounts, 20), + outboundAgentStreamTypesTop: this.getTopCounts(this.outboundAgentStreamCounts, 20), + outboundAgentStreamAgentsTop: this.getTopCounts(this.outboundAgentStreamByAgentCounts, 20), + outboundBinaryFrameTypesTop: this.getTopCounts(this.outboundBinaryFrameCounts, 12), + bufferedAmount: bufferedAmountStats, runtime: sessionMetrics, latency: latencyStats, agents: agentSnapshot, @@ -1434,6 +1505,12 @@ export class VoiceAssistantWebSocketServer { } this.inboundMessageCounts.clear(); this.inboundSessionRequestCounts.clear(); + this.outboundMessageCounts.clear(); + this.outboundSessionMessageCounts.clear(); + this.outboundAgentStreamCounts.clear(); + this.outboundAgentStreamByAgentCounts.clear(); + this.outboundBinaryFrameCounts.clear(); + this.bufferedAmountSamples.length = 0; this.requestLatencies.clear(); this.runtimeWindowStartedAt = now; }