@assistant-ui/ai-sdk 0.0.7 → 0.0.9

package/src/converters/convertMessage.ts

@@ -24,6 +24,7 @@ import {
   type McpAppMetadata,
   type MessagePartStreamStatus,
   type RespondToToolApprovalOptions,
+  type Unstable_ToolInteractionLog,
 } from "@assistant-ui/core";
 import { stableStringifyToolArgs } from "@assistant-ui/core/internal";
 import {
@@ -49,7 +50,7 @@ const THREAD_METADATA_KEYS = new Set([
 const toThreadMetadata = (metadata: unknown): MessageMetadata => {
   if (!metadata || typeof metadata !== "object") return undefined;
   const result: Record<string, unknown> = {};
-  const extra: Record<string, unknown> = {};
+  const extra = Object.create(null) as Record<string, unknown>;
   for (const [key, value] of Object.entries(metadata)) {
     (THREAD_METADATA_KEYS.has(key) ? result : extra)[key] = value;
   }
@@ -61,16 +62,30 @@ const toThreadMetadata = (metadata: unknown): MessageMetadata => {
 export type AISDKMessageConverterMetadata =
   useExternalMessageConverter.Metadata & {
     toolArgsKeyOrderCache?: Map<string, Map<string, string[]>>;
+    /**
+     * Frozen `argsText` keyed weakly by a settled tool call's input object, then
+     * by call, since the text carries the call's streamed key order. A known
+     * call/input pair skips serialization; the entries become collectible once
+     * the input is unreachable. A fresh input object re-serializes in its own
+     * deterministic key order.
+     */
+    toolArgsTextCache?: WeakMap<ReadonlyJSONObject, Map<string, string>>;
     toolLastInputCache?: Map<string, ReadonlyJSONObject>;
     mcpAppMetadataCache?: Map<string, McpAppMetadata>;
+    toolArtifacts?: ReadonlyMap<string, unknown>;
+    toolInteractions?: ReadonlyMap<string, Unstable_ToolInteractionLog>;
     supportsRichToolApprovalResponses?: boolean;
     toolApprovalResponses?: ReadonlyMap<string, RespondToToolApprovalOptions>;
+    cancelledToolApprovalIds?: ReadonlySet<string>;
+    cancelledStatusMessageIds?: ReadonlySet<string>;
     /** Id of the currently-streaming message, flagged optimistic (#4037). */
     optimisticMessageId?: string | undefined;
   };
 
 function stripClosingDelimiters(json: string): string {
-  return json.replace(/[}\]"]+$/, "");
+  let end = json.length;
+  while (end > 0 && '}]"'.includes(json[end - 1]!)) end--;
+  return json.slice(0, end);
 }
 
 const MCP_APP_METADATA_CACHE_MAX = 100;
@@ -206,6 +221,38 @@ const normalizeToolApprovalOptions = (
   });
 };
 
+const APPROVAL_DESCRIPTOR_FIELDS = [
+  "prompt",
+  "display",
+  "allowFreeform",
+  "dismissible",
+  "options",
+  "optionId",
+  "text",
+  "resolution",
+] as const;
+
+// The AI SDK's approval object declares none of the core request and answer
+// fields and `validateUIMessages` strips unknown ones, so a host streams or
+// persists them inside the opaque `approvalDescriptor`. Only those fields are
+// read from it: a descriptor cannot approve its own request.
+const readApprovalDescriptor = (
+  descriptor: unknown,
+): Record<string, unknown> => {
+  if (
+    !descriptor ||
+    typeof descriptor !== "object" ||
+    Array.isArray(descriptor)
+  )
+    return {};
+  const fields: Record<string, unknown> = {};
+  for (const key of APPROVAL_DESCRIPTOR_FIELDS) {
+    if (Object.hasOwn(descriptor, key))
+      fields[key] = (descriptor as Record<string, unknown>)[key];
+  }
+  return fields;
+};
+
 function getToolApprovalAndInterrupt(
   part: {
     approval?: Record<string, unknown> | undefined;
@@ -215,17 +262,29 @@ function getToolApprovalAndInterrupt(
   toolApprovalResponses:
     | ReadonlyMap<string, RespondToToolApprovalOptions>
     | undefined,
+  cancelledToolApprovalIds: ReadonlySet<string> | undefined,
 ): {
   approval?: NonNullable<ToolCallMessagePart["approval"]>;
   interrupt?: NonNullable<ToolCallMessagePart["interrupt"]>;
 } {
   if (part.approval) {
+    const rawApproval = {
+      ...readApprovalDescriptor(part.approval.descriptor),
+      ...part.approval,
+    };
+    const approval =
+      typeof rawApproval.id === "string" &&
+      rawApproval.approved === undefined &&
+      rawApproval.resolution === undefined &&
+      cancelledToolApprovalIds?.has(rawApproval.id)
+        ? { ...rawApproval, resolution: "cancelled" as const }
+        : rawApproval;
     const response =
-      typeof part.approval.id === "string" &&
-      part.approval.approved === undefined &&
-      part.approval.resolution !== "cancelled" &&
-      part.approval.resolution !== "expired"
-        ? toolApprovalResponses?.get(part.approval.id)
+      typeof approval.id === "string" &&
+      approval.approved === undefined &&
+      approval.resolution !== "cancelled" &&
+      approval.resolution !== "expired"
+        ? toolApprovalResponses?.get(approval.id)
         : undefined;
     // The built-in AI SDK channel sends only id, approved and reason back to
     // the server, so a request shape promising any other answer would render
@@ -239,19 +298,20 @@ function getToolApprovalAndInterrupt(
       resolution,
       display,
       allowFreeform,
+      dismissible,
       options,
       optionId,
       text,
       ...additionalApprovalFields
     } = response
       ? {
-          ...part.approval,
+          ...approval,
           approved: response.approved,
           ...(response.reason != null && { reason: response.reason }),
           ...(response.optionId != null && { optionId: response.optionId }),
           ...(response.text != null && { text: response.text }),
         }
-      : part.approval;
+      : approval;
     const normalizedOptions = supportsRichToolApprovalResponses
       ? normalizeToolApprovalOptions(options)
       : undefined;
@@ -274,6 +334,7 @@ function getToolApprovalAndInterrupt(
               display === "select" ||
               display === "text") && { display }),
             ...(typeof allowFreeform === "boolean" && { allowFreeform }),
+            ...(typeof dismissible === "boolean" && { dismissible }),
             ...(normalizedOptions && { options: normalizedOptions }),
             ...(typeof optionId === "string" && { optionId }),
             ...(typeof text === "string" && { text }),
@@ -354,7 +415,14 @@ function convertParts(
         const toolCallId = part.toolCallId;
         const argsKeyOrderCacheKey = `${message.id}:${toolCallId}`;
 
-        const rawInput = part.input as ReadonlyJSONObject | null | undefined;
+        // A tool call that streamed complete arguments then failed schema
+        // validation keeps them in `rawInput`, not `input`; reading `input`
+        // alone would convert the error snapshot to `{}` and hide the input.
+        const rawInput = (part.input ??
+          ("rawInput" in part ? part.rawInput : undefined)) as
+          | ReadonlyJSONObject
+          | null
+          | undefined;
         let args: ReadonlyJSONObject;
         if (
           rawInput != null &&
@@ -387,12 +455,13 @@ function convertParts(
           };
         }
 
-        let argsText = stableStringifyToolArgs(
-          metadata.toolArgsKeyOrderCache,
-          argsKeyOrderCacheKey,
-          args,
-        );
+        let argsText: string;
         if (part.state === "input-streaming") {
+          argsText = stableStringifyToolArgs(
+            metadata.toolArgsKeyOrderCache,
+            argsKeyOrderCacheKey,
+            args,
+          );
           // strip closing delimiters added by the AI SDK's fix-json
           argsText = stripClosingDelimiters(argsText);
           // Re-parse so args carries the partial-JSON meta that marks which
@@ -401,6 +470,27 @@ function convertParts(
           // of the stripped text is the streaming frontier.
           args = parsePartialJsonObject(argsText) ?? args;
         } else {
+          // A settled part is re-converted whenever its message or the converter
+          // metadata changes; the text frozen on its input object skips
+          // re-serializing large args while the call keeps that input. Arrival
+          // order only matters while args stream, so the key-order entry is
+          // released.
+          const frozen =
+            metadata.toolArgsTextCache?.get(args) ?? new Map<string, string>();
+          const frozenText = frozen.get(argsKeyOrderCacheKey);
+          if (frozenText !== undefined) {
+            argsText = frozenText;
+          } else {
+            argsText = stableStringifyToolArgs(
+              metadata.toolArgsKeyOrderCache,
+              argsKeyOrderCacheKey,
+              args,
+            );
+            metadata.toolArgsTextCache?.set(
+              args,
+              frozen.set(argsKeyOrderCacheKey, argsText),
+            );
+          }
           metadata.toolArgsKeyOrderCache?.delete(argsKeyOrderCacheKey);
           if (
             part.state === "output-available" ||
@@ -416,6 +506,8 @@ function convertParts(
           part,
           metadata.mcpAppMetadataCache,
         );
+        const artifact = metadata.toolArtifacts?.get(toolCallId);
+        const interactions = metadata.toolInteractions?.get(toolCallId);
         return {
           type: "tool-call",
           toolName,
@@ -424,6 +516,12 @@ function convertParts(
           args,
           result,
           isError,
+          ...(artifact !== undefined && { artifact }),
+          ...(interactions !== undefined && {
+            unstable_interactions: interactions,
+          }),
+          ...(part.state === "output-available" &&
+            part.preliminary === true && { isPreliminary: true }),
           ...(modelContent !== undefined && { modelContent }),
           ...(mcpApp && { mcp: { app: mcpApp } }),
           ...(part.callProviderMetadata != null
@@ -437,6 +535,7 @@ function convertParts(
             toolStatus,
             metadata.supportsRichToolApprovalResponses === true,
             metadata.toolApprovalResponses,
+            metadata.cancelledToolApprovalIds,
           ),
         } satisfies ToolCallMessagePart;
       }
@@ -575,6 +674,15 @@ export const AISDKMessageConverter = unstable_createMessageConverter(
           id: message.id,
           createdAt,
           content,
+          ...(message.role === "assistant" &&
+          metadata.cancelledStatusMessageIds?.has(message.id)
+            ? {
+                status: {
+                  type: "incomplete" as const,
+                  reason: "cancelled" as const,
+                },
+              }
+            : undefined),
           metadata: {
             ...toThreadMetadata(message.metadata),
             ...(timing && { timing }),

package/src/index.native.ts

@@ -1,4 +1,4 @@
-/// <reference types="@assistant-ui/core/react" />
+/// <reference types="@assistant-ui/core/react" preserve="true" />
 
 export { useAISDKRuntime } from "./runtime/useAISDKRuntime";
 export { useChatRuntime } from "./runtime/useChatRuntime";

package/src/index.ts

@@ -1,4 +1,4 @@
-/// <reference types="@assistant-ui/core/react" />
+/// <reference types="@assistant-ui/core/react" preserve="true" />
 
 export { useAISDKRuntime } from "./runtime/useAISDKRuntime";
 export { useChatRuntime } from "./runtime/useChatRuntime";

package/src/runtime/AISDKChat.ts

@@ -18,10 +18,6 @@ const useAISDKChat = <UI_MESSAGE extends UIMessage = UIMessage>(
   options?: AISDKChatOptions<UI_MESSAGE>,
 ) => {
   const [id] = useState(() => options?.id ?? generateId());
-  // The transport resolves the request id from the thread list item, falling
-  // back to the runtime's main item, whose id here is the external store's
-  // placeholder constant. The single thread of this entry is the chat itself,
-  // so the handed-over item initializes to the chat id.
   const [threadListItem] = useState(() => ({
     initialize: async () => ({ remoteId: id, externalId: undefined }),
   }));

package/src/runtime/AISDKThreads.cloud.test.ts

@@ -1,8 +1,15 @@
 // @vitest-environment jsdom
 
+import { getEventListeners } from "node:events";
 import { describe, expect, it, vi } from "vitest";
-import { flushTapSync } from "@assistant-ui/tap";
-import { AuiConfig, createAssistantClient } from "@assistant-ui/store/client";
+import { flushTapSync, resource, useResource } from "@assistant-ui/tap";
+import {
+  attachTransformScopes,
+  AuiConfig,
+  createAssistantClient,
+} from "@assistant-ui/store/client";
+import { useAssistantClientDestroySignal } from "@assistant-ui/store/internal";
+import { inMemoryThreadListTransformScopes } from "@assistant-ui/core/store";
 import type { AssistantCloud } from "assistant-cloud";
 import type { RemoteThreadListAdapter } from "@assistant-ui/core";
 import type { ThreadHistoryAdapter } from "@assistant-ui/core";
@@ -152,4 +159,55 @@ describe("AISDKThreads cloud", () => {
       handle.destroy();
     }
   });
+
+  it("keeps cloud threads off the client destroy signal and stops them on destroy", async () => {
+    const chat = createCancellableTransport();
+    let destroySignal: AbortSignal | undefined;
+    function useThreads() {
+      destroySignal = useAssistantClientDestroySignal();
+      return useResource(
+        AISDKThreads({
+          cloud: {} as AssistantCloud,
+          threadId: "t1",
+          transport: () => chat.transport,
+        }),
+      );
+    }
+    attachTransformScopes(useThreads, inMemoryThreadListTransformScopes);
+    const handle = createAssistantClient(
+      AuiConfig({ threads: resource(useThreads)() }),
+    );
+    handle.subscribe(() => {});
+    const listeners = () => getEventListeners(destroySignal!, "abort").length;
+    try {
+      await handle.getClient().threads.getLoadThreadsPromise();
+      await vi.waitFor(() => {
+        expect(handle.getClient().threads.getState().mainThreadId).toBe("t1");
+      });
+      await vi.waitFor(() => {
+        expect(handle.getClient().thread.getState().isLoading).toBe(false);
+      });
+      const withFirstThread = listeners();
+
+      flushTapSync(() => handle.getClient().threads.switchToThread("t2"));
+      await vi.waitFor(() => {
+        expect(handle.getClient().threads.getState().mainThreadId).toBe("t2");
+      });
+      await vi.waitFor(() => {
+        expect(handle.getClient().thread.getState().isLoading).toBe(false);
+      });
+      expect(listeners()).toBe(withFirstThread);
+
+      flushTapSync(() => handle.getClient().composer.setText("stream me"));
+      flushTapSync(() => handle.getClient().composer.send());
+      await vi.waitFor(() => {
+        expect(handle.getClient().thread.getState().isRunning).toBe(true);
+      });
+    } finally {
+      handle.destroy();
+    }
+    await vi.waitFor(() => {
+      expect(chat.getCancelCount()).toBe(1);
+    });
+  });
 });

package/src/runtime/AISDKThreads.test.ts

@@ -28,7 +28,60 @@ const threadText = (aui: ReturnType<typeof createAssistantClient>) =>
       m.content.map((part) => (part.type === "text" ? part.text : "")).join(""),
     );
 
+const createLiveHandle = (
+  getOptions: () => Parameters<typeof AISDKThreads>[0],
+) => {
+  const listeners = new Set<() => void>();
+  const handle = createAssistantClient({
+    getConfig: () => AuiConfig({ threads: AISDKThreads(getOptions()) }),
+    subscribe: (listener) => {
+      listeners.add(listener);
+      return () => listeners.delete(listener);
+    },
+  });
+  handle.subscribe(() => {});
+  return {
+    handle,
+    rerender: () =>
+      flushTapSync(() => listeners.forEach((listener) => listener())),
+  };
+};
+
+const chatCallbacks = () => ({
+  onToolCall: vi.fn(),
+  onData: vi.fn(),
+  onFinish: vi.fn(),
+  onError: vi.fn(),
+  sendAutomaticallyWhen: vi.fn(() => false),
+});
+
 describe("AISDKThreads", () => {
+  it("emits the selected thread id when its composer sends", async () => {
+    const { transport } = createControlledTransport();
+    const handle = createAssistantClient(
+      AuiConfig({ threads: AISDKThreads({ transport }) }),
+    );
+    try {
+      handle.subscribe(() => {});
+      const aui = handle.getClient();
+      const sent = vi.fn();
+      aui.on({ scope: "thread", event: "composer.send" }, sent);
+
+      flushTapSync(() => aui.composer.setText("hello"));
+      flushTapSync(() => aui.composer.send());
+
+      await vi.waitFor(() => {
+        expect(sent).toHaveBeenCalledWith(
+          expect.objectContaining({
+            threadId: aui.threads.getState().mainThreadId,
+          }),
+        );
+      });
+    } finally {
+      handle.destroy();
+    }
+  });
+
   it("runs one chat per thread and keeps histories isolated across switches", async () => {
     const { transport, emit, close } = createControlledTransport();
     const handle = createAssistantClient(
@@ -253,25 +306,19 @@ describe("AISDKThreads", () => {
 
   it("forwards ChatInit callbacks to each thread's chat from the latest render", async () => {
     const { transport, emit, close } = createControlledTransport();
-    const onFinishA = vi.fn();
-    const onFinishB = vi.fn();
-    let onFinish = onFinishA;
-    const listeners = new Set<() => void>();
-    const handle = createAssistantClient({
-      getConfig: () =>
-        AuiConfig({
-          threads: AISDKThreads({ transport: () => transport, onFinish }),
-        }),
-      subscribe: (listener) => {
-        listeners.add(listener);
-        return () => listeners.delete(listener);
-      },
-    });
-    handle.subscribe(() => {});
+    const sendMessages = vi.spyOn(transport, "sendMessages");
+    const mounted = chatCallbacks();
+    const latest = chatCallbacks();
+    latest.sendAutomaticallyWhen.mockReturnValueOnce(true);
+    let callbacks = mounted;
+    const { handle, rerender } = createLiveHandle(() => ({
+      transport: () => transport,
+      ...callbacks,
+    }));
     const aui = handle.getClient();
 
-    onFinish = onFinishB;
-    flushTapSync(() => listeners.forEach((listener) => listener()));
+    callbacks = latest;
+    rerender();
 
     flushTapSync(() => aui.composer.setText("hi"));
     flushTapSync(() => aui.composer.send());
@@ -280,10 +327,36 @@ describe("AISDKThreads", () => {
         handle.getClient().thread.getState().messages.length,
       ).toBeGreaterThan(0);
     });
-    emit(...textReply("done"));
+    emit(
+      { type: "start" },
+      {
+        type: "tool-input-available",
+        toolCallId: "call-1",
+        toolName: "lookup",
+        input: {},
+      },
+      { type: "data-note", data: "note" },
+      { type: "finish" },
+    );
     close();
-    await vi.waitFor(() => expect(onFinishB).toHaveBeenCalledTimes(1));
-    expect(onFinishA).not.toHaveBeenCalled();
+    await vi.waitFor(() => expect(sendMessages).toHaveBeenCalledTimes(2));
+    emit({ type: "start" }, { type: "error", errorText: "boom" });
+    close();
+    await vi.waitFor(() => expect(latest.onFinish).toHaveBeenCalledTimes(2));
+
+    expect(latest.onToolCall).toHaveBeenCalledExactlyOnceWith({
+      toolCall: expect.objectContaining({ toolCallId: "call-1" }),
+    });
+    expect(latest.onData).toHaveBeenCalledExactlyOnceWith(
+      expect.objectContaining({ type: "data-note", data: "note" }),
+    );
+    expect(latest.onError).toHaveBeenCalledExactlyOnceWith(
+      expect.objectContaining({ message: "boom" }),
+    );
+    expect(latest.sendAutomaticallyWhen).toHaveBeenCalledOnce();
+    for (const callback of Object.values(mounted)) {
+      expect(callback).not.toHaveBeenCalled();
+    }
 
     handle.destroy();
   });
@@ -293,18 +366,10 @@ describe("AISDKThreads", () => {
     const onFinishA = vi.fn();
     const onFinishB = vi.fn();
     let onFinish = onFinishA;
-    const listeners = new Set<() => void>();
-    const handle = createAssistantClient({
-      getConfig: () =>
-        AuiConfig({
-          threads: AISDKThreads({ transport: () => transport, onFinish }),
-        }),
-      subscribe: (listener) => {
-        listeners.add(listener);
-        return () => listeners.delete(listener);
-      },
-    });
-    handle.subscribe(() => {});
+    const { handle, rerender } = createLiveHandle(() => ({
+      transport: () => transport,
+      onFinish,
+    }));
     const aui = handle.getClient();
 
     flushTapSync(() => aui.composer.setText("stream me"));
@@ -322,7 +387,7 @@ describe("AISDKThreads", () => {
 
     flushTapSync(() => aui.threads.switchToNewThread());
     onFinish = onFinishB;
-    flushTapSync(() => listeners.forEach((listener) => listener()));
+    rerender();
 
     emit({ type: "text-end", id: "t1" }, { type: "finish" });
     close();
@@ -361,18 +426,12 @@ describe("AISDKThreads", () => {
     const onFinishA = vi.fn();
     const onFinishB = vi.fn();
     let onFinish = onFinishA;
-    const listeners = new Set<() => void>();
-    const handle = createAssistantClient({
-      getConfig: () =>
-        AuiConfig({
-          threads: AISDKThreads({ cloud, threadId: "t1", transport, onFinish }),
-        }),
-      subscribe: (listener) => {
-        listeners.add(listener);
-        return () => listeners.delete(listener);
-      },
-    });
-    handle.subscribe(() => {});
+    const { handle, rerender } = createLiveHandle(() => ({
+      cloud,
+      threadId: "t1",
+      transport,
+      onFinish,
+    }));
     try {
       await handle.getClient().threads.getLoadThreadsPromise();
       await vi.waitFor(() => {
@@ -383,7 +442,7 @@ describe("AISDKThreads", () => {
       });
 
       onFinish = onFinishB;
-      flushTapSync(() => listeners.forEach((listener) => listener()));
+      rerender();
 
       flushTapSync(() => handle.getClient().composer.setText("hi"));
       flushTapSync(() => handle.getClient().composer.send());

package/src/runtime/AISDKThreads.ts

@@ -162,7 +162,7 @@ const useAISDKChatThread = <UI_MESSAGE extends UIMessage = UIMessage>({
           : fallbackItem,
       chat,
       messageRepositoryInstance: repository,
-      stopOnClientDestroy: cloud,
+      stopOnClientDestroy: false,
     },
   );