@ai-sdk/react 0.0.22 → 0.0.23

package/.turbo/turbo-build.log

@@ -1,5 +1,5 @@
 
-> @ai-sdk/react@0.0.22 build /home/runner/work/ai/ai/packages/react
+> @ai-sdk/react@0.0.23 build /home/runner/work/ai/ai/packages/react
 > tsup
 
 CLI Building entry: src/index.ts
@@ -11,11 +11,11 @@
 ESM Build start
 CJS dist/index.js     26.77 KB
 CJS dist/index.js.map 53.33 KB
-CJS ⚡️ Build success in 50ms
+CJS ⚡️ Build success in 61ms
 ESM dist/index.mjs     24.29 KB
 ESM dist/index.mjs.map 53.26 KB
-ESM ⚡️ Build success in 52ms
+ESM ⚡️ Build success in 69ms
 DTS Build start
-DTS ⚡️ Build success in 5293ms
+DTS ⚡️ Build success in 5498ms
 DTS dist/index.d.ts  10.42 KB
 DTS dist/index.d.mts 10.42 KB

package/.turbo/turbo-clean.log

@@ -1,4 +1,4 @@
 
-> @ai-sdk/react@0.0.22 clean /home/runner/work/ai/ai/packages/react
+> @ai-sdk/react@0.0.23 clean /home/runner/work/ai/ai/packages/react
 > rm -rf dist
 

package/CHANGELOG.md

@@ -1,5 +1,12 @@
 # @ai-sdk/react
 
+## 0.0.23
+
+### Patch Changes
+
+- Updated dependencies [1f67fe49]
+  - @ai-sdk/ui-utils@0.0.15
+
 ## 0.0.22
 
 ### Patch Changes

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@ai-sdk/react",
-  "version": "0.0.22",
+  "version": "0.0.23",
   "license": "Apache-2.0",
   "sideEffects": false,
   "main": "./dist/index.js",
@@ -16,7 +16,7 @@
   },
   "dependencies": {
     "@ai-sdk/provider-utils": "1.0.2",
-    "@ai-sdk/ui-utils": "0.0.14",
+    "@ai-sdk/ui-utils": "0.0.15",
     "swr": "2.2.0"
   },
   "devDependencies": {

package/src/use-chat.ui.test.tsx

@@ -1,21 +1,18 @@
 /* eslint-disable @next/next/no-img-element */
-import {
-  mockFetchDataStream,
-  mockFetchDataStreamWithGenerator,
-  mockFetchError,
-} from '@ai-sdk/ui-utils/test';
+import { withTestServer } from '@ai-sdk/provider-utils/test';
+import { formatStreamPart, getTextFromDataUrl } from '@ai-sdk/ui-utils';
 import '@testing-library/jest-dom/vitest';
 import {
-  act,
+  RenderResult,
   cleanup,
   findByText,
   render,
   screen,
+  waitFor,
 } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
 import React, { useRef, useState } from 'react';
 import { useChat } from './use-chat';
-import { formatStreamPart, getTextFromDataUrl } from '@ai-sdk/ui-utils';
 
 describe('stream data stream', () => {
   const TestComponent = () => {
@@ -59,110 +56,127 @@ describe('stream data stream', () => {
     cleanup();
   });
 
-  it('should show streamed response', async () => {
-    mockFetchDataStream({
-      url: 'https://example.com/api/chat',
-      chunks: ['0:"Hello"\n', '0:","\n', '0:" world"\n', '0:"."\n'],
-    });
-
-    await userEvent.click(screen.getByTestId('do-append'));
-
-    await screen.findByTestId('message-0');
-    expect(screen.getByTestId('message-0')).toHaveTextContent('User: hi');
-
-    await screen.findByTestId('message-1');
-    expect(screen.getByTestId('message-1')).toHaveTextContent(
-      'AI: Hello, world.',
-    );
-  });
-
-  it('should show streamed response with data', async () => {
-    mockFetchDataStream({
-      url: 'https://example.com/api/chat',
-      chunks: ['2:[{"t1":"v1"}]\n', '0:"Hello"\n'],
-    });
-
-    await userEvent.click(screen.getByTestId('do-append'));
-
-    await screen.findByTestId('data');
-    expect(screen.getByTestId('data')).toHaveTextContent('[{"t1":"v1"}]');
+  it(
+    'should show streamed response',
+    withTestServer(
+      {
+        type: 'stream-values',
+        url: '/api/chat',
+        content: ['0:"Hello"\n', '0:","\n', '0:" world"\n', '0:"."\n'],
+      },
+      async () => {
+        await userEvent.click(screen.getByTestId('do-append'));
 
-    await screen.findByTestId('message-1');
-    expect(screen.getByTestId('message-1')).toHaveTextContent('AI: Hello');
-  });
+        await screen.findByTestId('message-0');
+        expect(screen.getByTestId('message-0')).toHaveTextContent('User: hi');
 
-  it('should show error response', async () => {
-    mockFetchError({ statusCode: 404, errorMessage: 'Not found' });
+        await screen.findByTestId('message-1');
+        expect(screen.getByTestId('message-1')).toHaveTextContent(
+          'AI: Hello, world.',
+        );
+      },
+    ),
+  );
+
+  it(
+    'should show streamed response with data',
+    withTestServer(
+      {
+        type: 'stream-values',
+        url: '/api/chat',
+        content: ['2:[{"t1":"v1"}]\n', '0:"Hello"\n'],
+      },
+      async () => {
+        await userEvent.click(screen.getByTestId('do-append'));
 
-    await userEvent.click(screen.getByTestId('do-append'));
+        await screen.findByTestId('data');
+        expect(screen.getByTestId('data')).toHaveTextContent('[{"t1":"v1"}]');
 
-    await screen.findByTestId('error');
-    expect(screen.getByTestId('error')).toHaveTextContent('Error: Not found');
-  });
+        await screen.findByTestId('message-1');
+        expect(screen.getByTestId('message-1')).toHaveTextContent('AI: Hello');
+      },
+    ),
+  );
+
+  it(
+    'should show error response',
+    withTestServer(
+      { type: 'error', url: '/api/chat', status: 404, content: 'Not found' },
+      async () => {
+        await userEvent.click(screen.getByTestId('do-append'));
+
+        await screen.findByTestId('error');
+        expect(screen.getByTestId('error')).toHaveTextContent(
+          'Error: Not found',
+        );
+      },
+    ),
+  );
 
   describe('loading state', () => {
-    it('should show loading state', async () => {
-      let finishGeneration: ((value?: unknown) => void) | undefined;
-      const finishGenerationPromise = new Promise(resolve => {
-        finishGeneration = resolve;
-      });
-
-      mockFetchDataStreamWithGenerator({
-        url: 'https://example.com/api/chat',
-        chunkGenerator: (async function* generate() {
-          const encoder = new TextEncoder();
-          yield encoder.encode('0:"Hello"\n');
-          await finishGenerationPromise;
-        })(),
-      });
+    it(
+      'should show loading state',
+      withTestServer(
+        { url: '/api/chat', type: 'controlled-stream' },
+        async ({ streamController }) => {
+          streamController.enqueue('0:"Hello"\n');
 
-      await userEvent.click(screen.getByTestId('do-append'));
+          await userEvent.click(screen.getByTestId('do-append'));
 
-      await screen.findByTestId('loading');
-      expect(screen.getByTestId('loading')).toHaveTextContent('true');
+          await screen.findByTestId('loading');
+          expect(screen.getByTestId('loading')).toHaveTextContent('true');
 
-      finishGeneration?.();
+          streamController.close();
 
-      await findByText(await screen.findByTestId('loading'), 'false');
-      expect(screen.getByTestId('loading')).toHaveTextContent('false');
-    });
-
-    it('should reset loading state on error', async () => {
-      mockFetchError({ statusCode: 404, errorMessage: 'Not found' });
+          await findByText(await screen.findByTestId('loading'), 'false');
+          expect(screen.getByTestId('loading')).toHaveTextContent('false');
+        },
+      ),
+    );
 
-      await userEvent.click(screen.getByTestId('do-append'));
+    it(
+      'should reset loading state on error',
+      withTestServer(
+        { type: 'error', url: '/api/chat', status: 404, content: 'Not found' },
+        async () => {
+          await userEvent.click(screen.getByTestId('do-append'));
 
-      await screen.findByTestId('loading');
-      expect(screen.getByTestId('loading')).toHaveTextContent('false');
-    });
+          await screen.findByTestId('loading');
+          expect(screen.getByTestId('loading')).toHaveTextContent('false');
+        },
+      ),
+    );
   });
 
   describe('id', () => {
-    it('should clear out messages when the id changes', async () => {
-      mockFetchDataStream({
-        url: 'https://example.com/api/chat',
-        chunks: ['0:"Hello"\n', '0:","\n', '0:" world"\n', '0:"."\n'],
-      });
-
-      await userEvent.click(screen.getByTestId('do-append'));
+    it(
+      'should clear out messages when the id changes',
+      withTestServer(
+        {
+          url: '/api/chat',
+          type: 'stream-values',
+          content: ['0:"Hello"\n', '0:","\n', '0:" world"\n', '0:"."\n'],
+        },
+        async () => {
+          await userEvent.click(screen.getByTestId('do-append'));
 
-      await screen.findByTestId('message-1');
-      expect(screen.getByTestId('message-1')).toHaveTextContent(
-        'AI: Hello, world.',
-      );
+          await screen.findByTestId('message-1');
+          expect(screen.getByTestId('message-1')).toHaveTextContent(
+            'AI: Hello, world.',
+          );
 
-      await userEvent.click(screen.getByTestId('do-change-id'));
+          await userEvent.click(screen.getByTestId('do-change-id'));
 
-      expect(screen.queryByTestId('message-0')).not.toBeInTheDocument();
-    });
+          expect(screen.queryByTestId('message-0')).not.toBeInTheDocument();
+        },
+      ),
+    );
   });
 });
 
 describe('text stream', () => {
   const TestComponent = () => {
-    const { messages, append } = useChat({
-      streamMode: 'text',
-    });
+    const { messages, append } = useChat({ streamMode: 'text' });
 
     return (
       <div>
@@ -192,38 +206,35 @@ describe('text stream', () => {
     cleanup();
   });
 
-  it('should show streamed response', async () => {
-    mockFetchDataStream({
-      url: 'https://example.com/api/chat',
-      chunks: ['Hello', ',', ' world', '.'],
-    });
-
-    await userEvent.click(screen.getByTestId('do-append-text-stream'));
-
-    await screen.findByTestId('message-0-text-stream');
-    expect(screen.getByTestId('message-0-text-stream')).toHaveTextContent(
-      'User: hi',
-    );
-
-    await screen.findByTestId('message-1-text-stream');
-    expect(screen.getByTestId('message-1-text-stream')).toHaveTextContent(
-      'AI: Hello, world.',
-    );
-  });
+  it(
+    'should show streamed response',
+    withTestServer(
+      {
+        url: '/api/chat',
+        type: 'stream-values',
+        content: ['Hello', ',', ' world', '.'],
+      },
+      async () => {
+        await userEvent.click(screen.getByTestId('do-append-text-stream'));
+
+        await screen.findByTestId('message-0-text-stream');
+        expect(screen.getByTestId('message-0-text-stream')).toHaveTextContent(
+          'User: hi',
+        );
+
+        await screen.findByTestId('message-1-text-stream');
+        expect(screen.getByTestId('message-1-text-stream')).toHaveTextContent(
+          'AI: Hello, world.',
+        );
+      },
+    ),
+  );
 });
 
 describe('form actions', () => {
   const TestComponent = () => {
-    const {
-      messages,
-      append,
-      handleSubmit,
-      handleInputChange,
-      isLoading,
-      input,
-    } = useChat({
-      streamMode: 'text',
-    });
+    const { messages, handleSubmit, handleInputChange, isLoading, input } =
+      useChat({ streamMode: 'text' });
 
     return (
       <div>
@@ -256,37 +267,44 @@ describe('form actions', () => {
     cleanup();
   });
 
-  it('should show streamed response using handleSubmit', async () => {
-    mockFetchDataStream({
-      url: 'https://example.com/api/chat',
-      chunks: ['Hello', ',', ' world', '.'],
-    });
-
-    const firstInput = screen.getByTestId('do-input');
-    await userEvent.type(firstInput, 'hi');
-    await userEvent.keyboard('{Enter}');
-
-    await screen.findByTestId('message-0');
-    expect(screen.getByTestId('message-0')).toHaveTextContent('User: hi');
-
-    await screen.findByTestId('message-1');
-    expect(screen.getByTestId('message-1')).toHaveTextContent(
-      'AI: Hello, world.',
-    );
-
-    mockFetchDataStream({
-      url: 'https://example.com/api/chat',
-      chunks: ['How', ' can', ' I', ' help', ' you', '?'],
-    });
-
-    const secondInput = screen.getByTestId('do-input');
-    await userEvent.type(secondInput, '{Enter}');
-
-    await screen.findByTestId('message-2');
-    expect(screen.getByTestId('message-2')).toHaveTextContent(
-      'AI: How can I help you?',
-    );
-  });
+  it(
+    'should show streamed response using handleSubmit',
+    withTestServer(
+      [
+        {
+          url: '/api/chat',
+          type: 'stream-values',
+          content: ['Hello', ',', ' world', '.'],
+        },
+        {
+          url: '/api/chat',
+          type: 'stream-values',
+          content: ['How', ' can', ' I', ' help', ' you', '?'],
+        },
+      ],
+      async () => {
+        const firstInput = screen.getByTestId('do-input');
+        await userEvent.type(firstInput, 'hi');
+        await userEvent.keyboard('{Enter}');
+
+        await screen.findByTestId('message-0');
+        expect(screen.getByTestId('message-0')).toHaveTextContent('User: hi');
+
+        await screen.findByTestId('message-1');
+        expect(screen.getByTestId('message-1')).toHaveTextContent(
+          'AI: Hello, world.',
+        );
+
+        const secondInput = screen.getByTestId('do-input');
+        await userEvent.type(secondInput, '{Enter}');
+
+        await screen.findByTestId('message-2');
+        expect(screen.getByTestId('message-2')).toHaveTextContent(
+          'AI: How can I help you?',
+        );
+      },
+    ),
+  );
 });
 
 describe('prepareRequestBody', () => {
@@ -336,30 +354,35 @@ describe('prepareRequestBody', () => {
     cleanup();
   });
 
-  it('should show streamed response', async () => {
-    const { requestBody } = mockFetchDataStream({
-      url: 'https://example.com/api/chat',
-      chunks: ['0:"Hello"\n', '0:","\n', '0:" world"\n', '0:"."\n'],
-    });
-
-    await userEvent.click(screen.getByTestId('do-append'));
+  it(
+    'should show streamed response',
+    withTestServer(
+      {
+        url: '/api/chat',
+        type: 'stream-values',
+        content: ['0:"Hello"\n', '0:","\n', '0:" world"\n', '0:"."\n'],
+      },
+      async ({ call }) => {
+        await userEvent.click(screen.getByTestId('do-append'));
 
-    await screen.findByTestId('message-0');
-    expect(screen.getByTestId('message-0')).toHaveTextContent('User: hi');
+        await screen.findByTestId('message-0');
+        expect(screen.getByTestId('message-0')).toHaveTextContent('User: hi');
 
-    expect(bodyOptions).toStrictEqual({
-      messages: [{ role: 'user', content: 'hi', id: expect.any(String) }],
-      requestData: { 'test-data-key': 'test-data-value' },
-      requestBody: { 'request-body-key': 'request-body-value' },
-    });
+        expect(bodyOptions).toStrictEqual({
+          messages: [{ role: 'user', content: 'hi', id: expect.any(String) }],
+          requestData: { 'test-data-key': 'test-data-value' },
+          requestBody: { 'request-body-key': 'request-body-value' },
+        });
 
-    expect(await requestBody).toBe('"test-request-body"');
+        expect(await call(0).getRequestBodyJson()).toBe('test-request-body');
 
-    await screen.findByTestId('message-1');
-    expect(screen.getByTestId('message-1')).toHaveTextContent(
-      'AI: Hello, world.',
-    );
-  });
+        await screen.findByTestId('message-1');
+        expect(screen.getByTestId('message-1')).toHaveTextContent(
+          'AI: Hello, world.',
+        );
+      },
+    ),
+  );
 });
 
 describe('onToolCall', () => {
@@ -405,36 +428,203 @@ describe('onToolCall', () => {
     cleanup();
   });
 
-  it("should invoke onToolCall when a tool call is received from the server's response", async () => {
-    mockFetchDataStream({
-      url: 'https://example.com/api/chat',
-      chunks: [
-        formatStreamPart('tool_call', {
-          toolCallId: 'tool-call-0',
-          toolName: 'test-tool',
-          args: { testArg: 'test-value' },
-        }),
-      ],
-    });
+  it(
+    "should invoke onToolCall when a tool call is received from the server's response",
+    withTestServer(
+      {
+        url: '/api/chat',
+        type: 'stream-values',
+        content: [
+          formatStreamPart('tool_call', {
+            toolCallId: 'tool-call-0',
+            toolName: 'test-tool',
+            args: { testArg: 'test-value' },
+          }),
+        ],
+      },
+      async () => {
+        await userEvent.click(screen.getByTestId('do-append'));
 
-    await userEvent.click(screen.getByTestId('do-append'));
+        await screen.findByTestId('message-1');
+        expect(screen.getByTestId('message-1')).toHaveTextContent(
+          'test-tool-response: test-tool tool-call-0 {"testArg":"test-value"}',
+        );
+      },
+    ),
+  );
+});
+
+describe('tool invocations', () => {
+  let rerender: RenderResult['rerender'];
+
+  const TestComponent = () => {
+    const { messages, append } = useChat();
 
-    await screen.findByTestId('message-1');
-    expect(screen.getByTestId('message-1')).toHaveTextContent(
-      'test-tool-response: test-tool tool-call-0 {"testArg":"test-value"}',
+    return (
+      <div>
+        {messages.map((m, idx) => (
+          <div data-testid={`message-${idx}`} key={m.id}>
+            {m.toolInvocations?.map((toolInvocation, toolIdx) => {
+              return (
+                <div key={toolIdx} data-testid={`tool-invocation-${toolIdx}`}>
+                  {JSON.stringify(toolInvocation)}
+                </div>
+              );
+            })}
+          </div>
+        ))}
+
+        <button
+          data-testid="do-append"
+          onClick={() => {
+            append({ role: 'user', content: 'hi' });
+          }}
+        />
+      </div>
     );
+  };
+
+  beforeEach(() => {
+    const result = render(<TestComponent />);
+    rerender = result.rerender;
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+    cleanup();
   });
+
+  it(
+    'should display partial tool call, tool call, and tool result',
+    withTestServer(
+      { url: '/api/chat', type: 'controlled-stream' },
+      async ({ streamController }) => {
+        await userEvent.click(screen.getByTestId('do-append'));
+
+        streamController.enqueue(
+          formatStreamPart('tool_call_streaming_start', {
+            toolCallId: 'tool-call-0',
+            toolName: 'test-tool',
+          }),
+        );
+
+        await waitFor(() => {
+          expect(screen.getByTestId('message-1')).toHaveTextContent(
+            '{"state":"partial-call","toolCallId":"tool-call-0","toolName":"test-tool"}',
+          );
+        });
+
+        streamController.enqueue(
+          formatStreamPart('tool_call_delta', {
+            toolCallId: 'tool-call-0',
+            argsTextDelta: '{"testArg":"t',
+          }),
+        );
+
+        await waitFor(() => {
+          rerender(<TestComponent />);
+          expect(screen.getByTestId('message-1')).toHaveTextContent(
+            '{"state":"partial-call","toolCallId":"tool-call-0","toolName":"test-tool","args":{"testArg":"t"}}',
+          );
+        });
+
+        streamController.enqueue(
+          formatStreamPart('tool_call_delta', {
+            toolCallId: 'tool-call-0',
+            argsTextDelta: 'est-value"}}',
+          }),
+        );
+
+        await waitFor(() => {
+          rerender(<TestComponent />);
+          expect(screen.getByTestId('message-1')).toHaveTextContent(
+            '{"state":"partial-call","toolCallId":"tool-call-0","toolName":"test-tool","args":{"testArg":"test-value"}}',
+          );
+        });
+
+        streamController.enqueue(
+          formatStreamPart('tool_call', {
+            toolCallId: 'tool-call-0',
+            toolName: 'test-tool',
+            args: { testArg: 'test-value' },
+          }),
+        );
+
+        await waitFor(() => {
+          rerender(<TestComponent />);
+          expect(screen.getByTestId('message-1')).toHaveTextContent(
+            '{"state":"call","toolCallId":"tool-call-0","toolName":"test-tool","args":{"testArg":"test-value"}}',
+          );
+        });
+
+        streamController.enqueue(
+          formatStreamPart('tool_result', {
+            toolCallId: 'tool-call-0',
+            toolName: 'test-tool',
+            args: { testArg: 'test-value' },
+            result: 'test-result',
+          }),
+        );
+        streamController.close();
+
+        await waitFor(() => {
+          expect(screen.getByTestId('message-1')).toHaveTextContent(
+            '{"state":"result","toolCallId":"tool-call-0","toolName":"test-tool","args":{"testArg":"test-value"},"result":"test-result"}',
+          );
+        });
+      },
+    ),
+  );
+
+  it(
+    'should display partial tool call and tool result (when there is no tool call streaming)',
+    withTestServer(
+      { url: '/api/chat', type: 'controlled-stream' },
+      async ({ streamController }) => {
+        await userEvent.click(screen.getByTestId('do-append'));
+
+        streamController.enqueue(
+          formatStreamPart('tool_call', {
+            toolCallId: 'tool-call-0',
+            toolName: 'test-tool',
+            args: { testArg: 'test-value' },
+          }),
+        );
+
+        await waitFor(() => {
+          expect(screen.getByTestId('message-1')).toHaveTextContent(
+            '{"state":"call","toolCallId":"tool-call-0","toolName":"test-tool","args":{"testArg":"test-value"}}',
+          );
+        });
+
+        streamController.enqueue(
+          formatStreamPart('tool_result', {
+            toolCallId: 'tool-call-0',
+            toolName: 'test-tool',
+            args: { testArg: 'test-value' },
+            result: 'test-result',
+          }),
+        );
+        streamController.close();
+
+        await waitFor(() => {
+          expect(screen.getByTestId('message-1')).toHaveTextContent(
+            '{"state":"result","toolCallId":"tool-call-0","toolName":"test-tool","args":{"testArg":"test-value"},"result":"test-result"}',
+          );
+        });
+      },
+    ),
+  );
 });
 
 describe('maxToolRoundtrips', () => {
   describe('single automatic tool roundtrip', () => {
+    let onToolCallInvoked = false;
+
     const TestComponent = () => {
       const { messages, append } = useChat({
         async onToolCall({ toolCall }) {
-          mockFetchDataStream({
-            url: 'https://example.com/api/chat',
-            chunks: [formatStreamPart('text', 'final result')],
-          });
+          onToolCallInvoked = true;
 
           return `test-tool-response: ${toolCall.toolName} ${
             toolCall.toolCallId
@@ -463,6 +653,7 @@ describe('maxToolRoundtrips', () => {
 
     beforeEach(() => {
       render(<TestComponent />);
+      onToolCallInvoked = false;
     });
 
     afterEach(() => {
@@ -470,35 +661,48 @@ describe('maxToolRoundtrips', () => {
       cleanup();
     });
 
-    it('should automatically call api when tool call gets executed via onToolCall', async () => {
-      mockFetchDataStream({
-        url: 'https://example.com/api/chat',
-        chunks: [
-          formatStreamPart('tool_call', {
-            toolCallId: 'tool-call-0',
-            toolName: 'test-tool',
-            args: { testArg: 'test-value' },
-          }),
+    it(
+      'should automatically call api when tool call gets executed via onToolCall',
+      withTestServer(
+        [
+          {
+            url: '/api/chat',
+            type: 'stream-values',
+            content: [
+              formatStreamPart('tool_call', {
+                toolCallId: 'tool-call-0',
+                toolName: 'test-tool',
+                args: { testArg: 'test-value' },
+              }),
+            ],
+          },
+          {
+            url: '/api/chat',
+            type: 'stream-values',
+            content: [formatStreamPart('text', 'final result')],
+          },
         ],
-      });
+        async () => {
+          await userEvent.click(screen.getByTestId('do-append'));
 
-      await userEvent.click(screen.getByTestId('do-append'));
+          expect(onToolCallInvoked).toBe(true);
 
-      await screen.findByTestId('message-2');
-      expect(screen.getByTestId('message-2')).toHaveTextContent('final result');
-    });
+          await screen.findByTestId('message-2');
+          expect(screen.getByTestId('message-2')).toHaveTextContent(
+            'final result',
+          );
+        },
+      ),
+    );
   });
 
   describe('single roundtrip with error response', () => {
+    let onToolCallCounter = 0;
+
     const TestComponent = () => {
       const { messages, append, error } = useChat({
         async onToolCall({ toolCall }) {
-          mockFetchDataStream({
-            url: 'https://example.com/api/chat',
-            chunks: [formatStreamPart('error', 'some failure')],
-            maxCalls: 1,
-          });
-
+          onToolCallCounter++;
           return `test-tool-response: ${toolCall.toolName} ${
             toolCall.toolCallId
           } ${JSON.stringify(toolCall.args)}`;
@@ -534,6 +738,7 @@ describe('maxToolRoundtrips', () => {
 
     beforeEach(() => {
       render(<TestComponent />);
+      onToolCallCounter = 0;
     });
 
     afterEach(() => {
@@ -541,34 +746,47 @@ describe('maxToolRoundtrips', () => {
       cleanup();
     });
 
-    it('should automatically call api when tool call gets executed via onToolCall', async () => {
-      mockFetchDataStream({
-        url: 'https://example.com/api/chat',
-        chunks: [
-          formatStreamPart('tool_call', {
-            toolCallId: 'tool-call-0',
-            toolName: 'test-tool',
-            args: { testArg: 'test-value' },
-          }),
+    it(
+      'should automatically call api when tool call gets executed via onToolCall',
+      withTestServer(
+        [
+          {
+            url: '/api/chat',
+            type: 'stream-values',
+            content: [
+              formatStreamPart('tool_call', {
+                toolCallId: 'tool-call-0',
+                toolName: 'test-tool',
+                args: { testArg: 'test-value' },
+              }),
+            ],
+          },
+          {
+            url: '/api/chat',
+            type: 'error',
+            status: 400,
+            content: 'call failure',
+          },
         ],
-      });
+        async () => {
+          await userEvent.click(screen.getByTestId('do-append'));
 
-      await userEvent.click(screen.getByTestId('do-append'));
+          await screen.findByTestId('error');
+          expect(screen.getByTestId('error')).toHaveTextContent(
+            'Error: call failure',
+          );
 
-      await screen.findByTestId('error');
-      expect(screen.getByTestId('error')).toHaveTextContent(
-        'Error: Too many calls',
-      );
-    });
+          expect(onToolCallCounter).toBe(1);
+        },
+      ),
+    );
   });
 });
 
 describe('file attachments with data url', () => {
   const TestComponent = () => {
     const { messages, handleSubmit, handleInputChange, isLoading, input } =
-      useChat({
-        api: '/api/stream-chat',
-      });
+      useChat();
 
     const [attachments, setAttachments] = useState<FileList | undefined>(
       undefined,
@@ -648,102 +866,124 @@ describe('file attachments with data url', () => {
     cleanup();
   });
 
-  it('should handle text file attachment and submission', async () => {
-    const file = new File(['test file content'], 'test.txt', {
-      type: 'text/plain',
-    });
-
-    const { requestBody } = mockFetchDataStream({
-      url: '/api/stream-chat',
-      chunks: ['0:"Response to message with text attachment"\n'],
-    });
-
-    const fileInput = screen.getByTestId('file-input');
-    await userEvent.upload(fileInput, file);
-
-    const messageInput = screen.getByTestId('message-input');
-    await userEvent.type(messageInput, 'Message with text attachment');
-
-    const submitButton = screen.getByTestId('submit-button');
-    await userEvent.click(submitButton);
-
-    const sentBody = JSON.parse((await requestBody) as string);
-    expect(sentBody.messages[0].content).toBe('Message with text attachment');
-    expect(sentBody.messages[0].experimental_attachments).toBeDefined();
-    expect(sentBody.messages[0].experimental_attachments.length).toBe(1);
-    expect(sentBody.messages[0].experimental_attachments[0].name).toBe(
-      'test.txt',
-    );
-
-    await screen.findByTestId('message-0');
-    expect(screen.getByTestId('message-0')).toHaveTextContent(
-      'User: Message with text attachment',
-    );
-
-    await screen.findByTestId('attachment-0');
-    expect(screen.getByTestId('attachment-0')).toHaveTextContent(
-      'test file content',
-    );
-
-    await screen.findByTestId('message-1');
-    expect(screen.getByTestId('message-1')).toHaveTextContent(
-      'AI: Response to message with text attachment',
-    );
-  });
-
-  // image file
-
-  it('should handle image file attachment and submission', async () => {
-    const file = new File(['test image content'], 'test.png', {
-      type: 'image/png',
-    });
+  it(
+    'should handle text file attachment and submission',
+    withTestServer(
+      {
+        url: '/api/chat',
+        type: 'stream-values',
+        content: ['0:"Response to message with text attachment"\n'],
+      },
+      async ({ call }) => {
+        const file = new File(['test file content'], 'test.txt', {
+          type: 'text/plain',
+        });
 
-    const { requestBody } = mockFetchDataStream({
-      url: '/api/stream-chat',
-      chunks: ['0:"Response to message with image attachment"\n'],
-    });
+        const fileInput = screen.getByTestId('file-input');
+        await userEvent.upload(fileInput, file);
 
-    const fileInput = screen.getByTestId('file-input');
-    await userEvent.upload(fileInput, file);
+        const messageInput = screen.getByTestId('message-input');
+        await userEvent.type(messageInput, 'Message with text attachment');
 
-    const messageInput = screen.getByTestId('message-input');
-    await userEvent.type(messageInput, 'Message with image attachment');
+        const submitButton = screen.getByTestId('submit-button');
+        await userEvent.click(submitButton);
 
-    const submitButton = screen.getByTestId('submit-button');
-    await userEvent.click(submitButton);
+        expect(await call(0).getRequestBodyJson()).toStrictEqual({
+          messages: [
+            {
+              role: 'user',
+              content: 'Message with text attachment',
+              experimental_attachments: [
+                {
+                  name: 'test.txt',
+                  contentType: 'text/plain',
+                  url: 'data:text/plain;base64,dGVzdCBmaWxlIGNvbnRlbnQ=',
+                },
+              ],
+            },
+          ],
+        });
+
+        await screen.findByTestId('message-0');
+        expect(screen.getByTestId('message-0')).toHaveTextContent(
+          'User: Message with text attachment',
+        );
+
+        await screen.findByTestId('attachment-0');
+        expect(screen.getByTestId('attachment-0')).toHaveTextContent(
+          'test file content',
+        );
+
+        await screen.findByTestId('message-1');
+        expect(screen.getByTestId('message-1')).toHaveTextContent(
+          'AI: Response to message with text attachment',
+        );
+      },
+    ),
+  );
+
+  it(
+    'should handle image file attachment and submission',
+    withTestServer(
+      {
+        url: '/api/chat',
+        type: 'stream-values',
+        content: ['0:"Response to message with image attachment"\n'],
+      },
+      async ({ call }) => {
+        const file = new File(['test image content'], 'test.png', {
+          type: 'image/png',
+        });
 
-    const sentBody = JSON.parse((await requestBody) as string);
-    expect(sentBody.messages[0].content).toBe('Message with image attachment');
-    expect(sentBody.messages[0].experimental_attachments).toBeDefined();
-    expect(sentBody.messages[0].experimental_attachments.length).toBe(1);
-    expect(sentBody.messages[0].experimental_attachments[0].name).toBe(
-      'test.png',
-    );
+        const fileInput = screen.getByTestId('file-input');
+        await userEvent.upload(fileInput, file);
 
-    await screen.findByTestId('message-0');
-    expect(screen.getByTestId('message-0')).toHaveTextContent(
-      'User: Message with image attachment',
-    );
+        const messageInput = screen.getByTestId('message-input');
+        await userEvent.type(messageInput, 'Message with image attachment');
 
-    await screen.findByTestId('attachment-0');
-    expect(screen.getByTestId('attachment-0')).toHaveAttribute(
-      'src',
-      expect.stringContaining('data:image/png;base64'),
-    );
+        const submitButton = screen.getByTestId('submit-button');
+        await userEvent.click(submitButton);
 
-    await screen.findByTestId('message-1');
-    expect(screen.getByTestId('message-1')).toHaveTextContent(
-      'AI: Response to message with image attachment',
-    );
-  });
+        expect(await call(0).getRequestBodyJson()).toStrictEqual({
+          messages: [
+            {
+              role: 'user',
+              content: 'Message with image attachment',
+              experimental_attachments: [
+                {
+                  name: 'test.png',
+                  contentType: 'image/png',
+                  url: 'data:image/png;base64,dGVzdCBpbWFnZSBjb250ZW50',
+                },
+              ],
+            },
+          ],
+        });
+
+        await screen.findByTestId('message-0');
+        expect(screen.getByTestId('message-0')).toHaveTextContent(
+          'User: Message with image attachment',
+        );
+
+        await screen.findByTestId('attachment-0');
+        expect(screen.getByTestId('attachment-0')).toHaveAttribute(
+          'src',
+          expect.stringContaining('data:image/png;base64'),
+        );
+
+        await screen.findByTestId('message-1');
+        expect(screen.getByTestId('message-1')).toHaveTextContent(
+          'AI: Response to message with image attachment',
+        );
+      },
+    ),
+  );
 });
 
 describe('file attachments with url', () => {
   const TestComponent = () => {
     const { messages, handleSubmit, handleInputChange, isLoading, input } =
-      useChat({
-        api: '/api/stream-chat',
-      });
+      useChat();
 
     return (
       <div>
@@ -812,40 +1052,53 @@ describe('file attachments with url', () => {
     cleanup();
   });
 
-  it('should handle image file attachment and submission', async () => {
-    const { requestBody } = mockFetchDataStream({
-      url: '/api/stream-chat',
-      chunks: ['0:"Response to message with image attachment"\n'],
-    });
-
-    const messageInput = screen.getByTestId('message-input');
-    await userEvent.type(messageInput, 'Message with image attachment');
-
-    const submitButton = screen.getByTestId('submit-button');
-    await userEvent.click(submitButton);
-
-    const sentBody = JSON.parse((await requestBody) as string);
-    expect(sentBody.messages[0].content).toBe('Message with image attachment');
-    expect(sentBody.messages[0].experimental_attachments).toBeDefined();
-    expect(sentBody.messages[0].experimental_attachments.length).toBe(1);
-    expect(sentBody.messages[0].experimental_attachments[0].name).toBe(
-      'test.png',
-    );
-
-    await screen.findByTestId('message-0');
-    expect(screen.getByTestId('message-0')).toHaveTextContent(
-      'User: Message with image attachment',
-    );
-
-    await screen.findByTestId('attachment-0');
-    expect(screen.getByTestId('attachment-0')).toHaveAttribute(
-      'src',
-      expect.stringContaining('https://example.com/image.png'),
-    );
-
-    await screen.findByTestId('message-1');
-    expect(screen.getByTestId('message-1')).toHaveTextContent(
-      'AI: Response to message with image attachment',
-    );
-  });
+  it(
+    'should handle image file attachment and submission',
+    withTestServer(
+      {
+        url: '/api/chat',
+        type: 'stream-values',
+        content: ['0:"Response to message with image attachment"\n'],
+      },
+      async ({ call }) => {
+        const messageInput = screen.getByTestId('message-input');
+        await userEvent.type(messageInput, 'Message with image attachment');
+
+        const submitButton = screen.getByTestId('submit-button');
+        await userEvent.click(submitButton);
+
+        expect(await call(0).getRequestBodyJson()).toStrictEqual({
+          messages: [
+            {
+              role: 'user',
+              content: 'Message with image attachment',
+              experimental_attachments: [
+                {
+                  name: 'test.png',
+                  contentType: 'image/png',
+                  url: 'https://example.com/image.png',
+                },
+              ],
+            },
+          ],
+        });
+
+        await screen.findByTestId('message-0');
+        expect(screen.getByTestId('message-0')).toHaveTextContent(
+          'User: Message with image attachment',
+        );
+
+        await screen.findByTestId('attachment-0');
+        expect(screen.getByTestId('attachment-0')).toHaveAttribute(
+          'src',
+          expect.stringContaining('https://example.com/image.png'),
+        );
+
+        await screen.findByTestId('message-1');
+        expect(screen.getByTestId('message-1')).toHaveTextContent(
+          'AI: Response to message with image attachment',
+        );
+      },
+    ),
+  );
 });