@assistant-ui/core 0.3.14 → 0.3.15

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

@@ -9,6 +9,45 @@ describe("AssistantFrameProvider", () => {
   let messageHandler: ((event: MessageEvent) => void) | undefined;
   let parentWindow: Window;
 
+  const dispatchToolCall = (
+    origin: string,
+    source: Window = parentWindow,
+    id = "tool-call-1",
+  ) => {
+    messageHandler?.(
+      new MessageEvent("message", {
+        data: {
+          channel: FRAME_MESSAGE_CHANNEL,
+          message: {
+            type: "tool-call",
+            id,
+            toolName: "sensitiveTool",
+            args: {},
+          },
+        },
+        origin,
+        source,
+      }),
+    );
+  };
+
+  const dispatchToolCancel = (
+    origin: string,
+    source: Window = parentWindow,
+    id = "tool-call-1",
+  ) => {
+    messageHandler?.(
+      new MessageEvent("message", {
+        data: {
+          channel: FRAME_MESSAGE_CHANNEL,
+          message: { type: "tool-cancel", id },
+        },
+        origin,
+        source,
+      }),
+    );
+  };
+
   beforeEach(() => {
     parentWindow = {
       postMessage: vi.fn(),
@@ -47,37 +86,325 @@ describe("AssistantFrameProvider", () => {
       "https://parent.example",
     );
 
-    const toolCall = {
-      channel: FRAME_MESSAGE_CHANNEL,
-      message: {
-        type: "tool-call",
-        id: "tool-call-1",
-        toolName: "sensitiveTool",
-        args: {},
-      },
-    };
     const otherWindow = {
       postMessage: vi.fn(),
     } as unknown as Window;
 
-    messageHandler?.(
-      new MessageEvent("message", {
-        data: toolCall,
-        origin: "https://parent.example",
-        source: otherWindow,
+    dispatchToolCall("https://parent.example", otherWindow);
+
+    expect(execute).not.toHaveBeenCalled();
+
+    dispatchToolCall("https://parent.example");
+
+    await vi.waitFor(() => expect(execute).toHaveBeenCalledOnce());
+  });
+
+  it("aborts in-flight tool calls when the parent cancels them", async () => {
+    let toolSignal: AbortSignal | undefined;
+    const execute = vi.fn(
+      async (_args: unknown, context: { abortSignal: AbortSignal }) => {
+        toolSignal = context.abortSignal;
+        await new Promise<never>((_resolve, reject) => {
+          context.abortSignal.addEventListener(
+            "abort",
+            () => reject(context.abortSignal.reason),
+            { once: true },
+          );
+        });
+      },
+    );
+    AssistantFrameProvider.addModelContextProvider({
+      getModelContext: () => ({
+        tools: { sensitiveTool: { execute } },
+      }),
+    });
+
+    dispatchToolCall("*");
+    await vi.waitFor(() => expect(toolSignal).toBeDefined());
+
+    dispatchToolCancel("*");
+
+    expect(toolSignal?.aborted).toBe(true);
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    expect(parentWindow.postMessage).not.toHaveBeenCalledWith(
+      expect.objectContaining({
+        message: expect.objectContaining({ type: "tool-result" }),
       }),
+      expect.anything(),
+    );
+  });
+
+  it("cancels only the matching in-flight tool call", async () => {
+    const signals = new Map<string, AbortSignal>();
+    const execute = vi.fn(
+      async (
+        _args: unknown,
+        context: { toolCallId: string; abortSignal: AbortSignal },
+      ) => {
+        signals.set(context.toolCallId, context.abortSignal);
+        await new Promise<never>((_resolve, reject) => {
+          context.abortSignal.addEventListener(
+            "abort",
+            () => reject(context.abortSignal.reason),
+            { once: true },
+          );
+        });
+      },
+    );
+    AssistantFrameProvider.addModelContextProvider({
+      getModelContext: () => ({ tools: { sensitiveTool: { execute } } }),
+    });
+
+    dispatchToolCall("*", parentWindow, "tool-a");
+    dispatchToolCall("*", parentWindow, "tool-b");
+    await vi.waitFor(() => expect(signals.size).toBe(2));
+
+    dispatchToolCancel("*", parentWindow, "tool-a");
+
+    expect(signals.get("tool-a")?.aborted).toBe(true);
+    expect(signals.get("tool-b")?.aborted).toBe(false);
+  });
+
+  it("aborts an earlier call when a duplicate ID arrives", async () => {
+    const signals: AbortSignal[] = [];
+    const execute = vi.fn(
+      async (_args: unknown, context: { abortSignal: AbortSignal }) => {
+        signals.push(context.abortSignal);
+        await new Promise<never>((_resolve, reject) => {
+          context.abortSignal.addEventListener(
+            "abort",
+            () => reject(context.abortSignal.reason),
+            { once: true },
+          );
+        });
+      },
+    );
+    AssistantFrameProvider.addModelContextProvider({
+      getModelContext: () => ({ tools: { sensitiveTool: { execute } } }),
+    });
+
+    dispatchToolCall("*", parentWindow, "duplicate");
+    await vi.waitFor(() => expect(signals).toHaveLength(1));
+    dispatchToolCall("*", parentWindow, "duplicate");
+    await vi.waitFor(() => expect(signals).toHaveLength(2));
+
+    expect(signals[0]?.aborted).toBe(true);
+    expect(signals[1]?.aborted).toBe(false);
+  });
+
+  it("aborts in-flight tool calls when the provider is disposed", async () => {
+    let toolSignal: AbortSignal | undefined;
+    const execute = vi.fn(
+      async (_args: unknown, context: { abortSignal: AbortSignal }) => {
+        toolSignal = context.abortSignal;
+        await new Promise<never>((_resolve, reject) => {
+          context.abortSignal.addEventListener(
+            "abort",
+            () => reject(context.abortSignal.reason),
+            { once: true },
+          );
+        });
+      },
+    );
+    AssistantFrameProvider.addModelContextProvider({
+      getModelContext: () => ({ tools: { sensitiveTool: { execute } } }),
+    });
+
+    dispatchToolCall("*");
+    await vi.waitFor(() => expect(toolSignal).toBeDefined());
+
+    AssistantFrameProvider.dispose();
+
+    expect(toolSignal?.aborted).toBe(true);
+    expect(parentWindow.postMessage).toHaveBeenCalledWith(
+      {
+        channel: FRAME_MESSAGE_CHANNEL,
+        message: {
+          type: "tool-result",
+          id: "tool-call-1",
+          error: "AssistantFrameProvider has been disposed",
+        },
+      },
+      { targetOrigin: "*" },
+    );
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    const toolResults = vi
+      .mocked(parentWindow.postMessage)
+      .mock.calls.filter(
+        ([data]) =>
+          (data as { message?: { type?: string } }).message?.type ===
+          "tool-result",
+      );
+    expect(toolResults).toHaveLength(1);
+  });
+
+  it("upgrades a wildcard origin policy when a strict provider registers", async () => {
+    AssistantFrameProvider.addModelContextProvider(
+      { getModelContext: () => ({}) },
+      "*",
     );
 
+    const execute = vi.fn(async () => "result");
+    AssistantFrameProvider.addModelContextProvider(
+      {
+        getModelContext: () => ({
+          tools: {
+            sensitiveTool: { execute },
+          },
+        }),
+      },
+      "https://parent.example",
+    );
+
+    expect(parentWindow.postMessage).toHaveBeenLastCalledWith(
+      expect.anything(),
+      "https://parent.example",
+    );
+
+    dispatchToolCall("https://untrusted.example");
+
     expect(execute).not.toHaveBeenCalled();
 
-    messageHandler?.(
-      new MessageEvent("message", {
-        data: toolCall,
-        origin: "https://parent.example",
-        source: parentWindow,
-      }),
+    dispatchToolCall("https://parent.example");
+
+    await vi.waitFor(() => expect(execute).toHaveBeenCalledOnce());
+  });
+
+  it("does not downgrade a strict origin policy for a wildcard provider", async () => {
+    const execute = vi.fn(async () => "result");
+    AssistantFrameProvider.addModelContextProvider(
+      {
+        getModelContext: () => ({
+          tools: {
+            sensitiveTool: { execute },
+          },
+        }),
+      },
+      "https://parent.example",
+    );
+    AssistantFrameProvider.addModelContextProvider(
+      { getModelContext: () => ({}) },
+      "*",
+    );
+
+    dispatchToolCall("https://untrusted.example");
+
+    expect(execute).not.toHaveBeenCalled();
+
+    dispatchToolCall("https://parent.example");
+
+    await vi.waitFor(() => expect(execute).toHaveBeenCalledOnce());
+  });
+
+  it("rejects conflicting strict origin policies", () => {
+    AssistantFrameProvider.addModelContextProvider(
+      { getModelContext: () => ({}) },
+      "https://first.example",
+    );
+
+    expect(() =>
+      AssistantFrameProvider.addModelContextProvider(
+        { getModelContext: () => ({}) },
+        "https://second.example",
+      ),
+    ).toThrow(
+      'AssistantFrameProvider cannot register conflicting target origins: "https://first.example" and "https://second.example"',
+    );
+  });
+
+  it("resets the origin policy after every provider unsubscribes", () => {
+    const unsubscribe = AssistantFrameProvider.addModelContextProvider(
+      { getModelContext: () => ({}) },
+      "https://first.example",
+    );
+
+    unsubscribe();
+
+    expect(parentWindow.postMessage).toHaveBeenLastCalledWith(
+      expect.anything(),
+      "*",
+    );
+
+    expect(() =>
+      AssistantFrameProvider.addModelContextProvider(
+        { getModelContext: () => ({}) },
+        "https://second.example",
+      ),
+    ).not.toThrow();
+    expect(parentWindow.postMessage).toHaveBeenLastCalledWith(
+      expect.anything(),
+      "https://second.example",
     );
+  });
+
+  it("treats a second unsubscribe as a no-op", () => {
+    const unsubscribe = AssistantFrameProvider.addModelContextProvider(
+      { getModelContext: () => ({}) },
+      "https://first.example",
+    );
+
+    unsubscribe();
+    unsubscribe();
+
+    const unsubscribeSecond = AssistantFrameProvider.addModelContextProvider(
+      { getModelContext: () => ({}) },
+      "https://second.example",
+    );
+    expect(parentWindow.postMessage).toHaveBeenLastCalledWith(
+      expect.anything(),
+      "https://second.example",
+    );
+
+    unsubscribeSecond();
+    unsubscribeSecond();
+
+    expect(parentWindow.postMessage).toHaveBeenLastCalledWith(
+      expect.anything(),
+      "*",
+    );
+  });
+
+  it("recomputes the origin policy from providers that remain", () => {
+    AssistantFrameProvider.addModelContextProvider(
+      { getModelContext: () => ({}) },
+      "*",
+    );
+    const unsubscribeStrict = AssistantFrameProvider.addModelContextProvider(
+      { getModelContext: () => ({}) },
+      "https://parent.example",
+    );
+
+    unsubscribeStrict();
+
+    expect(parentWindow.postMessage).toHaveBeenLastCalledWith(
+      expect.anything(),
+      "*",
+    );
+  });
+
+  it("keeps a shared strict origin after one provider unsubscribes", async () => {
+    const unsubscribeFirst = AssistantFrameProvider.addModelContextProvider(
+      { getModelContext: () => ({}) },
+      "https://parent.example",
+    );
+    const execute = vi.fn(async () => "result");
+    AssistantFrameProvider.addModelContextProvider(
+      {
+        getModelContext: () => ({
+          tools: {
+            sensitiveTool: { execute },
+          },
+        }),
+      },
+      "https://parent.example",
+    );
+
+    unsubscribeFirst();
+
+    dispatchToolCall("https://untrusted.example");
+    expect(execute).not.toHaveBeenCalled();
 
+    dispatchToolCall("https://parent.example");
     await vi.waitFor(() => expect(execute).toHaveBeenCalledOnce());
   });
 });

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

@@ -37,7 +37,12 @@ export class AssistantFrameProvider {
     ModelContextProvider,
     Unsubscribe | undefined
   >();
+  private _activeToolCalls = new Map<
+    string,
+    { abortController: AbortController; event: MessageEvent }
+  >();
   private _targetOrigin: string;
+  private _strictRegistrations = 0;
 
   private constructor(targetOrigin: string = "*") {
     this._targetOrigin = targetOrigin;
@@ -52,10 +57,25 @@ export class AssistantFrameProvider {
       AssistantFrameProvider._instance = new AssistantFrameProvider(
         targetOrigin,
       );
+    } else {
+      AssistantFrameProvider._instance.reconcileTargetOrigin(targetOrigin);
     }
     return AssistantFrameProvider._instance;
   }
 
+  private reconcileTargetOrigin(targetOrigin: string = "*") {
+    if (targetOrigin === "*" || targetOrigin === this._targetOrigin) return;
+
+    if (this._targetOrigin === "*") {
+      this._targetOrigin = targetOrigin;
+      return;
+    }
+
+    throw new Error(
+      `AssistantFrameProvider cannot register conflicting target origins: "${this._targetOrigin}" and "${targetOrigin}"`,
+    );
+  }
+
   private handleMessage(event: MessageEvent) {
     if (this._targetOrigin !== "*" && event.origin !== this._targetOrigin)
       return;
@@ -75,6 +95,10 @@ export class AssistantFrameProvider {
       case "tool-call":
         this.handleToolCall(message, event);
         break;
+
+      case "tool-cancel":
+        this.cancelToolCall(message.id);
+        break;
     }
   }
 
@@ -83,6 +107,10 @@ export class AssistantFrameProvider {
     event: MessageEvent,
   ) {
     const tool = this.getModelContext().tools?.[message.toolName];
+    const abortController = new AbortController();
+    this._activeToolCalls.get(message.id)?.abortController.abort();
+    const activeCall = { abortController, event };
+    this._activeToolCalls.set(message.id, activeCall);
 
     let result: any;
     let error: string | undefined;
@@ -94,7 +122,7 @@ export class AssistantFrameProvider {
         result = tool.execute
           ? await tool.execute(message.args, {
               toolCallId: message.id,
-              abortSignal: new AbortController().signal,
+              abortSignal: abortController.signal,
               human: async () => {
                 throw new Error(
                   "Tool human input is not supported in frame context",
@@ -107,6 +135,9 @@ 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,
@@ -114,6 +145,13 @@ export class AssistantFrameProvider {
     });
   }
 
+  private cancelToolCall(id: string) {
+    const activeCall = this._activeToolCalls.get(id);
+    if (!activeCall) return;
+    this._activeToolCalls.delete(id);
+    activeCall.abortController.abort();
+  }
+
   private sendMessage(event: MessageEvent, message: FrameMessage) {
     event.source?.postMessage(
       { channel: FRAME_MESSAGE_CHANNEL, message },
@@ -157,8 +195,10 @@ export class AssistantFrameProvider {
     provider: ModelContextProvider,
     targetOrigin?: string,
   ): Unsubscribe {
-    const instance = AssistantFrameProvider.getInstance(targetOrigin);
+    const origin = targetOrigin ?? "*";
+    const instance = AssistantFrameProvider.getInstance(origin);
     instance._providers.add(provider);
+    if (origin !== "*") instance._strictRegistrations += 1;
 
     const unsubscribe = provider.subscribe?.(() => instance.broadcastUpdate());
     if (unsubscribe) {
@@ -167,10 +207,17 @@ export class AssistantFrameProvider {
 
     instance.broadcastUpdate();
 
+    let released = false;
     return () => {
+      if (released) return;
+      released = true;
       instance._providers.delete(provider);
       instance._providerUnsubscribes.get(provider)?.();
       instance._providerUnsubscribes.delete(provider);
+      if (origin !== "*") {
+        instance._strictRegistrations -= 1;
+        if (instance._strictRegistrations === 0) instance._targetOrigin = "*";
+      }
       instance.broadcastUpdate();
     };
   }
@@ -183,6 +230,15 @@ export class AssistantFrameProvider {
       instance._providerUnsubscribes.forEach((unsubscribe) => unsubscribe?.());
       instance._providerUnsubscribes.clear();
       instance._providers.clear();
+      instance._activeToolCalls.forEach(({ abortController, event }, id) => {
+        abortController.abort();
+        instance.sendMessage(event, {
+          type: "tool-result",
+          id,
+          error: "AssistantFrameProvider has been disposed",
+        });
+      });
+      instance._activeToolCalls.clear();
 
       AssistantFrameProvider._instance = null;
     }

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

@@ -14,6 +14,7 @@ export type FrameMessageType =
   | "model-context-request"
   | "model-context-update"
   | "tool-call"
+  | "tool-cancel"
   | "tool-result";
 
 export type FrameMessage =
@@ -30,6 +31,10 @@ export type FrameMessage =
       toolName: string;
       args: unknown;
     }
+  | {
+      type: "tool-cancel";
+      id: string;
+    }
   | {
       type: "tool-result";
       id: string;