@assistant-ui/core 0.3.17 → 0.3.18

package/src/internal.ts

@@ -67,6 +67,10 @@ export {
   resolveFilePartSource,
 } from "./utils/data-url";
 export { invokeUserCallback } from "./utils/invoke-user-callback";
+export {
+  abortableIterable,
+  openAbortableIterable,
+} from "./utils/abortable-iterable";
 export { detectImageMediaType } from "./utils/image-media-type";
 export {
   resolveFileMediaType,

package/src/model-context/frame/host.test.ts

@@ -133,6 +133,30 @@ describe("AssistantFrameHost", () => {
     host.dispose();
   });
 
+  it("cleans up tool calls when posting the request fails", async () => {
+    const { execute, host, postMessage } = createHost();
+    const error = new Error("postMessage failed");
+    const abortController = new AbortController();
+    const removeEventListener = vi.spyOn(
+      abortController.signal,
+      "removeEventListener",
+    );
+    postMessage.mockImplementation((data) => {
+      if (data.message.type === "tool-call") throw error;
+    });
+
+    await expect(
+      execute({}, { ...executionContext, abortSignal: abortController.signal }),
+    ).rejects.toBe(error);
+
+    expect(vi.getTimerCount()).toBe(0);
+    expect(removeEventListener).toHaveBeenCalledWith(
+      "abort",
+      expect.any(Function),
+    );
+    host.dispose();
+  });
+
   it("rejects tool results that carry an empty error message", async () => {
     const { dispatchMessage, execute, getToolCallId, host } = createHost();
     const result = Promise.resolve(

package/src/model-context/frame/host.ts

@@ -195,10 +195,16 @@ export class AssistantFrameHost implements ModelContextProvider {
       }, timeout);
       abortSignal?.addEventListener("abort", onAbort, { once: true });
 
-      this._iframeWindow.postMessage(
-        { channel: FRAME_MESSAGE_CHANNEL, message },
-        this._targetOrigin,
-      );
+      try {
+        this._iframeWindow.postMessage(
+          { channel: FRAME_MESSAGE_CHANNEL, message },
+          this._targetOrigin,
+        );
+      } catch (error) {
+        const pending = this._pendingRequests.get(message.id);
+        this._pendingRequests.delete(message.id);
+        pending?.reject(error);
+      }
     });
   }
 

package/src/model-context/frame/provider.test.ts

@@ -120,6 +120,42 @@ describe("AssistantFrameProvider", () => {
     await vi.waitFor(() => expect(execute).toHaveBeenCalledOnce());
   });
 
+  it("does not broadcast from a disposed provider", async () => {
+    vi.useFakeTimers();
+
+    try {
+      AssistantFrameProvider.addModelContextProvider({
+        getModelContext: () => ({ system: "disposed context" }),
+      });
+      AssistantFrameProvider.dispose();
+
+      AssistantFrameProvider.addModelContextProvider({
+        getModelContext: () => ({ system: "current context" }),
+      });
+      vi.mocked(parentWindow.postMessage).mockClear();
+
+      await vi.runAllTimersAsync();
+
+      expect(parentWindow.postMessage).toHaveBeenCalledOnce();
+      expect(parentWindow.postMessage).toHaveBeenCalledWith(
+        {
+          channel: FRAME_MESSAGE_CHANNEL,
+          message: {
+            type: "model-context-update",
+            context: {
+              system: "current context",
+              tools: {},
+            },
+          },
+        },
+        window.location.origin,
+      );
+    } finally {
+      AssistantFrameProvider.dispose();
+      vi.useRealTimers();
+    }
+  });
+
   it("reports a failure even when the thrown error has an empty message", async () => {
     const execute = vi.fn(async () => {
       throw new Error();
@@ -146,6 +182,41 @@ describe("AssistantFrameProvider", () => {
     });
   });
 
+  it("reports tool results that cannot cross the frame boundary", async () => {
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => undefined);
+    const execute = vi.fn(async () => () => undefined);
+    vi.mocked(parentWindow.postMessage).mockImplementation((data) => {
+      structuredClone(data);
+    });
+    AssistantFrameProvider.addModelContextProvider({
+      getModelContext: () => ({
+        tools: { sensitiveTool: { execute } },
+      }),
+    });
+
+    dispatchToolCall(window.location.origin);
+
+    await vi.waitFor(() => {
+      expect(consoleError).toHaveBeenCalledWith(
+        "[assistant-ui] AssistantFrame tool result could not be sent.",
+        expect.objectContaining({ name: "DataCloneError" }),
+      );
+      expect(parentWindow.postMessage).toHaveBeenCalledWith(
+        {
+          channel: FRAME_MESSAGE_CHANNEL,
+          message: {
+            type: "tool-result",
+            id: "tool-call-1",
+            error: "Tool result could not be sent across the frame boundary",
+          },
+        },
+        { targetOrigin: window.location.origin },
+      );
+    });
+  });
+
   it("aborts in-flight tool calls when the parent cancels them", async () => {
     let toolSignal: AbortSignal | undefined;
     const execute = vi.fn(

package/src/model-context/frame/provider.ts

@@ -47,13 +47,18 @@ export class AssistantFrameProvider {
   private _targetOrigin: string;
   private _strictRegistrations = 0;
   private _wildcardRegistrations = 0;
+  private _startupTimer: ReturnType<typeof setTimeout> | undefined;
+  private _disposed = false;
 
   private constructor(targetOrigin: string = getDefaultTargetOrigin()) {
     this._targetOrigin = targetOrigin;
     this.handleMessage = this.handleMessage.bind(this);
     window.addEventListener("message", this.handleMessage);
 
-    setTimeout(() => this.broadcastUpdate(), 0);
+    this._startupTimer = setTimeout(() => {
+      this._startupTimer = undefined;
+      this.broadcastUpdate();
+    }, 0);
   }
 
   private static getInstance(targetOrigin?: string): AssistantFrameProvider {
@@ -106,7 +111,12 @@ export class AssistantFrameProvider {
         break;
 
       case "tool-call":
-        this.handleToolCall(message, event);
+        void this.handleToolCall(message, event).catch((error: unknown) => {
+          console.error(
+            "[assistant-ui] AssistantFrame tool call failed.",
+            error,
+          );
+        });
         break;
 
       case "tool-cancel":
@@ -156,11 +166,25 @@ export class AssistantFrameProvider {
     if (this._activeToolCalls.get(message.id) !== activeCall) return;
     this._activeToolCalls.delete(message.id);
 
-    this.sendMessage(event, {
-      type: "tool-result",
-      id: message.id,
-      ...(error !== undefined ? { error } : { result }),
-    });
+    try {
+      this.sendMessage(event, {
+        type: "tool-result",
+        id: message.id,
+        ...(error !== undefined ? { error } : { result }),
+      });
+    } catch (sendError) {
+      if (error !== undefined) throw sendError;
+
+      console.error(
+        "[assistant-ui] AssistantFrame tool result could not be sent.",
+        sendError,
+      );
+      this.sendMessage(event, {
+        type: "tool-result",
+        id: message.id,
+        error: "Tool result could not be sent across the frame boundary",
+      });
+    }
   }
 
   private cancelToolCall(id: string) {
@@ -244,6 +268,7 @@ export class AssistantFrameProvider {
   }
 
   private broadcastUpdate() {
+    if (this._disposed) return;
     if (window.parent && window.parent !== window) {
       const updateMessage: FrameMessage = {
         type: "model-context-update",
@@ -375,6 +400,11 @@ export class AssistantFrameProvider {
   static dispose() {
     if (AssistantFrameProvider._instance) {
       const instance = AssistantFrameProvider._instance;
+      instance._disposed = true;
+      if (instance._startupTimer !== undefined) {
+        clearTimeout(instance._startupTimer);
+        instance._startupTimer = undefined;
+      }
       window.removeEventListener("message", instance.handleMessage);
 
       let cleanupFailed = false;

package/src/model-context/interactable-composer-metadata.test.ts

@@ -131,6 +131,34 @@ describe("shallowMergeInteractableState", () => {
     });
   });
 
+  it("uses the first valid patch for each array item id", () => {
+    expect(
+      shallowMergeInteractableState(
+        { tasks: [{ id: 1, title: "Original", done: false }] },
+        {
+          tasks: {
+            update: [
+              null,
+              { title: "Missing id" },
+              { id: "1", title: "String id" },
+              { id: 1, title: "First", done: true },
+              { id: 1, title: "Second" },
+            ],
+          },
+        },
+      ),
+    ).toEqual({ tasks: [{ id: 1, title: "First", done: true }] });
+  });
+
+  it("does not match NaN array item ids", () => {
+    expect(
+      shallowMergeInteractableState(
+        { tasks: [{ id: Number.NaN, title: "Original" }] },
+        { tasks: { update: [{ id: Number.NaN, title: "Patched" }] } },
+      ),
+    ).toEqual({ tasks: [{ id: Number.NaN, title: "Original" }] });
+  });
+
   it("keeps raw array replacement semantics", () => {
     expect(
       shallowMergeInteractableState(

package/src/model-context/interactable-composer-metadata.ts

@@ -94,12 +94,18 @@ function applyArrayUpdate(
 
   const patches = update.update;
   if (Array.isArray(patches) && patches.length > 0) {
+    const patchesById = new Map<string | number, Record<string, unknown>>();
+    for (const candidate of patches) {
+      const id = getArrayItemId(candidate);
+      if (id !== undefined && !Number.isNaN(id) && !patchesById.has(id)) {
+        patchesById.set(id, candidate as Record<string, unknown>);
+      }
+    }
+
     next = next.map((item) => {
       const id = getArrayItemId(item);
       if (id === undefined || !isRecord(item)) return item;
-      const patch = patches.find(
-        (candidate) => isRecord(candidate) && candidate.id === id,
-      );
+      const patch = patchesById.get(id);
       return patch ? { ...item, ...patch } : item;
     });
   }

package/src/model-context/registry.test.ts

@@ -2,6 +2,18 @@ import { describe, expect, it, vi } from "vitest";
 import { ModelContextRegistry } from "./registry";
 
 describe("ModelContextRegistry", () => {
+  it.each(["__proto__", "constructor", "toString"])(
+    "retains a tool named %s",
+    (toolName) => {
+      const registry = new ModelContextRegistry();
+      registry.addTool({ toolName });
+
+      const tools = registry.getModelContext().tools!;
+      expect(Object.hasOwn(tools, toolName)).toBe(true);
+      expect(Object.keys(tools)).toEqual([toolName]);
+    },
+  );
+
   it("notifies every subscriber and rethrows when a registration subscriber throws", () => {
     const registry = new ModelContextRegistry();
     const error = new Error("subscriber failed");

package/src/model-context/registry.ts

@@ -13,6 +13,7 @@ import type {
   ModelContextRegistryInstructionHandle,
   ModelContextRegistryProviderHandle,
 } from "./registry-handles";
+import { nullProtoRecord } from "../utils/record";
 
 export class ModelContextRegistry implements ModelContextProvider {
   private _tools = new Map<symbol, AssistantToolProps<any, any>>();
@@ -34,7 +35,7 @@ export class ModelContextRegistry implements ModelContextProvider {
     const system =
       instructions.length > 0 ? instructions.join("\n\n") : undefined;
 
-    const tools: Record<string, Tool<any, any>> = {};
+    const tools = nullProtoRecord<Tool<any, any>>();
     for (const toolProps of this._tools.values()) {
       const { toolName, render, ...tool } = toolProps;
       tools[toolName] = tool;
@@ -54,7 +55,7 @@ export class ModelContextRegistry implements ModelContextProvider {
     }
 
     if (providerContexts.tools) {
-      context.tools = { ...(context.tools || {}), ...providerContexts.tools };
+      context.tools = nullProtoRecord(context.tools, providerContexts.tools);
     }
 
     if (providerContexts.callSettings) {

package/src/model-context/types.test.ts

@@ -0,0 +1,25 @@
+import { describe, expect, it } from "vitest";
+import { mergeModelContexts } from "./types";
+
+const provider = (tools: Record<string, any>, priority?: number) =>
+  ({ getModelContext: () => ({ tools, priority }) }) as any;
+
+describe("mergeModelContexts", () => {
+  it.each(["__proto__", "constructor", "toString"])(
+    "retains and prioritizes a tool named %s",
+    (name) => {
+      const highPriorityTool = { description: "high", parameters: {} };
+      const lowPriorityTool = { description: "low", parameters: {} };
+      const otherTool = { description: "other", parameters: {} };
+      const result = mergeModelContexts(
+        new Set([
+          provider({ [name]: highPriorityTool, ok: otherTool }, 1),
+          provider({ [name]: lowPriorityTool }, 0),
+        ]),
+      );
+
+      expect(Object.keys(result.tools ?? {})).toEqual([name, "ok"]);
+      expect(result.tools?.[name]?.description).toBe("high");
+    },
+  );
+});

package/src/model-context/types.ts

@@ -1,5 +1,6 @@
 import type { Unsubscribe } from "../types/unsubscribe";
 import type { Tool } from "assistant-stream";
+import { nullProtoRecord } from "../utils/record";
 
 export type LanguageModelV1CallSettings = {
   maxTokens?: number;
@@ -68,7 +69,7 @@ export const mergeModelContexts = (
     .map((c) => c.getModelContext())
     .sort((a, b) => (b.priority ?? 0) - (a.priority ?? 0));
 
-  const toolPriorities: Record<string, number> = {};
+  const toolPriorities = nullProtoRecord<number>();
 
   return configs.reduce((acc, config) => {
     const priority = config.priority ?? 0;
@@ -81,7 +82,9 @@ export const mergeModelContexts = (
     }
     if (config.tools) {
       for (const [name, tool] of Object.entries(config.tools)) {
-        const existing = acc.tools?.[name];
+        const hasExisting =
+          acc.tools !== undefined && Object.hasOwn(acc.tools, name);
+        const existing = hasExisting ? acc.tools![name] : undefined;
         if (existing && existing !== tool) {
           const existingPriority = toolPriorities[name]!;
           if (existingPriority === priority) {
@@ -106,9 +109,11 @@ export const mergeModelContexts = (
           continue;
         }
 
-        if (!acc.tools) acc.tools = {};
+        if (!acc.tools) acc.tools = nullProtoRecord();
         acc.tools[name] = stripOverwrite(tool);
-        toolPriorities[name] ??= priority;
+        if (!Object.hasOwn(toolPriorities, name)) {
+          toolPriorities[name] = priority;
+        }
       }
     }
     if (config.config) {

package/src/react/RuntimeAdapter.test.ts

@@ -69,6 +69,40 @@ describe("RuntimeAdapter via the neutral store entry", () => {
     handle.destroy();
   });
 
+  it("emits composer.send for a thread.append, flagging a suggestion", async () => {
+    const { runtime, onNew } = createRuntime({
+      suggestions: [{ prompt: "hi there" }],
+    });
+    const handle = createAssistantClient(
+      AuiConfig({ threads: RuntimeAdapter(runtime) }),
+    );
+    handle.subscribe(() => {});
+    const aui = handle.getClient();
+    const sent = vi.fn();
+    aui.on({ scope: "thread", event: "composer.send" }, sent);
+
+    flushTapSync(() =>
+      aui.thread.append({ content: [{ type: "text", text: "hi there" }] }),
+    );
+    flushTapSync(() => aui.thread.append("plain text"));
+    await new Promise((resolve) => setTimeout(resolve));
+
+    expect(onNew).toHaveBeenCalledTimes(2);
+    expect(sent).toHaveBeenCalledTimes(2);
+    expect(sent.mock.calls[0]![0]).toMatchObject({
+      chars: 8,
+      attachments: 0,
+      suggestion: true,
+    });
+    expect(sent.mock.calls[1]![0]).toMatchObject({
+      chars: 10,
+      attachments: 0,
+    });
+    expect(sent.mock.calls[1]![0]).not.toHaveProperty("suggestion");
+
+    handle.destroy();
+  });
+
   it("registers the client's model context on the runtime", () => {
     const { runtime } = createRuntime();
     const register = vi.spyOn(runtime, "registerModelContextProvider");

package/src/react/client/InMemoryThreadList.ts

@@ -209,6 +209,7 @@ const useInMemoryThreadList = (
       mainThreadId,
       newThreadId: null,
       isLoading: false,
+      loadError: undefined,
       isLoadingMore: false,
       hasMore: false,
       threadIds: regularThreads.map((t) => t.id),

package/src/react/client/Interactables.test.ts

@@ -210,6 +210,34 @@ describe("Interactables registration", () => {
     expect(stateOf(root, "n1")).toEqual({ v: 7 });
   });
 
+  // `definitions` is read by bare key with ids that reach the runtime from
+  // model tool calls, so it must stay prototype-free through every transition.
+  it.each(["__proto__", "constructor", "toString"])(
+    "registers, updates and unregisters an interactable named %s",
+    (id) => {
+      root = mount();
+      const unregister = root.getValue().register(reg(id));
+      expect(stateOf(root, id)).toEqual({ v: 0 });
+
+      root.getValue().setState(id, () => ({ v: 1 }));
+      expect(stateOf(root, id)).toEqual({ v: 1 });
+      expect(Object.keys(root.getValue().getState().definitions)).toEqual([id]);
+
+      unregister();
+      expect(Object.keys(root.getValue().getState().definitions)).toEqual([]);
+    },
+  );
+
+  it("keeps the state records prototype-free", () => {
+    root = mount();
+    root.getValue().register(reg("n1"));
+    root.getValue().setState("n1", () => ({ v: 1 }));
+
+    const state = root.getValue().getState();
+    expect(Object.getPrototypeOf(state.definitions)).toBeNull();
+    expect(Object.getPrototypeOf(state.persistence)).toBeNull();
+  });
+
   it("restores detached state when an instance re-registers in-session", async () => {
     root = mount();
     const unregister = root.getValue().register(reg("n1"));

package/src/react/client/Interactables.ts

@@ -25,6 +25,7 @@ import {
 } from "../../model-context/interactable-composer-metadata";
 import { notifySubscribers as notifyStateSubscribers } from "../../subscribable/subscribable";
 import { useInteractablePersistenceQueue } from "../interactables-shared/useInteractablePersistenceQueue";
+import { nullProtoRecord } from "../../utils/record";
 
 type RestorePersistedStateOptions = {
   stash: Map<string, unknown>;
@@ -77,8 +78,8 @@ const useInteractablesResource = ({
   persistence,
 }: Unstable_InteractablesConfig = {}): ClientOutput<"unstable_interactables"> => {
   const [state, setState] = useState<Unstable_InteractablesState>(() => ({
-    definitions: {},
-    persistence: {},
+    definitions: nullProtoRecord(),
+    persistence: nullProtoRecord(),
   }));
 
   const clientRef = useAssistantClientRef();
@@ -123,7 +124,8 @@ const useInteractablesResource = ({
   );
 
   const exportState = useCallback((): Unstable_InteractablePersistedState => {
-    const result: Unstable_InteractablePersistedState = {};
+    const result =
+      nullProtoRecord<Unstable_InteractablePersistedState[string]>();
     for (const [id, def] of Object.entries(stateRef.current.definitions)) {
       if (def.scope === "thread") continue; // thread items persist via snapshot, not the adapter
       result[id] = { name: def.name, state: def.state };
@@ -167,7 +169,7 @@ const useInteractablesResource = ({
       }
       setStateAndRef((prev) => {
         let changed = false;
-        const definitions = { ...prev.definitions };
+        const definitions = nullProtoRecord(prev.definitions);
         for (const [id, entry] of Object.entries(saved)) {
           const def = definitions[id];
           if (!def || !shouldApply(id, def)) continue;
@@ -260,10 +262,9 @@ const useInteractablesResource = ({
         if (!existing) return prev;
         return {
           ...prev,
-          definitions: {
-            ...prev.definitions,
+          definitions: nullProtoRecord(prev.definitions, {
             [id]: { ...existing, state: updater(existing.state) },
-          },
+          }),
         };
       });
       if (stateRef.current.definitions[id]?.scope !== "thread") {
@@ -450,8 +451,7 @@ const useInteractablesResource = ({
 
       setStateAndRef((prev) => ({
         ...prev,
-        definitions: {
-          ...prev.definitions,
+        definitions: nullProtoRecord(prev.definitions, {
           [def.id]: {
             id: def.id,
             name: def.name,
@@ -466,7 +466,7 @@ const useInteractablesResource = ({
               loaded ??
               def.initialState,
           },
-        },
+        }),
       }));
 
       return () => {
@@ -498,9 +498,11 @@ const useInteractablesResource = ({
             }
           }
           partialSchemaCacheRef.current.delete(def.id);
-          const { [def.id]: _, ...rest } = prev.definitions;
-          const { [def.id]: __, ...restPersistence } = prev.persistence;
-          return { ...prev, definitions: rest, persistence: restPersistence };
+          const definitions = nullProtoRecord(prev.definitions);
+          const persistence = nullProtoRecord(prev.persistence);
+          delete definitions[def.id];
+          delete persistence[def.id];
+          return { ...prev, definitions, persistence };
         });
       };
     },