Fix/tool race conditions

.changeset/tool-race-conditions.md

@@ -0,0 +1,19 @@
+---
+'@tanstack/ai': patch
+'@tanstack/ai-client': patch
+---
+
+fix: improve tool execution reliability and prevent race conditions
+
+- Fix client tool execution race conditions by tracking pending tool executions
+- Prevent duplicate continuation attempts with continuationPending flag
+- Guard against concurrent stream processing in streamResponse
+- Add approval info to ToolCall type for server-side decision tracking
+- Include approval info in model message conversion for approval workflows
+- Check ModelMessage format for approval info extraction in chat activity
+
+This change improves the reliability of tool execution, especially for:
+
+- Client tools with async execute functions
+- Approval-based tool workflows
+- Sequential tool execution scenarios

.gitignore

@@ -52,4 +52,8 @@ vite.config.ts.timestamp-*
 test-traces
 **/adapters/output
 .nitro
-.output
+.output
+
+# Playwright
+playwright-report
+test-results

packages/typescript/ai-client/src/chat-client.ts

@@ -29,6 +29,10 @@ export class ChatClient {
   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: {
@@ -133,31 +137,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: {
@@ -238,6 +252,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()
 
@@ -298,21 +318,35 @@ 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()
   }
 
   /**
    * 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()
@@ -325,12 +359,13 @@ export class ChatClient {
 
       // Connect and stream
       const stream = this.connection.connect(
-        modelMessages,
+        messages,
         bodyWithConversationId,
         this.abortController.signal,
       )
 
       await this.processStream(stream)
+      streamCompletedSuccessfully = true
     } catch (err) {
       if (err instanceof Error) {
         if (err.name === 'AbortError') {
@@ -346,6 +381,20 @@ 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)
+      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)
+          }
+        }
+      }
     }
   }
 
@@ -490,8 +539,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
+      }
     }
   }
 

packages/typescript/ai-client/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)
     },
   }
 }