@assistant-ui/core 0.3.15 → 0.3.16

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

@@ -99,6 +99,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 +165,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 +201,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 +229,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 +256,7 @@ describe("AssistantFrameProvider", () => {
       getModelContext: () => ({ tools: { sensitiveTool: { execute } } }),
     });
 
-    dispatchToolCall("*");
+    dispatchToolCall(window.location.origin);
     await vi.waitFor(() => expect(toolSignal).toBeDefined());
 
     AssistantFrameProvider.dispose();
@@ -225,7 +271,7 @@ describe("AssistantFrameProvider", () => {
           error: "AssistantFrameProvider has been disposed",
         },
       },
-      { targetOrigin: "*" },
+      { targetOrigin: window.location.origin },
     );
     await new Promise((resolve) => setTimeout(resolve, 0));
     const toolResults = vi
@@ -312,7 +358,185 @@ 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("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 +546,7 @@ describe("AssistantFrameProvider", () => {
 
     expect(parentWindow.postMessage).toHaveBeenLastCalledWith(
       expect.anything(),
-      "*",
+      window.location.origin,
     );
 
     expect(() =>
@@ -360,7 +584,7 @@ describe("AssistantFrameProvider", () => {
 
     expect(parentWindow.postMessage).toHaveBeenLastCalledWith(
       expect.anything(),
-      "*",
+      window.location.origin,
     );
   });
 
@@ -382,6 +606,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: () => ({}) },

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

@@ -29,22 +29,22 @@ const serializeModelContext = (
   }),
 });
 
+const getDefaultTargetOrigin = () => window.location.origin;
+
 export class AssistantFrameProvider {
   private static _instance: AssistantFrameProvider | null = null;
 
-  private _providers = new Set<ModelContextProvider>();
-  private _providerUnsubscribes = new Map<
-    ModelContextProvider,
-    Unsubscribe | undefined
-  >();
+  private _providers = new Map<symbol, ModelContextProvider>();
+  private _providerUnsubscribes = new Map<symbol, Unsubscribe | undefined>();
   private _activeToolCalls = new Map<
     string,
     { abortController: AbortController; event: MessageEvent }
   >();
   private _targetOrigin: string;
   private _strictRegistrations = 0;
+  private _wildcardRegistrations = 0;
 
-  private constructor(targetOrigin: string = "*") {
+  private constructor(targetOrigin: string = getDefaultTargetOrigin()) {
     this._targetOrigin = targetOrigin;
     this.handleMessage = this.handleMessage.bind(this);
     window.addEventListener("message", this.handleMessage);
@@ -63,8 +63,17 @@ export class AssistantFrameProvider {
     return AssistantFrameProvider._instance;
   }
 
-  private reconcileTargetOrigin(targetOrigin: string = "*") {
-    if (targetOrigin === "*" || targetOrigin === this._targetOrigin) return;
+  private reconcileTargetOrigin(
+    targetOrigin: string = getDefaultTargetOrigin(),
+  ) {
+    if (targetOrigin === this._targetOrigin) return;
+
+    if (this._providers.size === 0) {
+      this._targetOrigin = targetOrigin;
+      return;
+    }
+
+    if (targetOrigin === "*") return;
 
     if (this._targetOrigin === "*") {
       this._targetOrigin = targetOrigin;
@@ -141,7 +150,7 @@ export class AssistantFrameProvider {
     this.sendMessage(event, {
       type: "tool-result",
       id: message.id,
-      ...(error ? { error } : { result }),
+      ...(error !== undefined ? { error } : { result }),
     });
   }
 
@@ -160,7 +169,7 @@ export class AssistantFrameProvider {
   }
 
   private getModelContext(): ModelContext {
-    const contexts = Array.from(this._providers).map((p) =>
+    const contexts = Array.from(new Set(this._providers.values())).map((p) =>
       p.getModelContext(),
     );
 
@@ -191,34 +200,87 @@ export class AssistantFrameProvider {
     }
   }
 
+  private removeProvider(id: symbol, origin: string): Unsubscribe | undefined {
+    this._providers.delete(id);
+    const unsubscribe = this._providerUnsubscribes.get(id);
+    this._providerUnsubscribes.delete(id);
+    if (origin === "*") {
+      this._wildcardRegistrations -= 1;
+      if (
+        this._wildcardRegistrations === 0 &&
+        this._strictRegistrations === 0
+      ) {
+        this._targetOrigin = getDefaultTargetOrigin();
+      }
+    } else {
+      this._strictRegistrations -= 1;
+      if (this._strictRegistrations === 0) {
+        this._targetOrigin =
+          this._wildcardRegistrations > 0 ? "*" : getDefaultTargetOrigin();
+      }
+    }
+    return unsubscribe;
+  }
+
   static addModelContextProvider(
     provider: ModelContextProvider,
     targetOrigin?: string,
   ): Unsubscribe {
-    const origin = targetOrigin ?? "*";
+    const origin = targetOrigin ?? getDefaultTargetOrigin();
     const instance = AssistantFrameProvider.getInstance(origin);
-    instance._providers.add(provider);
-    if (origin !== "*") instance._strictRegistrations += 1;
-
-    const unsubscribe = provider.subscribe?.(() => instance.broadcastUpdate());
-    if (unsubscribe) {
-      instance._providerUnsubscribes.set(provider, unsubscribe);
+    const id = Symbol();
+    instance._providers.set(id, provider);
+    if (origin === "*") {
+      instance._wildcardRegistrations += 1;
+    } else {
+      instance._strictRegistrations += 1;
     }
 
-    instance.broadcastUpdate();
+    try {
+      const unsubscribe = provider.subscribe?.(() =>
+        instance.broadcastUpdate(),
+      );
+      if (unsubscribe) {
+        instance._providerUnsubscribes.set(id, unsubscribe);
+      }
+
+      instance.broadcastUpdate();
+    } catch (error) {
+      const unsubscribe = instance.removeProvider(id, origin);
+      // Rollback failures must not replace the registration error.
+      try {
+        unsubscribe?.();
+      } catch (unsubscribeError) {
+        console.error(unsubscribeError);
+      }
+      try {
+        instance.broadcastUpdate();
+      } catch (broadcastError) {
+        console.error(broadcastError);
+      }
+      throw error;
+    }
 
     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 = "*";
+      const unsubscribe = instance.removeProvider(id, origin);
+      let unsubscribeFailed = false;
+      let unsubscribeError: unknown;
+      try {
+        unsubscribe?.();
+      } catch (error) {
+        unsubscribeFailed = true;
+        unsubscribeError = error;
       }
-      instance.broadcastUpdate();
+      try {
+        instance.broadcastUpdate();
+      } catch (error) {
+        if (!unsubscribeFailed) throw error;
+        console.error(error);
+      }
+      if (unsubscribeFailed) throw unsubscribeError;
     };
   }
 
@@ -227,20 +289,40 @@ export class AssistantFrameProvider {
       const instance = AssistantFrameProvider._instance;
       window.removeEventListener("message", instance.handleMessage);
 
-      instance._providerUnsubscribes.forEach((unsubscribe) => unsubscribe?.());
+      let cleanupFailed = false;
+      let cleanupError: unknown;
+      const runCleanup = (cleanup: () => void) => {
+        try {
+          cleanup();
+        } catch (error) {
+          if (cleanupFailed) {
+            console.error(error);
+          } else {
+            cleanupFailed = true;
+            cleanupError = error;
+          }
+        }
+      };
+
+      instance._providerUnsubscribes.forEach((unsubscribe) => {
+        if (unsubscribe) runCleanup(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",
+        runCleanup(() => {
+          abortController.abort();
+          instance.sendMessage(event, {
+            type: "tool-result",
+            id,
+            error: "AssistantFrameProvider has been disposed",
+          });
         });
       });
       instance._activeToolCalls.clear();
 
       AssistantFrameProvider._instance = null;
+      if (cleanupFailed) throw cleanupError;
     }
   }
 }

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

@@ -33,7 +33,7 @@ type SnapshotCarrierMessage = {
  * `metadata.custom.interactables`, or `undefined` if none. This is the read
  * half of the snapshot channel — integrations use it to surface interactable
  * state to the model (see `unstable_injectInteractableContext` in
- * `@assistant-ui/react-ai-sdk` for the AI SDK implementation).
+ * `@assistant-ui/ai-sdk` for the AI SDK implementation).
  *
  * @deprecated Unstable / Experimental (not actually removed).
  */

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

@@ -128,4 +128,48 @@ describe("ModelContextRegistry", () => {
     expect(() => publishUpdate()).toThrow(error);
     expect(laterSubscriber).toHaveBeenCalledTimes(1);
   });
+
+  it("notifies subscribers exactly once when a provider changes", () => {
+    const registry = new ModelContextRegistry();
+    const callbacks = new Set<() => void>();
+    const subscriber = vi.fn();
+
+    registry.addProvider({
+      getModelContext: () => ({ system: "provider instructions" }),
+      subscribe: (callback) => {
+        callbacks.add(callback);
+        return () => callbacks.delete(callback);
+      },
+    });
+    registry.subscribe(subscriber);
+
+    for (const callback of callbacks) {
+      callback();
+    }
+
+    expect(subscriber).toHaveBeenCalledTimes(1);
+  });
+
+  it("removes a provider's context and subscription through the addProvider handle", () => {
+    const registry = new ModelContextRegistry();
+    const unsubscribe = vi.fn();
+
+    const handle = registry.addProvider({
+      getModelContext: () => ({ system: "provider instructions" }),
+      subscribe: () => unsubscribe,
+    });
+    expect(registry.getModelContext().system).toContain(
+      "provider instructions",
+    );
+
+    const subscriber = vi.fn();
+    registry.subscribe(subscriber);
+    handle.remove();
+
+    expect(unsubscribe).toHaveBeenCalledTimes(1);
+    expect(registry.getModelContext().system ?? "").not.toContain(
+      "provider instructions",
+    );
+    expect(subscriber).toHaveBeenCalledTimes(1);
+  });
 });