@assistant-ui/core 0.3.14 → 0.3.16

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

@@ -8,7 +8,9 @@ const executionContext = {
   human: async () => undefined,
 };
 
-const createHost = () => {
+const DEFAULT_ORIGIN = "https://host.example";
+
+const createHost = (targetOrigin?: string) => {
   let handleMessage: ((event: MessageEvent) => void) | undefined;
   const addEventListener = vi.fn(
     (_type: string, listener: EventListenerOrEventListenerObject) => {
@@ -16,15 +18,23 @@ const createHost = () => {
     },
   );
   const removeEventListener = vi.fn();
-  vi.stubGlobal("window", { addEventListener, removeEventListener });
+  vi.stubGlobal("window", {
+    addEventListener,
+    removeEventListener,
+    location: { origin: DEFAULT_ORIGIN },
+  });
 
   const postMessage = vi.fn();
   const iframeWindow = { postMessage } as unknown as Window;
-  const host = new AssistantFrameHost(iframeWindow);
+  const host = new AssistantFrameHost(iframeWindow, targetOrigin);
 
-  const dispatchMessage = (message: FrameMessage) =>
+  const dispatchMessage = (
+    message: FrameMessage,
+    origin = targetOrigin ?? DEFAULT_ORIGIN,
+  ) =>
     handleMessage?.({
       source: iframeWindow,
+      origin,
       data: {
         channel: FRAME_MESSAGE_CHANNEL,
         message,
@@ -46,7 +56,15 @@ const createHost = () => {
   const execute = host.getModelContext().tools?.search?.execute;
   if (!execute) throw new Error("Expected the search tool to be available");
 
-  return { dispatchMessage, execute, host, postMessage };
+  const getToolCallId = () => {
+    const call = postMessage.mock.calls.find(
+      ([data]) => data.message.type === "tool-call",
+    );
+    if (!call) throw new Error("Expected a tool call to be posted");
+    return call[0].message.id as string;
+  };
+
+  return { dispatchMessage, execute, getToolCallId, host, postMessage };
 };
 
 beforeEach(() => {
@@ -59,15 +77,54 @@ afterEach(() => {
 });
 
 describe("AssistantFrameHost", () => {
+  it("defaults to the current origin", () => {
+    const { host, postMessage } = createHost();
+
+    expect(postMessage).toHaveBeenCalledWith(expect.anything(), DEFAULT_ORIGIN);
+
+    host.dispose();
+  });
+
+  it("ignores context updates from another origin", () => {
+    const { dispatchMessage, host } = createHost();
+
+    dispatchMessage(
+      {
+        type: "model-context-update",
+        context: { system: "untrusted instructions" },
+      },
+      "https://untrusted.example",
+    );
+
+    expect(host.getModelContext().system).toBeUndefined();
+    host.dispose();
+  });
+
+  it("allows an explicit wildcard origin", () => {
+    const { dispatchMessage, host, postMessage } = createHost("*");
+
+    dispatchMessage(
+      {
+        type: "model-context-update",
+        context: { system: "cross-origin instructions" },
+      },
+      "https://iframe.example",
+    );
+
+    expect(postMessage).toHaveBeenCalledWith(expect.anything(), "*");
+    expect(host.getModelContext().system).toBe("cross-origin instructions");
+    host.dispose();
+  });
+
   it("resolves tool calls from frame results", async () => {
-    const { dispatchMessage, execute, host } = createHost();
+    const { dispatchMessage, execute, getToolCallId, host } = createHost();
     const result = Promise.resolve(
       execute({ query: "weather" }, executionContext),
     );
 
     dispatchMessage({
       type: "tool-result",
-      id: "tool-0",
+      id: getToolCallId(),
       result: "sunny",
     });
 
@@ -76,20 +133,63 @@ describe("AssistantFrameHost", () => {
     host.dispose();
   });
 
+  it("rejects tool results that carry an empty error message", async () => {
+    const { dispatchMessage, execute, getToolCallId, host } = createHost();
+    const result = Promise.resolve(
+      execute({ query: "weather" }, executionContext),
+    );
+
+    dispatchMessage({
+      type: "tool-result",
+      id: getToolCallId(),
+      error: "",
+    });
+
+    await expect(result).rejects.toThrow();
+    expect(vi.getTimerCount()).toBe(0);
+    host.dispose();
+  });
+
+  it("resolves results from guests that serialize an absent error as null", async () => {
+    const { dispatchMessage, execute, getToolCallId, host } = createHost();
+    const result = Promise.resolve(
+      execute({ query: "weather" }, executionContext),
+    );
+
+    dispatchMessage({
+      type: "tool-result",
+      id: getToolCallId(),
+      result: "sunny",
+      error: null,
+    } as unknown as FrameMessage);
+
+    await expect(result).resolves.toBe("sunny");
+    expect(vi.getTimerCount()).toBe(0);
+    host.dispose();
+  });
+
   it("rejects pending tool calls when disposed", async () => {
-    const { execute, host } = createHost();
+    const { execute, getToolCallId, host, postMessage } = createHost();
     const result = Promise.resolve(execute({}, executionContext));
+    const toolCallId = getToolCallId();
 
     host.dispose();
 
     await expect(result).rejects.toThrow(
       "AssistantFrameHost has been disposed",
     );
+    expect(postMessage).toHaveBeenLastCalledWith(
+      {
+        channel: FRAME_MESSAGE_CHANNEL,
+        message: { type: "tool-cancel", id: toolCallId },
+      },
+      DEFAULT_ORIGIN,
+    );
     expect(vi.getTimerCount()).toBe(0);
   });
 
   it("rejects pending tool calls when execution is aborted", async () => {
-    const { execute, host } = createHost();
+    const { execute, getToolCallId, host, postMessage } = createHost();
     const abortController = new AbortController();
     const abortError = new Error("Run cancelled");
     abortError.name = "AbortError";
@@ -104,15 +204,58 @@ describe("AssistantFrameHost", () => {
     );
     const onRejected = vi.fn();
     void result.catch(onRejected);
+    const toolCallId = getToolCallId();
 
     abortController.abort(abortError);
     await Promise.resolve();
 
     expect(onRejected).toHaveBeenCalledWith(abortError);
+    expect(postMessage).toHaveBeenLastCalledWith(
+      {
+        channel: FRAME_MESSAGE_CHANNEL,
+        message: { type: "tool-cancel", id: toolCallId },
+      },
+      DEFAULT_ORIGIN,
+    );
+    expect(vi.getTimerCount()).toBe(0);
+    host.dispose();
+  });
+
+  it("cancels tool calls when they time out", async () => {
+    const { execute, getToolCallId, host, postMessage } = createHost();
+    const result = Promise.resolve(execute({}, executionContext));
+    const toolCallId = getToolCallId();
+    const rejection = expect(result).rejects.toThrow(
+      'Tool call "search" timed out',
+    );
+
+    await vi.advanceTimersByTimeAsync(30000);
+
+    await rejection;
+    expect(postMessage).toHaveBeenLastCalledWith(
+      {
+        channel: FRAME_MESSAGE_CHANNEL,
+        message: { type: "tool-cancel", id: toolCallId },
+      },
+      DEFAULT_ORIGIN,
+    );
     expect(vi.getTimerCount()).toBe(0);
     host.dispose();
   });
 
+  it("uses unique tool IDs across host instances", () => {
+    const first = createHost();
+    const second = createHost();
+
+    void first.execute({}, executionContext).catch(() => undefined);
+    void second.execute({}, executionContext).catch(() => undefined);
+
+    expect(first.getToolCallId()).not.toBe(second.getToolCallId());
+
+    first.host.dispose();
+    second.host.dispose();
+  });
+
   it("does not post tool calls when execution is already aborted", async () => {
     const { execute, host, postMessage } = createHost();
     const abortController = new AbortController();

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

@@ -2,6 +2,7 @@ import type { ModelContextProvider, ModelContext } from "../types";
 import type { Unsubscribe } from "../../types/unsubscribe";
 import type { Tool } from "assistant-stream";
 import { notifySubscribers as notifyStateSubscribers } from "../../subscribable/subscribable";
+import { generateId } from "../../utils/id";
 import {
   type FrameMessage,
   FRAME_MESSAGE_CHANNEL,
@@ -9,6 +10,8 @@ import {
   type SerializedTool,
 } from "./types";
 
+const getDefaultTargetOrigin = () => window.location.origin;
+
 /**
  * Deserializes tools from JSON Schema format back to Tool objects
  */
@@ -58,12 +61,14 @@ export class AssistantFrameHost implements ModelContextProvider {
       reject: (error: any) => void;
     }
   >();
-  private _requestCounter = 0;
   private _iframeWindow: Window;
   private _targetOrigin: string;
   private _disposed = false;
 
-  constructor(iframeWindow: Window, targetOrigin: string = "*") {
+  constructor(
+    iframeWindow: Window,
+    targetOrigin: string = getDefaultTargetOrigin(),
+  ) {
     this._iframeWindow = iframeWindow;
     this._targetOrigin = targetOrigin;
 
@@ -90,7 +95,7 @@ export class AssistantFrameHost implements ModelContextProvider {
       case "tool-result": {
         const pending = this._pendingRequests.get(message.id);
         if (pending) {
-          if (message.error) {
+          if (typeof message.error === "string") {
             pending.reject(new Error(message.error));
           } else {
             pending.resolve(message.result);
@@ -130,7 +135,7 @@ export class AssistantFrameHost implements ModelContextProvider {
     return this.sendRequest(
       {
         type: "tool-call",
-        id: `tool-${this._requestCounter++}`,
+        id: `tool-${generateId()}`,
         toolName,
         args,
       },
@@ -159,6 +164,7 @@ export class AssistantFrameHost implements ModelContextProvider {
         if (!abortSignal) return;
         const pending = this._pendingRequests.get(message.id);
         if (pending) {
+          this.cancelToolCall(message.id);
           pending.reject(getAbortReason(abortSignal));
           this._pendingRequests.delete(message.id);
         }
@@ -182,6 +188,7 @@ export class AssistantFrameHost implements ModelContextProvider {
       timeoutId = setTimeout(() => {
         const pending = this._pendingRequests.get(message.id);
         if (pending) {
+          this.cancelToolCall(message.id);
           pending.reject(new Error(timeoutMessage));
           this._pendingRequests.delete(message.id);
         }
@@ -195,6 +202,16 @@ 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 requestContext() {
     this._iframeWindow.postMessage(
       {
@@ -225,7 +242,8 @@ export class AssistantFrameHost implements ModelContextProvider {
     window.removeEventListener("message", this.handleMessage);
     this._subscribers.clear();
     const error = new Error("AssistantFrameHost has been disposed");
-    for (const pending of this._pendingRequests.values()) {
+    for (const [id, pending] of this._pendingRequests) {
+      this.cancelToolCall(id);
       pending.reject(error);
     }
     this._pendingRequests.clear();