@tanstack/ai-client 0.14.1 → 0.15.0

package/src/chat-client.ts

@@ -2,7 +2,9 @@ import {
   StreamProcessor,
   convertSchemaToJsonSchema,
   generateMessageId,
+  isStandardSchema,
   normalizeToUIMessage,
+  parseWithStandardSchema,
 } from '@tanstack/ai/client'
 import { createNoOpChatDevtoolsBridge } from './devtools-noop'
 import {
@@ -39,6 +41,37 @@ import type {
   UIMessage,
 } from './types'
 
+type ChatClientUpdateOptionsWithoutContext<
+  TTools extends ReadonlyArray<AnyClientTool>,
+> = {
+  connection?: ConnectionAdapter
+  fetcher?: ChatFetcher
+  /** @deprecated Use `forwardedProps` instead. */
+  body?: Record<string, any>
+  forwardedProps?: Record<string, any>
+  tools?: TTools
+  onResponse?: (response?: Response) => void | Promise<void>
+  onChunk?: (chunk: StreamChunk) => void
+  onFinish?: (message: UIMessage) => void
+  onError?: (error: Error) => void
+  onSubscriptionChange?: (isSubscribed: boolean) => void
+  onConnectionStatusChange?: (status: ConnectionStatus) => void
+  onSessionGeneratingChange?: (isGenerating: boolean) => void
+  onCustomEvent?: (
+    eventType: string,
+    data: unknown,
+    context: { toolCallId?: string },
+  ) => void
+}
+
+type ClientToolResult = {
+  toolCallId: string
+  tool: string
+  output: any
+  state?: 'output-available' | 'output-error'
+  errorText?: string
+}
+
 function resolveTransport(transport: {
   connection?: ConnectionAdapter
   fetcher?: ChatFetcher
@@ -54,7 +87,10 @@ function resolveTransport(transport: {
   throw new Error('ChatClient: either `connection` or `fetcher` is required.')
 }
 
-export class ChatClient {
+export class ChatClient<
+  TTools extends ReadonlyArray<AnyClientTool> = any,
+  TContext = unknown,
+> {
   private readonly processor: StreamProcessor
   private connection: SubscribeConnectionAdapter
   private readonly uniqueId: string
@@ -65,6 +101,7 @@ export class ChatClient {
   // merged on every send, with `forwardedProps` winning on key collision.
   private bodyOption: Record<string, any> = {}
   private forwardedPropsOption: Record<string, any> = {}
+  private context: TContext | undefined = undefined
   private pendingMessageBody: Record<string, any> | undefined = undefined
   private isLoading = false
   private isSubscribed = false
@@ -86,6 +123,8 @@ export class ChatClient {
   private readonly postStreamActions: Array<() => Promise<void>> = []
   // Track pending client tool executions to await them before stream finalization
   private readonly pendingToolExecutions: Map<string, Promise<void>> = new Map()
+  private activeClientTools: Map<string, AnyClientTool> | null = null
+  private activeContext: TContext | undefined = undefined
   // Flag to deduplicate continuation checks during action draining
   private continuationPending = false
   private subscriptionAbortController: AbortController | null = null
@@ -121,7 +160,7 @@ export class ChatClient {
     }
   }
 
-  constructor(options: ChatClientOptions) {
+  constructor(options: ChatClientOptions<TTools, TContext>) {
     this.uniqueId = options.id || this.generateUniqueId('chat')
     this.threadId = options.threadId || this.generateUniqueId('thread')
     // Both `body` (deprecated) and `forwardedProps` populate the AG-UI
@@ -131,6 +170,7 @@ export class ChatClient {
     // winning on key collision.
     this.bodyOption = options.body || {}
     this.forwardedPropsOption = options.forwardedProps || {}
+    this.context = options.context
     this.connection = normalizeConnectionAdapter(resolveTransport(options))
 
     // Build client tools map
@@ -288,7 +328,9 @@ export class ChatClient {
           input: any
         }) => {
           // Handle client-side tool execution automatically
-          const clientTool = this.clientToolsRef.current.get(args.toolName)
+          const clientTools =
+            this.activeClientTools ?? this.clientToolsRef.current
+          const clientTool = clientTools.get(args.toolName)
           const executeFunc = clientTool?.execute
           if (executeFunc) {
             // Capture the run context at execution-start so a tool whose
@@ -300,18 +342,27 @@ export class ChatClient {
             // Create and track the execution promise
             const executionPromise = (async () => {
               try {
-                const output = await executeFunc(args.input)
-                await this.addToolResultInternal(
+                const context =
+                  this.activeClientTools === null
+                    ? this.context
+                    : this.activeContext
+                const output = await executeFunc(args.input, {
+                  toolCallId: args.toolCallId,
+                  context: context as TContext,
+                  emitCustomEvent: () => {},
+                })
+                await this.addToolResultForClientTool(
                   {
                     toolCallId: args.toolCallId,
                     tool: args.toolName,
                     output,
                     state: 'output-available',
                   },
+                  clientTool,
                   runEventContext,
                 )
               } catch (error: any) {
-                await this.addToolResultInternal(
+                await this.addToolResultForClientTool(
                   {
                     toolCallId: args.toolCallId,
                     tool: args.toolName,
@@ -319,6 +370,7 @@ export class ChatClient {
                     state: 'output-error',
                     errorText: error.message,
                   },
+                  clientTool,
                   runEventContext,
                 )
               } finally {
@@ -761,6 +813,8 @@ export class ChatClient {
     try {
       // Get UIMessages with parts (preserves approval state and client tool results)
       const messages = this.processor.getMessages()
+      const clientTools = new Map(this.clientToolsRef.current)
+      const runtimeContext = this.context
 
       // Call onResponse callback
       await this.callbacksRef.current.onResponse()
@@ -796,6 +850,8 @@ export class ChatClient {
       this.currentStreamId = this.generateUniqueId('stream')
       this.devtoolsBridge.setCurrentStreamId(this.currentStreamId)
       this.currentMessageId = null
+      this.activeClientTools = clientTools
+      this.activeContext = runtimeContext
 
       // Reset processor stream state for new response — prevents stale
       // messageStates entries (from a previous stream) from blocking
@@ -819,15 +875,13 @@ export class ChatClient {
       const runContext = {
         threadId: this.threadId,
         runId: `run-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`,
-        clientTools: Array.from(this.clientToolsRef.current.values()).map(
-          (t) => ({
-            name: t.name,
-            description: t.description,
-            parameters: t.inputSchema
-              ? convertSchemaToJsonSchema(t.inputSchema)
-              : { type: 'object' },
-          }),
-        ),
+        clientTools: Array.from(clientTools.values()).map((t) => ({
+          name: t.name,
+          description: t.description,
+          parameters: t.inputSchema
+            ? convertSchemaToJsonSchema(t.inputSchema)
+            : { type: 'object' },
+        })),
         forwardedProps: { ...mergedBody },
       }
       this.devtoolsBridge.beginRun(runContext.runId, this.threadId)
@@ -913,6 +967,8 @@ export class ChatClient {
         this.currentStreamId = null
         this.devtoolsBridge.setCurrentStreamId(null)
         this.currentMessageId = null
+        this.activeClientTools = null
+        this.activeContext = undefined
         this.abortController = null
         this.setIsLoading(false)
         this.pendingMessageBody = undefined // Ensure it's cleared even on error
@@ -1044,26 +1100,32 @@ export class ChatClient {
   /**
    * Add the result of a client-side tool execution
    */
-  async addToolResult(result: {
-    toolCallId: string
-    tool: string
-    output: any
-    state?: 'output-available' | 'output-error'
-    errorText?: string
-  }): Promise<void> {
-    await this.addToolResultInternal(result)
+  async addToolResult(result: ClientToolResult): Promise<void> {
+    const clientTool = this.clientToolsRef.current.get(result.tool)
+    await this.addToolResultForClientTool(result, clientTool)
   }
 
-  private async addToolResultInternal(
-    result: {
-      toolCallId: string
-      tool: string
-      output: any
-      state?: 'output-available' | 'output-error'
-      errorText?: string
-    },
+  private async addToolResultForClientTool(
+    result: ClientToolResult,
+    clientTool: AnyClientTool | undefined,
     context?: ChatClientRunEventContext,
   ): Promise<void> {
+    if (clientTool && result.state !== 'output-error') {
+      try {
+        result = {
+          ...result,
+          output: this.validateClientToolOutput(clientTool, result.output),
+        }
+      } catch (error: any) {
+        result = {
+          ...result,
+          output: null,
+          state: 'output-error',
+          errorText: error.message,
+        }
+      }
+    }
+
     this.events.toolResultAdded(
       result.toolCallId,
       result.tool,
@@ -1088,6 +1150,17 @@ export class ChatClient {
     await this.checkForContinuation()
   }
 
+  private validateClientToolOutput(
+    clientTool: AnyClientTool,
+    output: any,
+  ): any {
+    if (clientTool.outputSchema && isStandardSchema(clientTool.outputSchema)) {
+      return parseWithStandardSchema(clientTool.outputSchema, output)
+    }
+
+    return output
+  }
+
   /**
    * Respond to a tool approval request
    */
@@ -1206,8 +1279,8 @@ export class ChatClient {
   /**
    * Get current messages
    */
-  getMessages(): Array<UIMessage> {
-    return this.processor.getMessages()
+  getMessages(): Array<UIMessage<TTools>> {
+    return this.processor.getMessages() as Array<UIMessage<TTools>>
   }
 
   /**
@@ -1258,7 +1331,7 @@ export class ChatClient {
   /**
    * Manually set messages
    */
-  setMessagesManually(messages: Array<UIMessage>): void {
+  setMessagesManually(messages: Array<UIMessage<TTools>>): void {
     this.processor.setMessages(messages)
     this.devtoolsBridge.emitSnapshot()
   }
@@ -1266,26 +1339,16 @@ export class ChatClient {
   /**
    * Update options refs (for use in React hooks to avoid recreating client)
    */
-  updateOptions(options: {
-    connection?: ConnectionAdapter
-    fetcher?: ChatFetcher
-    /** @deprecated Use `forwardedProps` instead. */
-    body?: Record<string, any>
-    forwardedProps?: Record<string, any>
-    tools?: ReadonlyArray<AnyClientTool>
-    onResponse?: (response?: Response) => void | Promise<void>
-    onChunk?: (chunk: StreamChunk) => void
-    onFinish?: (message: UIMessage) => void
-    onError?: (error: Error) => void
-    onSubscriptionChange?: (isSubscribed: boolean) => void
-    onConnectionStatusChange?: (status: ConnectionStatus) => void
-    onSessionGeneratingChange?: (isGenerating: boolean) => void
-    onCustomEvent?: (
-      eventType: string,
-      data: unknown,
-      context: { toolCallId?: string },
-    ) => void
-  }): void {
+  updateOptions(options: ChatClientUpdateOptionsWithoutContext<TTools>): void
+  updateOptions(
+    options: ChatClientUpdateOptionsWithoutContext<TTools> &
+      Pick<ChatClientOptions<TTools, TContext>, 'context'>,
+  ): void
+  updateOptions(
+    options: ChatClientUpdateOptionsWithoutContext<TTools> & {
+      context?: TContext | undefined
+    },
+  ): void {
     if (options.connection !== undefined || options.fetcher !== undefined) {
       const wasSubscribed = this.isSubscribed
 
@@ -1312,15 +1375,19 @@ export class ChatClient {
         this.subscribe()
       }
     }
-    // Replace each slot independently so callers can update one without
-    // wiping the other. (Passing `undefined` for either field is a "leave
-    // unchanged" signal - to clear a slot, pass an empty object `{}`.)
+    // Replace each wire-payload slot independently so callers can update one
+    // without wiping the other. Passing `undefined` for `body` or
+    // `forwardedProps` leaves that slot unchanged; context is cleared when the
+    // key is present with an `undefined` value.
     if (options.body !== undefined) {
       this.bodyOption = options.body
     }
     if (options.forwardedProps !== undefined) {
       this.forwardedPropsOption = options.forwardedProps
     }
+    if ('context' in options) {
+      this.context = options.context
+    }
     if (options.tools !== undefined) {
       this.clientToolsRef.current = new Map()
       for (const tool of options.tools) {

package/src/devtools.ts

@@ -1219,7 +1219,9 @@ function hydrateToolCallOutputs(
         candidate.output === undefined,
     )
     if (toolCall) {
-      toolCall.output = parseFixtureResultContent(part.content)
+      toolCall.output = Array.isArray(part.content)
+        ? part.content
+        : parseFixtureResultContent(part.content)
     }
   }
 

package/src/index.ts

@@ -13,8 +13,10 @@ export type {
   StructuredOutputPart,
   // Client configuration types
   ChatClientOptions,
+  ClientContextOptionFromTools,
   ChatRequestBody,
   InferChatMessages,
+  InferredClientContext,
   ChatClientState,
   ConnectionStatus,
   ChatFetcher,

package/src/types.ts

@@ -220,7 +220,7 @@ export type ToolCallPart<TTools extends ReadonlyArray<AnyClientTool> = any> =
 export interface ToolResultPart {
   type: 'tool-result'
   toolCallId: string
-  content: string
+  content: string | Array<ContentPart>
   state: ToolResultState
   error?: string // Error message if state is "error"
 }
@@ -266,14 +266,91 @@ export interface UIMessage<
   createdAt?: Date
 }
 
+type IsUnknown<T> = unknown extends T
+  ? [T] extends [unknown]
+    ? true
+    : false
+  : false
+
+type KnownContext<T> = IsUnknown<T> extends true ? never : T
+
+type MergeContext<TLeft, TRight> = [TLeft] extends [never]
+  ? TRight
+  : [TRight] extends [never]
+    ? TLeft
+    : TLeft & TRight
+
+type UnionToIntersection<T> = [T] extends [never]
+  ? never
+  : (T extends unknown ? (value: T) => void : never) extends (
+        value: infer TIntersection,
+      ) => void
+    ? TIntersection
+    : never
+
+type DefinedContext<T> = Exclude<T, undefined>
+
+type ContextFromExecute<T> = T extends (...args: any) => any
+  ? NonNullable<Parameters<T>[1]> extends { context: infer TContext }
+    ? KnownContext<TContext>
+    : never
+  : never
+
+type ContextFromClientTool<T> = T extends AnyClientTool
+  ? T extends { execute?: infer TExecute }
+    ? ContextFromExecute<TExecute>
+    : never
+  : never
+
+type RequiredContextFromClientToolUnion<T> = T extends unknown
+  ? undefined extends ContextFromClientTool<T>
+    ? never
+    : ContextFromClientTool<T>
+  : never
+
+type ContextFromClientToolUnion<T> = [
+  UnionToIntersection<DefinedContext<ContextFromClientTool<T>>>,
+] extends [never]
+  ? never
+  : [RequiredContextFromClientToolUnion<T>] extends [never]
+    ? UnionToIntersection<DefinedContext<ContextFromClientTool<T>>> | undefined
+    : UnionToIntersection<DefinedContext<ContextFromClientTool<T>>>
+
+type ContextFromClientTools<TTools> =
+  IsUnknown<TTools> extends true
+    ? never
+    : TTools extends readonly [infer THead, ...infer TTail]
+      ? MergeContext<
+          ContextFromClientTool<THead>,
+          ContextFromClientTools<TTail>
+        >
+      : TTools extends ReadonlyArray<infer TItem>
+        ? ContextFromClientToolUnion<TItem>
+        : never
+
+export type InferredClientContext<TTools> = [
+  ContextFromClientTools<TTools>,
+] extends [never]
+  ? unknown
+  : ContextFromClientTools<TTools>
+
+export type ClientContextOptionFromTools<TTools, TContext> = [
+  ContextFromClientTools<TTools>,
+] extends [never]
+  ? { context?: TContext }
+  : undefined extends ContextFromClientTools<TTools>
+    ? { context?: TContext & ContextFromClientTools<TTools> }
+    : { context: TContext & ContextFromClientTools<TTools> }
+
 /**
- * Options for `ChatClient`. Exactly one of `connection` or `fetcher` must be
- * provided — the type-level XOR is enforced via `ChatTransport`.
+ * Base options for `ChatClient`, excluding the transport (`connection` or
+ * `fetcher`) which is supplied separately via `ChatTransport` so the XOR
+ * is preserved when composing the final `ChatClientOptions` type.
  */
-export type ChatClientOptions<
+export interface ChatClientBaseOptions<
   TTools extends ReadonlyArray<AnyClientTool> = any,
   TContext = unknown,
-> = {
+> {
   /**
    * Initial messages to populate the chat
    */
@@ -311,6 +388,14 @@ export type ChatClientOptions<
    */
   body?: Record<string, any>
 
+  /**
+   * Client-local runtime context passed to client tool implementations.
+   *
+   * This value is not serialized to the server. Use `forwardedProps` for
+   * explicit client-to-server handoff of serializable values.
+   */
+  context?: TContext
+
   /**
    * Callback when a response is received
    */
@@ -390,11 +475,6 @@ export type ChatClientOptions<
    */
   tools?: TTools
 
-  /**
-   * Client-local context passed to client-side tool execution.
-   */
-  context?: TContext
-
   /**
    * Devtools hook metadata for this client instance.
    */
@@ -420,7 +500,18 @@ export type ChatClientOptions<
      */
     chunkStrategy?: ChunkStrategy
   }
-} & ChatTransport
+}
+
+/**
+ * Options for `ChatClient`. Exactly one of `connection` or `fetcher` must be
+ * provided — the type-level XOR is enforced via `ChatTransport`.
+ */
+export type ChatClientOptions<
+  TTools extends ReadonlyArray<AnyClientTool> = any,
+  TContext = InferredClientContext<TTools>,
+> = DistributedOmit<ChatClientBaseOptions<TTools, TContext>, 'context'> &
+  ClientContextOptionFromTools<TTools, TContext> &
+  ChatTransport
 
 export interface ChatRequestBody {
   messages: Array<ModelMessage>
@@ -466,7 +557,7 @@ export function clientTools<const T extends Array<AnyClientTool>>(
  */
 export function createChatClientOptions<
   const TTools extends ReadonlyArray<AnyClientTool>,
-  TContext = unknown,
+  TContext = InferredClientContext<TTools>,
 >(
   options: ChatClientOptions<TTools, TContext>,
 ): ChatClientOptions<TTools, TContext> {