diff --git a/.changeset/websocket-transport.md b/.changeset/websocket-transport.md new file mode 100644 index 000000000..f4ca05f40 --- /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. diff --git a/docs/chat/connection-adapters.md b/docs/chat/connection-adapters.md index 96ff035f9..0e19f194d 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 55a365282..8ef58756b 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 640e172b4..b8081c077 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 01f2b856f..614b01387 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 000000000..d418df586 --- /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. diff --git a/examples/ts-react-chat/package.json b/examples/ts-react-chat/package.json index 67b58ad43..6439e7282 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 9c71d1577..205487525 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 000000000..f716a5345 --- /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 0f05c32c4..d8f99d305 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 000000000..55d7da2c5 --- /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 9bbd20764..24052e99a 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-angular/src/index.ts b/packages/ai-angular/src/index.ts index 4f9565c58..fb2093f62 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-client/src/connection-adapters.ts b/packages/ai-client/src/connection-adapters.ts index 0883ae77a..0258e3d14 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 @@ -1560,6 +1602,316 @@ 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 +} + +/** 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. + * + * 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), + options: WebSocketConnectionOptions = {}, +): SubscribeConnectionAdapter & { + joinRun: ( + runId: string, + abortSignal?: AbortSignal, + ) => AsyncIterable +} { + 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() + 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 + const ws = options.protocols + ? new Impl(target, options.protocols) + : new Impl(target) + socket = ws + 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. + openPromise.catch(() => {}) + ws.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 envelopeId = isNdjsonEnvelope(parsed) ? parsed.id : undefined + const chunk = isNdjsonEnvelope(parsed) + ? parsed.chunk + : (parsed as StreamChunk) + + // 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 + // memoized promise (set once in `openOnce`), so none of them clobber + // another's onopen/onerror handler. + return openPromise ?? Promise.resolve() + } + + 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 } + 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) + } + })() + }, + 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, + }) + 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> = [] + 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* () { + 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) + } + })() + }, + } +} + /** * 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 41605186c..f99e73b5f 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-resumable.test.ts b/packages/ai-client/tests/connection-adapters-resumable.test.ts index d60602b8f..59414100f 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 + }) +}) 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 000000000..553f33420 --- /dev/null +++ b/packages/ai-client/tests/connection-adapters-websocket.test.ts @@ -0,0 +1,425 @@ +import { describe, expect, it } from 'vitest' +import { + StreamReconnectLimitError, + 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) + } + })() +} + +/** 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 = [] + 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']) + }) + + 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' }, + ) + + 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) + 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') + }) +}) + +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', + ]) + }) +}) + +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) + }) +}) diff --git a/packages/ai-react/src/index.ts b/packages/ai-react/src/index.ts index 773c05456..c563a9c28 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/packages/ai-solid/src/index.ts b/packages/ai-solid/src/index.ts index 1d0c517b3..0d16a53a5 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 151fdab5d..0d9004d03 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 1d0c517b3..0d16a53a5 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, diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 6c3cf5a27..6c0e8b220 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-response.ts b/packages/ai/src/stream-to-response.ts index 09fda5b98..4f4674a94 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 new file mode 100644 index 000000000..009cd92bf --- /dev/null +++ b/packages/ai/src/stream-to-websocket.ts @@ -0,0 +1,358 @@ +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' + +/** + * 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 } + | { 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 } +} + +/** Per-turn context for one inbound `run` frame on a conversation-scoped socket. */ +export interface WsRunContext { + messages: 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 }) +} + +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 { + 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(() => { + 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( + () => { + // 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), + ) + + 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 + 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(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, + runId: params.runId, + forwardedProps: params.forwardedProps, + resumeOffset: null, + request: buildTurnRequest(request, params.runId, null), + signal: turnAbort.signal, + } + 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 { + // 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 { + const logger = resolveDebugOption(options.debug) + 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, + }, + ) + void (async () => { + 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 { + socket.close(1011, 'resume failed') + } catch { + // socket already closing/closed — nothing to do + } + }) +} + +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. + * + * @example + * ```ts + * resumeWebSocketResponse({ adapter: memoryStream(request) }) + * ``` + */ +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 new file mode 100644 index 000000000..e93da4c2b --- /dev/null +++ b/packages/ai/tests/stream-to-websocket.test.ts @@ -0,0 +1,374 @@ +import { describe, expect, it } from 'vitest' +import { + buildTurnRequest, + decodeWsFrame, + encodeWsFrame, + resumeWebSocketResponse, + resumeWebSocketStream, + toWebSocketResponse, + 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 { StreamDurability } from '../src/stream-durability' +import type { StreamChunk } from '../src/types' + +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' }) + }) +}) + +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']) + }) +}) + +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') + }) +}) + +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 + }) +}) + +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', + ]) + }) +}) + +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']) + }) +}) + +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']) + // 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 () => { + 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) + }) + + 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) + }) +}) + +describe('toWebSocketResponse', () => { + it('throws a directive error when WebSocketPair is unavailable', () => { + expect(() => + toWebSocketResponse(new Request('https://x/api/chat'), { + onRun: () => (async function* () {})(), + }), + ).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/) + }) +}) diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 5dccceec4..5e8ded4f8 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 @@ -2621,6 +2627,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 +2646,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)) @@ -16028,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'} @@ -16220,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) @@ -16309,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) @@ -17979,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 @@ -21973,7 +21902,7 @@ snapshots: '@tanstack/devtools-event-bus@0.4.1': dependencies: - ws: 8.19.0 + ws: 8.21.0 transitivePeerDependencies: - bufferutil - utf-8-validate @@ -23890,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: @@ -31816,7 +31744,6 @@ snapshots: jsdom: 27.3.0(postcss@8.5.15) transitivePeerDependencies: - msw - optional: true vlq@1.0.1: {} @@ -32031,8 +31958,6 @@ snapshots: ws@8.18.3: {} - ws@8.19.0: {} - ws@8.20.1: {} ws@8.21.0: {} diff --git a/testing/e2e/package.json b/testing/e2e/package.json index b761eb58a..92ee1fc2b 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 000000000..43d50bb5f --- /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 fb32f9519..980f2649f 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 000000000..c7e891504 --- /dev/null +++ b/testing/e2e/tests/websocket.spec.ts @@ -0,0 +1,173 @@ +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 e522fe93b..416f80f0d 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'],