diff --git a/src/gateway/server-chat-state.test.ts b/src/gateway/server-chat-state.test.ts index 7f00a67da16b..a98b83e60627 100644 --- a/src/gateway/server-chat-state.test.ts +++ b/src/gateway/server-chat-state.test.ts @@ -17,7 +17,6 @@ describe("createChatRunState", () => { planSnapshot: { steps: [{ step: "Inspect", status: "in_progress" }] }, bufferUpdatedAt: 1, deltaSentAt: 2, - deltaLastBroadcastLen: 9, deltaLastBroadcastText: "projected", agentText: { assistant: { lastSentAt: 3 } }, abortMarker: createChatAbortMarker(4), diff --git a/src/gateway/server-chat-state.ts b/src/gateway/server-chat-state.ts index 8ab751e3b3e6..696abff4a899 100644 --- a/src/gateway/server-chat-state.ts +++ b/src/gateway/server-chat-state.ts @@ -104,6 +104,11 @@ type ChatRunToolRecipientState = { finalizedAt?: number; }; +type PendingChatDeltaFlush = { + timer: NodeJS.Timeout; + flush: () => void; +}; + type ChatRunRecord = { registrations?: ChatRunEntry[]; rawBuffer?: string; @@ -126,6 +131,11 @@ type ChatRunRecord = { toolRecipient?: ChatRunToolRecipientState; }; +type InternalChatRunRecord = ChatRunRecord & { + /** Fixed-deadline trailing wake-up owned by this run's buffered state. */ + pendingDeltaFlush?: PendingChatDeltaFlush; +}; + type ChatRunRecordStore = { runs: Map; getOrCreate: (runId: string) => ChatRunRecord; @@ -153,6 +163,19 @@ function createChatRunRecordStore(): ChatRunRecordStore { return { runs, getOrCreate, releaseIfEmpty }; } +function internalChatRunRecord(record: ChatRunRecord): InternalChatRunRecord { + return record; +} + +function clearPendingChatDeltaFlush(record: ChatRunRecord): void { + const internal = internalChatRunRecord(record); + if (!internal.pendingDeltaFlush) { + return; + } + clearTimeout(internal.pendingDeltaFlush.timer); + delete internal.pendingDeltaFlush; +} + export type ChatRunRegistry = { add: (sessionId: string, entry: ChatRunRegistration) => void; peek: (sessionId: string) => ChatRunEntry | undefined; @@ -274,11 +297,15 @@ export function createChatRunState(): ChatRunState { delete record.deltaSentAt; delete record.deltaLastBroadcastLen; delete record.deltaLastBroadcastText; + clearPendingChatDeltaFlush(record); delete record.agentText; store.releaseIfEmpty(runId); }; const clear = () => { + for (const record of store.runs.values()) { + clearPendingChatDeltaFlush(record); + } store.runs.clear(); }; diff --git a/src/gateway/server-chat.agent-events.test.ts b/src/gateway/server-chat.agent-events.test.ts index 37511b5e0f69..f42bb64fba15 100644 --- a/src/gateway/server-chat.agent-events.test.ts +++ b/src/gateway/server-chat.agent-events.test.ts @@ -239,6 +239,13 @@ describe("agent event handler", () => { return broadcast.mock.calls.filter(([event]) => event === "chat"); } + function chatDeltaTexts(broadcast: ReturnType) { + return chatBroadcastCalls(broadcast) + .map(([, payload]) => payload as { state?: string; deltaText?: string }) + .filter((payload) => payload.state === "delta") + .map((payload) => payload.deltaText); + } + function agentBroadcastCalls(broadcast: ReturnType) { return broadcast.mock.calls.filter(([event]) => event === "agent"); } @@ -1440,6 +1447,73 @@ describe("agent event handler", () => { nowSpy.mockRestore(); }); + it("delivers a throttled delta when the window expires without another event", () => { + vi.useFakeTimers(); + let now = 12_000; + const nowSpy = vi.spyOn(Date, "now").mockImplementation(() => now); + const { broadcast, chatRunState, handler } = createHarness(); + registerNamedChatRun(chatRunState, "trailing"); + + emitAgentEvent(handler, "run-trailing", "assistant", { text: "Hello" }); + now = 12_020; + emitAgentEvent(handler, "run-trailing", "assistant", { text: "Hello world" }); + + expect(chatDeltaTexts(broadcast)).toEqual(["Hello"]); + now = 12_150; + vi.advanceTimersByTime(130); + + expect(chatDeltaTexts(broadcast)).toEqual(["Hello", " world"]); + expect(vi.getTimerCount()).toBe(0); + nowSpy.mockRestore(); + }); + + it("cancels the trailing delta before the terminal frame", () => { + vi.useFakeTimers(); + let now = 13_000; + const nowSpy = vi.spyOn(Date, "now").mockImplementation(() => now); + const { broadcast, chatRunState, handler } = createHarness(); + registerNamedChatRun(chatRunState, "terminal-trailing"); + + emitAgentEvent(handler, "run-terminal-trailing", "assistant", { text: "Hello" }); + now = 13_020; + emitAgentEvent(handler, "run-terminal-trailing", "assistant", { text: "Hello world" }); + expect(vi.getTimerCount()).toBe(1); + + emitLifecycleEnd(handler, "run-terminal-trailing"); + expect(vi.getTimerCount()).toBe(0); + const statesAtTerminal = chatBroadcastCalls(broadcast).map( + ([, payload]) => (payload as { state?: string }).state, + ); + expect(statesAtTerminal).toEqual(["delta", "delta", "final"]); + + now = 13_500; + vi.advanceTimersByTime(1_000); + expect( + chatBroadcastCalls(broadcast).map(([, payload]) => (payload as { state?: string }).state), + ).toEqual(statesAtTerminal); + nowSpy.mockRestore(); + }); + + it("cancels trailing deltas when gateway chat state is cleared", () => { + vi.useFakeTimers(); + let now = 14_000; + const nowSpy = vi.spyOn(Date, "now").mockImplementation(() => now); + const { broadcast, chatRunState, handler } = createHarness(); + registerNamedChatRun(chatRunState, "shutdown-trailing"); + + emitAgentEvent(handler, "run-shutdown-trailing", "assistant", { text: "Hello" }); + now = 14_020; + emitAgentEvent(handler, "run-shutdown-trailing", "assistant", { text: "Hello world" }); + expect(vi.getTimerCount()).toBe(1); + + chatRunState.clear(); + expect(vi.getTimerCount()).toBe(0); + now = 14_500; + vi.advanceTimersByTime(1_000); + expect(chatDeltaTexts(broadcast)).toEqual(["Hello"]); + nowSpy.mockRestore(); + }); + it("does not emit a delta when a repeated assistant snapshot is unchanged", () => { let now = 11_250; const nowSpy = vi.spyOn(Date, "now").mockImplementation(() => now); diff --git a/src/gateway/server-chat.ts b/src/gateway/server-chat.ts index b858f1713869..2053d6c06299 100644 --- a/src/gateway/server-chat.ts +++ b/src/gateway/server-chat.ts @@ -221,6 +221,7 @@ function normalizeHeartbeatChatFinalText(params: { * do not finalize a run before fallback or retry reuses the same runId. */ const AGENT_LIFECYCLE_ERROR_RETRY_GRACE_MS = 15_000; +const CHAT_DELTA_THROTTLE_MS = 150; export type ChatEventBroadcast = GatewayBroadcastFn; @@ -359,6 +360,16 @@ type AgentEventHandler = ((event: AgentEventPayload) => void) & { dispose: () => void; }; +type InternalChatRunRecord = ReturnType & { + pendingDeltaFlush?: { timer: NodeJS.Timeout; flush: () => void }; +}; + +function internalChatRunRecord( + record: ReturnType, +): InternalChatRunRecord { + return record; +} + function roundedChatSendTimingMs(value: number): number { return Math.max(0, Math.round(value * 1000) / 1000); } @@ -421,6 +432,16 @@ export function createAgentEventHandler({ const agentTextThrottleStreams = ["assistant", "thinking"] as const; + const cancelPendingChatDeltaFlush = (clientRunId: string) => { + const record = chatRunState.runs.get(clientRunId); + const pending = record ? internalChatRunRecord(record).pendingDeltaFlush : undefined; + if (!pending || !record) { + return; + } + clearTimeout(pending.timer); + delete internalChatRunRecord(record).pendingDeltaFlush; + }; + const clearBufferedChatState = (clientRunId: string) => { chatRunState.clearRun(clientRunId); }; @@ -887,6 +908,89 @@ export function createAgentEventHandler({ pendingTerminalLifecycleErrors.set(evt.runId, { timer, event: evt, opts }); }; + const broadcastChatDelta = ( + sessionKey: string, + agentId: string | undefined, + clientRunId: string, + sourceRunId: string, + seq: number, + text: string, + opts?: { controlUiVisible?: boolean; firstAssistantTimingEntry?: ChatRunEntry }, + ) => { + cancelPendingChatDeltaFlush(clientRunId); + const run = chatRunState.getOrCreate(clientRunId); + const broadcastDelta = resolveBroadcastDelta({ + text, + previousBroadcastText: run.deltaLastBroadcastText, + }); + if (!broadcastDelta) { + return; + } + const now = Date.now(); + run.deltaSentAt = now; + run.deltaLastBroadcastText = text; + const spawnedBy = resolveSpawnedBy(sessionKey); + const payload = { + runId: clientRunId, + sessionKey, + ...(agentId ? { agentId } : {}), + ...(spawnedBy && { spawnedBy }), + seq, + state: "delta" as const, + deltaText: broadcastDelta.deltaText, + ...(broadcastDelta.replace ? { replace: true as const } : {}), + message: { + role: "assistant", + content: [{ type: "text", text }], + timestamp: now, + }, + }; + emitFirstAssistantChatSendTiming( + opts?.firstAssistantTimingEntry ?? chatRunState.registry.peek(sourceRunId), + ); + sendChatPayload(sessionKey, payload, { + agentId, + controlUiVisible: opts?.controlUiVisible ?? true, + dropIfSlow: true, + }); + }; + + const scheduleChatDeltaFlush = ( + sessionKey: string, + agentId: string | undefined, + clientRunId: string, + sourceRunId: string, + seq: number, + delayMs: number, + controlUiVisible: boolean | undefined, + ) => { + const run = internalChatRunRecord(chatRunState.getOrCreate(clientRunId)); + const flush = () => { + const projected = chatRunState.resolveBuffer(clientRunId); + if (projected.suppress || shouldHideHeartbeatChatOutput(clientRunId, sourceRunId)) { + return; + } + broadcastChatDelta(sessionKey, agentId, clientRunId, sourceRunId, seq, projected.text, { + controlUiVisible, + }); + }; + const existing = run.pendingDeltaFlush; + if (existing) { + existing.flush = flush; + return; + } + const timer = setSafeTimeout(() => { + const pending = run.pendingDeltaFlush; + if (!pending || pending.timer !== timer) { + return; + } + cancelPendingChatDeltaFlush(clientRunId); + pending.flush(); + }, delayMs); + timer.unref?.(); + run.pendingDeltaFlush = { timer, flush }; + }; + const emitChatDelta = ( sessionKey: string, agentId: string | undefined, @@ -910,8 +1014,17 @@ export function createAgentEventHandler({ const now = Date.now(); run.rawBuffer = mergedRawText; run.bufferUpdatedAt = now; - const last = run.deltaSentAt ?? 0; - if (now - last < 150) { + const waitedMs = now - (run.deltaSentAt ?? 0); + if (waitedMs < CHAT_DELTA_THROTTLE_MS) { + scheduleChatDeltaFlush( + sessionKey, + agentId, + clientRunId, + sourceRunId, + seq, + CHAT_DELTA_THROTTLE_MS - waitedMs, + opts?.controlUiVisible, + ); return; } const projected = chatRunState.resolveBuffer(clientRunId); @@ -919,38 +1032,7 @@ export function createAgentEventHandler({ if (projected.suppress || shouldHideHeartbeatChatOutput(clientRunId, sourceRunId)) { return; } - const broadcastDelta = resolveBroadcastDelta({ - text: mergedText, - previousBroadcastText: run.deltaLastBroadcastText, - }); - if (!broadcastDelta) { - return; - } - run.deltaSentAt = now; - run.deltaLastBroadcastLen = mergedText.length; - run.deltaLastBroadcastText = mergedText; - const spawnedBy = resolveSpawnedBy(sessionKey); - const payload = { - runId: clientRunId, - sessionKey, - ...(agentId ? { agentId } : {}), - ...(spawnedBy && { spawnedBy }), - seq, - state: "delta" as const, - deltaText: broadcastDelta.deltaText, - ...(broadcastDelta.replace ? { replace: true as const } : {}), - message: { - role: "assistant", - content: [{ type: "text", text: mergedText }], - timestamp: now, - }, - }; - emitFirstAssistantChatSendTiming(chatRunState.registry.peek(sourceRunId)); - sendChatPayload(sessionKey, payload, { - agentId, - controlUiVisible: opts?.controlUiVisible ?? true, - dropIfSlow: true, - }); + broadcastChatDelta(sessionKey, agentId, clientRunId, sourceRunId, seq, mergedText, opts); }; const resolveBufferedChatTextState = ( @@ -981,6 +1063,7 @@ export function createAgentEventHandler({ seq: number, opts?: { controlUiVisible?: boolean; firstAssistantTimingEntry?: ChatRunEntry }, ) => { + cancelPendingChatDeltaFlush(clientRunId); const { text, shouldSuppressSilent } = resolveBufferedChatTextState(clientRunId, sourceRunId, { suppressLeadFragments: true, }); @@ -992,42 +1075,7 @@ export function createAgentEventHandler({ return; } - const now = Date.now(); - const run = chatRunState.getOrCreate(clientRunId); - const delta = resolveBroadcastDelta({ - text, - previousBroadcastText: run.deltaLastBroadcastText, - }); - if (!delta) { - return; - } - const spawnedBy = resolveSpawnedBy(sessionKey); - const flushPayload = { - runId: clientRunId, - sessionKey, - ...(agentId ? { agentId } : {}), - ...(spawnedBy && { spawnedBy }), - seq, - state: "delta" as const, - deltaText: delta.deltaText, - ...(delta.replace ? { replace: true as const } : {}), - message: { - role: "assistant", - content: [{ type: "text", text }], - timestamp: now, - }, - }; - emitFirstAssistantChatSendTiming( - opts?.firstAssistantTimingEntry ?? chatRunState.registry.peek(sourceRunId), - ); - sendChatPayload(sessionKey, flushPayload, { - agentId, - controlUiVisible: opts?.controlUiVisible ?? true, - dropIfSlow: true, - }); - run.deltaLastBroadcastLen = text.length; - run.deltaLastBroadcastText = text; - run.deltaSentAt = now; + broadcastChatDelta(sessionKey, agentId, clientRunId, sourceRunId, seq, text, opts); }; const sendChatPayload = ( diff --git a/src/gateway/server-node-events.test.ts b/src/gateway/server-node-events.test.ts index a0d960772dfa..4a9a6d1269a4 100644 --- a/src/gateway/server-node-events.test.ts +++ b/src/gateway/server-node-events.test.ts @@ -1324,8 +1324,7 @@ describe("voice transcript events", () => { const [runId, runMetadata] = mockCall(addChatRun) ?? []; expect(runId).toBe(optsRecord.runId); const clientRunId = (runMetadata as { clientRunId?: unknown } | undefined)?.clientRunId; - expect(typeof clientRunId).toBe("string"); - expect(clientRunId).toMatch(/^voice-/); + expect(clientRunId).toBe(runId); }); it("does not block agent dispatch when session-store touch fails", async () => { diff --git a/src/gateway/server-node-events.ts b/src/gateway/server-node-events.ts index fff486cb831b..b433f1cd1f06 100644 --- a/src/gateway/server-node-events.ts +++ b/src/gateway/server-node-events.ts @@ -607,11 +607,11 @@ export const handleNodeEvent = async ( isConnectionCurrent: opts?.isConnectionCurrent, }); - // Ensure chat UI clients refresh when this run completes (even though it wasn't started via chat.send). - // This maps agent bus events (keyed by per-turn runId) to chat events (keyed by clientRunId). + // Voice now has a unique per-turn run id, so it is also the stable + // client identity for chat streaming and abort lifecycle ownership. ctx.addChatRun(runId, { sessionKey: canonicalKey, - clientRunId: `voice-${randomUUID()}`, + clientRunId: runId, }); }, }); diff --git a/test/e2e/qa-lab/runtime/agent-session-streaming.e2e.test.ts b/test/e2e/qa-lab/runtime/agent-session-streaming.e2e.test.ts index b6ae8d766048..d7171651159a 100644 --- a/test/e2e/qa-lab/runtime/agent-session-streaming.e2e.test.ts +++ b/test/e2e/qa-lab/runtime/agent-session-streaming.e2e.test.ts @@ -10,7 +10,7 @@ import { const TEST_TIMEOUT_MS = 120_000; const REQUEST_TIMEOUT_MS = 20_000; -const STREAM_INTERVAL_MS = 225; +const STREAM_INTERVAL_MS = 20; const MODEL_REF = "mock-openai/gpt-5.6-luna"; const SESSION_KEY = "agent:qa:qa:session-streaming"; const IDEMPOTENCY_KEY = "qa-session-streaming"; @@ -27,6 +27,12 @@ type GatewayEvent = { event: string; payload?: unknown; }; +type ChatEventPayload = { + runId?: string; + sessionKey?: string; + state?: string; + deltaText?: string; +}; type AgentEvent = { runId?: string; sessionKey?: string; @@ -41,6 +47,14 @@ type AgentEvent = { const cleanups: Array<() => Promise> = []; +function createDeferred() { + let resolve = () => {}; + const promise = new Promise((settle) => { + resolve = settle; + }); + return { promise, resolve }; +} + afterEach(async () => { const errors: unknown[] = []; for (const cleanup of cleanups.splice(0).toReversed()) { @@ -62,7 +76,11 @@ function writeEvent(response: ServerResponse, event: unknown): void { response.write(`data: ${JSON.stringify(event)}\n\n`); } -async function writeStreamingResponse(response: ServerResponse): Promise { +async function writeStreamingResponse( + response: ServerResponse, + deltasSent: ReturnType, + terminalRelease: ReturnType, +): Promise { const message = { type: "message", id: "qa-session-streaming-message", @@ -91,6 +109,8 @@ async function writeStreamingResponse(response: ServerResponse): Promise { }); await delay(STREAM_INTERVAL_MS); } + deltasSent.resolve(); + await terminalRelease.promise; writeEvent(response, { type: "response.output_text.done", item_id: message.id, @@ -114,6 +134,13 @@ async function writeStreamingResponse(response: ServerResponse): Promise { async function startStreamingProvider() { const providerRequests: Array> = []; const transportRequests: string[] = []; + const deltasSent = createDeferred(); + const terminalRelease = createDeferred(); + let terminalReleased = false; + const releaseTerminal = () => { + terminalReleased = true; + terminalRelease.resolve(); + }; const server = createServer((request, response) => { void (async () => { if (request.method === "GET" && request.url === "/v1/models") { @@ -133,7 +160,7 @@ async function startStreamingProvider() { providerRequests.push( JSON.parse(Buffer.concat(chunks).toString("utf8")) as Record, ); - await writeStreamingResponse(response); + await writeStreamingResponse(response, deltasSent, terminalRelease); return; } transportRequests.push(`${request.method ?? "UNKNOWN"} ${request.url ?? ""}`); @@ -153,9 +180,13 @@ async function startStreamingProvider() { } return { baseUrl: `http://127.0.0.1:${address.port}`, + deltasSent: deltasSent.promise, + isTerminalReleased: () => terminalReleased, + releaseTerminal, providerRequests, transportRequests, stop: async () => { + releaseTerminal(); server.closeAllConnections(); await new Promise((resolve, reject) => { server.close((error) => (error ? reject(error) : resolve())); @@ -216,6 +247,12 @@ function asAgentEvent(event: GatewayEvent): AgentEvent | undefined { : undefined; } +function asChatEvent(event: GatewayEvent): ChatEventPayload | undefined { + return event.event === "chat" && event.payload && typeof event.payload === "object" + ? (event.payload as ChatEventPayload) + : undefined; +} + function messageRole(message: unknown): string | undefined { const role = message && typeof message === "object" ? (message as { role?: unknown }).role : null; return typeof role === "string" ? role : undefined; @@ -275,6 +312,7 @@ describe("agent session streaming", () => { const gatewayEvents: GatewayEvent[] = []; const client = await connectOperator(gateway, gatewayEvents); cleanups.push(() => client.stopAndWait({ timeoutMs: 1_000 })); + await client.request("sessions.messages.subscribe", { key: SESSION_KEY }); const accepted = await client.request("agent", { sessionKey: SESSION_KEY, message: REQUEST_MESSAGE, @@ -286,6 +324,21 @@ describe("agent session streaming", () => { runId: IDEMPOTENCY_KEY, }); + await provider.deltasSent; + await vi.waitFor( + () => { + expect(provider.isTerminalReleased()).toBe(false); + const streamedChatText = gatewayEvents + .map(asChatEvent) + .filter((event) => event?.runId === IDEMPOTENCY_KEY && event.state === "delta") + .map((event) => event?.deltaText ?? "") + .join(""); + expect(streamedChatText).toBe(TERMINAL_TEXT); + }, + { interval: 20, timeout: REQUEST_TIMEOUT_MS }, + ); + provider.releaseTerminal(); + const terminal = await client.request( "agent.wait", { runId: IDEMPOTENCY_KEY, timeoutMs: 30_000 }, @@ -345,6 +398,32 @@ describe("agent session streaming", () => { }); expect(terminalEvents[0]?.seq).toBeGreaterThan(deltaSeqs.at(-1) ?? 0); + await vi.waitFor( + () => { + const runChatEvents = gatewayEvents + .map(asChatEvent) + .filter((event) => event?.runId === IDEMPOTENCY_KEY); + expect(runChatEvents.filter((event) => event?.state === "final")).toHaveLength(1); + }, + { interval: 20, timeout: REQUEST_TIMEOUT_MS }, + ); + const runChatEventsAtTerminal = gatewayEvents + .map(asChatEvent) + .filter((event) => event?.runId === IDEMPOTENCY_KEY); + const finalIndex = runChatEventsAtTerminal.findIndex((event) => event?.state === "final"); + expect(finalIndex).toBeGreaterThan(-1); + expect( + runChatEventsAtTerminal.slice(finalIndex + 1).some((event) => event?.state === "delta"), + ).toBe(false); + await delay(250); + expect( + gatewayEvents + .map(asChatEvent) + .filter((event) => event?.runId === IDEMPOTENCY_KEY) + .slice(finalIndex + 1) + .some((event) => event?.state === "delta"), + ).toBe(false); + const streamedText = assistantDeltas.map((event) => event.data?.delta ?? "").join(""); expect(streamedText).toBe(TERMINAL_TEXT);