@tanstack/ai-client 0.26.0 → 0.29.0

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,
@@ -12,6 +17,7 @@ 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,
@@ -120,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
@@ -156,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
@@ -265,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
@@ -289,7 +311,7 @@ export class GenerationClient<
           [],
           mergedData,
           signal,
-          this.createRunContext(runId),
+          this.createRunContext(runId, headers),
         )
         await this.processStream(stream, runId, signal)
       } else {
@@ -312,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)
@@ -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

@@ -5,6 +5,7 @@ import type {
   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 {
@@ -271,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>
 }
 
 /**
@@ -338,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>
 

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/message-date-normalizer.ts

@@ -0,0 +1,29 @@
+import type { UIMessage } from './types'
+
+function validDate(value: unknown): Date | undefined {
+  if (value instanceof Date)
+    return Number.isNaN(value.getTime()) ? undefined : value
+  if (typeof value !== 'string') return undefined
+  const date = new Date(value)
+  return Number.isNaN(date.getTime()) ? undefined : date
+}
+
+export function normalizeMessageDates(message: UIMessage): UIMessage {
+  const messageDate = validDate(message.createdAt)
+  const parts = message.parts.map((part) => {
+    if (part.type !== 'tool-result') return part
+    const date = validDate(part.createdAt)
+    const { createdAt: _ignored, ...rest } = part
+    return date ? { ...rest, createdAt: date } : rest
+  })
+  const { createdAt: _ignored, ...rest } = message
+  return messageDate
+    ? { ...rest, parts, createdAt: messageDate }
+    : { ...rest, parts }
+}
+
+export function normalizeMessagesDates(
+  messages: Array<UIMessage>,
+): Array<UIMessage> {
+  return messages.map(normalizeMessageDates)
+}

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}`,
     )

package/src/storage-adapters.ts

@@ -1,4 +1,5 @@
-import type { ChatPersistedState, ChatStorageAdapter } from './types'
+import type { ChatPersistedState, ChatStorageAdapter, UIMessage } from './types'
+import { normalizeMessagesDates } from './message-date-normalizer'
 
 export interface WebStoragePersistenceOptions {
   keyPrefix?: string
@@ -45,6 +46,16 @@ function stringifyJson(value: ChatPersistedState): string {
   return serialized
 }
 
+function reviveMessageCreatedAt(message: UIMessage): UIMessage {
+  return normalizeMessagesDates([message]).at(0) ?? message
+}
+
+function revivePersistedState(state: ChatPersistedState): ChatPersistedState {
+  if (state == null || typeof state !== 'object') return state
+  if (!Array.isArray(state.messages)) return state
+  return { ...state, messages: state.messages.map(reviveMessageCreatedAt) }
+}
+
 function createWebStoragePersistence(
   storageName: 'localStorage' | 'sessionStorage',
   options: WebStoragePersistenceOptions,
@@ -69,7 +80,7 @@ function createWebStoragePersistence(
   return {
     getItem(id) {
       const item = getStorage().getItem(key(id))
-      return item === null ? null : deserialize(item)
+      return item === null ? null : revivePersistedState(deserialize(item))
     },
     setItem(id, value) {
       getStorage().setItem(key(id), serialize(value))

package/src/types.ts

@@ -25,6 +25,7 @@ import type {
   UIResourcePart,
   VideoPart,
 } from '@tanstack/ai/client'
+import type { ByokClient } from './byok'
 import type { ConnectionAdapter } from './connection-adapters'
 import type { AIDevtoolsClientMetadata } from './devtools'
 import type { ChatDevtoolsBridgeFactory } from './devtools-noop'
@@ -291,6 +292,8 @@ export interface ChatFetcherInput {
 export interface ChatFetcherOptions {
   /** Fires when `stop()` is called or the request is superseded. */
   signal: AbortSignal
+  /** Extra request headers for this run (e.g. BYOK keys). */
+  headers?: Record<string, string>
 }
 
 /**
@@ -474,9 +477,9 @@ export interface QueueConfig {
  * function form (no `batch` via function). Per-call `sendOptions.whenBusy`
  * overrides the strategy for that send.
  *
- * Actions match {@link WhenBusy}: `'queue' | 'drop' | 'interrupt'`. Concurrent
- * streams are not supported. `pending.id` is the id that will be stored if the
- * action is `'queue'` (safe to pass to `cancelQueued`).
+ * Actions match {@link WhenBusy}. Concurrent streams are not supported.
+ * `pending.id` is the id that will be stored if the action is `'queue'`
+ * (safe to pass to `cancelQueued`).
  */
 export type QueueStrategy = (ctx: {
   pending: QueuedMessage
@@ -491,6 +494,16 @@ export type QueueOption = WhenBusy | QueueConfig | QueueStrategy
 export interface SendMessageOptions {
   /** Overrides the configured `whenBusy` for this one send. */
   whenBusy?: WhenBusy
+  /**
+   * Extra JSON merged into this request's wire `forwardedProps`.
+   * Shallow merge: `{ ...chatBody, ...positionalBody, ...body }`.
+   * This field wins on key collisions.
+   *
+   * Framework hooks (`useChat`, `injectChat`, `createChat`) expose
+   * `sendMessage(content, options)` with no positional body, so this field
+   * is the per-call body channel on those surfaces.
+   */
+  body?: Record<string, any>
 }
 
 /**
@@ -582,10 +595,14 @@ export type ToolCallPart<TTools extends ReadonlyArray<AnyClientTool> = any> =
 
 export interface ToolResultPart {
   type: 'tool-result'
+  id?: string
+  name?: string
   toolCallId: string
   content: string | Array<ContentPart>
   state: ToolResultState
   error?: string // Error message if state is "error"
+  metadata?: Record<string, unknown>
+  createdAt?: Date
 }
 
 export interface ThinkingPart {
@@ -626,6 +643,7 @@ export interface UIMessage<
 > {
   id: string
   role: 'system' | 'user' | 'assistant'
+  name?: string
   parts: Array<MessagePart<TTools, TData>>
   createdAt?: Date
   /**
@@ -862,6 +880,21 @@ export interface ChatClientBaseOptions<
    */
   body?: Record<string, any>
 
+  /**
+   * Optional BYOK keyring. On each send 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 chat. If it returns a provider slug,
+   * only that key is prepared and sent. Otherwise the merged `provider`
+   * from `forwardedProps`, `body`, and per-call `sendMessage` `body` is
+   * used. Later sources win. If no slug resolves, the send throws
+   * instead of attaching every stored key.
+   */
+  byokProvider?: () => string | undefined
+
   /**
    * Client-local runtime context passed to client tool implementations.
    *

package/src/video-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,
@@ -12,6 +17,7 @@ import { createNoOpVideoDevtoolsBridge } 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,
@@ -122,6 +128,8 @@ export class VideoGenerationClient<TOutput = VideoGenerateResult> {
   // 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: VideoGenerateInput | null = null
@@ -156,6 +164,8 @@ export class VideoGenerationClient<TOutput = VideoGenerateResult> {
     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
@@ -258,15 +268,24 @@ export class VideoGenerationClient<TOutput = VideoGenerateResult> {
     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) {
-        await this.generateWithFetcher(input, signal, runId)
+        await this.generateWithFetcher(input, signal, runId, headers)
       } else if (this.connection) {
         const mergedData = { ...this.body, ...input }
         const stream = this.connection.connect(
           [],
           mergedData,
           signal,
-          this.createRunContext(runId),
+          this.createRunContext(runId, headers),
         )
         await this.processStream(stream, runId, signal)
       } else {
@@ -284,6 +303,12 @@ export class VideoGenerationClient<TOutput = VideoGenerateResult> {
     } 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)
@@ -309,11 +334,15 @@ export class VideoGenerationClient<TOutput = VideoGenerateResult> {
     input: VideoGenerateInput,
     signal: AbortSignal,
     runId: string,
+    headers?: Record<string, string>,
   ): Promise<void> {
     if (!this.fetcher) return
 
     // Fetcher returns a completed result directly, or a Response with SSE body
-    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) {
@@ -470,6 +499,8 @@ export class VideoGenerationClient<TOutput = VideoGenerateResult> {
       Pick<
         VideoGenerationClientOptions<TOutput>,
         | 'body'
+        | 'byok'
+        | 'byokProvider'
         | 'onResult'
         | 'onError'
         | 'onProgress'
@@ -482,6 +513,12 @@ export class VideoGenerationClient<TOutput = VideoGenerateResult> {
     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
     }
@@ -694,10 +731,14 @@ export class VideoGenerationClient<TOutput = VideoGenerateResult> {
     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 } : {}),
     }
   }