@tanstack/ai-client 0.3.0 → 0.4.1

package/src/chat-client.ts

@@ -4,29 +4,42 @@ import {
   normalizeToUIMessage,
 } from '@tanstack/ai'
 import { DefaultChatClientEventEmitter } from './events'
+import type {
+  AnyClientTool,
+  ContentPart,
+  ModelMessage,
+  StreamChunk,
+} from '@tanstack/ai'
+import type { ConnectionAdapter } from './connection-adapters'
+import type { ChatClientEventEmitter } from './events'
 import type {
   ChatClientOptions,
+  ChatClientState,
   MessagePart,
+  MultimodalContent,
   ToolCallPart,
   UIMessage,
 } from './types'
-import type { AnyClientTool, ModelMessage, StreamChunk } from '@tanstack/ai'
-import type { ConnectionAdapter } from './connection-adapters'
-import type { ChatClientEventEmitter } from './events'
 
 export class ChatClient {
   private processor: StreamProcessor
   private connection: ConnectionAdapter
   private uniqueId: string
   private body: Record<string, any> = {}
+  private pendingMessageBody: Record<string, any> | undefined = undefined
   private isLoading = false
   private error: Error | undefined = undefined
+  private status: ChatClientState = 'ready'
   private abortController: AbortController | null = null
   private events: ChatClientEventEmitter
   private clientToolsRef: { current: Map<string, AnyClientTool> }
   private currentStreamId: string | null = null
   private currentMessageId: string | null = null
   private postStreamActions: Array<() => Promise<void>> = []
+  // Track pending client tool executions to await them before stream finalization
+  private pendingToolExecutions: Map<string, Promise<void>> = new Map()
+  // Flag to deduplicate continuation checks during action draining
+  private continuationPending = false
 
   private callbacksRef: {
     current: {
@@ -37,6 +50,7 @@ export class ChatClient {
       onMessagesChange: (messages: Array<UIMessage>) => void
       onLoadingChange: (isLoading: boolean) => void
       onErrorChange: (error: Error | undefined) => void
+      onStatusChange: (status: ChatClientState) => void
     }
   }
 
@@ -63,6 +77,7 @@ export class ChatClient {
         onMessagesChange: options.onMessagesChange || (() => {}),
         onLoadingChange: options.onLoadingChange || (() => {}),
         onErrorChange: options.onErrorChange || (() => {}),
+        onStatusChange: options.onStatusChange || (() => {}),
       },
     }
 
@@ -75,13 +90,15 @@ export class ChatClient {
           this.callbacksRef.current.onMessagesChange(messages)
         },
         onStreamStart: () => {
-          // Stream started
+          this.setStatus('streaming')
         },
         onStreamEnd: (message: UIMessage) => {
           this.callbacksRef.current.onFinish(message)
+          this.setStatus('ready')
         },
         onError: (error: Error) => {
           this.setError(error)
+          this.setStatus('error')
           this.callbacksRef.current.onError(error)
         },
         onTextUpdate: (messageId: string, content: string) => {
@@ -127,31 +144,41 @@ export class ChatClient {
             )
           }
         },
-        onToolCall: async (args: {
+        onToolCall: (args: {
           toolCallId: string
           toolName: string
           input: any
         }) => {
           // Handle client-side tool execution automatically
           const clientTool = this.clientToolsRef.current.get(args.toolName)
-          if (clientTool?.execute) {
-            try {
-              const output = await clientTool.execute(args.input)
-              await this.addToolResult({
-                toolCallId: args.toolCallId,
-                tool: args.toolName,
-                output,
-                state: 'output-available',
-              })
-            } catch (error: any) {
-              await this.addToolResult({
-                toolCallId: args.toolCallId,
-                tool: args.toolName,
-                output: null,
-                state: 'output-error',
-                errorText: error.message,
-              })
-            }
+          const executeFunc = clientTool?.execute
+          if (executeFunc) {
+            // Create and track the execution promise
+            const executionPromise = (async () => {
+              try {
+                const output = await executeFunc(args.input)
+                await this.addToolResult({
+                  toolCallId: args.toolCallId,
+                  tool: args.toolName,
+                  output,
+                  state: 'output-available',
+                })
+              } catch (error: any) {
+                await this.addToolResult({
+                  toolCallId: args.toolCallId,
+                  tool: args.toolName,
+                  output: null,
+                  state: 'output-error',
+                  errorText: error.message,
+                })
+              } finally {
+                // Remove from pending when complete
+                this.pendingToolExecutions.delete(args.toolCallId)
+              }
+            })()
+
+            // Track the pending execution
+            this.pendingToolExecutions.set(args.toolCallId, executionPromise)
           }
         },
         onApprovalRequest: (args: {
@@ -187,6 +214,11 @@ export class ChatClient {
     this.events.loadingChanged(isLoading)
   }
 
+  private setStatus(status: ChatClientState): void {
+    this.status = status
+    this.callbacksRef.current.onStatusChange(status)
+  }
+
   private setError(error: Error | undefined): void {
     this.error = error
     this.callbacksRef.current.onErrorChange(error)
@@ -227,6 +259,12 @@ export class ChatClient {
       await new Promise((resolve) => setTimeout(resolve, 0))
     }
 
+    // Wait for all pending tool executions to complete before finalizing
+    // This ensures client tools finish before we check for continuation
+    if (this.pendingToolExecutions.size > 0) {
+      await Promise.all(this.pendingToolExecutions.values())
+    }
+
     // Finalize the stream
     this.processor.finalizeStream()
 
@@ -251,20 +289,82 @@ export class ChatClient {
   }
 
   /**
-   * Send a message and stream the response
+   * Send a message and stream the response.
+   * Supports both simple string content and multimodal content (images, audio, video, documents).
+   *
+   * @param content - The message content. Can be:
+   *   - A simple string for text-only messages
+   *   - 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.
+   *
+   * @example
+   * ```ts
+   * // Simple text message
+   * await client.sendMessage('Hello!')
+   *
+   * // Text message with custom body params
+   * await client.sendMessage('Hello!', { temperature: 0.7 })
+   *
+   * // Multimodal message with image
+   * await client.sendMessage({
+   *   content: [
+   *     { type: 'text', content: 'What is in this image?' },
+   *     { type: 'image', source: { type: 'url', value: 'https://example.com/photo.jpg' } }
+   *   ]
+   * })
+   *
+   * // Multimodal message with custom ID and body params
+   * await client.sendMessage(
+   *   {
+   *     content: [
+   *       { type: 'text', content: 'Describe this audio' },
+   *       { type: 'audio', source: { type: 'data', value: 'base64...' } }
+   *     ],
+   *     id: 'custom-message-id'
+   *   },
+   *   { model: 'gpt-4-audio' }
+   * )
+   * ```
    */
-  async sendMessage(content: string): Promise<void> {
-    if (!content.trim() || this.isLoading) {
+  async sendMessage(
+    content: string | MultimodalContent,
+    body?: Record<string, any>,
+  ): Promise<void> {
+    const emptyMessage = typeof content === 'string' && !content.trim()
+    if (emptyMessage || this.isLoading) {
       return
     }
+    // Normalize input to extract content, id, and validate
+    const normalizedContent = this.normalizeMessageInput(content)
+
+    // Store the per-message body for use in streamResponse
+    this.pendingMessageBody = body
 
     // Add user message via processor
-    const userMessage = this.processor.addUserMessage(content.trim())
-    this.events.messageSent(userMessage.id, content.trim())
+    const userMessage = this.processor.addUserMessage(
+      normalizedContent.content,
+      normalizedContent.id,
+    )
+    this.events.messageSent(userMessage.id, normalizedContent.content)
 
     await this.streamResponse()
   }
 
+  /**
+   * Normalize the message input to extract content and optional id.
+   * Trims string content automatically.
+   */
+  private normalizeMessageInput(input: string | MultimodalContent): {
+    content: string | Array<ContentPart>
+    id?: string
+  } {
+    if (typeof input === 'string') {
+      return { content: input.trim() }
+    }
+    return { content: input.content, id: input.id }
+  }
+
   /**
    * Append a message and stream the response
    */
@@ -287,6 +387,12 @@ export class ChatClient {
     const messages = this.processor.getMessages()
     this.processor.setMessages([...messages, uiMessage])
 
+    // If stream is in progress, queue the response for after it ends
+    if (this.isLoading) {
+      this.queuePostStreamAction(() => this.streamResponse())
+      return
+    }
+
     await this.streamResponse()
   }
 
@@ -294,45 +400,76 @@ export class ChatClient {
    * Stream a response from the LLM
    */
   private async streamResponse(): Promise<void> {
+    // Guard against concurrent streams - if already loading, skip
+    if (this.isLoading) {
+      return
+    }
+
     this.setIsLoading(true)
+    this.setStatus('submitted')
     this.setError(undefined)
     this.abortController = new AbortController()
+    // Reset pending tool executions for the new stream
+    this.pendingToolExecutions.clear()
+    let streamCompletedSuccessfully = false
 
     try {
-      // Get model messages for the LLM
-      const modelMessages = this.processor.toModelMessages()
+      // Get UIMessages with parts (preserves approval state and client tool results)
+      const messages = this.processor.getMessages()
 
       // Call onResponse callback
       await this.callbacksRef.current.onResponse()
 
-      // Include conversationId in the body for server-side event correlation
-      const bodyWithConversationId = {
+      // Merge body: base body + per-message body (per-message takes priority)
+      // Include conversationId for server-side event correlation
+      const mergedBody = {
         ...this.body,
+        ...this.pendingMessageBody,
         conversationId: this.uniqueId,
       }
 
+      // Clear the pending message body after use
+      this.pendingMessageBody = undefined
+
       // Connect and stream
       const stream = this.connection.connect(
-        modelMessages,
-        bodyWithConversationId,
+        messages,
+        mergedBody,
         this.abortController.signal,
       )
 
       await this.processStream(stream)
+      streamCompletedSuccessfully = true
     } catch (err) {
       if (err instanceof Error) {
         if (err.name === 'AbortError') {
           return
         }
         this.setError(err)
+        this.setStatus('error')
         this.callbacksRef.current.onError(err)
       }
     } finally {
       this.abortController = null
       this.setIsLoading(false)
+      this.pendingMessageBody = undefined // Ensure it's cleared even on error
 
       // 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)
+      if (streamCompletedSuccessfully) {
+        const messages = this.processor.getMessages()
+        const lastPart = messages.at(-1)?.parts.at(-1)
+
+        if (lastPart?.type === 'tool-result' && this.shouldAutoSend()) {
+          try {
+            await this.checkForContinuation()
+          } catch (error) {
+            console.error('Failed to continue flow after tool result:', error)
+          }
+        }
+      }
     }
   }
 
@@ -368,6 +505,7 @@ export class ChatClient {
       this.abortController = null
     }
     this.setIsLoading(false)
+    this.setStatus('ready')
     this.events.stopped()
   }
 
@@ -476,8 +614,18 @@ export class ChatClient {
    * Check if we should continue the flow and do so if needed
    */
   private async checkForContinuation(): Promise<void> {
+    // Prevent duplicate continuation attempts
+    if (this.continuationPending || this.isLoading) {
+      return
+    }
+
     if (this.shouldAutoSend()) {
-      await this.streamResponse()
+      this.continuationPending = true
+      try {
+        await this.streamResponse()
+      } finally {
+        this.continuationPending = false
+      }
     }
   }
 
@@ -502,6 +650,13 @@ export class ChatClient {
     return this.isLoading
   }
 
+  /**
+   * Get current status
+   */
+  getStatus(): ChatClientState {
+    return this.status
+  }
+
   /**
    * Get current error
    */

package/src/connection-adapters.ts

@@ -1,4 +1,3 @@
-import { convertMessagesToModelMessages } from '@tanstack/ai'
 import type { ModelMessage, StreamChunk, UIMessage } from '@tanstack/ai'
 
 /**
@@ -138,16 +137,15 @@ export function fetchServerSentEvents(
       const resolvedOptions =
         typeof options === 'function' ? await options() : options
 
-      const modelMessages = convertMessagesToModelMessages(messages)
-
       const requestHeaders: Record<string, string> = {
         'Content-Type': 'application/json',
         ...mergeHeaders(resolvedOptions.headers),
       }
 
-      // Merge body from options with messages and data
+      // Send messages as-is (UIMessages with parts preserved)
+      // Server-side TextEngine handles conversion to ModelMessages
       const requestBody = {
-        messages: modelMessages,
+        messages,
         data,
         ...resolvedOptions.body,
       }
@@ -238,17 +236,15 @@ export function fetchHttpStream(
       const resolvedOptions =
         typeof options === 'function' ? await options() : options
 
-      // Convert UIMessages to ModelMessages if needed
-      const modelMessages = convertMessagesToModelMessages(messages)
-
       const requestHeaders: Record<string, string> = {
         'Content-Type': 'application/json',
         ...mergeHeaders(resolvedOptions.headers),
       }
 
-      // Merge body from options with messages and data
+      // Send messages as-is (UIMessages with parts preserved)
+      // Server-side TextEngine handles conversion to ModelMessages
       const requestBody = {
-        messages: modelMessages,
+        messages,
         data,
         ...resolvedOptions.body,
       }
@@ -302,14 +298,15 @@ export function fetchHttpStream(
  */
 export function stream(
   streamFactory: (
-    messages: Array<ModelMessage>,
+    messages: Array<UIMessage> | Array<ModelMessage>,
     data?: Record<string, any>,
   ) => AsyncIterable<StreamChunk>,
 ): ConnectionAdapter {
   return {
     async *connect(messages, data) {
-      const modelMessages = convertMessagesToModelMessages(messages)
-      yield* streamFactory(modelMessages, data)
+      // Pass messages as-is (UIMessages with parts preserved)
+      // Server-side chat() handles conversion to ModelMessages
+      yield* streamFactory(messages, data)
     },
   }
 }
@@ -332,16 +329,15 @@ export function stream(
  */
 export function rpcStream(
   rpcCall: (
-    messages: Array<ModelMessage>,
+    messages: Array<UIMessage> | Array<ModelMessage>,
     data?: Record<string, any>,
   ) => AsyncIterable<StreamChunk>,
 ): ConnectionAdapter {
   return {
     async *connect(messages, data) {
-      const modelMessages = convertMessagesToModelMessages(messages)
-      // Simply yield from the RPC call
-      // The RPC layer handles WebSocket transport
-      yield* rpcCall(modelMessages, data)
+      // Pass messages as-is (UIMessages with parts preserved)
+      // Server-side chat() handles conversion to ModelMessages
+      yield* rpcCall(messages, data)
     },
   }
 }

package/src/events.ts

@@ -1,4 +1,5 @@
 import { aiEventClient } from '@tanstack/ai/event-client'
+import type { ContentPart } from '@tanstack/ai'
 import type { UIMessage } from './types'
 
 /**
@@ -137,19 +138,36 @@ export abstract class ChatClientEventEmitter {
   }
 
   /**
-   * Emit message sent event
+   * Emit message sent event.
+   * Supports both simple string content and multimodal content arrays.
+   *
+   * @param messageId - The ID of the sent message
+   * @param content - The message content (string or array of ContentPart for multimodal)
    */
-  messageSent(messageId: string, content: string): void {
+  messageSent(messageId: string, content: string | Array<ContentPart>): void {
+    // For text content, extract it; for multimodal, provide the array
+    const textContent =
+      typeof content === 'string'
+        ? content
+        : content
+            .filter((part) => part.type === 'text')
+            .map((part) => (part as { type: 'text'; content: string }).content)
+            .join('')
+
     this.emitEvent('text:message:created', {
       messageId,
       role: 'user',
-      content,
+      content: textContent,
+      // Include full content for multimodal messages
+      ...(Array.isArray(content) && { parts: content }),
     })
 
     this.emitEvent('text:message:user', {
       messageId,
       role: 'user',
-      content,
+      content: textContent,
+      // Include full content for multimodal messages
+      ...(Array.isArray(content) && { parts: content }),
     })
   }
 
@@ -161,7 +179,6 @@ export abstract class ChatClientEventEmitter {
       fromMessageIndex,
     })
   }
-
   /**
    * Emit stopped event
    */

package/src/index.ts

@@ -11,6 +11,9 @@ export type {
   ChatClientOptions,
   ChatRequestBody,
   InferChatMessages,
+  ChatClientState,
+  // Multimodal content input type
+  MultimodalContent,
 } from './types'
 export { clientTools, createChatClientOptions } from './types'
 export type {

package/src/types.ts

@@ -1,10 +1,15 @@
 import type {
   AnyClientTool,
+  AudioPart,
   ChunkStrategy,
+  ContentPart,
+  DocumentPart,
+  ImagePart,
   InferToolInput,
   InferToolOutput,
   ModelMessage,
   StreamChunk,
+  VideoPart,
 } from '@tanstack/ai'
 import type { ConnectionAdapter } from './connection-adapters'
 
@@ -26,6 +31,40 @@ export type ToolResultState =
   | 'complete' // Result is complete
   | 'error' // Error occurred
 
+/**
+ * ChatClient state - track the lifecycle of a chat
+ */
+export type ChatClientState = 'ready' | 'submitted' | 'streaming' | 'error'
+
+/**
+ * Multimodal content input for sending messages with rich media.
+ * Allows sending text, images, audio, video, and documents to the LLM.
+ *
+ * @example
+ * ```ts
+ * // Send an image with a question
+ * client.sendMessage({
+ *   content: [
+ *     { type: 'text', content: 'What is in this image?' },
+ *     { type: 'image', source: { type: 'url', value: 'https://example.com/photo.jpg' } }
+ *   ],
+ *   id: 'custom-message-id' // optional
+ * })
+ * ```
+ */
+export interface MultimodalContent {
+  /**
+   * The content of the message.
+   * Can be a simple string or an array of content parts for multimodal messages.
+   */
+  content: string | Array<ContentPart>
+  /**
+   * Optional custom ID for the message.
+   * If not provided, a unique ID will be generated.
+   */
+  id?: string
+}
+
 /**
  * Message parts - building blocks of UIMessage
  */
@@ -116,6 +155,10 @@ export interface ThinkingPart {
 
 export type MessagePart<TTools extends ReadonlyArray<AnyClientTool> = any> =
   | TextPart
+  | ImagePart
+  | AudioPart
+  | VideoPart
+  | DocumentPart
   | ToolCallPart<TTools>
   | ToolResultPart
   | ThinkingPart
@@ -191,6 +234,11 @@ export interface ChatClientOptions<
    */
   onErrorChange?: (error: Error | undefined) => void
 
+  /**
+   * Callback when chat status changes
+   */
+  onStatusChange?: (status: ChatClientState) => void
+
   /**
    * Client-side tools with execution logic
    * When provided, tools with execute functions will be called automatically