@assistant-ui/ai-sdk 0.0.6 → 0.0.7

package/src/runtime/useAISDKRuntime.ts

@@ -14,7 +14,7 @@ import type {
   CreateUIMessage,
   UseChatHelpers,
 } from "@ai-sdk/react";
-import { isToolUIPart, generateId } from "ai";
+import { isToolUIPart, generateId, getToolName } from "ai";
 import {
   useExternalStoreRuntime,
   useRuntimeAdapters,
@@ -37,6 +37,7 @@ import type {
   AppendMessage,
   RunConfig,
   McpAppMetadata,
+  RespondToToolApprovalOptions,
 } from "@assistant-ui/core";
 import {
   getExternalStoreMessages,
@@ -52,7 +53,10 @@ import { sliceMessagesUntil } from "../utils/sliceMessagesUntil";
 import { toCreateMessage } from "../converters/toCreateMessage";
 import { vercelAttachmentAdapter } from "../adapters/vercelAttachmentAdapter";
 import { getVercelAIMessages } from "../utils/getVercelAIMessages";
-import { AISDKMessageConverter } from "../converters/convertMessage";
+import {
+  AISDKMessageConverter,
+  type AISDKMessageConverterMetadata,
+} from "../converters/convertMessage";
 import { wrapModelContentEnvelope } from "../converters/modelContentEnvelope";
 import {
   type AISDKStorageFormat,
@@ -81,6 +85,23 @@ const toUIMessage = <UI_MESSAGE extends UIMessage>(
     role: createMessage.role ?? fallbackRole,
   }) as UI_MESSAGE;
 
+const toVoiceTranscriptUIMessage = <UI_MESSAGE extends UIMessage>(
+  message: ThreadMessage,
+): UI_MESSAGE =>
+  ({
+    id: message.id,
+    role: message.role,
+    parts: message.content
+      .filter((part) => part.type === "text")
+      .map((part) => ({ type: "text", text: part.text })),
+    metadata: {
+      modality: "voice",
+      ...(Object.keys(message.metadata.custom).length > 0 && {
+        custom: message.metadata.custom,
+      }),
+    },
+  }) as UI_MESSAGE;
+
 export type AISDKRuntimeAdapter<UI_MESSAGE extends UIMessage = UIMessage> =
   ExternalStoreSharedOptions & {
     adapters?:
@@ -115,6 +136,24 @@ export type AISDKRuntimeAdapter<UI_MESSAGE extends UIMessage = UIMessage> =
      * Provide this to bridge resume-tool-call invocations into a custom handler.
      */
     onResumeToolCall?: ExternalStoreAdapter["onResumeToolCall"];
+    /**
+     * Answers tool approval requests through a host-owned channel instead of the AI SDK's `addToolApprovalResponse`.
+     *
+     * Called for every approval request in the thread with the complete response, including option and free-form answers. Hand requests the host does not own to `respondViaAISDK`, which is what runs when this option is omitted. The answer applies to the approval when the handler starts and is removed if it throws. It is never written into the `useChat` messages, so `sendAutomaticallyWhen` cannot forward it, and it lasts as long as this runtime: until then a second response to the same request rejects, and a runtime mounted again over the same chat shows the request open until the resumed run records its resolution in the chat.
+     *
+     * While a handler is set, an approval's `display`, `allowFreeform` and `options` reach the renderer, because the handler can receive answers the AI SDK cannot carry.
+     */
+    onRespondToToolApproval?:
+      | ((
+          response: RespondToToolApprovalOptions,
+          context: {
+            toolCallId: string;
+            toolName: string;
+            /** Sends this response through the AI SDK's `addToolApprovalResponse`, which carries only `approved` and `reason`. */
+            respondViaAISDK: () => Promise<void>;
+          },
+        ) => Promise<void> | void)
+      | undefined;
     /**
      * How consecutive assistant messages are rendered.
      *
@@ -220,6 +259,11 @@ const useGeneratedSuggestions = (
 
 const NO_CANCELLED_MESSAGE_IDS: ReadonlySet<string> = new Set();
 
+const NO_TOOL_APPROVAL_RESPONSES: ReadonlyMap<
+  string,
+  RespondToToolApprovalOptions
+> = new Map();
+
 const toChatError = (error: Error): AssistantError => {
   const code = (error as { code?: unknown }).code;
   return {
@@ -243,6 +287,7 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     cancelPendingToolCallsOnSend = true,
     onResume,
     onResumeToolCall,
+    onRespondToToolApproval: customOnRespondToToolApproval,
     joinStrategy,
     messageRepository,
     unstable_onBranchChange,
@@ -256,6 +301,10 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     chatId: string;
     ids: ReadonlySet<string>;
   } | null>(null);
+  const [toolApprovalResponses, setToolApprovalResponses] = useState<
+    ReadonlyMap<string, RespondToToolApprovalOptions>
+  >(NO_TOOL_APPROVAL_RESPONSES);
+  const hostApprovalIdsRef = useRef(new Set<string>());
   const toolArgsKeyOrderCacheRef = useRef<Map<string, Map<string, string[]>>>(
     new Map(),
   );
@@ -285,6 +334,22 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     cancelledMessages?.chatId === chatHelpers.id
       ? cancelledMessages.ids
       : NO_CANCELLED_MESSAGE_IDS;
+  const supportsRichToolApprovalResponses =
+    customOnRespondToToolApproval != null;
+
+  const toThreadMessages = useCallback(
+    (sourceMessages: UI_MESSAGE[]) => {
+      const metadata: AISDKMessageConverterMetadata = {
+        supportsRichToolApprovalResponses,
+      };
+      return AISDKMessageConverter.toThreadMessages(
+        sourceMessages,
+        false,
+        metadata,
+      );
+    },
+    [supportsRichToolApprovalResponses],
+  );
 
   const retractCancellation = useCallback(
     (chatId: string, messageId: string) => {
@@ -321,18 +386,20 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     isRunning,
     messages: chatHelpers.messages,
     joinStrategy,
-    metadata: useMemo(
+    metadata: useMemo<AISDKMessageConverterMetadata>(
       () => ({
         toolStatuses,
         messageTiming,
         toolArgsKeyOrderCache: toolArgsKeyOrderCacheRef.current,
         toolLastInputCache: toolLastInputCacheRef.current,
         mcpAppMetadataCache: mcpAppMetadataCacheRef.current,
+        supportsRichToolApprovalResponses,
         ...(optimisticMessageId && { optimisticMessageId }),
         ...(chatHelpers.error && {
           error: toChatError(chatHelpers.error),
         }),
         ...(cancelledMessageIds.size > 0 && { cancelledMessageIds }),
+        ...(toolApprovalResponses.size > 0 && { toolApprovalResponses }),
       }),
       [
         toolStatuses,
@@ -340,6 +407,8 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
         optimisticMessageId,
         chatHelpers.error,
         cancelledMessageIds,
+        toolApprovalResponses,
+        supportsRichToolApprovalResponses,
       ],
     ),
   });
@@ -347,13 +416,11 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
   const exportedMessageRepository = useMemo(() => {
     if (!messageRepository) return undefined;
     const converted = toExportedMessageRepository(
-      AISDKMessageConverter.toThreadMessages as (
-        messages: UI_MESSAGE[],
-      ) => ThreadMessage[],
+      toThreadMessages,
       messageRepository,
     );
     return converted.messages.length > 0 ? converted : undefined;
-  }, [messageRepository]);
+  }, [messageRepository, toThreadMessages]);
 
   const generatedSuggestions = useGeneratedSuggestions(
     suggestionAdapter,
@@ -370,9 +437,7 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
   const { isLoading, deleteMessage: deleteHistoryMessage } = useExternalHistory(
     runtimeRef,
     adapters?.history ?? contextAdapters?.history,
-    AISDKMessageConverter.toThreadMessages as (
-      messages: UI_MESSAGE[],
-    ) => ThreadMessage[],
+    toThreadMessages,
     aiSDKV6FormatAdapter as MessageFormatAdapter<
       UI_MESSAGE,
       AISDKStorageFormat
@@ -434,6 +499,69 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     });
   };
 
+  const respondViaAISDK = ({
+    approvalId,
+    approved,
+    reason,
+  }: RespondToToolApprovalOptions) =>
+    Promise.resolve(
+      chatHelpers.addToolApprovalResponse({
+        id: approvalId,
+        approved,
+        ...(reason != null && { reason }),
+        options: { metadata: lastRunConfigRef.current },
+      }),
+    );
+
+  const respondViaHost = async (
+    onRespond: NonNullable<AISDKRuntimeAdapter["onRespondToToolApproval"]>,
+    response: RespondToToolApprovalOptions,
+  ) => {
+    const { approvalId } = response;
+    const requested = chatHelpers.messages
+      .flatMap((message) => message.parts)
+      .filter(isToolUIPart)
+      .find(
+        (part) =>
+          part.state === "approval-requested" &&
+          part.approval.id === approvalId,
+      );
+    if (!requested || hostApprovalIdsRef.current.has(approvalId))
+      throw new Error(
+        `Tool approval ${approvalId} is not waiting for a response.`,
+      );
+
+    // A host answer stays out of the useChat messages, where sendAutomaticallyWhen would forward it to the chat route.
+    const applyResponse = (applied: boolean) => {
+      if (applied) hostApprovalIdsRef.current.add(approvalId);
+      else hostApprovalIdsRef.current.delete(approvalId);
+      setToolApprovalResponses((prev) => {
+        const responses = new Map(prev);
+        if (applied) responses.set(approvalId, response);
+        else responses.delete(approvalId);
+        return responses;
+      });
+    };
+
+    applyResponse(true);
+    try {
+      await onRespond(response, {
+        toolCallId: requested.toolCallId,
+        toolName: getToolName(requested),
+        respondViaAISDK: async () => {
+          try {
+            await respondViaAISDK(response);
+          } finally {
+            applyResponse(false);
+          }
+        },
+      });
+    } catch (error) {
+      if (hostApprovalIdsRef.current.has(approvalId)) applyResponse(false);
+      throw error;
+    }
+  };
+
   const hasSeededRepositoryRef = useRef(false);
   const shouldFeedRepository =
     exportedMessageRepository != null &&
@@ -461,6 +589,11 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
           .filter(Boolean)
           .flat(),
       ),
+    onVoiceTranscript: (message: ThreadMessage) =>
+      chatHelpers.setMessages((current) => [
+        ...current,
+        toVoiceTranscriptUIMessage<UI_MESSAGE>(message),
+      ]),
     onExportExternalState: (): MessageFormatRepository<UI_MESSAGE> => {
       const exported = runtimeRef.current.thread.export();
 
@@ -501,10 +634,7 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
     },
     onLoadExternalState: (repo: MessageFormatRepository<UI_MESSAGE>) => {
       // Convert MessageFormatRepository to ExportedMessageRepository
-      const exportedRepo = toExportedMessageRepository(
-        AISDKMessageConverter.toThreadMessages,
-        repo,
-      );
+      const exportedRepo = toExportedMessageRepository(toThreadMessages, repo);
 
       // Import into the thread's MessageRepository
       runtimeRef.current.thread.import(exportedRepo);
@@ -634,15 +764,9 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
         );
       }
     },
-    onRespondToToolApproval: ({ approvalId, approved, reason }) =>
-      Promise.resolve(
-        chatHelpers.addToolApprovalResponse({
-          id: approvalId,
-          approved,
-          ...(reason != null && { reason }),
-          options: { metadata: lastRunConfigRef.current },
-        }),
-      ),
+    onRespondToToolApproval: customOnRespondToToolApproval
+      ? (response) => respondViaHost(customOnRespondToToolApproval, response)
+      : respondViaAISDK,
     ...pickExternalStoreSharedOptions(adapter),
     ...(adapter.unstable_messageRepositoryInstance && {
       unstable_messageRepositoryInstance:

package/src/runtime/useAISDKRuntime.voice.test.tsx

@@ -0,0 +1,270 @@
+// @vitest-environment jsdom
+
+import { act, renderHook, waitFor } from "@testing-library/react";
+import {
+  getExternalStoreMessages,
+  type RealtimeVoiceAdapter,
+  type ThreadHistoryAdapter,
+  type ThreadMessage,
+} from "@assistant-ui/core";
+import type { UIMessage } from "ai";
+import { describe, expect, it, vi } from "vitest";
+
+import { useAISDKRuntime } from "./useAISDKRuntime";
+
+const createChatHelpers = (initialMessages: UIMessage[] = []) => {
+  let currentMessages: UIMessage[] = initialMessages;
+  const chatHelpers: any = {
+    id: "chat-1",
+    status: "ready",
+    error: null,
+    messages: currentMessages,
+    setMessages: vi.fn(
+      (next: UIMessage[] | ((current: UIMessage[]) => UIMessage[])) => {
+        currentMessages =
+          typeof next === "function" ? next(currentMessages) : [...next];
+        chatHelpers.messages = currentMessages;
+        return currentMessages;
+      },
+    ),
+    sendMessage: vi.fn().mockResolvedValue(undefined),
+    regenerate: vi.fn().mockResolvedValue(undefined),
+    addToolResult: vi.fn(),
+    addToolOutput: vi.fn(),
+    stop: vi.fn(),
+  };
+  return chatHelpers;
+};
+
+const createVoiceAdapter = () => {
+  let transcriptCallback:
+    | ((transcript: RealtimeVoiceAdapter.TranscriptItem) => void)
+    | undefined;
+  const session: RealtimeVoiceAdapter.Session = {
+    status: { type: "running" },
+    isMuted: false,
+    disconnect: vi.fn(),
+    mute: vi.fn(),
+    unmute: vi.fn(),
+    onStatusChange: () => () => {},
+    onTranscript: (callback) => {
+      transcriptCallback = callback;
+      return () => {
+        transcriptCallback = undefined;
+      };
+    },
+    onModeChange: () => () => {},
+    onVolumeChange: () => () => {},
+  };
+  return {
+    adapter: { connect: () => session } satisfies RealtimeVoiceAdapter,
+    emitTranscript: (transcript: RealtimeVoiceAdapter.TranscriptItem) =>
+      transcriptCallback?.(transcript),
+  };
+};
+
+const createHistoryAdapter = () => {
+  const append = vi.fn(async () => {});
+  const adapter: ThreadHistoryAdapter = {
+    load: vi.fn(),
+    append: vi.fn(),
+    withFormat: vi.fn().mockReturnValue({
+      load: vi.fn().mockResolvedValue({ headId: null, messages: [] }),
+      append,
+    }),
+  };
+  return { adapter, append };
+};
+
+const textOf = (message: ThreadMessage) =>
+  message.content
+    .map((part) => (part.type === "text" ? part.text : ""))
+    .join("");
+
+const renderVoiceRuntime = async (
+  chat: ReturnType<typeof createChatHelpers>,
+  adapters: { voice: RealtimeVoiceAdapter; history?: ThreadHistoryAdapter },
+) => {
+  const rendered = renderHook(() => useAISDKRuntime(chat, { adapters }));
+  await waitFor(() => {
+    expect(rendered.result.current.thread.getState().capabilities.voice).toBe(
+      true,
+    );
+  });
+  return rendered;
+};
+
+describe("useAISDKRuntime voice transcripts", () => {
+  it("persists finalized transcripts through the useChat messages", async () => {
+    const chat = createChatHelpers();
+    const voice = createVoiceAdapter();
+    const { result, rerender } = await renderVoiceRuntime(chat, {
+      voice: voice.adapter,
+    });
+
+    act(() => {
+      result.current.thread.connectVoice();
+      voice.emitTranscript({
+        role: "assistant",
+        text: "Spoken reply",
+        isFinal: true,
+      });
+    });
+
+    const transcript = chat.messages[0];
+    expect(transcript).toEqual({
+      id: expect.any(String),
+      role: "assistant",
+      parts: [{ type: "text", text: "Spoken reply" }],
+      metadata: { modality: "voice" },
+    });
+
+    rerender();
+
+    await waitFor(() => {
+      const messages = result.current.thread
+        .getState()
+        .messages.filter((message) => message.id === transcript.id);
+      expect(messages).toHaveLength(1);
+      expect(messages[0]?.metadata.modality).toBe("voice");
+      expect(getExternalStoreMessages(messages[0]!)).toEqual([transcript]);
+    });
+  });
+
+  it("keeps a transcript spoken after a typed reply as its own message", async () => {
+    const chat = createChatHelpers([
+      {
+        id: "typed-user",
+        role: "user",
+        parts: [{ type: "text", text: "Hi" }],
+      },
+      {
+        id: "typed-assistant",
+        role: "assistant",
+        parts: [{ type: "text", text: "Typed reply" }],
+      },
+    ]);
+    const voice = createVoiceAdapter();
+    const { result, rerender } = await renderVoiceRuntime(chat, {
+      voice: voice.adapter,
+    });
+
+    act(() => {
+      result.current.thread.connectVoice();
+      voice.emitTranscript({
+        role: "assistant",
+        text: "Spoken reply",
+        isFinal: true,
+      });
+    });
+    const transcript = chat.messages[2];
+    rerender();
+
+    await waitFor(() => {
+      const spoken = result.current.thread.getState().messages.at(-1)!;
+      expect(getExternalStoreMessages(spoken)).toEqual([transcript]);
+    });
+    const messages = result.current.thread.getState().messages;
+    expect(messages.map((message) => message.id)).toEqual([
+      "typed-user",
+      "typed-assistant",
+      transcript.id,
+    ]);
+    expect(textOf(messages[1]!)).toBe("Typed reply");
+    expect(messages[1]?.metadata.modality).toBeUndefined();
+    expect(textOf(messages[2]!)).toBe("Spoken reply");
+    expect(messages[2]?.metadata.modality).toBe("voice");
+  });
+
+  it("keeps consecutive assistant transcripts as separate messages", async () => {
+    const chat = createChatHelpers();
+    const voice = createVoiceAdapter();
+    const { result, rerender } = await renderVoiceRuntime(chat, {
+      voice: voice.adapter,
+    });
+
+    act(() => {
+      result.current.thread.connectVoice();
+      voice.emitTranscript({ role: "assistant", text: "First", isFinal: true });
+      voice.emitTranscript({
+        role: "assistant",
+        text: "Second",
+        isFinal: true,
+      });
+    });
+    const [first, second] = chat.messages;
+    rerender();
+
+    await waitFor(() => {
+      const messages = result.current.thread.getState().messages;
+      expect(
+        messages.map((message) => getExternalStoreMessages(message)),
+      ).toEqual([[first], [second]]);
+    });
+    expect(result.current.thread.getState().messages.map(textOf)).toEqual([
+      "First",
+      "Second",
+    ]);
+  });
+
+  it("keeps the transcript in the thread after the session disconnects", async () => {
+    const chat = createChatHelpers();
+    const voice = createVoiceAdapter();
+    const { result, rerender } = await renderVoiceRuntime(chat, {
+      voice: voice.adapter,
+    });
+
+    act(() => {
+      result.current.thread.connectVoice();
+      voice.emitTranscript({
+        role: "assistant",
+        text: "Spoken reply",
+        isFinal: true,
+      });
+    });
+    const transcript = chat.messages[0];
+    rerender();
+    await waitFor(() => {
+      const messages = result.current.thread.getState().messages;
+      expect(getExternalStoreMessages(messages[0]!)).toEqual([transcript]);
+    });
+
+    act(() => {
+      result.current.thread.disconnectVoice();
+    });
+    rerender();
+
+    const messages = result.current.thread.getState().messages;
+    expect(messages).toHaveLength(1);
+    expect(messages[0]?.id).toBe(transcript.id);
+    expect(messages[0]?.metadata.modality).toBe("voice");
+    expect(result.current.thread.getState().voice).toBeUndefined();
+  });
+
+  it("appends a finalized transcript to the history adapter without a text run", async () => {
+    const chat = createChatHelpers();
+    const voice = createVoiceAdapter();
+    const history = createHistoryAdapter();
+    const { result, rerender } = await renderVoiceRuntime(chat, {
+      voice: voice.adapter,
+      history: history.adapter,
+    });
+
+    act(() => {
+      result.current.thread.connectVoice();
+      voice.emitTranscript({
+        role: "assistant",
+        text: "Spoken reply",
+        isFinal: true,
+      });
+    });
+    const transcript = chat.messages[0];
+    rerender();
+
+    await waitFor(() => expect(history.append).toHaveBeenCalledTimes(1));
+    expect(history.append).toHaveBeenCalledWith({
+      parentId: null,
+      message: transcript,
+    });
+  });
+});

package/src/runtime/useChatRuntime.test.ts

@@ -53,6 +53,11 @@ const mocks = vi.hoisted(() => {
 
 vi.mock("@ai-sdk/react", () => ({
   useChat: mocks.useChat,
+  Chat: class MockChat {
+    constructor(config: unknown) {
+      Object.assign(this, config);
+    }
+  },
 }));
 
 vi.mock("@assistant-ui/core/react", async (importOriginal) => ({
@@ -95,6 +100,51 @@ describe("useChatRuntime", () => {
     window.sessionStorage.clear();
   });
 
+  it("forwards a callback through a ref, so a later render's callback fires instead of the mounted one", () => {
+    mocks.useChat.mockReturnValue({
+      resumeStream: vi.fn(),
+      status: "ready",
+    });
+
+    const onToolCallA = vi.fn();
+    const onToolCallB = vi.fn();
+
+    const { rerender } = renderHook(
+      ({ onToolCall }: { onToolCall: typeof onToolCallA }) =>
+        useChatRuntime({ onToolCall }),
+      { initialProps: { onToolCall: onToolCallA } },
+    );
+
+    const chat = mocks.useChat.mock.calls[0]?.[0]?.chat as {
+      onToolCall?: (arg: unknown) => void;
+      sendAutomaticallyWhen?: (arg: unknown) => boolean;
+    };
+
+    chat.onToolCall?.("first");
+    expect(onToolCallA).toHaveBeenCalledExactlyOnceWith("first");
+
+    rerender({ onToolCall: onToolCallB });
+    chat.onToolCall?.("second");
+
+    expect(onToolCallB).toHaveBeenCalledExactlyOnceWith("second");
+    expect(onToolCallA).toHaveBeenCalledOnce();
+  });
+
+  it("coerces an unset sendAutomaticallyWhen to false, matching useChat's own default", () => {
+    mocks.useChat.mockReturnValue({
+      resumeStream: vi.fn(),
+      status: "ready",
+    });
+
+    renderHook(() => useChatRuntime());
+
+    const chat = mocks.useChat.mock.calls[0]?.[0]?.chat as {
+      sendAutomaticallyWhen?: (arg: unknown) => boolean;
+    };
+
+    expect(chat.sendAutomaticallyWhen?.({})).toBe(false);
+  });
+
   it("forwards a defined chat update throttle to useChat", () => {
     mocks.useChat.mockReturnValue({
       resumeStream: vi.fn(),
@@ -111,6 +161,24 @@ describe("useChatRuntime", () => {
     expect(mocks.useChat.mock.calls[1]?.[0]).not.toHaveProperty("throttle");
   });
 
+  it("forwards a custom approval handler to the runtime only", () => {
+    const onRespondToToolApproval = vi.fn();
+    mocks.useChat.mockReturnValue({
+      resumeStream: vi.fn(),
+      status: "ready",
+    });
+
+    renderHook(() => useChatRuntime({ onRespondToToolApproval }));
+
+    expect(mocks.useAISDKRuntime).toHaveBeenCalledWith(
+      expect.anything(),
+      expect.objectContaining({ onRespondToToolApproval }),
+    );
+    expect(mocks.useChat.mock.calls[0]?.[0]).not.toHaveProperty(
+      "onRespondToToolApproval",
+    );
+  });
+
   it("waits for external history to load before resuming a stream", async () => {
     mocks.state.isLoadingHistory = true;
     const resumeStream = vi.fn().mockResolvedValue(undefined);
@@ -283,8 +351,8 @@ describe("useChatRuntime", () => {
       resumeStream: vi.fn().mockResolvedValue(undefined),
       status: "streaming",
     };
-    mocks.useChat.mockImplementation(({ id }: { id: string }) =>
-      id === "thread-a" ? threadA : threadB,
+    mocks.useChat.mockImplementation(({ chat }: { chat: { id: string } }) =>
+      chat.id === "thread-a" ? threadA : threadB,
     );
 
     mocks.state.threadId = "thread-a";
@@ -321,8 +389,8 @@ describe("useChatRuntime", () => {
       resumeStream: vi.fn().mockResolvedValue(undefined),
       status: "ready",
     };
-    mocks.useChat.mockImplementation(({ id }: { id: string }) =>
-      id === "__LOCALID_background" ? backgroundThread : mainThread,
+    mocks.useChat.mockImplementation(({ chat }: { chat: { id: string } }) =>
+      chat.id === "__LOCALID_background" ? backgroundThread : mainThread,
     );
     const transport = {
       getResumableAdapter: () => ({