@assistant-ui/core 0.3.14 → 0.3.16

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

@@ -29,17 +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);
@@ -52,10 +57,34 @@ export class AssistantFrameProvider {
       AssistantFrameProvider._instance = new AssistantFrameProvider(
         targetOrigin,
       );
+    } else {
+      AssistantFrameProvider._instance.reconcileTargetOrigin(targetOrigin);
     }
     return AssistantFrameProvider._instance;
   }
 
+  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;
+      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 +104,10 @@ export class AssistantFrameProvider {
       case "tool-call":
         this.handleToolCall(message, event);
         break;
+
+      case "tool-cancel":
+        this.cancelToolCall(message.id);
+        break;
     }
   }
 
@@ -83,6 +116,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 +131,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,13 +144,23 @@ 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,
-      ...(error ? { error } : { result }),
+      ...(error !== undefined ? { error } : { result }),
     });
   }
 
+  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 },
@@ -122,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(),
     );
 
@@ -153,25 +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 instance = AssistantFrameProvider.getInstance(targetOrigin);
-    instance._providers.add(provider);
-
-    const unsubscribe = provider.subscribe?.(() => instance.broadcastUpdate());
-    if (unsubscribe) {
-      instance._providerUnsubscribes.set(provider, unsubscribe);
+    const origin = targetOrigin ?? getDefaultTargetOrigin();
+    const instance = AssistantFrameProvider.getInstance(origin);
+    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);
+      }
 
-    return () => {
-      instance._providers.delete(provider);
-      instance._providerUnsubscribes.get(provider)?.();
-      instance._providerUnsubscribes.delete(provider);
       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;
+      const unsubscribe = instance.removeProvider(id, origin);
+      let unsubscribeFailed = false;
+      let unsubscribeError: unknown;
+      try {
+        unsubscribe?.();
+      } catch (error) {
+        unsubscribeFailed = true;
+        unsubscribeError = error;
+      }
+      try {
+        instance.broadcastUpdate();
+      } catch (error) {
+        if (!unsubscribeFailed) throw error;
+        console.error(error);
+      }
+      if (unsubscribeFailed) throw unsubscribeError;
     };
   }
 
@@ -180,11 +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) => {
+        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/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;

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);
+  });
 });

package/src/model-context/registry.ts

@@ -2,12 +2,12 @@ import type { Tool } from "assistant-stream";
 import {
   type ModelContext,
   type ModelContextProvider,
-  mergeModelContexts,
   type AssistantToolProps,
   type AssistantInstructionsConfig,
 } from "./types";
-import type { Unsubscribe } from "../types/unsubscribe";
 import { notifySubscribers as notifyStateSubscribers } from "../subscribable/subscribable";
+import { CompositeContextProvider } from "../utils/composite-context-provider";
+import type { Unsubscribe } from "../types/unsubscribe";
 import type {
   ModelContextRegistryToolHandle,
   ModelContextRegistryInstructionHandle,
@@ -17,9 +17,14 @@ import type {
 export class ModelContextRegistry implements ModelContextProvider {
   private _tools = new Map<symbol, AssistantToolProps<any, any>>();
   private _instructions = new Map<symbol, string>();
-  private _providers = new Map<symbol, ModelContextProvider>();
+  private _contextProviders = new CompositeContextProvider();
   private _subscribers = new Set<() => void>();
-  private _providerUnsubscribes = new Map<symbol, Unsubscribe | undefined>();
+
+  constructor() {
+    this._contextProviders.subscribe(() => {
+      this.notifySubscribers();
+    });
+  }
 
   getModelContext(): ModelContext {
     const instructions = Array.from(this._instructions.values()).filter(
@@ -35,9 +40,7 @@ export class ModelContextRegistry implements ModelContextProvider {
       tools[toolName] = tool;
     }
 
-    const providerContexts = mergeModelContexts(
-      new Set(this._providers.values()),
-    );
+    const providerContexts = this._contextProviders.getModelContext();
 
     const context: ModelContext = {
       system,
@@ -139,36 +142,11 @@ export class ModelContextRegistry implements ModelContextProvider {
   addProvider(
     provider: ModelContextProvider,
   ): ModelContextRegistryProviderHandle {
-    const id = Symbol();
-
-    this._providers.set(id, provider);
-
-    let unsubscribe: Unsubscribe | undefined;
-    try {
-      unsubscribe = provider.subscribe?.(() => {
-        this.notifySubscribers();
-      });
-    } catch (error) {
-      this._providers.delete(id);
-      try {
-        this.notifySubscribers();
-      } catch (notifyError) {
-        console.error(notifyError);
-      }
-      throw error;
-    }
-    this._providerUnsubscribes.set(id, unsubscribe);
-
-    this.notifySubscribers();
+    const unregister =
+      this._contextProviders.registerModelContextProvider(provider);
 
     return {
-      remove: () => {
-        this._providers.delete(id);
-        const unsubscribe = this._providerUnsubscribes.get(id);
-        unsubscribe?.();
-        this._providerUnsubscribes.delete(id);
-        this.notifySubscribers();
-      },
+      remove: unregister,
     };
   }
 }

package/src/react/client/InMemoryThreadList.test.tsx

@@ -146,3 +146,52 @@ describe("InMemoryThreadList delete", () => {
     expect(state.mainThreadId).not.toBe("main");
   });
 });
+
+describe("InMemoryThreadList item index selectors", () => {
+  it("resolves index selectors within the archived and regular subsets", async () => {
+    const { getAui } = setup();
+    await act(async () => {});
+
+    await act(async () => {
+      getAui().threads.switchToNewThread();
+    });
+    await act(async () => {});
+    const b = getAui().threads.getState().mainThreadId;
+
+    await act(async () => {
+      getAui().threads.switchToNewThread();
+    });
+    await act(async () => {});
+
+    await act(async () => {
+      getAui().threads.item({ id: b }).archive();
+    });
+    await act(async () => {});
+
+    const state = getAui().threads.getState();
+    expect(state.archivedThreadIds).toEqual([b]);
+    expect(state.threadIds).toHaveLength(2);
+
+    for (const [index, id] of state.archivedThreadIds.entries()) {
+      expect(
+        getAui().threads.item({ index, archived: true }).getState().id,
+      ).toBe(id);
+    }
+    for (const [index, id] of state.threadIds.entries()) {
+      expect(getAui().threads.item({ index }).getState().id).toBe(id);
+      expect(
+        getAui().threads.item({ index, archived: false }).getState().id,
+      ).toBe(id);
+    }
+
+    expect(() =>
+      getAui().threads.item({ index: state.threadIds.length }),
+    ).toThrow("out of bounds");
+    expect(() =>
+      getAui().threads.item({
+        index: state.archivedThreadIds.length,
+        archived: true,
+      }),
+    ).toThrow("out of bounds");
+  });
+});

package/src/react/client/InMemoryThreadList.ts

@@ -236,7 +236,13 @@ const useInMemoryThreadList = (
         const index = threads.findIndex((t) => t.id === selector.id);
         return threadListItems.get({ index });
       }
-      return threadListItems.get(selector);
+      // The lookup is keyed over the combined thread array, while index
+      // selectors address the regular/archived subset the primitives render.
+      const ids = selector.archived ? state.archivedThreadIds : state.threadIds;
+      const id = ids[selector.index];
+      if (id === undefined) return threadListItems.get({ index: -1 });
+      const index = threads.findIndex((t) => t.id === id);
+      return threadListItems.get({ index });
     },
     thread: () => mainThreadClient.methods,
   };