@assistant-ui/core 0.3.18 → 0.3.20

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

@@ -74,9 +74,31 @@ beforeEach(() => {
 afterEach(() => {
   vi.useRealTimers();
   vi.unstubAllGlobals();
+  vi.restoreAllMocks();
 });
 
 describe("AssistantFrameHost", () => {
+  it("does not install its message listener when initialization fails", () => {
+    const addEventListener = vi.fn();
+    const removeEventListener = vi.fn();
+    vi.stubGlobal("window", {
+      addEventListener,
+      removeEventListener,
+      location: { origin: DEFAULT_ORIGIN },
+    });
+    const error = new Error("postMessage failed");
+    const iframeWindow = {
+      postMessage: vi.fn(() => {
+        throw error;
+      }),
+    } as unknown as Window;
+
+    expect(() => new AssistantFrameHost(iframeWindow)).toThrow(error);
+
+    expect(addEventListener).not.toHaveBeenCalled();
+    expect(removeEventListener).not.toHaveBeenCalled();
+  });
+
   it("defaults to the current origin", () => {
     const { host, postMessage } = createHost();
 
@@ -212,8 +234,83 @@ describe("AssistantFrameHost", () => {
     expect(vi.getTimerCount()).toBe(0);
   });
 
+  it("settles every pending tool call when cancellation posting fails", async () => {
+    const { execute, host, postMessage } = createHost();
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => undefined);
+    const firstAbortController = new AbortController();
+    const secondAbortController = new AbortController();
+    const firstRemoveEventListener = vi.spyOn(
+      firstAbortController.signal,
+      "removeEventListener",
+    );
+    const secondRemoveEventListener = vi.spyOn(
+      secondAbortController.signal,
+      "removeEventListener",
+    );
+    const firstRejected = vi.fn();
+    const secondRejected = vi.fn();
+
+    void execute(
+      {},
+      { ...executionContext, abortSignal: firstAbortController.signal },
+    ).catch(firstRejected);
+    void execute(
+      {},
+      { ...executionContext, abortSignal: secondAbortController.signal },
+    ).catch(secondRejected);
+
+    const firstTransportError = new Error("first tool cancellation failed");
+    const secondTransportError = new Error("second tool cancellation failed");
+    let cancellationCount = 0;
+    postMessage.mockImplementation((data) => {
+      if (data.message.type !== "tool-cancel") return;
+      cancellationCount += 1;
+      throw cancellationCount === 1
+        ? firstTransportError
+        : secondTransportError;
+    });
+
+    expect(() => host.dispose()).toThrow(firstTransportError);
+    await Promise.resolve();
+
+    expect(firstRejected).toHaveBeenCalledWith(
+      expect.objectContaining({
+        message: "AssistantFrameHost has been disposed",
+      }),
+    );
+    expect(secondRejected).toHaveBeenCalledWith(
+      expect.objectContaining({
+        message: "AssistantFrameHost has been disposed",
+      }),
+    );
+    expect(
+      postMessage.mock.calls.filter(
+        ([data]) => data.message.type === "tool-cancel",
+      ),
+    ).toHaveLength(2);
+    expect(firstRemoveEventListener).toHaveBeenCalledWith(
+      "abort",
+      expect.any(Function),
+    );
+    expect(secondRemoveEventListener).toHaveBeenCalledWith(
+      "abort",
+      expect.any(Function),
+    );
+    expect(vi.getTimerCount()).toBe(0);
+    expect(consoleError).toHaveBeenCalledWith(
+      "[assistant-ui] AssistantFrameHost tool cancellation could not be sent.",
+      secondTransportError,
+    );
+    expect(() => host.dispose()).not.toThrow();
+  });
+
   it("rejects pending tool calls when execution is aborted", async () => {
     const { execute, getToolCallId, host, postMessage } = createHost();
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => undefined);
     const abortController = new AbortController();
     const abortError = new Error("Run cancelled");
     abortError.name = "AbortError";
@@ -229,8 +326,13 @@ describe("AssistantFrameHost", () => {
     const onRejected = vi.fn();
     void result.catch(onRejected);
     const toolCallId = getToolCallId();
+    postMessage.mockImplementation((data) => {
+      if (data.message.type === "tool-cancel") {
+        throw new Error("tool cancellation failed");
+      }
+    });
 
-    abortController.abort(abortError);
+    expect(() => abortController.abort(abortError)).not.toThrow();
     await Promise.resolve();
 
     expect(onRejected).toHaveBeenCalledWith(abortError);
@@ -241,17 +343,26 @@ describe("AssistantFrameHost", () => {
       },
       DEFAULT_ORIGIN,
     );
+    expect(consoleError).toHaveBeenCalledOnce();
     expect(vi.getTimerCount()).toBe(0);
     host.dispose();
   });
 
   it("cancels tool calls when they time out", async () => {
     const { execute, getToolCallId, host, postMessage } = createHost();
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => undefined);
     const result = Promise.resolve(execute({}, executionContext));
     const toolCallId = getToolCallId();
     const rejection = expect(result).rejects.toThrow(
       'Tool call "search" timed out',
     );
+    postMessage.mockImplementation((data) => {
+      if (data.message.type === "tool-cancel") {
+        throw new Error("tool cancellation failed");
+      }
+    });
 
     await vi.advanceTimersByTimeAsync(30000);
 
@@ -263,6 +374,7 @@ describe("AssistantFrameHost", () => {
       },
       DEFAULT_ORIGIN,
     );
+    expect(consoleError).toHaveBeenCalledOnce();
     expect(vi.getTimerCount()).toBe(0);
     host.dispose();
   });

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

@@ -12,6 +12,13 @@ import {
 
 const getDefaultTargetOrigin = () => window.location.origin;
 
+const logCancellationFailure = (error: unknown) => {
+  console.error(
+    "[assistant-ui] AssistantFrameHost tool cancellation could not be sent.",
+    error,
+  );
+};
+
 /**
  * Deserializes tools from JSON Schema format back to Tool objects
  */
@@ -73,9 +80,10 @@ export class AssistantFrameHost implements ModelContextProvider {
     this._targetOrigin = targetOrigin;
 
     this.handleMessage = this.handleMessage.bind(this);
-    window.addEventListener("message", this.handleMessage);
-
+    // A posted message cannot be delivered before this constructor returns, so
+    // requesting first means a failed request leaves no listener to clean up.
     this.requestContext();
+    window.addEventListener("message", this.handleMessage);
   }
 
   private handleMessage(event: MessageEvent) {
@@ -208,14 +216,21 @@ export class AssistantFrameHost implements ModelContextProvider {
     });
   }
 
-  private cancelToolCall(id: string) {
-    this._iframeWindow.postMessage(
-      {
-        channel: FRAME_MESSAGE_CHANNEL,
-        message: { type: "tool-cancel", id } satisfies FrameMessage,
-      },
-      this._targetOrigin,
-    );
+  private cancelToolCall(
+    id: string,
+    onError: (error: unknown) => void = logCancellationFailure,
+  ) {
+    try {
+      this._iframeWindow.postMessage(
+        {
+          channel: FRAME_MESSAGE_CHANNEL,
+          message: { type: "tool-cancel", id } satisfies FrameMessage,
+        },
+        this._targetOrigin,
+      );
+    } catch (error) {
+      onError(error);
+    }
   }
 
   private requestContext() {
@@ -248,10 +263,22 @@ export class AssistantFrameHost implements ModelContextProvider {
     window.removeEventListener("message", this.handleMessage);
     this._subscribers.clear();
     const error = new Error("AssistantFrameHost has been disposed");
+    let cancellationFailed = false;
+    let cancellationError: unknown;
+
     for (const [id, pending] of this._pendingRequests) {
-      this.cancelToolCall(id);
+      this._pendingRequests.delete(id);
+      this.cancelToolCall(id, (error) => {
+        if (!cancellationFailed) {
+          cancellationFailed = true;
+          cancellationError = error;
+        } else {
+          logCancellationFailure(error);
+        }
+      });
       pending.reject(error);
     }
-    this._pendingRequests.clear();
+
+    if (cancellationFailed) throw cancellationError;
   }
 }

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

@@ -2,9 +2,24 @@
  * @vitest-environment jsdom
  */
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import type { Tool } from "assistant-stream";
 import { AssistantFrameProvider } from "./provider";
 import { FRAME_MESSAGE_CHANNEL } from "./types";
 
+const createTool = <TResult>(
+  execute: NonNullable<
+    Extract<
+      Tool<Record<string, unknown>, TResult>,
+      { type: "frontend" }
+    >["execute"]
+  >,
+) =>
+  ({
+    type: "frontend",
+    parameters: { type: "object", properties: {} },
+    execute,
+  }) satisfies Tool<Record<string, unknown>, TResult>;
+
 describe("AssistantFrameProvider", () => {
   let messageHandler: ((event: MessageEvent) => void) | undefined;
   let parentWindow: Window;
@@ -80,7 +95,7 @@ describe("AssistantFrameProvider", () => {
       {
         getModelContext: () => ({
           tools: {
-            sensitiveTool: { execute },
+            sensitiveTool: createTool(execute),
           },
         }),
       },
@@ -104,7 +119,7 @@ describe("AssistantFrameProvider", () => {
     const execute = vi.fn(async () => "result");
     AssistantFrameProvider.addModelContextProvider({
       getModelContext: () => ({
-        tools: { sensitiveTool: { execute } },
+        tools: { sensitiveTool: createTool(execute) },
       }),
     });
 
@@ -163,7 +178,7 @@ describe("AssistantFrameProvider", () => {
     AssistantFrameProvider.addModelContextProvider({
       getModelContext: () => ({
         tools: {
-          sensitiveTool: { execute },
+          sensitiveTool: createTool(execute),
         },
       }),
     });
@@ -192,7 +207,7 @@ describe("AssistantFrameProvider", () => {
     });
     AssistantFrameProvider.addModelContextProvider({
       getModelContext: () => ({
-        tools: { sensitiveTool: { execute } },
+        tools: { sensitiveTool: createTool(execute) },
       }),
     });
 
@@ -233,7 +248,7 @@ describe("AssistantFrameProvider", () => {
     );
     AssistantFrameProvider.addModelContextProvider({
       getModelContext: () => ({
-        tools: { sensitiveTool: { execute } },
+        tools: { sensitiveTool: createTool(execute) },
       }),
     });
 
@@ -270,7 +285,9 @@ describe("AssistantFrameProvider", () => {
       },
     );
     AssistantFrameProvider.addModelContextProvider({
-      getModelContext: () => ({ tools: { sensitiveTool: { execute } } }),
+      getModelContext: () => ({
+        tools: { sensitiveTool: createTool(execute) },
+      }),
     });
 
     dispatchToolCall(window.location.origin, parentWindow, "tool-a");
@@ -298,7 +315,9 @@ describe("AssistantFrameProvider", () => {
       },
     );
     AssistantFrameProvider.addModelContextProvider({
-      getModelContext: () => ({ tools: { sensitiveTool: { execute } } }),
+      getModelContext: () => ({
+        tools: { sensitiveTool: createTool(execute) },
+      }),
     });
 
     dispatchToolCall(window.location.origin, parentWindow, "duplicate");
@@ -325,7 +344,9 @@ describe("AssistantFrameProvider", () => {
       },
     );
     AssistantFrameProvider.addModelContextProvider({
-      getModelContext: () => ({ tools: { sensitiveTool: { execute } } }),
+      getModelContext: () => ({
+        tools: { sensitiveTool: createTool(execute) },
+      }),
     });
 
     dispatchToolCall(window.location.origin);
@@ -371,7 +392,9 @@ describe("AssistantFrameProvider", () => {
       },
     );
     const removeProvider = AssistantFrameProvider.addModelContextProvider({
-      getModelContext: () => ({ tools: { sensitiveTool: { execute } } }),
+      getModelContext: () => ({
+        tools: { sensitiveTool: createTool(execute) },
+      }),
     });
 
     dispatchToolCall(window.location.origin);
@@ -417,7 +440,9 @@ describe("AssistantFrameProvider", () => {
       },
     );
     const provider = {
-      getModelContext: () => ({ tools: { sensitiveTool: { execute } } }),
+      getModelContext: () => ({
+        tools: { sensitiveTool: createTool(execute) },
+      }),
     };
     const removeFirst =
       AssistantFrameProvider.addModelContextProvider(provider);
@@ -438,7 +463,7 @@ describe("AssistantFrameProvider", () => {
     const shadowedExecute = vi.fn(async () => "shadowed");
     const removeShadowed = AssistantFrameProvider.addModelContextProvider({
       getModelContext: () => ({
-        tools: { sensitiveTool: { execute: shadowedExecute } },
+        tools: { sensitiveTool: createTool(shadowedExecute) },
       }),
     });
 
@@ -457,7 +482,7 @@ describe("AssistantFrameProvider", () => {
     );
     const removeOwner = AssistantFrameProvider.addModelContextProvider({
       getModelContext: () => ({
-        tools: { sensitiveTool: { execute } },
+        tools: { sensitiveTool: createTool(execute) },
       }),
     });
 
@@ -490,10 +515,10 @@ describe("AssistantFrameProvider", () => {
       },
     );
     const removeFirst = AssistantFrameProvider.addModelContextProvider({
-      getModelContext: () => ({ tools: { firstTool: { execute } } }),
+      getModelContext: () => ({ tools: { firstTool: createTool(execute) } }),
     });
     const removeSecond = AssistantFrameProvider.addModelContextProvider({
-      getModelContext: () => ({ tools: { secondTool: { execute } } }),
+      getModelContext: () => ({ tools: { secondTool: createTool(execute) } }),
     });
 
     dispatchToolCall(
@@ -540,7 +565,7 @@ describe("AssistantFrameProvider", () => {
       {
         getModelContext: () => ({
           tools: {
-            sensitiveTool: { execute },
+            sensitiveTool: createTool(execute),
           },
         }),
       },
@@ -567,7 +592,7 @@ describe("AssistantFrameProvider", () => {
       {
         getModelContext: () => ({
           tools: {
-            sensitiveTool: { execute },
+            sensitiveTool: createTool(execute),
           },
         }),
       },
@@ -609,7 +634,7 @@ describe("AssistantFrameProvider", () => {
       AssistantFrameProvider.addModelContextProvider(
         {
           getModelContext: () => ({
-            tools: { sensitiveTool: { execute } },
+            tools: { sensitiveTool: createTool(execute) },
           }),
           subscribe: () => {
             throw new Error("subscribe failed");
@@ -648,7 +673,7 @@ describe("AssistantFrameProvider", () => {
     expect(() =>
       AssistantFrameProvider.addModelContextProvider({
         getModelContext: () => ({
-          tools: { sensitiveTool: { execute } },
+          tools: { sensitiveTool: createTool(execute) },
         }),
         subscribe: () => {
           dispatchToolCall(window.location.origin);
@@ -678,7 +703,7 @@ describe("AssistantFrameProvider", () => {
     let subscriptionCount = 0;
     const provider = {
       getModelContext: () => ({
-        tools: { sensitiveTool: { execute } },
+        tools: { sensitiveTool: createTool(execute) },
       }),
       subscribe: () => {
         subscriptionCount += 1;
@@ -934,7 +959,7 @@ describe("AssistantFrameProvider", () => {
       {
         getModelContext: () => ({
           tools: {
-            sensitiveTool: { execute },
+            sensitiveTool: createTool(execute),
           },
         }),
       },

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

@@ -104,6 +104,18 @@ describe("interactableToolName", () => {
 });
 
 describe("shallowMergeInteractableState", () => {
+  it("adds prototype-named fields as own properties", () => {
+    const value = { enabled: true };
+    const result = shallowMergeInteractableState(
+      { title: "Example" },
+      Object.fromEntries([["__proto__", value]]),
+    ) as Record<string, unknown>;
+
+    expect(Object.getPrototypeOf(result)).toBe(Object.prototype);
+    expect(Object.hasOwn(result, "__proto__")).toBe(true);
+    expect(result["__proto__"]).toBe(value);
+  });
+
   it("applies array operations from a baseline", () => {
     const prev = {
       tasks: [
@@ -453,6 +465,24 @@ describe("gateInteractableComposerMetadata", () => {
     expect(gated?.interactables).toEqual([entry("a", { v: 1 })]);
   });
 
+  it.each(["__proto__", "toString"])(
+    "treats a removed prototype-named field %s as a full snapshot",
+    (field) => {
+      const known = Object.fromEntries([
+        [field, { enabled: true }],
+        ["title", "draft"],
+        ["stable", true],
+      ]);
+      const current = { title: "edited", stable: true };
+      const meta = { interactables: [entry("a", current)] };
+      const history = [userMsg([entry("a", known)])];
+
+      const gated = gateInteractableComposerMetadata(meta, history);
+
+      expect(gated?.interactables).toEqual([entry("a", current)]);
+    },
+  );
+
   it("omits an interactable the model already knows via its own update_* call", () => {
     const meta = { interactables: [entry("a", { v: 2 }, "note")] };
     const history = [

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

@@ -149,7 +149,7 @@ export function shallowMergeInteractableState(
   const baseline = isRecord(options?.arrayBaseline)
     ? options.arrayBaseline
     : prev;
-  const next = { ...prev };
+  const next = Object.entries(prev);
   for (const [key, value] of Object.entries(partial)) {
     const baseValue = baseline[key];
     if (Array.isArray(baseValue) && isRecord(value)) {
@@ -158,12 +158,12 @@ export function shallowMergeInteractableState(
         (options.idKeyedFields === undefined || options.idKeyedFields.has(key))
           ? () => options.idFactory?.(key)
           : undefined;
-      next[key] = applyArrayUpdate(baseValue, value, mintId);
+      next.push([key, applyArrayUpdate(baseValue, value, mintId)]);
     } else {
-      next[key] = value;
+      next.push([key, value]);
     }
   }
-  return next;
+  return Object.fromEntries(next);
 }
 
 /**
@@ -179,17 +179,17 @@ function shallowDiffInteractableState(
 ): Record<string, unknown> | undefined {
   if (!isRecord(known) || !isRecord(next)) return undefined;
   for (const key of Object.keys(known)) {
-    if (!(key in next)) return undefined;
+    if (!Object.hasOwn(next, key)) return undefined;
   }
-  const diff: Record<string, unknown> = {};
+  const diff: [string, unknown][] = [];
   for (const [key, value] of Object.entries(next)) {
-    if (!(key in known) || !isJSONValueEqual(known[key], value)) {
-      diff[key] = value;
+    if (!Object.hasOwn(known, key) || !isJSONValueEqual(known[key], value)) {
+      diff.push([key, value]);
     }
   }
-  const changed = Object.keys(diff).length;
+  const changed = diff.length;
   if (changed === 0 || changed === Object.keys(next).length) return undefined;
-  return diff;
+  return Object.fromEntries(diff);
 }
 
 type ToolCallLikePart = {

package/src/react/RuntimeAdapter.test.ts

@@ -5,13 +5,16 @@ import { RuntimeAdapter } from "../store";
 import { ExternalStoreRuntimeCore } from "../runtimes/internal";
 import { AssistantRuntimeImpl } from "../runtime/internal";
 import type { ExternalStoreAdapter } from "../runtimes/external-store/external-store-adapter";
+import type { AppendMessage } from "../types/message";
 
 type DemoMessage = { role: "user" | "assistant"; text: string };
 
 const createRuntime = (
   overrides?: Partial<ExternalStoreAdapter<DemoMessage>>,
 ) => {
-  const onNew = vi.fn(async () => {});
+  const onNew = vi.fn<(message: AppendMessage) => Promise<void>>(
+    async () => {},
+  );
   const adapter: ExternalStoreAdapter<DemoMessage> = {
     messages: [{ role: "user", text: "hi" }],
     convertMessage: (message) => ({