From db8b0c0e4a993ec7d8bd9b46f091a6bdd2b369c6 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 20 Jul 2026 15:22:35 +0200 Subject: [PATCH 01/26] feat(ai): ws frame codec for websocket transport --- packages/ai/src/stream-to-websocket.ts | 37 +++++++++++++++++++ packages/ai/tests/stream-to-websocket.test.ts | 32 ++++++++++++++++ 2 files changed, 69 insertions(+) create mode 100644 packages/ai/src/stream-to-websocket.ts create mode 100644 packages/ai/tests/stream-to-websocket.test.ts diff --git a/packages/ai/src/stream-to-websocket.ts b/packages/ai/src/stream-to-websocket.ts new file mode 100644 index 0000000000..50c9fa73a7 --- /dev/null +++ b/packages/ai/src/stream-to-websocket.ts @@ -0,0 +1,37 @@ +import type { StreamChunk } from './types' + +/** One inbound WS text frame, after JSON parse + shape discrimination. */ +export type InboundFrame = + | { kind: 'run'; input: unknown } + | { kind: 'abort'; runId: string } + +/** + * Encode one server→client frame. Durable frames carry the opaque offset in an + * `{ id, chunk }` envelope (identical to the NDJSON wire); non-durable frames + * are the bare chunk. Unambiguous because a bare chunk always has a top-level + * `type` and the envelope never does. + */ +export function encodeWsFrame( + chunk: StreamChunk, + id: string | undefined, +): string { + return JSON.stringify(id === undefined ? chunk : { id, chunk }) +} + +/** + * Decode one client→server frame. An `{ type: 'abort', runId }` object is a + * control frame; anything else is treated as a `RunAgentInput` and validated + * downstream by `chatParamsFromRequestBody`. + */ +export function decodeWsFrame(data: string): InboundFrame { + const parsed: unknown = JSON.parse(data) + if ( + typeof parsed === 'object' && + parsed !== null && + (parsed as { type?: unknown }).type === 'abort' && + typeof (parsed as { runId?: unknown }).runId === 'string' + ) { + return { kind: 'abort', runId: (parsed as { runId: string }).runId } + } + return { kind: 'run', input: parsed } +} diff --git a/packages/ai/tests/stream-to-websocket.test.ts b/packages/ai/tests/stream-to-websocket.test.ts new file mode 100644 index 0000000000..df5208ac62 --- /dev/null +++ b/packages/ai/tests/stream-to-websocket.test.ts @@ -0,0 +1,32 @@ +import { describe, expect, it } from 'vitest' +import { decodeWsFrame, encodeWsFrame } from '../src/stream-to-websocket' +import { ev } from './test-utils' + +describe('ws frame codec', () => { + it('encodes a durable frame as an { id, chunk } envelope', () => { + const chunk = ev.textContent('hi') + expect(JSON.parse(encodeWsFrame(chunk, 'off-1'))).toEqual({ + id: 'off-1', + chunk, + }) + }) + + it('encodes a non-durable frame as a bare chunk', () => { + const chunk = ev.textContent('hi') + expect(JSON.parse(encodeWsFrame(chunk, undefined))).toEqual(chunk) + }) + + it('decodes a RunAgentInput frame as a run', () => { + const input = { threadId: 't', runId: 'r', messages: [] } + expect(decodeWsFrame(JSON.stringify(input))).toEqual({ + kind: 'run', + input, + }) + }) + + it('decodes an abort control frame', () => { + expect(decodeWsFrame(JSON.stringify({ type: 'abort', runId: 'r' }))).toEqual( + { kind: 'abort', runId: 'r' }, + ) + }) +}) From 22f2d7407e6fe551ea360f66a490b67f04b9bf6c Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 20 Jul 2026 15:29:41 +0200 Subject: [PATCH 02/26] feat(ai): WebSocketLike surface + fake socket test double --- packages/ai/src/stream-to-websocket.ts | 15 +++++++ packages/ai/tests/stream-to-websocket.test.ts | 41 +++++++++++++++++++ 2 files changed, 56 insertions(+) diff --git a/packages/ai/src/stream-to-websocket.ts b/packages/ai/src/stream-to-websocket.ts index 50c9fa73a7..44cb2ef46e 100644 --- a/packages/ai/src/stream-to-websocket.ts +++ b/packages/ai/src/stream-to-websocket.ts @@ -1,5 +1,20 @@ import type { StreamChunk } from './types' +/** + * The minimal WHATWG WebSocket surface the core needs. Cloudflare + * `WebSocketPair` server sockets and Deno sockets already satisfy it; `ws` + * (Node) and Bun `ServerWebSocket` get a ~10-line adapter at the call site. + */ +export interface WebSocketLike { + send: (data: string) => void + close: (code?: number, reason?: string) => void + addEventListener: ( + type: 'message' | 'close' | 'error', + handler: (ev: any) => void, + ) => void + readonly bufferedAmount?: number +} + /** One inbound WS text frame, after JSON parse + shape discrimination. */ export type InboundFrame = | { kind: 'run'; input: unknown } diff --git a/packages/ai/tests/stream-to-websocket.test.ts b/packages/ai/tests/stream-to-websocket.test.ts index df5208ac62..a28e8b10e7 100644 --- a/packages/ai/tests/stream-to-websocket.test.ts +++ b/packages/ai/tests/stream-to-websocket.test.ts @@ -1,6 +1,7 @@ import { describe, expect, it } from 'vitest' import { decodeWsFrame, encodeWsFrame } from '../src/stream-to-websocket' import { ev } from './test-utils' +import type { WebSocketLike } from '../src/stream-to-websocket' describe('ws frame codec', () => { it('encodes a durable frame as an { id, chunk } envelope', () => { @@ -30,3 +31,43 @@ describe('ws frame codec', () => { ) }) }) + +class FakeSocket implements WebSocketLike { + sent: Array = [] + closed = false + closeCode: number | undefined + bufferedAmount = 0 + private handlers: Record void>> = {} + send(data: string): void { + this.sent.push(data) + } + close(code?: number): void { + this.closed = true + this.closeCode = code + this.emit('close', { code }) + } + addEventListener(type: string, handler: (ev: any) => void): void { + ;(this.handlers[type] ??= []).push(handler) + } + emitMessage(data: string): void { + this.emit('message', { data }) + } + emitClose(): void { + this.emit('close', {}) + } + private emit(type: string, ev: any): void { + for (const h of this.handlers[type] ?? []) h(ev) + } +} + +describe('FakeSocket double', () => { + it('records sent frames and delivers messages to listeners', () => { + const socket = new FakeSocket() + const received: Array = [] + socket.addEventListener('message', (e) => received.push(e.data)) + socket.send('a') + socket.emitMessage('b') + expect(socket.sent).toEqual(['a']) + expect(received).toEqual(['b']) + }) +}) From cfbc42ae7bd51a1acc5932e7b0f8d572fd156deb Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 20 Jul 2026 15:36:12 +0200 Subject: [PATCH 03/26] feat(ai): WsRunContext + synthetic per-turn request builder --- packages/ai/src/stream-to-websocket.ts | 34 ++++++++++++++++++- packages/ai/tests/stream-to-websocket.test.ts | 28 ++++++++++++++- 2 files changed, 60 insertions(+), 2 deletions(-) diff --git a/packages/ai/src/stream-to-websocket.ts b/packages/ai/src/stream-to-websocket.ts index 44cb2ef46e..64017dadb6 100644 --- a/packages/ai/src/stream-to-websocket.ts +++ b/packages/ai/src/stream-to-websocket.ts @@ -1,4 +1,4 @@ -import type { StreamChunk } from './types' +import type { ModelMessage, StreamChunk, UIMessage } from './types' /** * The minimal WHATWG WebSocket surface the core needs. Cloudflare @@ -50,3 +50,35 @@ export function decodeWsFrame(data: string): InboundFrame { } return { kind: 'run', input: parsed } } + +/** Per-turn context for one inbound `run` frame on a conversation-scoped socket. */ +export interface WsRunContext { + messages: Array | Array + threadId: string + runId: string + forwardedProps?: Record + /** Non-null when this socket opened as a reconnect (carried `?offset`). */ + resumeOffset: string | null + /** Synthetic per-turn request carrying `?runId=` so durability keys correctly. */ + request: Request + /** Aborts on socket close or an `abort` control frame for this run. */ + signal: AbortSignal +} + +/** + * Build the synthetic per-turn request. A conversation-scoped socket multiplexes + * many runs; each turn's durability adapter must key on the frame's `runId`, + * which we carry in the URL query (`memoryStream`/`durableStream` already read + * `?runId` / `?offset` there). Headers are copied from the handshake so + * auth/cookies survive. + */ +export function buildTurnRequest( + handshake: Request, + runId: string, + offset: string | null, +): Request { + const url = new URL(handshake.url) + url.searchParams.set('runId', runId) + if (offset !== null) url.searchParams.set('offset', offset) + return new Request(url, { headers: handshake.headers }) +} diff --git a/packages/ai/tests/stream-to-websocket.test.ts b/packages/ai/tests/stream-to-websocket.test.ts index a28e8b10e7..0419f224c9 100644 --- a/packages/ai/tests/stream-to-websocket.test.ts +++ b/packages/ai/tests/stream-to-websocket.test.ts @@ -1,5 +1,9 @@ import { describe, expect, it } from 'vitest' -import { decodeWsFrame, encodeWsFrame } from '../src/stream-to-websocket' +import { + buildTurnRequest, + decodeWsFrame, + encodeWsFrame, +} from '../src/stream-to-websocket' import { ev } from './test-utils' import type { WebSocketLike } from '../src/stream-to-websocket' @@ -71,3 +75,25 @@ describe('FakeSocket double', () => { expect(received).toEqual(['b']) }) }) + +describe('buildTurnRequest', () => { + it('keys the synthetic request by runId and preserves headers', () => { + const handshake = new Request('https://x/api/chat', { + headers: { authorization: 'Bearer t' }, + }) + const req = buildTurnRequest(handshake, 'run-9', null) + const url = new URL(req.url) + expect(url.searchParams.get('runId')).toBe('run-9') + expect(url.searchParams.get('offset')).toBeNull() + expect(req.headers.get('authorization')).toBe('Bearer t') + }) + + it('adds ?offset on a reconnect', () => { + const req = buildTurnRequest( + new Request('https://x/api/chat'), + 'run-9', + 'off-3', + ) + expect(new URL(req.url).searchParams.get('offset')).toBe('off-3') + }) +}) From 4683e312c0deb0ddf444d06653abd41189b28819 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 20 Jul 2026 15:43:34 +0200 Subject: [PATCH 04/26] feat(ai): toWebSocketStream conversation loop (non-durable happy path) --- packages/ai/src/stream-to-websocket.ts | 58 +++++++++++++++++++ packages/ai/tests/stream-to-websocket.test.ts | 46 +++++++++++++++ 2 files changed, 104 insertions(+) diff --git a/packages/ai/src/stream-to-websocket.ts b/packages/ai/src/stream-to-websocket.ts index 64017dadb6..c6326e0e6e 100644 --- a/packages/ai/src/stream-to-websocket.ts +++ b/packages/ai/src/stream-to-websocket.ts @@ -1,3 +1,6 @@ +import { chatParamsFromRequestBody } from './utilities/chat-params' +import type { StreamDurability } from './stream-durability' +import type { DebugOption } from './logger/types' import type { ModelMessage, StreamChunk, UIMessage } from './types' /** @@ -82,3 +85,58 @@ export function buildTurnRequest( if (offset !== null) url.searchParams.set('offset', offset) return new Request(url, { headers: handshake.headers }) } + +export interface WebSocketStreamInit { + /** Build a fresh chat() stream for each inbound RunAgentInput frame. */ + onRun: (ctx: WsRunContext) => AsyncIterable + /** Per-TURN durability factory, keyed by the frame's runId via ctx.request. */ + durability?: (ctx: WsRunContext) => StreamDurability + /** Chunks buffered per durability append (default 32). */ + batch?: number + /** Heartbeat ping interval in ms (default 30_000). */ + heartbeatMs?: number + /** Close after this many ms without any inbound frame (default 300_000). */ + idleTimeoutMs?: number + debug?: DebugOption +} + +/** + * Run a full-duplex, conversation-scoped chat over an already-accepted server + * socket. Each inbound RunAgentInput frame starts one chat() turn (via onRun) + * whose chunks are pumped back as frames; the socket stays open across turns + * (pending client-tool resubmit, next user message) until close/abort/idle. + */ +export function toWebSocketStream( + socket: WebSocketLike, + request: Request, + init: WebSocketStreamInit, +): void { + socket.addEventListener('message', (event: { data: unknown }) => { + if (typeof event.data !== 'string') return + void handleInbound(String(event.data)) + }) + + async function handleInbound(data: string): Promise { + const frame = decodeWsFrame(data) + if (frame.kind !== 'run') return // abort handled in a later task + const params = await chatParamsFromRequestBody(frame.input) + const turnAbort = new AbortController() + const ctx: WsRunContext = { + // chatParamsFromRequestBody returns a single array that may mix + // UIMessage/ModelMessage per-element; WsRunContext expects one + // homogeneous array type. The runtime shape is already consistent + // (all elements come from the same parsed request), so this is a + // safe narrowing, matching the existing `chat/index.ts` precedent. + messages: params.messages as Array | Array, + threadId: params.threadId, + runId: params.runId, + forwardedProps: params.forwardedProps, + resumeOffset: null, + request: buildTurnRequest(request, params.runId, null), + signal: turnAbort.signal, + } + for await (const chunk of init.onRun(ctx)) { + socket.send(encodeWsFrame(chunk, undefined)) + } + } +} diff --git a/packages/ai/tests/stream-to-websocket.test.ts b/packages/ai/tests/stream-to-websocket.test.ts index 0419f224c9..47c20794eb 100644 --- a/packages/ai/tests/stream-to-websocket.test.ts +++ b/packages/ai/tests/stream-to-websocket.test.ts @@ -3,9 +3,11 @@ import { buildTurnRequest, decodeWsFrame, encodeWsFrame, + toWebSocketStream, } from '../src/stream-to-websocket' import { ev } from './test-utils' import type { WebSocketLike } from '../src/stream-to-websocket' +import type { StreamChunk } from '../src/types' describe('ws frame codec', () => { it('encodes a durable frame as an { id, chunk } envelope', () => { @@ -97,3 +99,47 @@ describe('buildTurnRequest', () => { expect(new URL(req.url).searchParams.get('offset')).toBe('off-3') }) }) + +function inputFrame(runId: string): string { + return JSON.stringify({ + threadId: 'thread-1', + runId, + messages: [{ id: 'u1', role: 'user', content: 'hi' }], + tools: [], + context: [], + forwardedProps: {}, + state: {}, + }) +} + +async function flush(): Promise { + await new Promise((r) => setTimeout(r, 0)) + await new Promise((r) => setTimeout(r, 0)) +} + +describe('toWebSocketStream (non-durable)', () => { + it('pumps onRun chunks as bare frames and keeps the socket open', async () => { + const socket = new FakeSocket() + toWebSocketStream(socket, new Request('https://x/api/chat'), { + onRun: ({ runId, threadId }): AsyncIterable => + (async function* () { + yield ev.runStarted(runId, threadId) + yield ev.textContent('a') + yield { + type: 'RUN_FINISHED', + runId, + threadId, + model: 'm', + finishReason: 'stop', + timestamp: Date.now(), + } as StreamChunk + })(), + }) + socket.emitMessage(inputFrame('run-1')) + await flush() + + const types = socket.sent.map((s) => JSON.parse(s).type) + expect(types).toEqual(['RUN_STARTED', 'TEXT_MESSAGE_CONTENT', 'RUN_FINISHED']) + expect(socket.closed).toBe(false) // conversation-scoped: stays open + }) +}) From ee44d8f7c4a8ab5c3345a5b525457dac0093b864 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 20 Jul 2026 15:52:24 +0200 Subject: [PATCH 05/26] fix(ai): widen WsRunContext.messages to drop unsound cast --- packages/ai/src/stream-to-websocket.ts | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/packages/ai/src/stream-to-websocket.ts b/packages/ai/src/stream-to-websocket.ts index c6326e0e6e..04b30c7f94 100644 --- a/packages/ai/src/stream-to-websocket.ts +++ b/packages/ai/src/stream-to-websocket.ts @@ -56,7 +56,7 @@ export function decodeWsFrame(data: string): InboundFrame { /** Per-turn context for one inbound `run` frame on a conversation-scoped socket. */ export interface WsRunContext { - messages: Array | Array + messages: Array threadId: string runId: string forwardedProps?: Record @@ -122,12 +122,7 @@ export function toWebSocketStream( const params = await chatParamsFromRequestBody(frame.input) const turnAbort = new AbortController() const ctx: WsRunContext = { - // chatParamsFromRequestBody returns a single array that may mix - // UIMessage/ModelMessage per-element; WsRunContext expects one - // homogeneous array type. The runtime shape is already consistent - // (all elements come from the same parsed request), so this is a - // safe narrowing, matching the existing `chat/index.ts` precedent. - messages: params.messages as Array | Array, + messages: params.messages, threadId: params.threadId, runId: params.runId, forwardedProps: params.forwardedProps, From df45ae62f1c4e9c5bd790953624a8d33ede390f2 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 20 Jul 2026 15:56:28 +0200 Subject: [PATCH 06/26] feat(ai): durable ws pumping via shared durableStreamSource --- packages/ai/src/stream-to-response.ts | 2 +- packages/ai/src/stream-to-websocket.ts | 18 ++++++++-- packages/ai/tests/stream-to-websocket.test.ts | 33 +++++++++++++++++++ 3 files changed, 50 insertions(+), 3 deletions(-) diff --git a/packages/ai/src/stream-to-response.ts b/packages/ai/src/stream-to-response.ts index 09fda5b98d..4f4674a948 100644 --- a/packages/ai/src/stream-to-response.ts +++ b/packages/ai/src/stream-to-response.ts @@ -263,7 +263,7 @@ function isDurabilityFlushBoundary(chunk: StreamChunk): boolean { * The returned `getId` maps each forwarded chunk to the exact opaque offset * returned by the durability adapter for the SSE `id:` line. */ -function durableStreamSource( +export function durableStreamSource( stream: AsyncIterable, durability: StreamDurability, options: { diff --git a/packages/ai/src/stream-to-websocket.ts b/packages/ai/src/stream-to-websocket.ts index 04b30c7f94..9f0cb55680 100644 --- a/packages/ai/src/stream-to-websocket.ts +++ b/packages/ai/src/stream-to-websocket.ts @@ -1,4 +1,6 @@ import { chatParamsFromRequestBody } from './utilities/chat-params' +import { durableStreamSource } from './stream-to-response' +import { resolveDebugOption } from './logger/resolve' import type { StreamDurability } from './stream-durability' import type { DebugOption } from './logger/types' import type { ModelMessage, StreamChunk, UIMessage } from './types' @@ -130,8 +132,20 @@ export function toWebSocketStream( request: buildTurnRequest(request, params.runId, null), signal: turnAbort.signal, } - for await (const chunk of init.onRun(ctx)) { - socket.send(encodeWsFrame(chunk, undefined)) + if (init.durability) { + const adapter = init.durability(ctx) + const { source, getId } = durableStreamSource(init.onRun(ctx), adapter, { + abortController: turnAbort, + ...(init.batch === undefined ? {} : { batch: init.batch }), + logger: resolveDebugOption(init.debug), + }) + for await (const chunk of source) { + socket.send(encodeWsFrame(chunk, getId(chunk))) + } + } else { + for await (const chunk of init.onRun(ctx)) { + socket.send(encodeWsFrame(chunk, undefined)) + } } } } diff --git a/packages/ai/tests/stream-to-websocket.test.ts b/packages/ai/tests/stream-to-websocket.test.ts index 47c20794eb..f637193317 100644 --- a/packages/ai/tests/stream-to-websocket.test.ts +++ b/packages/ai/tests/stream-to-websocket.test.ts @@ -5,6 +5,7 @@ import { encodeWsFrame, toWebSocketStream, } from '../src/stream-to-websocket' +import { memoryStream } from '../src/stream-durability' import { ev } from './test-utils' import type { WebSocketLike } from '../src/stream-to-websocket' import type { StreamChunk } from '../src/types' @@ -143,3 +144,35 @@ describe('toWebSocketStream (non-durable)', () => { expect(socket.closed).toBe(false) // conversation-scoped: stays open }) }) + +describe('toWebSocketStream (durable)', () => { + it('tags each frame with an { id, chunk } envelope from the durability log', async () => { + const socket = new FakeSocket() + toWebSocketStream(socket, new Request('https://x/api/chat'), { + durability: (ctx) => memoryStream(ctx.request), + onRun: ({ runId, threadId }): AsyncIterable => + (async function* () { + yield ev.textContent('a') + yield { + type: 'RUN_FINISHED', + runId, + threadId, + model: 'm', + finishReason: 'stop', + timestamp: Date.now(), + } as StreamChunk + })(), + }) + socket.emitMessage(inputFrame('run-2')) + await flush() + + const frames = socket.sent.map((s) => JSON.parse(s)) + expect(frames.every((f) => typeof f.id === 'string' && 'chunk' in f)).toBe( + true, + ) + expect(frames.map((f) => f.chunk.type)).toEqual([ + 'TEXT_MESSAGE_CONTENT', + 'RUN_FINISHED', + ]) + }) +}) From 18329129e843ba9eaf08a3dbb06d6d0281e4a91a Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 20 Jul 2026 16:04:11 +0200 Subject: [PATCH 07/26] feat(ai): ws abort frame + close/idle/heartbeat lifecycle + malformed-frame guard --- packages/ai/src/stream-to-websocket.ts | 86 +++++++++++++++---- packages/ai/tests/stream-to-websocket.test.ts | 70 +++++++++++++++ 2 files changed, 138 insertions(+), 18 deletions(-) diff --git a/packages/ai/src/stream-to-websocket.ts b/packages/ai/src/stream-to-websocket.ts index 9f0cb55680..f6c9170bb3 100644 --- a/packages/ai/src/stream-to-websocket.ts +++ b/packages/ai/src/stream-to-websocket.ts @@ -113,16 +113,62 @@ export function toWebSocketStream( request: Request, init: WebSocketStreamInit, ): void { + const logger = resolveDebugOption(init.debug) + const activeTurns = new Map() + const heartbeatMs = init.heartbeatMs ?? 30_000 + const idleTimeoutMs = init.idleTimeoutMs ?? 300_000 + let lastActivity = Date.now() + + const heartbeat = setInterval(() => { + socket.send(JSON.stringify({ type: 'ping' })) + }, heartbeatMs) + const idle = setInterval( + () => { + if (Date.now() - lastActivity > idleTimeoutMs) socket.close(1000, 'idle') + }, + Math.min(idleTimeoutMs, 30_000), + ) + + socket.addEventListener('close', () => { + for (const controller of activeTurns.values()) controller.abort() + activeTurns.clear() + clearInterval(heartbeat) + clearInterval(idle) + }) + socket.addEventListener('message', (event: { data: unknown }) => { if (typeof event.data !== 'string') return - void handleInbound(String(event.data)) + lastActivity = Date.now() + + // Inbound frames are client-controlled: a malformed frame (bad JSON, or + // valid JSON that isn't an AG-UI RunAgentInput/abort shape) must be + // dropped, not crash the socket or leak an unhandled rejection. + let frame: InboundFrame + try { + frame = decodeWsFrame(event.data) + } catch (error) { + logger.errors('Failed to decode inbound WS frame; dropping it', { + error, + }) + return + } + + if (frame.kind === 'abort') { + activeTurns.get(frame.runId)?.abort() + return + } + + handleInbound(frame.input).catch((error: unknown) => { + logger.errors('Failed to handle inbound WS run frame; dropping it', { + error, + }) + }) }) - async function handleInbound(data: string): Promise { - const frame = decodeWsFrame(data) - if (frame.kind !== 'run') return // abort handled in a later task - const params = await chatParamsFromRequestBody(frame.input) + async function handleInbound(input: unknown): Promise { + const params = await chatParamsFromRequestBody(input) const turnAbort = new AbortController() + activeTurns.set(params.runId, turnAbort) const ctx: WsRunContext = { messages: params.messages, threadId: params.threadId, @@ -132,20 +178,24 @@ export function toWebSocketStream( request: buildTurnRequest(request, params.runId, null), signal: turnAbort.signal, } - if (init.durability) { - const adapter = init.durability(ctx) - const { source, getId } = durableStreamSource(init.onRun(ctx), adapter, { - abortController: turnAbort, - ...(init.batch === undefined ? {} : { batch: init.batch }), - logger: resolveDebugOption(init.debug), - }) - for await (const chunk of source) { - socket.send(encodeWsFrame(chunk, getId(chunk))) - } - } else { - for await (const chunk of init.onRun(ctx)) { - socket.send(encodeWsFrame(chunk, undefined)) + try { + if (init.durability) { + const adapter = init.durability(ctx) + const { source, getId } = durableStreamSource(init.onRun(ctx), adapter, { + abortController: turnAbort, + ...(init.batch === undefined ? {} : { batch: init.batch }), + logger, + }) + for await (const chunk of source) { + socket.send(encodeWsFrame(chunk, getId(chunk))) + } + } else { + for await (const chunk of init.onRun(ctx)) { + socket.send(encodeWsFrame(chunk, undefined)) + } } + } finally { + activeTurns.delete(params.runId) } } } diff --git a/packages/ai/tests/stream-to-websocket.test.ts b/packages/ai/tests/stream-to-websocket.test.ts index f637193317..6a41a56b76 100644 --- a/packages/ai/tests/stream-to-websocket.test.ts +++ b/packages/ai/tests/stream-to-websocket.test.ts @@ -176,3 +176,73 @@ describe('toWebSocketStream (durable)', () => { ]) }) }) + +describe('toWebSocketStream lifecycle', () => { + it('aborts the matching turn on an abort control frame', async () => { + const socket = new FakeSocket() + let aborted = false + toWebSocketStream(socket, new Request('https://x/api/chat'), { + onRun: ({ runId, threadId, signal }): AsyncIterable => + (async function* () { + yield ev.runStarted(runId, threadId) + await new Promise((resolve) => { + signal.addEventListener('abort', () => { + aborted = true + resolve() + }) + }) + })(), + }) + socket.emitMessage(inputFrame('run-3')) + await flush() + socket.emitMessage(JSON.stringify({ type: 'abort', runId: 'run-3' })) + await flush() + expect(aborted).toBe(true) + }) + + it('aborts the active turn when the socket closes', async () => { + const socket = new FakeSocket() + let aborted = false + toWebSocketStream(socket, new Request('https://x/api/chat'), { + onRun: ({ signal }): AsyncIterable => + // eslint-disable-next-line require-yield + (async function* () { + await new Promise((resolve) => { + signal.addEventListener('abort', () => { + aborted = true + resolve() + }) + }) + })(), + }) + socket.emitMessage(inputFrame('run-4')) + await flush() + socket.emitClose() + await flush() + expect(aborted).toBe(true) + }) + + it('drops a malformed inbound frame without crashing the socket, and still processes a subsequent valid frame', async () => { + const socket = new FakeSocket() + toWebSocketStream(socket, new Request('https://x/api/chat'), { + onRun: ({ runId, threadId }): AsyncIterable => + (async function* () { + yield ev.runStarted(runId, threadId) + })(), + debug: false, + }) + + expect(() => socket.emitMessage('{')).not.toThrow() + await flush() + expect(() => + socket.emitMessage(JSON.stringify({ foo: 'bar' })), + ).not.toThrow() + await flush() + expect(socket.closed).toBe(false) + + socket.emitMessage(inputFrame('run-5')) + await flush() + const types = socket.sent.map((s) => JSON.parse(s).type) + expect(types).toEqual(['RUN_STARTED']) + }) +}) From 9d9196f7f46c94e17df515e8d180d828e29924e4 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 20 Jul 2026 16:12:32 +0200 Subject: [PATCH 08/26] feat(ai): resumeWebSocketStream read-only log replay --- packages/ai/src/stream-to-websocket.ts | 54 ++++++++++++++++++- packages/ai/tests/stream-to-websocket.test.ts | 44 +++++++++++++++ 2 files changed, 97 insertions(+), 1 deletion(-) diff --git a/packages/ai/src/stream-to-websocket.ts b/packages/ai/src/stream-to-websocket.ts index f6c9170bb3..260633fb20 100644 --- a/packages/ai/src/stream-to-websocket.ts +++ b/packages/ai/src/stream-to-websocket.ts @@ -195,7 +195,59 @@ export function toWebSocketStream( } } } finally { - activeTurns.delete(params.runId) + // Only delete if this turn still owns the entry: a duplicate in-flight + // runId (e.g. a client resubmitting before the first turn finished) + // would otherwise let the OLDER turn's cleanup delete the NEWER turn's + // still-active controller (TOCTOU). + if (activeTurns.get(params.runId) === turnAbort) { + activeTurns.delete(params.runId) + } } } } + +/** + * A resume is served entirely from the durability log, so there is no + * producer to iterate. This empty source satisfies `durableStreamSource`'s + * signature; on a resume it replays from the log and never touches this. + * Mirrors the private helper of the same name in `stream-to-response.ts`. + */ +function emptyDurableSource(): AsyncIterable { + return (async function* () {})() +} + +/** + * Read-only replay of a run's durability log over a socket (mirrors + * `resumeServerSentEventsResponse`). The adapter captures the offset from the + * request (`?offset`/`Last-Event-ID`); no model runs. Closes 1008 when there + * is nothing to resume. + */ +export function resumeWebSocketStream( + socket: WebSocketLike, + options: { + adapter: StreamDurability + batch?: number + debug?: DebugOption + }, +): void { + if (options.adapter.resumeFrom() === null) { + socket.close(1008, 'no resume offset') + return + } + const abortController = new AbortController() + socket.addEventListener('close', () => abortController.abort()) + const { source, getId } = durableStreamSource( + emptyDurableSource(), + options.adapter, + { + abortController, + ...(options.batch === undefined ? {} : { batch: options.batch }), + logger: resolveDebugOption(options.debug), + }, + ) + void (async () => { + for await (const chunk of source) { + socket.send(encodeWsFrame(chunk, getId(chunk))) + } + })() +} diff --git a/packages/ai/tests/stream-to-websocket.test.ts b/packages/ai/tests/stream-to-websocket.test.ts index 6a41a56b76..3b7f34ad08 100644 --- a/packages/ai/tests/stream-to-websocket.test.ts +++ b/packages/ai/tests/stream-to-websocket.test.ts @@ -3,6 +3,7 @@ import { buildTurnRequest, decodeWsFrame, encodeWsFrame, + resumeWebSocketStream, toWebSocketStream, } from '../src/stream-to-websocket' import { memoryStream } from '../src/stream-durability' @@ -246,3 +247,46 @@ describe('toWebSocketStream lifecycle', () => { expect(types).toEqual(['RUN_STARTED']) }) }) + +describe('resumeWebSocketStream', () => { + it('replays a completed run from the log without running a model', async () => { + // Produce a run into the shared memory log first. + const producer = new FakeSocket() + toWebSocketStream(producer, new Request('https://x/api/chat'), { + durability: (ctx) => memoryStream(ctx.request), + onRun: ({ runId, threadId }): AsyncIterable => + (async function* () { + yield ev.textContent('a') + yield { + type: 'RUN_FINISHED', + runId, + threadId, + model: 'm', + finishReason: 'stop', + timestamp: Date.now(), + } as StreamChunk + })(), + }) + producer.emitMessage(inputFrame('run-join')) + await flush() + + // Join from the start via a read-only replay socket. + const joiner = new FakeSocket() + const joinReq = new Request('https://x/api/chat?runId=run-join&offset=-1') + resumeWebSocketStream(joiner, { adapter: memoryStream(joinReq) }) + await flush() + + const chunkTypes = joiner.sent.map((s) => JSON.parse(s).chunk.type) + expect(chunkTypes).toEqual(['TEXT_MESSAGE_CONTENT', 'RUN_FINISHED']) + }) + + it('closes with 1008 when no offset is present', async () => { + const joiner = new FakeSocket() + resumeWebSocketStream(joiner, { + adapter: memoryStream(new Request('https://x/api/chat')), + }) + await flush() + expect(joiner.closed).toBe(true) + expect(joiner.closeCode).toBe(1008) + }) +}) From 6d51a4d6b8438aa4f6dae9c173910bf89f06702f Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 20 Jul 2026 16:18:56 +0200 Subject: [PATCH 09/26] fix(ai): guard resumeWebSocketStream replay pump against unhandled rejection --- packages/ai/src/stream-to-websocket.ts | 12 ++++++++-- packages/ai/tests/stream-to-websocket.test.ts | 22 +++++++++++++++++++ 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/packages/ai/src/stream-to-websocket.ts b/packages/ai/src/stream-to-websocket.ts index 260633fb20..e6aa5e752c 100644 --- a/packages/ai/src/stream-to-websocket.ts +++ b/packages/ai/src/stream-to-websocket.ts @@ -230,6 +230,7 @@ export function resumeWebSocketStream( debug?: DebugOption }, ): void { + const logger = resolveDebugOption(options.debug) if (options.adapter.resumeFrom() === null) { socket.close(1008, 'no resume offset') return @@ -242,12 +243,19 @@ export function resumeWebSocketStream( { abortController, ...(options.batch === undefined ? {} : { batch: options.batch }), - logger: resolveDebugOption(options.debug), + logger, }, ) void (async () => { for await (const chunk of source) { socket.send(encodeWsFrame(chunk, getId(chunk))) } - })() + })().catch((error: unknown) => { + logger.errors('resume websocket replay failed', { error }) + try { + socket.close(1011, 'resume failed') + } catch { + // socket already closing/closed — nothing to do + } + }) } diff --git a/packages/ai/tests/stream-to-websocket.test.ts b/packages/ai/tests/stream-to-websocket.test.ts index 3b7f34ad08..e06994caf0 100644 --- a/packages/ai/tests/stream-to-websocket.test.ts +++ b/packages/ai/tests/stream-to-websocket.test.ts @@ -9,6 +9,7 @@ import { import { memoryStream } from '../src/stream-durability' import { ev } from './test-utils' import type { WebSocketLike } from '../src/stream-to-websocket' +import type { StreamDurability } from '../src/stream-durability' import type { StreamChunk } from '../src/types' describe('ws frame codec', () => { @@ -289,4 +290,25 @@ describe('resumeWebSocketStream', () => { expect(joiner.closed).toBe(true) expect(joiner.closeCode).toBe(1008) }) + + it('closes with 1011 instead of leaking an unhandled rejection when the replay log throws', async () => { + const failingAdapter: StreamDurability = { + resumeFrom: () => 'off-1', + append: () => Promise.resolve([]), + // eslint-disable-next-line require-yield + read: async function* () { + throw new Error('boom') + }, + close: () => Promise.resolve(), + } + const joiner = new FakeSocket() + + expect(() => + resumeWebSocketStream(joiner, { adapter: failingAdapter, debug: false }), + ).not.toThrow() + await flush() + + expect(joiner.closed).toBe(true) + expect(joiner.closeCode).toBe(1011) + }) }) From 3cbd47d7e7aec5c5f0fbb96044aa6ec66168019e Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 20 Jul 2026 16:23:29 +0200 Subject: [PATCH 10/26] feat(ai): toWebSocketResponse/resumeWebSocketResponse CF wrappers + exports --- packages/ai/src/index.ts | 16 +++++ packages/ai/src/stream-to-websocket.ts | 70 +++++++++++++++++++ packages/ai/tests/stream-to-websocket.test.ts | 11 +++ 3 files changed, 97 insertions(+) diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 6c3cf5a277..6c0e8b2204 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -85,6 +85,22 @@ export { export { memoryStream } from './stream-durability' export type { MemoryStreamOptions, StreamDurability } from './stream-durability' +// WebSocket transport utilities +export { + toWebSocketStream, + toWebSocketResponse, + resumeWebSocketStream, + resumeWebSocketResponse, + encodeWsFrame, + decodeWsFrame, +} from './stream-to-websocket' +export type { + WebSocketLike, + WsRunContext, + WebSocketStreamInit, + InboundFrame, +} from './stream-to-websocket' + // Tool call management export { ToolCallManager } from './activities/chat/tools/tool-calls' diff --git a/packages/ai/src/stream-to-websocket.ts b/packages/ai/src/stream-to-websocket.ts index e6aa5e752c..5017e76b8d 100644 --- a/packages/ai/src/stream-to-websocket.ts +++ b/packages/ai/src/stream-to-websocket.ts @@ -259,3 +259,73 @@ export function resumeWebSocketStream( } }) } + +interface WebSocketPairCtor { + new (): { 0: unknown; 1: WebSocketLike & { accept?: () => void } } +} + +function upgradeOrThrow(helper: string): { + client: unknown + server: WebSocketLike +} { + const Pair = (globalThis as { WebSocketPair?: WebSocketPairCtor }) + .WebSocketPair + if (!Pair) { + throw new Error( + `${helper} requires a runtime with WebSocketPair (Cloudflare Workers/Durable Objects). ` + + `On other runtimes upgrade the socket yourself and call ${helper.replace('Response', 'Stream')}.`, + ) + } + const pair = new Pair() + const server = pair[1] + server.accept?.() + return { client: pair[0], server } +} + +function upgradeResponse(client: unknown): Response { + return new Response(null, { + status: 101, + // Cloudflare-specific field; typed loosely to avoid a DOM lib dependency. + webSocket: client, + } as ResponseInit & { webSocket: unknown }) +} + +/** + * Cloudflare/Deno wrapper: creates a `WebSocketPair`, accepts the server + * socket, delegates to {@link toWebSocketStream}, and returns the 101 + * upgrade `Response` carrying the client socket. Throws when the runtime has + * no `WebSocketPair` (e.g. Node) — upgrade the socket yourself and call + * {@link toWebSocketStream} directly there. + */ +export function toWebSocketResponse( + request: Request, + init: WebSocketStreamInit, +): Response { + const { client, server } = upgradeOrThrow('toWebSocketResponse') + toWebSocketStream(server, request, init) + return upgradeResponse(client) +} + +/** + * Cloudflare/Deno wrapper: creates a `WebSocketPair`, accepts the server + * socket, delegates to {@link resumeWebSocketStream}, and returns the 101 + * upgrade `Response` carrying the client socket. Throws when the runtime has + * no `WebSocketPair` (e.g. Node) — upgrade the socket yourself and call + * {@link resumeWebSocketStream} directly there. + */ +export function resumeWebSocketResponse( + // Unused: kept so the signature mirrors `toWebSocketResponse` (both take + // the upgrade `Request` first) even though the resume path only needs the + // pre-built adapter — the caller already captured the request when + // constructing it (e.g. `memoryStream(request)`). + _request: Request, + options: { + adapter: StreamDurability + batch?: number + debug?: DebugOption + }, +): Response { + const { client, server } = upgradeOrThrow('resumeWebSocketResponse') + resumeWebSocketStream(server, options) + return upgradeResponse(client) +} diff --git a/packages/ai/tests/stream-to-websocket.test.ts b/packages/ai/tests/stream-to-websocket.test.ts index e06994caf0..dcf842750e 100644 --- a/packages/ai/tests/stream-to-websocket.test.ts +++ b/packages/ai/tests/stream-to-websocket.test.ts @@ -4,6 +4,7 @@ import { decodeWsFrame, encodeWsFrame, resumeWebSocketStream, + toWebSocketResponse, toWebSocketStream, } from '../src/stream-to-websocket' import { memoryStream } from '../src/stream-durability' @@ -312,3 +313,13 @@ describe('resumeWebSocketStream', () => { expect(joiner.closeCode).toBe(1011) }) }) + +describe('toWebSocketResponse', () => { + it('throws a directive error when WebSocketPair is unavailable', () => { + expect(() => + toWebSocketResponse(new Request('https://x/api/chat'), { + onRun: () => (async function* () {})(), + }), + ).toThrow(/WebSocketPair/) + }) +}) From a43820a49c17c5349bea2ad1645f3d806966f57e Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 20 Jul 2026 16:29:44 +0200 Subject: [PATCH 11/26] fix(ai): drop unused request param from resumeWebSocketResponse --- packages/ai/src/stream-to-websocket.ts | 10 +++++----- packages/ai/tests/stream-to-websocket.test.ts | 13 +++++++++++++ 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/packages/ai/src/stream-to-websocket.ts b/packages/ai/src/stream-to-websocket.ts index 5017e76b8d..f370f8a4ae 100644 --- a/packages/ai/src/stream-to-websocket.ts +++ b/packages/ai/src/stream-to-websocket.ts @@ -312,13 +312,13 @@ export function toWebSocketResponse( * upgrade `Response` carrying the client socket. Throws when the runtime has * no `WebSocketPair` (e.g. Node) — upgrade the socket yourself and call * {@link resumeWebSocketStream} directly there. + * + * @example + * ```ts + * resumeWebSocketResponse({ adapter: memoryStream(request) }) + * ``` */ export function resumeWebSocketResponse( - // Unused: kept so the signature mirrors `toWebSocketResponse` (both take - // the upgrade `Request` first) even though the resume path only needs the - // pre-built adapter — the caller already captured the request when - // constructing it (e.g. `memoryStream(request)`). - _request: Request, options: { adapter: StreamDurability batch?: number diff --git a/packages/ai/tests/stream-to-websocket.test.ts b/packages/ai/tests/stream-to-websocket.test.ts index dcf842750e..03c7394c18 100644 --- a/packages/ai/tests/stream-to-websocket.test.ts +++ b/packages/ai/tests/stream-to-websocket.test.ts @@ -3,6 +3,7 @@ import { buildTurnRequest, decodeWsFrame, encodeWsFrame, + resumeWebSocketResponse, resumeWebSocketStream, toWebSocketResponse, toWebSocketStream, @@ -323,3 +324,15 @@ describe('toWebSocketResponse', () => { ).toThrow(/WebSocketPair/) }) }) + +describe('resumeWebSocketResponse', () => { + it('throws a directive error when WebSocketPair is unavailable', () => { + expect(() => + resumeWebSocketResponse({ + adapter: memoryStream( + new Request('https://x/api/chat?runId=r&offset=-1'), + ), + }), + ).toThrow(/WebSocketPair/) + }) +}) From f8802abce899cb49f3ce5a050da032c7ae2490ff Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 20 Jul 2026 16:38:22 +0200 Subject: [PATCH 12/26] feat(ai-client): webSocket() subscribe/send connection adapter --- packages/ai-client/src/connection-adapters.ts | 128 ++++++++++++++++++ packages/ai-client/src/index.ts | 2 + .../connection-adapters-websocket.test.ts | 77 +++++++++++ 3 files changed, 207 insertions(+) create mode 100644 packages/ai-client/tests/connection-adapters-websocket.test.ts diff --git a/packages/ai-client/src/connection-adapters.ts b/packages/ai-client/src/connection-adapters.ts index 0883ae77ad..3933881a4b 100644 --- a/packages/ai-client/src/connection-adapters.ts +++ b/packages/ai-client/src/connection-adapters.ts @@ -1560,6 +1560,134 @@ export function xhrHttpStream( } } +export interface WebSocketConnectionOptions { + protocols?: string | Array + body?: Record + reconnect?: ReconnectOptions + /** Override the WebSocket implementation (tests / non-browser runtimes). */ + WebSocketImpl?: typeof WebSocket +} + +function runIdQuery(url: string, runId: string | undefined): string { + return runId ? withSearchParams(url, { runId }) : url +} + +/** + * Full-duplex, conversation-scoped WebSocket connection adapter. Pairs with the + * server `toWebSocketResponse` / `toWebSocketStream`. `send()` writes a + * RunAgentInput frame; `subscribe()` yields inbound chunks. Durable runs are + * offset-tagged (`{ id, chunk }`); reconnect is added in a later task. + */ +export function webSocket( + url: string | (() => string), + options: WebSocketConnectionOptions = {}, +): SubscribeConnectionAdapter & { + joinRun: (runId: string, abortSignal?: AbortSignal) => AsyncIterable +} { + const Impl = options.WebSocketImpl ?? WebSocket + let socket: WebSocket | undefined + const listeners = new Set<(chunk: StreamChunk) => void>() + + function openOnce(target: string): WebSocket { + if (socket && socket.readyState <= 1) return socket + socket = options.protocols + ? new Impl(target, options.protocols) + : new Impl(target) + socket.onmessage = (event: MessageEvent) => { + const parsed: unknown = JSON.parse(String(event.data)) + if ( + typeof parsed === 'object' && + parsed !== null && + (parsed as { type?: unknown }).type === 'ping' + ) { + return + } + const chunk = isNdjsonEnvelope(parsed) + ? parsed.chunk + : (parsed as StreamChunk) + for (const l of listeners) l(chunk) + } + return socket + } + + function waitOpen(ws: WebSocket): Promise { + if (ws.readyState === 1) return Promise.resolve() + return new Promise((resolve, reject) => { + ws.onopen = () => resolve() + ws.onerror = (e) => reject(new StreamReadError(e)) + }) + } + + return { + subscribe(abortSignal?: AbortSignal): AsyncIterable { + const queue: Array = [] + const waiters: Array<(c: StreamChunk | null) => void> = [] + const push = (c: StreamChunk) => { + const w = waiters.shift() + if (w) w(c) + else queue.push(c) + } + listeners.add(push) + abortSignal?.addEventListener('abort', () => { + const w = waiters.shift() + if (w) w(null) + }) + return (async function* () { + try { + while (!abortSignal?.aborted) { + const buffered = queue.shift() + const chunk = + buffered ?? + (await new Promise((r) => waiters.push(r))) + if (chunk === null) return + yield chunk + } + } finally { + listeners.delete(push) + } + })() + }, + async send(messages, data, _abortSignal, runContext) { + const target = typeof url === 'function' ? url() : url + const ws = openOnce(runIdQuery(target, runContext?.runId)) + await waitOpen(ws) + const body = buildRunAgentInputBody(messages, data, runContext, { + body: options.body, + }) + ws.send(JSON.stringify(body)) + }, + joinRun(runId, abortSignal): AsyncIterable { + const target = withSearchParams( + typeof url === 'function' ? url() : url, + { offset: '-1', runId }, + ) + openOnce(target) + const chunks: Array = [] + const waiters: Array<(c: StreamChunk | null) => void> = [] + const push = (c: StreamChunk) => { + const w = waiters.shift() + if (w) w(c) + else chunks.push(c) + } + listeners.add(push) + abortSignal?.addEventListener('abort', () => waiters.shift()?.(null)) + return (async function* () { + try { + while (!abortSignal?.aborted) { + const c = + chunks.shift() ?? + (await new Promise((r) => waiters.push(r))) + if (c === null) return + yield c + } + } finally { + listeners.delete(push) + } + })() + }, + } +} + /** * Create a direct stream connection adapter (for server functions or direct streams) * diff --git a/packages/ai-client/src/index.ts b/packages/ai-client/src/index.ts index 41605186c2..f99e73b5f8 100644 --- a/packages/ai-client/src/index.ts +++ b/packages/ai-client/src/index.ts @@ -94,6 +94,7 @@ export { xhrHttpStream, stream, rpcStream, + webSocket, StreamTruncatedError, DurableStreamIncompleteError, StreamReconnectLimitError, @@ -104,6 +105,7 @@ export { type ResumableConnectConnectionAdapter, type RunAgentInputContext, type SubscribeConnectionAdapter, + type WebSocketConnectionOptions, type XhrConnectionOptions, } from './connection-adapters' diff --git a/packages/ai-client/tests/connection-adapters-websocket.test.ts b/packages/ai-client/tests/connection-adapters-websocket.test.ts new file mode 100644 index 0000000000..71fb450034 --- /dev/null +++ b/packages/ai-client/tests/connection-adapters-websocket.test.ts @@ -0,0 +1,77 @@ +import { describe, expect, it } from 'vitest' +import { webSocket } from '../src/connection-adapters' + +// Minimal fake WebSocket (constructor-compatible with the WHATWG shape). +class FakeWebSocket { + static instances: Array = [] + onopen: (() => void) | null = null + onmessage: ((ev: { data: string }) => void) | null = null + onclose: (() => void) | null = null + onerror: ((ev: unknown) => void) | null = null + sent: Array = [] + readyState = 0 + url: string + constructor(url: string) { + this.url = url + FakeWebSocket.instances.push(this) + queueMicrotask(() => { + this.readyState = 1 + this.onopen?.() + }) + } + send(data: string): void { + this.sent.push(data) + } + close(): void { + this.readyState = 3 + this.onclose?.() + } + emit(chunk: unknown, id?: string): void { + this.onmessage?.({ + data: JSON.stringify(id === undefined ? chunk : { id, chunk }), + }) + } +} + +function drain( + iter: AsyncIterable, + sink: Array, + signal: AbortSignal, +): Promise { + return (async () => { + for await (const c of iter) { + if (signal.aborted) break + sink.push(c) + } + })() +} + +describe('webSocket() subscribe/send', () => { + it('sends a RunAgentInput frame and yields inbound chunks', async () => { + FakeWebSocket.instances = [] + const conn = webSocket('wss://x/api/chat', { + WebSocketImpl: FakeWebSocket as unknown as typeof WebSocket, + }) + const ac = new AbortController() + const received: Array = [] + const sub = drain(conn.subscribe(ac.signal), received, ac.signal) + await conn.send([{ role: 'user', content: 'hi' } as any], undefined, ac.signal, { + threadId: 't', + runId: 'r', + }) + const ws = FakeWebSocket.instances[0] + if (!ws) throw new Error('expected a FakeWebSocket instance to have been created') + await new Promise((r) => setTimeout(r, 0)) + expect(ws.sent.length).toBe(1) + const sentFrame = ws.sent[0] + if (!sentFrame) throw new Error('expected a sent frame') + expect(JSON.parse(sentFrame).runId).toBe('r') + + ws.emit({ type: 'TEXT_MESSAGE_CONTENT', delta: 'a', timestamp: 0 }, 'off-1') + ws.emit({ type: 'ping' }) // must be ignored + await new Promise((r) => setTimeout(r, 0)) + ac.abort() + await sub + expect(received.map((c) => c.type)).toEqual(['TEXT_MESSAGE_CONTENT']) + }) +}) From b62f4b2aecb70e2406c15bf74af9550dad28af99 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 20 Jul 2026 16:47:29 +0200 Subject: [PATCH 13/26] fix(ai-client): memoize webSocket open-promise + cover bare-chunk/joinRun --- packages/ai-client/src/connection-adapters.ts | 27 ++++-- .../connection-adapters-websocket.test.ts | 88 +++++++++++++++++++ 2 files changed, 108 insertions(+), 7 deletions(-) diff --git a/packages/ai-client/src/connection-adapters.ts b/packages/ai-client/src/connection-adapters.ts index 3933881a4b..a3c687040a 100644 --- a/packages/ai-client/src/connection-adapters.ts +++ b/packages/ai-client/src/connection-adapters.ts @@ -1586,14 +1586,27 @@ export function webSocket( } { const Impl = options.WebSocketImpl ?? WebSocket let socket: WebSocket | undefined + // Memoized per-socket open promise. `openOnce` sets `onopen`/`onerror` + // exactly ONCE, at socket-creation time, and stores the resulting promise + // here. Without this, `waitOpen` assigning `onopen`/`onerror` on every call + // would clobber a still-pending prior caller's handlers: `openOnce` reuses + // the same in-flight socket for concurrent callers (`readyState <= 1`), so a + // second `send()` issued before the handshake completes would overwrite the + // first call's handlers and leave its promise permanently unresolved. + let openPromise: Promise | undefined const listeners = new Set<(chunk: StreamChunk) => void>() function openOnce(target: string): WebSocket { if (socket && socket.readyState <= 1) return socket - socket = options.protocols + const ws = options.protocols ? new Impl(target, options.protocols) : new Impl(target) - socket.onmessage = (event: MessageEvent) => { + socket = ws + openPromise = new Promise((resolve, reject) => { + ws.onopen = () => resolve() + ws.onerror = (e) => reject(new StreamReadError(e)) + }) + ws.onmessage = (event: MessageEvent) => { const parsed: unknown = JSON.parse(String(event.data)) if ( typeof parsed === 'object' && @@ -1607,15 +1620,15 @@ export function webSocket( : (parsed as StreamChunk) for (const l of listeners) l(chunk) } - return socket + return ws } function waitOpen(ws: WebSocket): Promise { if (ws.readyState === 1) return Promise.resolve() - return new Promise((resolve, reject) => { - ws.onopen = () => resolve() - ws.onerror = (e) => reject(new StreamReadError(e)) - }) + // Concurrent callers awaiting the SAME in-flight socket share the SAME + // memoized promise (set once in `openOnce`), so none of them clobber + // another's onopen/onerror handler. + return openPromise ?? Promise.resolve() } return { diff --git a/packages/ai-client/tests/connection-adapters-websocket.test.ts b/packages/ai-client/tests/connection-adapters-websocket.test.ts index 71fb450034..ded44336c9 100644 --- a/packages/ai-client/tests/connection-adapters-websocket.test.ts +++ b/packages/ai-client/tests/connection-adapters-websocket.test.ts @@ -74,4 +74,92 @@ describe('webSocket() subscribe/send', () => { await sub expect(received.map((c) => c.type)).toEqual(['TEXT_MESSAGE_CONTENT']) }) + + it('delivers a bare chunk (no id envelope) as-is to subscribe()', async () => { + FakeWebSocket.instances = [] + const conn = webSocket('wss://x/api/chat', { + WebSocketImpl: FakeWebSocket as unknown as typeof WebSocket, + }) + const ac = new AbortController() + const received: Array = [] + const sub = drain(conn.subscribe(ac.signal), received, ac.signal) + await conn.send([{ role: 'user', content: 'hi' } as any], undefined, ac.signal, { + threadId: 't', + runId: 'r', + }) + const ws = FakeWebSocket.instances[0] + if (!ws) throw new Error('expected a FakeWebSocket instance to have been created') + await new Promise((r) => setTimeout(r, 0)) + + // Bare chunk (no id) — not wrapped in an {id, chunk} envelope. + ws.emit({ type: 'TEXT_MESSAGE_CONTENT', delta: 'x', timestamp: 0 }) + await new Promise((r) => setTimeout(r, 0)) + ac.abort() + await sub + expect(received.map((c) => c.type)).toEqual(['TEXT_MESSAGE_CONTENT']) + expect(received[0].delta).toBe('x') + }) + + it('joinRun opens a socket with offset=-1 and runId, and yields emitted chunks', async () => { + FakeWebSocket.instances = [] + const conn = webSocket('wss://x/api/chat', { + WebSocketImpl: FakeWebSocket as unknown as typeof WebSocket, + }) + const ac = new AbortController() + const received: Array = [] + const sub = drain(conn.joinRun('run-j', ac.signal), received, ac.signal) + await new Promise((r) => setTimeout(r, 0)) + + const ws = FakeWebSocket.instances[0] + if (!ws) throw new Error('expected a FakeWebSocket instance to have been created') + expect(ws.url).toContain('offset=-1') + expect(ws.url).toContain('runId=run-j') + + ws.emit({ type: 'TEXT_MESSAGE_CONTENT', delta: 'joined', timestamp: 0 }, 'off-1') + await new Promise((r) => setTimeout(r, 0)) + ac.abort() + await sub + expect(received.map((c) => c.type)).toEqual(['TEXT_MESSAGE_CONTENT']) + expect(received[0].delta).toBe('joined') + }) + + it('does not hang when two send()s race the handshake (waitOpen memoization)', async () => { + FakeWebSocket.instances = [] + const conn = webSocket('wss://x/api/chat', { + WebSocketImpl: FakeWebSocket as unknown as typeof WebSocket, + }) + const ac = new AbortController() + + // Issue two send()s back-to-back BEFORE the fake socket's queued + // microtask flips readyState/fires onopen. Both `send()` calls resolve + // openOnce() to the SAME in-flight socket (readyState 0), so both must + // await the same memoized open-promise rather than clobbering each + // other's onopen handler. + const send1 = conn.send( + [{ role: 'user', content: 'first' } as any], + undefined, + ac.signal, + { threadId: 't1', runId: 'r1' }, + ) + const send2 = conn.send( + [{ role: 'user', content: 'second' } as any], + undefined, + ac.signal, + { threadId: 't2', runId: 'r2' }, + ) + + await Promise.race([ + Promise.all([send1, send2]), + new Promise((_, reject) => + setTimeout(() => reject(new Error('send() calls hung')), 1000), + ), + ]) + + expect(FakeWebSocket.instances.length).toBe(1) + const ws = FakeWebSocket.instances[0] + if (!ws) throw new Error('expected a FakeWebSocket instance to have been created') + expect(ws.sent.length).toBe(2) + expect(JSON.parse(ws.sent[0]!).runId).toBe('r1') + expect(JSON.parse(ws.sent[1]!).runId).toBe('r2') + }) }) From 229c7d92d2b2e4e285ebe0ea5916a42bc6fd2927 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 20 Jul 2026 16:55:03 +0200 Subject: [PATCH 14/26] fix(ai-client): silence eager open-promise rejection + test tidy --- packages/ai-client/src/connection-adapters.ts | 3 +++ .../connection-adapters-websocket.test.ts | 25 +++++++++++++------ 2 files changed, 20 insertions(+), 8 deletions(-) diff --git a/packages/ai-client/src/connection-adapters.ts b/packages/ai-client/src/connection-adapters.ts index a3c687040a..7242a0a923 100644 --- a/packages/ai-client/src/connection-adapters.ts +++ b/packages/ai-client/src/connection-adapters.ts @@ -1606,6 +1606,9 @@ export function webSocket( ws.onopen = () => resolve() ws.onerror = (e) => reject(new StreamReadError(e)) }) + // Attach a no-op handler so a socket nobody awaits (e.g. joinRun) can't raise + // an unhandled rejection if it errors. Awaiters of openPromise still see the rejection. + openPromise.catch(() => {}) ws.onmessage = (event: MessageEvent) => { const parsed: unknown = JSON.parse(String(event.data)) if ( diff --git a/packages/ai-client/tests/connection-adapters-websocket.test.ts b/packages/ai-client/tests/connection-adapters-websocket.test.ts index ded44336c9..42fe8b882a 100644 --- a/packages/ai-client/tests/connection-adapters-websocket.test.ts +++ b/packages/ai-client/tests/connection-adapters-websocket.test.ts @@ -148,18 +148,27 @@ describe('webSocket() subscribe/send', () => { { threadId: 't2', runId: 'r2' }, ) - await Promise.race([ - Promise.all([send1, send2]), - new Promise((_, reject) => - setTimeout(() => reject(new Error('send() calls hung')), 1000), - ), - ]) + let timeoutId: ReturnType | undefined + try { + await Promise.race([ + Promise.all([send1, send2]), + new Promise((_, reject) => { + timeoutId = setTimeout(() => reject(new Error('send() calls hung')), 1000) + }), + ]) + } finally { + clearTimeout(timeoutId) + } expect(FakeWebSocket.instances.length).toBe(1) const ws = FakeWebSocket.instances[0] if (!ws) throw new Error('expected a FakeWebSocket instance to have been created') expect(ws.sent.length).toBe(2) - expect(JSON.parse(ws.sent[0]!).runId).toBe('r1') - expect(JSON.parse(ws.sent[1]!).runId).toBe('r2') + const sent0 = ws.sent[0] + if (sent0 === undefined) throw new Error('expected a sent frame') + const sent1 = ws.sent[1] + if (sent1 === undefined) throw new Error('expected a sent frame') + expect(JSON.parse(sent0).runId).toBe('r1') + expect(JSON.parse(sent1).runId).toBe('r2') }) }) From cdfe2bf443d27441ae0d51e93a3503a45d50441c Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 20 Jul 2026 17:00:39 +0200 Subject: [PATCH 15/26] refactor(ai-client): extract createReconnectTracker from resumableStream --- packages/ai-client/src/connection-adapters.ts | 132 ++++++++++++------ .../connection-adapters-resumable.test.ts | 13 ++ 2 files changed, 100 insertions(+), 45 deletions(-) diff --git a/packages/ai-client/src/connection-adapters.ts b/packages/ai-client/src/connection-adapters.ts index 7242a0a923..2c8ce741cb 100644 --- a/packages/ai-client/src/connection-adapters.ts +++ b/packages/ai-client/src/connection-adapters.ts @@ -128,6 +128,84 @@ function resolveReconnectOptions( return { maxAttempts, delayMs } } +/** + * Reconnect bookkeeping shared by every resumable-stream driver: de-dupes + * offsets, tracks the last acknowledged offset, honors the SSE empty-id reset + * convention, and bounds consecutive no-progress reconnects behind a + * throttling delay. Extracted out of {@link resumableStream} so a WebSocket + * reconnect driver can reuse the exact same semantics. + */ +export interface ReconnectTracker { + /** The most recently accepted (non-duplicate, non-empty) offset, if any. */ + readonly lastEventId: string | undefined + /** + * Record an incoming offset. Returns `'reset'` for an empty id (SSE's + * resume-cursor reset — clears the de-dupe set and `lastEventId`), + * `'duplicate'` for an already-seen id, and `'new'` otherwise (including + * `undefined`, which is untracked — no offset to remember). + */ + note: (id: string | undefined) => 'new' | 'duplicate' | 'reset' + /** + * Throttle before a reconnect attempt. Resets the no-progress counter when + * `madeProgress` is true; otherwise increments it and throws + * {@link StreamReconnectLimitError} once it exceeds the configured ceiling. + */ + waitBeforeReconnect: ( + madeProgress: boolean, + signal?: AbortSignal, + ) => Promise +} + +/** Create a {@link ReconnectTracker} bound to the given reconnect bounds. */ +export function createReconnectTracker( + options?: ReconnectOptions, +): ReconnectTracker { + const reconnect = resolveReconnectOptions(options) + // Retains every delivered offset for the run's lifetime. Intentionally + // bounded by run length (not evicted): a conforming server replays strictly + // after the acknowledged offset, so this only needs to catch the single + // boundary event on reconnect, but keeping the full set keeps de-dup + // correct even if a server replays a wider overlap. + const seen = new Set() + let lastEventId: string | undefined + let reconnectAttempts = 0 + return { + get lastEventId() { + return lastEventId + }, + note(id) { + if (id === undefined) return 'new' + if (id === '') { + // SSE spec: an empty `id:` resets the resume cursor. Drop the last + // offset and clear the de-dupe set; the chunk itself still delivers. + lastEventId = undefined + seen.clear() + return 'reset' + } + if (seen.has(id)) return 'duplicate' + seen.add(id) + lastEventId = id + return 'new' + }, + // Bound only CONSECUTIVE no-progress reconnects. A reconnect that made + // forward progress resets the counter, so a healthy long run (even one + // whose socket rolls after every event) never approaches the ceiling; it + // fires only when the run is genuinely stuck — reconnecting repeatedly + // with nothing new. + async waitBeforeReconnect(madeProgress, signal) { + if (madeProgress) { + reconnectAttempts = 0 + } else { + reconnectAttempts += 1 + if (reconnectAttempts > reconnect.maxAttempts) { + throw new StreamReconnectLimitError(reconnect.maxAttempts) + } + } + await abortableDelay(reconnect.delayMs, signal) + }, + } +} + /** Resolve after `ms`, or immediately once `signal` aborts. Never rejects. */ function abortableDelay(ms: number, signal?: AbortSignal): Promise { if (ms <= 0 || signal?.aborted) return Promise.resolve() @@ -495,39 +573,14 @@ async function* resumableStream( abortSignal?: AbortSignal, reconnectOptions?: ReconnectOptions, ): AsyncGenerator { - // Retains every delivered offset for the run's lifetime. Intentionally bounded - // by run length (not evicted): a conforming server replays strictly after the - // acknowledged offset, so this only needs to catch the single boundary event - // on reconnect, but keeping the full set keeps de-dup correct even if a server - // replays a wider overlap. - const seen = new Set() - let lastEventId: string | undefined - const reconnect = resolveReconnectOptions(reconnectOptions) - let reconnectAttempts = 0 - - // Throttle before re-issuing the request, and enforce the total ceiling so a - // producer that keeps dropping after each event is bounded rather than - // reconnecting forever. - // Bound only CONSECUTIVE no-progress reconnects. A reconnect that made forward - // progress resets the counter, so a healthy long run (even one whose socket - // rolls after every event) never approaches the ceiling; it fires only when - // the run is genuinely stuck — reconnecting repeatedly with nothing new. - async function waitBeforeReconnect(madeProgress: boolean): Promise { - if (madeProgress) { - reconnectAttempts = 0 - } else { - reconnectAttempts += 1 - if (reconnectAttempts > reconnect.maxAttempts) { - throw new StreamReconnectLimitError(reconnect.maxAttempts) - } - } - await abortableDelay(reconnect.delayMs, abortSignal) - } + const tracker = createReconnectTracker(reconnectOptions) for (;;) { if (abortSignal?.aborted) return const extraHeaders: Record = - lastEventId !== undefined ? { 'Last-Event-ID': lastEventId } : {} + tracker.lastEventId !== undefined + ? { 'Last-Event-ID': tracker.lastEventId } + : {} let sawTerminal = false let progressed = false @@ -536,18 +589,7 @@ async function* resumableStream( extraHeaders, abortSignal, )) { - if (id !== undefined) { - if (id === '') { - // SSE spec: an empty `id:` resets the resume cursor. Drop the last - // offset and clear the de-dupe set; the chunk itself still delivers. - lastEventId = undefined - seen.clear() - } else { - if (seen.has(id)) continue - seen.add(id) - lastEventId = id - } - } + if (tracker.note(id) === 'duplicate') continue progressed = true if (chunk.type === 'RUN_FINISHED' || chunk.type === 'RUN_ERROR') { sawTerminal = true @@ -574,9 +616,9 @@ async function* resumableStream( if ( (error instanceof StreamTruncatedError || error instanceof StreamReadError) && - lastEventId !== undefined + tracker.lastEventId !== undefined ) { - await waitBeforeReconnect(progressed) + await tracker.waitBeforeReconnect(progressed, abortSignal) continue } throw error @@ -590,14 +632,14 @@ async function* resumableStream( // would re-open past the final offset and see an empty window. if (sawTerminal) return - if (lastEventId !== undefined) { + if (tracker.lastEventId !== undefined) { // A durable (id-tagged) run. if (progressed) { // Clean end WITHOUT a terminal event but we advanced — the producer is // still going (or the socket rolled over). Reconnect from the last // offset (backing off to avoid a hot loop against the origin). Progress // resets the no-progress ceiling. - await waitBeforeReconnect(true) + await tracker.waitBeforeReconnect(true, abortSignal) continue } // Ended without a terminal event AND made no forward progress on this diff --git a/packages/ai-client/tests/connection-adapters-resumable.test.ts b/packages/ai-client/tests/connection-adapters-resumable.test.ts index d60602b8f5..59414100f3 100644 --- a/packages/ai-client/tests/connection-adapters-resumable.test.ts +++ b/packages/ai-client/tests/connection-adapters-resumable.test.ts @@ -3,6 +3,7 @@ import { EventType } from '@tanstack/ai/client' import { DurableStreamIncompleteError, StreamReconnectLimitError, + createReconnectTracker, fetchServerSentEvents, } from '../src/connection-adapters' import type { StreamChunk } from '@tanstack/ai/client' @@ -582,3 +583,15 @@ describe('resumable SSE connection adapter', () => { expect(fetchClient).toHaveBeenCalledTimes(2) }) }) + +describe('createReconnectTracker', () => { + it('de-dupes ids, tracks lastEventId, and resets on empty id', () => { + const t = createReconnectTracker() + expect(t.note('a')).toBe('new') + expect(t.lastEventId).toBe('a') + expect(t.note('a')).toBe('duplicate') + expect(t.note('')).toBe('reset') + expect(t.lastEventId).toBeUndefined() + expect(t.note('a')).toBe('new') // seen was cleared by reset + }) +}) From 0092106cf16a287641c7f52b7849c557401cedb4 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 20 Jul 2026 17:15:26 +0200 Subject: [PATCH 16/26] feat(ai-client): auto-resume webSocket subscribe on drop via ?offset --- packages/ai-client/src/connection-adapters.ts | 150 ++++++++++++++++-- .../connection-adapters-websocket.test.ts | 47 ++++++ 2 files changed, 186 insertions(+), 11 deletions(-) diff --git a/packages/ai-client/src/connection-adapters.ts b/packages/ai-client/src/connection-adapters.ts index 2c8ce741cb..28fdad3474 100644 --- a/packages/ai-client/src/connection-adapters.ts +++ b/packages/ai-client/src/connection-adapters.ts @@ -1614,11 +1614,39 @@ function runIdQuery(url: string, runId: string | undefined): string { return runId ? withSearchParams(url, { runId }) : url } +/** A subscribe()/joinRun() consumer's registration: receives chunks or a fatal error. */ +interface WebSocketChunkSink { + push: (chunk: StreamChunk) => void + fail: (error: unknown) => void +} + +/** + * The send()-driven run currently owning auto-reconnect for a `webSocket()` + * connection. `undefined` when no `send()` has established a run yet (e.g. a + * `joinRun()`-only connection), so a drop there is never auto-resumed — Task + * 11 scopes reconnect to the run `subscribe()`/`send()` are driving. + */ +interface WebSocketRunSession { + runId: string | undefined + readonly tracker: ReconnectTracker + sawTerminal: boolean + /** Made forward progress (a new, non-duplicate chunk) since the last (re)connect. */ + progressed: boolean + signal: AbortSignal | undefined +} + /** * Full-duplex, conversation-scoped WebSocket connection adapter. Pairs with the * server `toWebSocketResponse` / `toWebSocketStream`. `send()` writes a - * RunAgentInput frame; `subscribe()` yields inbound chunks. Durable runs are - * offset-tagged (`{ id, chunk }`); reconnect is added in a later task. + * RunAgentInput frame; `subscribe()` yields inbound chunks. + * + * Resumable: `send()` establishes a run session backed by a + * {@link createReconnectTracker}. If the socket closes before a terminal + * (`RUN_FINISHED`/`RUN_ERROR`) chunk is seen and the run is durable + * (offset-tagged `{ id, chunk }` envelopes), the socket is reopened at + * `?runId=&offset=`, de-duping the replayed boundary. A drop with + * no offset ever observed (non-durable) surfaces {@link StreamReadError} + * instead of reconnecting — there is nothing to resume from. */ export function webSocket( url: string | (() => string), @@ -1636,7 +1664,12 @@ export function webSocket( // second `send()` issued before the handshake completes would overwrite the // first call's handlers and leave its promise permanently unresolved. let openPromise: Promise | undefined - const listeners = new Set<(chunk: StreamChunk) => void>() + const listeners = new Set() + let currentSession: WebSocketRunSession | undefined + + function failAll(error: unknown): void { + for (const l of listeners) l.fail(error) + } function openOnce(target: string): WebSocket { if (socket && socket.readyState <= 1) return socket @@ -1660,14 +1693,73 @@ export function webSocket( ) { return } + const envelopeId = isNdjsonEnvelope(parsed) ? parsed.id : undefined const chunk = isNdjsonEnvelope(parsed) ? parsed.chunk : (parsed as StreamChunk) - for (const l of listeners) l(chunk) + + // Thread durable chunks through the active run session's tracker (if + // any) so a later reconnect knows the last offset and can skip a + // replayed boundary. A socket with no active session (e.g. joinRun() + // only) dispatches chunks as-is, exactly as before Task 11. + const session = currentSession + if (session) { + if (session.tracker.note(envelopeId) === 'duplicate') return + session.progressed = true + if (session.runId === undefined) { + session.runId = getChunkRunId(chunk) + } + if (chunk.type === 'RUN_FINISHED' || chunk.type === 'RUN_ERROR') { + session.sawTerminal = true + } + } + for (const l of listeners) l.push(chunk) + } + ws.onclose = () => { + const session = currentSession + if (!session || session.signal?.aborted || session.sawTerminal) return + const lastEventId = session.tracker.lastEventId + if (lastEventId === undefined) { + // Non-durable run (no offset ever observed) — nothing to resume + // from. Surface a hard failure rather than silently reconnecting + // forever against a server that never tags its events. + currentSession = undefined + failAll( + new StreamReadError(new Error('WebSocket connection closed')), + ) + return + } + void reconnect(session, lastEventId) } return ws } + async function reconnect( + session: WebSocketRunSession, + offset: string, + ): Promise { + try { + // Bounded by the shared tracker's consecutive-no-progress ceiling — + // mirrors resumableStream so a flapping server can't reconnect forever. + await session.tracker.waitBeforeReconnect( + session.progressed, + session.signal, + ) + } catch (error) { + currentSession = undefined + failAll(error) + return + } + if (session.signal?.aborted) return + session.progressed = false + const base = typeof url === 'function' ? url() : url + const target = withSearchParams(base, { + ...(session.runId !== undefined ? { runId: session.runId } : {}), + offset, + }) + openOnce(target) + } + function waitOpen(ws: WebSocket): Promise { if (ws.readyState === 1) return Promise.resolve() // Concurrent callers awaiting the SAME in-flight socket share the SAME @@ -1680,12 +1772,19 @@ export function webSocket( subscribe(abortSignal?: AbortSignal): AsyncIterable { const queue: Array = [] const waiters: Array<(c: StreamChunk | null) => void> = [] + let failure: unknown const push = (c: StreamChunk) => { const w = waiters.shift() if (w) w(c) else queue.push(c) } - listeners.add(push) + const fail = (e: unknown) => { + failure = e + const w = waiters.shift() + if (w) w(null) + } + const sink: WebSocketChunkSink = { push, fail } + listeners.add(sink) abortSignal?.addEventListener('abort', () => { const w = waiters.shift() if (w) w(null) @@ -1697,18 +1796,37 @@ export function webSocket( const chunk = buffered ?? (await new Promise((r) => waiters.push(r))) - if (chunk === null) return + if (chunk === null) { + if (failure !== undefined) throw failure + return + } yield chunk } } finally { - listeners.delete(push) + listeners.delete(sink) } })() }, - async send(messages, data, _abortSignal, runContext) { + async send(messages, data, abortSignal, runContext) { const target = typeof url === 'function' ? url() : url const ws = openOnce(runIdQuery(target, runContext?.runId)) await waitOpen(ws) + // Establish (or continue) the run session this socket is driving, so + // an unterminated drop can auto-resume it. A distinct runId starts a + // fresh tracker (a new run's offsets are unrelated to the last one's); + // the same runId reuses the tracker so a repeat send() on an + // already-tracked run doesn't lose its de-dupe/offset state. + if (!currentSession || currentSession.runId !== runContext?.runId) { + currentSession = { + runId: runContext?.runId, + tracker: createReconnectTracker(options.reconnect), + sawTerminal: false, + progressed: false, + signal: abortSignal, + } + } else { + currentSession.signal = abortSignal + } const body = buildRunAgentInputBody(messages, data, runContext, { body: options.body, }) @@ -1722,12 +1840,19 @@ export function webSocket( openOnce(target) const chunks: Array = [] const waiters: Array<(c: StreamChunk | null) => void> = [] + let failure: unknown const push = (c: StreamChunk) => { const w = waiters.shift() if (w) w(c) else chunks.push(c) } - listeners.add(push) + const fail = (e: unknown) => { + failure = e + const w = waiters.shift() + if (w) w(null) + } + const sink: WebSocketChunkSink = { push, fail } + listeners.add(sink) abortSignal?.addEventListener('abort', () => waiters.shift()?.(null)) return (async function* () { try { @@ -1735,11 +1860,14 @@ export function webSocket( const c = chunks.shift() ?? (await new Promise((r) => waiters.push(r))) - if (c === null) return + if (c === null) { + if (failure !== undefined) throw failure + return + } yield c } } finally { - listeners.delete(push) + listeners.delete(sink) } })() }, diff --git a/packages/ai-client/tests/connection-adapters-websocket.test.ts b/packages/ai-client/tests/connection-adapters-websocket.test.ts index 42fe8b882a..2ac9d99d81 100644 --- a/packages/ai-client/tests/connection-adapters-websocket.test.ts +++ b/packages/ai-client/tests/connection-adapters-websocket.test.ts @@ -172,3 +172,50 @@ describe('webSocket() subscribe/send', () => { expect(JSON.parse(sent1).runId).toBe('r2') }) }) + +describe('webSocket() reconnect', () => { + it('reopens with ?runId&offset after a drop and de-dupes the overlap', async () => { + FakeWebSocket.instances = [] + const conn = webSocket('wss://x/api/chat', { + WebSocketImpl: FakeWebSocket as unknown as typeof WebSocket, + reconnect: { delayMs: 0, maxAttempts: 5 }, + }) + const ac = new AbortController() + const received: Array = [] + const sub = drain(conn.subscribe(ac.signal), received, ac.signal) + await conn.send([{ role: 'user', content: 'hi' } as any], undefined, ac.signal, { + threadId: 't', + runId: 'run-x', + }) + const ws1 = FakeWebSocket.instances[0] + if (!ws1) throw new Error('expected a FakeWebSocket instance to have been created') + await new Promise((r) => setTimeout(r, 0)) + ws1.emit({ type: 'TEXT_MESSAGE_CONTENT', delta: 'a', timestamp: 0 }, 'off-1') + await new Promise((r) => setTimeout(r, 0)) + ws1.close() // drop mid-run + + await new Promise((r) => setTimeout(r, 5)) + const ws2 = FakeWebSocket.instances[1] + expect(ws2).toBeDefined() + if (!ws2) throw new Error('expected a second FakeWebSocket instance to have been created') + const url2 = new URL(ws2.url) + expect(url2.searchParams.get('runId')).toBe('run-x') + expect(url2.searchParams.get('offset')).toBe('off-1') + + // Server replays the de-duped boundary + a new chunk + terminal. + ws2.emit({ type: 'TEXT_MESSAGE_CONTENT', delta: 'a', timestamp: 0 }, 'off-1') + ws2.emit({ type: 'TEXT_MESSAGE_CONTENT', delta: 'b', timestamp: 0 }, 'off-2') + ws2.emit( + { type: 'RUN_FINISHED', runId: 'run-x', threadId: 't', model: 'm', finishReason: 'stop', timestamp: 0 }, + 'off-3', + ) + await new Promise((r) => setTimeout(r, 0)) + ac.abort() + await sub + expect(received.map((c) => c.delta ?? c.type)).toEqual([ + 'a', + 'b', + 'RUN_FINISHED', + ]) + }) +}) From 2721eb70d1560076f5d912867e8243483e5cb533 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 20 Jul 2026 17:38:12 +0200 Subject: [PATCH 17/26] fix(ai-client): surface fatal ws drop to consumer instead of hanging + tests --- packages/ai-client/src/connection-adapters.ts | 57 +++++-- .../connection-adapters-websocket.test.ts | 152 +++++++++++++++++- 2 files changed, 196 insertions(+), 13 deletions(-) diff --git a/packages/ai-client/src/connection-adapters.ts b/packages/ai-client/src/connection-adapters.ts index 28fdad3474..70b263ed1c 100644 --- a/packages/ai-client/src/connection-adapters.ts +++ b/packages/ai-client/src/connection-adapters.ts @@ -1792,14 +1792,30 @@ export function webSocket( return (async function* () { try { while (!abortSignal?.aborted) { + // Drain buffered chunks before ever awaiting a new promise — a + // fatal drop that lands while chunks are still queued (fail() + // finds no pending waiter, since the consumer hasn't caught up + // to its buffer yet) must not be lost. const buffered = queue.shift() - const chunk = - buffered ?? - (await new Promise((r) => waiters.push(r))) - if (chunk === null) { - if (failure !== undefined) throw failure - return + if (buffered !== undefined) { + yield buffered + continue } + // Buffer exhausted: surface a failure recorded while we were + // draining, rather than awaiting a promise that will never + // resolve (the connection is dead — no future push/fail). + if (failure !== undefined) throw failure + const chunk = await new Promise((r) => + waiters.push(r), + ) + // The wait resolved because fail() woke us — surface the error + // instead of treating the null sentinel as a clean end. TS narrows + // `failure` to `undefined` from the check above and doesn't know + // the `fail()` closure can reassign it while we were awaiting — + // this check is very much still reachable. + // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition + if (failure !== undefined) throw failure + if (chunk === null) return yield chunk } } finally { @@ -1857,13 +1873,30 @@ export function webSocket( return (async function* () { try { while (!abortSignal?.aborted) { - const c = - chunks.shift() ?? - (await new Promise((r) => waiters.push(r))) - if (c === null) { - if (failure !== undefined) throw failure - return + // Drain buffered chunks before ever awaiting a new promise — a + // fatal drop that lands while chunks are still queued (fail() + // finds no pending waiter, since the consumer hasn't caught up + // to its buffer yet) must not be lost. + const buffered = chunks.shift() + if (buffered !== undefined) { + yield buffered + continue } + // Buffer exhausted: surface a failure recorded while we were + // draining, rather than awaiting a promise that will never + // resolve (the connection is dead — no future push/fail). + if (failure !== undefined) throw failure + const c = await new Promise((r) => + waiters.push(r), + ) + // The wait resolved because fail() woke us — surface the error + // instead of treating the null sentinel as a clean end. TS narrows + // `failure` to `undefined` from the check above and doesn't know + // the `fail()` closure can reassign it while we were awaiting — + // this check is very much still reachable. + // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition + if (failure !== undefined) throw failure + if (c === null) return yield c } } finally { diff --git a/packages/ai-client/tests/connection-adapters-websocket.test.ts b/packages/ai-client/tests/connection-adapters-websocket.test.ts index 2ac9d99d81..466ba976b7 100644 --- a/packages/ai-client/tests/connection-adapters-websocket.test.ts +++ b/packages/ai-client/tests/connection-adapters-websocket.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from 'vitest' -import { webSocket } from '../src/connection-adapters' +import { StreamReconnectLimitError, webSocket } from '../src/connection-adapters' // Minimal fake WebSocket (constructor-compatible with the WHATWG shape). class FakeWebSocket { @@ -46,6 +46,42 @@ function drain( })() } +/** Let queued microtasks (and delayMs: 0 reconnect delays) settle. */ +function tick(): Promise { + return new Promise((r) => setTimeout(r, 0)) +} + +/** + * Drain an async iterator to completion, collecting yielded values. Used + * (instead of `drain()`) where the test needs the returned promise to REJECT + * when the generator throws, rather than swallowing it. + */ +async function drainToEnd( + iter: AsyncIterator, +): Promise> { + const received: Array = [] + for (;;) { + const { value, done } = await iter.next() + if (done) return received + received.push(value) + } +} + +/** + * Race a promise against a bounded timeout so a regression that hangs the + * generator (instead of rejecting it) fails the test loudly rather than + * hanging the whole suite. + */ +function withTimeout(promise: Promise, message: string): Promise { + let timeoutId: ReturnType | undefined + return Promise.race([ + promise.finally(() => clearTimeout(timeoutId)), + new Promise((_, reject) => { + timeoutId = setTimeout(() => reject(new Error(message)), 1000) + }), + ]) +} + describe('webSocket() subscribe/send', () => { it('sends a RunAgentInput frame and yields inbound chunks', async () => { FakeWebSocket.instances = [] @@ -219,3 +255,117 @@ describe('webSocket() reconnect', () => { ]) }) }) + +describe('webSocket() fatal drop surfacing', () => { + it('a non-durable drop (no offset ever seen) rejects subscribe() with StreamReadError and does not reconnect', async () => { + FakeWebSocket.instances = [] + const conn = webSocket('wss://x/api/chat', { + WebSocketImpl: FakeWebSocket as unknown as typeof WebSocket, + }) + const ac = new AbortController() + // Use the raw iterator (not the `drain` helper) so the promise driving + // consumption REJECTS when the generator throws — `drain`'s for-await + // would also reject the same way, but we want an explicit iterator to + // control exactly when `.next()` is pulled relative to the drop below. + const iter = conn.subscribe(ac.signal)[Symbol.asyncIterator]() + + await conn.send( + [{ role: 'user', content: 'hi' } as any], + undefined, + ac.signal, + { threadId: 't', runId: 'r' }, + ) + const ws = FakeWebSocket.instances[0] + if (!ws) throw new Error('expected a FakeWebSocket instance to have been created') + await tick() + + // Bare chunk — NO id envelope, so no durable offset is ever tracked. + // This chunk lands in the sink's buffer (queue) BEFORE the consumer has + // pulled anything via .next(), reproducing the real chat-client shape + // (chunks can arrive well before the consumer's next await). + ws.emit({ type: 'TEXT_MESSAGE_CONTENT', delta: 'a', timestamp: 0 }) + await tick() + + // Fatal drop: no lastEventId was ever observed, so `onclose` fails the + // run immediately instead of attempting a reconnect. + ws.close() + + const result = await withTimeout( + drainToEnd(iter), + 'subscribe() hung instead of rejecting after a non-durable drop — the ' + + 'fatal failure was recorded but never surfaced to the consumer', + ) + .then((received) => ({ received, error: undefined as unknown })) + .catch((error: unknown) => ({ received: undefined, error })) + + // The buffered chunk delivered before the drop must still be observed — + // the fix drains the buffer before checking for a recorded failure. + expect(result.error).toBeDefined() + expect((result.error as Error).name).toBe('StreamReadError') + + // No reconnect: a non-durable stream has no offset to resume from. + expect(FakeWebSocket.instances.length).toBe(1) + }) + + it('exceeding the reconnect ceiling on a durable run with no forward progress rejects subscribe() with StreamReconnectLimitError', async () => { + FakeWebSocket.instances = [] + const conn = webSocket('wss://x/api/chat', { + WebSocketImpl: FakeWebSocket as unknown as typeof WebSocket, + reconnect: { delayMs: 0, maxAttempts: 2 }, + }) + const ac = new AbortController() + const iter = conn.subscribe(ac.signal)[Symbol.asyncIterator]() + + await conn.send( + [{ role: 'user', content: 'hi' } as any], + undefined, + ac.signal, + { threadId: 't', runId: 'run-stuck' }, + ) + const ws1 = FakeWebSocket.instances[0] + if (!ws1) throw new Error('expected a FakeWebSocket instance to have been created') + await tick() + + // One durable chunk establishes an offset (and makes progress, resetting + // the no-progress counter), then the socket drops — this reconnect is + // legitimate and does not count against the ceiling. + ws1.emit({ type: 'TEXT_MESSAGE_CONTENT', delta: 'a', timestamp: 0 }, 'off-1') + await tick() + ws1.close() + await tick() + + // From here on, every reconnect opens a socket that drops immediately + // with NO new chunk — no forward progress — which is exactly what the + // ceiling bounds. With maxAttempts: 2, the 3rd consecutive no-progress + // close must exceed it. + const ws2 = FakeWebSocket.instances[1] + if (!ws2) throw new Error('expected a second FakeWebSocket instance (post-drop reconnect)') + ws2.close() // no-progress attempt 1 (1 <= 2, reconnects again) + await tick() + + const ws3 = FakeWebSocket.instances[2] + if (!ws3) throw new Error('expected a third FakeWebSocket instance') + ws3.close() // no-progress attempt 2 (2 <= 2, reconnects again) + await tick() + + const ws4 = FakeWebSocket.instances[3] + if (!ws4) throw new Error('expected a fourth FakeWebSocket instance') + ws4.close() // no-progress attempt 3 (3 > 2, ceiling exceeded) + await tick() + + const result = await withTimeout( + drainToEnd(iter), + 'subscribe() hung instead of rejecting once the reconnect ceiling was exceeded', + ) + .then((received) => ({ received, error: undefined as unknown })) + .catch((error: unknown) => ({ received: undefined, error })) + + expect(result.error).toBeDefined() + expect(result.error).toBeInstanceOf(StreamReconnectLimitError) + + // Every reconnect attempt actually opened a fresh socket (4 total: the + // original send() connection plus 3 reconnects), confirming the ceiling + // was reached via genuine no-progress reconnects, not a shortcut. + expect(FakeWebSocket.instances.length).toBe(4) + }) +}) From 0b116f050eabc6844023fc79688d4f00f94f6cd6 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 20 Jul 2026 18:15:09 +0200 Subject: [PATCH 18/26] test(e2e): websocket transport delivery + reconnect-resume --- pnpm-lock.yaml | 8 +- testing/e2e/package.json | 2 + .../e2e/src/lib/durable-delivery-ws-plugin.ts | 71 ++++++++ .../e2e/src/routes/api.durable-delivery.ts | 5 +- testing/e2e/tests/websocket.spec.ts | 171 ++++++++++++++++++ testing/e2e/vite.config.ts | 2 + 6 files changed, 257 insertions(+), 2 deletions(-) create mode 100644 testing/e2e/src/lib/durable-delivery-ws-plugin.ts create mode 100644 testing/e2e/tests/websocket.spec.ts diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 5dccceec4c..651752183a 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -2621,6 +2621,9 @@ importers: vite-tsconfig-paths: specifier: ^5.1.4 version: 5.1.4(typescript@5.9.3)(vite@7.3.3(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2)) + ws: + specifier: ^8.18.3 + version: 8.21.0 zod: specifier: ^4.2.0 version: 4.3.6 @@ -2637,6 +2640,9 @@ importers: '@types/react-dom': specifier: ^19.2.3 version: 19.2.3(@types/react@19.2.7) + '@types/ws': + specifier: ^8.18.1 + version: 8.18.1 '@vitejs/plugin-react': specifier: ^5.1.2 version: 5.1.2(vite@7.3.3(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2)) @@ -21973,7 +21979,7 @@ snapshots: '@tanstack/devtools-event-bus@0.4.1': dependencies: - ws: 8.19.0 + ws: 8.21.0 transitivePeerDependencies: - bufferutil - utf-8-validate diff --git a/testing/e2e/package.json b/testing/e2e/package.json index b761eb58a0..92ee1fc2bd 100644 --- a/testing/e2e/package.json +++ b/testing/e2e/package.json @@ -49,6 +49,7 @@ "remark-gfm": "^4.0.1", "tailwindcss": "^4.1.18", "vite-tsconfig-paths": "^5.1.4", + "ws": "^8.18.3", "zod": "^4.2.0" }, "devDependencies": { @@ -56,6 +57,7 @@ "@types/node": "^24.10.1", "@types/react": "^19.2.7", "@types/react-dom": "^19.2.3", + "@types/ws": "^8.18.1", "@vitejs/plugin-react": "^5.1.2", "miniflare": "^4.20260609.0", "typescript": "5.9.3", diff --git a/testing/e2e/src/lib/durable-delivery-ws-plugin.ts b/testing/e2e/src/lib/durable-delivery-ws-plugin.ts new file mode 100644 index 0000000000..43d50bb5f3 --- /dev/null +++ b/testing/e2e/src/lib/durable-delivery-ws-plugin.ts @@ -0,0 +1,71 @@ +import { WebSocketServer } from 'ws' +import { + memoryStream, + resumeWebSocketStream, + toWebSocketStream, +} from '@tanstack/ai' +import { fixedRun } from '../routes/api.durable-delivery' +import type { Plugin } from 'vite' +import type { WebSocketLike } from '@tanstack/ai' + +/** + * A Vite dev-server plugin exposing a WebSocket arm of the provider-free + * delivery-durability harness (see `api.durable-delivery.ts` for the SSE/ + * NDJSON arms — same `fixedRun` sequence, same `memoryStream` durability + * sink). The e2e app is served by Vite/Nitro on plain Node, which has no + * `WebSocketPair`, so `toWebSocketResponse`/`resumeWebSocketResponse` (the + * CF/Deno wrapper) can't be used here. Instead this hooks the underlying + * Node http server's `upgrade` event directly — the same pattern + * `examples/ts-group-chat/chat-server/vite-plugin.ts` uses for its Cap'n Web + * socket — obtains a server socket via `ws`'s `WebSocketServer({ noServer: + * true })`, and hands it to `toWebSocketStream`/`resumeWebSocketStream`. + * + * Exempt from the aimock policy: this streams the same fixed AG-UI sequence + * as the SSE/NDJSON arms and never reaches an LLM provider's HTTP layer. + * + * Wire protocol (mirrors `stream-to-websocket.ts`): + * - Connect to `/api/durable-delivery-ws?runId=` and send one + * `RunAgentInput`-shaped JSON frame to start a fresh run — frames come back + * as `{ id, chunk }` envelopes (durable) or bare chunks, plus periodic + * `{ type: 'ping' }` heartbeats to ignore. + * - Connect to `/api/durable-delivery-ws?runId=&offset=` to + * resume: no input frame is sent, the log replays strictly after `offset`. + */ +const WS_PATH = '/api/durable-delivery-ws' + +export function durableDeliveryWebSocketPlugin(): Plugin { + return { + name: 'durable-delivery-websocket-plugin', + enforce: 'pre', + configureServer(server) { + if (!server.httpServer) return + const wss = new WebSocketServer({ noServer: true }) + + server.httpServer.on('upgrade', (req, socket, head) => { + const url = new URL( + req.url ?? '/', + `http://${req.headers.host ?? 'localhost'}`, + ) + if (url.pathname !== WS_PATH) return + + wss.handleUpgrade(req, socket, head, (ws) => { + const request = new Request(url) + // `ws`'s WebSocket implements the WHATWG send/close/addEventListener/ + // bufferedAmount surface `WebSocketLike` needs. + const socketLike = ws as unknown as WebSocketLike + + if (url.searchParams.get('offset') !== null) { + resumeWebSocketStream(socketLike, { + adapter: memoryStream(request), + }) + } else { + toWebSocketStream(socketLike, request, { + durability: (ctx) => memoryStream(ctx.request), + onRun: ({ runId, threadId }) => fixedRun(threadId, runId), + }) + } + }) + }) + }, + } +} diff --git a/testing/e2e/src/routes/api.durable-delivery.ts b/testing/e2e/src/routes/api.durable-delivery.ts index fb32f9519a..980f2649f4 100644 --- a/testing/e2e/src/routes/api.durable-delivery.ts +++ b/testing/e2e/src/routes/api.durable-delivery.ts @@ -28,7 +28,10 @@ import type { StreamChunk } from '@tanstack/ai' // bracketing: this harness deliberately exercises raw chunk delivery + resume, // not UIMessage reassembly. The durability layer terminalizes on RUN_FINISHED // (emitted below), which is all resume/join needs. -function fixedRun(threadId: string, runId: string): AsyncIterable { +export function fixedRun( + threadId: string, + runId: string, +): AsyncIterable { return (async function* () { yield { type: 'RUN_STARTED', diff --git a/testing/e2e/tests/websocket.spec.ts b/testing/e2e/tests/websocket.spec.ts new file mode 100644 index 0000000000..de547b60e5 --- /dev/null +++ b/testing/e2e/tests/websocket.spec.ts @@ -0,0 +1,171 @@ +import { expect, test } from '@playwright/test' + +/** + * Delivery durability (transport layer) over a real WebSocket. + * + * Mirrors `delivery-durability.spec.ts` (SSE/NDJSON) but exercises the WS arm + * of the same provider-free harness: `api.durable-delivery-ws` streams the + * same fixed AG-UI sequence (`fixedRun`) through the same `memoryStream` + * durability sink, so join/reconnect assertions stay deterministic. See + * `testing/e2e/src/lib/durable-delivery-ws-plugin.ts` for the server side — + * the Vite/Nitro dev server runs on plain Node (no `WebSocketPair`), so the + * plugin hooks the underlying http server's `upgrade` event directly and + * drives `toWebSocketStream`/`resumeWebSocketStream` over a `ws` socket. + * + * Exempt from the aimock policy: this harness streams a fixed sequence and + * never reaches an LLM provider's HTTP layer, so there is nothing to mock. + */ + +const FIXED_TYPES = [ + 'RUN_STARTED', + 'TEXT_MESSAGE_CONTENT', + 'TEXT_MESSAGE_CONTENT', + 'TEXT_MESSAGE_CONTENT', + 'TEXT_MESSAGE_CONTENT', + 'TEXT_MESSAGE_CONTENT', + 'RUN_FINISHED', +] + +function runInput(runId: string): Record { + return { + threadId: 'thread-durable', + runId, + messages: [], + tools: [], + context: [], + forwardedProps: {}, + state: {}, + } +} + +test.describe('delivery durability (websocket)', () => { + test('streams the fixed run in order over a single connection', async ({ + page, + }) => { + await page.goto('/') + + const result = await page.evaluate( + async ({ runId, input }) => { + const wsUrl = `${location.origin.replace(/^http/, 'ws')}/api/durable-delivery-ws?runId=${encodeURIComponent(runId)}` + const ws = new WebSocket(wsUrl) + await new Promise((resolve, reject) => { + ws.onopen = () => resolve() + ws.onerror = () => reject(new Error('ws open failed')) + }) + + const received: Array<{ type: string; id?: string }> = [] + const done = new Promise((resolve) => { + ws.onmessage = (e) => { + const frame = JSON.parse(e.data as string) as { + type?: string + id?: string + chunk?: { type: string } + } + const chunk = frame.chunk ?? (frame as { type: string }) + if (chunk.type === 'ping') return + received.push({ type: chunk.type, id: frame.id }) + if (chunk.type === 'RUN_FINISHED') resolve() + } + }) + ws.send(JSON.stringify(input)) + await done + ws.close() + return received + }, + { runId: 'e2e-ws-1', input: runInput('e2e-ws-1') }, + ) + + expect(result.map((r) => r.type)).toEqual(FIXED_TYPES) + // Durable frames are tagged with an opaque offset id. + expect(result.every((r) => typeof r.id === 'string')).toBeTruthy() + }) + + test('reconnect with ?offset resumes the remainder exactly once, in order', async ({ + page, + }) => { + await page.goto('/') + + const result = await page.evaluate( + async ({ runId, input }) => { + const base = `${location.origin.replace(/^http/, 'ws')}/api/durable-delivery-ws` + + // First connection: read only the first two chunks, then disconnect. + const first = new WebSocket(`${base}?runId=${encodeURIComponent(runId)}`) + await new Promise((resolve, reject) => { + first.onopen = () => resolve() + first.onerror = () => reject(new Error('ws open failed')) + }) + + const firstChunks: Array<{ type: string; id?: string }> = [] + const gotTwo = new Promise((resolve) => { + first.onmessage = (e) => { + const frame = JSON.parse(e.data as string) as { + id?: string + chunk?: { type: string } + type?: string + } + const chunk = frame.chunk ?? (frame as { type: string }) + if (chunk.type === 'ping') return + firstChunks.push({ type: chunk.type, id: frame.id }) + if (firstChunks.length === 2) resolve() + } + }) + first.send(JSON.stringify(input)) + await gotTwo + first.close() + // Let the close propagate server-side before reconnecting. + await new Promise((r) => setTimeout(r, 50)) + + const lastId = firstChunks[1]!.id! + + // Second connection: resume from the last acknowledged offset. + const second = new WebSocket( + `${base}?runId=${encodeURIComponent(runId)}&offset=${encodeURIComponent(lastId)}`, + ) + await new Promise((resolve, reject) => { + second.onopen = () => resolve() + second.onerror = () => reject(new Error('ws reopen failed')) + }) + + const resumedChunks: Array<{ type: string; id?: string }> = [] + const finished = new Promise((resolve) => { + second.onmessage = (e) => { + const frame = JSON.parse(e.data as string) as { + id?: string + chunk?: { type: string } + type?: string + } + const chunk = frame.chunk ?? (frame as { type: string }) + if (chunk.type === 'ping') return + resumedChunks.push({ type: chunk.type, id: frame.id }) + if (chunk.type === 'RUN_FINISHED') resolve() + } + }) + await finished + second.close() + + return { firstChunks, resumedChunks, lastId } + }, + { runId: 'e2e-ws-2', input: runInput('e2e-ws-2') }, + ) + + // First connection saw RUN_STARTED + the first content delta. + expect(result.firstChunks.map((c) => c.type)).toEqual([ + 'RUN_STARTED', + 'TEXT_MESSAGE_CONTENT', + ]) + + // Resume delivers exactly the remaining 4 content deltas + RUN_FINISHED, + // in order, with no duplicates of anything already delivered. + expect(result.resumedChunks.map((c) => c.type)).toEqual([ + 'TEXT_MESSAGE_CONTENT', + 'TEXT_MESSAGE_CONTENT', + 'TEXT_MESSAGE_CONTENT', + 'TEXT_MESSAGE_CONTENT', + 'RUN_FINISHED', + ]) + const resumedIds = result.resumedChunks.map((c) => c.id) + expect(resumedIds).not.toContain(result.lastId) + expect(new Set(resumedIds).size).toBe(resumedIds.length) + }) +}) diff --git a/testing/e2e/vite.config.ts b/testing/e2e/vite.config.ts index e522fe93b7..416f80f0d6 100644 --- a/testing/e2e/vite.config.ts +++ b/testing/e2e/vite.config.ts @@ -4,6 +4,7 @@ import viteReact from '@vitejs/plugin-react' import viteTsConfigPaths from 'vite-tsconfig-paths' import tailwindcss from '@tailwindcss/vite' import { nitroV2Plugin } from '@tanstack/nitro-v2-vite-plugin' +import { durableDeliveryWebSocketPlugin } from './src/lib/durable-delivery-ws-plugin' const config = defineConfig({ // Server-side only fix. @elevenlabs/elevenlabs-js ships a top-level @@ -17,6 +18,7 @@ const config = defineConfig({ external: ['@elevenlabs/elevenlabs-js'], }, plugins: [ + durableDeliveryWebSocketPlugin(), nitroV2Plugin({ externals: { external: ['@elevenlabs/elevenlabs-js'], From 580f7b143bddfc6352ac3ab42769432108914eb2 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 20 Jul 2026 18:31:27 +0200 Subject: [PATCH 19/26] feat(examples): websocket chat route in ts-react-chat Add a full-duplex WebSocket chat demo to ts-react-chat: a Vite dev-server plugin (websocket-chat-plugin.ts) hooks the Node http server's `upgrade` event and wires toWebSocketStream/resumeWebSocketStream around chat() with gpt-5.5, mirroring the pattern in testing/e2e's durable-delivery-ws-plugin (no WebSocketPair on Node/Nitro). The /websocket-chat route uses useChat with the webSocket() connection adapter, linked from the nav. @tanstack/ai-react didn't re-export webSocket/WebSocketConnectionOptions from @tanstack/ai-client yet (only the other connection adapters were), so this adds that re-export as it's required for the example's `import { useChat, webSocket } from '@tanstack/ai-react'`. --- examples/ts-react-chat/package.json | 2 + .../ts-react-chat/src/components/Header.tsx | 14 +++ .../src/lib/websocket-chat-plugin.ts | 90 ++++++++++++++ examples/ts-react-chat/src/routeTree.gen.ts | 21 ++++ .../src/routes/websocket-chat.tsx | 115 ++++++++++++++++++ examples/ts-react-chat/vite.config.ts | 2 + packages/ai-react/src/index.ts | 2 + pnpm-lock.yaml | 105 ++-------------- 8 files changed, 258 insertions(+), 93 deletions(-) create mode 100644 examples/ts-react-chat/src/lib/websocket-chat-plugin.ts create mode 100644 examples/ts-react-chat/src/routes/websocket-chat.tsx diff --git a/examples/ts-react-chat/package.json b/examples/ts-react-chat/package.json index 67b58ad431..6439e72828 100644 --- a/examples/ts-react-chat/package.json +++ b/examples/ts-react-chat/package.json @@ -61,6 +61,7 @@ "remark-gfm": "^4.0.1", "tailwindcss": "^4.1.18", "vite-tsconfig-paths": "^5.1.4", + "ws": "^8.18.3", "zod": "^4.2.0" }, "devDependencies": { @@ -71,6 +72,7 @@ "@types/node": "^24.10.1", "@types/react": "^19.2.7", "@types/react-dom": "^19.2.3", + "@types/ws": "^8.18.1", "@vitejs/plugin-react": "^5.1.2", "concurrently": "^9.1.2", "jsdom": "^27.2.0", diff --git a/examples/ts-react-chat/src/components/Header.tsx b/examples/ts-react-chat/src/components/Header.tsx index 9c71d1577a..2054875251 100644 --- a/examples/ts-react-chat/src/components/Header.tsx +++ b/examples/ts-react-chat/src/components/Header.tsx @@ -22,6 +22,7 @@ import { Server, Sparkles, Video, + Wifi, X, } from 'lucide-react' @@ -312,6 +313,19 @@ export default function Header() { Resumable Streams + setIsOpen(false)} + className="flex items-center gap-3 p-3 rounded-lg hover:bg-gray-800 transition-colors mb-2" + activeProps={{ + className: + 'flex items-center gap-3 p-3 rounded-lg bg-cyan-600 hover:bg-cyan-700 transition-colors mb-2', + }} + > + + WebSocket Chat + + setIsOpen(false)} diff --git a/examples/ts-react-chat/src/lib/websocket-chat-plugin.ts b/examples/ts-react-chat/src/lib/websocket-chat-plugin.ts new file mode 100644 index 0000000000..f716a53455 --- /dev/null +++ b/examples/ts-react-chat/src/lib/websocket-chat-plugin.ts @@ -0,0 +1,90 @@ +import { WebSocketServer } from 'ws' +import { + chat, + memoryStream, + resumeWebSocketStream, + toWebSocketStream, +} from '@tanstack/ai' +import { openaiText } from '@tanstack/ai-openai' +import type { Plugin } from 'vite' +import type { WebSocketLike } from '@tanstack/ai' + +/** + * A Vite dev-server plugin exposing a full-duplex WebSocket chat endpoint. + * This app is served by Vite/Nitro on plain Node, which has no + * `WebSocketPair`, so `toWebSocketResponse`/`resumeWebSocketResponse` (the + * CF/Deno wrapper) can't be used here. Instead this hooks the underlying + * Node http server's `upgrade` event directly — obtains a server socket via + * `ws`'s `WebSocketServer({ noServer: true })`, and hands it to + * `toWebSocketStream`/`resumeWebSocketStream`. Mirrors + * `testing/e2e/src/lib/durable-delivery-ws-plugin.ts`, which uses the same + * pattern for a provider-free harness. + * + * Wire protocol (client side is `webSocket('/api/chat-ws')` from + * `@tanstack/ai-react`): + * - Connect to `/api/chat-ws` and send one `RunAgentInput`-shaped JSON frame + * per turn to run the model — frames come back as `{ id, chunk }` + * envelopes (durable) plus periodic `{ type: 'ping' }` heartbeats. + * - Connect to `/api/chat-ws?offset=` to resume a dropped run: no + * input frame is sent, the durability log replays strictly after `offset`. + */ +const WS_PATH = '/api/chat-ws' + +/** + * `chat()` cancels via an `AbortController` it can read `.signal` off of, but + * `WsRunContext` (the per-turn context `toWebSocketStream` hands `onRun`) + * only carries the `AbortSignal` half — it aborts on socket close or an + * `{ type: 'abort', runId }` control frame. Bridge the two by mirroring the + * signal onto a fresh controller. + */ +function abortControllerFromSignal(signal: AbortSignal): AbortController { + const controller = new AbortController() + if (signal.aborted) controller.abort() + else + signal.addEventListener('abort', () => controller.abort(), { once: true }) + return controller +} + +export function webSocketChatPlugin(): Plugin { + return { + name: 'websocket-chat-plugin', + enforce: 'pre', + configureServer(server) { + if (!server.httpServer) return + const wss = new WebSocketServer({ noServer: true }) + + server.httpServer.on('upgrade', (req, socket, head) => { + const url = new URL( + req.url ?? '/', + `http://${req.headers.host ?? 'localhost'}`, + ) + if (url.pathname !== WS_PATH) return + + wss.handleUpgrade(req, socket, head, (ws) => { + const request = new Request(url) + // `ws`'s WebSocket implements the WHATWG send/close/addEventListener/ + // bufferedAmount surface `WebSocketLike` needs. + const socketLike = ws as unknown as WebSocketLike + + if (url.searchParams.get('offset') !== null) { + resumeWebSocketStream(socketLike, { + adapter: memoryStream(request), + }) + } else { + toWebSocketStream(socketLike, request, { + durability: (ctx) => memoryStream(ctx.request), + onRun: ({ messages, threadId, runId, signal }) => + chat({ + adapter: openaiText('gpt-5.5'), + messages, + threadId, + runId, + abortController: abortControllerFromSignal(signal), + }), + }) + } + }) + }) + }, + } +} diff --git a/examples/ts-react-chat/src/routeTree.gen.ts b/examples/ts-react-chat/src/routeTree.gen.ts index 0f05c32c42..d8f99d3051 100644 --- a/examples/ts-react-chat/src/routeTree.gen.ts +++ b/examples/ts-react-chat/src/routeTree.gen.ts @@ -9,6 +9,7 @@ // Additionally, you should also exclude this file from your linter and/or formatter to prevent it from being checked or modified. import { Route as rootRouteImport } from './routes/__root' +import { Route as WebsocketChatRouteImport } from './routes/websocket-chat' import { Route as TypesafeToolsRouteImport } from './routes/typesafe-tools' import { Route as ThreadsRouteImport } from './routes/threads' import { Route as ServerFnChatRouteImport } from './routes/server-fn-chat' @@ -58,6 +59,11 @@ import { Route as ApiGenerateSpeechRouteImport } from './routes/api.generate.spe import { Route as ApiGenerateImageRouteImport } from './routes/api.generate.image' import { Route as ApiGenerateAudioRouteImport } from './routes/api.generate.audio' +const WebsocketChatRoute = WebsocketChatRouteImport.update({ + id: '/websocket-chat', + path: '/websocket-chat', + getParentRoute: () => rootRouteImport, +} as any) const TypesafeToolsRoute = TypesafeToolsRouteImport.update({ id: '/typesafe-tools', path: '/typesafe-tools', @@ -318,6 +324,7 @@ export interface FileRoutesByFullPath { '/server-fn-chat': typeof ServerFnChatRoute '/threads': typeof ThreadsRoute '/typesafe-tools': typeof TypesafeToolsRoute + '/websocket-chat': typeof WebsocketChatRoute '/api/capability-demo': typeof ApiCapabilityDemoRoute '/api/image-gen': typeof ApiImageGenRoute '/api/image-tool-repro': typeof ApiImageToolReproRoute @@ -368,6 +375,7 @@ export interface FileRoutesByTo { '/server-fn-chat': typeof ServerFnChatRoute '/threads': typeof ThreadsRoute '/typesafe-tools': typeof TypesafeToolsRoute + '/websocket-chat': typeof WebsocketChatRoute '/api/capability-demo': typeof ApiCapabilityDemoRoute '/api/image-gen': typeof ApiImageGenRoute '/api/image-tool-repro': typeof ApiImageToolReproRoute @@ -419,6 +427,7 @@ export interface FileRoutesById { '/server-fn-chat': typeof ServerFnChatRoute '/threads': typeof ThreadsRoute '/typesafe-tools': typeof TypesafeToolsRoute + '/websocket-chat': typeof WebsocketChatRoute '/api/capability-demo': typeof ApiCapabilityDemoRoute '/api/image-gen': typeof ApiImageGenRoute '/api/image-tool-repro': typeof ApiImageToolReproRoute @@ -471,6 +480,7 @@ export interface FileRouteTypes { | '/server-fn-chat' | '/threads' | '/typesafe-tools' + | '/websocket-chat' | '/api/capability-demo' | '/api/image-gen' | '/api/image-tool-repro' @@ -521,6 +531,7 @@ export interface FileRouteTypes { | '/server-fn-chat' | '/threads' | '/typesafe-tools' + | '/websocket-chat' | '/api/capability-demo' | '/api/image-gen' | '/api/image-tool-repro' @@ -571,6 +582,7 @@ export interface FileRouteTypes { | '/server-fn-chat' | '/threads' | '/typesafe-tools' + | '/websocket-chat' | '/api/capability-demo' | '/api/image-gen' | '/api/image-tool-repro' @@ -622,6 +634,7 @@ export interface RootRouteChildren { ServerFnChatRoute: typeof ServerFnChatRoute ThreadsRoute: typeof ThreadsRoute TypesafeToolsRoute: typeof TypesafeToolsRoute + WebsocketChatRoute: typeof WebsocketChatRoute ApiCapabilityDemoRoute: typeof ApiCapabilityDemoRoute ApiImageGenRoute: typeof ApiImageGenRoute ApiImageToolReproRoute: typeof ApiImageToolReproRoute @@ -659,6 +672,13 @@ export interface RootRouteChildren { declare module '@tanstack/react-router' { interface FileRoutesByPath { + '/websocket-chat': { + id: '/websocket-chat' + path: '/websocket-chat' + fullPath: '/websocket-chat' + preLoaderRoute: typeof WebsocketChatRouteImport + parentRoute: typeof rootRouteImport + } '/typesafe-tools': { id: '/typesafe-tools' path: '/typesafe-tools' @@ -1014,6 +1034,7 @@ const rootRouteChildren: RootRouteChildren = { ServerFnChatRoute: ServerFnChatRoute, ThreadsRoute: ThreadsRoute, TypesafeToolsRoute: TypesafeToolsRoute, + WebsocketChatRoute: WebsocketChatRoute, ApiCapabilityDemoRoute: ApiCapabilityDemoRoute, ApiImageGenRoute: ApiImageGenRoute, ApiImageToolReproRoute: ApiImageToolReproRoute, diff --git a/examples/ts-react-chat/src/routes/websocket-chat.tsx b/examples/ts-react-chat/src/routes/websocket-chat.tsx new file mode 100644 index 0000000000..55d7da2c51 --- /dev/null +++ b/examples/ts-react-chat/src/routes/websocket-chat.tsx @@ -0,0 +1,115 @@ +import { createFileRoute } from '@tanstack/react-router' +import { useState } from 'react' +import { useChat, webSocket } from '@tanstack/ai-react' + +export const Route = createFileRoute('/websocket-chat')({ + component: WebSocketChatPage, +}) + +// A full-duplex WebSocket connection. `/api/chat-ws` (see +// `src/lib/websocket-chat-plugin.ts`) records every chunk to a durability +// log and serves `?offset=` resumes over the same socket protocol, so a +// dropped connection reconnects and resumes automatically — nothing on the +// client opts in beyond using `useChat` with `webSocket()`. +const connection = webSocket('/api/chat-ws') + +function WebSocketChatPage() { + const { messages, sendMessage, isLoading, connectionStatus } = useChat({ + connection, + }) + const [input, setInput] = useState('Write a short haiku about WebSockets.') + + const handleSubmit = (e: React.FormEvent) => { + e.preventDefault() + const text = input.trim() + if (!text || isLoading) return + setInput('') + void sendMessage(text) + } + + return ( +
+

WebSocket chat

+

+ This chat talks to the model over a single full-duplex WebSocket instead + of HTTP/SSE. The server durably logs each chunk, so a dropped connection + reconnects with ?offset= and resumes the run rather than + restarting the model — all through useChat with the{' '} + webSocket() connection adapter. +

+ +
+ connection: {connectionStatus} +
+ +
+ {messages.map((message) => ( +
+
{message.role}
+ {message.parts.map((part, index) => + part.type === 'text' && part.content ? ( +

+ {part.content} +

+ ) : null, + )} +
+ ))} +
+ +
+ setInput(e.target.value)} + placeholder="Ask something…" + style={{ flex: 1, padding: 8 }} + /> + +
+
+ ) +} + +const page: React.CSSProperties = { + maxWidth: 720, + margin: '0 auto', + padding: 24, + fontFamily: 'system-ui, sans-serif', +} + +const bubble: React.CSSProperties = { + borderRadius: 8, + padding: '10px 14px', + maxWidth: '85%', +} + +const userBubble: React.CSSProperties = { + ...bubble, + alignSelf: 'flex-end', + background: '#eef2ff', +} + +const assistantBubble: React.CSSProperties = { + ...bubble, + alignSelf: 'flex-start', + background: '#f6f6f6', +} + +const roleLabel: React.CSSProperties = { + fontSize: 11, + textTransform: 'uppercase', + letterSpacing: '0.05em', + color: '#999', + marginBottom: 4, +} diff --git a/examples/ts-react-chat/vite.config.ts b/examples/ts-react-chat/vite.config.ts index 9bbd207642..24052e99a2 100644 --- a/examples/ts-react-chat/vite.config.ts +++ b/examples/ts-react-chat/vite.config.ts @@ -5,6 +5,7 @@ import viteTsConfigPaths from 'vite-tsconfig-paths' import tailwindcss from '@tailwindcss/vite' import { nitroV2Plugin } from '@tanstack/nitro-v2-vite-plugin' import { devtools } from '@tanstack/devtools-vite' +import { webSocketChatPlugin } from './src/lib/websocket-chat-plugin' // `dockerode` is a server-only dependency that pulls in optional native addons // (`ssh2` → `cpu-features`, a `.node` binary that this install does not compile). @@ -47,6 +48,7 @@ const config = defineConfig({ build: { rollupOptions: { external: SERVER_ONLY_NATIVE } }, plugins: [ devtools(), + webSocketChatPlugin(), nitroV2Plugin({ externals: { external: ['@elevenlabs/elevenlabs-js', ...SERVER_ONLY_NATIVE], diff --git a/packages/ai-react/src/index.ts b/packages/ai-react/src/index.ts index 773c05456a..c563a9c28e 100644 --- a/packages/ai-react/src/index.ts +++ b/packages/ai-react/src/index.ts @@ -74,6 +74,7 @@ export { xhrHttpStream, stream, rpcStream, + webSocket, createChatClientOptions, createMcpAppBridge, type McpAppBridge, @@ -87,6 +88,7 @@ export { type RunAgentInputContext, type FetchConnectionOptions, type XhrConnectionOptions, + type WebSocketConnectionOptions, type InferChatMessages, type GenerationClientState, type ImageGenerateInput, diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 651752183a..5e8ded4f81 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -297,7 +297,7 @@ importers: version: 1.159.5(react-dom@19.2.3(react@19.2.3))(react@19.2.3) '@tanstack/react-start': specifier: ^1.159.0 - version: 1.159.5(crossws@0.4.6(srvx@0.10.1))(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(vite-plugin-solid@2.11.10(solid-js@1.9.10)(vite@7.3.3(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2)))(vite@7.3.3(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2)) + version: 1.159.5(crossws@0.4.6(srvx@0.11.17))(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(vite-plugin-solid@2.11.10(solid-js@1.9.10)(vite@7.3.3(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2)))(vite@7.3.3(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2)) lucide-react: specifier: ^0.561.0 version: 0.561.0(react@19.2.3) @@ -573,7 +573,7 @@ importers: version: 1.159.5(@tanstack/react-router@1.159.5(react-dom@19.2.3(react@19.2.3))(react@19.2.3))(@tanstack/router-core@1.159.4)(csstype@3.2.3)(react-dom@19.2.3(react@19.2.3))(react@19.2.3) '@tanstack/react-start': specifier: ^1.159.0 - version: 1.159.5(crossws@0.4.6(srvx@0.11.17))(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(vite-plugin-solid@2.11.10(solid-js@1.9.10)(vite@7.3.3(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2)))(vite@7.3.3(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2)) + version: 1.159.5(crossws@0.4.6(srvx@0.10.1))(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(vite-plugin-solid@2.11.10(solid-js@1.9.10)(vite@7.3.3(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2)))(vite@7.3.3(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2)) '@tanstack/router-plugin': specifier: ^1.158.4 version: 1.159.5(@tanstack/react-router@1.159.5(react-dom@19.2.3(react@19.2.3))(react@19.2.3))(vite-plugin-solid@2.11.10(solid-js@1.9.10)(vite@7.3.3(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2)))(vite@7.3.3(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2)) @@ -791,6 +791,9 @@ importers: vite-tsconfig-paths: specifier: ^5.1.4 version: 5.1.4(typescript@5.9.3)(vite@7.3.3(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2)) + ws: + specifier: ^8.18.3 + version: 8.21.0 zod: specifier: ^4.2.0 version: 4.2.1 @@ -816,6 +819,9 @@ importers: '@types/react-dom': specifier: ^19.2.3 version: 19.2.3(@types/react@19.2.7) + '@types/ws': + specifier: ^8.18.1 + version: 8.18.1 '@vitejs/plugin-react': specifier: ^5.1.2 version: 5.1.2(vite@7.3.3(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2)) @@ -1410,10 +1416,10 @@ importers: devDependencies: '@analogjs/vite-plugin-angular': specifier: ^2.6.2 - version: 2.6.2(@angular/build@21.2.15(60b70a9cad7e6bc3ab5f2203f7c1389b))(@emnapi/core@1.11.1)(@emnapi/runtime@1.11.1)(vite@7.3.3(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2)) + version: 2.6.2(@angular/build@21.2.15(5bb6eeec973d49398ca24eade340c469))(@emnapi/core@1.11.1)(@emnapi/runtime@1.11.1)(vite@7.3.3(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.97.3)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2)) '@angular/build': specifier: ^21.2.0 - version: 21.2.15(60b70a9cad7e6bc3ab5f2203f7c1389b) + version: 21.2.15(5bb6eeec973d49398ca24eade340c469) '@angular/common': specifier: ^21.2.0 version: 21.2.17(@angular/core@21.2.17(@angular/compiler@21.2.17)(rxjs@7.8.2)(zone.js@0.15.1))(rxjs@7.8.2) @@ -1446,7 +1452,7 @@ importers: version: 5.9.3 vitest: specifier: ^4.0.14 - version: 4.1.4(@opentelemetry/api@1.9.1)(@types/node@24.10.3)(@vitest/coverage-v8@4.0.14)(happy-dom@20.0.11)(jsdom@27.3.0(postcss@8.5.15))(vite@7.3.3(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2)) + version: 4.1.4(@opentelemetry/api@1.9.1)(@types/node@24.10.3)(happy-dom@20.0.11)(jsdom@27.3.0(postcss@8.5.15))(vite@7.3.3(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.97.3)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2)) zod: specifier: ^4.2.0 version: 4.3.6 @@ -16034,18 +16040,6 @@ packages: utf-8-validate: optional: true - ws@8.19.0: - resolution: {integrity: sha512-blAT2mjOEIi0ZzruJfIhb3nps74PRWTCz1IjglWEEpQl5XS/UNama6u2/rjFkDDouqr4L67ry+1aGIALViWjDg==} - engines: {node: '>=10.0.0'} - peerDependencies: - bufferutil: ^4.0.1 - utf-8-validate: '>=5.0.2' - peerDependenciesMeta: - bufferutil: - optional: true - utf-8-validate: - optional: true - ws@8.20.1: resolution: {integrity: sha512-It4dO0K5v//JtTXuPkfEOaI3uUN87iYPnqo/ZzqCoG3g8uhA66QUMs/SrM0YK7/NAu+r4LMh/9dq2A7k+rHs+w==} engines: {node: '>=10.0.0'} @@ -16226,20 +16220,6 @@ snapshots: - '@emnapi/core' - '@emnapi/runtime' - '@analogjs/vite-plugin-angular@2.6.2(@angular/build@21.2.15(60b70a9cad7e6bc3ab5f2203f7c1389b))(@emnapi/core@1.11.1)(@emnapi/runtime@1.11.1)(vite@7.3.3(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2))': - dependencies: - magic-string: 0.30.21 - obug: 2.1.1 - oxc-parser: 0.121.0(@emnapi/core@1.11.1)(@emnapi/runtime@1.11.1) - tinyglobby: 0.2.16 - ts-morph: 21.0.1 - optionalDependencies: - '@angular/build': 21.2.15(60b70a9cad7e6bc3ab5f2203f7c1389b) - vite: 7.3.3(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2) - transitivePeerDependencies: - - '@emnapi/core' - - '@emnapi/runtime' - '@angular-devkit/architect@0.2102.15(chokidar@5.0.0)': dependencies: '@angular-devkit/core': 21.2.15(chokidar@5.0.0) @@ -16315,63 +16295,6 @@ snapshots: - tsx - yaml - '@angular/build@21.2.15(60b70a9cad7e6bc3ab5f2203f7c1389b)': - dependencies: - '@ampproject/remapping': 2.3.0 - '@angular-devkit/architect': 0.2102.15(chokidar@5.0.0) - '@angular/compiler': 21.2.17 - '@angular/compiler-cli': 21.2.17(@angular/compiler@21.2.17)(typescript@5.9.3) - '@babel/core': 7.29.0 - '@babel/helper-annotate-as-pure': 7.27.3 - '@babel/helper-split-export-declaration': 7.24.7 - '@inquirer/confirm': 5.1.21(@types/node@24.10.3) - '@vitejs/plugin-basic-ssl': 2.1.4(vite@7.3.2(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.97.3)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2)) - beasties: 0.4.1 - browserslist: 4.28.1 - esbuild: 0.27.3 - https-proxy-agent: 7.0.6 - istanbul-lib-instrument: 6.0.3 - jsonc-parser: 3.3.1 - listr2: 9.0.5 - magic-string: 0.30.21 - mrmime: 2.0.1 - parse5-html-rewriting-stream: 8.0.0 - picomatch: 4.0.4 - piscina: 5.1.4 - rolldown: 1.0.0-rc.4(@emnapi/core@1.11.1)(@emnapi/runtime@1.11.1) - sass: 1.97.3 - semver: 7.7.4 - source-map-support: 0.5.21 - tinyglobby: 0.2.15 - tslib: 2.8.1 - typescript: 5.9.3 - undici: 7.24.4 - vite: 7.3.2(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.97.3)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2) - watchpack: 2.5.1 - optionalDependencies: - '@angular/core': 21.2.17(@angular/compiler@21.2.17)(rxjs@7.8.2)(zone.js@0.15.1) - '@angular/platform-browser': 21.2.17(@angular/common@21.2.17(@angular/core@21.2.17(@angular/compiler@21.2.17)(rxjs@7.8.2)(zone.js@0.15.1))(rxjs@7.8.2))(@angular/core@21.2.17(@angular/compiler@21.2.17)(rxjs@7.8.2)(zone.js@0.15.1)) - less: 4.6.6 - lmdb: 3.5.1 - ng-packagr: 21.2.5(@angular/compiler-cli@21.2.17(@angular/compiler@21.2.17)(typescript@5.9.3))(tailwindcss@4.1.18)(tslib@2.8.1)(typescript@5.9.3) - postcss: 8.5.15 - tailwindcss: 4.1.18 - vitest: 4.1.4(@opentelemetry/api@1.9.1)(@types/node@24.10.3)(@vitest/coverage-v8@4.0.14)(happy-dom@20.0.11)(jsdom@27.3.0(postcss@8.5.15))(vite@7.3.3(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2)) - transitivePeerDependencies: - - '@emnapi/core' - - '@emnapi/runtime' - - '@types/node' - - chokidar - - jiti - - lightningcss - - sass-embedded - - stylus - - sugarss - - supports-color - - terser - - tsx - - yaml - '@angular/common@21.2.17(@angular/core@21.2.17(@angular/compiler@21.2.17)(rxjs@7.8.2)(zone.js@0.15.1))(rxjs@7.8.2)': dependencies: '@angular/core': 21.2.17(@angular/compiler@21.2.17)(rxjs@7.8.2)(zone.js@0.15.1) @@ -17985,7 +17908,7 @@ snapshots: dependencies: command-exists: 1.2.9 node-fetch: 2.7.0 - ws: 8.19.0 + ws: 8.21.0 transitivePeerDependencies: - bufferutil - encoding @@ -23896,7 +23819,6 @@ snapshots: magic-string: 0.30.21 optionalDependencies: vite: 7.3.3(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.97.3)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2) - optional: true '@vitest/pretty-format@4.0.14': dependencies: @@ -31822,7 +31744,6 @@ snapshots: jsdom: 27.3.0(postcss@8.5.15) transitivePeerDependencies: - msw - optional: true vlq@1.0.1: {} @@ -32037,8 +31958,6 @@ snapshots: ws@8.18.3: {} - ws@8.19.0: {} - ws@8.20.1: {} ws@8.21.0: {} From b8513d752bd8654d5fbac0749065b12cced7bad0 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 20 Jul 2026 18:35:09 +0200 Subject: [PATCH 20/26] feat(framework): re-export webSocket connection adapter from solid/vue/svelte/angular --- packages/ai-angular/src/index.ts | 2 ++ packages/ai-solid/src/index.ts | 2 ++ packages/ai-svelte/src/index.ts | 2 ++ packages/ai-vue/src/index.ts | 2 ++ 4 files changed, 8 insertions(+) diff --git a/packages/ai-angular/src/index.ts b/packages/ai-angular/src/index.ts index 4f9565c58c..fb2093f625 100644 --- a/packages/ai-angular/src/index.ts +++ b/packages/ai-angular/src/index.ts @@ -82,6 +82,7 @@ export { xhrHttpStream, stream, rpcStream, + webSocket, createChatClientOptions, type ConnectionAdapter, type ConnectConnectionAdapter, @@ -89,6 +90,7 @@ export { type RunAgentInputContext, type FetchConnectionOptions, type XhrConnectionOptions, + type WebSocketConnectionOptions, type InferChatMessages, type GenerationClientState, type ImageGenerateInput, diff --git a/packages/ai-solid/src/index.ts b/packages/ai-solid/src/index.ts index 1d0c517b39..0d16a53a56 100644 --- a/packages/ai-solid/src/index.ts +++ b/packages/ai-solid/src/index.ts @@ -67,6 +67,7 @@ export { xhrHttpStream, stream, rpcStream, + webSocket, createChatClientOptions, type ConnectionAdapter, type ConnectConnectionAdapter, @@ -74,6 +75,7 @@ export { type RunAgentInputContext, type FetchConnectionOptions, type XhrConnectionOptions, + type WebSocketConnectionOptions, type InferChatMessages, type GenerationClientState, type ImageGenerateInput, diff --git a/packages/ai-svelte/src/index.ts b/packages/ai-svelte/src/index.ts index 151fdab5db..0d9004d030 100644 --- a/packages/ai-svelte/src/index.ts +++ b/packages/ai-svelte/src/index.ts @@ -70,6 +70,7 @@ export { xhrHttpStream, stream, rpcStream, + webSocket, createChatClientOptions, clientTools, type ConnectionAdapter, @@ -78,6 +79,7 @@ export { type RunAgentInputContext, type FetchConnectionOptions, type XhrConnectionOptions, + type WebSocketConnectionOptions, type InferChatMessages, type GenerationClientState, type ImageGenerateInput, diff --git a/packages/ai-vue/src/index.ts b/packages/ai-vue/src/index.ts index 1d0c517b39..0d16a53a56 100644 --- a/packages/ai-vue/src/index.ts +++ b/packages/ai-vue/src/index.ts @@ -67,6 +67,7 @@ export { xhrHttpStream, stream, rpcStream, + webSocket, createChatClientOptions, type ConnectionAdapter, type ConnectConnectionAdapter, @@ -74,6 +75,7 @@ export { type RunAgentInputContext, type FetchConnectionOptions, type XhrConnectionOptions, + type WebSocketConnectionOptions, type InferChatMessages, type GenerationClientState, type ImageGenerateInput, From 1522e6a45e50cfbcc80bb38fb8323badac6cc4c5 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 20 Jul 2026 18:52:14 +0200 Subject: [PATCH 21/26] docs(resumable-streams): websocket transport Document the new WebSocket transport (toWebSocketStream/toWebSocketResponse, resumeWebSocketStream/resumeWebSocketResponse, client webSocket()): a new standalone WebSockets page, a short intro + snippet in the Overview, a reconnect/lifecycle summary in Advanced, and a pointer from Connection Adapters to the built-in adapter instead of the old hand-rolled example. --- docs/chat/connection-adapters.md | 23 +- docs/config.json | 13 +- docs/resumable-streams/advanced.md | 44 ++++ docs/resumable-streams/overview.md | 39 ++++ docs/resumable-streams/websockets.md | 311 +++++++++++++++++++++++++++ 5 files changed, 424 insertions(+), 6 deletions(-) create mode 100644 docs/resumable-streams/websockets.md diff --git a/docs/chat/connection-adapters.md b/docs/chat/connection-adapters.md index 96ff035f95..0e19f194db 100644 --- a/docs/chat/connection-adapters.md +++ b/docs/chat/connection-adapters.md @@ -32,7 +32,8 @@ This page covers every supported transport, when to pick which, and how to build | Code that **synchronously** returns an `AsyncIterable` (in-process `chat()`, an RSC stream, tests) | [`stream`](#server-functions-and-direct-async-iterables) | | An **async** call — a TanStack Start server function or any `Promise`-returning function — resolving to a `Response` or an `AsyncIterable` | [`fetcher`](#server-functions-via-fetcher) | | An RPC framework like Cap'n Web, gRPC-Web, or tRPC | [`rpcStream`](#rpc-streams) | -| A single long-lived WebSocket (or BroadcastChannel, postMessage, shared worker) serving many runs | [Custom `subscribe` / `send` adapter](#persistent-transports-websockets-and-friends) | +| A single long-lived, resumable WebSocket serving many runs | [`webSocket`](#websockets) | +| BroadcastChannel, postMessage, a shared worker, or another persistent transport | [Custom `subscribe` / `send` adapter](#persistent-transports-websockets-and-friends) | | Standard SSE but with custom fetch wrapping (auth refresh, retries) | [`fetchServerSentEvents` with `fetchClient`](#custom-fetch-client) | | Something else entirely (HTTP/3, Server-Sent Events over a different protocol, etc.) | [Custom `connect` adapter](#custom-request-scoped-adapters) | @@ -295,11 +296,25 @@ const { messages } = useChat({ }); ``` +## WebSockets + +For a persistent, resumable WebSocket, use the built-in `webSocket()` adapter instead of hand-rolling a `SubscribeConnectionAdapter`. It opens one socket for the whole conversation, reconnects a dropped durable run automatically, and pairs with the server's `toWebSocketStream` / `toWebSocketResponse`: + +```typescript +import { useChat, webSocket } from "@tanstack/ai-react"; + +const connection = webSocket("/api/chat-ws"); + +const { messages, sendMessage } = useChat({ connection }); +``` + +See [WebSockets](../resumable-streams/websockets) for the server side, the wire protocol, reconnect details, and hosting on Node vs Cloudflare. + ## Persistent Transports (WebSockets and Friends) A persistent transport — WebSocket, BroadcastChannel, postMessage between iframes, a shared worker — is fundamentally different from request/response. You open the channel **once**, then send and receive over it for the lifetime of the client. `stream()`/`connect()` can't model this cleanly because they assume one async iterable per request. -For these cases, implement the `SubscribeConnectionAdapter` interface directly. The shape (full definition in [The Adapter Interface](#the-adapter-interface)): +The built-in `webSocket()` adapter above covers the common resumable WebSocket case. For anything else persistent, implement the `SubscribeConnectionAdapter` interface directly. The shape (full definition in [The Adapter Interface](#the-adapter-interface)): ```typescript import type { SubscribeConnectionAdapter } from "@tanstack/ai-react"; @@ -313,7 +328,9 @@ import type { SubscribeConnectionAdapter } from "@tanstack/ai-react"; The runtime correlates them: chunks emitted on the subscription queue between `send()` and the next terminal event (`RUN_FINISHED` / `RUN_ERROR`) are attributed to that run. -### WebSocket example +### Custom WebSocket example + +Building your own protocol instead of the built-in `webSocket()` adapter (a different wire format, no resume support needed, or a server you don't control)? Implement `SubscribeConnectionAdapter` by hand: ```typescript import { useChat, type SubscribeConnectionAdapter } from "@tanstack/ai-react"; diff --git a/docs/config.json b/docs/config.json index 55a3652822..8ef58756b8 100644 --- a/docs/config.json +++ b/docs/config.json @@ -168,7 +168,7 @@ "label": "Connection Adapters", "to": "chat/connection-adapters", "addedAt": "2026-04-15", - "updatedAt": "2026-07-17" + "updatedAt": "2026-07-20" }, { "label": "Thinking & Reasoning", @@ -188,12 +188,19 @@ { "label": "Overview", "to": "resumable-streams/overview", - "addedAt": "2026-07-17" + "addedAt": "2026-07-17", + "updatedAt": "2026-07-20" }, { "label": "Advanced", "to": "resumable-streams/advanced", - "addedAt": "2026-07-17" + "addedAt": "2026-07-17", + "updatedAt": "2026-07-20" + }, + { + "label": "WebSockets", + "to": "resumable-streams/websockets", + "addedAt": "2026-07-20" }, { "label": "Custom Durability Adapter", diff --git a/docs/resumable-streams/advanced.md b/docs/resumable-streams/advanced.md index 640e172b4c..b8081c077c 100644 --- a/docs/resumable-streams/advanced.md +++ b/docs/resumable-streams/advanced.md @@ -165,6 +165,50 @@ read failure it retries from the last valid position, capping consecutive failures (`reconnect: { maxReadFailures: 10, delayMs: 250 }`). Normal long-poll advancement is never throttled. +## WebSocket transport + +Everything above describes SSE and NDJSON: one connection per turn. A +WebSocket instead stays open for the whole conversation, so a few things work +differently. See [WebSockets](./websockets) for the full protocol and code; +this section covers only what changes relative to the rest of this page. + +**Reconnect rides in the URL, not a header.** SSE and NDJSON resume with +`Last-Event-ID` because `fetch`/XHR can set arbitrary headers before opening a +request. A browser's `WebSocket` constructor can't set custom headers on the +handshake, so the offset instead rides in the URL: `?runId=&offset=`. +`webSocket()` reopens at that URL automatically on a durable run's drop, the +same de-dupe guarantee as `fetchServerSentEvents`, just carried differently on +the wire. + +**The socket is conversation-scoped, not run-scoped.** One socket carries many +turns. `toWebSocketStream`'s `durability` option is therefore a factory keyed +by each turn's `runId` (via a synthetic per-turn request), not a single value +built once. The socket itself closes on client close, an idle timeout, or +process shutdown, not when one turn's `RUN_FINISHED` arrives. + +**Heartbeat and idle timeout replace the transport-level keepalive SSE gets +for free.** `toWebSocketStream` pings every `heartbeatMs` (default 30s) and +closes the socket after `idleTimeoutMs` (default 5 minutes) with no inbound +frame. + +**An abort frame targets one turn, not the socket.** `{ type: 'abort', runId }` +aborts only that turn's `onRun` iteration; the socket stays open for the next +turn. Closing the socket aborts every turn still in flight on it. + +**Hosting differs by runtime.** On Cloudflare Workers or Durable Objects, +`toWebSocketResponse` uses the global `WebSocketPair` and needs no manual +upgrade. On Node (or anywhere without `WebSocketPair`), you upgrade the +connection yourself, typically with `ws`'s `WebSocketServer({ noServer: true })` +hooked into the HTTP server's `upgrade` event, and hand the resulting socket to +`toWebSocketStream` / `resumeWebSocketStream` directly. + +**The producer-vs-socket caveat from [memoryStream in +production](#memorystream-in-production) applies unchanged.** With +`memoryStream`, a dropped WebSocket aborts the `chat()` call backing it just +like a dropped SSE connection does; reconnecting replays the log rather than +resuming a still-running model call. `durableStream` decouples the two exactly +as it does for SSE and NDJSON. + ## Offset ownership `StreamDurability` owns its offset format. Core only passes returned diff --git a/docs/resumable-streams/overview.md b/docs/resumable-streams/overview.md index 01f2b856fc..614b01387c 100644 --- a/docs/resumable-streams/overview.md +++ b/docs/resumable-streams/overview.md @@ -138,6 +138,45 @@ For NDJSON, swap `fetchServerSentEvents` for `fetchHttpStream` (with the server on `toHttpResponse`). The XHR adapters (`xhrServerSentEvents`, `xhrHttpStream`) work the same way, for runtimes without streaming `fetch`. +## Or go full-duplex: WebSockets + +SSE and NDJSON open one connection per turn. If you want a single persistent +socket that carries the whole conversation instead, wrap it with +`toWebSocketStream` (or `toWebSocketResponse` on Cloudflare) and pair it with +the client's `webSocket()` adapter: + +```ts +import { chat, memoryStream, toWebSocketStream } from '@tanstack/ai' +import { openaiText } from '@tanstack/ai-openai' +import type { WebSocketLike } from '@tanstack/ai' + +// `socket` is a server socket you already accepted (see WebSockets for how, +// on Node vs Cloudflare) and `request` is the handshake request. +function handleChatSocket(socket: WebSocketLike, request: Request) { + toWebSocketStream(socket, request, { + durability: (ctx) => memoryStream(ctx.request), + onRun: ({ messages, threadId, runId }) => + chat({ adapter: openaiText('gpt-5.5'), messages, threadId, runId }), + }) +} +``` + +```tsx +import { useChat, webSocket } from '@tanstack/ai-react' + +const connection = webSocket('/api/chat-ws') + +export function Chat() { + const { messages, sendMessage } = useChat({ connection }) + return +} +``` + +The socket resumes the same way SSE and NDJSON do: a dropped connection +reopens with the last offset and replays only what's missing. See +[WebSockets](./websockets) for the wire protocol, reconnect details, and +hosting on Node vs Cloudflare. + That covers the common case. For the durability contract, terminal and error handling, reconnect tuning, attaching to a run by id, Cloudflare deployment, and production process-death concerns, see [Advanced](./advanced). diff --git a/docs/resumable-streams/websockets.md b/docs/resumable-streams/websockets.md new file mode 100644 index 0000000000..d418df5866 --- /dev/null +++ b/docs/resumable-streams/websockets.md @@ -0,0 +1,311 @@ +--- +title: WebSockets +id: websockets +description: "Run resumable streams over one full-duplex WebSocket instead of a request per turn: toWebSocketStream/toWebSocketResponse on the server, webSocket() on the client, resume via a URL query, and hosting on Node or Cloudflare." +keywords: + - websocket transport + - toWebSocketStream + - toWebSocketResponse + - resumeWebSocketStream + - resumeWebSocketResponse + - full-duplex chat + - conversation-scoped socket + - websocket reconnect + - websocket heartbeat +--- + +# WebSockets + +SSE and NDJSON open one connection per turn. A WebSocket is different: one +socket stays open for the whole conversation, carries every turn, and lets the +server push chunks without waiting on a request. Reach for it when you want a +persistent channel instead of a request-per-message model: a WebSocket +gateway you already run, a mobile client that wants to avoid repeated +handshakes, or a UI that needs the server to push outside of a reply (this page +only covers the request/reply case; server-initiated pushes need your own +framing on top). + +By the end of this page you can run a chat turn over a socket that resumes +after a drop, without re-running the model. + +## The shape of it + +One socket, many turns: + +1. The client opens the socket once. +2. Each user message becomes one JSON frame sent over the socket (the same + shape as the SSE/NDJSON POST body). +3. The server runs one `chat()` turn per frame and streams the chunks back + over the same socket. +4. The socket stays open after `RUN_FINISHED` (waiting for the next frame) + until the client closes it, an idle timeout fires, or the process aborts. + +## Server: `onRun` + +Pair the socket you already accepted with `toWebSocketStream`. It decodes +inbound frames, starts one `chat()` turn per frame through `onRun`, and pumps +the resulting chunks back out: + +```ts +import { chat, memoryStream, toWebSocketStream } from '@tanstack/ai' +import { openaiText } from '@tanstack/ai-openai' +import type { WebSocketLike } from '@tanstack/ai' + +// Bridge the per-turn AbortSignal WsRunContext hands you into the +// AbortController chat() expects. +function abortControllerFromSignal(signal: AbortSignal): AbortController { + const controller = new AbortController() + if (signal.aborted) controller.abort() + else signal.addEventListener('abort', () => controller.abort(), { once: true }) + return controller +} + +// `socket` is a WHATWG-shaped server socket you already accepted (see +// "Hosting" below for where it comes from on Node vs Cloudflare) and +// `request` is the original handshake request. +export function handleChatSocket(socket: WebSocketLike, request: Request) { + toWebSocketStream(socket, request, { + // Per-turn durability, keyed by the frame's runId (see "Durability is + // per turn" below). + durability: (ctx) => memoryStream(ctx.request), + onRun: ({ messages, threadId, runId, signal }) => + chat({ + adapter: openaiText('gpt-5.5'), + messages, + threadId, + runId, + abortController: abortControllerFromSignal(signal), + }), + }) +} +``` + +`onRun` receives one `WsRunContext` per inbound frame: + +| Field | What it is | +| --- | --- | +| `messages` | The turn's `UIMessage[]` / `ModelMessage[]`, decoded from the frame. | +| `threadId` / `runId` | The AG-UI identifiers for this turn. | +| `forwardedProps` | Any extra data the client sent with the frame. | +| `resumeOffset` | Always `null` for a fresh turn on this socket (a resume goes through `resumeWebSocketStream` instead, see below). | +| `request` | A synthetic per-turn `Request` carrying `?runId=` in its URL, for keying a durability adapter. | +| `signal` | Aborts when the socket closes, or when this turn receives an `abort` frame (see below). It does not abort other turns on the same socket. | + +Skip `durability` entirely for a socket that doesn't need to survive a drop. +`toWebSocketStream` still pumps chunks, it just can't replay them on +reconnect. + +## Client: `webSocket()` + +`webSocket()` is a `SubscribeConnectionAdapter` for `useChat` (also exported +from `@tanstack/ai-solid`, `@tanstack/ai-vue`, `@tanstack/ai-svelte`, and +`@tanstack/ai-angular`). It opens the socket lazily on the first +`sendMessage`, and reuses it for every later turn in the conversation: + +```tsx +import { useChat, webSocket } from '@tanstack/ai-react' + +const connection = webSocket('/api/chat-ws') + +export function Chat() { + const { messages, sendMessage } = useChat({ connection }) + + return +} +``` + +Nothing else changes: `messages`, `sendMessage`, `stop()`, tool calls, and +persistence all work the same as with any other connection adapter. See +[Connection Adapters](../chat/connection-adapters) for the full option set +(`protocols`, `body`, `reconnect`, `WebSocketImpl`). + +## Wire protocol + +| Direction | Frame | Meaning | +| --- | --- | --- | +| Client → server | A `RunAgentInput`-shaped JSON object (same shape as the SSE/NDJSON POST body) | Start one `chat()` turn. | +| Client → server | `{ "type": "abort", "runId": "…" }` | Abort one in-flight turn (see below). | +| Server → client | `{ "id": "…", "chunk": }` | One chunk, tagged with a durability offset (only when `durability` is configured). | +| Server → client | A bare `StreamChunk` | One chunk, untagged (no `durability` configured). | +| Server → client | `{ "type": "ping" }` | Heartbeat. `webSocket()` drops these automatically; a hand-rolled client should ignore anything with `type: "ping"`. | + +The two server→client shapes are unambiguous: the envelope never has a +top-level `type`, and every bare `StreamChunk` does. + +## Resume rides in the URL, not a header + +SSE and NDJSON resume with the `Last-Event-ID` header, because a `fetch`/XHR +request can set arbitrary headers before it opens. A browser's `WebSocket` +constructor cannot set custom headers on the handshake, so the offset instead +rides in the URL: `?runId=&offset=`. + +`webSocket()` handles this for you. If the socket drops before a run's +terminal chunk (`RUN_FINISHED` / `RUN_ERROR`) and that run was durable +(offset-tagged envelopes), it reopens at `?runId=&offset=` and de-dupes the +replayed boundary, the same reconnect guarantee `fetchServerSentEvents` gives +you, just carried differently on the wire. A run that never emitted an offset +(no `durability` configured) has nothing to resume from: the drop surfaces as +a connection error instead of retrying forever. + +On the server, a URL carrying `?offset=` is a resume, not a fresh turn. Route +it to `resumeWebSocketStream`, a read-only replay of the durability log with +no model call: + +```ts +import { memoryStream, resumeWebSocketStream } from '@tanstack/ai' +import type { WebSocketLike } from '@tanstack/ai' + +export function handleResumeSocket(socket: WebSocketLike, request: Request) { + resumeWebSocketStream(socket, { adapter: memoryStream(request) }) +} +``` + +`resumeWebSocketStream` closes the socket with code `1008` if the request +carries no resume offset: there's nothing to replay. + +## Durability is per turn + +A conversation-scoped socket carries many turns, so its durability adapter +can't be built once for the whole connection. Each turn needs its own log, +keyed by that turn's `runId`. That's why `durability` in `toWebSocketStream` +is a factory, not a value: it receives the per-turn `ctx` and reads +`ctx.request`, whose URL already carries that turn's `?runId=` +(`memoryStream`/`durableStream` key off it automatically). + +## Abort a turn + +An `{ type: 'abort', runId }` frame aborts only that turn's `onRun` iteration +(`ctx.signal` fires); the socket itself stays open for the next turn. This is +a protocol-level primitive: the built-in `webSocket()` client adapter does not +send this frame on `stop()` today, so `stop()` only stops the client from +processing further chunks for that run locally, without telling the server to +cancel the model call. If you need the server to actually stop generating, +send the frame yourself over a reference to the socket, or drive the protocol +directly the way the e2e suite's raw-`WebSocket` tests do. + +Closing the whole socket aborts every turn still in flight on it. + +## Heartbeat and idle timeout + +`toWebSocketStream` sends a `{ type: 'ping' }` frame every `heartbeatMs` +(default 30 seconds) to keep the connection alive through proxies that drop +idle sockets. It closes the socket if no inbound frame arrives for +`idleTimeoutMs` (default 5 minutes); a heartbeat itself doesn't count as +activity, only a client-sent frame does: + +```ts +import { chat, toWebSocketStream } from '@tanstack/ai' +import { openaiText } from '@tanstack/ai-openai' +import type { WebSocketLike } from '@tanstack/ai' + +function handleChatSocket(socket: WebSocketLike, request: Request) { + toWebSocketStream(socket, request, { + onRun: ({ messages, threadId, runId }) => + chat({ adapter: openaiText('gpt-5.5'), messages, threadId, runId }), + heartbeatMs: 15_000, + idleTimeoutMs: 60_000, + }) +} +``` + +## Hosting on Node + +Plain Node (and anywhere else without a global `WebSocketPair`) has no +built-in way to accept a WebSocket upgrade, so you do it yourself and hand the +result to `toWebSocketStream`. The pattern: hook the HTTP server's `upgrade` +event, accept the socket with [`ws`](https://github.com/websockets/ws)'s +`WebSocketServer({ noServer: true })`, and pass the resulting socket straight +through. `ws`'s socket already implements the `send`/`close`/ +`addEventListener`/`bufferedAmount` surface `WebSocketLike` needs: + +```ts ignore +import { WebSocketServer } from 'ws' +import { memoryStream, resumeWebSocketStream, toWebSocketStream } from '@tanstack/ai' +import type { Plugin } from 'vite' +import type { WebSocketLike } from '@tanstack/ai' +import { handleChatSocket } from './handle-chat-socket' + +const WS_PATH = '/api/chat-ws' + +export function webSocketChatPlugin(): Plugin { + return { + name: 'websocket-chat-plugin', + configureServer(server) { + if (!server.httpServer) return + const wss = new WebSocketServer({ noServer: true }) + + server.httpServer.on('upgrade', (req, socket, head) => { + const url = new URL(req.url ?? '/', `http://${req.headers.host}`) + if (url.pathname !== WS_PATH) return + + wss.handleUpgrade(req, socket, head, (ws) => { + const request = new Request(url) + const socketLike = ws as unknown as WebSocketLike + + if (url.searchParams.get('offset') !== null) { + resumeWebSocketStream(socketLike, { adapter: memoryStream(request) }) + } else { + handleChatSocket(socketLike, request) + } + }) + }) + }, + } +} +``` + +This is the same pattern used by the working +[`examples/ts-react-chat`](https://github.com/TanStack/ai/blob/main/examples/ts-react-chat/src/lib/websocket-chat-plugin.ts) +WebSocket example and by the e2e suite's +[`durable-delivery-ws-plugin.ts`](https://github.com/TanStack/ai/blob/main/testing/e2e/src/lib/durable-delivery-ws-plugin.ts). +Bun's `ServerWebSocket` and Deno's `Deno.upgradeWebSocket` need the same shape +of adapter: accept the platform's socket, then call `toWebSocketStream` / +`resumeWebSocketStream` with it. + +## Hosting on Cloudflare + +Cloudflare Workers (and Durable Objects) expose a global `WebSocketPair`, so +you don't upgrade anything by hand. `toWebSocketResponse` creates the pair, +accepts the server half, wires it to `toWebSocketStream`, and returns the 101 +upgrade `Response` for you: + +```ts +import { chat, toWebSocketResponse } from '@tanstack/ai' +import { openaiText } from '@tanstack/ai-openai' + +export default { + fetch(request: Request): Response { + return toWebSocketResponse(request, { + onRun: ({ messages, threadId, runId }) => + chat({ adapter: openaiText('gpt-5.5'), messages, threadId, runId }), + }) + }, +} +``` + +`resumeWebSocketResponse({ adapter })` is the matching read-only wrapper for a +`?offset=` resume. Both throw with a message pointing at their `*Stream` +counterpart if called somewhere without `WebSocketPair` (Node, for instance), +so there's no way to accidentally ship the Cloudflare wrapper to a runtime +that can't upgrade a socket itself. + +## Producer vs. socket lifecycle + +The same caveat that applies to SSE and NDJSON applies here: with +`memoryStream`, the run's producer and the delivery socket live in the same +process. If that socket drops, the `chat()` call backing it aborts too, so +reconnecting only replays what was already logged rather than resuming a run +still in progress. `durableStream` decouples the two (the producer runs +against a backend, not the client's socket), so a drop there can reconnect to +a run that's still actively producing. See +[memoryStream in production](./advanced#memorystream-in-production) for the +full explanation. It applies to WebSockets exactly as written there. + +## Next steps + +- [Overview](./overview): the durability adapter contract this transport + reuses. +- [Advanced](./advanced): reconnection bounding, offset ownership, and + Cloudflare Durable Streams deployment, shared across every transport. +- [Connection Adapters](../chat/connection-adapters): where `webSocket()` + fits among the other client adapters. From cc6cefa7f418612cd27af4fad5cfc0be52a55f03 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 20 Jul 2026 18:54:33 +0200 Subject: [PATCH 22/26] chore: changeset for websocket transport --- .changeset/websocket-transport.md | 42 +++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) create mode 100644 .changeset/websocket-transport.md diff --git a/.changeset/websocket-transport.md b/.changeset/websocket-transport.md new file mode 100644 index 0000000000..f4ca05f40e --- /dev/null +++ b/.changeset/websocket-transport.md @@ -0,0 +1,42 @@ +--- +'@tanstack/ai': minor +'@tanstack/ai-client': minor +'@tanstack/ai-react': minor +'@tanstack/ai-solid': minor +'@tanstack/ai-vue': minor +'@tanstack/ai-svelte': minor +'@tanstack/ai-angular': minor +--- + +WebSocket transport: a full-duplex, resumable third transport alongside SSE and +NDJSON, reusing the same delivery-durability seam. + +On the server, `@tanstack/ai` adds `toWebSocketStream(socket, request, { onRun, +durability, batch, heartbeatMs, idleTimeoutMs, debug })` — a portable core that +pumps a conversation over an already-accepted WHATWG `WebSocketLike` server +socket (Node via `ws`, Bun, etc.), and `toWebSocketResponse(request, { onRun, +… })`, a thin wrapper that upgrades via `WebSocketPair` and returns a 101 +`Response` on Cloudflare Workers/Durable Objects (it throws elsewhere, pointing +you at `toWebSocketStream`). Because one socket outlives many `chat()` turns +(client-tool resubmits, follow-up user messages), you pass an `onRun(ctx) => +AsyncIterable` factory instead of a prebuilt stream — the helper +calls it per inbound `RunAgentInput` frame. The socket is conversation-scoped: +it stays open across turns and closes on client close, an `{ type: 'abort', +runId }` control frame (which aborts only that turn), or the idle timeout, with a +periodic `{ type: 'ping' }` heartbeat. Durability is keyed per turn and reuses +the existing `durableStreamSource`, so server→client frames carry the same +`{ id, chunk }` envelope as NDJSON. `resumeWebSocketStream(socket, { adapter })` +and `resumeWebSocketResponse({ adapter })` replay a run read-only from the +durability log (no model call). + +On the client, `webSocket(url, options)` (in `@tanstack/ai-client`, re-exported +from `@tanstack/ai-react`, `-solid`, `-vue`, `-svelte`, and `-angular`) is a +full-duplex `subscribe` + `send` connection adapter for `useChat`. `send()` +writes a `RunAgentInput` frame; `subscribe()` yields inbound chunks, ignores +heartbeats, unwraps durable envelopes, and auto-reconnects a dropped durable run +by reopening with `?runId=&offset=` (browsers can't set a `Last-Event-ID` +handshake header, so the offset rides in the URL). The reconnect bookkeeping +(offset de-dupe, no-progress ceiling → `StreamReconnectLimitError`) is shared +with the HTTP adapters via the new `createReconnectTracker`, and a fatal drop +surfaces to the consumer (`StreamReadError` / `StreamReconnectLimitError`) +instead of hanging. From 1433ca307eb8a60cf52166c5ed95f9ceda9205b9 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Mon, 20 Jul 2026 19:31:21 +0200 Subject: [PATCH 23/26] fix(ai,ai-client): close resume ws on log exhaustion; guard idle-reap/heartbeat; remove abort listener leak --- packages/ai-client/src/connection-adapters.ts | 10 ++++-- packages/ai/src/stream-to-websocket.ts | 30 +++++++++++++++-- packages/ai/tests/stream-to-websocket.test.ts | 32 +++++++++++++++++++ 3 files changed, 67 insertions(+), 5 deletions(-) diff --git a/packages/ai-client/src/connection-adapters.ts b/packages/ai-client/src/connection-adapters.ts index 70b263ed1c..e2242d8975 100644 --- a/packages/ai-client/src/connection-adapters.ts +++ b/packages/ai-client/src/connection-adapters.ts @@ -1785,10 +1785,11 @@ export function webSocket( } const sink: WebSocketChunkSink = { push, fail } listeners.add(sink) - abortSignal?.addEventListener('abort', () => { + const onAbort = () => { const w = waiters.shift() if (w) w(null) - }) + } + abortSignal?.addEventListener('abort', onAbort) return (async function* () { try { while (!abortSignal?.aborted) { @@ -1820,6 +1821,7 @@ export function webSocket( } } finally { listeners.delete(sink) + abortSignal?.removeEventListener('abort', onAbort) } })() }, @@ -1869,7 +1871,8 @@ export function webSocket( } const sink: WebSocketChunkSink = { push, fail } listeners.add(sink) - abortSignal?.addEventListener('abort', () => waiters.shift()?.(null)) + const onAbort = () => waiters.shift()?.(null) + abortSignal?.addEventListener('abort', onAbort) return (async function* () { try { while (!abortSignal?.aborted) { @@ -1901,6 +1904,7 @@ export function webSocket( } } finally { listeners.delete(sink) + abortSignal?.removeEventListener('abort', onAbort) } })() }, diff --git a/packages/ai/src/stream-to-websocket.ts b/packages/ai/src/stream-to-websocket.ts index f370f8a4ae..9fd7e6e951 100644 --- a/packages/ai/src/stream-to-websocket.ts +++ b/packages/ai/src/stream-to-websocket.ts @@ -120,11 +120,25 @@ export function toWebSocketStream( let lastActivity = Date.now() const heartbeat = setInterval(() => { - socket.send(JSON.stringify({ type: 'ping' })) + try { + socket.send(JSON.stringify({ type: 'ping' })) + } catch { + // Socket is CLOSING/CLOSED between ticks — the close handler below + // clears this interval; swallow so the timer callback doesn't throw + // uncaught in the meantime. + } }, heartbeatMs) const idle = setInterval( () => { - if (Date.now() - lastActivity > idleTimeoutMs) socket.close(1000, 'idle') + // Never idle-reap while a turn is in flight: a long single onRun + // iteration (agentic loop / >5-min generation) sends no INBOUND + // frames, so idle would otherwise fire and kill live work. + if ( + activeTurns.size === 0 && + Date.now() - lastActivity > idleTimeoutMs + ) { + socket.close(1000, 'idle') + } }, Math.min(idleTimeoutMs, 30_000), ) @@ -250,6 +264,18 @@ export function resumeWebSocketStream( for await (const chunk of source) { socket.send(encodeWsFrame(chunk, getId(chunk))) } + // Source exhausted = the durability log is complete/terminal; nothing more + // will arrive on this read-only socket. Close so the client's reconnect + // loop sees onclose and terminates (bounded) instead of awaiting a chunk + // that never comes. Safe across durability models: a live decoupled + // producer (e.g. durableStream) keeps `read` parked until the terminal, + // so the source doesn't exhaust until the run truly ends; a completed + // in-process log closes immediately. + try { + socket.close(1000) + } catch { + // socket already closing/closed — nothing to do + } })().catch((error: unknown) => { logger.errors('resume websocket replay failed', { error }) try { diff --git a/packages/ai/tests/stream-to-websocket.test.ts b/packages/ai/tests/stream-to-websocket.test.ts index 03c7394c18..6e078b1a28 100644 --- a/packages/ai/tests/stream-to-websocket.test.ts +++ b/packages/ai/tests/stream-to-websocket.test.ts @@ -281,6 +281,38 @@ describe('resumeWebSocketStream', () => { const chunkTypes = joiner.sent.map((s) => JSON.parse(s).chunk.type) expect(chunkTypes).toEqual(['TEXT_MESSAGE_CONTENT', 'RUN_FINISHED']) + // Regression: the replay pump must close the socket once the durability + // log is exhausted, even on the success path (terminal already present). + // Otherwise a client that auto-reconnects to `?runId&offset` after a + // drop would see no terminal frame, no close, no error — just a hang. + expect(joiner.closed).toBe(true) + expect(joiner.closeCode).toBe(1000) + }) + + it('closes the socket (1000) after replaying a complete log that has NO terminal event', async () => { + // Simulates a mid-generation drop: the durability log terminalized + // (log.complete === true) but the producer's turn was aborted by the + // socket closing before RUN_FINISHED was ever appended. `read` yields + // whatever was produced and returns with no terminal chunk. + const noTerminalAdapter: StreamDurability = { + resumeFrom: () => '-1', + append: () => Promise.resolve([]), + read: async function* () { + yield { offset: 'off-1', chunk: ev.textContent('partial') } + }, + close: () => Promise.resolve(), + } + const joiner = new FakeSocket() + + resumeWebSocketStream(joiner, { adapter: noTerminalAdapter }) + await flush() + + const chunkTypes = joiner.sent.map((s) => JSON.parse(s).chunk.type) + expect(chunkTypes).toEqual(['TEXT_MESSAGE_CONTENT']) + // The proof: no terminal event was ever produced, yet the socket still + // closes once the log exhausts — no consumer hang, no leaked socket. + expect(joiner.closed).toBe(true) + expect(joiner.closeCode).toBe(1000) }) it('closes with 1008 when no offset is present', async () => { From 2225a78d5b934a0d6ebe67c72805fffc8b7a02de Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Mon, 20 Jul 2026 17:37:14 +0000 Subject: [PATCH 24/26] ci: apply automated fixes --- packages/ai-client/src/connection-adapters.ts | 17 +-- .../connection-adapters-websocket.test.ts | 118 +++++++++++++----- packages/ai/src/stream-to-websocket.ts | 33 ++--- packages/ai/tests/stream-to-websocket.test.ts | 12 +- testing/e2e/tests/websocket.spec.ts | 4 +- 5 files changed, 123 insertions(+), 61 deletions(-) diff --git a/packages/ai-client/src/connection-adapters.ts b/packages/ai-client/src/connection-adapters.ts index e2242d8975..0258e3d14e 100644 --- a/packages/ai-client/src/connection-adapters.ts +++ b/packages/ai-client/src/connection-adapters.ts @@ -1652,7 +1652,10 @@ export function webSocket( url: string | (() => string), options: WebSocketConnectionOptions = {}, ): SubscribeConnectionAdapter & { - joinRun: (runId: string, abortSignal?: AbortSignal) => AsyncIterable + joinRun: ( + runId: string, + abortSignal?: AbortSignal, + ) => AsyncIterable } { const Impl = options.WebSocketImpl ?? WebSocket let socket: WebSocket | undefined @@ -1724,9 +1727,7 @@ export function webSocket( // from. Surface a hard failure rather than silently reconnecting // forever against a server that never tags its events. currentSession = undefined - failAll( - new StreamReadError(new Error('WebSocket connection closed')), - ) + failAll(new StreamReadError(new Error('WebSocket connection closed'))) return } void reconnect(session, lastEventId) @@ -1851,10 +1852,10 @@ export function webSocket( ws.send(JSON.stringify(body)) }, joinRun(runId, abortSignal): AsyncIterable { - const target = withSearchParams( - typeof url === 'function' ? url() : url, - { offset: '-1', runId }, - ) + const target = withSearchParams(typeof url === 'function' ? url() : url, { + offset: '-1', + runId, + }) openOnce(target) const chunks: Array = [] const waiters: Array<(c: StreamChunk | null) => void> = [] diff --git a/packages/ai-client/tests/connection-adapters-websocket.test.ts b/packages/ai-client/tests/connection-adapters-websocket.test.ts index 466ba976b7..553f33420d 100644 --- a/packages/ai-client/tests/connection-adapters-websocket.test.ts +++ b/packages/ai-client/tests/connection-adapters-websocket.test.ts @@ -1,5 +1,8 @@ import { describe, expect, it } from 'vitest' -import { StreamReconnectLimitError, webSocket } from '../src/connection-adapters' +import { + StreamReconnectLimitError, + webSocket, +} from '../src/connection-adapters' // Minimal fake WebSocket (constructor-compatible with the WHATWG shape). class FakeWebSocket { @@ -56,9 +59,7 @@ function tick(): Promise { * (instead of `drain()`) where the test needs the returned promise to REJECT * when the generator throws, rather than swallowing it. */ -async function drainToEnd( - iter: AsyncIterator, -): Promise> { +async function drainToEnd(iter: AsyncIterator): Promise> { const received: Array = [] for (;;) { const { value, done } = await iter.next() @@ -91,12 +92,18 @@ describe('webSocket() subscribe/send', () => { const ac = new AbortController() const received: Array = [] const sub = drain(conn.subscribe(ac.signal), received, ac.signal) - await conn.send([{ role: 'user', content: 'hi' } as any], undefined, ac.signal, { - threadId: 't', - runId: 'r', - }) + await conn.send( + [{ role: 'user', content: 'hi' } as any], + undefined, + ac.signal, + { + threadId: 't', + runId: 'r', + }, + ) const ws = FakeWebSocket.instances[0] - if (!ws) throw new Error('expected a FakeWebSocket instance to have been created') + if (!ws) + throw new Error('expected a FakeWebSocket instance to have been created') await new Promise((r) => setTimeout(r, 0)) expect(ws.sent.length).toBe(1) const sentFrame = ws.sent[0] @@ -119,12 +126,18 @@ describe('webSocket() subscribe/send', () => { const ac = new AbortController() const received: Array = [] const sub = drain(conn.subscribe(ac.signal), received, ac.signal) - await conn.send([{ role: 'user', content: 'hi' } as any], undefined, ac.signal, { - threadId: 't', - runId: 'r', - }) + await conn.send( + [{ role: 'user', content: 'hi' } as any], + undefined, + ac.signal, + { + threadId: 't', + runId: 'r', + }, + ) const ws = FakeWebSocket.instances[0] - if (!ws) throw new Error('expected a FakeWebSocket instance to have been created') + if (!ws) + throw new Error('expected a FakeWebSocket instance to have been created') await new Promise((r) => setTimeout(r, 0)) // Bare chunk (no id) — not wrapped in an {id, chunk} envelope. @@ -147,11 +160,15 @@ describe('webSocket() subscribe/send', () => { await new Promise((r) => setTimeout(r, 0)) const ws = FakeWebSocket.instances[0] - if (!ws) throw new Error('expected a FakeWebSocket instance to have been created') + if (!ws) + throw new Error('expected a FakeWebSocket instance to have been created') expect(ws.url).toContain('offset=-1') expect(ws.url).toContain('runId=run-j') - ws.emit({ type: 'TEXT_MESSAGE_CONTENT', delta: 'joined', timestamp: 0 }, 'off-1') + ws.emit( + { type: 'TEXT_MESSAGE_CONTENT', delta: 'joined', timestamp: 0 }, + 'off-1', + ) await new Promise((r) => setTimeout(r, 0)) ac.abort() await sub @@ -189,7 +206,10 @@ describe('webSocket() subscribe/send', () => { await Promise.race([ Promise.all([send1, send2]), new Promise((_, reject) => { - timeoutId = setTimeout(() => reject(new Error('send() calls hung')), 1000) + timeoutId = setTimeout( + () => reject(new Error('send() calls hung')), + 1000, + ) }), ]) } finally { @@ -198,7 +218,8 @@ describe('webSocket() subscribe/send', () => { expect(FakeWebSocket.instances.length).toBe(1) const ws = FakeWebSocket.instances[0] - if (!ws) throw new Error('expected a FakeWebSocket instance to have been created') + if (!ws) + throw new Error('expected a FakeWebSocket instance to have been created') expect(ws.sent.length).toBe(2) const sent0 = ws.sent[0] if (sent0 === undefined) throw new Error('expected a sent frame') @@ -219,30 +240,55 @@ describe('webSocket() reconnect', () => { const ac = new AbortController() const received: Array = [] const sub = drain(conn.subscribe(ac.signal), received, ac.signal) - await conn.send([{ role: 'user', content: 'hi' } as any], undefined, ac.signal, { - threadId: 't', - runId: 'run-x', - }) + await conn.send( + [{ role: 'user', content: 'hi' } as any], + undefined, + ac.signal, + { + threadId: 't', + runId: 'run-x', + }, + ) const ws1 = FakeWebSocket.instances[0] - if (!ws1) throw new Error('expected a FakeWebSocket instance to have been created') + if (!ws1) + throw new Error('expected a FakeWebSocket instance to have been created') await new Promise((r) => setTimeout(r, 0)) - ws1.emit({ type: 'TEXT_MESSAGE_CONTENT', delta: 'a', timestamp: 0 }, 'off-1') + ws1.emit( + { type: 'TEXT_MESSAGE_CONTENT', delta: 'a', timestamp: 0 }, + 'off-1', + ) await new Promise((r) => setTimeout(r, 0)) ws1.close() // drop mid-run await new Promise((r) => setTimeout(r, 5)) const ws2 = FakeWebSocket.instances[1] expect(ws2).toBeDefined() - if (!ws2) throw new Error('expected a second FakeWebSocket instance to have been created') + if (!ws2) + throw new Error( + 'expected a second FakeWebSocket instance to have been created', + ) const url2 = new URL(ws2.url) expect(url2.searchParams.get('runId')).toBe('run-x') expect(url2.searchParams.get('offset')).toBe('off-1') // Server replays the de-duped boundary + a new chunk + terminal. - ws2.emit({ type: 'TEXT_MESSAGE_CONTENT', delta: 'a', timestamp: 0 }, 'off-1') - ws2.emit({ type: 'TEXT_MESSAGE_CONTENT', delta: 'b', timestamp: 0 }, 'off-2') ws2.emit( - { type: 'RUN_FINISHED', runId: 'run-x', threadId: 't', model: 'm', finishReason: 'stop', timestamp: 0 }, + { type: 'TEXT_MESSAGE_CONTENT', delta: 'a', timestamp: 0 }, + 'off-1', + ) + ws2.emit( + { type: 'TEXT_MESSAGE_CONTENT', delta: 'b', timestamp: 0 }, + 'off-2', + ) + ws2.emit( + { + type: 'RUN_FINISHED', + runId: 'run-x', + threadId: 't', + model: 'm', + finishReason: 'stop', + timestamp: 0, + }, 'off-3', ) await new Promise((r) => setTimeout(r, 0)) @@ -276,7 +322,8 @@ describe('webSocket() fatal drop surfacing', () => { { threadId: 't', runId: 'r' }, ) const ws = FakeWebSocket.instances[0] - if (!ws) throw new Error('expected a FakeWebSocket instance to have been created') + if (!ws) + throw new Error('expected a FakeWebSocket instance to have been created') await tick() // Bare chunk — NO id envelope, so no durable offset is ever tracked. @@ -323,13 +370,17 @@ describe('webSocket() fatal drop surfacing', () => { { threadId: 't', runId: 'run-stuck' }, ) const ws1 = FakeWebSocket.instances[0] - if (!ws1) throw new Error('expected a FakeWebSocket instance to have been created') + if (!ws1) + throw new Error('expected a FakeWebSocket instance to have been created') await tick() // One durable chunk establishes an offset (and makes progress, resetting // the no-progress counter), then the socket drops — this reconnect is // legitimate and does not count against the ceiling. - ws1.emit({ type: 'TEXT_MESSAGE_CONTENT', delta: 'a', timestamp: 0 }, 'off-1') + ws1.emit( + { type: 'TEXT_MESSAGE_CONTENT', delta: 'a', timestamp: 0 }, + 'off-1', + ) await tick() ws1.close() await tick() @@ -339,7 +390,10 @@ describe('webSocket() fatal drop surfacing', () => { // ceiling bounds. With maxAttempts: 2, the 3rd consecutive no-progress // close must exceed it. const ws2 = FakeWebSocket.instances[1] - if (!ws2) throw new Error('expected a second FakeWebSocket instance (post-drop reconnect)') + if (!ws2) + throw new Error( + 'expected a second FakeWebSocket instance (post-drop reconnect)', + ) ws2.close() // no-progress attempt 1 (1 <= 2, reconnects again) await tick() diff --git a/packages/ai/src/stream-to-websocket.ts b/packages/ai/src/stream-to-websocket.ts index 9fd7e6e951..009cd92bf1 100644 --- a/packages/ai/src/stream-to-websocket.ts +++ b/packages/ai/src/stream-to-websocket.ts @@ -133,10 +133,7 @@ export function toWebSocketStream( // Never idle-reap while a turn is in flight: a long single onRun // iteration (agentic loop / >5-min generation) sends no INBOUND // frames, so idle would otherwise fire and kill live work. - if ( - activeTurns.size === 0 && - Date.now() - lastActivity > idleTimeoutMs - ) { + if (activeTurns.size === 0 && Date.now() - lastActivity > idleTimeoutMs) { socket.close(1000, 'idle') } }, @@ -195,11 +192,15 @@ export function toWebSocketStream( try { if (init.durability) { const adapter = init.durability(ctx) - const { source, getId } = durableStreamSource(init.onRun(ctx), adapter, { - abortController: turnAbort, - ...(init.batch === undefined ? {} : { batch: init.batch }), - logger, - }) + const { source, getId } = durableStreamSource( + init.onRun(ctx), + adapter, + { + abortController: turnAbort, + ...(init.batch === undefined ? {} : { batch: init.batch }), + logger, + }, + ) for await (const chunk of source) { socket.send(encodeWsFrame(chunk, getId(chunk))) } @@ -344,13 +345,13 @@ export function toWebSocketResponse( * resumeWebSocketResponse({ adapter: memoryStream(request) }) * ``` */ -export function resumeWebSocketResponse( - options: { - adapter: StreamDurability - batch?: number - debug?: DebugOption - }, -): Response { +export function resumeWebSocketResponse< + TOffset extends string = string, +>(options: { + adapter: StreamDurability + batch?: number + debug?: DebugOption +}): Response { const { client, server } = upgradeOrThrow('resumeWebSocketResponse') resumeWebSocketStream(server, options) return upgradeResponse(client) diff --git a/packages/ai/tests/stream-to-websocket.test.ts b/packages/ai/tests/stream-to-websocket.test.ts index 6e078b1a28..e93da4c2b8 100644 --- a/packages/ai/tests/stream-to-websocket.test.ts +++ b/packages/ai/tests/stream-to-websocket.test.ts @@ -37,9 +37,9 @@ describe('ws frame codec', () => { }) it('decodes an abort control frame', () => { - expect(decodeWsFrame(JSON.stringify({ type: 'abort', runId: 'r' }))).toEqual( - { kind: 'abort', runId: 'r' }, - ) + expect( + decodeWsFrame(JSON.stringify({ type: 'abort', runId: 'r' })), + ).toEqual({ kind: 'abort', runId: 'r' }) }) }) @@ -144,7 +144,11 @@ describe('toWebSocketStream (non-durable)', () => { await flush() const types = socket.sent.map((s) => JSON.parse(s).type) - expect(types).toEqual(['RUN_STARTED', 'TEXT_MESSAGE_CONTENT', 'RUN_FINISHED']) + expect(types).toEqual([ + 'RUN_STARTED', + 'TEXT_MESSAGE_CONTENT', + 'RUN_FINISHED', + ]) expect(socket.closed).toBe(false) // conversation-scoped: stays open }) }) diff --git a/testing/e2e/tests/websocket.spec.ts b/testing/e2e/tests/websocket.spec.ts index de547b60e5..c7e891504c 100644 --- a/testing/e2e/tests/websocket.spec.ts +++ b/testing/e2e/tests/websocket.spec.ts @@ -90,7 +90,9 @@ test.describe('delivery durability (websocket)', () => { const base = `${location.origin.replace(/^http/, 'ws')}/api/durable-delivery-ws` // First connection: read only the first two chunks, then disconnect. - const first = new WebSocket(`${base}?runId=${encodeURIComponent(runId)}`) + const first = new WebSocket( + `${base}?runId=${encodeURIComponent(runId)}`, + ) await new Promise((resolve, reject) => { first.onopen = () => resolve() first.onerror = () => reject(new Error('ws open failed')) From ec7b58b7bc054311c8318d71b8c8374756f0ec7f Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 18 Aug 2026 15:03:19 +0200 Subject: [PATCH 25/26] fix(ai,ai-client): address websocket transport review feedback --- docs/chat/connection-adapters.md | 24 +++++++- docs/config.json | 3 +- docs/resumable-streams/websockets.md | 23 +++++-- .../src/lib/websocket-chat-plugin.ts | 27 +++++--- packages/ai-client/src/connection-adapters.ts | 16 ++++- .../connection-adapters-websocket.test.ts | 61 +++++++++++++++++++ packages/ai/src/stream-to-websocket.ts | 4 ++ packages/ai/tests/stream-to-websocket.test.ts | 28 +++++++++ .../e2e/src/lib/durable-delivery-ws-plugin.ts | 21 +++++-- 9 files changed, 186 insertions(+), 21 deletions(-) diff --git a/docs/chat/connection-adapters.md b/docs/chat/connection-adapters.md index 42e195cb49..aae69fddb0 100644 --- a/docs/chat/connection-adapters.md +++ b/docs/chat/connection-adapters.md @@ -327,7 +327,29 @@ const connection = webSocket("/api/chat-ws"); const { messages, sendMessage } = useChat({ connection }); ``` -See [WebSockets](../resumable-streams/websockets) for the server side, the wire protocol, reconnect details, and hosting on Node vs Cloudflare. +On Cloudflare Workers or Durable Objects, pair that client with `toWebSocketResponse`. Elsewhere, accept the socket yourself and pass it to `toWebSocketStream` (see [WebSockets](../resumable-streams/websockets)): + +```typescript +import { chat, memoryStream, toWebSocketResponse } from "@tanstack/ai"; +import { openaiText } from "@tanstack/ai-openai"; + +export default { + fetch(request: Request): Response { + return toWebSocketResponse(request, { + durability: (ctx) => memoryStream(ctx.request), + onRun: ({ messages, threadId, runId }) => + chat({ + adapter: openaiText("gpt-5.5"), + messages, + threadId, + runId, + }), + }); + }, +}; +``` + +See [WebSockets](../resumable-streams/websockets) for the wire protocol, reconnect details, and hosting on Node vs Cloudflare. ## Persistent Transports (WebSockets and Friends) diff --git a/docs/config.json b/docs/config.json index c311970a13..7f2556fa5c 100644 --- a/docs/config.json +++ b/docs/config.json @@ -232,7 +232,8 @@ { "label": "WebSockets", "to": "resumable-streams/websockets", - "addedAt": "2026-07-20" + "addedAt": "2026-07-20", + "updatedAt": "2026-08-18" }, { "label": "Custom Durability Adapter", diff --git a/docs/resumable-streams/websockets.md b/docs/resumable-streams/websockets.md index d418df5866..e3b093d89c 100644 --- a/docs/resumable-streams/websockets.md +++ b/docs/resumable-streams/websockets.md @@ -55,8 +55,11 @@ import type { WebSocketLike } from '@tanstack/ai' // AbortController chat() expects. function abortControllerFromSignal(signal: AbortSignal): AbortController { const controller = new AbortController() - if (signal.aborted) controller.abort() - else signal.addEventListener('abort', () => controller.abort(), { once: true }) + if (signal.aborted) controller.abort(signal.reason) + else + signal.addEventListener('abort', () => controller.abort(signal.reason), { + once: true, + }) return controller } @@ -227,6 +230,16 @@ import { handleChatSocket } from './handle-chat-socket' const WS_PATH = '/api/chat-ws' +function isWebSocketLike(value: unknown): value is WebSocketLike { + if (typeof value !== 'object' || value === null) return false + if (!('send' in value) || typeof value.send !== 'function') return false + if (!('close' in value) || typeof value.close !== 'function') return false + if (!('addEventListener' in value) || typeof value.addEventListener !== 'function') { + return false + } + return true +} + export function webSocketChatPlugin(): Plugin { return { name: 'websocket-chat-plugin', @@ -240,12 +253,12 @@ export function webSocketChatPlugin(): Plugin { wss.handleUpgrade(req, socket, head, (ws) => { const request = new Request(url) - const socketLike = ws as unknown as WebSocketLike + if (!isWebSocketLike(ws)) return if (url.searchParams.get('offset') !== null) { - resumeWebSocketStream(socketLike, { adapter: memoryStream(request) }) + resumeWebSocketStream(ws, { adapter: memoryStream(request) }) } else { - handleChatSocket(socketLike, request) + handleChatSocket(ws, request) } }) }) diff --git a/examples/ts-react-chat/src/lib/websocket-chat-plugin.ts b/examples/ts-react-chat/src/lib/websocket-chat-plugin.ts index f716a53455..45d89129f3 100644 --- a/examples/ts-react-chat/src/lib/websocket-chat-plugin.ts +++ b/examples/ts-react-chat/src/lib/websocket-chat-plugin.ts @@ -30,6 +30,19 @@ import type { WebSocketLike } from '@tanstack/ai' */ const WS_PATH = '/api/chat-ws' +function isWebSocketLike(value: unknown): value is WebSocketLike { + if (typeof value !== 'object' || value === null) return false + if (!('send' in value) || typeof value.send !== 'function') return false + if (!('close' in value) || typeof value.close !== 'function') return false + if ( + !('addEventListener' in value) || + typeof value.addEventListener !== 'function' + ) { + return false + } + return true +} + /** * `chat()` cancels via an `AbortController` it can read `.signal` off of, but * `WsRunContext` (the per-turn context `toWebSocketStream` hands `onRun`) @@ -39,9 +52,11 @@ const WS_PATH = '/api/chat-ws' */ function abortControllerFromSignal(signal: AbortSignal): AbortController { const controller = new AbortController() - if (signal.aborted) controller.abort() + if (signal.aborted) controller.abort(signal.reason) else - signal.addEventListener('abort', () => controller.abort(), { once: true }) + signal.addEventListener('abort', () => controller.abort(signal.reason), { + once: true, + }) return controller } @@ -62,16 +77,14 @@ export function webSocketChatPlugin(): Plugin { wss.handleUpgrade(req, socket, head, (ws) => { const request = new Request(url) - // `ws`'s WebSocket implements the WHATWG send/close/addEventListener/ - // bufferedAmount surface `WebSocketLike` needs. - const socketLike = ws as unknown as WebSocketLike + if (!isWebSocketLike(ws)) return if (url.searchParams.get('offset') !== null) { - resumeWebSocketStream(socketLike, { + resumeWebSocketStream(ws, { adapter: memoryStream(request), }) } else { - toWebSocketStream(socketLike, request, { + toWebSocketStream(ws, request, { durability: (ctx) => memoryStream(ctx.request), onRun: ({ messages, threadId, runId, signal }) => chat({ diff --git a/packages/ai-client/src/connection-adapters.ts b/packages/ai-client/src/connection-adapters.ts index 80f58357a7..776e6e6a6b 100644 --- a/packages/ai-client/src/connection-adapters.ts +++ b/packages/ai-client/src/connection-adapters.ts @@ -2026,7 +2026,13 @@ export function webSocket( // an unhandled rejection if it errors. Awaiters of openPromise still see the rejection. openPromise.catch(() => {}) ws.onmessage = (event: MessageEvent) => { - const parsed: unknown = JSON.parse(String(event.data)) + let parsed: unknown + try { + parsed = JSON.parse(String(event.data)) + } catch (error) { + failAll(new StreamReadError(error)) + return + } if ( typeof parsed === 'object' && parsed !== null && @@ -2058,7 +2064,13 @@ export function webSocket( } ws.onclose = () => { const session = currentSession - if (!session || session.signal?.aborted || session.sawTerminal) return + if (!session) { + // joinRun() never creates a session. Surface the drop so those + // iterators do not stay parked on a dead socket. + failAll(new StreamReadError(new Error('WebSocket connection closed'))) + return + } + if (session.signal?.aborted || session.sawTerminal) return const lastEventId = session.tracker.lastEventId if (lastEventId === undefined) { // Non-durable run (no offset ever observed) — nothing to resume diff --git a/packages/ai-client/tests/connection-adapters-websocket.test.ts b/packages/ai-client/tests/connection-adapters-websocket.test.ts index 553f33420d..38fdc55524 100644 --- a/packages/ai-client/tests/connection-adapters-websocket.test.ts +++ b/packages/ai-client/tests/connection-adapters-websocket.test.ts @@ -34,6 +34,9 @@ class FakeWebSocket { data: JSON.stringify(id === undefined ? chunk : { id, chunk }), }) } + emitRaw(data: string): void { + this.onmessage?.({ data }) + } } function drain( @@ -422,4 +425,62 @@ describe('webSocket() fatal drop surfacing', () => { // was reached via genuine no-progress reconnects, not a shortcut. expect(FakeWebSocket.instances.length).toBe(4) }) + + it('a malformed inbound frame rejects subscribe() with StreamReadError', async () => { + FakeWebSocket.instances = [] + const conn = webSocket('wss://x/api/chat', { + WebSocketImpl: FakeWebSocket as unknown as typeof WebSocket, + }) + const ac = new AbortController() + const iter = conn.subscribe(ac.signal)[Symbol.asyncIterator]() + + await conn.send( + [{ role: 'user', content: 'hi' } as any], + undefined, + ac.signal, + { threadId: 't', runId: 'r' }, + ) + const ws = FakeWebSocket.instances[0] + if (!ws) + throw new Error('expected a FakeWebSocket instance to have been created') + await tick() + + ws.emitRaw('{') + + const result = await withTimeout( + drainToEnd(iter), + 'subscribe() hung instead of rejecting after a malformed frame', + ) + .then((received) => ({ received, error: undefined as unknown })) + .catch((error: unknown) => ({ received: undefined, error })) + + expect(result.error).toBeDefined() + expect((result.error as Error).name).toBe('StreamReadError') + }) + + it('a joinRun-only socket drop rejects the iterator instead of hanging', async () => { + FakeWebSocket.instances = [] + const conn = webSocket('wss://x/api/chat', { + WebSocketImpl: FakeWebSocket as unknown as typeof WebSocket, + }) + const ac = new AbortController() + const iter = conn.joinRun('run-j', ac.signal)[Symbol.asyncIterator]() + await tick() + + const ws = FakeWebSocket.instances[0] + if (!ws) + throw new Error('expected a FakeWebSocket instance to have been created') + ws.close() + + const result = await withTimeout( + drainToEnd(iter), + 'joinRun() hung instead of rejecting after the socket closed', + ) + .then((received) => ({ received, error: undefined as unknown })) + .catch((error: unknown) => ({ received: undefined, error })) + + expect(result.error).toBeDefined() + expect((result.error as Error).name).toBe('StreamReadError') + expect(FakeWebSocket.instances.length).toBe(1) + }) }) diff --git a/packages/ai/src/stream-to-websocket.ts b/packages/ai/src/stream-to-websocket.ts index 009cd92bf1..845eba9639 100644 --- a/packages/ai/src/stream-to-websocket.ts +++ b/packages/ai/src/stream-to-websocket.ts @@ -179,6 +179,10 @@ export function toWebSocketStream( async function handleInbound(input: unknown): Promise { const params = await chatParamsFromRequestBody(input) const turnAbort = new AbortController() + // A second inbound frame with the same runId (client resubmit) must + // abort the earlier turn. Otherwise the old controller is overwritten + // and close/abort frames can no longer reach it. + activeTurns.get(params.runId)?.abort() activeTurns.set(params.runId, turnAbort) const ctx: WsRunContext = { messages: params.messages, diff --git a/packages/ai/tests/stream-to-websocket.test.ts b/packages/ai/tests/stream-to-websocket.test.ts index 57126aa4ef..b23d888ad5 100644 --- a/packages/ai/tests/stream-to-websocket.test.ts +++ b/packages/ai/tests/stream-to-websocket.test.ts @@ -233,6 +233,34 @@ describe('toWebSocketStream lifecycle', () => { expect(aborted).toBe(true) }) + it('aborts an in-flight turn when a second run frame reuses the same runId', async () => { + const socket = new FakeSocket() + const signals: Array = [] + toWebSocketStream(socket, new Request('https://x/api/chat'), { + onRun: ({ runId, threadId, signal }): AsyncIterable => + (async function* () { + signals.push(signal) + yield ev.runStarted(runId, threadId) + await new Promise((resolve) => { + if (signal.aborted) { + resolve() + return + } + signal.addEventListener('abort', () => resolve(), { once: true }) + }) + })(), + }) + socket.emitMessage(inputFrame('run-dup')) + await flush() + socket.emitMessage(inputFrame('run-dup')) + await flush() + + expect(signals).toHaveLength(2) + expect(signals[0]?.aborted).toBe(true) + expect(signals[1]?.aborted).toBe(false) + expect(socket.closed).toBe(false) + }) + it('drops a malformed inbound frame without crashing the socket, and still processes a subsequent valid frame', async () => { const socket = new FakeSocket() toWebSocketStream(socket, new Request('https://x/api/chat'), { diff --git a/testing/e2e/src/lib/durable-delivery-ws-plugin.ts b/testing/e2e/src/lib/durable-delivery-ws-plugin.ts index 43d50bb5f3..63a1eaa87d 100644 --- a/testing/e2e/src/lib/durable-delivery-ws-plugin.ts +++ b/testing/e2e/src/lib/durable-delivery-ws-plugin.ts @@ -33,6 +33,19 @@ import type { WebSocketLike } from '@tanstack/ai' */ const WS_PATH = '/api/durable-delivery-ws' +function isWebSocketLike(value: unknown): value is WebSocketLike { + if (typeof value !== 'object' || value === null) return false + if (!('send' in value) || typeof value.send !== 'function') return false + if (!('close' in value) || typeof value.close !== 'function') return false + if ( + !('addEventListener' in value) || + typeof value.addEventListener !== 'function' + ) { + return false + } + return true +} + export function durableDeliveryWebSocketPlugin(): Plugin { return { name: 'durable-delivery-websocket-plugin', @@ -50,16 +63,14 @@ export function durableDeliveryWebSocketPlugin(): Plugin { wss.handleUpgrade(req, socket, head, (ws) => { const request = new Request(url) - // `ws`'s WebSocket implements the WHATWG send/close/addEventListener/ - // bufferedAmount surface `WebSocketLike` needs. - const socketLike = ws as unknown as WebSocketLike + if (!isWebSocketLike(ws)) return if (url.searchParams.get('offset') !== null) { - resumeWebSocketStream(socketLike, { + resumeWebSocketStream(ws, { adapter: memoryStream(request), }) } else { - toWebSocketStream(socketLike, request, { + toWebSocketStream(ws, request, { durability: (ctx) => memoryStream(ctx.request), onRun: ({ runId, threadId }) => fixedRun(threadId, runId), }) From 1449ea378e38afe516ca7911e6faf39d93c1ed07 Mon Sep 17 00:00:00 2001 From: Tom Beckenham <34339192+tombeckenham@users.noreply.github.com> Date: Wed, 19 Aug 2026 21:34:27 +1000 Subject: [PATCH 26/26] fix(ai,ai-client): address WebSocket transport review findings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Server (toWebSocketStream): - Surface turn failures as a live RUN_ERROR frame instead of logging and dropping them — the conversation socket stays open, so a swallowed error left the client hanging forever (parity with the HTTP transports' synthesized terminal). - Register an error listener on run and resume sockets: on ws an unhandled error event is an uncaught exception, and an errored socket leaked its heartbeat/idle timers and never aborted its turns. - Close the register-after-close race: a turn validated while the socket closed no longer starts unabortable work, and an abort frame that races ahead of its run's registration is honored instead of dropped. - buildTurnRequest scrubs a handshake ?offset so a mis-routed resume can't silently take the durability replay branch; the always-null WsRunContext.resumeOffset field is removed. - WebSocketLike: overloaded addEventListener replaces the (ev: any) handler, and the never-read bufferedAmount field is dropped. ws sockets now satisfy the type directly, so the plugins' type-laundering guard is deleted. Client (webSocket()): - openOnce is keyed by run/resume mode: joinRun owns a dedicated replay socket (its ?offset handshake was silently discarded when a conversation socket existed, and a later send() wrote run frames onto a listener-less resume socket); a send() during reconnect backoff supersedes the pending resume instead of racing it. - A same-runId resubmit resets sawTerminal/progressed, so a drop during a client-tool continuation reconnects instead of hanging. - stop() now sends the { type: 'abort', runId } frame so the server cancels the turn instead of generating to completion. - The duplicated subscribe/joinRun iterator pump is extracted into a shared createChunkPipe; body option is Record. Tests: server multi-turn multiplexing with per-turn durability logs, idle timeout (incl. the in-flight-turn guard), error-frame surfacing, socket-error teardown; client new-runId session reset, same-runId resubmit reconnect, abort-frame emission, run/resume socket isolation, clean joinRun completion; new e2e driving the real webSocket() adapter through a forced mid-run drop and auto-resume. Docs updated to match (abort semantics, idle-timeout exemption, runtime support claims, adapter options), config.json dates fixed. Co-Authored-By: Claude Fable 5 --- .changeset/websocket-transport.md | 13 +- docs/config.json | 5 +- docs/resumable-streams/advanced.md | 3 +- docs/resumable-streams/websockets.md | 58 ++- .../src/lib/websocket-chat-plugin.ts | 20 +- packages/ai-client/src/connection-adapters.ts | 334 +++++++++++------- .../connection-adapters-websocket.test.ts | 261 +++++++++++++- packages/ai/src/stream-to-response.ts | 2 +- packages/ai/src/stream-to-websocket.ts | 140 +++++--- packages/ai/tests/stream-to-websocket.test.ts | 194 ++++++++-- .../e2e/src/lib/durable-delivery-ws-plugin.ts | 61 +++- testing/e2e/src/routeTree.gen.ts | 21 ++ testing/e2e/src/routes/websocket-adapter.tsx | 81 +++++ testing/e2e/tests/websocket.spec.ts | 24 ++ 14 files changed, 951 insertions(+), 266 deletions(-) create mode 100644 testing/e2e/src/routes/websocket-adapter.tsx diff --git a/.changeset/websocket-transport.md b/.changeset/websocket-transport.md index f4ca05f40e..ed3cc5cbb4 100644 --- a/.changeset/websocket-transport.md +++ b/.changeset/websocket-transport.md @@ -21,9 +21,11 @@ you at `toWebSocketStream`). Because one socket outlives many `chat()` turns (client-tool resubmits, follow-up user messages), you pass an `onRun(ctx) => AsyncIterable` factory instead of a prebuilt stream — the helper calls it per inbound `RunAgentInput` frame. The socket is conversation-scoped: -it stays open across turns and closes on client close, an `{ type: 'abort', -runId }` control frame (which aborts only that turn), or the idle timeout, with a -periodic `{ type: 'ping' }` heartbeat. Durability is keyed per turn and reuses +it stays open across turns and closes on client close or the idle timeout +(which never fires while a turn is still streaming), with a periodic +`{ type: 'ping' }` heartbeat. An `{ type: 'abort', runId }` control frame +aborts only that turn, leaving the socket open. A turn that fails is surfaced +to the client as a live `RUN_ERROR` frame, mirroring the HTTP transports. Durability is keyed per turn and reuses the existing `durableStreamSource`, so server→client frames carry the same `{ id, chunk }` envelope as NDJSON. `resumeWebSocketStream(socket, { adapter })` and `resumeWebSocketResponse({ adapter })` replay a run read-only from the @@ -39,4 +41,7 @@ handshake header, so the offset rides in the URL). The reconnect bookkeeping (offset de-dupe, no-progress ceiling → `StreamReconnectLimitError`) is shared with the HTTP adapters via the new `createReconnectTracker`, and a fatal drop surfaces to the consumer (`StreamReadError` / `StreamReconnectLimitError`) -instead of hanging. +instead of hanging. Aborting a run (`stop()` in `useChat`) sends the +`{ type: 'abort', runId }` frame so the server cancels the turn instead of +generating to completion, and `joinRun()` opens its own replay socket so a +rejoin never collides with the live conversation socket. diff --git a/docs/config.json b/docs/config.json index a106644d6f..3f50e9e7f6 100644 --- a/docs/config.json +++ b/docs/config.json @@ -227,13 +227,12 @@ "label": "Advanced", "to": "resumable-streams/advanced", "addedAt": "2026-08-04", - "updatedAt": "2026-08-18" + "updatedAt": "2026-08-19" }, { "label": "WebSockets", "to": "resumable-streams/websockets", - "addedAt": "2026-07-20", - "updatedAt": "2026-08-18" + "addedAt": "2026-08-19" }, { "label": "Custom Durability Adapter", diff --git a/docs/resumable-streams/advanced.md b/docs/resumable-streams/advanced.md index 1f56749f27..0475af75f2 100644 --- a/docs/resumable-streams/advanced.md +++ b/docs/resumable-streams/advanced.md @@ -213,7 +213,8 @@ process shutdown, not when one turn's `RUN_FINISHED` arrives. **Heartbeat and idle timeout replace the transport-level keepalive SSE gets for free.** `toWebSocketStream` pings every `heartbeatMs` (default 30s) and closes the socket after `idleTimeoutMs` (default 5 minutes) with no inbound -frame. +frame — but never while a turn is still streaming, so a long single +generation is safe. **An abort frame targets one turn, not the socket.** `{ type: 'abort', runId }` aborts only that turn's `onRun` iteration; the socket stays open for the next diff --git a/docs/resumable-streams/websockets.md b/docs/resumable-streams/websockets.md index e3b093d89c..7afa95aba1 100644 --- a/docs/resumable-streams/websockets.md +++ b/docs/resumable-streams/websockets.md @@ -90,7 +90,6 @@ export function handleChatSocket(socket: WebSocketLike, request: Request) { | `messages` | The turn's `UIMessage[]` / `ModelMessage[]`, decoded from the frame. | | `threadId` / `runId` | The AG-UI identifiers for this turn. | | `forwardedProps` | Any extra data the client sent with the frame. | -| `resumeOffset` | Always `null` for a fresh turn on this socket (a resume goes through `resumeWebSocketStream` instead, see below). | | `request` | A synthetic per-turn `Request` carrying `?runId=` in its URL, for keying a durability adapter. | | `signal` | Aborts when the socket closes, or when this turn receives an `abort` frame (see below). It does not abort other turns on the same socket. | @@ -118,9 +117,17 @@ export function Chat() { ``` Nothing else changes: `messages`, `sendMessage`, `stop()`, tool calls, and -persistence all work the same as with any other connection adapter. See -[Connection Adapters](../chat/connection-adapters) for the full option set -(`protocols`, `body`, `reconnect`, `WebSocketImpl`). +persistence all work the same as with any other connection adapter. Options: + +| Option | What it does | +| --- | --- | +| `protocols` | WebSocket subprotocol(s), passed straight to the `WebSocket` constructor. | +| `body` | Static extra fields merged into every outgoing `RunAgentInput` frame (same as the HTTP adapters' `body`). | +| `reconnect` | Reconnect bounds (`maxAttempts`, `delayMs`), shared semantics with `fetchServerSentEvents` — see [Advanced](./advanced). | +| `WebSocketImpl` | Override the `WebSocket` implementation (tests, non-browser runtimes). | + +See [Connection Adapters](../chat/connection-adapters) for where `webSocket()` +fits among the other adapters. ## Wire protocol @@ -178,13 +185,11 @@ is a factory, not a value: it receives the per-turn `ctx` and reads ## Abort a turn An `{ type: 'abort', runId }` frame aborts only that turn's `onRun` iteration -(`ctx.signal` fires); the socket itself stays open for the next turn. This is -a protocol-level primitive: the built-in `webSocket()` client adapter does not -send this frame on `stop()` today, so `stop()` only stops the client from -processing further chunks for that run locally, without telling the server to -cancel the model call. If you need the server to actually stop generating, -send the frame yourself over a reference to the socket, or drive the protocol -directly the way the e2e suite's raw-`WebSocket` tests do. +(`ctx.signal` fires); the socket itself stays open for the next turn. The +built-in `webSocket()` client adapter sends this frame when the run's abort +signal fires (`stop()` in `useChat`), so the server actually stops generating +instead of running the model call to completion against a client that stopped +listening. A hand-rolled client should send the same frame to cancel a turn. Closing the whole socket aborts every turn still in flight on it. @@ -194,7 +199,9 @@ Closing the whole socket aborts every turn still in flight on it. (default 30 seconds) to keep the connection alive through proxies that drop idle sockets. It closes the socket if no inbound frame arrives for `idleTimeoutMs` (default 5 minutes); a heartbeat itself doesn't count as -activity, only a client-sent frame does: +activity, only a client-sent frame does. The idle timeout never fires while a +turn is still streaming, so a long single generation (an agentic loop, a +turn longer than the timeout) is never cut off: ```ts import { chat, toWebSocketStream } from '@tanstack/ai' @@ -219,27 +226,16 @@ result to `toWebSocketStream`. The pattern: hook the HTTP server's `upgrade` event, accept the socket with [`ws`](https://github.com/websockets/ws)'s `WebSocketServer({ noServer: true })`, and pass the resulting socket straight through. `ws`'s socket already implements the `send`/`close`/ -`addEventListener`/`bufferedAmount` surface `WebSocketLike` needs: +`addEventListener` surface `WebSocketLike` needs: ```ts ignore import { WebSocketServer } from 'ws' -import { memoryStream, resumeWebSocketStream, toWebSocketStream } from '@tanstack/ai' +import { memoryStream, resumeWebSocketStream } from '@tanstack/ai' import type { Plugin } from 'vite' -import type { WebSocketLike } from '@tanstack/ai' import { handleChatSocket } from './handle-chat-socket' const WS_PATH = '/api/chat-ws' -function isWebSocketLike(value: unknown): value is WebSocketLike { - if (typeof value !== 'object' || value === null) return false - if (!('send' in value) || typeof value.send !== 'function') return false - if (!('close' in value) || typeof value.close !== 'function') return false - if (!('addEventListener' in value) || typeof value.addEventListener !== 'function') { - return false - } - return true -} - export function webSocketChatPlugin(): Plugin { return { name: 'websocket-chat-plugin', @@ -252,9 +248,9 @@ export function webSocketChatPlugin(): Plugin { if (url.pathname !== WS_PATH) return wss.handleUpgrade(req, socket, head, (ws) => { + // `ws`'s socket satisfies WebSocketLike structurally — pass it + // straight through, no adapter needed. const request = new Request(url) - if (!isWebSocketLike(ws)) return - if (url.searchParams.get('offset') !== null) { resumeWebSocketStream(ws, { adapter: memoryStream(request) }) } else { @@ -271,9 +267,11 @@ This is the same pattern used by the working [`examples/ts-react-chat`](https://github.com/TanStack/ai/blob/main/examples/ts-react-chat/src/lib/websocket-chat-plugin.ts) WebSocket example and by the e2e suite's [`durable-delivery-ws-plugin.ts`](https://github.com/TanStack/ai/blob/main/testing/e2e/src/lib/durable-delivery-ws-plugin.ts). -Bun's `ServerWebSocket` and Deno's `Deno.upgradeWebSocket` need the same shape -of adapter: accept the platform's socket, then call `toWebSocketStream` / -`resumeWebSocketStream` with it. +The same wiring pattern applies elsewhere: accept the platform's socket, then +call `toWebSocketStream` / `resumeWebSocketStream` with it. A socket from +Deno's `Deno.upgradeWebSocket` is WHATWG-shaped and satisfies `WebSocketLike` +directly; Bun's `ServerWebSocket` (handler-object API) needs a small adapter +first. ## Hosting on Cloudflare diff --git a/examples/ts-react-chat/src/lib/websocket-chat-plugin.ts b/examples/ts-react-chat/src/lib/websocket-chat-plugin.ts index 45d89129f3..ca36a27f95 100644 --- a/examples/ts-react-chat/src/lib/websocket-chat-plugin.ts +++ b/examples/ts-react-chat/src/lib/websocket-chat-plugin.ts @@ -7,13 +7,12 @@ import { } from '@tanstack/ai' import { openaiText } from '@tanstack/ai-openai' import type { Plugin } from 'vite' -import type { WebSocketLike } from '@tanstack/ai' /** * A Vite dev-server plugin exposing a full-duplex WebSocket chat endpoint. * This app is served by Vite/Nitro on plain Node, which has no * `WebSocketPair`, so `toWebSocketResponse`/`resumeWebSocketResponse` (the - * CF/Deno wrapper) can't be used here. Instead this hooks the underlying + * Cloudflare wrapper) can't be used here. Instead this hooks the underlying * Node http server's `upgrade` event directly — obtains a server socket via * `ws`'s `WebSocketServer({ noServer: true })`, and hands it to * `toWebSocketStream`/`resumeWebSocketStream`. Mirrors @@ -30,19 +29,6 @@ import type { WebSocketLike } from '@tanstack/ai' */ const WS_PATH = '/api/chat-ws' -function isWebSocketLike(value: unknown): value is WebSocketLike { - if (typeof value !== 'object' || value === null) return false - if (!('send' in value) || typeof value.send !== 'function') return false - if (!('close' in value) || typeof value.close !== 'function') return false - if ( - !('addEventListener' in value) || - typeof value.addEventListener !== 'function' - ) { - return false - } - return true -} - /** * `chat()` cancels via an `AbortController` it can read `.signal` off of, but * `WsRunContext` (the per-turn context `toWebSocketStream` hands `onRun`) @@ -76,9 +62,9 @@ export function webSocketChatPlugin(): Plugin { if (url.pathname !== WS_PATH) return wss.handleUpgrade(req, socket, head, (ws) => { + // `ws`'s socket satisfies WebSocketLike structurally — no adapter + // (and no runtime guard) needed. const request = new Request(url) - if (!isWebSocketLike(ws)) return - if (url.searchParams.get('offset') !== null) { resumeWebSocketStream(ws, { adapter: memoryStream(request), diff --git a/packages/ai-client/src/connection-adapters.ts b/packages/ai-client/src/connection-adapters.ts index 776e6e6a6b..6fb3592f4e 100644 --- a/packages/ai-client/src/connection-adapters.ts +++ b/packages/ai-client/src/connection-adapters.ts @@ -1942,7 +1942,7 @@ export function xhrHttpStream( export interface WebSocketConnectionOptions { protocols?: string | Array - body?: Record + body?: Record reconnect?: ReconnectOptions /** Override the WebSocket implementation (tests / non-browser runtimes). */ WebSocketImpl?: typeof WebSocket @@ -1952,17 +1952,97 @@ function runIdQuery(url: string, runId: string | undefined): string { return runId ? withSearchParams(url, { runId }) : url } -/** A subscribe()/joinRun() consumer's registration: receives chunks or a fatal error. */ +function isPingFrame(parsed: unknown): boolean { + return ( + typeof parsed === 'object' && + parsed !== null && + (parsed as { type?: unknown }).type === 'ping' + ) +} + +/** A subscribe() consumer's registration: receives chunks or a fatal error. */ interface WebSocketChunkSink { push: (chunk: StreamChunk) => void fail: (error: unknown) => void } +/** + * A push→pull bridge from socket callbacks to an async iterable: chunks queue + * until the consumer pulls, a recorded failure rejects the iterator, and + * `end()` (or the abort signal) finishes it cleanly. Shared by `webSocket()`'s + * `subscribe()` and `joinRun()`. + */ +function createChunkPipe( + abortSignal: AbortSignal | undefined, + onFinally: () => void, +): { + push: (chunk: StreamChunk) => void + fail: (error: unknown) => void + end: () => void + iterable: AsyncIterable +} { + const queue: Array = [] + const waiters: Array<(c: StreamChunk | null) => void> = [] + let failure: unknown + let ended = false + const wake = () => waiters.shift()?.(null) + const push = (chunk: StreamChunk) => { + const w = waiters.shift() + if (w) w(chunk) + else queue.push(chunk) + } + const fail = (error: unknown) => { + failure = error + wake() + } + const end = () => { + ended = true + wake() + } + const onAbort = () => wake() + abortSignal?.addEventListener('abort', onAbort) + const iterable = (async function* () { + try { + while (!abortSignal?.aborted) { + // Drain buffered chunks before ever awaiting a new promise — a + // fatal drop that lands while chunks are still queued (fail() + // finds no pending waiter, since the consumer hasn't caught up + // to its buffer yet) must not be lost. + const buffered = queue.shift() + if (buffered !== undefined) { + yield buffered + continue + } + // Buffer exhausted: surface a failure recorded while we were + // draining, rather than awaiting a promise that will never + // resolve (the connection is dead — no future push/fail). + if (failure !== undefined) throw failure + if (ended) return + const chunk = await new Promise((r) => + waiters.push(r), + ) + // The wait resolved because fail() woke us — surface the error + // instead of treating the null sentinel as a clean end. TS narrows + // `failure` to `undefined` from the check above and doesn't know + // the `fail()` closure can reassign it while we were awaiting — + // this check is very much still reachable. + // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition + if (failure !== undefined) throw failure + if (chunk === null) return + yield chunk + } + } finally { + abortSignal?.removeEventListener('abort', onAbort) + onFinally() + } + })() + return { push, fail, end, iterable } +} + /** * The send()-driven run currently owning auto-reconnect for a `webSocket()` - * connection. `undefined` when no `send()` has established a run yet (e.g. a - * `joinRun()`-only connection), so a drop there is never auto-resumed — Task - * 11 scopes reconnect to the run `subscribe()`/`send()` are driving. + * connection: reconnect is scoped to the run `send()` is driving, so a drop + * with no established run is surfaced to subscribers rather than auto-resumed. */ interface WebSocketRunSession { runId: string | undefined @@ -1997,6 +2077,10 @@ export function webSocket( } { const Impl = options.WebSocketImpl ?? WebSocket let socket: WebSocket | undefined + // Whether the current socket is the conversation socket ('run') or a + // read-only replay connection opened by a reconnect ('resume'). Only the + // conversation socket accepts run frames server-side. + let socketMode: 'run' | 'resume' | undefined // Memoized per-socket open promise. `openOnce` sets `onopen`/`onerror` // exactly ONCE, at socket-creation time, and stores the resulting promise // here. Without this, `waitOpen` assigning `onopen`/`onerror` on every call @@ -2012,20 +2096,38 @@ export function webSocket( for (const l of listeners) l.fail(error) } - function openOnce(target: string): WebSocket { - if (socket && socket.readyState <= 1) return socket + function openOnce(target: string, mode: 'run' | 'resume'): WebSocket { + // Only the conversation socket is reused — it multiplexes many turns. A + // 'resume' handshake carries ?offset and must reach the server as its own + // connection (reusing any open socket would discard that query, so no + // replay would ever be requested), and a run frame must never be written + // to a read-only resume socket (the server registers no message listener + // there, so the frame would be silently ignored). + if ( + socket && + socket.readyState <= 1 && + mode === 'run' && + socketMode === 'run' + ) { + return socket + } + const prior = socket const ws = options.protocols ? new Impl(target, options.protocols) : new Impl(target) socket = ws + socketMode = mode openPromise = new Promise((resolve, reject) => { ws.onopen = () => resolve() ws.onerror = (e) => reject(new StreamReadError(e)) }) - // Attach a no-op handler so a socket nobody awaits (e.g. joinRun) can't raise - // an unhandled rejection if it errors. Awaiters of openPromise still see the rejection. + // Attach a no-op handler so a socket nobody awaits can't raise an + // unhandled rejection if it errors. Awaiters of openPromise still see the rejection. openPromise.catch(() => {}) ws.onmessage = (event: MessageEvent) => { + // A retired socket (a newer connection took over below) must not keep + // feeding the shared listeners. + if (ws !== socket) return let parsed: unknown try { parsed = JSON.parse(String(event.data)) @@ -2033,13 +2135,7 @@ export function webSocket( failAll(new StreamReadError(error)) return } - if ( - typeof parsed === 'object' && - parsed !== null && - (parsed as { type?: unknown }).type === 'ping' - ) { - return - } + if (isPingFrame(parsed)) return const envelopeId = isNdjsonEnvelope(parsed) ? parsed.id : undefined const chunk = isNdjsonEnvelope(parsed) ? parsed.chunk @@ -2047,8 +2143,8 @@ export function webSocket( // Thread durable chunks through the active run session's tracker (if // any) so a later reconnect knows the last offset and can skip a - // replayed boundary. A socket with no active session (e.g. joinRun() - // only) dispatches chunks as-is, exactly as before Task 11. + // replayed boundary. A socket with no active session dispatches chunks + // as-is. const session = currentSession if (session) { if (session.tracker.note(envelopeId) === 'duplicate') return @@ -2063,10 +2159,12 @@ export function webSocket( for (const l of listeners) l.push(chunk) } ws.onclose = () => { + // Retired deliberately in favor of a newer connection — not a drop. + if (ws !== socket) return const session = currentSession if (!session) { - // joinRun() never creates a session. Surface the drop so those - // iterators do not stay parked on a dead socket. + // No run session (never established, or cleared by a prior failure). + // Surface the drop so subscribers do not stay parked on a dead socket. failAll(new StreamReadError(new Error('WebSocket connection closed'))) return } @@ -2082,6 +2180,11 @@ export function webSocket( } void reconnect(session, lastEventId) } + // Retire a superseded socket (e.g. a lingering resume socket when send() + // opens the next conversation socket) so two sockets never feed the + // shared listeners at once. Its handlers see it is no longer current and + // ignore the close. + if (prior && prior.readyState <= 1) prior.close() return ws } @@ -2097,18 +2200,25 @@ export function webSocket( session.signal, ) } catch (error) { - currentSession = undefined + if (currentSession === session) currentSession = undefined failAll(error) return } if (session.signal?.aborted) return + // A send() issued during the backoff supersedes this resume: a newer run + // (or a resubmit of this one) already owns a fresh conversation socket, + // and its turn re-delivers from the durability log — the tracker de-dupes + // any overlap. Opening the resume socket anyway would retire that live + // conversation socket. + if (currentSession !== session) return + if (socket && socket.readyState <= 1) return session.progressed = false const base = typeof url === 'function' ? url() : url const target = withSearchParams(base, { ...(session.runId !== undefined ? { runId: session.runId } : {}), offset, }) - openOnce(target) + openOnce(target, 'resume') } function waitOpen(ws: WebSocket): Promise { @@ -2121,64 +2231,14 @@ export function webSocket( return { subscribe(abortSignal?: AbortSignal): AsyncIterable { - const queue: Array = [] - const waiters: Array<(c: StreamChunk | null) => void> = [] - let failure: unknown - const push = (c: StreamChunk) => { - const w = waiters.shift() - if (w) w(c) - else queue.push(c) - } - const fail = (e: unknown) => { - failure = e - const w = waiters.shift() - if (w) w(null) - } - const sink: WebSocketChunkSink = { push, fail } + const pipe = createChunkPipe(abortSignal, () => listeners.delete(sink)) + const sink: WebSocketChunkSink = { push: pipe.push, fail: pipe.fail } listeners.add(sink) - const onAbort = () => { - const w = waiters.shift() - if (w) w(null) - } - abortSignal?.addEventListener('abort', onAbort) - return (async function* () { - try { - while (!abortSignal?.aborted) { - // Drain buffered chunks before ever awaiting a new promise — a - // fatal drop that lands while chunks are still queued (fail() - // finds no pending waiter, since the consumer hasn't caught up - // to its buffer yet) must not be lost. - const buffered = queue.shift() - if (buffered !== undefined) { - yield buffered - continue - } - // Buffer exhausted: surface a failure recorded while we were - // draining, rather than awaiting a promise that will never - // resolve (the connection is dead — no future push/fail). - if (failure !== undefined) throw failure - const chunk = await new Promise((r) => - waiters.push(r), - ) - // The wait resolved because fail() woke us — surface the error - // instead of treating the null sentinel as a clean end. TS narrows - // `failure` to `undefined` from the check above and doesn't know - // the `fail()` closure can reassign it while we were awaiting — - // this check is very much still reachable. - // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition - if (failure !== undefined) throw failure - if (chunk === null) return - yield chunk - } - } finally { - listeners.delete(sink) - abortSignal?.removeEventListener('abort', onAbort) - } - })() + return pipe.iterable }, async send(messages, data, abortSignal, runContext) { const target = typeof url === 'function' ? url() : url - const ws = openOnce(runIdQuery(target, runContext?.runId)) + const ws = openOnce(runIdQuery(target, runContext?.runId), 'run') await waitOpen(ws) // Establish (or continue) the run session this socket is driving, so // an unterminated drop can auto-resume it. A distinct runId starts a @@ -2194,8 +2254,40 @@ export function webSocket( signal: abortSignal, } } else { + // Same-runId resubmit (e.g. a client-tool continuation): keep the + // tracker, but this is a NEW turn — with the previous turn's + // `sawTerminal` left set, a drop during the resubmitted turn would + // neither reconnect nor surface an error. currentSession.signal = abortSignal + currentSession.sawTerminal = false + currentSession.progressed = false } + const session = currentSession + // stop() must reach the server: the conversation socket outlives the + // turn, so without an abort frame the model keeps generating (and + // billing) server-side. The frame aborts only this run's turn. + abortSignal?.addEventListener( + 'abort', + () => { + const abortRunId = session.runId + const live = socket + if ( + abortRunId === undefined || + session.sawTerminal || + socketMode !== 'run' || + live === undefined || + live.readyState !== 1 + ) { + return + } + try { + live.send(JSON.stringify({ type: 'abort', runId: abortRunId })) + } catch { + // Socket is CLOSING/CLOSED — the server aborts the turn on close. + } + }, + { once: true }, + ) const body = buildRunAgentInputBody(messages, data, runContext, { body: options.body, }) @@ -2206,58 +2298,48 @@ export function webSocket( offset: '-1', runId, }) - openOnce(target) - const chunks: Array = [] - const waiters: Array<(c: StreamChunk | null) => void> = [] - let failure: unknown - const push = (c: StreamChunk) => { - const w = waiters.shift() - if (w) w(c) - else chunks.push(c) - } - const fail = (e: unknown) => { - failure = e - const w = waiters.shift() - if (w) w(null) - } - const sink: WebSocketChunkSink = { push, fail } - listeners.add(sink) - const onAbort = () => waiters.shift()?.(null) - abortSignal?.addEventListener('abort', onAbort) - return (async function* () { + // A replay handshake must reach the server as its own connection: + // reusing the conversation socket would discard the ?offset query (no + // replay ever requested), and the conversation socket must not be + // replaced by a read-only replay socket. So joinRun owns a dedicated + // socket and never touches the shared socket or run session. + const ws = options.protocols + ? new Impl(target, options.protocols) + : new Impl(target) + const pipe = createChunkPipe(abortSignal, () => { + if (ws.readyState <= 1) ws.close() + }) + ws.onmessage = (event: MessageEvent) => { + let parsed: unknown try { - while (!abortSignal?.aborted) { - // Drain buffered chunks before ever awaiting a new promise — a - // fatal drop that lands while chunks are still queued (fail() - // finds no pending waiter, since the consumer hasn't caught up - // to its buffer yet) must not be lost. - const buffered = chunks.shift() - if (buffered !== undefined) { - yield buffered - continue - } - // Buffer exhausted: surface a failure recorded while we were - // draining, rather than awaiting a promise that will never - // resolve (the connection is dead — no future push/fail). - if (failure !== undefined) throw failure - const c = await new Promise((r) => - waiters.push(r), - ) - // The wait resolved because fail() woke us — surface the error - // instead of treating the null sentinel as a clean end. TS narrows - // `failure` to `undefined` from the check above and doesn't know - // the `fail()` closure can reassign it while we were awaiting — - // this check is very much still reachable. - // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition - if (failure !== undefined) throw failure - if (c === null) return - yield c - } - } finally { - listeners.delete(sink) - abortSignal?.removeEventListener('abort', onAbort) + parsed = JSON.parse(String(event.data)) + } catch (error) { + pipe.fail(new StreamReadError(error)) + return } - })() + if (isPingFrame(parsed)) return + pipe.push( + isNdjsonEnvelope(parsed) ? parsed.chunk : (parsed as StreamChunk), + ) + } + ws.onclose = (event?: CloseEvent) => { + // 1000 = the server finished replaying the log and closed cleanly. + // Anything else is a drop or a policy refusal (e.g. 1008 "no resume + // offset") and must surface — a joinRun socket never auto-reconnects. + if (event?.code === 1000) { + pipe.end() + return + } + const detail = event + ? `${event.code}${event.reason ? `: ${event.reason}` : ''}` + : 'unknown' + pipe.fail( + new StreamReadError( + new Error(`WebSocket connection closed (${detail})`), + ), + ) + } + return pipe.iterable }, } } diff --git a/packages/ai-client/tests/connection-adapters-websocket.test.ts b/packages/ai-client/tests/connection-adapters-websocket.test.ts index 38fdc55524..9b206e0ddb 100644 --- a/packages/ai-client/tests/connection-adapters-websocket.test.ts +++ b/packages/ai-client/tests/connection-adapters-websocket.test.ts @@ -9,7 +9,7 @@ class FakeWebSocket { static instances: Array = [] onopen: (() => void) | null = null onmessage: ((ev: { data: string }) => void) | null = null - onclose: (() => void) | null = null + onclose: ((ev?: { code?: number; reason?: string }) => void) | null = null onerror: ((ev: unknown) => void) | null = null sent: Array = [] readyState = 0 @@ -29,6 +29,11 @@ class FakeWebSocket { this.readyState = 3 this.onclose?.() } + /** Simulate the server closing the socket with an explicit close code. */ + serverClose(code: number, reason = ''): void { + this.readyState = 3 + this.onclose?.({ code, reason }) + } emit(chunk: unknown, id?: string): void { this.onmessage?.({ data: JSON.stringify(id === undefined ? chunk : { id, chunk }), @@ -303,6 +308,171 @@ describe('webSocket() reconnect', () => { 'RUN_FINISHED', ]) }) + + it('a second send() with a new runId resets the run session (fresh offsets, reconnect re-armed)', async () => { + FakeWebSocket.instances = [] + const conn = webSocket('wss://x/api/chat', { + WebSocketImpl: FakeWebSocket as unknown as typeof WebSocket, + reconnect: { delayMs: 0, maxAttempts: 5 }, + }) + const ac = new AbortController() + const received: Array = [] + const sub = drain(conn.subscribe(ac.signal), received, ac.signal) + + // Run 1: streams to terminal — sawTerminal is now set for the session. + await conn.send( + [{ role: 'user', content: 'hi' } as any], + undefined, + undefined, + { threadId: 't', runId: 'run-1' }, + ) + const ws1 = FakeWebSocket.instances[0] + if (!ws1) throw new Error('expected a first FakeWebSocket instance') + await tick() + ws1.emit( + { type: 'TEXT_MESSAGE_CONTENT', delta: 'a', timestamp: 0 }, + 'r1-off-1', + ) + ws1.emit( + { + type: 'RUN_FINISHED', + runId: 'run-1', + threadId: 't', + model: 'm', + finishReason: 'stop', + timestamp: 0, + }, + 'r1-off-2', + ) + await tick() + + // Run 2 on the same conversation socket. If the session were reused, + // run-1's sawTerminal would suppress this reconnect entirely, and run-1's + // lastEventId would be sent as run-2's resume offset. + await conn.send( + [{ role: 'user', content: 'again' } as any], + undefined, + undefined, + { threadId: 't', runId: 'run-2' }, + ) + await tick() + expect(FakeWebSocket.instances.length).toBe(1) // conversation socket reused + ws1.emit( + { type: 'TEXT_MESSAGE_CONTENT', delta: 'b', timestamp: 0 }, + 'r2-off-1', + ) + await tick() + ws1.close() // drop mid-run-2 + + await new Promise((r) => setTimeout(r, 5)) + const ws2 = FakeWebSocket.instances[1] + if (!ws2) throw new Error('expected a reconnect FakeWebSocket instance') + const url2 = new URL(ws2.url) + expect(url2.searchParams.get('runId')).toBe('run-2') + expect(url2.searchParams.get('offset')).toBe('r2-off-1') + + ws2.emit( + { + type: 'RUN_FINISHED', + runId: 'run-2', + threadId: 't', + model: 'm', + finishReason: 'stop', + timestamp: 0, + }, + 'r2-off-2', + ) + await tick() + ac.abort() + await sub + expect(received.map((c) => c.delta ?? c.type)).toEqual([ + 'a', + 'RUN_FINISHED', + 'b', + 'RUN_FINISHED', + ]) + }) + + it('a same-runId resubmit clears sawTerminal so a drop during the new turn still reconnects', async () => { + FakeWebSocket.instances = [] + const conn = webSocket('wss://x/api/chat', { + WebSocketImpl: FakeWebSocket as unknown as typeof WebSocket, + reconnect: { delayMs: 0, maxAttempts: 5 }, + }) + const ac = new AbortController() + const received: Array = [] + const sub = drain(conn.subscribe(ac.signal), received, ac.signal) + + await conn.send( + [{ role: 'user', content: 'hi' } as any], + undefined, + undefined, + { threadId: 't', runId: 'run-x' }, + ) + const ws1 = FakeWebSocket.instances[0] + if (!ws1) throw new Error('expected a first FakeWebSocket instance') + await tick() + ws1.emit( + { + type: 'RUN_FINISHED', + runId: 'run-x', + threadId: 't', + model: 'm', + finishReason: 'tool_calls', + timestamp: 0, + }, + 'off-1', + ) + await tick() + + // Client-tool continuation: resubmit the SAME runId, then drop mid-turn. + await conn.send( + [{ role: 'user', content: 'tool result' } as any], + undefined, + undefined, + { threadId: 't', runId: 'run-x' }, + ) + await tick() + ws1.emit( + { type: 'TEXT_MESSAGE_CONTENT', delta: 'more', timestamp: 0 }, + 'off-2', + ) + await tick() + ws1.close() + + await new Promise((r) => setTimeout(r, 5)) + // The stale terminal from the first turn must not suppress this reconnect. + const ws2 = FakeWebSocket.instances[1] + if (!ws2) throw new Error('expected a reconnect FakeWebSocket instance') + expect(new URL(ws2.url).searchParams.get('offset')).toBe('off-2') + ac.abort() + await sub + }) + + it('aborting a send() emits an { type: "abort", runId } frame on the conversation socket', async () => { + FakeWebSocket.instances = [] + const conn = webSocket('wss://x/api/chat', { + WebSocketImpl: FakeWebSocket as unknown as typeof WebSocket, + }) + const ac = new AbortController() + await conn.send( + [{ role: 'user', content: 'hi' } as any], + undefined, + ac.signal, + { threadId: 't', runId: 'run-stop' }, + ) + const ws = FakeWebSocket.instances[0] + if (!ws) throw new Error('expected a FakeWebSocket instance') + await tick() + expect(ws.sent.length).toBe(1) + + ac.abort() + expect(ws.sent.length).toBe(2) + expect(JSON.parse(ws.sent[1]!)).toEqual({ + type: 'abort', + runId: 'run-stop', + }) + }) }) describe('webSocket() fatal drop surfacing', () => { @@ -483,4 +653,93 @@ describe('webSocket() fatal drop surfacing', () => { expect((result.error as Error).name).toBe('StreamReadError') expect(FakeWebSocket.instances.length).toBe(1) }) + + it('a joinRun socket the server closes cleanly (1000, replay complete) ends the iterator without error', async () => { + FakeWebSocket.instances = [] + const conn = webSocket('wss://x/api/chat', { + WebSocketImpl: FakeWebSocket as unknown as typeof WebSocket, + }) + const iter = conn.joinRun('run-j')[Symbol.asyncIterator]() + await tick() + + const ws = FakeWebSocket.instances[0] + if (!ws) throw new Error('expected a FakeWebSocket instance') + ws.emit( + { type: 'TEXT_MESSAGE_CONTENT', delta: 'replayed', timestamp: 0 }, + 'off-1', + ) + ws.serverClose(1000) + + const received = await withTimeout( + drainToEnd(iter), + 'joinRun() hung instead of ending after the server closed the replay', + ) + expect(received.map((c) => c.delta)).toEqual(['replayed']) + }) +}) + +describe('webSocket() run/resume socket isolation', () => { + it('joinRun during an open conversation socket opens its own replay socket, leaving the conversation socket intact', async () => { + FakeWebSocket.instances = [] + const conn = webSocket('wss://x/api/chat', { + WebSocketImpl: FakeWebSocket as unknown as typeof WebSocket, + }) + const ac = new AbortController() + await conn.send( + [{ role: 'user', content: 'hi' } as any], + undefined, + ac.signal, + { threadId: 't', runId: 'run-live' }, + ) + const conversation = FakeWebSocket.instances[0] + if (!conversation) throw new Error('expected a conversation socket') + await tick() + + const joinAc = new AbortController() + void conn.joinRun('run-old', joinAc.signal) + await tick() + + // The replay handshake reached the server as its OWN connection carrying + // the ?offset query — reusing the conversation socket would have silently + // discarded it and no replay would ever be requested. + expect(FakeWebSocket.instances.length).toBe(2) + const join = FakeWebSocket.instances[1] + if (!join) throw new Error('expected a joinRun socket') + expect(new URL(join.url).searchParams.get('offset')).toBe('-1') + expect(new URL(join.url).searchParams.get('runId')).toBe('run-old') + expect(conversation.readyState).toBe(1) + joinAc.abort() + ac.abort() + }) + + it('send() after joinRun opens a conversation socket instead of writing the run frame onto the replay socket', async () => { + FakeWebSocket.instances = [] + const conn = webSocket('wss://x/api/chat', { + WebSocketImpl: FakeWebSocket as unknown as typeof WebSocket, + }) + const joinAc = new AbortController() + void conn.joinRun('run-old', joinAc.signal) + await tick() + const join = FakeWebSocket.instances[0] + if (!join) throw new Error('expected a joinRun socket') + + // Page-reload rejoin flow: catch up on the old run, then send the next + // message. The server registers no message listener on a replay socket, + // so a run frame written there would be silently dropped. + await conn.send( + [{ role: 'user', content: 'next' } as any], + undefined, + undefined, + { threadId: 't', runId: 'run-new' }, + ) + await tick() + expect(FakeWebSocket.instances.length).toBe(2) + const conversation = FakeWebSocket.instances[1] + if (!conversation) throw new Error('expected a conversation socket') + expect(join.sent.length).toBe(0) + expect(conversation.sent.length).toBe(1) + expect(JSON.parse(conversation.sent[0]!).runId).toBe('run-new') + expect(new URL(conversation.url).searchParams.get('offset')).toBeNull() + joinAc.abort() + }) }) diff --git a/packages/ai/src/stream-to-response.ts b/packages/ai/src/stream-to-response.ts index 1e5e3bd849..6518c4fab8 100644 --- a/packages/ai/src/stream-to-response.ts +++ b/packages/ai/src/stream-to-response.ts @@ -78,7 +78,7 @@ function combineFailures( ) } -function runErrorChunk( +export function runErrorChunk( error: unknown, ): Extract { const payload = toRunErrorPayload(error) diff --git a/packages/ai/src/stream-to-websocket.ts b/packages/ai/src/stream-to-websocket.ts index 845eba9639..30b4ecabff 100644 --- a/packages/ai/src/stream-to-websocket.ts +++ b/packages/ai/src/stream-to-websocket.ts @@ -1,5 +1,5 @@ import { chatParamsFromRequestBody } from './utilities/chat-params' -import { durableStreamSource } from './stream-to-response' +import { durableStreamSource, runErrorChunk } from './stream-to-response' import { resolveDebugOption } from './logger/resolve' import type { StreamDurability } from './stream-durability' import type { DebugOption } from './logger/types' @@ -7,17 +7,17 @@ import type { ModelMessage, StreamChunk, UIMessage } from './types' /** * The minimal WHATWG WebSocket surface the core needs. Cloudflare - * `WebSocketPair` server sockets and Deno sockets already satisfy it; `ws` - * (Node) and Bun `ServerWebSocket` get a ~10-line adapter at the call site. + * `WebSocketPair` server sockets, Deno's upgraded sockets, and `ws` (Node) + * sockets already satisfy it; Bun's `ServerWebSocket` (handler-object API) + * gets a ~10-line adapter at the call site. */ export interface WebSocketLike { send: (data: string) => void close: (code?: number, reason?: string) => void - addEventListener: ( - type: 'message' | 'close' | 'error', - handler: (ev: any) => void, - ) => void - readonly bufferedAmount?: number + addEventListener: { + (type: 'message', handler: (ev: { data: unknown }) => void): void + (type: 'close' | 'error', handler: () => void): void + } } /** One inbound WS text frame, after JSON parse + shape discrimination. */ @@ -62,8 +62,6 @@ export interface WsRunContext { threadId: string runId: string forwardedProps?: Record - /** Non-null when this socket opened as a reconnect (carried `?offset`). */ - resumeOffset: string | null /** Synthetic per-turn request carrying `?runId=` so durability keys correctly. */ request: Request /** Aborts on socket close or an `abort` control frame for this run. */ @@ -75,16 +73,15 @@ export interface WsRunContext { * many runs; each turn's durability adapter must key on the frame's `runId`, * which we carry in the URL query (`memoryStream`/`durableStream` already read * `?runId` / `?offset` there). Headers are copied from the handshake so - * auth/cookies survive. + * auth/cookies survive. A handshake carrying `?offset` is a resume and never + * reaches a fresh turn (`resumeWebSocketStream` serves it), so the offset is + * scrubbed here — otherwise a mis-routed resume handshake would make the turn's + * durability adapter silently take the replay branch instead of running onRun. */ -export function buildTurnRequest( - handshake: Request, - runId: string, - offset: string | null, -): Request { +export function buildTurnRequest(handshake: Request, runId: string): Request { const url = new URL(handshake.url) url.searchParams.set('runId', runId) - if (offset !== null) url.searchParams.set('offset', offset) + url.searchParams.delete('offset') return new Request(url, { headers: handshake.headers }) } @@ -97,7 +94,11 @@ export interface WebSocketStreamInit { batch?: number /** Heartbeat ping interval in ms (default 30_000). */ heartbeatMs?: number - /** Close after this many ms without any inbound frame (default 300_000). */ + /** + * Close after this many ms without any inbound frame (default 300_000). + * Never fires while a turn is still streaming, so a long single generation + * (agentic loop, >5-min turn) is safe. + */ idleTimeoutMs?: number debug?: DebugOption } @@ -106,7 +107,8 @@ export interface WebSocketStreamInit { * Run a full-duplex, conversation-scoped chat over an already-accepted server * socket. Each inbound RunAgentInput frame starts one chat() turn (via onRun) * whose chunks are pumped back as frames; the socket stays open across turns - * (pending client-tool resubmit, next user message) until close/abort/idle. + * (pending client-tool resubmit, next user message) until the client closes it + * or the idle timeout fires. An abort control frame aborts only its turn. */ export function toWebSocketStream( socket: WebSocketLike, @@ -115,17 +117,22 @@ export function toWebSocketStream( ): void { const logger = resolveDebugOption(init.debug) const activeTurns = new Map() + // Abort frames that raced ahead of their run's registration: `handleInbound` + // awaits body validation before it registers into `activeTurns`, so an abort + // arriving inside that window would otherwise be silently discarded. + const earlyAborts = new Set() const heartbeatMs = init.heartbeatMs ?? 30_000 const idleTimeoutMs = init.idleTimeoutMs ?? 300_000 let lastActivity = Date.now() + let closed = false const heartbeat = setInterval(() => { try { socket.send(JSON.stringify({ type: 'ping' })) } catch { - // Socket is CLOSING/CLOSED between ticks — the close handler below - // clears this interval; swallow so the timer callback doesn't throw - // uncaught in the meantime. + // Socket is CLOSING/CLOSED between ticks — teardown below clears this + // interval; swallow so the timer callback doesn't throw uncaught in the + // meantime. } }, heartbeatMs) const idle = setInterval( @@ -140,11 +147,26 @@ export function toWebSocketStream( Math.min(idleTimeoutMs, 30_000), ) - socket.addEventListener('close', () => { + function teardown(): void { + closed = true for (const controller of activeTurns.values()) controller.abort() activeTurns.clear() clearInterval(heartbeat) clearInterval(idle) + } + + socket.addEventListener('close', teardown) + // Without this, an errored socket whose `close` never follows would leak + // both intervals and never abort its turns — and on `ws` (an EventEmitter) + // an `error` event with no listener is thrown as an uncaught exception. + socket.addEventListener('error', () => { + logger.errors('WebSocket errored; aborting its turns') + teardown() + try { + socket.close(1011, 'socket error') + } catch { + // socket already closing/closed — nothing to do + } }) socket.addEventListener('message', (event: { data: unknown }) => { @@ -165,32 +187,56 @@ export function toWebSocketStream( } if (frame.kind === 'abort') { - activeTurns.get(frame.runId)?.abort() + const turn = activeTurns.get(frame.runId) + if (turn) turn.abort() + else earlyAborts.add(frame.runId) return } - handleInbound(frame.input).catch((error: unknown) => { - logger.errors('Failed to handle inbound WS run frame; dropping it', { - error, - }) - }) + void handleInbound(frame.input) }) + /** + * Surface a turn failure to the client as a live `RUN_ERROR` frame. The + * socket is conversation-scoped and stays open, so without this frame the + * client would see neither a terminal chunk nor a close — a permanent hang. + * Mirrors the HTTP transports, which synthesize the live `RUN_ERROR` when + * the producer rethrows (see `durableStreamSource`'s terminal contract). + */ + function sendRunError(error: unknown): void { + try { + socket.send(encodeWsFrame(runErrorChunk(error), undefined)) + } catch { + // Socket is CLOSING/CLOSED — the client sees onclose instead. + } + } + async function handleInbound(input: unknown): Promise { - const params = await chatParamsFromRequestBody(input) + let params: Awaited> + try { + params = await chatParamsFromRequestBody(input) + } catch (error) { + logger.errors('Invalid inbound WS run frame; dropping it', { error }) + sendRunError(error) + return + } + // The socket may have closed (or errored) during the await above — the + // teardown that drains `activeTurns` already ran, so registering now + // would start a turn nothing can ever abort. + if (closed) return const turnAbort = new AbortController() // A second inbound frame with the same runId (client resubmit) must // abort the earlier turn. Otherwise the old controller is overwritten // and close/abort frames can no longer reach it. activeTurns.get(params.runId)?.abort() activeTurns.set(params.runId, turnAbort) + if (earlyAborts.delete(params.runId)) turnAbort.abort() const ctx: WsRunContext = { messages: params.messages, threadId: params.threadId, runId: params.runId, forwardedProps: params.forwardedProps, - resumeOffset: null, - request: buildTurnRequest(request, params.runId, null), + request: buildTurnRequest(request, params.runId), signal: turnAbort.signal, } try { @@ -213,6 +259,13 @@ export function toWebSocketStream( socket.send(encodeWsFrame(chunk, undefined)) } } + } catch (error) { + // An aborted turn (socket close, abort frame, same-runId resubmit) is + // expected teardown, not a turn failure — nothing to report. + if (!turnAbort.signal.aborted) { + logger.errors('WS turn failed', { error }) + sendRunError(error) + } } finally { // Only delete if this turn still owns the entry: a duplicate in-flight // runId (e.g. a client resubmitting before the first turn finished) @@ -256,6 +309,9 @@ export function resumeWebSocketStream( } const abortController = new AbortController() socket.addEventListener('close', () => abortController.abort()) + // An `error` with no listener is an uncaught exception on `ws`; abort the + // replay so the pump below stops instead of writing to a dead socket. + socket.addEventListener('error', () => abortController.abort()) const { source, getId } = durableStreamSource( emptyDurableSource(), options.adapter, @@ -322,11 +378,11 @@ function upgradeResponse(client: unknown): Response { } /** - * Cloudflare/Deno wrapper: creates a `WebSocketPair`, accepts the server - * socket, delegates to {@link toWebSocketStream}, and returns the 101 - * upgrade `Response` carrying the client socket. Throws when the runtime has - * no `WebSocketPair` (e.g. Node) — upgrade the socket yourself and call - * {@link toWebSocketStream} directly there. + * Cloudflare wrapper (Workers/Durable Objects): creates a `WebSocketPair`, + * accepts the server socket, delegates to {@link toWebSocketStream}, and + * returns the 101 upgrade `Response` carrying the client socket. Throws when + * the runtime has no `WebSocketPair` (Node, Deno, Bun) — upgrade the socket + * yourself and call {@link toWebSocketStream} directly there. */ export function toWebSocketResponse( request: Request, @@ -338,11 +394,11 @@ export function toWebSocketResponse( } /** - * Cloudflare/Deno wrapper: creates a `WebSocketPair`, accepts the server - * socket, delegates to {@link resumeWebSocketStream}, and returns the 101 - * upgrade `Response` carrying the client socket. Throws when the runtime has - * no `WebSocketPair` (e.g. Node) — upgrade the socket yourself and call - * {@link resumeWebSocketStream} directly there. + * Cloudflare wrapper (Workers/Durable Objects): creates a `WebSocketPair`, + * accepts the server socket, delegates to {@link resumeWebSocketStream}, and + * returns the 101 upgrade `Response` carrying the client socket. Throws when + * the runtime has no `WebSocketPair` (Node, Deno, Bun) — upgrade the socket + * yourself and call {@link resumeWebSocketStream} directly there. * * @example * ```ts diff --git a/packages/ai/tests/stream-to-websocket.test.ts b/packages/ai/tests/stream-to-websocket.test.ts index b23d888ad5..6a620ff1de 100644 --- a/packages/ai/tests/stream-to-websocket.test.ts +++ b/packages/ai/tests/stream-to-websocket.test.ts @@ -48,61 +48,55 @@ class FakeSocket implements WebSocketLike { sent: Array = [] closed = false closeCode: number | undefined - bufferedAmount = 0 - private handlers: Record void>> = {} + private handlers: Record void>> = {} send(data: string): void { + if (this.closed) throw new Error('socket is closed') this.sent.push(data) } close(code?: number): void { this.closed = true this.closeCode = code - this.emit('close', { code }) + this.emit('close', { data: undefined }) } - addEventListener(type: string, handler: (ev: any) => void): void { + addEventListener( + type: 'message' | 'close' | 'error', + handler: (ev: { data: unknown }) => void, + ): void { ;(this.handlers[type] ??= []).push(handler) } emitMessage(data: string): void { this.emit('message', { data }) } emitClose(): void { - this.emit('close', {}) + this.closed = true + this.emit('close', { data: undefined }) + } + emitError(): void { + this.emit('error', { data: undefined }) } - private emit(type: string, ev: any): void { + private emit(type: string, ev: { data: unknown }): void { for (const h of this.handlers[type] ?? []) h(ev) } } -describe('FakeSocket double', () => { - it('records sent frames and delivers messages to listeners', () => { - const socket = new FakeSocket() - const received: Array = [] - socket.addEventListener('message', (e) => received.push(e.data)) - socket.send('a') - socket.emitMessage('b') - expect(socket.sent).toEqual(['a']) - expect(received).toEqual(['b']) - }) -}) - describe('buildTurnRequest', () => { it('keys the synthetic request by runId and preserves headers', () => { const handshake = new Request('https://x/api/chat', { headers: { authorization: 'Bearer t' }, }) - const req = buildTurnRequest(handshake, 'run-9', null) + const req = buildTurnRequest(handshake, 'run-9') const url = new URL(req.url) expect(url.searchParams.get('runId')).toBe('run-9') expect(url.searchParams.get('offset')).toBeNull() expect(req.headers.get('authorization')).toBe('Bearer t') }) - it('adds ?offset on a reconnect', () => { + it('scrubs a handshake ?offset so a mis-routed resume cannot hit the replay branch', () => { const req = buildTurnRequest( - new Request('https://x/api/chat'), + new Request('https://x/api/chat?offset=off-3'), 'run-9', - 'off-3', ) - expect(new URL(req.url).searchParams.get('offset')).toBe('off-3') + expect(new URL(req.url).searchParams.get('offset')).toBeNull() }) }) @@ -273,6 +267,8 @@ describe('toWebSocketStream lifecycle', () => { expect(() => socket.emitMessage('{')).not.toThrow() await flush() + // Valid JSON with an invalid RunAgentInput shape surfaces as RUN_ERROR + // (see the dedicated test below) rather than crashing the socket. expect(() => socket.emitMessage(JSON.stringify({ foo: 'bar' })), ).not.toThrow() @@ -282,7 +278,157 @@ describe('toWebSocketStream lifecycle', () => { socket.emitMessage(inputFrame('run-5')) await flush() const types = socket.sent.map((s) => JSON.parse(s).type) - expect(types).toEqual(['RUN_STARTED']) + expect(types).toEqual(['RUN_ERROR', 'RUN_STARTED']) + }) + + it('surfaces a turn failure as a live RUN_ERROR frame and keeps the socket open', async () => { + const socket = new FakeSocket() + toWebSocketStream(socket, new Request('https://x/api/chat'), { + onRun: ({ runId, threadId }): AsyncIterable => + (async function* () { + yield ev.runStarted(runId, threadId) + throw new Error('provider exploded') + })(), + debug: false, + }) + socket.emitMessage(inputFrame('run-err')) + await flush() + + // The conversation socket stays open, so the failure MUST reach the + // client as a terminal frame — otherwise the consumer hangs forever. + const frames = socket.sent.map((s) => JSON.parse(s)) + expect(frames.map((f) => f.type)).toEqual(['RUN_STARTED', 'RUN_ERROR']) + expect(frames[1].message).toContain('provider exploded') + expect(socket.closed).toBe(false) + }) + + it('surfaces an invalid RunAgentInput body as a RUN_ERROR frame', async () => { + const socket = new FakeSocket() + toWebSocketStream(socket, new Request('https://x/api/chat'), { + onRun: () => (async function* () {})(), + debug: false, + }) + // Valid JSON, but not a RunAgentInput (missing threadId/runId/messages). + socket.emitMessage(JSON.stringify({ messages: 'nope' })) + await flush() + const types = socket.sent.map((s) => JSON.parse(s).type) + expect(types).toEqual(['RUN_ERROR']) + expect(socket.closed).toBe(false) + }) + + it('multiplexes two sequential turns over one socket with separate per-turn durability logs', async () => { + const socket = new FakeSocket() + const durabilityRunIds: Array = [] + toWebSocketStream(socket, new Request('https://x/api/chat'), { + durability: (ctx) => { + durabilityRunIds.push( + new URL(ctx.request.url).searchParams.get('runId'), + ) + return memoryStream(ctx.request) + }, + onRun: ({ runId, threadId }): AsyncIterable => + (async function* () { + yield ev.textContent(runId) + yield { + type: 'RUN_FINISHED', + runId, + threadId, + model: 'm', + finishReason: 'stop', + timestamp: Date.now(), + } as StreamChunk + })(), + }) + + socket.emitMessage(inputFrame('run-a')) + await flush() + socket.emitMessage(inputFrame('run-b')) + await flush() + + // Both turns streamed to completion over the same open socket, and each + // got its own durability log keyed by its own runId. + const frames = socket.sent.map((s) => JSON.parse(s)) + expect(frames.every((f) => typeof f.id === 'string' && 'chunk' in f)).toBe( + true, + ) + expect(frames.map((f) => f.chunk.type)).toEqual([ + 'CUSTOM', + 'TEXT_MESSAGE_CONTENT', + 'RUN_FINISHED', + 'CUSTOM', + 'TEXT_MESSAGE_CONTENT', + 'RUN_FINISHED', + ]) + expect(durabilityRunIds).toEqual(['run-a', 'run-b']) + // Each turn's replay log is independent: joining run-a replays only run-a. + const joinA = memoryStream( + new Request('https://x/api/chat?runId=run-a&offset=-1'), + ) + const replayed: Array = [] + for await (const entry of joinA.read('-1')) { + replayed.push(entry.chunk.type) + } + expect(replayed).toEqual(['CUSTOM', 'TEXT_MESSAGE_CONTENT', 'RUN_FINISHED']) + expect(socket.closed).toBe(false) + }) + + it('idle-closes a quiet socket (1000) but never one with a turn still in flight', async () => { + // Quiet socket: no inbound frame within idleTimeoutMs → reaped. + const quiet = new FakeSocket() + toWebSocketStream(quiet, new Request('https://x/api/chat'), { + onRun: () => (async function* () {})(), + idleTimeoutMs: 5, + heartbeatMs: 60_000, + }) + await new Promise((r) => setTimeout(r, 30)) + expect(quiet.closed).toBe(true) + expect(quiet.closeCode).toBe(1000) + + // Socket with a parked in-flight turn: a long generation sends no inbound + // frames, but the idle reaper must not kill live work. + const busy = new FakeSocket() + let release: (() => void) | undefined + toWebSocketStream(busy, new Request('https://x/api/chat'), { + onRun: ({ runId, threadId }): AsyncIterable => + (async function* () { + yield ev.runStarted(runId, threadId) + await new Promise((r) => { + release = r + }) + })(), + idleTimeoutMs: 5, + heartbeatMs: 60_000, + }) + busy.emitMessage(inputFrame('run-busy')) + await new Promise((r) => setTimeout(r, 30)) + expect(busy.closed).toBe(false) + release?.() + busy.emitClose() + }) + + it('tears down and aborts turns when the socket errors', async () => { + const socket = new FakeSocket() + let aborted = false + toWebSocketStream(socket, new Request('https://x/api/chat'), { + onRun: ({ signal }): AsyncIterable => + // eslint-disable-next-line require-yield + (async function* () { + await new Promise((resolve) => { + signal.addEventListener('abort', () => { + aborted = true + resolve() + }) + }) + })(), + debug: false, + }) + socket.emitMessage(inputFrame('run-oops')) + await flush() + socket.emitError() + await flush() + expect(aborted).toBe(true) + expect(socket.closed).toBe(true) + expect(socket.closeCode).toBe(1011) }) }) diff --git a/testing/e2e/src/lib/durable-delivery-ws-plugin.ts b/testing/e2e/src/lib/durable-delivery-ws-plugin.ts index 63a1eaa87d..b8a35cf59f 100644 --- a/testing/e2e/src/lib/durable-delivery-ws-plugin.ts +++ b/testing/e2e/src/lib/durable-delivery-ws-plugin.ts @@ -14,7 +14,7 @@ import type { WebSocketLike } from '@tanstack/ai' * NDJSON arms — same `fixedRun` sequence, same `memoryStream` durability * sink). The e2e app is served by Vite/Nitro on plain Node, which has no * `WebSocketPair`, so `toWebSocketResponse`/`resumeWebSocketResponse` (the - * CF/Deno wrapper) can't be used here. Instead this hooks the underlying + * Cloudflare wrapper) can't be used here. Instead this hooks the underlying * Node http server's `upgrade` event directly — the same pattern * `examples/ts-group-chat/chat-server/vite-plugin.ts` uses for its Cap'n Web * socket — obtains a server socket via `ws`'s `WebSocketServer({ noServer: @@ -30,20 +30,42 @@ import type { WebSocketLike } from '@tanstack/ai' * `{ type: 'ping' }` heartbeats to ignore. * - Connect to `/api/durable-delivery-ws?runId=&offset=` to * resume: no input frame is sent, the log replays strictly after `offset`. + * - `?dropAfter=` (fresh runs only) simulates a mid-run connection drop for + * the client-adapter reconnect e2e: the server closes the socket after + * forwarding n frames while the run keeps producing into the durability + * log, so the client's auto-reconnect can replay the remainder. */ const WS_PATH = '/api/durable-delivery-ws' -function isWebSocketLike(value: unknown): value is WebSocketLike { - if (typeof value !== 'object' || value === null) return false - if (!('send' in value) || typeof value.send !== 'function') return false - if (!('close' in value) || typeof value.close !== 'function') return false - if ( - !('addEventListener' in value) || - typeof value.addEventListener !== 'function' - ) { - return false +/** + * Wrap a socket so the client sees the connection drop after `dropAfter` + * outbound frames. Later frames are swallowed BY DESIGN (they're the "frames + * lost to a dead connection" half of the simulation) while the producer keeps + * appending them to the durability log; once the terminal frame has been + * produced (the log is complete) the socket closes, so the client's + * auto-reconnect deterministically replays the full remainder from the log. + */ +function dropAfterSocket(ws: WebSocketLike, dropAfter: number): WebSocketLike { + let sent = 0 + const addEventListener: WebSocketLike['addEventListener'] = ( + type: 'message' | 'close' | 'error', + handler: (ev: { data: unknown }) => void, + ) => { + if (type === 'message') ws.addEventListener(type, handler) + else ws.addEventListener(type, () => handler({ data: undefined })) + } + return { + send: (data) => { + sent += 1 + if (sent <= dropAfter) { + ws.send(data) + return + } + if (data.includes('RUN_FINISHED')) ws.close() + }, + close: (code, reason) => ws.close(code, reason), + addEventListener, } - return true } export function durableDeliveryWebSocketPlugin(): Plugin { @@ -62,18 +84,23 @@ export function durableDeliveryWebSocketPlugin(): Plugin { if (url.pathname !== WS_PATH) return wss.handleUpgrade(req, socket, head, (ws) => { + // `ws`'s socket satisfies WebSocketLike structurally — no adapter + // (and no runtime guard) needed. const request = new Request(url) - if (!isWebSocketLike(ws)) return - if (url.searchParams.get('offset') !== null) { resumeWebSocketStream(ws, { adapter: memoryStream(request), }) } else { - toWebSocketStream(ws, request, { - durability: (ctx) => memoryStream(ctx.request), - onRun: ({ runId, threadId }) => fixedRun(threadId, runId), - }) + const dropAfter = Number(url.searchParams.get('dropAfter')) + toWebSocketStream( + dropAfter > 0 ? dropAfterSocket(ws, dropAfter) : ws, + request, + { + durability: (ctx) => memoryStream(ctx.request), + onRun: ({ runId, threadId }) => fixedRun(threadId, runId), + }, + ) } }) }) diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index c6d9853490..00af916dcc 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -9,6 +9,7 @@ // Additionally, you should also exclude this file from your linter and/or formatter to prevent it from being checked or modified. import { Route as rootRouteImport } from './routes/__root' +import { Route as WebsocketAdapterRouteImport } from './routes/websocket-adapter' import { Route as ToolsTestRouteImport } from './routes/tools-test' import { Route as PersistenceDurabilityRouteImport } from './routes/persistence-durability' import { Route as MiddlewareTestRouteImport } from './routes/middleware-test' @@ -77,6 +78,11 @@ import { Route as ApiTranscriptionStreamRouteImport } from './routes/api.transcr import { Route as ApiImageStreamRouteImport } from './routes/api.image.stream' import { Route as ApiAudioStreamRouteImport } from './routes/api.audio.stream' +const WebsocketAdapterRoute = WebsocketAdapterRouteImport.update({ + id: '/websocket-adapter', + path: '/websocket-adapter', + getParentRoute: () => rootRouteImport, +} as any) const ToolsTestRoute = ToolsTestRouteImport.update({ id: '/tools-test', path: '/tools-test', @@ -441,6 +447,7 @@ export interface FileRoutesByFullPath { '/middleware-test': typeof MiddlewareTestRoute '/persistence-durability': typeof PersistenceDurabilityRoute '/tools-test': typeof ToolsTestRoute + '/websocket-adapter': typeof WebsocketAdapterRoute '/$provider/$feature': typeof ProviderFeatureRoute '/api/anthropic-bug-test': typeof ApiAnthropicBugTestRoute '/api/anthropic-skills-wire': typeof ApiAnthropicSkillsWireRoute @@ -510,6 +517,7 @@ export interface FileRoutesByTo { '/middleware-test': typeof MiddlewareTestRoute '/persistence-durability': typeof PersistenceDurabilityRoute '/tools-test': typeof ToolsTestRoute + '/websocket-adapter': typeof WebsocketAdapterRoute '/$provider/$feature': typeof ProviderFeatureRoute '/api/anthropic-bug-test': typeof ApiAnthropicBugTestRoute '/api/anthropic-skills-wire': typeof ApiAnthropicSkillsWireRoute @@ -580,6 +588,7 @@ export interface FileRoutesById { '/middleware-test': typeof MiddlewareTestRoute '/persistence-durability': typeof PersistenceDurabilityRoute '/tools-test': typeof ToolsTestRoute + '/websocket-adapter': typeof WebsocketAdapterRoute '/$provider/$feature': typeof ProviderFeatureRoute '/api/anthropic-bug-test': typeof ApiAnthropicBugTestRoute '/api/anthropic-skills-wire': typeof ApiAnthropicSkillsWireRoute @@ -651,6 +660,7 @@ export interface FileRouteTypes { | '/middleware-test' | '/persistence-durability' | '/tools-test' + | '/websocket-adapter' | '/$provider/$feature' | '/api/anthropic-bug-test' | '/api/anthropic-skills-wire' @@ -720,6 +730,7 @@ export interface FileRouteTypes { | '/middleware-test' | '/persistence-durability' | '/tools-test' + | '/websocket-adapter' | '/$provider/$feature' | '/api/anthropic-bug-test' | '/api/anthropic-skills-wire' @@ -789,6 +800,7 @@ export interface FileRouteTypes { | '/middleware-test' | '/persistence-durability' | '/tools-test' + | '/websocket-adapter' | '/$provider/$feature' | '/api/anthropic-bug-test' | '/api/anthropic-skills-wire' @@ -859,6 +871,7 @@ export interface RootRouteChildren { MiddlewareTestRoute: typeof MiddlewareTestRoute PersistenceDurabilityRoute: typeof PersistenceDurabilityRoute ToolsTestRoute: typeof ToolsTestRoute + WebsocketAdapterRoute: typeof WebsocketAdapterRoute ProviderFeatureRoute: typeof ProviderFeatureRoute ApiAnthropicBugTestRoute: typeof ApiAnthropicBugTestRoute ApiAnthropicSkillsWireRoute: typeof ApiAnthropicSkillsWireRoute @@ -908,6 +921,13 @@ export interface RootRouteChildren { declare module '@tanstack/react-router' { interface FileRoutesByPath { + '/websocket-adapter': { + id: '/websocket-adapter' + path: '/websocket-adapter' + fullPath: '/websocket-adapter' + preLoaderRoute: typeof WebsocketAdapterRouteImport + parentRoute: typeof rootRouteImport + } '/tools-test': { id: '/tools-test' path: '/tools-test' @@ -1456,6 +1476,7 @@ const rootRouteChildren: RootRouteChildren = { MiddlewareTestRoute: MiddlewareTestRoute, PersistenceDurabilityRoute: PersistenceDurabilityRoute, ToolsTestRoute: ToolsTestRoute, + WebsocketAdapterRoute: WebsocketAdapterRoute, ProviderFeatureRoute: ProviderFeatureRoute, ApiAnthropicBugTestRoute: ApiAnthropicBugTestRoute, ApiAnthropicSkillsWireRoute: ApiAnthropicSkillsWireRoute, diff --git a/testing/e2e/src/routes/websocket-adapter.tsx b/testing/e2e/src/routes/websocket-adapter.tsx new file mode 100644 index 0000000000..3a04d94991 --- /dev/null +++ b/testing/e2e/src/routes/websocket-adapter.tsx @@ -0,0 +1,81 @@ +import { useEffect, useState } from 'react' +import { createFileRoute } from '@tanstack/react-router' +import { webSocket } from '@tanstack/ai-react' + +export const Route = createFileRoute('/websocket-adapter')({ + component: WebSocketAdapterPage, +}) + +/** + * Client-adapter arm of the WebSocket delivery-durability harness: drives the + * REAL `webSocket()` connection adapter (not a hand-rolled socket) against + * `/api/durable-delivery-ws`, so the adapter's envelope unwrapping, ping + * filtering, offset de-dupe, and auto-reconnect are what the e2e asserts. + * + * "run" streams the fixed sequence over one socket. "run with drop" adds + * `?dropAfter=2`, making the server drop the socket after two frames — the + * adapter must then reconnect at `?runId=&offset=` and resume the remainder + * from the durability log, exactly once and in order. + */ +function WebSocketAdapterPage() { + const [result, setResult] = useState('') + // Buttons enable only after hydration so a Playwright click can't land + // before React has attached the handlers (click auto-waits for enabled). + const [ready, setReady] = useState(false) + useEffect(() => setReady(true), []) + + async function run(dropAfter: number | null) { + const runId = `e2e-ws-adapter-${dropAfter === null ? 'plain' : 'drop'}-${Date.now()}` + const base = `${location.origin.replace(/^http/, 'ws')}/api/durable-delivery-ws` + const connection = webSocket( + dropAfter === null ? base : `${base}?dropAfter=${dropAfter}`, + { reconnect: { delayMs: 25, maxAttempts: 5 } }, + ) + const controller = new AbortController() + const received: Array<{ type: string; delta?: string }> = [] + const consume = (async () => { + for await (const chunk of connection.subscribe(controller.signal)) { + received.push({ + type: chunk.type, + ...('delta' in chunk && typeof chunk.delta === 'string' + ? { delta: chunk.delta } + : {}), + }) + if (chunk.type === 'RUN_FINISHED' || chunk.type === 'RUN_ERROR') { + controller.abort() + } + } + })() + try { + await connection.send([], undefined, undefined, { + threadId: 'thread-durable', + runId, + }) + await consume + setResult(JSON.stringify({ runId, received })) + } catch (error) { + setResult(JSON.stringify({ runId, received, error: String(error) })) + } + } + + return ( +
+

webSocket() adapter harness

+ + +
{result}
+
+ ) +} diff --git a/testing/e2e/tests/websocket.spec.ts b/testing/e2e/tests/websocket.spec.ts index ca5a2b2379..a212a9986d 100644 --- a/testing/e2e/tests/websocket.spec.ts +++ b/testing/e2e/tests/websocket.spec.ts @@ -174,4 +174,28 @@ test.describe('delivery durability (websocket)', () => { expect(resumedIds).not.toContain(result.lastId) expect(new Set(resumedIds).size).toBe(resumedIds.length) }) + + test('the webSocket() client adapter auto-resumes a server-side drop exactly once, in order', async ({ + page, + }) => { + // Unlike the raw-WebSocket tests above, this drives the REAL client + // adapter (see /websocket-adapter): the server drops the socket after two + // frames (?dropAfter=2), and the adapter's own reconnect machinery must + // reopen at ?runId=&offset= and resume the remainder from the log. + await page.goto('/websocket-adapter') + await page.getByTestId('ws-adapter-run-drop').click() + + const resultText = page.getByTestId('ws-adapter-result') + await expect(resultText).toContainText('RUN_FINISHED', { timeout: 15_000 }) + const { received } = JSON.parse( + (await resultText.textContent()) ?? '{}', + ) as { received: Array<{ type: string; delta?: string }> } + + // The full fixed sequence arrived despite the mid-run drop — no gap where + // the connection died, no duplicate around the resume boundary. + expect(received.map((c) => c.type)).toEqual(FIXED_TYPES) + expect( + received.filter((c) => c.delta !== undefined).map((c) => c.delta), + ).toEqual(['1', '2', '3', '4', '5']) + }) })