fix (ui): prevent early addToolResult submission (#5427)

lgrammel · web-flow · commit a043b14da143 · 2025-03-31T18:25:15.000+02:00
diff --git a/.changeset/empty-dolphins-ring.md b/.changeset/empty-dolphins-ring.md
@@ -0,0 +1,8 @@
+---
+'@ai-sdk/svelte': patch
+'@ai-sdk/react': patch
+'@ai-sdk/solid': patch
+'@ai-sdk/vue': patch
+---
+
+fix (ui): prevent early addToolResult submission
diff --git a/examples/sveltekit-openai/src/routes/chat/+page.svelte b/examples/sveltekit-openai/src/routes/chat/+page.svelte
@@ -4,7 +4,21 @@
   import { Textarea } from '$lib/components/ui/textarea/index.js';
   import { Chat } from '@ai-sdk/svelte';
 
-  const chat = new Chat();
+  const chat = new Chat({
+    maxSteps: 5,
+
+    // run client-side tools that are automatically executed:
+    async onToolCall({ toolCall }) {
+      // artificial 2 second delay
+      await new Promise(resolve => setTimeout(resolve, 2000));
+
+      if (toolCall.toolName === 'getLocation') {
+        const cities = ['New York', 'Los Angeles', 'Chicago', 'San Francisco'];
+        return cities[Math.floor(Math.random() * cities.length)];
+      }
+    },
+  });
+
   const disabled = $derived(chat.status !== 'ready');
 
   function mapRoleToClass(role: string) {
diff --git a/packages/react/src/use-chat.ts b/packages/react/src/use-chat.ts
@@ -548,15 +548,27 @@ By default, it's set to 1, which means that only a single LLM call is made.
         toolResult: result,
       });
 
-      mutate(currentMessages, false);
+      // array mutation is required to trigger a re-render
+      mutate(
+        [
+          ...currentMessages.slice(0, currentMessages.length - 1),
+          { ...currentMessages[currentMessages.length - 1] },
+        ],
+        false,
+      );
+
+      // when the request is ongoing, the auto-submit will be triggered after the request is finished
+      if (status === 'submitted' || status === 'streaming') {
+        return;
+      }
 
       // auto-submit when all tool calls in the last assistant message have results:
       const lastMessage = currentMessages[currentMessages.length - 1];
       if (isAssistantMessageWithCompletedToolCalls(lastMessage)) {
         triggerRequest({ messages: currentMessages });
       }
     },
-    [mutate, triggerRequest],
+    [mutate, status, triggerRequest],
   );
 
   return {
diff --git a/packages/react/src/use-chat.ui.test.tsx b/packages/react/src/use-chat.ui.test.tsx
@@ -387,9 +387,6 @@ describe('text stream', () => {
     server.urls['/api/chat'].response = {
       type: 'controlled-stream',
       controller,
-      headers: {
-        'Content-Type': 'text/event-stream',
-      },
     };
 
     await userEvent.click(screen.getByTestId('do-append-text-stream'));
@@ -735,7 +732,9 @@ describe('onToolCall', () => {
 
 describe('tool invocations', () => {
   setupTestComponent(() => {
-    const { messages, append, addToolResult } = useChat();
+    const { messages, append, addToolResult } = useChat({
+      maxSteps: 5,
+    });
 
     return (
       <div>
@@ -852,15 +851,11 @@ describe('tool invocations', () => {
     });
   });
 
-  it('should display partial tool call and tool result (when there is no tool call streaming)', async () => {
+  it('should display tool call and tool result (when there is no tool call streaming)', async () => {
     const controller = new TestResponseController();
-
     server.urls['/api/chat'].response = {
       type: 'controlled-stream',
       controller,
-      headers: {
-        'Content-Type': 'text/event-stream',
-      },
     };
 
     await userEvent.click(screen.getByTestId('do-append'));
@@ -922,6 +917,73 @@ describe('tool invocations', () => {
       );
     });
   });
+
+  it('should delay tool result submission until the stream is finished', async () => {
+    const controller1 = new TestResponseController();
+    const controller2 = new TestResponseController();
+
+    server.urls['/api/chat'].response = [
+      { type: 'controlled-stream', controller: controller1 },
+      { type: 'controlled-stream', controller: controller2 },
+    ];
+
+    await userEvent.click(screen.getByTestId('do-append'));
+
+    // start stream
+    controller1.write(
+      formatDataStreamPart('start_step', {
+        messageId: '1234',
+      }),
+    );
+
+    // tool call
+    controller1.write(
+      formatDataStreamPart('tool_call', {
+        toolCallId: 'tool-call-0',
+        toolName: 'test-tool',
+        args: { testArg: 'test-value' },
+      }),
+    );
+
+    await waitFor(() => {
+      expect(screen.getByTestId('message-1')).toHaveTextContent(
+        '{"state":"call","step":0,"toolCallId":"tool-call-0","toolName":"test-tool","args":{"testArg":"test-value"}}',
+      );
+    });
+
+    // user submits the tool result
+    await userEvent.click(screen.getByTestId('add-result-0'));
+
+    // UI should show the tool result
+    await waitFor(() => {
+      expect(screen.getByTestId('message-1')).toHaveTextContent(
+        '{"state":"result","step":0,"toolCallId":"tool-call-0","toolName":"test-tool","args":{"testArg":"test-value"},"result":"test-result"}',
+      );
+    });
+
+    // should not have called the API yet
+    expect(server.calls.length).toBe(1);
+
+    // finish stream
+    controller1.write(
+      formatDataStreamPart('finish_step', {
+        isContinued: false,
+        finishReason: 'tool-calls',
+      }),
+    );
+    controller1.write(
+      formatDataStreamPart('finish_message', {
+        finishReason: 'tool-calls',
+      }),
+    );
+
+    await controller1.close();
+
+    // 2nd call should happen after the stream is finished
+    await waitFor(() => {
+      expect(server.calls.length).toBe(2);
+    });
+  });
 });
 
 describe('maxSteps', () => {
diff --git a/packages/solid/src/use-chat.ts b/packages/solid/src/use-chat.ts
@@ -516,6 +516,11 @@ export function useChat(
 
     mutate(currentMessages);
 
+    // when the request is ongoing, the auto-submit will be triggered after the request is finished
+    if (status() === 'submitted' || status() === 'streaming') {
+      return;
+    }
+
     // auto-submit when all tool calls in the last assistant message have results:
     const lastMessage = currentMessages[currentMessages.length - 1];
     if (isAssistantMessageWithCompletedToolCalls(lastMessage)) {
diff --git a/packages/svelte/src/chat.svelte.test.ts b/packages/svelte/src/chat.svelte.test.ts
@@ -535,7 +535,9 @@ describe('tool invocations', () => {
   let chat: Chat;
 
   beforeEach(() => {
-    chat = new Chat();
+    chat = new Chat({
+      maxSteps: 5,
+    });
   });
 
   it('should display partial tool call, tool call, and tool result', async () => {
@@ -774,6 +776,98 @@ describe('tool invocations', () => {
       );
     });
   });
+
+  it('should delay tool result submission until the stream is finished', async () => {
+    const controller1 = new TestResponseController();
+    const controller2 = new TestResponseController();
+
+    server.urls['/api/chat'].response = [
+      { type: 'controlled-stream', controller: controller1 },
+      { type: 'controlled-stream', controller: controller2 },
+    ];
+
+    chat.append({ role: 'user', content: 'hi' });
+
+    // start stream
+    controller1.write(
+      formatDataStreamPart('start_step', {
+        messageId: '1234',
+      }),
+    );
+
+    // tool call
+    controller1.write(
+      formatDataStreamPart('tool_call', {
+        toolCallId: 'tool-call-0',
+        toolName: 'test-tool',
+        args: { testArg: 'test-value' },
+      }),
+    );
+
+    await vi.waitFor(() => {
+      expect(chat.messages.at(1)).toStrictEqual(
+        expect.objectContaining({
+          toolInvocations: [
+            {
+              state: 'call',
+              step: 0,
+              toolCallId: 'tool-call-0',
+              toolName: 'test-tool',
+              args: { testArg: 'test-value' },
+            },
+          ],
+        }),
+      );
+    });
+
+    // user submits the tool result
+    chat.addToolResult({
+      toolCallId: 'tool-call-0',
+      result: 'test-result',
+    });
+
+    // UI should show the tool result
+    await vi.waitFor(() => {
+      expect(chat.messages.at(1)).toStrictEqual(
+        expect.objectContaining({
+          toolInvocations: [
+            {
+              state: 'result',
+              step: 0,
+              toolCallId: 'tool-call-0',
+              toolName: 'test-tool',
+              args: { testArg: 'test-value' },
+              result: 'test-result',
+            },
+          ],
+        }),
+      );
+    });
+
+    // should not have called the API yet
+    expect(server.calls.length).toBe(1);
+
+    // finish stream
+    controller1.write(
+      formatDataStreamPart('finish_step', {
+        isContinued: false,
+        finishReason: 'tool-calls',
+      }),
+    );
+
+    controller1.write(
+      formatDataStreamPart('finish_message', {
+        finishReason: 'tool-calls',
+      }),
+    );
+
+    await controller1.close();
+
+    // 2nd call should happen after the stream is finished
+    await vi.waitFor(() => {
+      expect(server.calls.length).toBe(2);
+    });
+  });
 });
 
 describe('maxSteps', () => {
diff --git a/packages/svelte/src/chat.svelte.ts b/packages/svelte/src/chat.svelte.ts
@@ -220,6 +220,14 @@ export class Chat {
       toolResult: result,
     });
 
+    // when the request is ongoing, the auto-submit will be triggered after the request is finished
+    if (
+      this.#store.status === 'submitted' ||
+      this.#store.status === 'streaming'
+    ) {
+      return;
+    }
+
     const lastMessage = this.messages[this.messages.length - 1];
     if (isAssistantMessageWithCompletedToolCalls(lastMessage)) {
       await this.#triggerRequest({ messages: this.messages });
diff --git a/packages/vue/src/TestChatToolInvocationsComponent.vue b/packages/vue/src/TestChatToolInvocationsComponent.vue
@@ -1,7 +1,9 @@
 <script setup lang="ts">
 import { useChat } from './use-chat';
 
-const { messages, append, addToolResult } = useChat();
+const { messages, append, addToolResult } = useChat({
+  maxSteps: 5,
+});
 </script>
 
 <template>
@@ -27,6 +29,9 @@ const { messages, append, addToolResult } = useChat();
           "
         />
       </div>
+      <div :data-testid="`text-${idx}`">
+        {{ m.content }}
+      </div>
     </div>
 
     <button
diff --git a/packages/vue/src/use-chat.ts b/packages/vue/src/use-chat.ts
@@ -424,6 +424,11 @@ export function useChat(
 
     mutate(currentMessages);
 
+    // when the request is ongoing, the auto-submit will be triggered after the request is finished
+    if (status.value === 'submitted' || status.value === 'streaming') {
+      return;
+    }
+
     // auto-submit when all tool calls in the last assistant message have results:
     const lastMessage = currentMessages[currentMessages.length - 1];
     if (isAssistantMessageWithCompletedToolCalls(lastMessage)) {
diff --git a/packages/vue/src/use-chat.ui.test.tsx b/packages/vue/src/use-chat.ui.test.tsx
