@tanstack/ai-client 0.26.0 → 0.28.0

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>
 }
 
 /**
@@ -862,6 +875,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 } : {}),
     }
   }