@tanstack/ai-client 0.25.2 → 0.28.0

package/src/chat-client.ts

@@ -3,9 +3,21 @@ import {
   convertSchemaToJsonSchema,
   generateMessageId,
   isStandardSchema,
+  mergeMetadata,
   normalizeToUIMessage,
   parseWithStandardSchema,
+  restoreInboundChunk,
+  tanstackMetadata,
 } from '@tanstack/ai/client'
+import {
+  ByokBlockedError,
+  ByokMissingError,
+  ByokUnresolvedProviderError,
+} from '@tanstack/ai/byok'
+import {
+  prepareResolvedByokHeaders,
+  resolveByokProviderId,
+} from './byok/resolve'
 import { createNoOpChatDevtoolsBridge } from './devtools-noop'
 import {
   fetcherToConnectionAdapter,
@@ -24,6 +36,7 @@ import type {
   RunAgentResumeItem,
   StreamChunk,
 } from '@tanstack/ai/client'
+import type { ByokClient } from './byok'
 import type {
   ChatHydrationResult,
   ConnectionAdapter,
@@ -94,6 +107,8 @@ type ChatClientUpdateOptionsWithoutContext<
   /** @deprecated Use `forwardedProps` instead. */
   body?: Record<string, any>
   forwardedProps?: Record<string, any>
+  byok?: ByokClient
+  byokProvider?: () => string | undefined
   tools?: TTools
   interrupts?: TInterrupts
   queue?: QueueOption
@@ -189,6 +204,7 @@ export function normalizeQueueOption(
  * Merge a run of queued messages into a single send for `drain: 'batch'`.
  * All-string content is joined with newlines; mixed/multimodal content is
  * flattened into a single `ContentPart` array. The last item's `body` wins.
+ * Object-form metadata is merged last-write-wins per key.
  */
 function mergeQueuedMessages(items: Array<InternalQueuedMessage>): {
   content: string | MultimodalContent
@@ -209,17 +225,24 @@ function mergeQueuedMessages(items: Array<InternalQueuedMessage>): {
     }
   }
   const parts: Array<ContentPart> = []
+  let metadata: Record<string, any> | undefined
   for (const item of items) {
     if (typeof item.content === 'string') {
       parts.push({ type: 'text', content: item.content })
-    } else if (typeof item.content.content === 'string') {
+      continue
+    }
+    if (typeof item.content.content === 'string') {
       parts.push({ type: 'text', content: item.content.content })
     } else {
       parts.push(...item.content.content)
     }
+    metadata = mergeMetadata(metadata, item.content.metadata)
   }
   return {
-    content: { content: parts },
+    content: {
+      content: parts,
+      ...(metadata !== undefined ? { metadata } : {}),
+    },
     ...(body !== undefined ? { body } : {}),
   }
 }
@@ -337,6 +360,8 @@ export class ChatClient<
   // merged on every send, with `forwardedProps` winning on key collision.
   private bodyOption: Record<string, any> = {}
   private forwardedPropsOption: Record<string, any> = {}
+  private byok: ByokClient | undefined
+  private byokProvider: (() => string | undefined) | undefined
   private context: TContext | undefined = undefined
   private pendingMessageBody: Record<string, any> | undefined = undefined
   private queueConfig: NormalizedQueueConfig
@@ -483,6 +508,8 @@ export class ChatClient<
     // winning on key collision.
     this.bodyOption = options.body || {}
     this.forwardedPropsOption = options.forwardedProps || {}
+    this.byok = options.byok
+    this.byokProvider = options.byokProvider
     this.context = options.context
     this.queueConfig = normalizeQueueOption(options.queue)
     this.connection = normalizeConnectionAdapter(resolveTransport(options))
@@ -1747,12 +1774,14 @@ export class ChatClient<
     chunk: StreamChunk,
     options?: { defer?: boolean },
   ): Promise<void> {
+    chunk = restoreInboundChunk(chunk)
     if (
       chunk.type === 'RUN_ERROR' &&
       this.isActiveInterruptSubmissionFailure(chunk)
     ) {
+      const interruptErrors = tanstackMetadata(chunk)?.interruptErrors
       this.interruptSubmissionFailure = {
-        errors: chunk['tanstack:interruptErrors'] ?? [],
+        errors: Array.isArray(interruptErrors) ? interruptErrors : [],
       }
     }
     if (this.connectionStatus === 'connecting') {
@@ -1792,8 +1821,10 @@ export class ChatClient<
     chunk: Extract<StreamChunk, { type: 'RUN_ERROR' }>,
   ): boolean {
     const submission = this.activeInterruptSubmission
-    const errors = chunk['tanstack:interruptErrors']
-    if (!submission || !errors || errors.length === 0) return false
+    const errors = tanstackMetadata(chunk)?.interruptErrors
+    if (!submission || !Array.isArray(errors) || errors.length === 0) {
+      return false
+    }
     const runId = getChunkRunId(chunk)
     if (runId !== undefined && runId !== this.currentRunId) return false
     if (
@@ -1802,12 +1833,22 @@ export class ChatClient<
     ) {
       return false
     }
-    return errors.every(
-      (error) =>
+    return errors.every((error) => {
+      if (
+        error == null ||
+        typeof error !== 'object' ||
+        typeof error.threadId !== 'string' ||
+        typeof error.interruptedRunId !== 'string' ||
+        typeof error.generation !== 'number'
+      ) {
+        return false
+      }
+      return (
         error.threadId === submission.threadId &&
         error.interruptedRunId === submission.interruptedRunId &&
-        error.generation === submission.generation,
-    )
+        error.generation === submission.generation
+      )
+    })
   }
 
   private resolveJoinedRun(chunk: StreamChunk): void {
@@ -1857,8 +1898,12 @@ export class ChatClient<
    *   - A MultimodalContent object with content array and optional custom ID
    * @param body - Optional body parameters to merge with the client's base body for this request.
    *               Uses shallow merge with per-message body taking priority.
-   * @param sendOptions - Per-call overrides, e.g. `{ whenBusy: 'interrupt' }` to
-   *                      override the configured queue policy for this one send.
+   * @param sendOptions - Per-call overrides. `{ whenBusy }` overrides the
+   *                      queue policy for this one send. `{ body }`
+   *                      shallow-merges with `body` and with the chat-level
+   *                      `body` / `forwardedProps`. `sendOptions.body` wins
+   *                      on key collisions. Framework hooks forward this
+   *                      object as their second argument.
    *
    * @example
    * ```ts
@@ -1868,9 +1913,13 @@ export class ChatClient<
    * // Text message with custom body params
    * await client.sendMessage('Hello!', { temperature: 0.7 })
    *
-   * // Per-call whenBusy override (body must still be the 2nd arg on ChatClient)
+   * // Per-call whenBusy override
    * await client.sendMessage('Urgent', undefined, { whenBusy: 'interrupt' })
    *
+   * // Per-call body via options. Same effect as the positional arg.
+   * // This is the shape the framework hooks (`useChat`, `injectChat`) forward.
+   * await client.sendMessage('Hello!', undefined, { body: { temperature: 0.7 } })
+   *
    * // Multimodal message with image
    * await client.sendMessage({
    *   content: [
@@ -1908,13 +1957,15 @@ export class ChatClient<
       )
     }
 
+    const resolvedBody = { ...body, ...sendOptions?.body }
+
     if (this.isSendBusy()) {
       const { action, id } = this.decideWhenBusy(content, sendOptions)
       if (action === 'drop') {
         return
       }
       if (action === 'queue') {
-        this.enqueueMessage(content, body, id)
+        this.enqueueMessage(content, resolvedBody, id)
         return
       }
       // 'interrupt': abort the current stream, then send now.
@@ -1931,7 +1982,7 @@ export class ChatClient<
     }
 
     try {
-      await this.deliverMessage(content, body)
+      await this.deliverMessage(content, resolvedBody)
     } finally {
       this.sendInFlight = false
     }
@@ -1982,6 +2033,7 @@ export class ChatClient<
       const userMessage = this.processor.addUserMessage(
         normalizedContent.content,
         normalizedContent.id,
+        normalizedContent.metadata,
       )
       this.events.messageSent(userMessage.id, normalizedContent.content)
       return await this.streamResponse()
@@ -2041,17 +2093,22 @@ export class ChatClient<
   }
 
   /**
-   * Normalize the message input to extract content and optional id.
-   * Trims string content automatically.
+   * Normalize the message input to extract content, optional id, and
+   * optional metadata. String form has no metadata. Trims string content.
    */
   private normalizeMessageInput(input: string | MultimodalContent): {
     content: string | Array<ContentPart>
     id?: string
+    metadata?: Record<string, any>
   } {
     if (typeof input === 'string') {
       return { content: input.trim() }
     }
-    return { content: input.content, id: input.id }
+    return {
+      content: input.content,
+      id: input.id,
+      ...(input.metadata !== undefined ? { metadata: input.metadata } : {}),
+    }
   }
 
   /**
@@ -2160,7 +2217,7 @@ export class ChatClient<
       // order (later spreads win):
       //   1. Legacy `body` option (deprecated).
       //   2. Canonical `forwardedProps` option (wins over `body`).
-      //   3. Per-message `body` arg passed to `sendMessage` (highest).
+      //   3. Per-call body (`pendingMessageBody`: positional + sendOptions.body).
       // The AG-UI standard `threadId` is sent at the wire's top level for
       // run/conversation correlation, so we no longer auto-emit a separate
       // `conversationId` here — `chat({ threadId })` server-side covers the
@@ -2200,6 +2257,15 @@ export class ChatClient<
       // AG-UI servers consuming `RunAgentInput.tools[].parameters` expect
       // JSON Schema; sending a Standard Schema instance directly would
       // serialize to an unusable shape.
+      let byokHeaders: Record<string, string> | undefined
+      if (this.byok) {
+        const provider = resolveByokProviderId(
+          this.byokProvider,
+          mergedBody.provider,
+        )
+        byokHeaders = await prepareResolvedByokHeaders(this.byok, provider)
+      }
+
       const runContext = {
         threadId: resumeThreadId ?? this.threadId,
         runId,
@@ -2215,6 +2281,7 @@ export class ChatClient<
         })),
         forwardedProps: { ...mergedBody },
         ...(resumeItems ? { resume: resumeItems } : {}),
+        ...(byokHeaders ? { headers: byokHeaders } : {}),
       }
       this.devtoolsBridge.beginRun(runContext.runId, runContext.threadId)
       activeDevtoolsRunId = runContext.runId
@@ -2270,32 +2337,45 @@ export class ChatClient<
       // Finalize (idempotent — may already be done by RUN_FINISHED handler)
       this.processor.finalizeStream()
       streamCompletedSuccessfully = true
-    } catch (err) {
-      if (err instanceof Error) {
-        if (err.name === 'AbortError') {
-          if (activeDevtoolsRunId) {
-            this.devtoolsBridge.emitRunLifecycle(
-              'run:cancelled',
-              activeDevtoolsRunId,
-              'cancelled',
-            )
-            runTerminalEventEmitted = true
-          }
-          return false
+    } catch (err: unknown) {
+      const error = err instanceof Error ? err : new Error(String(err))
+      if (error.name === 'AbortError') {
+        if (activeDevtoolsRunId) {
+          this.devtoolsBridge.emitRunLifecycle(
+            'run:cancelled',
+            activeDevtoolsRunId,
+            'cancelled',
+          )
+          runTerminalEventEmitted = true
         }
-        if (generation === this.streamGeneration) {
-          this.reportStreamError(err)
-          if (activeDevtoolsRunId) {
-            this.devtoolsBridge.emitRunLifecycle(
-              'run:errored',
-              activeDevtoolsRunId,
-              'errored',
-              { error: err.message },
-            )
-            runTerminalEventEmitted = true
-          }
+        return false
+      }
+      if (error instanceof ByokMissingError) {
+        this.byok?.request(error.provider, 'missing')
+      }
+      if (error instanceof ByokBlockedError && error.reason === 'locked') {
+        this.byok?.request(error.provider, 'locked')
+      }
+      if (generation === this.streamGeneration) {
+        this.reportStreamError(error)
+        if (activeDevtoolsRunId) {
+          this.devtoolsBridge.emitRunLifecycle(
+            'run:errored',
+            activeDevtoolsRunId,
+            'errored',
+            { error: error.message },
+          )
+          runTerminalEventEmitted = true
         }
       }
+      if (
+        generation === this.streamGeneration &&
+        (error instanceof ByokMissingError ||
+          error instanceof ByokBlockedError ||
+          error instanceof ByokUnresolvedProviderError)
+      ) {
+        throw error
+      }
     } finally {
       // Only clean up if this is still the active stream.
       // A superseded stream (e.g. reload() started a new one) must not
@@ -2330,43 +2410,18 @@ export class ChatClient<
         // Drain any actions that were queued while the stream was in progress
         await this.drainPostStreamActions()
 
-        // Continue conversation if the stream ended with a tool result (server tool completed)
-        // but ONLY if the model indicated it wants to continue (finishReason !== 'stop').
-        // When finishReason is 'stop', the model is done — don't re-send.
         if (streamCompletedSuccessfully) {
-          const messages = this.processor.getMessages()
-          const lastPart = messages.at(-1)?.parts.at(-1)
-          const { finishReason } = this.processor.getState()
-
-          if (
-            lastPart?.type === 'tool-result' &&
-            finishReason !== 'stop' &&
-            this.shouldAutoSend()
-          ) {
-            try {
-              await this.checkForContinuation()
-            } catch (error) {
-              console.error('Failed to continue flow after tool result:', error)
-              // Continuation failed without starting a new stream — don't
-              // leave queued user messages stranded forever. (isLoading is
-              // already false in this finally block.)
-              await this.drainQueue()
-            }
-          } else {
-            if (this.status !== 'ready') {
-              // Terminal run, but onStreamEnd never fired: the processor had
-              // no assistant message to emit it for (e.g. a bare
-              // RUN_FINISHED{stop}, #421). The normal path already set
-              // 'ready', so this is a no-op.
-              this.setStatus('ready')
-            }
-            // Auto-send queued messages once the run fully settles. When a
-            // continuation runs instead (tool-result branch above), that
-            // continuation's own finally drains the queue. Skip if a drain
-            // loop is already walking the queue (avoids nested re-entry).
-            if (!this.messageQueueDraining) {
-              await this.drainQueue()
-            }
+          if (this.status !== 'ready') {
+            // Terminal run, but onStreamEnd never fired: the processor had
+            // no assistant message to emit it for (e.g. a bare
+            // RUN_FINISHED{stop}, #421). The normal path already set
+            // 'ready', so this is a no-op.
+            this.setStatus('ready')
+          }
+          // Auto-send queued messages once the run fully settles. Skip if a
+          // drain loop is already walking the queue (avoids nested re-entry).
+          if (!this.messageQueueDraining) {
+            await this.drainQueue()
           }
         } else {
           // Error/abort settle for the active generation: don't strand or
@@ -2948,6 +3003,12 @@ export class ChatClient<
     if (options.forwardedProps !== undefined) {
       this.forwardedPropsOption = options.forwardedProps
     }
+    if (options.byok !== undefined) {
+      this.byok = options.byok
+    }
+    if (options.byokProvider !== undefined) {
+      this.byokProvider = options.byokProvider
+    }
     if ('context' in options) {
       this.context = options.context
     }