@tanstack/ai-client 0.25.2 → 0.28.0

package/src/connection-adapters.ts

@@ -1,4 +1,12 @@
-import { EventType, uiMessagesToWire } from '@tanstack/ai/client'
+import {
+  EventType,
+  getChunkRunId as getNormalizedChunkRunId,
+  restoreInboundChunk,
+  tanstackMetadata,
+  uiMessagesToWire,
+  withTanstackMetadata,
+} from '@tanstack/ai/client'
+import { ByokMissingError, isByokMissingBody } from '@tanstack/ai/byok'
 import {
   createResponseStreamTextDecoder,
   getResponseStreamReader,
@@ -8,7 +16,6 @@ import type {
   ModelMessage,
   RunAgentResumeItem,
   RunErrorEvent,
-  RunFinishedEvent,
   StreamChunk,
   UIMessage,
 } from '@tanstack/ai/client'
@@ -34,12 +41,7 @@ export function getChunkRunId(chunk: StreamChunk): string | undefined {
   // the client's run identity to win when a provider stamps its own id; for
   // resumable reconnect/join the two ids match, so precedence is moot there.
   const requestRunId = chunkRunIds.get(chunk)
-  return (
-    requestRunId ??
-    ('runId' in chunk && typeof chunk.runId === 'string'
-      ? chunk.runId
-      : undefined)
-  )
+  return requestRunId ?? getNormalizedChunkRunId(chunk)
 }
 
 /**
@@ -385,6 +387,22 @@ function isNdjsonEnvelope(
   )
 }
 
+/** Rebuild pre-wire extras after SSE/NDJSON ingest. */
+function restoreInboundUsage(chunk: StreamChunk): StreamChunk {
+  return restoreInboundChunk(chunk)
+}
+
+function sseChunkModel(chunk: StreamChunk): string | undefined {
+  const tanstackModel = tanstackMetadata(chunk)?.model
+  if (typeof tanstackModel === 'string') return tanstackModel
+  const usage = 'usage' in chunk ? chunk.usage : undefined
+  if (Array.isArray(usage)) {
+    const model = (usage[0] as { model?: unknown } | undefined)?.model
+    if (typeof model === 'string') return model
+  }
+  return undefined
+}
+
 /**
  * Parse SSE-format lines into stream events, pairing each chunk with the `id:`
  * offset of the event it arrived on. Shared by the fetch- and XHR-backed SSE
@@ -431,27 +449,31 @@ async function* linesToSSEEvents(
     }
     const data = parseSseDataLine(line)
     if (data === '[DONE]') {
-      const synthetic: RunFinishedEvent = {
-        type: EventType.RUN_FINISHED,
-        threadId: lastThreadId ?? fallbackIds?.threadId ?? '',
-        runId: lastRunId ?? fallbackIds?.runId ?? '',
-        model: lastModel ?? '',
-        timestamp: Date.now(),
-        finishReason: 'stop',
+      yield {
+        chunk: withTanstackMetadata(
+          {
+            type: EventType.RUN_FINISHED,
+            threadId: lastThreadId ?? fallbackIds?.threadId ?? '',
+            runId: lastRunId ?? fallbackIds?.runId ?? '',
+            timestamp: Date.now(),
+          },
+          {
+            finishReason: 'stop',
+            ...(lastModel !== undefined ? { model: lastModel } : {}),
+          },
+        ) as StreamChunk,
       }
-      yield { chunk: synthetic }
       return
     }
-    const chunk = JSON.parse(data) as StreamChunk
+    const chunk = restoreInboundUsage(JSON.parse(data) as StreamChunk)
     if ('threadId' in chunk && typeof chunk.threadId === 'string') {
       lastThreadId = chunk.threadId
     }
     if ('runId' in chunk && typeof chunk.runId === 'string') {
       lastRunId = chunk.runId
     }
-    if ('model' in chunk && typeof chunk.model === 'string') {
-      lastModel = chunk.model
-    }
+    const model = sseChunkModel(chunk)
+    if (model !== undefined) lastModel = model
     const id = pendingId
     pendingId = undefined
     yield { chunk, ...(id !== undefined ? { id } : {}) }
@@ -470,19 +492,42 @@ async function* linesToNdjsonEvents(
   for await (const line of lines) {
     const parsed = JSON.parse(line) as unknown
     if (isNdjsonEnvelope(parsed)) {
-      yield { chunk: parsed.chunk, id: parsed.id }
+      yield { chunk: restoreInboundUsage(parsed.chunk), id: parsed.id }
     } else {
-      yield { chunk: parsed as StreamChunk }
+      yield { chunk: restoreInboundUsage(parsed as StreamChunk) }
     }
   }
 }
 
-function assertResponseOk(response: Response): void {
-  if (!response.ok) {
-    throw new Error(
-      `HTTP error! status: ${response.status} ${response.statusText}`,
-    )
+async function assertResponseOk(response: Response): Promise<void> {
+  if (response.ok) return
+  if (response.status === 401) {
+    const body: unknown = await response
+      .clone()
+      .json()
+      .catch(() => null)
+    if (isByokMissingBody(body)) {
+      throw new ByokMissingError(body.error.provider)
+    }
+  }
+  throw new Error(
+    `HTTP error! status: ${response.status} ${response.statusText}`,
+  )
+}
+
+function errorFromXhrStatus(xhr: XMLHttpRequest): Error {
+  if (xhr.status === 401) {
+    let parsed: unknown = null
+    try {
+      parsed = JSON.parse(xhr.responseText)
+    } catch {
+      parsed = null
+    }
+    if (isByokMissingBody(parsed)) {
+      return new ByokMissingError(parsed.error.provider)
+    }
   }
+  return new Error(`XHR error! status: ${xhr.status} ${xhr.statusText}`)
 }
 
 /**
@@ -504,7 +549,7 @@ async function fetchThreadHydration(
     headers: { Accept: 'application/json', ...headers },
     credentials,
   })
-  assertResponseOk(response)
+  await assertResponseOk(response)
   const data = (await response.json()) as {
     messages?: Array<UIMessage>
     activeRun?: { runId?: unknown } | null
@@ -553,7 +598,7 @@ async function fetchGenerationHydration(
     headers: { Accept: 'application/json', ...headers },
     credentials,
   })
-  assertResponseOk(response)
+  await assertResponseOk(response)
   const raw: unknown = await response.json()
   // A 200 carrying `null` is a legitimate hydration miss — the server has no
   // record for this thread — and reading `.activeRun` off `null` would throw.
@@ -588,7 +633,7 @@ async function* responseToSSEEvents(
   abortSignal?: AbortSignal,
   fallbackIds?: { threadId?: string; runId?: string },
 ): AsyncGenerator<StreamEvent> {
-  assertResponseOk(response)
+  await assertResponseOk(response)
   const reader = getResponseStreamReader(response)
   yield* linesToSSEEvents(readStreamLines(reader, abortSignal), fallbackIds)
 }
@@ -598,7 +643,7 @@ async function* responseToNdjsonEvents(
   response: Response,
   abortSignal?: AbortSignal,
 ): AsyncGenerator<StreamEvent> {
-  assertResponseOk(response)
+  await assertResponseOk(response)
   const reader = getResponseStreamReader(response)
   yield* linesToNdjsonEvents(readStreamLines(reader, abortSignal))
 }
@@ -788,6 +833,8 @@ export interface RunAgentInputContext {
   }>
   /** Arbitrary user-controlled passthrough data. */
   forwardedProps?: Record<string, unknown>
+  /** Extra request headers for this run (e.g. BYOK keys). POST only. */
+  headers?: Record<string, string>
 }
 
 export interface ConnectConnectionAdapter {
@@ -1064,21 +1111,24 @@ export function normalizeConnectionAdapter(
         // observed, but stamp the caller's request runId so getChunkRunId()
         // correlates to activeRunIds / currentRunId (same as real stream chunks).
         if (!abortSignal?.aborted && !hasTerminalEvent) {
-          const synthetic: RunFinishedEvent = {
-            type: EventType.RUN_FINISHED,
-            threadId: requireSyntheticId(
-              upstreamThreadId ?? runContext?.threadId,
-              'threadId',
-            ),
-            runId: requireSyntheticId(
-              upstreamRunId ?? runContext?.runId,
-              'runId',
-            ),
-            model: 'connect-wrapper',
-            timestamp: Date.now(),
-            finishReason: 'stop',
-          }
-          push(synthetic, runContext?.runId)
+          push(
+            withTanstackMetadata(
+              {
+                type: EventType.RUN_FINISHED,
+                threadId: requireSyntheticId(
+                  upstreamThreadId ?? runContext?.threadId,
+                  'threadId',
+                ),
+                runId: requireSyntheticId(
+                  upstreamRunId ?? runContext?.runId,
+                  'runId',
+                ),
+                timestamp: Date.now(),
+              },
+              { finishReason: 'stop', model: 'connect-wrapper' },
+            ) as StreamChunk,
+            runContext?.runId,
+          )
         }
       } catch (err) {
         if (!abortSignal?.aborted && !hasTerminalEvent) {
@@ -1248,6 +1298,7 @@ export function fetchServerSentEvents(
       const requestHeaders: Record<string, string> = {
         'Content-Type': 'application/json',
         ...mergeHeaders(resolvedOptions.headers),
+        ...mergeHeaders(runContext?.headers),
         ...runIdHeader(runContext?.runId),
       }
 
@@ -1419,6 +1470,7 @@ export function fetchHttpStream(
       const requestHeaders: Record<string, string> = {
         'Content-Type': 'application/json',
         ...mergeHeaders(resolvedOptions.headers),
+        ...mergeHeaders(runContext?.headers),
         ...runIdHeader(runContext?.runId),
       }
 
@@ -1572,7 +1624,7 @@ function readXhrLines(
 
   const enqueueDelta = () => {
     if (xhr.status !== 0 && (xhr.status < 200 || xhr.status >= 300)) {
-      error = new Error(`XHR error! status: ${xhr.status} ${xhr.statusText}`)
+      error = errorFromXhrStatus(xhr)
       done = true
       return
     }
@@ -1601,7 +1653,7 @@ function readXhrLines(
     // is an error, but status 0 here is not — treat the trailing buffer as a
     // truncation check instead of fabricating a bogus "status: 0" error.
     if (xhr.status !== 0 && (xhr.status < 200 || xhr.status >= 300)) {
-      error = new Error(`XHR error! status: ${xhr.status} ${xhr.statusText}`)
+      error = errorFromXhrStatus(xhr)
     } else if (buffer.trim() && !aborted) {
       error = new StreamTruncatedError()
     }
@@ -1697,6 +1749,7 @@ function createConfiguredXhrRequest(
   const requestHeaders: Record<string, string> = {
     'Content-Type': 'application/json',
     ...mergeHeaders(options.headers),
+    ...mergeHeaders(method === 'POST' ? runContext?.headers : undefined),
     // Client-chosen run id for durability (POST only; the GET join carries it
     // in the query instead).
     ...(method === 'POST' ? runIdHeader(runContext?.runId) : {}),
@@ -2143,9 +2196,9 @@ export function webSocket(
       }
       if (isPingFrame(parsed)) return
       const envelopeId = isNdjsonEnvelope(parsed) ? parsed.id : undefined
-      const chunk = isNdjsonEnvelope(parsed)
-        ? parsed.chunk
-        : (parsed as StreamChunk)
+      const chunk = restoreInboundUsage(
+        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
@@ -2325,7 +2378,9 @@ export function webSocket(
         }
         if (isPingFrame(parsed)) return
         pipe.push(
-          isNdjsonEnvelope(parsed) ? parsed.chunk : (parsed as StreamChunk),
+          restoreInboundUsage(
+            isNdjsonEnvelope(parsed) ? parsed.chunk : (parsed as StreamChunk),
+          ),
         )
       }
       ws.onclose = (event?: CloseEvent) => {
@@ -2468,7 +2523,7 @@ export function fetcherToConnectionAdapter(
             ? { resume: runContext.resume }
             : {}),
         },
-        { signal: abortSignal },
+        { signal: abortSignal, headers: runContext.headers },
       )
       if (result instanceof Response) {
         yield* responseToSSEChunks(result, abortSignal)

package/src/generation-client.ts

@@ -1,3 +1,8 @@
+import { ByokBlockedError, ByokMissingError } from '@tanstack/ai/byok'
+import {
+  prepareResolvedByokHeaders,
+  resolveByokProviderId,
+} from './byok/resolve'
 import {
   GENERATION_EVENTS,
   GENERATION_STREAM_TRUNCATED_MESSAGE,
@@ -10,7 +15,9 @@ import {
 } from './generation-types'
 import { createNoOpGenerationDevtoolsBridge } from './devtools-noop'
 import { parseSSEResponse } from './sse-parser'
+import { restoreInboundChunk } from '@tanstack/ai/client'
 import type { StreamChunk } from '@tanstack/ai/client'
+import type { ByokClient } from './byok'
 import type {
   ConnectConnectionAdapter,
   GenerationHydrationResult,
@@ -119,6 +126,8 @@ export class GenerationClient<
   // the client hydrates the last generation for `threadId` from the server.
   private readonly serverDriven: boolean = false
   private body: Record<string, any>
+  private byok: ByokClient | undefined
+  private byokProvider: (() => string | undefined) | undefined
   private result: TOutput | null = null
   private input: TInput | null = null
   private progress: AIDevtoolsGenerationProgress | null = null
@@ -155,6 +164,8 @@ export class GenerationClient<
     this.hydrateGenerationHandler = options.hydrateGeneration
     this.joinRunHandler = options.joinRun
     this.body = options.body ?? {}
+    this.byok = options.byok
+    this.byokProvider = options.byokProvider
     // `persistence` is `false`/omitted (ephemeral) or `true` (server-driven:
     // hydrate the last generation for `threadId` from the server on mount).
     this.serverDriven = options.persistence === true
@@ -264,9 +275,21 @@ export class GenerationClient<
     const { signal } = abortController
 
     try {
+      let headers: Record<string, string> | undefined
+      if (this.byok) {
+        const provider = resolveByokProviderId(
+          this.byokProvider,
+          this.body.provider,
+        )
+        headers = await prepareResolvedByokHeaders(this.byok, provider)
+      }
+
       if (this.fetcher) {
         // Direct fetch path
-        const result = await this.fetcher(input, { signal })
+        const result = await this.fetcher(
+          input,
+          headers === undefined ? { signal } : { signal, headers },
+        )
         if (signal.aborted) return
         if (result instanceof Response) {
           // Server function returned SSE Response — parse stream
@@ -288,7 +311,7 @@ export class GenerationClient<
           [],
           mergedData,
           signal,
-          this.createRunContext(runId),
+          this.createRunContext(runId, headers),
         )
         await this.processStream(stream, runId, signal)
       } else {
@@ -311,6 +334,12 @@ export class GenerationClient<
     } catch (err: unknown) {
       if (signal.aborted) return
       const error = err instanceof Error ? err : new Error(String(err))
+      if (error instanceof ByokMissingError) {
+        this.byok?.request(error.provider, 'missing')
+      }
+      if (error instanceof ByokBlockedError && error.reason === 'locked') {
+        this.byok?.request(error.provider, 'locked')
+      }
       this.setError(error)
       this.setStatus('error')
       this.recordResumeSnapshotError(error)
@@ -349,9 +378,10 @@ export class GenerationClient<
     let streamRunId: string | undefined
     let sawTerminalChunk = false
 
-    for await (const chunk of source) {
+    for await (const raw of source) {
       if (signal.aborted) break
 
+      const chunk = restoreInboundChunk(raw)
       this.callbacksRef.onChunk?.(chunk)
       this.observeResumeSnapshot(chunk)
       const chunkRunId =
@@ -390,11 +420,9 @@ export class GenerationClient<
           this.devtoolsBridge.ensureRunStarted(
             chunkRunId ?? streamRunId ?? fallbackRunId,
           )
-          // Prefer spec `message`; fall back to deprecated `error.message`
+          // Spec RUN_ERROR message. Missing message uses this fallback.
           const msg =
-            (chunk.message as string | undefined) ||
-            chunk.error?.message ||
-            'An error occurred'
+            (chunk.message as string | undefined) || 'An error occurred'
           throw new Error(msg)
         }
         default:
@@ -462,13 +490,25 @@ export class GenerationClient<
     options: Partial<
       Pick<
         GenerationClientOptions<TInput, TResult, TOutput>,
-        'body' | 'onResult' | 'onError' | 'onProgress' | 'onChunk'
+        | 'body'
+        | 'byok'
+        | 'byokProvider'
+        | 'onResult'
+        | 'onError'
+        | 'onProgress'
+        | 'onChunk'
       >
     >,
   ): void {
     if (options.body !== undefined) {
       this.body = options.body ?? {}
     }
+    if (options.byok !== undefined) {
+      this.byok = options.byok
+    }
+    if (options.byokProvider !== undefined) {
+      this.byokProvider = options.byokProvider
+    }
     if (options.onResult !== undefined) {
       this.callbacksRef.onResult = options.onResult
     }
@@ -645,10 +685,14 @@ export class GenerationClient<
     return `${prefix}-${Date.now()}-${Math.random().toString(36).substring(7)}`
   }
 
-  private createRunContext(runId: string): RunAgentInputContext {
+  private createRunContext(
+    runId: string,
+    headers?: Record<string, string>,
+  ): RunAgentInputContext {
     return {
       threadId: this.threadId,
       runId,
+      ...(headers ? { headers } : {}),
     }
   }
 

package/src/generation-types.ts

@@ -1,9 +1,11 @@
+import { tanstackMetadata } from '@tanstack/ai/client'
 import type {
   MediaPrompt,
   PersistedArtifactRef,
   StreamChunk,
 } from '@tanstack/ai/client'
 import type { TokenUsage, TranscriptionResponseFormat } from '@tanstack/ai'
+import type { ByokClient } from './byok'
 import type { ConnectConnectionAdapter } from './connection-adapters'
 import type { AIDevtoolsClientMetadata } from './devtools'
 import type {
@@ -270,6 +272,8 @@ export const GENERATION_EVENTS = {
 export interface GenerationFetcherOptions {
   /** AbortSignal that is triggered when the user calls `stop()` */
   signal: AbortSignal
+  /** Extra request headers for this run (e.g. BYOK keys). */
+  headers?: Record<string, string>
 }
 
 /**
@@ -337,6 +341,19 @@ export interface GenerationClientOptions<_TInput, TResult, TOutput = TResult> {
   /** Additional body parameters to send with connect-based adapter requests */
   body?: Record<string, any>
 
+  /**
+   * Optional BYOK keyring. On each generate the client prepares the resolved
+   * provider and stamps `x-byok-*` request headers. Keys never go in the body.
+   */
+  byok?: ByokClient
+
+  /**
+   * Optional provider id for this generation. If it returns a provider slug,
+   * only that key is prepared and sent. Otherwise `body.provider` is used.
+   * If no slug resolves, generate throws instead of attaching every stored key.
+   */
+  byokProvider?: () => string | undefined
+
   /** Metadata used to register this generation hook with TanStack AI Devtools */
   devtools?: Partial<AIDevtoolsClientMetadata>
 
@@ -461,8 +478,13 @@ export function updateGenerationResumeSnapshot(
   previous: GenerationResumeSnapshot | null | undefined,
   chunk: StreamChunk,
 ): GenerationResumeSnapshot {
-  const threadId = stringField(chunk, 'threadId')
-  const runId = stringField(chunk, 'runId')
+  const tanstack = tanstackMetadata(chunk)
+  const threadId =
+    stringField(chunk, 'threadId') ??
+    (typeof tanstack?.threadId === 'string' ? tanstack.threadId : undefined)
+  const runId =
+    stringField(chunk, 'runId') ??
+    (typeof tanstack?.runId === 'string' ? tanstack.runId : undefined)
   const carried = chunk.type === 'RUN_STARTED' ? undefined : previous
   const previousArtifacts = carried?.pendingArtifacts ?? []
   const next: GenerationResumeSnapshot = {

package/src/interrupt-manager.ts

@@ -102,6 +102,10 @@ interface TransactionToken {
   active: boolean
 }
 
+interface SubmissionOperation {
+  submission: InterruptManagerSubmission
+}
+
 interface RuntimeInterruptCheckpoint {
   status: InterruptItemStatus
   resolution?: RunAgentResumeItem
@@ -511,6 +515,7 @@ export class InterruptManager<
     resuming: false,
   })
   private activeTransaction: TransactionToken | undefined
+  private activeSubmissionOperation: SubmissionOperation | undefined
   private retrySubmission: InterruptManagerSubmission | undefined
   private resuming = false
   private tools: TTools | undefined
@@ -544,6 +549,7 @@ export class InterruptManager<
     hydration: InterruptManagerHydration,
     source: InterruptManagerChangeSource = 'live',
   ): void {
+    this.activeSubmissionOperation = undefined
     this.hydration = {
       threadId: hydration.threadId,
       interruptedRunId: hydration.interruptedRunId,
@@ -597,6 +603,7 @@ export class InterruptManager<
     preserveRootErrors?: boolean
     source?: InterruptManagerChangeSource
   }): void {
+    this.activeSubmissionOperation = undefined
     this.hydration = undefined
     this.items = []
     this.snapshot = Object.freeze([])
@@ -1528,25 +1535,32 @@ export class InterruptManager<
   }
 
   private submitBatch(submission: InterruptManagerSubmission): void {
+    // Track ownership so a superseded submission cannot mutate current state.
+    const operation = { submission }
+    this.activeSubmissionOperation = operation
     this.resuming = true
     this.retrySubmission = undefined
     for (const item of this.items) {
       if (isClientOwnedInterrupt(item)) item.status = 'submitting'
     }
     this.publish()
-    void this.performSubmission(submission)
+    void this.performSubmission(operation)
   }
 
   private async performSubmission(
-    submission: InterruptManagerSubmission,
+    operation: SubmissionOperation,
   ): Promise<void> {
     try {
-      await this.options.submit(submission)
+      await this.options.submit(operation.submission)
     } catch (error) {
-      this.handleSubmissionFailure(error, submission)
+      if (this.activeSubmissionOperation !== operation) return
+      this.handleSubmissionFailure(error, operation.submission)
     } finally {
-      this.resuming = false
-      this.publish()
+      if (this.activeSubmissionOperation === operation) {
+        this.activeSubmissionOperation = undefined
+        this.resuming = false
+        this.publish()
+      }
     }
   }
 

package/src/sse-parser.ts

@@ -1,3 +1,4 @@
+import { ByokMissingError, isByokMissingBody } from '@tanstack/ai/byok'
 import {
   createResponseStreamTextDecoder,
   getResponseStreamReader,
@@ -51,6 +52,15 @@ export async function* parseSSEResponse(
   abortSignal?: AbortSignal,
 ): AsyncGenerator<StreamChunk> {
   if (!response.ok) {
+    if (response.status === 401) {
+      const body: unknown = await response
+        .clone()
+        .json()
+        .catch(() => null)
+      if (isByokMissingBody(body)) {
+        throw new ByokMissingError(body.error.provider)
+      }
+    }
     throw new Error(
       `HTTP error! status: ${response.status} ${response.statusText}`,
     )