use-assistant: fix missing message content. (#1425)

Author: lgrammel

.changeset/dirty-countries-learn.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+use-assistant: fix missing message content

examples/next-openai/app/assistant/page.tsx

@@ -14,9 +14,7 @@ const roleToColorMap: Record<Message['role'], string> = {
 
 export default function Chat() {
   const { status, messages, input, submitMessage, handleInputChange, error } =
-    useAssistant({
-      api: '/api/assistant',
-    });
+    useAssistant({ api: '/api/assistant' });
 
   // When status changes to accepting messages, focus the input:
   const inputRef = useRef<HTMLInputElement>(null);

packages/core/react/use-assistant.ts

@@ -153,26 +153,26 @@ export function useAssistant({
 
     setInput('');
 
-    const result = await fetch(api, {
-      method: 'POST',
-      credentials,
-      headers: { 'Content-Type': 'application/json', ...headers },
-      body: JSON.stringify({
-        ...body,
-        // always use user-provided threadId when available:
-        threadId: threadIdParam ?? threadId ?? null,
-        message: input,
-
-        // optional request data:
-        data: requestOptions?.data,
-      }),
-    });
-
-    if (result.body == null) {
-      throw new Error('The response body is empty.');
-    }
-
     try {
+      const result = await fetch(api, {
+        method: 'POST',
+        credentials,
+        headers: { 'Content-Type': 'application/json', ...headers },
+        body: JSON.stringify({
+          ...body,
+          // always use user-provided threadId when available:
+          threadId: threadIdParam ?? threadId ?? null,
+          message: message.content,
+
+          // optional request data:
+          data: requestOptions?.data,
+        }),
+      });
+
+      if (result.body == null) {
+        throw new Error('The response body is empty.');
+      }
+
       for await (const { type, value } of readDataStream(
         result.body.getReader(),
       )) {

packages/core/react/use-assistant.ui.test.tsx

@@ -0,0 +1,132 @@
+import '@testing-library/jest-dom/vitest';
+import { cleanup, findByText, render, screen } from '@testing-library/react';
+import userEvent from '@testing-library/user-event';
+import { formatStreamPart } from '../streams';
+import {
+  mockFetchDataStream,
+  mockFetchDataStreamWithGenerator,
+} from '../tests/utils/mock-fetch';
+import { useAssistant } from './use-assistant';
+
+describe('stream data stream', () => {
+  const TestComponent = () => {
+    const { status, messages, append } = useAssistant({
+      api: '/api/assistant',
+    });
+
+    return (
+      <div>
+        <div data-testid="status">{status}</div>
+        {messages.map((m, idx) => (
+          <div data-testid={`message-${idx}`} key={m.id}>
+            {m.role === 'user' ? 'User: ' : 'AI: '}
+            {m.content}
+          </div>
+        ))}
+
+        <button
+          data-testid="do-append"
+          onClick={() => {
+            append({ role: 'user', content: 'hi' });
+          }}
+        />
+      </div>
+    );
+  };
+
+  beforeEach(() => {
+    render(<TestComponent />);
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+    cleanup();
+  });
+
+  it('should show streamed response', async () => {
+    const { requestBody } = mockFetchDataStream({
+      url: 'https://example.com/api/assistant',
+      chunks: [
+        formatStreamPart('assistant_control_data', {
+          threadId: 't0',
+          messageId: 'm0',
+        }),
+        formatStreamPart('assistant_message', {
+          id: 'm0',
+          role: 'assistant',
+          content: [{ type: 'text', text: { value: '' } }],
+        }),
+        // text parts:
+        '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.',
+    );
+
+    // check that correct information was sent to the server:
+    expect(await requestBody).toStrictEqual(
+      JSON.stringify({
+        threadId: null,
+        message: 'hi',
+      }),
+    );
+  });
+
+  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(
+            formatStreamPart('assistant_control_data', {
+              threadId: 't0',
+              messageId: 'm1',
+            }),
+          );
+
+          yield encoder.encode(
+            formatStreamPart('assistant_message', {
+              id: 'm1',
+              role: 'assistant',
+              content: [{ type: 'text', text: { value: '' } }],
+            }),
+          );
+
+          yield encoder.encode('0:"Hello"\n');
+
+          await finishGenerationPromise;
+        })(),
+      });
+
+      await userEvent.click(screen.getByTestId('do-append'));
+
+      await screen.findByTestId('status');
+      expect(screen.getByTestId('status')).toHaveTextContent('in_progress');
+
+      finishGeneration?.();
+
+      await findByText(await screen.findByTestId('status'), 'awaiting_message');
+      expect(screen.getByTestId('status')).toHaveTextContent(
+        'awaiting_message',
+      );
+    });
+  });
+});

packages/core/tests/utils/mock-fetch.ts

@@ -51,7 +51,7 @@ export function mockFetchDataStream({
     }
   }
 
-  mockFetchDataStreamWithGenerator({
+  return mockFetchDataStreamWithGenerator({
     url,
     chunkGenerator: generateChunks(),
   });
@@ -64,7 +64,14 @@ export function mockFetchDataStreamWithGenerator({
   url: string;
   chunkGenerator: AsyncGenerator<Uint8Array, void, unknown>;
 }) {
-  vi.spyOn(global, 'fetch').mockImplementation(async () => {
+  let requestBodyResolve: ((value?: unknown) => void) | undefined;
+  const requestBodyPromise = new Promise(resolve => {
+    requestBodyResolve = resolve;
+  });
+
+  vi.spyOn(global, 'fetch').mockImplementation(async (url, init) => {
+    requestBodyResolve?.(init!.body as string);
+
     return {
       url,
       ok: true,
@@ -83,6 +90,10 @@ export function mockFetchDataStreamWithGenerator({
       },
     } as unknown as Response;
   });
+
+  return {
+    requestBody: requestBodyPromise,
+  };
 }
 
 export function mockFetchError({