@tanstack/ai 0.45.0 → 0.46.0

package/src/stream-to-websocket.ts

@@ -0,0 +1,418 @@
+import { chatParamsFromRequestBody } from './utilities/chat-params'
+import { durableStreamSource, runErrorChunk } 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, 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 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<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 function buildTurnRequest(handshake: Request, runId: string): Request {
+  const url = new URL(handshake.url)
+  url.searchParams.set('runId', runId)
+  url.searchParams.delete('offset')
+  return new Request(url, { headers: handshake.headers })
+}
+
+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 function toWebSocketStream<TOffset extends string = string>(
+  socket: WebSocketLike,
+  request: Request,
+  init: WebSocketStreamInit<TOffset>,
+): void {
+  const logger = resolveDebugOption(init.debug)
+  const activeTurns = new Map<string, AbortController>()
+  // Abort frames that raced ahead of their run's registration: `handleInbound`
+  // awaits body validation before it registers into `activeTurns`, so an abort
+  // arriving inside that window would otherwise be silently discarded.
+  const earlyAborts = new Set<string>()
+  const heartbeatMs = init.heartbeatMs ?? 30_000
+  const idleTimeoutMs = init.idleTimeoutMs ?? 300_000
+  let lastActivity = Date.now()
+  let closed = false
+
+  const heartbeat = setInterval(() => {
+    try {
+      socket.send(JSON.stringify({ type: 'ping' }))
+    } catch {
+      // Socket is CLOSING/CLOSED between ticks — teardown 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),
+  )
+
+  function teardown(): void {
+    closed = true
+    for (const controller of activeTurns.values()) controller.abort()
+    activeTurns.clear()
+    clearInterval(heartbeat)
+    clearInterval(idle)
+  }
+
+  socket.addEventListener('close', teardown)
+  // Without this, an errored socket whose `close` never follows would leak
+  // both intervals and never abort its turns — and on `ws` (an EventEmitter)
+  // an `error` event with no listener is thrown as an uncaught exception.
+  socket.addEventListener('error', () => {
+    logger.errors('WebSocket errored; aborting its turns')
+    teardown()
+    try {
+      socket.close(1011, 'socket error')
+    } catch {
+      // socket already closing/closed — nothing to do
+    }
+  })
+
+  socket.addEventListener('message', (event: { data: unknown }) => {
+    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') {
+      const turn = activeTurns.get(frame.runId)
+      if (turn) turn.abort()
+      else earlyAborts.add(frame.runId)
+      return
+    }
+
+    void handleInbound(frame.input)
+  })
+
+  /**
+   * Surface a turn failure to the client as a live `RUN_ERROR` frame. The
+   * socket is conversation-scoped and stays open, so without this frame the
+   * client would see neither a terminal chunk nor a close — a permanent hang.
+   * Mirrors the HTTP transports, which synthesize the live `RUN_ERROR` when
+   * the producer rethrows (see `durableStreamSource`'s terminal contract).
+   */
+  function sendRunError(error: unknown): void {
+    try {
+      socket.send(encodeWsFrame(runErrorChunk(error), undefined))
+    } catch {
+      // Socket is CLOSING/CLOSED — the client sees onclose instead.
+    }
+  }
+
+  async function handleInbound(input: unknown): Promise<void> {
+    let params: Awaited<ReturnType<typeof chatParamsFromRequestBody>>
+    try {
+      params = await chatParamsFromRequestBody(input)
+    } catch (error) {
+      logger.errors('Invalid inbound WS run frame; dropping it', { error })
+      sendRunError(error)
+      return
+    }
+    // The socket may have closed (or errored) during the await above — the
+    // teardown that drains `activeTurns` already ran, so registering now
+    // would start a turn nothing can ever abort.
+    if (closed) return
+    const turnAbort = new AbortController()
+    // A second inbound frame with the same runId (client resubmit) must
+    // abort the earlier turn. Otherwise the old controller is overwritten
+    // and close/abort frames can no longer reach it.
+    activeTurns.get(params.runId)?.abort()
+    activeTurns.set(params.runId, turnAbort)
+    if (earlyAborts.delete(params.runId)) turnAbort.abort()
+    const ctx: WsRunContext = {
+      messages: params.messages,
+      threadId: params.threadId,
+      runId: params.runId,
+      forwardedProps: params.forwardedProps,
+      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 === 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))
+        }
+      }
+    } catch (error) {
+      // An aborted turn (socket close, abort frame, same-runId resubmit) is
+      // expected teardown, not a turn failure — nothing to report.
+      if (!turnAbort.signal.aborted) {
+        logger.errors('WS turn failed', { error })
+        sendRunError(error)
+      }
+    } finally {
+      // Only delete if this turn still owns the entry: a duplicate in-flight
+      // runId (e.g. a client resubmitting before the first turn finished)
+      // 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<StreamChunk> {
+  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<TOffset extends string = string>(
+  socket: WebSocketLike,
+  options: {
+    adapter: StreamDurability<TOffset>
+    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())
+  // An `error` with no listener is an uncaught exception on `ws`; abort the
+  // replay so the pump below stops instead of writing to a dead socket.
+  socket.addEventListener('error', () => abortController.abort())
+  const { source, getId } = durableStreamSource(
+    emptyDurableSource(),
+    options.adapter,
+    {
+      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 wrapper (Workers/Durable Objects): creates a `WebSocketPair`,
+ * accepts the server socket, delegates to {@link toWebSocketStream}, and
+ * returns the 101 upgrade `Response` carrying the client socket. Throws when
+ * the runtime has no `WebSocketPair` (Node, Deno, Bun) — upgrade the socket
+ * yourself and call {@link toWebSocketStream} directly there.
+ */
+export function toWebSocketResponse<TOffset extends string = string>(
+  request: Request,
+  init: WebSocketStreamInit<TOffset>,
+): Response {
+  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) })
+ * ```
+ */
+export function resumeWebSocketResponse<
+  TOffset extends string = string,
+>(options: {
+  adapter: StreamDurability<TOffset>
+  batch?: number
+  debug?: DebugOption
+}): Response {
+  const { client, server } = upgradeOrThrow('resumeWebSocketResponse')
+  resumeWebSocketStream(server, options)
+  return upgradeResponse(client)
+}

package/src/types.ts

@@ -11,6 +11,8 @@ import type { ProviderTool } from './tools/provider-tool'
 // package (which `@tanstack/ai` already depends on) so there is a single source
 // of truth without a dependency cycle. They are re-exported below.
 import type {
+  BilledUsage,
+  BillingUnit,
   CompletionTokensDetails,
   PromptTokensDetails,
   ProviderUsageDetails,
@@ -1106,6 +1108,8 @@ export interface RunStartedEvent extends AGUIRunStartedEvent {
 // Re-export the canonical usage types (defined in `@tanstack/ai-event-client`)
 // so `@tanstack/ai` consumers keep importing them from here unchanged.
 export type {
+  BilledUsage,
+  BillingUnit,
   CompletionTokensDetails,
   PromptTokensDetails,
   ProviderUsageDetails,
@@ -2078,9 +2082,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
 }
@@ -2463,8 +2468,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 */