@assistant-ui/ai-sdk 0.0.1

package/src/runtime/useAISDKRuntime.tool-output.test.tsx

@@ -0,0 +1,113 @@
+// @vitest-environment jsdom
+
+import { renderHook } from "@testing-library/react";
+import type { ExternalStoreAdapter } from "@assistant-ui/core";
+import { describe, expect, it, vi } from "vitest";
+
+const mocks = vi.hoisted(() => ({
+  adapter: undefined as ExternalStoreAdapter | undefined,
+}));
+
+vi.mock("@assistant-ui/core/react", async (importOriginal) => {
+  const original =
+    await importOriginal<typeof import("@assistant-ui/core/react")>();
+  return {
+    ...original,
+    useExternalStoreRuntime: vi.fn((adapter: ExternalStoreAdapter) => {
+      mocks.adapter = adapter;
+      return {};
+    }),
+    useRuntimeAdapters: vi.fn(() => ({})),
+  };
+});
+
+vi.mock("./useExternalHistory", async (importOriginal) => {
+  const original =
+    await importOriginal<typeof import("./useExternalHistory")>();
+  return {
+    ...original,
+    useExternalHistory: vi.fn(() => ({
+      isLoading: false,
+      deleteMessage: vi.fn().mockResolvedValue(undefined),
+    })),
+  };
+});
+
+import { useAISDKRuntime } from "./useAISDKRuntime";
+
+const createChatHelpers = (addToolOutput: ReturnType<typeof vi.fn>) => ({
+  id: "chat-1",
+  status: "ready",
+  error: undefined,
+  messages: [],
+  setMessages: vi.fn(),
+  sendMessage: vi.fn(),
+  regenerate: vi.fn(),
+  addToolOutput,
+  addToolApprovalResponse: vi.fn(),
+  stop: vi.fn(),
+});
+
+const callOnAddToolResult = (
+  options: Parameters<NonNullable<ExternalStoreAdapter["onAddToolResult"]>>[0],
+) => {
+  const callback = mocks.adapter?.onAddToolResult;
+  expect(callback).toBeDefined();
+  return callback!(options);
+};
+
+describe("useAISDKRuntime tool output callbacks", () => {
+  it("returns the addToolOutput promise for successful results", () => {
+    const outputPromise = Promise.resolve();
+    const addToolOutput = vi.fn(() => outputPromise);
+
+    renderHook(() =>
+      useAISDKRuntime(createChatHelpers(addToolOutput) as never),
+    );
+
+    const callbackPromise = callOnAddToolResult({
+      messageId: "message-1",
+      toolCallId: "tool-call-1",
+      toolName: "weather",
+      result: { temperature: 72 },
+      isError: false,
+    });
+
+    expect(callbackPromise).toBe(outputPromise);
+    expect(addToolOutput).toHaveBeenCalledWith({
+      tool: "weather",
+      toolCallId: "tool-call-1",
+      output: { temperature: 72 },
+      options: { metadata: undefined },
+    });
+  });
+
+  it("returns the addToolOutput promise for error results", async () => {
+    const error = new Error("tool failed");
+    const outputPromise = Promise.reject(error);
+    void outputPromise.catch(() => undefined);
+    const addToolOutput = vi.fn(() => outputPromise);
+
+    renderHook(() =>
+      useAISDKRuntime(createChatHelpers(addToolOutput) as never),
+    );
+
+    const callbackPromise = callOnAddToolResult({
+      messageId: "message-1",
+      toolCallId: "tool-call-1",
+      toolName: "weather",
+      result: "tool failed",
+      isError: true,
+    });
+
+    expect(callbackPromise).toBe(outputPromise);
+    expect(addToolOutput).toHaveBeenCalledWith({
+      state: "output-error",
+      tool: "weather",
+      toolCallId: "tool-call-1",
+      errorText: "tool failed",
+      options: { metadata: undefined },
+    });
+    await expect(outputPromise).rejects.toBe(error);
+  });
+});

package/src/runtime/useAISDKRuntime.ts

@@ -0,0 +1,524 @@
+"use client";
+
+import { useEffect, useMemo, useRef, useState } from "react";
+import type {
+  UIMessage,
+  useChat,
+  CreateUIMessage,
+  UseChatHelpers,
+} from "@ai-sdk/react";
+import { isToolUIPart, generateId } from "ai";
+import {
+  useExternalStoreRuntime,
+  useRuntimeAdapters,
+  type JoinStrategy,
+} from "@assistant-ui/core/react";
+import type {
+  SuggestionAdapter,
+  ThreadSuggestion,
+  ToolExecutionStatus,
+} from "@assistant-ui/core";
+import type {
+  ExternalStoreAdapter,
+  ExternalStoreSharedOptions,
+  ThreadHistoryAdapter,
+  AssistantRuntime,
+  ThreadMessage,
+  MessageFormatAdapter,
+  MessageFormatItem,
+  MessageFormatRepository,
+  AppendMessage,
+  RunConfig,
+  McpAppMetadata,
+} from "@assistant-ui/core";
+import {
+  getExternalStoreMessages,
+  pickExternalStoreSharedOptions,
+} from "@assistant-ui/core";
+import { consumeSuggestionResult } from "@assistant-ui/core/internal";
+import type { ReadonlyJSONObject } from "assistant-stream/utils";
+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 { wrapModelContentEnvelope } from "../converters/modelContentEnvelope";
+import {
+  type AISDKStorageFormat,
+  aiSDKV6FormatAdapter,
+} from "../adapters/aiSDKFormatAdapter";
+import {
+  useExternalHistory,
+  toExportedMessageRepository,
+} from "./useExternalHistory";
+import { useStreamingTiming } from "./useStreamingTiming";
+import { aiSDKExtras } from "../aiSDKExtras";
+
+export type CustomToCreateMessageFunction = <
+  UI_MESSAGE extends UIMessage = UIMessage,
+>(
+  message: AppendMessage,
+) => CreateUIMessage<UI_MESSAGE>;
+
+const toUIMessage = <UI_MESSAGE extends UIMessage>(
+  createMessage: CreateUIMessage<UI_MESSAGE>,
+  fallbackRole: UI_MESSAGE["role"],
+): UI_MESSAGE =>
+  ({
+    ...createMessage,
+    id: createMessage.id ?? generateId(),
+    role: createMessage.role ?? fallbackRole,
+  }) as UI_MESSAGE;
+
+export type AISDKRuntimeAdapter = ExternalStoreSharedOptions & {
+  adapters?:
+    | (NonNullable<ExternalStoreAdapter["adapters"]> & {
+        history?: ThreadHistoryAdapter | undefined;
+        suggestion?: SuggestionAdapter | undefined;
+      })
+    | undefined;
+  toCreateMessage?: CustomToCreateMessageFunction;
+  /**
+   * Whether to automatically cancel pending interactive tool calls when the user sends a new message.
+   *
+   * When enabled (default), the pending tool calls will be marked as failed with an error message
+   * indicating the user cancelled the tool call by sending a new message.
+   *
+   * @default true
+   */
+  cancelPendingToolCallsOnSend?: boolean | undefined;
+  /**
+   * Called when `runtime.thread.resumeRun(config)` is invoked.
+   *
+   * When omitted, `resumeRun` throws `"Runtime does not support resuming runs."`.
+   * Provide this to bridge resume invocations into a custom replay channel
+   * (for example, an SSE reconnect endpoint keyed by turn id).
+   */
+  onResume?: ExternalStoreAdapter["onResume"];
+  /**
+   * Called when `runtime.thread.resumeToolCall(options)` is invoked for a tool call the in-process tracker does not own.
+   *
+   * When omitted, `resumeToolCall` throws `"Tool call ${toolCallId} is not waiting for resume."`.
+   * Provide this to bridge resume-tool-call invocations into a custom handler.
+   */
+  onResumeToolCall?: ExternalStoreAdapter["onResumeToolCall"];
+  /**
+   * How consecutive assistant messages are rendered.
+   *
+   * `"concat-content"` (the default) merges them into a single thread message.
+   * `"none"` keeps each assistant message as its own thread message, which is
+   * useful when a backend persists proactive or consecutive assistant messages
+   * as separate entries.
+   */
+  joinStrategy?: JoinStrategy | undefined;
+};
+
+const EMPTY_SUGGESTIONS: readonly ThreadSuggestion[] = [];
+
+const useGeneratedSuggestions = (
+  suggestionAdapter: SuggestionAdapter | undefined,
+  messages: readonly ThreadMessage[],
+  isRunning: boolean,
+): readonly ThreadSuggestion[] => {
+  const [suggestions, setSuggestions] =
+    useState<readonly ThreadSuggestion[]>(EMPTY_SUGGESTIONS);
+  const controllerRef = useRef<AbortController | null>(null);
+  const wasRunningRef = useRef(false);
+  const messagesRef = useRef(messages);
+  messagesRef.current = messages;
+  const adapterRef = useRef(suggestionAdapter);
+  adapterRef.current = suggestionAdapter;
+  const hasAdapter = suggestionAdapter != null;
+
+  useEffect(() => {
+    const clearSuggestions = () => {
+      controllerRef.current?.abort();
+      controllerRef.current = null;
+      setSuggestions((prev) => (prev.length === 0 ? prev : EMPTY_SUGGESTIONS));
+    };
+
+    const adapter = adapterRef.current;
+    if (!adapter) {
+      clearSuggestions();
+      wasRunningRef.current = isRunning;
+      return;
+    }
+
+    if (isRunning) {
+      if (!wasRunningRef.current) {
+        clearSuggestions();
+      }
+      wasRunningRef.current = true;
+      return;
+    }
+
+    if (!wasRunningRef.current) return;
+    wasRunningRef.current = false;
+
+    const currentMessages = messagesRef.current;
+    const last = currentMessages.at(-1);
+    if (last?.role !== "assistant") return;
+    if (last.status?.type === "requires-action") return;
+
+    const controller = new AbortController();
+    controllerRef.current = controller;
+    const { signal } = controller;
+
+    void (async () => {
+      try {
+        const promiseOrGenerator = adapter.generate({
+          messages: currentMessages,
+          signal,
+        });
+
+        await consumeSuggestionResult(promiseOrGenerator, {
+          signal,
+          onUpdate: setSuggestions,
+        });
+      } catch {}
+    })();
+  }, [hasAdapter, isRunning]);
+
+  useEffect(() => {
+    return () => {
+      controllerRef.current?.abort();
+    };
+  }, []);
+
+  return suggestions;
+};
+
+export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
+  chatHelpers: ReturnType<typeof useChat<UI_MESSAGE>>,
+  adapter: AISDKRuntimeAdapter = {},
+) => {
+  const {
+    adapters,
+    toCreateMessage: customToCreateMessage,
+    cancelPendingToolCallsOnSend = true,
+    onResume,
+    onResumeToolCall,
+    joinStrategy,
+  } = adapter;
+  const suggestionAdapter = adapters?.suggestion;
+  const contextAdapters = useRuntimeAdapters();
+  const [toolStatuses, setToolStatuses] = useState<
+    Record<string, ToolExecutionStatus>
+  >({});
+  const toolArgsKeyOrderCacheRef = useRef<Map<string, Map<string, string[]>>>(
+    new Map(),
+  );
+  const toolLastInputCacheRef = useRef<Map<string, ReadonlyJSONObject>>(
+    new Map(),
+  );
+  const mcpAppMetadataCacheRef = useRef<Map<string, McpAppMetadata>>(new Map());
+  const lastRunConfigRef = useRef<RunConfig | undefined>(undefined);
+
+  const hasExecutingTools = Object.values(toolStatuses).some(
+    (s) => s?.type === "executing",
+  );
+  const isRunning =
+    chatHelpers.status === "submitted" ||
+    chatHelpers.status === "streaming" ||
+    hasExecutingTools;
+
+  const messageTiming = useStreamingTiming(chatHelpers.messages, isRunning);
+
+  // Flag the streaming message optimistic: its id can be swapped for a server
+  // id mid-run, and the repository then drops the orphaned pre-swap id (#4037).
+  const lastMessage = chatHelpers.messages.at(-1);
+  const optimisticMessageId =
+    isRunning && lastMessage?.role === "assistant" ? lastMessage.id : undefined;
+
+  const messages = AISDKMessageConverter.useThreadMessages({
+    isRunning,
+    messages: chatHelpers.messages,
+    joinStrategy,
+    metadata: useMemo(
+      () => ({
+        toolStatuses,
+        messageTiming,
+        toolArgsKeyOrderCache: toolArgsKeyOrderCacheRef.current,
+        toolLastInputCache: toolLastInputCacheRef.current,
+        mcpAppMetadataCache: mcpAppMetadataCacheRef.current,
+        ...(optimisticMessageId && { optimisticMessageId }),
+        ...(chatHelpers.error && { error: chatHelpers.error.message }),
+      }),
+      [toolStatuses, messageTiming, optimisticMessageId, chatHelpers.error],
+    ),
+  });
+
+  const generatedSuggestions = useGeneratedSuggestions(
+    suggestionAdapter,
+    messages,
+    isRunning,
+  );
+
+  const [runtimeRef] = useState(() => ({
+    get current(): AssistantRuntime {
+      return runtime;
+    },
+  }));
+
+  const { isLoading, deleteMessage: deleteHistoryMessage } = useExternalHistory(
+    runtimeRef,
+    adapters?.history ?? contextAdapters?.history,
+    AISDKMessageConverter.toThreadMessages as (
+      messages: UI_MESSAGE[],
+    ) => ThreadMessage[],
+    aiSDKV6FormatAdapter as MessageFormatAdapter<
+      UI_MESSAGE,
+      AISDKStorageFormat
+    >,
+    (messages) => {
+      chatHelpers.setMessages(messages);
+    },
+  );
+
+  const {
+    id: chatId,
+    messages: chatMessages,
+    status: chatStatus,
+    error,
+  } = chatHelpers;
+  const extras = useMemo(
+    () =>
+      aiSDKExtras.provide({
+        chat: chatHelpers as unknown as UseChatHelpers<UIMessage>,
+        error,
+      }),
+    // oxlint-disable-next-line react/exhaustive-deps -- keyed on the chat's identity and reactive snapshots; useChat re-mints the helpers object every render while its remaining fields are instance-bound methods, and a render-stable extras identity is what lets the external-store core dedupe adapter updates
+    [chatId, chatMessages, chatStatus, error],
+  );
+
+  const completePendingToolCalls = async () => {
+    if (!cancelPendingToolCallsOnSend) return;
+
+    // The runtime auto-aborts in-flight tool invocations when a new run
+    // is dispatched (append() / startRun()). All we need to do here is
+    // mark any tool without a result as cancelled in the UI message list.
+
+    // Mark any tool without a result as cancelled (uses setMessages to avoid triggering sendAutomaticallyWhen)
+    chatHelpers.setMessages((messages) => {
+      const lastMessage = messages.at(-1);
+      if (lastMessage?.role !== "assistant") return messages;
+
+      let hasChanges = false;
+      const parts = lastMessage.parts?.map((part) => {
+        if (!isToolUIPart(part)) return part;
+        if (part.state === "output-available" || part.state === "output-error")
+          return part;
+
+        hasChanges = true;
+        const { approval: _approval, ...rest } = part;
+        return {
+          ...rest,
+          state: "output-error" as const,
+          errorText: "User cancelled tool call by sending a new message.",
+        };
+      });
+
+      if (!hasChanges) return messages;
+      return [...messages.slice(0, -1), { ...lastMessage, parts }];
+    });
+  };
+
+  const runtime = useExternalStoreRuntime({
+    isRunning,
+    messages,
+    unstable_enableToolInvocations: true,
+    setToolStatuses,
+    setMessages: (messages) =>
+      chatHelpers.setMessages(
+        messages
+          .map(getVercelAIMessages<UI_MESSAGE>)
+          .filter(Boolean)
+          .flat(),
+      ),
+    onImport: (messages) =>
+      chatHelpers.setMessages(
+        messages
+          .map(getVercelAIMessages<UI_MESSAGE>)
+          .filter(Boolean)
+          .flat(),
+      ),
+    onExportExternalState: (): MessageFormatRepository<UI_MESSAGE> => {
+      const exported = runtimeRef.current.thread.export();
+
+      const expandedMessages: MessageFormatItem<UI_MESSAGE>[] = [];
+      const lastInnerIdMap = new Map<string, string>();
+
+      for (const item of exported.messages) {
+        const innerMessages = getExternalStoreMessages<UI_MESSAGE>(
+          item.message,
+        );
+        let parentId =
+          item.parentId != null
+            ? (lastInnerIdMap.get(item.parentId) ?? item.parentId)
+            : null;
+        for (const innerMessage of innerMessages) {
+          expandedMessages.push({ parentId, message: innerMessage });
+          parentId = aiSDKV6FormatAdapter.getId(innerMessage as UIMessage);
+        }
+        if (innerMessages.length > 0) {
+          lastInnerIdMap.set(
+            item.message.id,
+            aiSDKV6FormatAdapter.getId(
+              innerMessages[innerMessages.length - 1]! as UIMessage,
+            ),
+          );
+        }
+      }
+
+      const result: MessageFormatRepository<UI_MESSAGE> = {
+        messages: expandedMessages,
+      };
+
+      if (exported.headId != null) {
+        result.headId = lastInnerIdMap.get(exported.headId) ?? exported.headId;
+      }
+
+      return result;
+    },
+    onLoadExternalState: (repo: MessageFormatRepository<UI_MESSAGE>) => {
+      // Convert MessageFormatRepository to ExportedMessageRepository
+      const exportedRepo = toExportedMessageRepository(
+        AISDKMessageConverter.toThreadMessages,
+        repo,
+      );
+
+      // Import into the thread's MessageRepository
+      runtimeRef.current.thread.import(exportedRepo);
+    },
+    onCancel: async () => {
+      try {
+        await chatHelpers.stop();
+      } catch (error) {
+        if (!(error instanceof Error && error.name === "AbortError")) {
+          throw error;
+        }
+      }
+    },
+    onNew: async (message) => {
+      const createMessage = (
+        customToCreateMessage ?? toCreateMessage
+      )<UI_MESSAGE>(message);
+
+      if (!(message.startRun ?? message.role === "user")) {
+        chatHelpers.setMessages((current) => [
+          ...current,
+          toUIMessage<UI_MESSAGE>(createMessage, message.role),
+        ]);
+        return;
+      }
+
+      lastRunConfigRef.current = message.runConfig;
+      await completePendingToolCalls();
+      await chatHelpers.sendMessage(createMessage, {
+        metadata: message.runConfig,
+      });
+    },
+    onEdit: async (message) => {
+      const createMessage = (
+        customToCreateMessage ?? toCreateMessage
+      )<UI_MESSAGE>(message);
+
+      if (!(message.startRun ?? message.role === "user")) {
+        chatHelpers.setMessages((current) => [
+          ...sliceMessagesUntil(current, message.parentId),
+          toUIMessage<UI_MESSAGE>(createMessage, message.role),
+        ]);
+        return;
+      }
+
+      lastRunConfigRef.current = message.runConfig;
+      chatHelpers.setMessages((current) =>
+        sliceMessagesUntil(current, message.parentId),
+      );
+      await chatHelpers.sendMessage(createMessage, {
+        metadata: message.runConfig,
+      });
+    },
+    onDelete: async (messageId) => {
+      const threadMessages = runtimeRef.current.thread.getState().messages;
+      const messageIndex = threadMessages.findIndex(
+        (message) => message.id === messageId,
+      );
+      if (messageIndex === -1) return;
+
+      await deleteHistoryMessage(messageId);
+
+      const deleteIds = new Set(
+        getExternalStoreMessages<UI_MESSAGE>(threadMessages[messageIndex]!).map(
+          (message) => message.id,
+        ),
+      );
+      chatHelpers.setMessages((current) =>
+        current.filter((message) => !deleteIds.has(message.id)),
+      );
+    },
+    onReload: async (parentId: string | null, config) => {
+      lastRunConfigRef.current = config.runConfig;
+      const newMessages = sliceMessagesUntil(chatHelpers.messages, parentId);
+      chatHelpers.setMessages(newMessages);
+
+      await chatHelpers.regenerate({ metadata: config.runConfig });
+    },
+    onAddToolResult: ({
+      toolCallId,
+      toolName,
+      result,
+      isError,
+      modelContent,
+    }) => {
+      const options = { metadata: lastRunConfigRef.current };
+      if (isError) {
+        return Promise.resolve(
+          chatHelpers.addToolOutput({
+            state: "output-error",
+            tool: toolName ?? toolCallId,
+            toolCallId,
+            errorText:
+              typeof result === "string" ? result : JSON.stringify(result),
+            options,
+          }),
+        );
+      } else {
+        const output =
+          modelContent !== undefined
+            ? wrapModelContentEnvelope(result, modelContent)
+            : result;
+        return Promise.resolve(
+          chatHelpers.addToolOutput({
+            tool: toolName,
+            toolCallId,
+            output,
+            options,
+          }),
+        );
+      }
+    },
+    onRespondToToolApproval: ({ approvalId, approved, reason }) =>
+      Promise.resolve(
+        chatHelpers.addToolApprovalResponse({
+          id: approvalId,
+          approved,
+          ...(reason != null && { reason }),
+          options: { metadata: lastRunConfigRef.current },
+        }),
+      ),
+    ...pickExternalStoreSharedOptions(adapter),
+    ...(suggestionAdapter ? { suggestions: generatedSuggestions } : {}),
+    ...(onResume && { onResume }),
+    ...(onResumeToolCall && { onResumeToolCall }),
+    adapters: {
+      attachments: vercelAttachmentAdapter,
+      ...contextAdapters,
+      ...adapters,
+    },
+    extras,
+    isLoading,
+  });
+
+  return runtime;
+};

package/src/runtime/useChatOnTap.test.ts

@@ -0,0 +1,58 @@
+import { describe, expect, it, vi } from "vitest";
+import { createTapRoot } from "@assistant-ui/tap";
+import { useChat } from "@ai-sdk/react";
+import { createControlledTransport } from "./__tests__/controlled-transport";
+
+describe("useChat hosted in a tap root", () => {
+  it("streams a full round trip without React", async () => {
+    const { transport, emit, close } = createControlledTransport();
+    const root = createTapRoot(function UseChatOnTap() {
+      return useChat({ transport });
+    });
+
+    try {
+      expect(root.getValue().status).toBe("ready");
+      expect(root.getValue().messages).toHaveLength(0);
+
+      const sendPromise = root.getValue().sendMessage({ text: "hi" });
+      await vi.waitFor(() => {
+        expect(root.getValue().status).toBe("submitted");
+        expect(root.getValue().messages).toHaveLength(1);
+      });
+
+      emit(
+        { type: "start" },
+        { type: "text-start", id: "t1" },
+        { type: "text-delta", id: "t1", delta: "hello " },
+      );
+      await vi.waitFor(() => {
+        expect(root.getValue().status).toBe("streaming");
+        const parts = root.getValue().messages.at(-1)?.parts;
+        expect(parts).toContainEqual(
+          expect.objectContaining({ type: "text", text: "hello " }),
+        );
+      });
+
+      emit(
+        { type: "text-delta", id: "t1", delta: "world" },
+        { type: "text-end", id: "t1" },
+        { type: "finish" },
+      );
+      close();
+      await sendPromise;
+
+      await vi.waitFor(() => {
+        expect(root.getValue().status).toBe("ready");
+        const { messages } = root.getValue();
+        expect(messages).toHaveLength(2);
+        expect(messages[0]!.role).toBe("user");
+        expect(messages[1]!.role).toBe("assistant");
+        expect(messages[1]!.parts).toContainEqual(
+          expect.objectContaining({ type: "text", text: "hello world" }),
+        );
+      });
+    } finally {
+      root.unmount();
+    }
+  });
+});