@assistant-ui/core 0.3.15 → 0.3.17

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

@@ -13,6 +13,7 @@ describe("AssistantFrameProvider", () => {
     origin: string,
     source: Window = parentWindow,
     id = "tool-call-1",
+    toolName = "sensitiveTool",
   ) => {
     messageHandler?.(
       new MessageEvent("message", {
@@ -21,7 +22,7 @@ describe("AssistantFrameProvider", () => {
           message: {
             type: "tool-call",
             id,
-            toolName: "sensitiveTool",
+            toolName,
             args: {},
           },
         },
@@ -99,6 +100,52 @@ describe("AssistantFrameProvider", () => {
     await vi.waitFor(() => expect(execute).toHaveBeenCalledOnce());
   });
 
+  it("defaults to the current origin", async () => {
+    const execute = vi.fn(async () => "result");
+    AssistantFrameProvider.addModelContextProvider({
+      getModelContext: () => ({
+        tools: { sensitiveTool: { execute } },
+      }),
+    });
+
+    expect(parentWindow.postMessage).toHaveBeenLastCalledWith(
+      expect.anything(),
+      window.location.origin,
+    );
+
+    dispatchToolCall("https://untrusted.example");
+    expect(execute).not.toHaveBeenCalled();
+
+    dispatchToolCall(window.location.origin);
+    await vi.waitFor(() => expect(execute).toHaveBeenCalledOnce());
+  });
+
+  it("reports a failure even when the thrown error has an empty message", async () => {
+    const execute = vi.fn(async () => {
+      throw new Error();
+    });
+    AssistantFrameProvider.addModelContextProvider({
+      getModelContext: () => ({
+        tools: {
+          sensitiveTool: { execute },
+        },
+      }),
+    });
+
+    dispatchToolCall(window.location.origin);
+
+    await vi.waitFor(() => {
+      const frame = (
+        parentWindow.postMessage as ReturnType<typeof vi.fn>
+      ).mock.calls
+        .map(([data]) => data)
+        .find((data) => data?.message?.type === "tool-result");
+      expect(frame).toBeDefined();
+      expect(frame.message).toHaveProperty("error");
+      expect(frame.message).not.toHaveProperty("result");
+    });
+  });
+
   it("aborts in-flight tool calls when the parent cancels them", async () => {
     let toolSignal: AbortSignal | undefined;
     const execute = vi.fn(
@@ -119,10 +166,10 @@ describe("AssistantFrameProvider", () => {
       }),
     });
 
-    dispatchToolCall("*");
+    dispatchToolCall(window.location.origin);
     await vi.waitFor(() => expect(toolSignal).toBeDefined());
 
-    dispatchToolCancel("*");
+    dispatchToolCancel(window.location.origin);
 
     expect(toolSignal?.aborted).toBe(true);
     await new Promise((resolve) => setTimeout(resolve, 0));
@@ -155,11 +202,11 @@ describe("AssistantFrameProvider", () => {
       getModelContext: () => ({ tools: { sensitiveTool: { execute } } }),
     });
 
-    dispatchToolCall("*", parentWindow, "tool-a");
-    dispatchToolCall("*", parentWindow, "tool-b");
+    dispatchToolCall(window.location.origin, parentWindow, "tool-a");
+    dispatchToolCall(window.location.origin, parentWindow, "tool-b");
     await vi.waitFor(() => expect(signals.size).toBe(2));
 
-    dispatchToolCancel("*", parentWindow, "tool-a");
+    dispatchToolCancel(window.location.origin, parentWindow, "tool-a");
 
     expect(signals.get("tool-a")?.aborted).toBe(true);
     expect(signals.get("tool-b")?.aborted).toBe(false);
@@ -183,9 +230,9 @@ describe("AssistantFrameProvider", () => {
       getModelContext: () => ({ tools: { sensitiveTool: { execute } } }),
     });
 
-    dispatchToolCall("*", parentWindow, "duplicate");
+    dispatchToolCall(window.location.origin, parentWindow, "duplicate");
     await vi.waitFor(() => expect(signals).toHaveLength(1));
-    dispatchToolCall("*", parentWindow, "duplicate");
+    dispatchToolCall(window.location.origin, parentWindow, "duplicate");
     await vi.waitFor(() => expect(signals).toHaveLength(2));
 
     expect(signals[0]?.aborted).toBe(true);
@@ -210,7 +257,7 @@ describe("AssistantFrameProvider", () => {
       getModelContext: () => ({ tools: { sensitiveTool: { execute } } }),
     });
 
-    dispatchToolCall("*");
+    dispatchToolCall(window.location.origin);
     await vi.waitFor(() => expect(toolSignal).toBeDefined());
 
     AssistantFrameProvider.dispose();
@@ -225,7 +272,53 @@ describe("AssistantFrameProvider", () => {
           error: "AssistantFrameProvider has been disposed",
         },
       },
-      { targetOrigin: "*" },
+      { targetOrigin: window.location.origin },
+    );
+    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("aborts in-flight tool calls when their provider is removed", 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 },
+          );
+        });
+      },
+    );
+    const removeProvider = AssistantFrameProvider.addModelContextProvider({
+      getModelContext: () => ({ tools: { sensitiveTool: { execute } } }),
+    });
+
+    dispatchToolCall(window.location.origin);
+    await vi.waitFor(() => expect(toolSignal).toBeDefined());
+
+    removeProvider();
+
+    expect(toolSignal?.aborted).toBe(true);
+    expect(parentWindow.postMessage).toHaveBeenCalledWith(
+      {
+        channel: FRAME_MESSAGE_CHANNEL,
+        message: {
+          type: "tool-result",
+          id: "tool-call-1",
+          error: "AssistantFrame tool provider has been removed",
+        },
+      },
+      { targetOrigin: window.location.origin },
     );
     await new Promise((resolve) => setTimeout(resolve, 0));
     const toolResults = vi
@@ -238,6 +331,133 @@ describe("AssistantFrameProvider", () => {
     expect(toolResults).toHaveLength(1);
   });
 
+  it("keeps tool calls active while the same provider remains registered", 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 },
+          );
+        });
+      },
+    );
+    const provider = {
+      getModelContext: () => ({ tools: { sensitiveTool: { execute } } }),
+    };
+    const removeFirst =
+      AssistantFrameProvider.addModelContextProvider(provider);
+    const removeSecond =
+      AssistantFrameProvider.addModelContextProvider(provider);
+
+    dispatchToolCall(window.location.origin);
+    await vi.waitFor(() => expect(toolSignal).toBeDefined());
+
+    removeFirst();
+    expect(toolSignal?.aborted).toBe(false);
+
+    removeSecond();
+    expect(toolSignal?.aborted).toBe(true);
+  });
+
+  it("cancels only calls owned by the removed provider", async () => {
+    const shadowedExecute = vi.fn(async () => "shadowed");
+    const removeShadowed = AssistantFrameProvider.addModelContextProvider({
+      getModelContext: () => ({
+        tools: { sensitiveTool: { execute: shadowedExecute } },
+      }),
+    });
+
+    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 },
+          );
+        });
+      },
+    );
+    const removeOwner = AssistantFrameProvider.addModelContextProvider({
+      getModelContext: () => ({
+        tools: { sensitiveTool: { execute } },
+      }),
+    });
+
+    dispatchToolCall(window.location.origin);
+    await vi.waitFor(() => expect(toolSignal).toBeDefined());
+    expect(shadowedExecute).not.toHaveBeenCalled();
+
+    removeShadowed();
+    expect(toolSignal?.aborted).toBe(false);
+
+    removeOwner();
+    expect(toolSignal?.aborted).toBe(true);
+  });
+
+  it("keeps other providers' tool calls active when one is removed", 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 },
+          );
+        });
+      },
+    );
+    const removeFirst = AssistantFrameProvider.addModelContextProvider({
+      getModelContext: () => ({ tools: { firstTool: { execute } } }),
+    });
+    const removeSecond = AssistantFrameProvider.addModelContextProvider({
+      getModelContext: () => ({ tools: { secondTool: { execute } } }),
+    });
+
+    dispatchToolCall(
+      window.location.origin,
+      parentWindow,
+      "first-call",
+      "firstTool",
+    );
+    dispatchToolCall(
+      window.location.origin,
+      parentWindow,
+      "second-call",
+      "secondTool",
+    );
+    await vi.waitFor(() => expect(signals.size).toBe(2));
+
+    removeFirst();
+
+    expect(signals.get("first-call")?.aborted).toBe(true);
+    expect(signals.get("second-call")?.aborted).toBe(false);
+    const toolResults = vi
+      .mocked(parentWindow.postMessage)
+      .mock.calls.filter(
+        ([data]) =>
+          (data as { message?: { type?: string } }).message?.type ===
+          "tool-result",
+      );
+    expect(toolResults).toHaveLength(1);
+    expect(toolResults[0]?.[0]).toMatchObject({
+      message: { id: "first-call" },
+    });
+
+    removeSecond();
+  });
+
   it("upgrades a wildcard origin policy when a strict provider registers", async () => {
     AssistantFrameProvider.addModelContextProvider(
       { getModelContext: () => ({}) },
@@ -312,7 +532,227 @@ describe("AssistantFrameProvider", () => {
     );
   });
 
-  it("resets the origin policy after every provider unsubscribes", () => {
+  it("rolls back a provider when registration fails", () => {
+    const execute = vi.fn(async () => "result");
+    expect(() =>
+      AssistantFrameProvider.addModelContextProvider(
+        {
+          getModelContext: () => ({
+            tools: { sensitiveTool: { execute } },
+          }),
+          subscribe: () => {
+            throw new Error("subscribe failed");
+          },
+        },
+        "https://first.example",
+      ),
+    ).toThrow("subscribe failed");
+
+    dispatchToolCall("https://first.example");
+    expect(execute).not.toHaveBeenCalled();
+
+    expect(() =>
+      AssistantFrameProvider.addModelContextProvider(
+        { getModelContext: () => ({}) },
+        "https://second.example",
+      ),
+    ).not.toThrow();
+  });
+
+  it("cancels reentrant tool calls when registration rolls back", 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 },
+          );
+        });
+      },
+    );
+
+    expect(() =>
+      AssistantFrameProvider.addModelContextProvider({
+        getModelContext: () => ({
+          tools: { sensitiveTool: { execute } },
+        }),
+        subscribe: () => {
+          dispatchToolCall(window.location.origin);
+          throw new Error("subscribe failed");
+        },
+      }),
+    ).toThrow("subscribe failed");
+
+    expect(toolSignal?.aborted).toBe(true);
+    expect(parentWindow.postMessage).toHaveBeenCalledWith(
+      {
+        channel: FRAME_MESSAGE_CHANNEL,
+        message: {
+          type: "tool-result",
+          id: "tool-call-1",
+          error: "AssistantFrame tool provider has been removed",
+        },
+      },
+      { targetOrigin: window.location.origin },
+    );
+    await new Promise((resolve) => setTimeout(resolve, 0));
+  });
+
+  it("keeps an existing registration when the same provider fails to register again", async () => {
+    const execute = vi.fn(async () => "result");
+    const firstUnsubscribe = vi.fn();
+    let subscriptionCount = 0;
+    const provider = {
+      getModelContext: () => ({
+        tools: { sensitiveTool: { execute } },
+      }),
+      subscribe: () => {
+        subscriptionCount += 1;
+        if (subscriptionCount === 1) return firstUnsubscribe;
+        throw new Error("second subscribe failed");
+      },
+    };
+    const releaseFirst = AssistantFrameProvider.addModelContextProvider(
+      provider,
+      "https://parent.example",
+    );
+
+    expect(() =>
+      AssistantFrameProvider.addModelContextProvider(
+        provider,
+        "https://parent.example",
+      ),
+    ).toThrow("second subscribe failed");
+    expect(firstUnsubscribe).not.toHaveBeenCalled();
+
+    dispatchToolCall("https://parent.example");
+
+    await vi.waitFor(() => expect(execute).toHaveBeenCalledOnce());
+    releaseFirst();
+    expect(firstUnsubscribe).toHaveBeenCalledOnce();
+  });
+
+  it("merges a provider registered more than once only once", () => {
+    const provider = {
+      getModelContext: () => ({ system: "shared system" }),
+    };
+
+    AssistantFrameProvider.addModelContextProvider(provider);
+    AssistantFrameProvider.addModelContextProvider(provider);
+
+    expect(parentWindow.postMessage).toHaveBeenLastCalledWith(
+      {
+        channel: FRAME_MESSAGE_CHANNEL,
+        message: {
+          type: "model-context-update",
+          context: {
+            system: "shared system",
+            tools: {},
+          },
+        },
+      },
+      window.location.origin,
+    );
+  });
+
+  it("releases a subscription when the initial broadcast fails", () => {
+    const unsubscribe = vi.fn();
+    expect(() =>
+      AssistantFrameProvider.addModelContextProvider(
+        {
+          getModelContext: () => {
+            throw new Error("context failed");
+          },
+          subscribe: () => unsubscribe,
+        },
+        "https://first.example",
+      ),
+    ).toThrow("context failed");
+
+    expect(unsubscribe).toHaveBeenCalledOnce();
+    expect(() =>
+      AssistantFrameProvider.addModelContextProvider(
+        { getModelContext: () => ({}) },
+        "https://second.example",
+      ),
+    ).not.toThrow();
+  });
+
+  it("reports rollback cleanup failures without replacing the original error", () => {
+    const contextError = new Error("context failed");
+    const unsubscribeError = new Error("unsubscribe failed");
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => undefined);
+
+    expect(() =>
+      AssistantFrameProvider.addModelContextProvider({
+        getModelContext: () => {
+          throw contextError;
+        },
+        subscribe: () => () => {
+          throw unsubscribeError;
+        },
+      }),
+    ).toThrow(contextError);
+
+    expect(consoleError).toHaveBeenCalledWith(unsubscribeError);
+  });
+
+  it("cleans up provider state when its unsubscribe throws", () => {
+    const unsubscribe = vi.fn(() => {
+      throw new Error("unsubscribe failed");
+    });
+    const release = AssistantFrameProvider.addModelContextProvider(
+      {
+        getModelContext: () => ({}),
+        subscribe: () => unsubscribe,
+      },
+      "https://first.example",
+    );
+
+    expect(release).toThrow("unsubscribe failed");
+    expect(unsubscribe).toHaveBeenCalledOnce();
+    expect(() =>
+      AssistantFrameProvider.addModelContextProvider(
+        { getModelContext: () => ({}) },
+        "https://second.example",
+      ),
+    ).not.toThrow();
+  });
+
+  it("finishes disposal when a provider unsubscribe throws", () => {
+    const error = new Error("unsubscribe failed");
+    const firstUnsubscribe = vi.fn(() => {
+      throw error;
+    });
+    const secondUnsubscribe = vi.fn();
+    AssistantFrameProvider.addModelContextProvider({
+      getModelContext: () => ({}),
+      subscribe: () => firstUnsubscribe,
+    });
+    AssistantFrameProvider.addModelContextProvider({
+      getModelContext: () => ({}),
+      subscribe: () => secondUnsubscribe,
+    });
+
+    expect(() => AssistantFrameProvider.dispose()).toThrow(error);
+    expect(firstUnsubscribe).toHaveBeenCalledOnce();
+    expect(secondUnsubscribe).toHaveBeenCalledOnce();
+
+    expect(() =>
+      AssistantFrameProvider.addModelContextProvider(
+        { getModelContext: () => ({}) },
+        "https://new.example",
+      ),
+    ).not.toThrow();
+    expect(window.addEventListener).toHaveBeenCalledTimes(2);
+  });
+
+  it("returns to the same-origin policy after every provider unsubscribes", () => {
     const unsubscribe = AssistantFrameProvider.addModelContextProvider(
       { getModelContext: () => ({}) },
       "https://first.example",
@@ -322,7 +762,7 @@ describe("AssistantFrameProvider", () => {
 
     expect(parentWindow.postMessage).toHaveBeenLastCalledWith(
       expect.anything(),
-      "*",
+      window.location.origin,
     );
 
     expect(() =>
@@ -360,7 +800,7 @@ describe("AssistantFrameProvider", () => {
 
     expect(parentWindow.postMessage).toHaveBeenLastCalledWith(
       expect.anything(),
-      "*",
+      window.location.origin,
     );
   });
 
@@ -382,6 +822,37 @@ describe("AssistantFrameProvider", () => {
     );
   });
 
+  it("returns to the same-origin policy after a wildcard provider unsubscribes", () => {
+    const unsubscribe = AssistantFrameProvider.addModelContextProvider(
+      { getModelContext: () => ({}) },
+      "*",
+    );
+
+    unsubscribe();
+
+    expect(parentWindow.postMessage).toHaveBeenLastCalledWith(
+      expect.anything(),
+      window.location.origin,
+    );
+  });
+
+  it("allows opting back into a wildcard policy after every provider unsubscribes", () => {
+    const unsubscribe = AssistantFrameProvider.addModelContextProvider({
+      getModelContext: () => ({}),
+    });
+    unsubscribe();
+
+    AssistantFrameProvider.addModelContextProvider(
+      { getModelContext: () => ({}) },
+      "*",
+    );
+
+    expect(parentWindow.postMessage).toHaveBeenLastCalledWith(
+      expect.anything(),
+      "*",
+    );
+  });
+
   it("keeps a shared strict origin after one provider unsubscribes", async () => {
     const unsubscribeFirst = AssistantFrameProvider.addModelContextProvider(
       { getModelContext: () => ({}) },