@tanstack/ai 0.45.0 → 0.46.0

package/dist/esm/stream-to-websocket.d.ts

@@ -0,0 +1,123 @@
+import { StreamDurability } from './stream-durability.js';
+import { DebugOption } from './logger/types.js';
+import { ModelMessage, StreamChunk, UIMessage } from './types.js';
+/**
+ * The minimal WHATWG WebSocket surface the core needs. Cloudflare
+ * `WebSocketPair` server sockets, Deno's upgraded sockets, and `ws` (Node)
+ * sockets already satisfy it; Bun's `ServerWebSocket` (handler-object API)
+ * gets a ~10-line adapter at the call site.
+ */
+export interface WebSocketLike {
+    send: (data: string) => void;
+    close: (code?: number, reason?: string) => void;
+    addEventListener: {
+        (type: 'message', handler: (ev: {
+            data: unknown;
+        }) => void): void;
+        (type: 'close' | 'error', handler: () => void): void;
+    };
+}
+/** 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 declare function encodeWsFrame(chunk: StreamChunk, id: string | undefined): string;
+/**
+ * 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 declare function decodeWsFrame(data: string): InboundFrame;
+/** Per-turn context for one inbound `run` frame on a conversation-scoped socket. */
+export interface WsRunContext {
+    messages: Array<UIMessage | ModelMessage>;
+    threadId: string;
+    runId: string;
+    forwardedProps?: Record<string, unknown>;
+    /** 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. A handshake carrying `?offset` is a resume and never
+ * reaches a fresh turn (`resumeWebSocketStream` serves it), so the offset is
+ * scrubbed here — otherwise a mis-routed resume handshake would make the turn's
+ * durability adapter silently take the replay branch instead of running onRun.
+ */
+export declare function buildTurnRequest(handshake: Request, runId: string): Request;
+export interface WebSocketStreamInit<TOffset extends string = string> {
+    /** Build a fresh chat() stream for each inbound RunAgentInput frame. */
+    onRun: (ctx: WsRunContext) => AsyncIterable<StreamChunk>;
+    /** Per-TURN durability factory, keyed by the frame's runId via ctx.request. */
+    durability?: (ctx: WsRunContext) => StreamDurability<TOffset>;
+    /** 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).
+     * Never fires while a turn is still streaming, so a long single generation
+     * (agentic loop, >5-min turn) is safe.
+     */
+    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 the client closes it
+ * or the idle timeout fires. An abort control frame aborts only its turn.
+ */
+export declare function toWebSocketStream<TOffset extends string = string>(socket: WebSocketLike, request: Request, init: WebSocketStreamInit<TOffset>): void;
+/**
+ * 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 declare function resumeWebSocketStream<TOffset extends string = string>(socket: WebSocketLike, options: {
+    adapter: StreamDurability<TOffset>;
+    batch?: number;
+    debug?: DebugOption;
+}): void;
+/**
+ * Cloudflare wrapper (Workers/Durable Objects): creates a `WebSocketPair`,
+ * accepts the server socket, delegates to {@link toWebSocketStream}, and
+ * returns the 101 upgrade `Response` carrying the client socket. Throws when
+ * the runtime has no `WebSocketPair` (Node, Deno, Bun) — upgrade the socket
+ * yourself and call {@link toWebSocketStream} directly there.
+ */
+export declare function toWebSocketResponse<TOffset extends string = string>(request: Request, init: WebSocketStreamInit<TOffset>): Response;
+/**
+ * Cloudflare wrapper (Workers/Durable Objects): creates a `WebSocketPair`,
+ * accepts the server socket, delegates to {@link resumeWebSocketStream}, and
+ * returns the 101 upgrade `Response` carrying the client socket. Throws when
+ * the runtime has no `WebSocketPair` (Node, Deno, Bun) — upgrade the socket
+ * yourself and call {@link resumeWebSocketStream} directly there.
+ *
+ * @example
+ * ```ts
+ * resumeWebSocketResponse({ adapter: memoryStream(request) })
+ * ```
+ */
+export declare function resumeWebSocketResponse<TOffset extends string = string>(options: {
+    adapter: StreamDurability<TOffset>;
+    batch?: number;
+    debug?: DebugOption;
+}): Response;

package/dist/esm/stream-to-websocket.js

@@ -0,0 +1,249 @@
+import { resolveDebugOption } from "./logger/resolve.js";
+import { durableStreamSource, runErrorChunk } from "./stream-to-response.js";
+import { chatParamsFromRequestBody } from "./utilities/chat-params.js";
+//#region src/stream-to-websocket.ts
+/**
+* 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.
+*/
+function encodeWsFrame(chunk, id) {
+	return JSON.stringify(id === void 0 ? 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`.
+*/
+function decodeWsFrame(data) {
+	const parsed = JSON.parse(data);
+	if (typeof parsed === "object" && parsed !== null && parsed.type === "abort" && typeof parsed.runId === "string") return {
+		kind: "abort",
+		runId: parsed.runId
+	};
+	return {
+		kind: "run",
+		input: parsed
+	};
+}
+/**
+* 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. A handshake carrying `?offset` is a resume and never
+* reaches a fresh turn (`resumeWebSocketStream` serves it), so the offset is
+* scrubbed here — otherwise a mis-routed resume handshake would make the turn's
+* durability adapter silently take the replay branch instead of running onRun.
+*/
+function buildTurnRequest(handshake, runId) {
+	const url = new URL(handshake.url);
+	url.searchParams.set("runId", runId);
+	url.searchParams.delete("offset");
+	return new Request(url, { headers: handshake.headers });
+}
+/**
+* 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 the client closes it
+* or the idle timeout fires. An abort control frame aborts only its turn.
+*/
+function toWebSocketStream(socket, request, init) {
+	const logger = resolveDebugOption(init.debug);
+	const activeTurns = /* @__PURE__ */ new Map();
+	const earlyAborts = /* @__PURE__ */ new Set();
+	const heartbeatMs = init.heartbeatMs ?? 3e4;
+	const idleTimeoutMs = init.idleTimeoutMs ?? 3e5;
+	let lastActivity = Date.now();
+	let closed = false;
+	const heartbeat = setInterval(() => {
+		try {
+			socket.send(JSON.stringify({ type: "ping" }));
+		} catch {}
+	}, heartbeatMs);
+	const idle = setInterval(() => {
+		if (activeTurns.size === 0 && Date.now() - lastActivity > idleTimeoutMs) socket.close(1e3, "idle");
+	}, Math.min(idleTimeoutMs, 3e4));
+	function teardown() {
+		closed = true;
+		for (const controller of activeTurns.values()) controller.abort();
+		activeTurns.clear();
+		clearInterval(heartbeat);
+		clearInterval(idle);
+	}
+	socket.addEventListener("close", teardown);
+	socket.addEventListener("error", () => {
+		logger.errors("WebSocket errored; aborting its turns");
+		teardown();
+		try {
+			socket.close(1011, "socket error");
+		} catch {}
+	});
+	socket.addEventListener("message", (event) => {
+		if (typeof event.data !== "string") return;
+		lastActivity = Date.now();
+		let frame;
+		try {
+			frame = decodeWsFrame(event.data);
+		} catch (error) {
+			logger.errors("Failed to decode inbound WS frame; dropping it", { error });
+			return;
+		}
+		if (frame.kind === "abort") {
+			const turn = activeTurns.get(frame.runId);
+			if (turn) turn.abort();
+			else earlyAborts.add(frame.runId);
+			return;
+		}
+		handleInbound(frame.input);
+	});
+	/**
+	* Surface a turn failure to the client as a live `RUN_ERROR` frame. The
+	* socket is conversation-scoped and stays open, so without this frame the
+	* client would see neither a terminal chunk nor a close — a permanent hang.
+	* Mirrors the HTTP transports, which synthesize the live `RUN_ERROR` when
+	* the producer rethrows (see `durableStreamSource`'s terminal contract).
+	*/
+	function sendRunError(error) {
+		try {
+			socket.send(encodeWsFrame(runErrorChunk(error), void 0));
+		} catch {}
+	}
+	async function handleInbound(input) {
+		let params;
+		try {
+			params = await chatParamsFromRequestBody(input);
+		} catch (error) {
+			logger.errors("Invalid inbound WS run frame; dropping it", { error });
+			sendRunError(error);
+			return;
+		}
+		if (closed) return;
+		const turnAbort = new AbortController();
+		activeTurns.get(params.runId)?.abort();
+		activeTurns.set(params.runId, turnAbort);
+		if (earlyAborts.delete(params.runId)) turnAbort.abort();
+		const ctx = {
+			messages: params.messages,
+			threadId: params.threadId,
+			runId: params.runId,
+			forwardedProps: params.forwardedProps,
+			request: buildTurnRequest(request, params.runId),
+			signal: turnAbort.signal
+		};
+		try {
+			if (init.durability) {
+				const adapter = init.durability(ctx);
+				const { source, getId } = durableStreamSource(init.onRun(ctx), adapter, {
+					abortController: turnAbort,
+					...init.batch === void 0 ? {} : { 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, void 0));
+		} catch (error) {
+			if (!turnAbort.signal.aborted) {
+				logger.errors("WS turn failed", { error });
+				sendRunError(error);
+			}
+		} finally {
+			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() {
+	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.
+*/
+function resumeWebSocketStream(socket, options) {
+	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());
+	socket.addEventListener("error", () => abortController.abort());
+	const { source, getId } = durableStreamSource(emptyDurableSource(), options.adapter, {
+		abortController,
+		...options.batch === void 0 ? {} : { batch: options.batch },
+		logger
+	});
+	(async () => {
+		for await (const chunk of source) socket.send(encodeWsFrame(chunk, getId(chunk)));
+		try {
+			socket.close(1e3);
+		} catch {}
+	})().catch((error) => {
+		logger.errors("resume websocket replay failed", { error });
+		try {
+			socket.close(1011, "resume failed");
+		} catch {}
+	});
+}
+function upgradeOrThrow(helper) {
+	const Pair = globalThis.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) {
+	return new Response(null, {
+		status: 101,
+		webSocket: client
+	});
+}
+/**
+* Cloudflare wrapper (Workers/Durable Objects): creates a `WebSocketPair`,
+* accepts the server socket, delegates to {@link toWebSocketStream}, and
+* returns the 101 upgrade `Response` carrying the client socket. Throws when
+* the runtime has no `WebSocketPair` (Node, Deno, Bun) — upgrade the socket
+* yourself and call {@link toWebSocketStream} directly there.
+*/
+function toWebSocketResponse(request, init) {
+	const { client, server } = upgradeOrThrow("toWebSocketResponse");
+	toWebSocketStream(server, request, init);
+	return upgradeResponse(client);
+}
+/**
+* Cloudflare wrapper (Workers/Durable Objects): creates a `WebSocketPair`,
+* accepts the server socket, delegates to {@link resumeWebSocketStream}, and
+* returns the 101 upgrade `Response` carrying the client socket. Throws when
+* the runtime has no `WebSocketPair` (Node, Deno, Bun) — upgrade the socket
+* yourself and call {@link resumeWebSocketStream} directly there.
+*
+* @example
+* ```ts
+* resumeWebSocketResponse({ adapter: memoryStream(request) })
+* ```
+*/
+function resumeWebSocketResponse(options) {
+	const { client, server } = upgradeOrThrow("resumeWebSocketResponse");
+	resumeWebSocketStream(server, options);
+	return upgradeResponse(client);
+}
+//#endregion
+export { buildTurnRequest, decodeWsFrame, encodeWsFrame, resumeWebSocketResponse, resumeWebSocketStream, toWebSocketResponse, toWebSocketStream };
+
+//# sourceMappingURL=stream-to-websocket.js.map

package/dist/esm/stream-to-websocket.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"stream-to-websocket.js","names":[],"sources":["../../src/stream-to-websocket.ts"],"sourcesContent":["import { chatParamsFromRequestBody } from './utilities/chat-params'\nimport { durableStreamSource, runErrorChunk } from './stream-to-response'\nimport { resolveDebugOption } from './logger/resolve'\nimport type { StreamDurability } from './stream-durability'\nimport type { DebugOption } from './logger/types'\nimport type { ModelMessage, StreamChunk, UIMessage } from './types'\n\n/**\n * The minimal WHATWG WebSocket surface the core needs. Cloudflare\n * `WebSocketPair` server sockets, Deno's upgraded sockets, and `ws` (Node)\n * sockets already satisfy it; Bun's `ServerWebSocket` (handler-object API)\n * gets a ~10-line adapter at the call site.\n */\nexport interface WebSocketLike {\n  send: (data: string) => void\n  close: (code?: number, reason?: string) => void\n  addEventListener: {\n    (type: 'message', handler: (ev: { data: unknown }) => void): void\n    (type: 'close' | 'error', handler: () => void): void\n  }\n}\n\n/** One inbound WS text frame, after JSON parse + shape discrimination. */\nexport type InboundFrame =\n  | { kind: 'run'; input: unknown }\n  | { kind: 'abort'; runId: string }\n\n/**\n * Encode one server→client frame. Durable frames carry the opaque offset in an\n * `{ id, chunk }` envelope (identical to the NDJSON wire); non-durable frames\n * are the bare chunk. Unambiguous because a bare chunk always has a top-level\n * `type` and the envelope never does.\n */\nexport function encodeWsFrame(\n  chunk: StreamChunk,\n  id: string | undefined,\n): string {\n  return JSON.stringify(id === undefined ? chunk : { id, chunk })\n}\n\n/**\n * Decode one client→server frame. An `{ type: 'abort', runId }` object is a\n * control frame; anything else is treated as a `RunAgentInput` and validated\n * downstream by `chatParamsFromRequestBody`.\n */\nexport function decodeWsFrame(data: string): InboundFrame {\n  const parsed: unknown = JSON.parse(data)\n  if (\n    typeof parsed === 'object' &&\n    parsed !== null &&\n    (parsed as { type?: unknown }).type === 'abort' &&\n    typeof (parsed as { runId?: unknown }).runId === 'string'\n  ) {\n    return { kind: 'abort', runId: (parsed as { runId: string }).runId }\n  }\n  return { kind: 'run', input: parsed }\n}\n\n/** Per-turn context for one inbound `run` frame on a conversation-scoped socket. */\nexport interface WsRunContext {\n  messages: Array<UIMessage | ModelMessage>\n  threadId: string\n  runId: string\n  forwardedProps?: Record<string, unknown>\n  /** Synthetic per-turn request carrying `?runId=` so durability keys correctly. */\n  request: Request\n  /** Aborts on socket close or an `abort` control frame for this run. */\n  signal: AbortSignal\n}\n\n/**\n * Build the synthetic per-turn request. A conversation-scoped socket multiplexes\n * many runs; each turn's durability adapter must key on the frame's `runId`,\n * which we carry in the URL query (`memoryStream`/`durableStream` already read\n * `?runId` / `?offset` there). Headers are copied from the handshake so\n * auth/cookies survive. A handshake carrying `?offset` is a resume and never\n * reaches a fresh turn (`resumeWebSocketStream` serves it), so the offset is\n * scrubbed here — otherwise a mis-routed resume handshake would make the turn's\n * durability adapter silently take the replay branch instead of running onRun.\n */\nexport function buildTurnRequest(handshake: Request, runId: string): Request {\n  const url = new URL(handshake.url)\n  url.searchParams.set('runId', runId)\n  url.searchParams.delete('offset')\n  return new Request(url, { headers: handshake.headers })\n}\n\nexport interface WebSocketStreamInit<TOffset extends string = string> {\n  /** Build a fresh chat() stream for each inbound RunAgentInput frame. */\n  onRun: (ctx: WsRunContext) => AsyncIterable<StreamChunk>\n  /** Per-TURN durability factory, keyed by the frame's runId via ctx.request. */\n  durability?: (ctx: WsRunContext) => StreamDurability<TOffset>\n  /** Chunks buffered per durability append (default 32). */\n  batch?: number\n  /** Heartbeat ping interval in ms (default 30_000). */\n  heartbeatMs?: number\n  /**\n   * Close after this many ms without any inbound frame (default 300_000).\n   * Never fires while a turn is still streaming, so a long single generation\n   * (agentic loop, >5-min turn) is safe.\n   */\n  idleTimeoutMs?: number\n  debug?: DebugOption\n}\n\n/**\n * Run a full-duplex, conversation-scoped chat over an already-accepted server\n * socket. Each inbound RunAgentInput frame starts one chat() turn (via onRun)\n * whose chunks are pumped back as frames; the socket stays open across turns\n * (pending client-tool resubmit, next user message) until the client closes it\n * or the idle timeout fires. An abort control frame aborts only its turn.\n */\nexport function toWebSocketStream<TOffset extends string = string>(\n  socket: WebSocketLike,\n  request: Request,\n  init: WebSocketStreamInit<TOffset>,\n): void {\n  const logger = resolveDebugOption(init.debug)\n  const activeTurns = new Map<string, AbortController>()\n  // Abort frames that raced ahead of their run's registration: `handleInbound`\n  // awaits body validation before it registers into `activeTurns`, so an abort\n  // arriving inside that window would otherwise be silently discarded.\n  const earlyAborts = new Set<string>()\n  const heartbeatMs = init.heartbeatMs ?? 30_000\n  const idleTimeoutMs = init.idleTimeoutMs ?? 300_000\n  let lastActivity = Date.now()\n  let closed = false\n\n  const heartbeat = setInterval(() => {\n    try {\n      socket.send(JSON.stringify({ type: 'ping' }))\n    } catch {\n      // Socket is CLOSING/CLOSED between ticks — teardown below clears this\n      // interval; swallow so the timer callback doesn't throw uncaught in the\n      // meantime.\n    }\n  }, heartbeatMs)\n  const idle = setInterval(\n    () => {\n      // Never idle-reap while a turn is in flight: a long single onRun\n      // iteration (agentic loop / >5-min generation) sends no INBOUND\n      // frames, so idle would otherwise fire and kill live work.\n      if (activeTurns.size === 0 && Date.now() - lastActivity > idleTimeoutMs) {\n        socket.close(1000, 'idle')\n      }\n    },\n    Math.min(idleTimeoutMs, 30_000),\n  )\n\n  function teardown(): void {\n    closed = true\n    for (const controller of activeTurns.values()) controller.abort()\n    activeTurns.clear()\n    clearInterval(heartbeat)\n    clearInterval(idle)\n  }\n\n  socket.addEventListener('close', teardown)\n  // Without this, an errored socket whose `close` never follows would leak\n  // both intervals and never abort its turns — and on `ws` (an EventEmitter)\n  // an `error` event with no listener is thrown as an uncaught exception.\n  socket.addEventListener('error', () => {\n    logger.errors('WebSocket errored; aborting its turns')\n    teardown()\n    try {\n      socket.close(1011, 'socket error')\n    } catch {\n      // socket already closing/closed — nothing to do\n    }\n  })\n\n  socket.addEventListener('message', (event: { data: unknown }) => {\n    if (typeof event.data !== 'string') return\n    lastActivity = Date.now()\n\n    // Inbound frames are client-controlled: a malformed frame (bad JSON, or\n    // valid JSON that isn't an AG-UI RunAgentInput/abort shape) must be\n    // dropped, not crash the socket or leak an unhandled rejection.\n    let frame: InboundFrame\n    try {\n      frame = decodeWsFrame(event.data)\n    } catch (error) {\n      logger.errors('Failed to decode inbound WS frame; dropping it', {\n        error,\n      })\n      return\n    }\n\n    if (frame.kind === 'abort') {\n      const turn = activeTurns.get(frame.runId)\n      if (turn) turn.abort()\n      else earlyAborts.add(frame.runId)\n      return\n    }\n\n    void handleInbound(frame.input)\n  })\n\n  /**\n   * Surface a turn failure to the client as a live `RUN_ERROR` frame. The\n   * socket is conversation-scoped and stays open, so without this frame the\n   * client would see neither a terminal chunk nor a close — a permanent hang.\n   * Mirrors the HTTP transports, which synthesize the live `RUN_ERROR` when\n   * the producer rethrows (see `durableStreamSource`'s terminal contract).\n   */\n  function sendRunError(error: unknown): void {\n    try {\n      socket.send(encodeWsFrame(runErrorChunk(error), undefined))\n    } catch {\n      // Socket is CLOSING/CLOSED — the client sees onclose instead.\n    }\n  }\n\n  async function handleInbound(input: unknown): Promise<void> {\n    let params: Awaited<ReturnType<typeof chatParamsFromRequestBody>>\n    try {\n      params = await chatParamsFromRequestBody(input)\n    } catch (error) {\n      logger.errors('Invalid inbound WS run frame; dropping it', { error })\n      sendRunError(error)\n      return\n    }\n    // The socket may have closed (or errored) during the await above — the\n    // teardown that drains `activeTurns` already ran, so registering now\n    // would start a turn nothing can ever abort.\n    if (closed) return\n    const turnAbort = new AbortController()\n    // A second inbound frame with the same runId (client resubmit) must\n    // abort the earlier turn. Otherwise the old controller is overwritten\n    // and close/abort frames can no longer reach it.\n    activeTurns.get(params.runId)?.abort()\n    activeTurns.set(params.runId, turnAbort)\n    if (earlyAborts.delete(params.runId)) turnAbort.abort()\n    const ctx: WsRunContext = {\n      messages: params.messages,\n      threadId: params.threadId,\n      runId: params.runId,\n      forwardedProps: params.forwardedProps,\n      request: buildTurnRequest(request, params.runId),\n      signal: turnAbort.signal,\n    }\n    try {\n      if (init.durability) {\n        const adapter = init.durability(ctx)\n        const { source, getId } = durableStreamSource(\n          init.onRun(ctx),\n          adapter,\n          {\n            abortController: turnAbort,\n            ...(init.batch === undefined ? {} : { batch: init.batch }),\n            logger,\n          },\n        )\n        for await (const chunk of source) {\n          socket.send(encodeWsFrame(chunk, getId(chunk)))\n        }\n      } else {\n        for await (const chunk of init.onRun(ctx)) {\n          socket.send(encodeWsFrame(chunk, undefined))\n        }\n      }\n    } catch (error) {\n      // An aborted turn (socket close, abort frame, same-runId resubmit) is\n      // expected teardown, not a turn failure — nothing to report.\n      if (!turnAbort.signal.aborted) {\n        logger.errors('WS turn failed', { error })\n        sendRunError(error)\n      }\n    } finally {\n      // Only delete if this turn still owns the entry: a duplicate in-flight\n      // runId (e.g. a client resubmitting before the first turn finished)\n      // would otherwise let the OLDER turn's cleanup delete the NEWER turn's\n      // still-active controller (TOCTOU).\n      if (activeTurns.get(params.runId) === turnAbort) {\n        activeTurns.delete(params.runId)\n      }\n    }\n  }\n}\n\n/**\n * A resume is served entirely from the durability log, so there is no\n * producer to iterate. This empty source satisfies `durableStreamSource`'s\n * signature; on a resume it replays from the log and never touches this.\n * Mirrors the private helper of the same name in `stream-to-response.ts`.\n */\nfunction emptyDurableSource(): AsyncIterable<StreamChunk> {\n  return (async function* () {})()\n}\n\n/**\n * Read-only replay of a run's durability log over a socket (mirrors\n * `resumeServerSentEventsResponse`). The adapter captures the offset from the\n * request (`?offset`/`Last-Event-ID`); no model runs. Closes 1008 when there\n * is nothing to resume.\n */\nexport function resumeWebSocketStream<TOffset extends string = string>(\n  socket: WebSocketLike,\n  options: {\n    adapter: StreamDurability<TOffset>\n    batch?: number\n    debug?: DebugOption\n  },\n): void {\n  const logger = resolveDebugOption(options.debug)\n  if (options.adapter.resumeFrom() === null) {\n    socket.close(1008, 'no resume offset')\n    return\n  }\n  const abortController = new AbortController()\n  socket.addEventListener('close', () => abortController.abort())\n  // An `error` with no listener is an uncaught exception on `ws`; abort the\n  // replay so the pump below stops instead of writing to a dead socket.\n  socket.addEventListener('error', () => abortController.abort())\n  const { source, getId } = durableStreamSource(\n    emptyDurableSource(),\n    options.adapter,\n    {\n      abortController,\n      ...(options.batch === undefined ? {} : { batch: options.batch }),\n      logger,\n    },\n  )\n  void (async () => {\n    for await (const chunk of source) {\n      socket.send(encodeWsFrame(chunk, getId(chunk)))\n    }\n    // Source exhausted = the durability log is complete/terminal; nothing more\n    // will arrive on this read-only socket. Close so the client's reconnect\n    // loop sees onclose and terminates (bounded) instead of awaiting a chunk\n    // that never comes. Safe across durability models: a live decoupled\n    // producer (e.g. durableStream) keeps `read` parked until the terminal,\n    // so the source doesn't exhaust until the run truly ends; a completed\n    // in-process log closes immediately.\n    try {\n      socket.close(1000)\n    } catch {\n      // socket already closing/closed — nothing to do\n    }\n  })().catch((error: unknown) => {\n    logger.errors('resume websocket replay failed', { error })\n    try {\n      socket.close(1011, 'resume failed')\n    } catch {\n      // socket already closing/closed — nothing to do\n    }\n  })\n}\n\ninterface WebSocketPairCtor {\n  new (): { 0: unknown; 1: WebSocketLike & { accept?: () => void } }\n}\n\nfunction upgradeOrThrow(helper: string): {\n  client: unknown\n  server: WebSocketLike\n} {\n  const Pair = (globalThis as { WebSocketPair?: WebSocketPairCtor })\n    .WebSocketPair\n  if (!Pair) {\n    throw new Error(\n      `${helper} requires a runtime with WebSocketPair (Cloudflare Workers/Durable Objects). ` +\n        `On other runtimes upgrade the socket yourself and call ${helper.replace('Response', 'Stream')}.`,\n    )\n  }\n  const pair = new Pair()\n  const server = pair[1]\n  server.accept?.()\n  return { client: pair[0], server }\n}\n\nfunction upgradeResponse(client: unknown): Response {\n  return new Response(null, {\n    status: 101,\n    // Cloudflare-specific field; typed loosely to avoid a DOM lib dependency.\n    webSocket: client,\n  } as ResponseInit & { webSocket: unknown })\n}\n\n/**\n * Cloudflare wrapper (Workers/Durable Objects): creates a `WebSocketPair`,\n * accepts the server socket, delegates to {@link toWebSocketStream}, and\n * returns the 101 upgrade `Response` carrying the client socket. Throws when\n * the runtime has no `WebSocketPair` (Node, Deno, Bun) — upgrade the socket\n * yourself and call {@link toWebSocketStream} directly there.\n */\nexport function toWebSocketResponse<TOffset extends string = string>(\n  request: Request,\n  init: WebSocketStreamInit<TOffset>,\n): Response {\n  const { client, server } = upgradeOrThrow('toWebSocketResponse')\n  toWebSocketStream(server, request, init)\n  return upgradeResponse(client)\n}\n\n/**\n * Cloudflare wrapper (Workers/Durable Objects): creates a `WebSocketPair`,\n * accepts the server socket, delegates to {@link resumeWebSocketStream}, and\n * returns the 101 upgrade `Response` carrying the client socket. Throws when\n * the runtime has no `WebSocketPair` (Node, Deno, Bun) — upgrade the socket\n * yourself and call {@link resumeWebSocketStream} directly there.\n *\n * @example\n * ```ts\n * resumeWebSocketResponse({ adapter: memoryStream(request) })\n * ```\n */\nexport function resumeWebSocketResponse<\n  TOffset extends string = string,\n>(options: {\n  adapter: StreamDurability<TOffset>\n  batch?: number\n  debug?: DebugOption\n}): Response {\n  const { client, server } = upgradeOrThrow('resumeWebSocketResponse')\n  resumeWebSocketStream(server, options)\n  return upgradeResponse(client)\n}\n"],"mappings":";;;;;;;;;;AAiCA,SAAgB,cACd,OACA,IACQ;CACR,OAAO,KAAK,UAAU,OAAO,KAAA,IAAY,QAAQ;EAAE;EAAI;CAAM,CAAC;AAChE;;;;;;AAOA,SAAgB,cAAc,MAA4B;CACxD,MAAM,SAAkB,KAAK,MAAM,IAAI;CACvC,IACE,OAAO,WAAW,YAClB,WAAW,QACV,OAA8B,SAAS,WACxC,OAAQ,OAA+B,UAAU,UAEjD,OAAO;EAAE,MAAM;EAAS,OAAQ,OAA6B;CAAM;CAErE,OAAO;EAAE,MAAM;EAAO,OAAO;CAAO;AACtC;;;;;;;;;;;AAwBA,SAAgB,iBAAiB,WAAoB,OAAwB;CAC3E,MAAM,MAAM,IAAI,IAAI,UAAU,GAAG;CACjC,IAAI,aAAa,IAAI,SAAS,KAAK;CACnC,IAAI,aAAa,OAAO,QAAQ;CAChC,OAAO,IAAI,QAAQ,KAAK,EAAE,SAAS,UAAU,QAAQ,CAAC;AACxD;;;;;;;;AA2BA,SAAgB,kBACd,QACA,SACA,MACM;CACN,MAAM,SAAS,mBAAmB,KAAK,KAAK;CAC5C,MAAM,8BAAc,IAAI,IAA6B;CAIrD,MAAM,8BAAc,IAAI,IAAY;CACpC,MAAM,cAAc,KAAK,eAAe;CACxC,MAAM,gBAAgB,KAAK,iBAAiB;CAC5C,IAAI,eAAe,KAAK,IAAI;CAC5B,IAAI,SAAS;CAEb,MAAM,YAAY,kBAAkB;EAClC,IAAI;GACF,OAAO,KAAK,KAAK,UAAU,EAAE,MAAM,OAAO,CAAC,CAAC;EAC9C,QAAQ,CAIR;CACF,GAAG,WAAW;CACd,MAAM,OAAO,kBACL;EAIJ,IAAI,YAAY,SAAS,KAAK,KAAK,IAAI,IAAI,eAAe,eACxD,OAAO,MAAM,KAAM,MAAM;CAE7B,GACA,KAAK,IAAI,eAAe,GAAM,CAChC;CAEA,SAAS,WAAiB;EACxB,SAAS;EACT,KAAK,MAAM,cAAc,YAAY,OAAO,GAAG,WAAW,MAAM;EAChE,YAAY,MAAM;EAClB,cAAc,SAAS;EACvB,cAAc,IAAI;CACpB;CAEA,OAAO,iBAAiB,SAAS,QAAQ;CAIzC,OAAO,iBAAiB,eAAe;EACrC,OAAO,OAAO,uCAAuC;EACrD,SAAS;EACT,IAAI;GACF,OAAO,MAAM,MAAM,cAAc;EACnC,QAAQ,CAER;CACF,CAAC;CAED,OAAO,iBAAiB,YAAY,UAA6B;EAC/D,IAAI,OAAO,MAAM,SAAS,UAAU;EACpC,eAAe,KAAK,IAAI;EAKxB,IAAI;EACJ,IAAI;GACF,QAAQ,cAAc,MAAM,IAAI;EAClC,SAAS,OAAO;GACd,OAAO,OAAO,kDAAkD,EAC9D,MACF,CAAC;GACD;EACF;EAEA,IAAI,MAAM,SAAS,SAAS;GAC1B,MAAM,OAAO,YAAY,IAAI,MAAM,KAAK;GACxC,IAAI,MAAM,KAAK,MAAM;QAChB,YAAY,IAAI,MAAM,KAAK;GAChC;EACF;EAEA,cAAmB,MAAM,KAAK;CAChC,CAAC;;;;;;;;CASD,SAAS,aAAa,OAAsB;EAC1C,IAAI;GACF,OAAO,KAAK,cAAc,cAAc,KAAK,GAAG,KAAA,CAAS,CAAC;EAC5D,QAAQ,CAER;CACF;CAEA,eAAe,cAAc,OAA+B;EAC1D,IAAI;EACJ,IAAI;GACF,SAAS,MAAM,0BAA0B,KAAK;EAChD,SAAS,OAAO;GACd,OAAO,OAAO,6CAA6C,EAAE,MAAM,CAAC;GACpE,aAAa,KAAK;GAClB;EACF;EAIA,IAAI,QAAQ;EACZ,MAAM,YAAY,IAAI,gBAAgB;EAItC,YAAY,IAAI,OAAO,KAAK,CAAC,EAAE,MAAM;EACrC,YAAY,IAAI,OAAO,OAAO,SAAS;EACvC,IAAI,YAAY,OAAO,OAAO,KAAK,GAAG,UAAU,MAAM;EACtD,MAAM,MAAoB;GACxB,UAAU,OAAO;GACjB,UAAU,OAAO;GACjB,OAAO,OAAO;GACd,gBAAgB,OAAO;GACvB,SAAS,iBAAiB,SAAS,OAAO,KAAK;GAC/C,QAAQ,UAAU;EACpB;EACA,IAAI;GACF,IAAI,KAAK,YAAY;IACnB,MAAM,UAAU,KAAK,WAAW,GAAG;IACnC,MAAM,EAAE,QAAQ,UAAU,oBACxB,KAAK,MAAM,GAAG,GACd,SACA;KACE,iBAAiB;KACjB,GAAI,KAAK,UAAU,KAAA,IAAY,CAAC,IAAI,EAAE,OAAO,KAAK,MAAM;KACxD;IACF,CACF;IACA,WAAW,MAAM,SAAS,QACxB,OAAO,KAAK,cAAc,OAAO,MAAM,KAAK,CAAC,CAAC;GAElD,OACE,WAAW,MAAM,SAAS,KAAK,MAAM,GAAG,GACtC,OAAO,KAAK,cAAc,OAAO,KAAA,CAAS,CAAC;EAGjD,SAAS,OAAO;GAGd,IAAI,CAAC,UAAU,OAAO,SAAS;IAC7B,OAAO,OAAO,kBAAkB,EAAE,MAAM,CAAC;IACzC,aAAa,KAAK;GACpB;EACF,UAAU;GAKR,IAAI,YAAY,IAAI,OAAO,KAAK,MAAM,WACpC,YAAY,OAAO,OAAO,KAAK;EAEnC;CACF;AACF;;;;;;;AAQA,SAAS,qBAAiD;CACxD,QAAQ,mBAAmB,CAAC,EAAA,CAAG;AACjC;;;;;;;AAQA,SAAgB,sBACd,QACA,SAKM;CACN,MAAM,SAAS,mBAAmB,QAAQ,KAAK;CAC/C,IAAI,QAAQ,QAAQ,WAAW,MAAM,MAAM;EACzC,OAAO,MAAM,MAAM,kBAAkB;EACrC;CACF;CACA,MAAM,kBAAkB,IAAI,gBAAgB;CAC5C,OAAO,iBAAiB,eAAe,gBAAgB,MAAM,CAAC;CAG9D,OAAO,iBAAiB,eAAe,gBAAgB,MAAM,CAAC;CAC9D,MAAM,EAAE,QAAQ,UAAU,oBACxB,mBAAmB,GACnB,QAAQ,SACR;EACE;EACA,GAAI,QAAQ,UAAU,KAAA,IAAY,CAAC,IAAI,EAAE,OAAO,QAAQ,MAAM;EAC9D;CACF,CACF;CACA,CAAM,YAAY;EAChB,WAAW,MAAM,SAAS,QACxB,OAAO,KAAK,cAAc,OAAO,MAAM,KAAK,CAAC,CAAC;EAShD,IAAI;GACF,OAAO,MAAM,GAAI;EACnB,QAAQ,CAER;CACF,EAAA,CAAG,CAAC,CAAC,OAAO,UAAmB;EAC7B,OAAO,OAAO,kCAAkC,EAAE,MAAM,CAAC;EACzD,IAAI;GACF,OAAO,MAAM,MAAM,eAAe;EACpC,QAAQ,CAER;CACF,CAAC;AACH;AAMA,SAAS,eAAe,QAGtB;CACA,MAAM,OAAQ,WACX;CACH,IAAI,CAAC,MACH,MAAM,IAAI,MACR,GAAG,OAAO,sIACkD,OAAO,QAAQ,YAAY,QAAQ,EAAE,EACnG;CAEF,MAAM,OAAO,IAAI,KAAK;CACtB,MAAM,SAAS,KAAK;CACpB,OAAO,SAAS;CAChB,OAAO;EAAE,QAAQ,KAAK;EAAI;CAAO;AACnC;AAEA,SAAS,gBAAgB,QAA2B;CAClD,OAAO,IAAI,SAAS,MAAM;EACxB,QAAQ;EAER,WAAW;CACb,CAA0C;AAC5C;;;;;;;;AASA,SAAgB,oBACd,SACA,MACU;CACV,MAAM,EAAE,QAAQ,WAAW,eAAe,qBAAqB;CAC/D,kBAAkB,QAAQ,SAAS,IAAI;CACvC,OAAO,gBAAgB,MAAM;AAC/B;;;;;;;;;;;;;AAcA,SAAgB,wBAEd,SAIW;CACX,MAAM,EAAE,QAAQ,WAAW,eAAe,yBAAyB;CACnE,sBAAsB,QAAQ,OAAO;CACrC,OAAO,gBAAgB,MAAM;AAC/B"}

package/dist/esm/types.d.ts

@@ -4,7 +4,7 @@ import { SystemPrompt } from './system-prompts.js';
 import { CapabilityContext } from './activities/chat/middleware/capabilities.js';
 import { InterruptSubmissionError } from './interrupts.js';
 import { ProviderTool } from './tools/provider-tool.js';
-import { CompletionTokensDetails, PromptTokensDetails, ProviderUsageDetails, TokenUsage, UsageCostBreakdown } from '@tanstack/ai-event-client';
+import { BilledUsage, BillingUnit, CompletionTokensDetails, PromptTokensDetails, ProviderUsageDetails, TokenUsage, UsageCostBreakdown } from '@tanstack/ai-event-client';
 import { BaseEvent as AGUIBaseEvent, CustomEvent as AGUICustomEvent, Interrupt as AGUIInterrupt, MessagesSnapshotEvent as AGUIMessagesSnapshotEvent, ReasoningEncryptedValueEvent as AGUIReasoningEncryptedValueEvent, ReasoningEndEvent as AGUIReasoningEndEvent, ReasoningMessageContentEvent as AGUIReasoningMessageContentEvent, ReasoningMessageEndEvent as AGUIReasoningMessageEndEvent, ReasoningMessageStartEvent as AGUIReasoningMessageStartEvent, ReasoningStartEvent as AGUIReasoningStartEvent, ResumeEntry as AGUIResumeEntry, RunErrorEvent as AGUIRunErrorEvent, RunFinishedEvent as AGUIRunFinishedEvent, RunFinishedOutcome as AGUIRunFinishedOutcome, RunStartedEvent as AGUIRunStartedEvent, StateDeltaEvent as AGUIStateDeltaEvent, StateSnapshotEvent as AGUIStateSnapshotEvent, StepFinishedEvent as AGUIStepFinishedEvent, StepStartedEvent as AGUIStepStartedEvent, TextMessageContentEvent as AGUITextMessageContentEvent, TextMessageEndEvent as AGUITextMessageEndEvent, TextMessageStartEvent as AGUITextMessageStartEvent, ToolCallArgsEvent as AGUIToolCallArgsEvent, ToolCallEndEvent as AGUIToolCallEndEvent, ToolCallResultEvent as AGUIToolCallResultEvent, ToolCallStartEvent as AGUIToolCallStartEvent, EventType } from '@ag-ui/core';
 export type { ProviderTool } from './tools/provider-tool.js';
 /**
@@ -891,7 +891,7 @@ export interface RunStartedEvent extends AGUIRunStartedEvent {
     /** Model identifier for multi-model support */
     model?: string;
 }
-export type { CompletionTokensDetails, PromptTokensDetails, ProviderUsageDetails, TokenUsage, UsageCostBreakdown, };
+export type { BilledUsage, BillingUnit, CompletionTokensDetails, PromptTokensDetails, ProviderUsageDetails, TokenUsage, UsageCostBreakdown, };
 /**
  * @deprecated Renamed to {@link TokenUsage}. Kept as an alias for backward
  * compatibility with `@tanstack/ai@0.23` and earlier; will be removed in a
@@ -1734,9 +1734,10 @@ export interface RerankResult<TDocument = string> {
     rerankedDocuments: Array<TDocument>;
     /**
      * Usage for the request. Rerank typically bills in provider-defined "search
-     * units" (`usage.unitsBilled`) rather than tokens. Some providers (e.g.
-     * OpenRouter) may also report `totalTokens` and `cost`; Cohere reports only
-     * search units and leaves the token counts at 0.
+     * units" (`usage.billed = { quantity, unit: 'units' }`) rather than tokens.
+     * Some providers (e.g. OpenRouter) may also report `totalTokens` and `cost`.
+     * Cohere reports only search units and leaves the token counts at 0.
+     * The deprecated `unitsBilled` field is still populated for compatibility.
      */
     usage: TokenUsage;
 }
@@ -2054,8 +2055,8 @@ export interface VideoUrlResult {
     expiresAt?: Date;
     /**
      * Usage information for the completed generation, when the adapter can report
-     * it. For usage-based providers (e.g. fal) this carries `unitsBilled` — the
-     * real billed quantity — so consumers can compute exact cost.
+     * it. For usage-based providers (e.g. fal) this carries `billed` — the real
+     * billed quantity paired with its unit — so consumers can compute exact cost.
      */
     usage?: TokenUsage;
     /** Persisted artifact references for generated assets, when available */

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@tanstack/ai",
-  "version": "0.45.0",
+  "version": "0.46.0",
   "description": "Type-safe TypeScript AI SDK for streaming chat, tool calling, agents, structured outputs, and multimodal generation.",
   "author": "Tanner Linsley",
   "license": "MIT",
@@ -80,7 +80,7 @@
     "@ag-ui/core": "0.1.1-canary.beta.0",
     "@standard-schema/spec": "^1.1.0",
     "partial-json": "^0.1.7",
-    "@tanstack/ai-event-client": "^0.8.0",
+    "@tanstack/ai-event-client": "^0.9.0",
     "@tanstack/ai-utils": "^0.4.0"
   },
   "peerDependencies": {

package/skills/ai-core/media-generation/SKILL.md

@@ -561,7 +561,7 @@ durations 4/8/12s, single `input_reference` image prompt part), `grokVideo(...)`
 outputs inherit the source clip's properties, so `size`/`aspect_ratio`/`resolution`
 throw in both modes and `duration` throws in edit mode — pass none of them there;
 generation uses the aspect-ratio size template like `'16:9_720p'` (1080p is 1.5-only),
-integer durations 1-15s, reports `usage.unitsBilled` seconds and exact `usage.cost`), `byteplusVideo(...)` (Seedance —
+integer durations 1-15s, reports `usage.billed` seconds ({ quantity, unit: 'seconds' }) and exact `usage.cost`), `byteplusVideo(...)` (Seedance —
 aspect-ratio size template like `'16:9_720p'`, durations 4-15s on the 2.0 family,
 4-12s on 1.5-pro, 2-12s on the 1.0-pro models; reads `ARK_API_KEY`),
 `openRouterVideo(...)` (OpenRouter's dedicated `POST /api/v1/videos` gateway),
@@ -620,10 +620,11 @@ const { generate, result, jobId, videoStatus, isLoading } = useGenerateVideo({
 
 fal bills media generation by usage-based units, not tokens. Every fal media
 adapter (`falImage`, `falAudio`, `falSpeech`, `falTranscription`, `falVideo`)
-surfaces the real billed quantity on the result as `usage.unitsBilled`, read
-from fal's `x-fal-billable-units` response header — no `fetch` interceptor
-needed. It rides on the canonical `TokenUsage` shape (token fields are `0` for
-media), mirroring how duration-billed transcription surfaces `durationSeconds`.
+surfaces the real billed quantity on the result as `usage.billed`
+({ quantity, unit: 'units' }), read from fal's `x-fal-billable-units` response
+header — no `fetch` interceptor needed. It rides on the canonical `TokenUsage`
+shape (token fields are `0` for media), mirroring how duration-billed
+transcription reports { quantity, unit: 'seconds' }.
 
 ```typescript
 import { generateImage } from '@tanstack/ai'
@@ -634,10 +635,10 @@ const result = await generateImage({
   prompt: 'a serene mountain lake',
 })
 
-// usage.unitsBilled is the priced quantity. Multiply by the endpoint unit
+// usage.billed.quantity is the priced quantity. Multiply by the endpoint unit
 // price (GET https://api.fal.ai/v1/models/pricing?endpoint_id=…) for exact cost.
-if (result.usage?.unitsBilled != null) {
-  const cost = result.usage.unitsBilled * unitPrice
+if (result.usage?.billed) {
+  const cost = result.usage.billed.quantity * unitPrice
 }
 ```