@tanstack/ai-react 0.18.1 → 0.19.0

package/src/use-chat.ts

@@ -5,11 +5,15 @@ import type {
   AnyClientTool,
   InferSchemaType,
   ModelMessage,
+  RunAgentResumeItem,
   SchemaInput,
   StreamChunk,
 } from '@tanstack/ai/client'
 import type {
   ChatClientState,
+  ChatInterrupt,
+  ChatInterruptState,
+  ChatResumeState,
   ConnectionStatus,
   InferredClientContext,
   QueuedMessage,
@@ -25,6 +29,9 @@ import type {
   UseChatReturn,
 } from './types'
 
+const EMPTY_INTERRUPTS = Object.freeze([])
+const EMPTY_INTERRUPT_ERRORS = Object.freeze([])
+
 export function useChat<
   const TTools extends ReadonlyArray<AnyClientTool> = any,
   TSchema extends SchemaInput | undefined = undefined,
@@ -32,8 +39,12 @@ export function useChat<
 >(
   options: UseChatOptions<TTools, TSchema, TContext>,
 ): UseChatReturn<TTools, TSchema> {
+  // The hook's identity is its `threadId` — also the persistence key, so a
+  // reload with the same `threadId` restores the same conversation. `hookId` is
+  // only a stable fallback for React's client-recreation keying when no
+  // `threadId` is given (an ephemeral chat), never a persistence key.
   const hookId = useId()
-  const clientId = options.id || hookId
+  const clientId = options.threadId ?? hookId
 
   const [messages, setMessages] = useState<Array<UIMessage<TTools>>>(
     options.initialMessages || [],
@@ -46,6 +57,15 @@ export function useChat<
     useState<ConnectionStatus>('disconnected')
   const [sessionGenerating, setSessionGenerating] = useState(false)
   const [queue, setQueue] = useState<Array<QueuedMessage>>([])
+  const [runId, setRunId] = useState<string | null>(null)
+  const [interruptState, setInterruptState] = useState<
+    ChatInterruptState<TTools>
+  >(() => ({
+    interrupts: EMPTY_INTERRUPTS,
+    pendingInterrupts: EMPTY_INTERRUPTS,
+    interruptErrors: EMPTY_INTERRUPT_ERRORS,
+    resuming: false,
+  }))
 
   type Partial = DeepPartial<InferSchemaType<NonNullable<TSchema>>>
   type Final = InferSchemaType<NonNullable<TSchema>>
@@ -55,10 +75,15 @@ export function useChat<
     options.initialMessages || [],
   )
   const isFirstMountRef = useRef(true)
+  const subscribedRef = useRef(false)
   const activeClientRef = useRef<ChatClient | null>(null)
   const cleanupInvalidationRef = useRef<ReturnType<typeof setTimeout> | null>(
     null,
   )
+  const cleanupDisposalRef = useRef<{
+    client: ChatClient
+    timeout: ReturnType<typeof setTimeout>
+  } | null>(null)
 
   // Update ref synchronously during render so it's always current when useMemo runs.
   messagesRef.current = messages
@@ -67,6 +92,12 @@ export function useChat<
   const optionsRef = useRef<UseChatOptions<TTools, TSchema, TContext>>(options)
   optionsRef.current = options
 
+  const syncResumeState = useCallback((target: ChatClient | null) => {
+    if (!target) return
+    setRunId(target.getCurrentRunId())
+    setInterruptState(target.getInterruptState())
+  }, [])
+
   // Create ChatClient instance with callbacks to sync state
   const client = useMemo(() => {
     const messagesToUse = options.initialMessages || []
@@ -81,10 +112,20 @@ export function useChat<
       ? { connection: initialOptions.connection }
       : { fetcher: initialOptions.fetcher }
 
+    const instanceHolder: {
+      current: ChatClient<TTools, TContext> | undefined
+    } = { current: undefined }
+    const getActiveInstance = () => {
+      const currentInstance = instanceHolder.current
+      if (!currentInstance || activeClientRef.current !== currentInstance) {
+        return undefined
+      }
+      return currentInstance
+    }
+    const pendingInitializationErrors: Array<Error> = []
     const instance = new ChatClient<TTools, TContext>({
       devtoolsBridgeFactory: createChatDevtoolsBridge,
       ...transport,
-      id: clientId,
       initialMessages: messagesToUse,
       ...(initialOptions.body !== undefined && { body: initialOptions.body }),
       ...(initialOptions.threadId !== undefined && {
@@ -96,6 +137,9 @@ export function useChat<
       ...(initialOptions.persistence !== undefined && {
         persistence: initialOptions.persistence,
       }),
+      ...(initialOptions.initialResumeSnapshot !== undefined && {
+        initialResumeSnapshot: initialOptions.initialResumeSnapshot,
+      }),
       ...(initialOptions.context !== undefined && {
         context: initialOptions.context,
       }),
@@ -106,57 +150,64 @@ export function useChat<
         outputKind: initialOptions.outputSchema ? 'structured' : 'chat',
       },
       onResponse: (response) => {
-        if (activeClientRef.current !== instance) return
+        if (!getActiveInstance()) return
         void optionsRef.current.onResponse?.(response)
       },
       onChunk: (chunk: StreamChunk) => {
-        if (activeClientRef.current !== instance) return
+        if (!getActiveInstance()) return
         optionsRef.current.onChunk?.(chunk)
       },
       onFinish: (message: UIMessage<TTools>) => {
-        if (activeClientRef.current !== instance) return
+        if (!getActiveInstance()) return
         optionsRef.current.onFinish?.(message)
       },
       onError: (error: Error) => {
-        if (activeClientRef.current !== instance) return
+        const currentInstance = instanceHolder.current
+        if (!currentInstance) {
+          pendingInitializationErrors.push(error)
+          return
+        }
+        if (activeClientRef.current !== currentInstance) return
         optionsRef.current.onError?.(error)
       },
       ...(initialOptions.tools !== undefined && {
         tools: initialOptions.tools,
       }),
       onCustomEvent: (eventType, data, context) => {
-        if (activeClientRef.current !== instance) return
+        if (!getActiveInstance()) return
         optionsRef.current.onCustomEvent?.(eventType, data, context)
       },
       ...(options.streamProcessor !== undefined && {
         streamProcessor: options.streamProcessor,
       }),
       onMessagesChange: (newMessages: Array<UIMessage<TTools>>) => {
-        if (activeClientRef.current !== instance) return
+        if (!getActiveInstance()) return
         setMessages(newMessages)
       },
       onLoadingChange: (newIsLoading: boolean) => {
-        if (activeClientRef.current !== instance) return
+        const currentInstance = getActiveInstance()
+        if (!currentInstance) return
         setIsLoading(newIsLoading)
+        syncResumeState(currentInstance)
       },
       onErrorChange: (newError: Error | undefined) => {
-        if (activeClientRef.current !== instance) return
+        if (!getActiveInstance()) return
         setError(newError)
       },
       onStatusChange: (status: ChatClientState) => {
-        if (activeClientRef.current !== instance) return
+        if (!getActiveInstance()) return
         setStatus(status)
       },
       onSubscriptionChange: (nextIsSubscribed: boolean) => {
-        if (activeClientRef.current !== instance) return
+        if (!getActiveInstance()) return
         setIsSubscribed(nextIsSubscribed)
       },
       onConnectionStatusChange: (nextStatus: ConnectionStatus) => {
-        if (activeClientRef.current !== instance) return
+        if (!getActiveInstance()) return
         setConnectionStatus(nextStatus)
       },
       onSessionGeneratingChange: (isGenerating: boolean) => {
-        if (activeClientRef.current !== instance) return
+        if (!getActiveInstance()) return
         setSessionGenerating(isGenerating)
       },
       ...(optionsRef.current.queue !== undefined && {
@@ -166,10 +217,32 @@ export function useChat<
         if (activeClientRef.current !== instance) return
         setQueue(nextQueue)
       },
+      onRunIdChange: (nextRunId) => {
+        if (!getActiveInstance()) return
+        setRunId(nextRunId)
+      },
+      onResumeStateChange: (_nextResumeState, nextPendingInterrupts) => {
+        if (!getActiveInstance()) return
+        setInterruptState((current) => ({
+          ...current,
+          interrupts: nextPendingInterrupts,
+          pendingInterrupts: nextPendingInterrupts,
+        }))
+      },
+      onInterruptStateChange: (nextInterruptState) => {
+        if (!getActiveInstance()) return
+        setInterruptState(nextInterruptState)
+        optionsRef.current.onInterruptStateChange?.(nextInterruptState)
+      },
     })
+    instanceHolder.current = instance
     activeClientRef.current = instance
+    for (const error of pendingInitializationErrors) {
+      if (activeClientRef.current !== instance) break
+      optionsRef.current.onError?.(error)
+    }
     return instance
-  }, [clientId])
+  }, [clientId, syncResumeState])
 
   useEffect(() => {
     const clientMessages = client.getMessages()
@@ -211,18 +284,55 @@ export function useChat<
   useEffect(() => {
     if (options.live) {
       client.subscribe()
-    } else {
+      subscribedRef.current = true
+    } else if (subscribedRef.current) {
+      // Only tear down a subscription we actually started. Calling
+      // `unsubscribe()` on initial mount (when `live` was never enabled) would
+      // abort an in-flight delivery resume — `resumeInFlightRun` is kicked off
+      // in the client constructor, and `unsubscribe()` cancels the shared
+      // in-flight stream — so a mid-stream reload would drop its rejoin before
+      // it delivers a single chunk. This is exactly why a reload froze instead
+      // of continuing.
       client.unsubscribe()
+      subscribedRef.current = false
     }
   }, [client, options.live])
 
+  // ONLY THE VIEW ON SCREEN HOLDS A STREAM.
+  //
+  // A page can own many chats — 40 sandboxes, 40 conversations — and a browser
+  // allows only ~6 connections per origin. One long-lived stream per chat reaches
+  // that ceiling after a handful of views, and every request after it QUEUES:
+  // measured, an in-page fetch took over two minutes while the same request from
+  // outside the browser took 17ms. So the connection follows the view.
+  //
+  // Immediate, not deferred: the deferred teardown below can be skipped when the
+  // same client remounts, which is right for disposal but useless for a
+  // connection. `attach` is idempotent and `detach` keeps the transcript and the
+  // resume pointer, so a Strict Mode remount is just detach-then-attach and the
+  // run is picked straight back up from the durable log.
+  useEffect(() => {
+    client.attach()
+    return () => {
+      client.detach()
+    }
+  }, [client])
+
   useEffect(() => {
+    if (cleanupDisposalRef.current?.client === client) {
+      clearTimeout(cleanupDisposalRef.current.timeout)
+      cleanupDisposalRef.current = null
+    }
     if (cleanupInvalidationRef.current) {
       clearTimeout(cleanupInvalidationRef.current)
       cleanupInvalidationRef.current = null
     }
     activeClientRef.current = client
     client.mountDevtools()
+    // Delivery-durability resume is transparent: the resumable SSE connection
+    // adapter re-attaches via the browser's native Last-Event-ID on reconnect.
+    // We only seed interrupt (state) resume from the client here.
+    syncResumeState(client)
 
     return () => {
       cleanupInvalidationRef.current = setTimeout(() => {
@@ -231,46 +341,71 @@ export function useChat<
         }
         cleanupInvalidationRef.current = null
       }, 0)
-      // Subscribe/unsubscribe on `options.live` is owned by the dedicated
-      // effect above. This cleanup only fires on unmount or client swap,
-      // so read `live` through the ref to avoid disposing the client every
-      // time `live` toggles.
-      if (optionsRef.current.live) {
-        client.unsubscribe()
-      } else {
-        client.stop()
+      // Soft cleanup only: do NOT stop/unsubscribe here. React Strict Mode
+      // remounts fire this cleanup then re-attach the same client one tick
+      // later; calling `stop()` would abort a constructor rejoin
+      // (`resumeInFlightRun`) and can wipe the durable resume pointer before
+      // the first chunk. Real teardown lives in the deferred dispose path
+      // below, which only runs when the client is not remounted.
+      // Subscribe/unsubscribe on `options.live` is still owned by the
+      // dedicated effect above for live toggles.
+      const disposal = {
+        client,
+        timeout: setTimeout(() => {
+          if (optionsRef.current.live) {
+            client.unsubscribe()
+          } else {
+            client.stop()
+          }
+          client.dispose()
+          if (cleanupDisposalRef.current === disposal) {
+            cleanupDisposalRef.current = null
+          }
+        }, 0),
       }
-      client.dispose()
+      cleanupDisposalRef.current = disposal
     }
-  }, [client])
+  }, [client, syncResumeState])
 
   const sendMessage = useCallback(
     async (
       content: string | MultimodalContent,
       sendOptions?: SendMessageOptions,
     ) => {
-      await client.sendMessage(content, undefined, sendOptions)
+      try {
+        await client.sendMessage(content, undefined, sendOptions)
+      } finally {
+        syncResumeState(client)
+      }
     },
-    [client],
+    [client, syncResumeState],
   )
 
   const cancelQueued = useCallback(
     (id: string) => {
       client.cancelQueued(id)
     },
-    [client],
+    [client, syncResumeState],
   )
 
   const append = useCallback(
     async (message: ModelMessage | UIMessage) => {
-      await client.append(message)
+      try {
+        await client.append(message)
+      } finally {
+        syncResumeState(client)
+      }
     },
-    [client],
+    [client, syncResumeState],
   )
 
   const reload = useCallback(async () => {
-    await client.reload()
-  }, [client])
+    try {
+      await client.reload()
+    } finally {
+      syncResumeState(client)
+    }
+  }, [client, syncResumeState])
 
   const stop = useCallback(() => {
     client.stop()
@@ -278,7 +413,8 @@ export function useChat<
 
   const clear = useCallback(() => {
     client.clear()
-  }, [client])
+    syncResumeState(client)
+  }, [client, syncResumeState])
 
   const setMessagesManually = useCallback(
     (newMessages: Array<UIMessage<TTools>>) => {
@@ -303,10 +439,50 @@ export function useChat<
   const addToolApprovalResponse = useCallback(
     async (response: { id: string; approved: boolean }) => {
       await client.addToolApprovalResponse(response)
+      syncResumeState(client)
+    },
+    [client, syncResumeState],
+  )
+
+  const resumeInterrupts = useCallback(
+    async (resumeItems: Array<RunAgentResumeItem>, state?: ChatResumeState) => {
+      const result = await client.resumeInterrupts(
+        resumeItems,
+        state ?? undefined,
+      )
+      syncResumeState(client)
+      return result
+    },
+    [client, syncResumeState],
+  )
+
+  const resolveInterrupts = useCallback(
+    (
+      resolution: boolean | ((interrupt: ChatInterrupt<TTools>) => undefined),
+    ) => {
+      if (typeof resolution === 'boolean') {
+        client.resolveInterrupts(resolution)
+      } else {
+        client.resolveInterrupts(resolution)
+      }
     },
     [client],
   )
 
+  const cancelInterrupts = useCallback(() => {
+    client.cancelInterrupts()
+  }, [client])
+
+  const retryInterrupts = useCallback(() => {
+    client.retryInterrupts()
+  }, [client])
+
+  const resumeInterruptsUnsafe = useCallback(
+    (resumeItems: Array<RunAgentResumeItem>, state?: ChatResumeState) =>
+      client.resumeInterruptsUnsafe(resumeItems, state),
+    [client],
+  )
+
   // The "active" structured-output part is the one on the assistant message
   // that follows the latest user message. No such message exists between
   // sendMessage() and the first chunk, so partial/final naturally read as
@@ -355,7 +531,7 @@ export function useChat<
   // The runtime shape unconditionally exposes partial/final; the public
   // return type hides them when no outputSchema was supplied. TS can't
   // structurally narrow across that conditional, so the `as` is the seam.
-  // eslint-disable-next-line no-restricted-syntax -- hook return shape diverges from generic UseChatReturn<TTools, TSchema> due to conditional type on TSchema; TS can't structurally narrow
+  // oxlint-disable-next-line eslint-js/no-restricted-syntax -- hook return shape diverges from generic UseChatReturn<TTools, TSchema> due to conditional type on TSchema; TS can't structurally narrow
   return {
     messages: renderedMessages,
     sendMessage,
@@ -374,6 +550,16 @@ export function useChat<
     addToolApprovalResponse,
     queue,
     cancelQueued,
+    runId,
+    interrupts: interruptState.interrupts,
+    pendingInterrupts: interruptState.pendingInterrupts,
+    interruptErrors: interruptState.interruptErrors,
+    resuming: interruptState.resuming,
+    resolveInterrupts,
+    cancelInterrupts,
+    retryInterrupts,
+    resumeInterruptsUnsafe,
+    resumeInterrupts,
     partial,
     final,
   } as unknown as UseChatReturn<TTools, TSchema>

package/src/use-generate-audio.ts

@@ -1,4 +1,5 @@
 import { useGeneration } from './use-generation'
+import { reconstructAudioResult } from '@tanstack/ai-client'
 import type { AudioGenerationResult, StreamChunk } from '@tanstack/ai'
 import type {
   AIDevtoolsDisplayOptions,
@@ -6,6 +7,7 @@ import type {
   ConnectConnectionAdapter,
   GenerationClientState,
   GenerationFetcher,
+  GenerationPersistenceOptions,
   InferGenerationOutputFromReturn,
 } from '@tanstack/ai-client'
 
@@ -19,12 +21,51 @@ export interface UseGenerateAudioOptions<TOutput = AudioGenerationResult> {
   connection?: ConnectConnectionAdapter
   /** Direct async function for audio generation */
   fetcher?: GenerationFetcher<AudioGenerateInput, AudioGenerationResult>
-  /** Unique identifier for this generation instance */
+  /**
+   * @deprecated Prefer `threadId`. Only allowed when `threadId` is omitted (see `GenerationPersistenceOptions`).
+   */
   id?: string
   /** Additional body parameters to send with connect-based adapter requests */
   body?: Record<string, any>
   /** Display options for TanStack AI Devtools. */
   devtools?: AIDevtoolsDisplayOptions
+  /**
+   * How this generation persists across reloads.
+   * - Omit / `false`: ephemeral, in-memory only.
+   * - `true`: server-driven — on mount the client hydrates the last generation
+   *   for its `threadId` from the server (needs a connection with a
+   *   `hydrateGeneration` handler) and repaints it; it never auto-starts a run.
+   */
+  persistence?: boolean
+  /**
+   * The **scope** this generation belongs to: a stable, app-chosen name for the
+   * slot successive runs fill — not a link to a chat conversation.
+   *
+   * The hook starts empty and produces many runs over its life; each gets its
+   * own `runId`, but all belong to one scope. Persistence keys on this, so
+   * derive it from your own domain and keep it identical across reloads (e.g.
+   * `` `video-${videoId}-start-frame` ``). It is also sent as the AG-UI thread
+   * id on the wire, which the protocol requires.
+   *
+   * **Required whenever `persistence` is set** — an app that cannot name the
+   * scope has nothing to restore to. Optional for ephemeral generations, where
+   * it falls back to `id` purely to satisfy the wire.
+   */
+  threadId?: string
+  /**
+   * Server-driven hydration handler for `persistence: true` when the
+   * connection doesn't carry one (e.g. alongside `fetcher`, or a `stream()` /
+   * `rpcStream()` adapter built without handlers) — typically a one-line
+   * server-function call. The connection's own handler takes precedence.
+   */
+  hydrateGeneration?: ConnectConnectionAdapter['hydrateGeneration']
+  /**
+   * Re-attach handler that replays a run still generating to completion on
+   * mount, when the connection doesn't carry one. Without it, a restored
+   * `running` snapshot surfaces as an (interrupted) error. The connection's
+   * own handler takes precedence.
+   */
+  joinRun?: ConnectConnectionAdapter['joinRun']
   /**
    * Callback when audio is generated. Can optionally return a transformed value.
    *
@@ -61,6 +102,13 @@ export interface UseGenerateAudioReturn<TOutput = AudioGenerationResult> {
   stop: () => void
   /** Clear result, error, and return to idle */
   reset: () => void
+  /**
+   * The id of the generation job currently running, or `null` when nothing is in
+   * flight. Each call to `generate` is one job with its own id. Pass it to your
+   * own endpoint to cancel or poll the provider job — `stop()` only aborts the
+   * local stream, it does not stop work already running on the provider.
+   */
+  runId: string | null
 }
 
 /**
@@ -94,9 +142,12 @@ export interface UseGenerateAudioReturn<TOutput = AudioGenerationResult> {
  * ```
  */
 export function useGenerateAudio<TTransformed = void>(
-  options: Omit<UseGenerateAudioOptions, 'onResult'> & {
+  options: Omit<
+    UseGenerateAudioOptions,
+    'onResult' | 'persistence' | 'threadId' | 'id'
+  > & {
     onResult?: (result: AudioGenerationResult) => TTransformed
-  },
+  } & GenerationPersistenceOptions,
 ): UseGenerateAudioReturn<
   InferGenerationOutputFromReturn<AudioGenerationResult, TTransformed>
 > {
@@ -106,19 +157,15 @@ export function useGenerateAudio<TTransformed = void>(
     hookName: 'useGenerateAudio',
     outputKind: 'audio' as const,
   }
-  const { generate, result, isLoading, error, status, stop, reset } =
-    useGeneration<AudioGenerateInput, AudioGenerationResult, TTransformed>({
-      ...options,
-      devtools,
-    })
+  const generation = useGeneration<
+    AudioGenerateInput,
+    AudioGenerationResult,
+    TTransformed
+  >({
+    ...options,
+    devtools,
+    reconstructResult: reconstructAudioResult,
+  })
 
-  return {
-    generate: generate as (input: AudioGenerateInput) => Promise<void>,
-    result,
-    isLoading,
-    error,
-    status,
-    stop,
-    reset,
-  }
+  return generation
 }