@tanstack/ai-client 0.26.0 → 0.29.0

package/src/chat-client.ts

@@ -9,6 +9,15 @@ import {
   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,
@@ -17,6 +26,7 @@ import {
 } from './connection-adapters'
 import { ChatPersistor } from './client-persistor'
 import { ClearedStreamTracker } from './cleared-stream-tracker'
+import { normalizeMessagesDates } from './message-date-normalizer'
 import { InterruptManager } from './interrupt-manager'
 import type {
   AnyClientTool,
@@ -27,6 +37,7 @@ import type {
   RunAgentResumeItem,
   StreamChunk,
 } from '@tanstack/ai/client'
+import type { ByokClient } from './byok'
 import type {
   ChatHydrationResult,
   ConnectionAdapter,
@@ -97,6 +108,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
@@ -348,6 +361,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
@@ -403,6 +418,11 @@ export class ChatClient<
   private processingResolve: (() => void) | null = null
   private errorReportedGeneration: number | null = null
   private streamGeneration = 0
+  private continuationGeneration = 0
+  // Generation of the run that opened the current stream. Public
+  // `addToolResult` must use this, not the live counter: `stop()` increments
+  // the live counter, so a post-stop call would otherwise look current.
+  private streamContinuationGeneration = 0
   // Tracks whether a queued checkForContinuation was skipped because
   // continuationPending was true (chained approval scenario)
   private continuationSkipped = false
@@ -494,6 +514,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))
@@ -721,6 +743,7 @@ export class ChatClient<
           const clientTool = clientTools.get(args.toolName)
           const executeFunc = clientTool?.execute
           if (executeFunc) {
+            const continuationGeneration = this.continuationGeneration
             // Capture the run context at execution-start so a tool whose
             // result lands AFTER the originating run finishes still reports
             // back against the originating run, not whatever run is
@@ -747,6 +770,7 @@ export class ChatClient<
                     state: 'output-available',
                   },
                   clientTool,
+                  continuationGeneration,
                   runEventContext,
                 )
               } catch (error: any) {
@@ -759,6 +783,7 @@ export class ChatClient<
                     errorText: error.message,
                   },
                   clientTool,
+                  continuationGeneration,
                   runEventContext,
                 )
               } finally {
@@ -1001,7 +1026,7 @@ export class ChatClient<
       // A send may have started while the fetch was in flight — don't stomp it.
       if (this.isLoading || this.abortController) return
       if (result.messages.length > 0) {
-        this.processor.setMessages(result.messages)
+        this.processor.setMessages(normalizeMessagesDates(result.messages))
       }
       if (result.interrupts && result.interrupts.pending.length > 0) {
         // Pending interrupt = the thread is paused awaiting a human decision, so
@@ -1288,6 +1313,21 @@ export class ChatClient<
   ): Promise<boolean> {
     const target = state ?? this.lastResume
     if (!target) return Promise.resolve(false)
+    return this.resumeInterruptsUnsafeForGeneration(
+      resume,
+      target,
+      this.continuationGeneration,
+    )
+  }
+
+  private resumeInterruptsUnsafeForGeneration(
+    resume: Array<RunAgentResumeItem>,
+    target: ChatResumeState,
+    continuationGeneration: number,
+  ): Promise<boolean> {
+    if (continuationGeneration !== this.continuationGeneration) {
+      return Promise.resolve(false)
+    }
     // Auto-executed client tools resolve during the parent stream's
     // `pendingToolExecutions` wait — while `isLoading` is still true.
     // Defer the child continuation until that stream settles so we do not
@@ -1296,7 +1336,13 @@ export class ChatClient<
       return new Promise<boolean>((resolve, reject) => {
         this.queuePostStreamAction(async () => {
           try {
-            resolve(await this.resumeInterruptsUnsafe(resume, target))
+            resolve(
+              await this.resumeInterruptsUnsafeForGeneration(
+                resume,
+                target,
+                continuationGeneration,
+              ),
+            )
           } catch (error) {
             reject(error)
           }
@@ -1320,6 +1366,7 @@ export class ChatClient<
   private async submitInterruptBatch(
     submission: InterruptManagerSubmission,
   ): Promise<void> {
+    const continuationGeneration = this.continuationGeneration
     this.activeInterruptSubmission = submission
     this.interruptSubmissionFailure = undefined
     // Reflect approval decisions in the local message tree immediately so a
@@ -1331,15 +1378,21 @@ export class ChatClient<
       const approvalId = resolution.interruptId
       this.processor.addToolApprovalResponse(approvalId, approved)
     }
-    const resumed = await this.resumeInterruptsUnsafe(
+    const resumed = await this.resumeInterruptsUnsafeForGeneration(
       [...submission.resolutions],
       {
         threadId: submission.threadId,
         runId: submission.interruptedRunId,
       },
+      continuationGeneration,
     ).finally(() => {
-      this.activeInterruptSubmission = undefined
+      // Only clear if this resume still owns the client: `stop()` may have
+      // invalidated it while the submission was settling.
+      if (this.activeInterruptSubmission === submission) {
+        this.activeInterruptSubmission = undefined
+      }
     })
+    if (continuationGeneration !== this.continuationGeneration) return
     const failure = this.takeInterruptSubmissionFailure()
     if (failure !== undefined) {
       throw { errors: failure.errors }
@@ -1662,6 +1715,7 @@ export class ChatClient<
     // persisted pointer with the provider id — so a SECOND reload would
     // `joinRun` an id the log isn't keyed by and never re-attach.
     this.lastResume = { threadId: this.threadId, runId }
+    this.streamContinuationGeneration = this.continuationGeneration
     this.setIsLoading(true)
     this.setStatus('streaming')
     void (async () => {
@@ -1882,8 +1936,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
@@ -1893,9 +1951,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: [
@@ -1933,13 +1995,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.
@@ -1956,7 +2020,7 @@ export class ChatClient<
     }
 
     try {
-      await this.deliverMessage(content, body)
+      await this.deliverMessage(content, resolvedBody)
     } finally {
       this.sendInFlight = false
     }
@@ -2137,6 +2201,7 @@ export class ChatClient<
 
     // Track generation so a superseded stream's cleanup doesn't clobber the new one
     const generation = ++this.streamGeneration
+    this.streamContinuationGeneration = this.continuationGeneration
     // Native interrupt continuation is a fresh child run. The interrupted run
     // is carried as parentRunId and the complete resolution batch as resume.
     const resumeThreadId = this.pendingResumeThreadId
@@ -2191,7 +2256,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
@@ -2231,6 +2296,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,
@@ -2246,6 +2320,7 @@ export class ChatClient<
         })),
         forwardedProps: { ...mergedBody },
         ...(resumeItems ? { resume: resumeItems } : {}),
+        ...(byokHeaders ? { headers: byokHeaders } : {}),
       }
       this.devtoolsBridge.beginRun(runContext.runId, runContext.threadId)
       activeDevtoolsRunId = runContext.runId
@@ -2301,32 +2376,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
@@ -2361,43 +2449,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
@@ -2482,9 +2545,14 @@ export class ChatClient<
    * Stop the current stream
    */
   stop(): void {
+    // Invalidate deferred work from the stopped continuation.
+    this.continuationGeneration++
     const hadLocalStream = this.abortController !== null
     this.cancelInFlightStream({ setReadyStatus: true })
     this.discardPendingSends()
+    this.lastResume = null
+    this.activeInterruptSubmission = undefined
+    this.interruptManager.reset()
     if (hadLocalStream) {
       this.resetSessionGenerating()
     }
@@ -2528,12 +2596,17 @@ export class ChatClient<
    */
   async addToolResult(result: ClientToolResult): Promise<void> {
     const clientTool = this.clientToolsRef.current.get(result.tool)
-    await this.addToolResultForClientTool(result, clientTool)
+    await this.addToolResultForClientTool(
+      result,
+      clientTool,
+      this.streamContinuationGeneration,
+    )
   }
 
   private async addToolResultForClientTool(
     result: ClientToolResult,
     clientTool: AnyClientTool | undefined,
+    continuationGeneration: number,
     context?: ChatClientRunEventContext,
   ): Promise<void> {
     if (clientTool && result.state !== 'output-error') {
@@ -2560,6 +2633,8 @@ export class ChatClient<
       context,
     )
 
+    if (continuationGeneration !== this.continuationGeneration) return
+
     // Always update local message state so the tool-call part is terminal in
     // the UI even when the AG-UI interrupt path owns server continuation.
     this.processor.addToolResult(
@@ -2585,7 +2660,11 @@ export class ChatClient<
 
     // If stream is in progress, queue continuation check for after it ends
     if (this.isLoading) {
-      this.queuePostStreamAction(() => this.checkForContinuation())
+      this.queuePostStreamAction(() =>
+        continuationGeneration === this.continuationGeneration
+          ? this.checkForContinuation()
+          : Promise.resolve(),
+      )
       return
     }
 
@@ -2665,7 +2744,11 @@ export class ChatClient<
    * Queue an action to be executed after the current stream ends
    */
   private queuePostStreamAction(action: () => Promise<void>): void {
-    this.postStreamActions.push(action)
+    const continuationGeneration = this.continuationGeneration
+    this.postStreamActions.push(async () => {
+      if (continuationGeneration !== this.continuationGeneration) return
+      await action()
+    })
   }
 
   /**
@@ -2688,6 +2771,10 @@ export class ChatClient<
    * Check if we should continue the flow and do so if needed
    */
   private async checkForContinuation(): Promise<void> {
+    // stop() bumps continuationGeneration without opening a new stream.
+    if (this.streamContinuationGeneration !== this.continuationGeneration) {
+      return
+    }
     if (this.hasPendingInterrupts()) return
 
     // Prevent duplicate continuation attempts
@@ -2979,6 +3066,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
     }

package/src/connection-adapters.ts

@@ -6,6 +6,7 @@ import {
   uiMessagesToWire,
   withTanstackMetadata,
 } from '@tanstack/ai/client'
+import { ByokMissingError, isByokMissingBody } from '@tanstack/ai/byok'
 import {
   createResponseStreamTextDecoder,
   getResponseStreamReader,
@@ -19,6 +20,7 @@ import type {
   UIMessage,
 } from '@tanstack/ai/client'
 import type { ChatFetcher, ChatPendingInterrupt } from './types'
+import { normalizeMessagesDates } from './message-date-normalizer'
 
 /**
  * Associates connect-wrapped chunks with the run they were produced under.
@@ -498,12 +500,35 @@ async function* linesToNdjsonEvents(
   }
 }
 
-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}`)
 }
 
 /**
@@ -525,7 +550,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
@@ -549,7 +574,9 @@ async function fetchThreadHydration(
         }
       : null
   return {
-    messages: Array.isArray(data.messages) ? data.messages : [],
+    messages: Array.isArray(data.messages)
+      ? normalizeMessagesDates(data.messages)
+      : [],
     activeRun,
     interrupts,
   }
@@ -574,7 +601,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.
@@ -609,7 +636,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)
 }
@@ -619,7 +646,7 @@ async function* responseToNdjsonEvents(
   response: Response,
   abortSignal?: AbortSignal,
 ): AsyncGenerator<StreamEvent> {
-  assertResponseOk(response)
+  await assertResponseOk(response)
   const reader = getResponseStreamReader(response)
   yield* linesToNdjsonEvents(readStreamLines(reader, abortSignal))
 }
@@ -809,6 +836,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 {
@@ -1198,7 +1227,7 @@ function buildRunAgentInputBody(
 ): Record<string, unknown> {
   // Precedence (later spreads win): static adapter `body` is the base,
   // overridden by `runContext.forwardedProps`, overridden by per-message `data`.
-  const wireMessages = uiMessagesToWire(messages as Array<UIMessage>)
+  const wireMessages = uiMessagesToWire(messages)
   const forwardedProps = {
     ...options.body,
     ...(runContext?.forwardedProps ?? {}),
@@ -1272,6 +1301,7 @@ export function fetchServerSentEvents(
       const requestHeaders: Record<string, string> = {
         'Content-Type': 'application/json',
         ...mergeHeaders(resolvedOptions.headers),
+        ...mergeHeaders(runContext?.headers),
         ...runIdHeader(runContext?.runId),
       }
 
@@ -1443,6 +1473,7 @@ export function fetchHttpStream(
       const requestHeaders: Record<string, string> = {
         'Content-Type': 'application/json',
         ...mergeHeaders(resolvedOptions.headers),
+        ...mergeHeaders(runContext?.headers),
         ...runIdHeader(runContext?.runId),
       }
 
@@ -1596,7 +1627,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
     }
@@ -1625,7 +1656,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()
     }
@@ -1721,6 +1752,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) : {}),
@@ -2494,7 +2526,7 @@ export function fetcherToConnectionAdapter(
             ? { resume: runContext.resume }
             : {}),
         },
-        { signal: abortSignal },
+        { signal: abortSignal, headers: runContext.headers },
       )
       if (result instanceof Response) {
         yield* responseToSSEChunks(result, abortSignal)