@tanstack/ai-client 0.26.0 → 0.28.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (49) hide show
  1. package/dist/esm/byok/client.d.ts +52 -0
  2. package/dist/esm/byok/client.js +300 -0
  3. package/dist/esm/byok/client.js.map +1 -0
  4. package/dist/esm/byok/passkey.d.ts +43 -0
  5. package/dist/esm/byok/passkey.js +251 -0
  6. package/dist/esm/byok/passkey.js.map +1 -0
  7. package/dist/esm/byok/resolve.d.ts +8 -0
  8. package/dist/esm/byok/resolve.js +20 -0
  9. package/dist/esm/byok/resolve.js.map +1 -0
  10. package/dist/esm/byok/storage.d.ts +15 -0
  11. package/dist/esm/byok/storage.js +16 -0
  12. package/dist/esm/byok/storage.js.map +1 -0
  13. package/dist/esm/byok.d.ts +8 -0
  14. package/dist/esm/byok.js +5 -0
  15. package/dist/esm/chat-client.d.ts +16 -3
  16. package/dist/esm/chat-client.js +49 -31
  17. package/dist/esm/chat-client.js.map +1 -1
  18. package/dist/esm/connection-adapters.d.ts +2 -0
  19. package/dist/esm/connection-adapters.js +33 -9
  20. package/dist/esm/connection-adapters.js.map +1 -1
  21. package/dist/esm/generation-client.d.ts +3 -1
  22. package/dist/esm/generation-client.js +23 -4
  23. package/dist/esm/generation-client.js.map +1 -1
  24. package/dist/esm/generation-types.d.ts +14 -0
  25. package/dist/esm/generation-types.js.map +1 -1
  26. package/dist/esm/interrupt-manager.d.ts +1 -0
  27. package/dist/esm/interrupt-manager.js +15 -6
  28. package/dist/esm/interrupt-manager.js.map +1 -1
  29. package/dist/esm/sse-parser.js +8 -1
  30. package/dist/esm/sse-parser.js.map +1 -1
  31. package/dist/esm/types.d.ts +29 -3
  32. package/dist/esm/types.js.map +1 -1
  33. package/dist/esm/video-generation-client.d.ts +3 -1
  34. package/dist/esm/video-generation-client.js +25 -6
  35. package/dist/esm/video-generation-client.js.map +1 -1
  36. package/package.json +6 -2
  37. package/src/byok/client.ts +360 -0
  38. package/src/byok/passkey.ts +399 -0
  39. package/src/byok/resolve.ts +30 -0
  40. package/src/byok/storage.ts +29 -0
  41. package/src/byok.ts +20 -0
  42. package/src/chat-client.ts +95 -65
  43. package/src/connection-adapters.ts +41 -12
  44. package/src/generation-client.ts +48 -4
  45. package/src/generation-types.ts +16 -0
  46. package/src/interrupt-manager.ts +20 -6
  47. package/src/sse-parser.ts +10 -0
  48. package/src/types.ts +31 -3
  49. package/src/video-generation-client.ts +45 -4
@@ -9,6 +9,15 @@ import {
9
9
  restoreInboundChunk,
10
10
  tanstackMetadata,
11
11
  } from '@tanstack/ai/client'
12
+ import {
13
+ ByokBlockedError,
14
+ ByokMissingError,
15
+ ByokUnresolvedProviderError,
16
+ } from '@tanstack/ai/byok'
17
+ import {
18
+ prepareResolvedByokHeaders,
19
+ resolveByokProviderId,
20
+ } from './byok/resolve'
12
21
  import { createNoOpChatDevtoolsBridge } from './devtools-noop'
13
22
  import {
14
23
  fetcherToConnectionAdapter,
@@ -27,6 +36,7 @@ import type {
27
36
  RunAgentResumeItem,
28
37
  StreamChunk,
29
38
  } from '@tanstack/ai/client'
39
+ import type { ByokClient } from './byok'
30
40
  import type {
31
41
  ChatHydrationResult,
32
42
  ConnectionAdapter,
@@ -97,6 +107,8 @@ type ChatClientUpdateOptionsWithoutContext<
97
107
  /** @deprecated Use `forwardedProps` instead. */
98
108
  body?: Record<string, any>
99
109
  forwardedProps?: Record<string, any>
110
+ byok?: ByokClient
111
+ byokProvider?: () => string | undefined
100
112
  tools?: TTools
101
113
  interrupts?: TInterrupts
102
114
  queue?: QueueOption
@@ -348,6 +360,8 @@ export class ChatClient<
348
360
  // merged on every send, with `forwardedProps` winning on key collision.
349
361
  private bodyOption: Record<string, any> = {}
350
362
  private forwardedPropsOption: Record<string, any> = {}
363
+ private byok: ByokClient | undefined
364
+ private byokProvider: (() => string | undefined) | undefined
351
365
  private context: TContext | undefined = undefined
352
366
  private pendingMessageBody: Record<string, any> | undefined = undefined
353
367
  private queueConfig: NormalizedQueueConfig
@@ -494,6 +508,8 @@ export class ChatClient<
494
508
  // winning on key collision.
495
509
  this.bodyOption = options.body || {}
496
510
  this.forwardedPropsOption = options.forwardedProps || {}
511
+ this.byok = options.byok
512
+ this.byokProvider = options.byokProvider
497
513
  this.context = options.context
498
514
  this.queueConfig = normalizeQueueOption(options.queue)
499
515
  this.connection = normalizeConnectionAdapter(resolveTransport(options))
@@ -1882,8 +1898,12 @@ export class ChatClient<
1882
1898
  * - A MultimodalContent object with content array and optional custom ID
1883
1899
  * @param body - Optional body parameters to merge with the client's base body for this request.
1884
1900
  * Uses shallow merge with per-message body taking priority.
1885
- * @param sendOptions - Per-call overrides, e.g. `{ whenBusy: 'interrupt' }` to
1886
- * override the configured queue policy for this one send.
1901
+ * @param sendOptions - Per-call overrides. `{ whenBusy }` overrides the
1902
+ * queue policy for this one send. `{ body }`
1903
+ * shallow-merges with `body` and with the chat-level
1904
+ * `body` / `forwardedProps`. `sendOptions.body` wins
1905
+ * on key collisions. Framework hooks forward this
1906
+ * object as their second argument.
1887
1907
  *
1888
1908
  * @example
1889
1909
  * ```ts
@@ -1893,9 +1913,13 @@ export class ChatClient<
1893
1913
  * // Text message with custom body params
1894
1914
  * await client.sendMessage('Hello!', { temperature: 0.7 })
1895
1915
  *
1896
- * // Per-call whenBusy override (body must still be the 2nd arg on ChatClient)
1916
+ * // Per-call whenBusy override
1897
1917
  * await client.sendMessage('Urgent', undefined, { whenBusy: 'interrupt' })
1898
1918
  *
1919
+ * // Per-call body via options. Same effect as the positional arg.
1920
+ * // This is the shape the framework hooks (`useChat`, `injectChat`) forward.
1921
+ * await client.sendMessage('Hello!', undefined, { body: { temperature: 0.7 } })
1922
+ *
1899
1923
  * // Multimodal message with image
1900
1924
  * await client.sendMessage({
1901
1925
  * content: [
@@ -1933,13 +1957,15 @@ export class ChatClient<
1933
1957
  )
1934
1958
  }
1935
1959
 
1960
+ const resolvedBody = { ...body, ...sendOptions?.body }
1961
+
1936
1962
  if (this.isSendBusy()) {
1937
1963
  const { action, id } = this.decideWhenBusy(content, sendOptions)
1938
1964
  if (action === 'drop') {
1939
1965
  return
1940
1966
  }
1941
1967
  if (action === 'queue') {
1942
- this.enqueueMessage(content, body, id)
1968
+ this.enqueueMessage(content, resolvedBody, id)
1943
1969
  return
1944
1970
  }
1945
1971
  // 'interrupt': abort the current stream, then send now.
@@ -1956,7 +1982,7 @@ export class ChatClient<
1956
1982
  }
1957
1983
 
1958
1984
  try {
1959
- await this.deliverMessage(content, body)
1985
+ await this.deliverMessage(content, resolvedBody)
1960
1986
  } finally {
1961
1987
  this.sendInFlight = false
1962
1988
  }
@@ -2191,7 +2217,7 @@ export class ChatClient<
2191
2217
  // order (later spreads win):
2192
2218
  // 1. Legacy `body` option (deprecated).
2193
2219
  // 2. Canonical `forwardedProps` option (wins over `body`).
2194
- // 3. Per-message `body` arg passed to `sendMessage` (highest).
2220
+ // 3. Per-call body (`pendingMessageBody`: positional + sendOptions.body).
2195
2221
  // The AG-UI standard `threadId` is sent at the wire's top level for
2196
2222
  // run/conversation correlation, so we no longer auto-emit a separate
2197
2223
  // `conversationId` here — `chat({ threadId })` server-side covers the
@@ -2231,6 +2257,15 @@ export class ChatClient<
2231
2257
  // AG-UI servers consuming `RunAgentInput.tools[].parameters` expect
2232
2258
  // JSON Schema; sending a Standard Schema instance directly would
2233
2259
  // serialize to an unusable shape.
2260
+ let byokHeaders: Record<string, string> | undefined
2261
+ if (this.byok) {
2262
+ const provider = resolveByokProviderId(
2263
+ this.byokProvider,
2264
+ mergedBody.provider,
2265
+ )
2266
+ byokHeaders = await prepareResolvedByokHeaders(this.byok, provider)
2267
+ }
2268
+
2234
2269
  const runContext = {
2235
2270
  threadId: resumeThreadId ?? this.threadId,
2236
2271
  runId,
@@ -2246,6 +2281,7 @@ export class ChatClient<
2246
2281
  })),
2247
2282
  forwardedProps: { ...mergedBody },
2248
2283
  ...(resumeItems ? { resume: resumeItems } : {}),
2284
+ ...(byokHeaders ? { headers: byokHeaders } : {}),
2249
2285
  }
2250
2286
  this.devtoolsBridge.beginRun(runContext.runId, runContext.threadId)
2251
2287
  activeDevtoolsRunId = runContext.runId
@@ -2301,32 +2337,45 @@ export class ChatClient<
2301
2337
  // Finalize (idempotent — may already be done by RUN_FINISHED handler)
2302
2338
  this.processor.finalizeStream()
2303
2339
  streamCompletedSuccessfully = true
2304
- } catch (err) {
2305
- if (err instanceof Error) {
2306
- if (err.name === 'AbortError') {
2307
- if (activeDevtoolsRunId) {
2308
- this.devtoolsBridge.emitRunLifecycle(
2309
- 'run:cancelled',
2310
- activeDevtoolsRunId,
2311
- 'cancelled',
2312
- )
2313
- runTerminalEventEmitted = true
2314
- }
2315
- return false
2340
+ } catch (err: unknown) {
2341
+ const error = err instanceof Error ? err : new Error(String(err))
2342
+ if (error.name === 'AbortError') {
2343
+ if (activeDevtoolsRunId) {
2344
+ this.devtoolsBridge.emitRunLifecycle(
2345
+ 'run:cancelled',
2346
+ activeDevtoolsRunId,
2347
+ 'cancelled',
2348
+ )
2349
+ runTerminalEventEmitted = true
2316
2350
  }
2317
- if (generation === this.streamGeneration) {
2318
- this.reportStreamError(err)
2319
- if (activeDevtoolsRunId) {
2320
- this.devtoolsBridge.emitRunLifecycle(
2321
- 'run:errored',
2322
- activeDevtoolsRunId,
2323
- 'errored',
2324
- { error: err.message },
2325
- )
2326
- runTerminalEventEmitted = true
2327
- }
2351
+ return false
2352
+ }
2353
+ if (error instanceof ByokMissingError) {
2354
+ this.byok?.request(error.provider, 'missing')
2355
+ }
2356
+ if (error instanceof ByokBlockedError && error.reason === 'locked') {
2357
+ this.byok?.request(error.provider, 'locked')
2358
+ }
2359
+ if (generation === this.streamGeneration) {
2360
+ this.reportStreamError(error)
2361
+ if (activeDevtoolsRunId) {
2362
+ this.devtoolsBridge.emitRunLifecycle(
2363
+ 'run:errored',
2364
+ activeDevtoolsRunId,
2365
+ 'errored',
2366
+ { error: error.message },
2367
+ )
2368
+ runTerminalEventEmitted = true
2328
2369
  }
2329
2370
  }
2371
+ if (
2372
+ generation === this.streamGeneration &&
2373
+ (error instanceof ByokMissingError ||
2374
+ error instanceof ByokBlockedError ||
2375
+ error instanceof ByokUnresolvedProviderError)
2376
+ ) {
2377
+ throw error
2378
+ }
2330
2379
  } finally {
2331
2380
  // Only clean up if this is still the active stream.
2332
2381
  // A superseded stream (e.g. reload() started a new one) must not
@@ -2361,43 +2410,18 @@ export class ChatClient<
2361
2410
  // Drain any actions that were queued while the stream was in progress
2362
2411
  await this.drainPostStreamActions()
2363
2412
 
2364
- // Continue conversation if the stream ended with a tool result (server tool completed)
2365
- // but ONLY if the model indicated it wants to continue (finishReason !== 'stop').
2366
- // When finishReason is 'stop', the model is done — don't re-send.
2367
2413
  if (streamCompletedSuccessfully) {
2368
- const messages = this.processor.getMessages()
2369
- const lastPart = messages.at(-1)?.parts.at(-1)
2370
- const { finishReason } = this.processor.getState()
2371
-
2372
- if (
2373
- lastPart?.type === 'tool-result' &&
2374
- finishReason !== 'stop' &&
2375
- this.shouldAutoSend()
2376
- ) {
2377
- try {
2378
- await this.checkForContinuation()
2379
- } catch (error) {
2380
- console.error('Failed to continue flow after tool result:', error)
2381
- // Continuation failed without starting a new stream — don't
2382
- // leave queued user messages stranded forever. (isLoading is
2383
- // already false in this finally block.)
2384
- await this.drainQueue()
2385
- }
2386
- } else {
2387
- if (this.status !== 'ready') {
2388
- // Terminal run, but onStreamEnd never fired: the processor had
2389
- // no assistant message to emit it for (e.g. a bare
2390
- // RUN_FINISHED{stop}, #421). The normal path already set
2391
- // 'ready', so this is a no-op.
2392
- this.setStatus('ready')
2393
- }
2394
- // Auto-send queued messages once the run fully settles. When a
2395
- // continuation runs instead (tool-result branch above), that
2396
- // continuation's own finally drains the queue. Skip if a drain
2397
- // loop is already walking the queue (avoids nested re-entry).
2398
- if (!this.messageQueueDraining) {
2399
- await this.drainQueue()
2400
- }
2414
+ if (this.status !== 'ready') {
2415
+ // Terminal run, but onStreamEnd never fired: the processor had
2416
+ // no assistant message to emit it for (e.g. a bare
2417
+ // RUN_FINISHED{stop}, #421). The normal path already set
2418
+ // 'ready', so this is a no-op.
2419
+ this.setStatus('ready')
2420
+ }
2421
+ // Auto-send queued messages once the run fully settles. Skip if a
2422
+ // drain loop is already walking the queue (avoids nested re-entry).
2423
+ if (!this.messageQueueDraining) {
2424
+ await this.drainQueue()
2401
2425
  }
2402
2426
  } else {
2403
2427
  // Error/abort settle for the active generation: don't strand or
@@ -2979,6 +3003,12 @@ export class ChatClient<
2979
3003
  if (options.forwardedProps !== undefined) {
2980
3004
  this.forwardedPropsOption = options.forwardedProps
2981
3005
  }
3006
+ if (options.byok !== undefined) {
3007
+ this.byok = options.byok
3008
+ }
3009
+ if (options.byokProvider !== undefined) {
3010
+ this.byokProvider = options.byokProvider
3011
+ }
2982
3012
  if ('context' in options) {
2983
3013
  this.context = options.context
2984
3014
  }
@@ -6,6 +6,7 @@ import {
6
6
  uiMessagesToWire,
7
7
  withTanstackMetadata,
8
8
  } from '@tanstack/ai/client'
9
+ import { ByokMissingError, isByokMissingBody } from '@tanstack/ai/byok'
9
10
  import {
10
11
  createResponseStreamTextDecoder,
11
12
  getResponseStreamReader,
@@ -498,12 +499,35 @@ async function* linesToNdjsonEvents(
498
499
  }
499
500
  }
500
501
 
501
- function assertResponseOk(response: Response): void {
502
- if (!response.ok) {
503
- throw new Error(
504
- `HTTP error! status: ${response.status} ${response.statusText}`,
505
- )
502
+ async function assertResponseOk(response: Response): Promise<void> {
503
+ if (response.ok) return
504
+ if (response.status === 401) {
505
+ const body: unknown = await response
506
+ .clone()
507
+ .json()
508
+ .catch(() => null)
509
+ if (isByokMissingBody(body)) {
510
+ throw new ByokMissingError(body.error.provider)
511
+ }
512
+ }
513
+ throw new Error(
514
+ `HTTP error! status: ${response.status} ${response.statusText}`,
515
+ )
516
+ }
517
+
518
+ function errorFromXhrStatus(xhr: XMLHttpRequest): Error {
519
+ if (xhr.status === 401) {
520
+ let parsed: unknown = null
521
+ try {
522
+ parsed = JSON.parse(xhr.responseText)
523
+ } catch {
524
+ parsed = null
525
+ }
526
+ if (isByokMissingBody(parsed)) {
527
+ return new ByokMissingError(parsed.error.provider)
528
+ }
506
529
  }
530
+ return new Error(`XHR error! status: ${xhr.status} ${xhr.statusText}`)
507
531
  }
508
532
 
509
533
  /**
@@ -525,7 +549,7 @@ async function fetchThreadHydration(
525
549
  headers: { Accept: 'application/json', ...headers },
526
550
  credentials,
527
551
  })
528
- assertResponseOk(response)
552
+ await assertResponseOk(response)
529
553
  const data = (await response.json()) as {
530
554
  messages?: Array<UIMessage>
531
555
  activeRun?: { runId?: unknown } | null
@@ -574,7 +598,7 @@ async function fetchGenerationHydration(
574
598
  headers: { Accept: 'application/json', ...headers },
575
599
  credentials,
576
600
  })
577
- assertResponseOk(response)
601
+ await assertResponseOk(response)
578
602
  const raw: unknown = await response.json()
579
603
  // A 200 carrying `null` is a legitimate hydration miss — the server has no
580
604
  // record for this thread — and reading `.activeRun` off `null` would throw.
@@ -609,7 +633,7 @@ async function* responseToSSEEvents(
609
633
  abortSignal?: AbortSignal,
610
634
  fallbackIds?: { threadId?: string; runId?: string },
611
635
  ): AsyncGenerator<StreamEvent> {
612
- assertResponseOk(response)
636
+ await assertResponseOk(response)
613
637
  const reader = getResponseStreamReader(response)
614
638
  yield* linesToSSEEvents(readStreamLines(reader, abortSignal), fallbackIds)
615
639
  }
@@ -619,7 +643,7 @@ async function* responseToNdjsonEvents(
619
643
  response: Response,
620
644
  abortSignal?: AbortSignal,
621
645
  ): AsyncGenerator<StreamEvent> {
622
- assertResponseOk(response)
646
+ await assertResponseOk(response)
623
647
  const reader = getResponseStreamReader(response)
624
648
  yield* linesToNdjsonEvents(readStreamLines(reader, abortSignal))
625
649
  }
@@ -809,6 +833,8 @@ export interface RunAgentInputContext {
809
833
  }>
810
834
  /** Arbitrary user-controlled passthrough data. */
811
835
  forwardedProps?: Record<string, unknown>
836
+ /** Extra request headers for this run (e.g. BYOK keys). POST only. */
837
+ headers?: Record<string, string>
812
838
  }
813
839
 
814
840
  export interface ConnectConnectionAdapter {
@@ -1272,6 +1298,7 @@ export function fetchServerSentEvents(
1272
1298
  const requestHeaders: Record<string, string> = {
1273
1299
  'Content-Type': 'application/json',
1274
1300
  ...mergeHeaders(resolvedOptions.headers),
1301
+ ...mergeHeaders(runContext?.headers),
1275
1302
  ...runIdHeader(runContext?.runId),
1276
1303
  }
1277
1304
 
@@ -1443,6 +1470,7 @@ export function fetchHttpStream(
1443
1470
  const requestHeaders: Record<string, string> = {
1444
1471
  'Content-Type': 'application/json',
1445
1472
  ...mergeHeaders(resolvedOptions.headers),
1473
+ ...mergeHeaders(runContext?.headers),
1446
1474
  ...runIdHeader(runContext?.runId),
1447
1475
  }
1448
1476
 
@@ -1596,7 +1624,7 @@ function readXhrLines(
1596
1624
 
1597
1625
  const enqueueDelta = () => {
1598
1626
  if (xhr.status !== 0 && (xhr.status < 200 || xhr.status >= 300)) {
1599
- error = new Error(`XHR error! status: ${xhr.status} ${xhr.statusText}`)
1627
+ error = errorFromXhrStatus(xhr)
1600
1628
  done = true
1601
1629
  return
1602
1630
  }
@@ -1625,7 +1653,7 @@ function readXhrLines(
1625
1653
  // is an error, but status 0 here is not — treat the trailing buffer as a
1626
1654
  // truncation check instead of fabricating a bogus "status: 0" error.
1627
1655
  if (xhr.status !== 0 && (xhr.status < 200 || xhr.status >= 300)) {
1628
- error = new Error(`XHR error! status: ${xhr.status} ${xhr.statusText}`)
1656
+ error = errorFromXhrStatus(xhr)
1629
1657
  } else if (buffer.trim() && !aborted) {
1630
1658
  error = new StreamTruncatedError()
1631
1659
  }
@@ -1721,6 +1749,7 @@ function createConfiguredXhrRequest(
1721
1749
  const requestHeaders: Record<string, string> = {
1722
1750
  'Content-Type': 'application/json',
1723
1751
  ...mergeHeaders(options.headers),
1752
+ ...mergeHeaders(method === 'POST' ? runContext?.headers : undefined),
1724
1753
  // Client-chosen run id for durability (POST only; the GET join carries it
1725
1754
  // in the query instead).
1726
1755
  ...(method === 'POST' ? runIdHeader(runContext?.runId) : {}),
@@ -2494,7 +2523,7 @@ export function fetcherToConnectionAdapter(
2494
2523
  ? { resume: runContext.resume }
2495
2524
  : {}),
2496
2525
  },
2497
- { signal: abortSignal },
2526
+ { signal: abortSignal, headers: runContext.headers },
2498
2527
  )
2499
2528
  if (result instanceof Response) {
2500
2529
  yield* responseToSSEChunks(result, abortSignal)
@@ -1,3 +1,8 @@
1
+ import { ByokBlockedError, ByokMissingError } from '@tanstack/ai/byok'
2
+ import {
3
+ prepareResolvedByokHeaders,
4
+ resolveByokProviderId,
5
+ } from './byok/resolve'
1
6
  import {
2
7
  GENERATION_EVENTS,
3
8
  GENERATION_STREAM_TRUNCATED_MESSAGE,
@@ -12,6 +17,7 @@ import { createNoOpGenerationDevtoolsBridge } from './devtools-noop'
12
17
  import { parseSSEResponse } from './sse-parser'
13
18
  import { restoreInboundChunk } from '@tanstack/ai/client'
14
19
  import type { StreamChunk } from '@tanstack/ai/client'
20
+ import type { ByokClient } from './byok'
15
21
  import type {
16
22
  ConnectConnectionAdapter,
17
23
  GenerationHydrationResult,
@@ -120,6 +126,8 @@ export class GenerationClient<
120
126
  // the client hydrates the last generation for `threadId` from the server.
121
127
  private readonly serverDriven: boolean = false
122
128
  private body: Record<string, any>
129
+ private byok: ByokClient | undefined
130
+ private byokProvider: (() => string | undefined) | undefined
123
131
  private result: TOutput | null = null
124
132
  private input: TInput | null = null
125
133
  private progress: AIDevtoolsGenerationProgress | null = null
@@ -156,6 +164,8 @@ export class GenerationClient<
156
164
  this.hydrateGenerationHandler = options.hydrateGeneration
157
165
  this.joinRunHandler = options.joinRun
158
166
  this.body = options.body ?? {}
167
+ this.byok = options.byok
168
+ this.byokProvider = options.byokProvider
159
169
  // `persistence` is `false`/omitted (ephemeral) or `true` (server-driven:
160
170
  // hydrate the last generation for `threadId` from the server on mount).
161
171
  this.serverDriven = options.persistence === true
@@ -265,9 +275,21 @@ export class GenerationClient<
265
275
  const { signal } = abortController
266
276
 
267
277
  try {
278
+ let headers: Record<string, string> | undefined
279
+ if (this.byok) {
280
+ const provider = resolveByokProviderId(
281
+ this.byokProvider,
282
+ this.body.provider,
283
+ )
284
+ headers = await prepareResolvedByokHeaders(this.byok, provider)
285
+ }
286
+
268
287
  if (this.fetcher) {
269
288
  // Direct fetch path
270
- const result = await this.fetcher(input, { signal })
289
+ const result = await this.fetcher(
290
+ input,
291
+ headers === undefined ? { signal } : { signal, headers },
292
+ )
271
293
  if (signal.aborted) return
272
294
  if (result instanceof Response) {
273
295
  // Server function returned SSE Response — parse stream
@@ -289,7 +311,7 @@ export class GenerationClient<
289
311
  [],
290
312
  mergedData,
291
313
  signal,
292
- this.createRunContext(runId),
314
+ this.createRunContext(runId, headers),
293
315
  )
294
316
  await this.processStream(stream, runId, signal)
295
317
  } else {
@@ -312,6 +334,12 @@ export class GenerationClient<
312
334
  } catch (err: unknown) {
313
335
  if (signal.aborted) return
314
336
  const error = err instanceof Error ? err : new Error(String(err))
337
+ if (error instanceof ByokMissingError) {
338
+ this.byok?.request(error.provider, 'missing')
339
+ }
340
+ if (error instanceof ByokBlockedError && error.reason === 'locked') {
341
+ this.byok?.request(error.provider, 'locked')
342
+ }
315
343
  this.setError(error)
316
344
  this.setStatus('error')
317
345
  this.recordResumeSnapshotError(error)
@@ -462,13 +490,25 @@ export class GenerationClient<
462
490
  options: Partial<
463
491
  Pick<
464
492
  GenerationClientOptions<TInput, TResult, TOutput>,
465
- 'body' | 'onResult' | 'onError' | 'onProgress' | 'onChunk'
493
+ | 'body'
494
+ | 'byok'
495
+ | 'byokProvider'
496
+ | 'onResult'
497
+ | 'onError'
498
+ | 'onProgress'
499
+ | 'onChunk'
466
500
  >
467
501
  >,
468
502
  ): void {
469
503
  if (options.body !== undefined) {
470
504
  this.body = options.body ?? {}
471
505
  }
506
+ if (options.byok !== undefined) {
507
+ this.byok = options.byok
508
+ }
509
+ if (options.byokProvider !== undefined) {
510
+ this.byokProvider = options.byokProvider
511
+ }
472
512
  if (options.onResult !== undefined) {
473
513
  this.callbacksRef.onResult = options.onResult
474
514
  }
@@ -645,10 +685,14 @@ export class GenerationClient<
645
685
  return `${prefix}-${Date.now()}-${Math.random().toString(36).substring(7)}`
646
686
  }
647
687
 
648
- private createRunContext(runId: string): RunAgentInputContext {
688
+ private createRunContext(
689
+ runId: string,
690
+ headers?: Record<string, string>,
691
+ ): RunAgentInputContext {
649
692
  return {
650
693
  threadId: this.threadId,
651
694
  runId,
695
+ ...(headers ? { headers } : {}),
652
696
  }
653
697
  }
654
698
 
@@ -5,6 +5,7 @@ import type {
5
5
  StreamChunk,
6
6
  } from '@tanstack/ai/client'
7
7
  import type { TokenUsage, TranscriptionResponseFormat } from '@tanstack/ai'
8
+ import type { ByokClient } from './byok'
8
9
  import type { ConnectConnectionAdapter } from './connection-adapters'
9
10
  import type { AIDevtoolsClientMetadata } from './devtools'
10
11
  import type {
@@ -271,6 +272,8 @@ export const GENERATION_EVENTS = {
271
272
  export interface GenerationFetcherOptions {
272
273
  /** AbortSignal that is triggered when the user calls `stop()` */
273
274
  signal: AbortSignal
275
+ /** Extra request headers for this run (e.g. BYOK keys). */
276
+ headers?: Record<string, string>
274
277
  }
275
278
 
276
279
  /**
@@ -338,6 +341,19 @@ export interface GenerationClientOptions<_TInput, TResult, TOutput = TResult> {
338
341
  /** Additional body parameters to send with connect-based adapter requests */
339
342
  body?: Record<string, any>
340
343
 
344
+ /**
345
+ * Optional BYOK keyring. On each generate the client prepares the resolved
346
+ * provider and stamps `x-byok-*` request headers. Keys never go in the body.
347
+ */
348
+ byok?: ByokClient
349
+
350
+ /**
351
+ * Optional provider id for this generation. If it returns a provider slug,
352
+ * only that key is prepared and sent. Otherwise `body.provider` is used.
353
+ * If no slug resolves, generate throws instead of attaching every stored key.
354
+ */
355
+ byokProvider?: () => string | undefined
356
+
341
357
  /** Metadata used to register this generation hook with TanStack AI Devtools */
342
358
  devtools?: Partial<AIDevtoolsClientMetadata>
343
359
 
@@ -102,6 +102,10 @@ interface TransactionToken {
102
102
  active: boolean
103
103
  }
104
104
 
105
+ interface SubmissionOperation {
106
+ submission: InterruptManagerSubmission
107
+ }
108
+
105
109
  interface RuntimeInterruptCheckpoint {
106
110
  status: InterruptItemStatus
107
111
  resolution?: RunAgentResumeItem
@@ -511,6 +515,7 @@ export class InterruptManager<
511
515
  resuming: false,
512
516
  })
513
517
  private activeTransaction: TransactionToken | undefined
518
+ private activeSubmissionOperation: SubmissionOperation | undefined
514
519
  private retrySubmission: InterruptManagerSubmission | undefined
515
520
  private resuming = false
516
521
  private tools: TTools | undefined
@@ -544,6 +549,7 @@ export class InterruptManager<
544
549
  hydration: InterruptManagerHydration,
545
550
  source: InterruptManagerChangeSource = 'live',
546
551
  ): void {
552
+ this.activeSubmissionOperation = undefined
547
553
  this.hydration = {
548
554
  threadId: hydration.threadId,
549
555
  interruptedRunId: hydration.interruptedRunId,
@@ -597,6 +603,7 @@ export class InterruptManager<
597
603
  preserveRootErrors?: boolean
598
604
  source?: InterruptManagerChangeSource
599
605
  }): void {
606
+ this.activeSubmissionOperation = undefined
600
607
  this.hydration = undefined
601
608
  this.items = []
602
609
  this.snapshot = Object.freeze([])
@@ -1528,25 +1535,32 @@ export class InterruptManager<
1528
1535
  }
1529
1536
 
1530
1537
  private submitBatch(submission: InterruptManagerSubmission): void {
1538
+ // Track ownership so a superseded submission cannot mutate current state.
1539
+ const operation = { submission }
1540
+ this.activeSubmissionOperation = operation
1531
1541
  this.resuming = true
1532
1542
  this.retrySubmission = undefined
1533
1543
  for (const item of this.items) {
1534
1544
  if (isClientOwnedInterrupt(item)) item.status = 'submitting'
1535
1545
  }
1536
1546
  this.publish()
1537
- void this.performSubmission(submission)
1547
+ void this.performSubmission(operation)
1538
1548
  }
1539
1549
 
1540
1550
  private async performSubmission(
1541
- submission: InterruptManagerSubmission,
1551
+ operation: SubmissionOperation,
1542
1552
  ): Promise<void> {
1543
1553
  try {
1544
- await this.options.submit(submission)
1554
+ await this.options.submit(operation.submission)
1545
1555
  } catch (error) {
1546
- this.handleSubmissionFailure(error, submission)
1556
+ if (this.activeSubmissionOperation !== operation) return
1557
+ this.handleSubmissionFailure(error, operation.submission)
1547
1558
  } finally {
1548
- this.resuming = false
1549
- this.publish()
1559
+ if (this.activeSubmissionOperation === operation) {
1560
+ this.activeSubmissionOperation = undefined
1561
+ this.resuming = false
1562
+ this.publish()
1563
+ }
1550
1564
  }
1551
1565
  }
1552
1566
 
package/src/sse-parser.ts CHANGED
@@ -1,3 +1,4 @@
1
+ import { ByokMissingError, isByokMissingBody } from '@tanstack/ai/byok'
1
2
  import {
2
3
  createResponseStreamTextDecoder,
3
4
  getResponseStreamReader,
@@ -51,6 +52,15 @@ export async function* parseSSEResponse(
51
52
  abortSignal?: AbortSignal,
52
53
  ): AsyncGenerator<StreamChunk> {
53
54
  if (!response.ok) {
55
+ if (response.status === 401) {
56
+ const body: unknown = await response
57
+ .clone()
58
+ .json()
59
+ .catch(() => null)
60
+ if (isByokMissingBody(body)) {
61
+ throw new ByokMissingError(body.error.provider)
62
+ }
63
+ }
54
64
  throw new Error(
55
65
  `HTTP error! status: ${response.status} ${response.statusText}`,
56
66
  )