@assistant-ui/core 0.3.16 → 0.3.18

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

@@ -38,18 +38,27 @@ export class AssistantFrameProvider {
   private _providerUnsubscribes = new Map<symbol, Unsubscribe | undefined>();
   private _activeToolCalls = new Map<
     string,
-    { abortController: AbortController; event: MessageEvent }
+    {
+      abortController: AbortController;
+      event: MessageEvent;
+      provider: ModelContextProvider | undefined;
+    }
   >();
   private _targetOrigin: string;
   private _strictRegistrations = 0;
   private _wildcardRegistrations = 0;
+  private _startupTimer: ReturnType<typeof setTimeout> | undefined;
+  private _disposed = false;
 
   private constructor(targetOrigin: string = getDefaultTargetOrigin()) {
     this._targetOrigin = targetOrigin;
     this.handleMessage = this.handleMessage.bind(this);
     window.addEventListener("message", this.handleMessage);
 
-    setTimeout(() => this.broadcastUpdate(), 0);
+    this._startupTimer = setTimeout(() => {
+      this._startupTimer = undefined;
+      this.broadcastUpdate();
+    }, 0);
   }
 
   private static getInstance(targetOrigin?: string): AssistantFrameProvider {
@@ -102,7 +111,12 @@ export class AssistantFrameProvider {
         break;
 
       case "tool-call":
-        this.handleToolCall(message, event);
+        void this.handleToolCall(message, event).catch((error: unknown) => {
+          console.error(
+            "[assistant-ui] AssistantFrame tool call failed.",
+            error,
+          );
+        });
         break;
 
       case "tool-cancel":
@@ -115,10 +129,15 @@ export class AssistantFrameProvider {
     message: Extract<FrameMessage, { type: "tool-call" }>,
     event: MessageEvent,
   ) {
-    const tool = this.getModelContext().tools?.[message.toolName];
+    const resolvedTool = this.getTool(message.toolName);
+    const tool = resolvedTool?.tool;
     const abortController = new AbortController();
     this._activeToolCalls.get(message.id)?.abortController.abort();
-    const activeCall = { abortController, event };
+    const activeCall = {
+      abortController,
+      event,
+      provider: resolvedTool?.provider,
+    };
     this._activeToolCalls.set(message.id, activeCall);
 
     let result: any;
@@ -147,11 +166,25 @@ 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 !== undefined ? { error } : { result }),
-    });
+    try {
+      this.sendMessage(event, {
+        type: "tool-result",
+        id: message.id,
+        ...(error !== undefined ? { error } : { result }),
+      });
+    } catch (sendError) {
+      if (error !== undefined) throw sendError;
+
+      console.error(
+        "[assistant-ui] AssistantFrame tool result could not be sent.",
+        sendError,
+      );
+      this.sendMessage(event, {
+        type: "tool-result",
+        id: message.id,
+        error: "Tool result could not be sent across the frame boundary",
+      });
+    }
   }
 
   private cancelToolCall(id: string) {
@@ -161,6 +194,37 @@ export class AssistantFrameProvider {
     activeCall.abortController.abort();
   }
 
+  private cancelToolCallsForProvider(provider: ModelContextProvider) {
+    const matchingCalls = Array.from(this._activeToolCalls).filter(
+      ([, activeCall]) => activeCall.provider === provider,
+    );
+    for (const [id, activeCall] of matchingCalls) {
+      this._activeToolCalls.delete(id);
+      activeCall.abortController.abort();
+    }
+
+    let sendFailed = false;
+    let sendError: unknown;
+    for (const [id, activeCall] of matchingCalls) {
+      try {
+        this.sendMessage(activeCall.event, {
+          type: "tool-result",
+          id,
+          error: "AssistantFrame tool provider has been removed",
+        });
+      } catch (error) {
+        if (sendFailed) {
+          console.error(error);
+        } else {
+          sendFailed = true;
+          sendError = error;
+        }
+      }
+    }
+
+    if (sendFailed) throw sendError;
+  }
+
   private sendMessage(event: MessageEvent, message: FrameMessage) {
     event.source?.postMessage(
       { channel: FRAME_MESSAGE_CHANNEL, message },
@@ -168,9 +232,26 @@ export class AssistantFrameProvider {
     );
   }
 
+  private getProviders() {
+    return Array.from(new Set(this._providers.values()));
+  }
+
+  private getTool(toolName: string) {
+    let resolved:
+      | { provider: ModelContextProvider; tool: Tool<any, any> }
+      | undefined;
+
+    for (const provider of this.getProviders()) {
+      const tool = provider.getModelContext().tools?.[toolName];
+      if (tool) resolved = { provider, tool };
+    }
+
+    return resolved;
+  }
+
   private getModelContext(): ModelContext {
-    const contexts = Array.from(new Set(this._providers.values())).map((p) =>
-      p.getModelContext(),
+    const contexts = this.getProviders().map((provider) =>
+      provider.getModelContext(),
     );
 
     return contexts.reduce(
@@ -187,6 +268,7 @@ export class AssistantFrameProvider {
   }
 
   private broadcastUpdate() {
+    if (this._disposed) return;
     if (window.parent && window.parent !== window) {
       const updateMessage: FrameMessage = {
         type: "model-context-update",
@@ -200,7 +282,14 @@ export class AssistantFrameProvider {
     }
   }
 
-  private removeProvider(id: symbol, origin: string): Unsubscribe | undefined {
+  private removeProvider(
+    id: symbol,
+    origin: string,
+  ): {
+    unsubscribe: Unsubscribe | undefined;
+    removedProvider: ModelContextProvider | undefined;
+  } {
+    const provider = this._providers.get(id);
     this._providers.delete(id);
     const unsubscribe = this._providerUnsubscribes.get(id);
     this._providerUnsubscribes.delete(id);
@@ -219,7 +308,11 @@ export class AssistantFrameProvider {
           this._wildcardRegistrations > 0 ? "*" : getDefaultTargetOrigin();
       }
     }
-    return unsubscribe;
+    const removedProvider =
+      provider && !this.getProviders().includes(provider)
+        ? provider
+        : undefined;
+    return { unsubscribe, removedProvider };
   }
 
   static addModelContextProvider(
@@ -246,8 +339,18 @@ export class AssistantFrameProvider {
 
       instance.broadcastUpdate();
     } catch (error) {
-      const unsubscribe = instance.removeProvider(id, origin);
+      const { unsubscribe, removedProvider } = instance.removeProvider(
+        id,
+        origin,
+      );
       // Rollback failures must not replace the registration error.
+      try {
+        if (removedProvider) {
+          instance.cancelToolCallsForProvider(removedProvider);
+        }
+      } catch (cancelError) {
+        console.error(cancelError);
+      }
       try {
         unsubscribe?.();
       } catch (unsubscribeError) {
@@ -265,28 +368,43 @@ export class AssistantFrameProvider {
     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);
+      const { unsubscribe, removedProvider } = instance.removeProvider(
+        id,
+        origin,
+      );
+      let cleanupFailed = false;
+      let cleanupError: unknown;
+      const runCleanup = (cleanup: () => void) => {
+        try {
+          cleanup();
+        } catch (error) {
+          if (cleanupFailed) {
+            console.error(error);
+          } else {
+            cleanupFailed = true;
+            cleanupError = error;
+          }
+        }
+      };
+
+      if (removedProvider) {
+        runCleanup(() => instance.cancelToolCallsForProvider(removedProvider));
       }
-      if (unsubscribeFailed) throw unsubscribeError;
+      if (unsubscribe) runCleanup(unsubscribe);
+      runCleanup(() => instance.broadcastUpdate());
+
+      if (cleanupFailed) throw cleanupError;
     };
   }
 
   static dispose() {
     if (AssistantFrameProvider._instance) {
       const instance = AssistantFrameProvider._instance;
+      instance._disposed = true;
+      if (instance._startupTimer !== undefined) {
+        clearTimeout(instance._startupTimer);
+        instance._startupTimer = undefined;
+      }
       window.removeEventListener("message", instance.handleMessage);
 
       let cleanupFailed = false;

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

@@ -131,6 +131,34 @@ describe("shallowMergeInteractableState", () => {
     });
   });
 
+  it("uses the first valid patch for each array item id", () => {
+    expect(
+      shallowMergeInteractableState(
+        { tasks: [{ id: 1, title: "Original", done: false }] },
+        {
+          tasks: {
+            update: [
+              null,
+              { title: "Missing id" },
+              { id: "1", title: "String id" },
+              { id: 1, title: "First", done: true },
+              { id: 1, title: "Second" },
+            ],
+          },
+        },
+      ),
+    ).toEqual({ tasks: [{ id: 1, title: "First", done: true }] });
+  });
+
+  it("does not match NaN array item ids", () => {
+    expect(
+      shallowMergeInteractableState(
+        { tasks: [{ id: Number.NaN, title: "Original" }] },
+        { tasks: { update: [{ id: Number.NaN, title: "Patched" }] } },
+      ),
+    ).toEqual({ tasks: [{ id: Number.NaN, title: "Original" }] });
+  });
+
   it("keeps raw array replacement semantics", () => {
     expect(
       shallowMergeInteractableState(

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

@@ -94,12 +94,18 @@ function applyArrayUpdate(
 
   const patches = update.update;
   if (Array.isArray(patches) && patches.length > 0) {
+    const patchesById = new Map<string | number, Record<string, unknown>>();
+    for (const candidate of patches) {
+      const id = getArrayItemId(candidate);
+      if (id !== undefined && !Number.isNaN(id) && !patchesById.has(id)) {
+        patchesById.set(id, candidate as Record<string, unknown>);
+      }
+    }
+
     next = next.map((item) => {
       const id = getArrayItemId(item);
       if (id === undefined || !isRecord(item)) return item;
-      const patch = patches.find(
-        (candidate) => isRecord(candidate) && candidate.id === id,
-      );
+      const patch = patchesById.get(id);
       return patch ? { ...item, ...patch } : item;
     });
   }

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

@@ -2,6 +2,18 @@ import { describe, expect, it, vi } from "vitest";
 import { ModelContextRegistry } from "./registry";
 
 describe("ModelContextRegistry", () => {
+  it.each(["__proto__", "constructor", "toString"])(
+    "retains a tool named %s",
+    (toolName) => {
+      const registry = new ModelContextRegistry();
+      registry.addTool({ toolName });
+
+      const tools = registry.getModelContext().tools!;
+      expect(Object.hasOwn(tools, toolName)).toBe(true);
+      expect(Object.keys(tools)).toEqual([toolName]);
+    },
+  );
+
   it("notifies every subscriber and rethrows when a registration subscriber throws", () => {
     const registry = new ModelContextRegistry();
     const error = new Error("subscriber failed");

package/src/model-context/registry.ts

@@ -13,6 +13,7 @@ import type {
   ModelContextRegistryInstructionHandle,
   ModelContextRegistryProviderHandle,
 } from "./registry-handles";
+import { nullProtoRecord } from "../utils/record";
 
 export class ModelContextRegistry implements ModelContextProvider {
   private _tools = new Map<symbol, AssistantToolProps<any, any>>();
@@ -34,7 +35,7 @@ export class ModelContextRegistry implements ModelContextProvider {
     const system =
       instructions.length > 0 ? instructions.join("\n\n") : undefined;
 
-    const tools: Record<string, Tool<any, any>> = {};
+    const tools = nullProtoRecord<Tool<any, any>>();
     for (const toolProps of this._tools.values()) {
       const { toolName, render, ...tool } = toolProps;
       tools[toolName] = tool;
@@ -54,7 +55,7 @@ export class ModelContextRegistry implements ModelContextProvider {
     }
 
     if (providerContexts.tools) {
-      context.tools = { ...(context.tools || {}), ...providerContexts.tools };
+      context.tools = nullProtoRecord(context.tools, providerContexts.tools);
     }
 
     if (providerContexts.callSettings) {

package/src/model-context/tool-call-text.test.ts

@@ -0,0 +1,67 @@
+import { describe, expect, it } from "vitest";
+import { resolveToolCallText } from "./tool-call-text";
+
+describe("resolveToolCallText", () => {
+  it("resolves literal text for running and complete tool calls", () => {
+    const text = {
+      running: "Searching...",
+      complete: "Done searching",
+    };
+
+    expect(
+      resolveToolCallText(text, {
+        args: {},
+        status: { type: "running" },
+      }),
+    ).toBe("Searching...");
+    expect(
+      resolveToolCallText(text, {
+        args: {},
+        status: { type: "complete" },
+      }),
+    ).toBe("Done searching");
+  });
+
+  it("passes args and results to text functions", () => {
+    const text = {
+      running: ({ args }: { args: { query: string } }) =>
+        `Searching ${args.query}...`,
+      complete: ({
+        args,
+        result,
+      }: {
+        args: { query: string };
+        result: number | undefined;
+      }) => `Found ${result ?? 0} results for ${args.query}`,
+    };
+
+    expect(
+      resolveToolCallText(text, {
+        args: { query: "docs" },
+        status: { type: "requires-action" },
+      }),
+    ).toBe("Searching docs...");
+    expect(
+      resolveToolCallText(text, {
+        args: { query: "docs" },
+        result: 3,
+        status: { type: "complete" },
+      }),
+    ).toBe("Found 3 results for docs");
+  });
+
+  it("returns null when the active status has no text", () => {
+    expect(
+      resolveToolCallText(
+        { running: "Searching..." },
+        { args: {}, status: { type: "complete" } },
+      ),
+    ).toBeNull();
+    expect(
+      resolveToolCallText(
+        { complete: "Done searching" },
+        { args: {}, status: { type: "running" } },
+      ),
+    ).toBeNull();
+  });
+});

package/src/model-context/tool-call-text.ts

@@ -0,0 +1,67 @@
+type ToolCallRunningText<TArgs extends Record<string, unknown>, TValue> =
+  | TValue
+  | undefined
+  | null
+  | ((options: { args: TArgs }) => TValue | undefined | null);
+
+type ToolCallCompleteText<
+  TArgs extends Record<string, unknown>,
+  TResult,
+  TValue,
+> =
+  | TValue
+  | undefined
+  | null
+  | ((options: {
+      args: TArgs;
+      result: TResult | undefined;
+    }) => TValue | undefined | null);
+
+export type ToolCallText<
+  TArgs extends Record<string, unknown>,
+  TResult,
+  TValue = string,
+> =
+  | {
+      running: ToolCallRunningText<TArgs, TValue>;
+      complete?: ToolCallCompleteText<TArgs, TResult, TValue> | undefined;
+    }
+  | {
+      running?: ToolCallRunningText<TArgs, TValue> | undefined;
+      complete: ToolCallCompleteText<TArgs, TResult, TValue>;
+    };
+
+type ToolCallTextPart<TArgs extends Record<string, unknown>, TResult> = {
+  readonly args: TArgs;
+  readonly result?: TResult | undefined;
+  readonly status?: { readonly type?: string | undefined } | undefined;
+};
+
+export const resolveToolCallText = <
+  TArgs extends Record<string, unknown>,
+  TResult,
+  TValue,
+>(
+  text: ToolCallText<TArgs, TResult, TValue>,
+  part: ToolCallTextPart<TArgs, TResult>,
+): TValue | undefined | null => {
+  const isRunning =
+    part.status?.type === "running" || part.status?.type === "requires-action";
+
+  if (!isRunning) {
+    const value = text.complete;
+    if (typeof value !== "function") return value ?? null;
+    return (
+      value as (options: {
+        args: TArgs;
+        result: TResult | undefined;
+      }) => TValue | undefined | null
+    )({ args: part.args, result: part.result });
+  }
+
+  const value = text.running;
+  if (typeof value !== "function") return value ?? null;
+  return (value as (options: { args: TArgs }) => TValue | undefined | null)({
+    args: part.args,
+  });
+};

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

@@ -0,0 +1,25 @@
+import { describe, expect, it } from "vitest";
+import { mergeModelContexts } from "./types";
+
+const provider = (tools: Record<string, any>, priority?: number) =>
+  ({ getModelContext: () => ({ tools, priority }) }) as any;
+
+describe("mergeModelContexts", () => {
+  it.each(["__proto__", "constructor", "toString"])(
+    "retains and prioritizes a tool named %s",
+    (name) => {
+      const highPriorityTool = { description: "high", parameters: {} };
+      const lowPriorityTool = { description: "low", parameters: {} };
+      const otherTool = { description: "other", parameters: {} };
+      const result = mergeModelContexts(
+        new Set([
+          provider({ [name]: highPriorityTool, ok: otherTool }, 1),
+          provider({ [name]: lowPriorityTool }, 0),
+        ]),
+      );
+
+      expect(Object.keys(result.tools ?? {})).toEqual([name, "ok"]);
+      expect(result.tools?.[name]?.description).toBe("high");
+    },
+  );
+});

package/src/model-context/types.ts

@@ -1,5 +1,6 @@
 import type { Unsubscribe } from "../types/unsubscribe";
 import type { Tool } from "assistant-stream";
+import { nullProtoRecord } from "../utils/record";
 
 export type LanguageModelV1CallSettings = {
   maxTokens?: number;
@@ -68,7 +69,7 @@ export const mergeModelContexts = (
     .map((c) => c.getModelContext())
     .sort((a, b) => (b.priority ?? 0) - (a.priority ?? 0));
 
-  const toolPriorities: Record<string, number> = {};
+  const toolPriorities = nullProtoRecord<number>();
 
   return configs.reduce((acc, config) => {
     const priority = config.priority ?? 0;
@@ -81,7 +82,9 @@ export const mergeModelContexts = (
     }
     if (config.tools) {
       for (const [name, tool] of Object.entries(config.tools)) {
-        const existing = acc.tools?.[name];
+        const hasExisting =
+          acc.tools !== undefined && Object.hasOwn(acc.tools, name);
+        const existing = hasExisting ? acc.tools![name] : undefined;
         if (existing && existing !== tool) {
           const existingPriority = toolPriorities[name]!;
           if (existingPriority === priority) {
@@ -106,9 +109,11 @@ export const mergeModelContexts = (
           continue;
         }
 
-        if (!acc.tools) acc.tools = {};
+        if (!acc.tools) acc.tools = nullProtoRecord();
         acc.tools[name] = stripOverwrite(tool);
-        toolPriorities[name] ??= priority;
+        if (!Object.hasOwn(toolPriorities, name)) {
+          toolPriorities[name] = priority;
+        }
       }
     }
     if (config.config) {

package/src/react/AssistantRuntimeProvider.test.tsx

@@ -1,8 +1,8 @@
 // @vitest-environment jsdom
 
 import type { ReactNode } from "react";
-import { useEffect } from "react";
-import { cleanup, render } from "@testing-library/react";
+import { useEffect, useState } from "react";
+import { act, cleanup, render } from "@testing-library/react";
 import { afterEach, describe, expect, it } from "vitest";
 import { resource } from "@assistant-ui/tap";
 import {
@@ -12,8 +12,10 @@ import {
   type AssistantClient,
 } from "@assistant-ui/store";
 import type { ChatModelAdapter } from "../runtime/utils/chat-model-adapter";
+import type { ThreadMessage } from "../types/message";
 import { AssistantRuntimeProvider } from "./AssistantRuntimeProvider";
 import { useLocalRuntime } from "./runtimes/useLocalRuntime";
+import { useExternalStoreRuntime } from "./runtimes/useExternalStoreRuntime";
 
 type AnyClient = Record<string, any>;
 
@@ -21,6 +23,20 @@ const chatModel: ChatModelAdapter = {
   run: async () => ({ content: [] }),
 };
 
+const hoistedConfig = AuiConfig({});
+
+const message = (id: string, role: "user" | "assistant"): ThreadMessage =>
+  ({
+    id,
+    role,
+    content: [{ type: "text", text: `text of ${id}` }],
+    createdAt: new Date(1718000000000),
+    ...(role === "assistant"
+      ? { status: { type: "complete", reason: "stop" } }
+      : { attachments: [] }),
+    metadata: { custom: {} },
+  }) as ThreadMessage;
+
 const makeCounterClient = (log?: string[]) => {
   const useCounterClient = () => {
     useEffect(() => {
@@ -152,4 +168,52 @@ describe("AssistantRuntimeProvider aui composition", () => {
 
     expect(aui.threads.getState().main).toBeDefined();
   });
+
+  const countPublicationsOverThreeRerenders = async (
+    makeMessages: () => ThreadMessage[],
+  ) => {
+    let advance!: () => void;
+    let notifications = 0;
+
+    const Subscriber = () => {
+      const aui = useAui();
+      useEffect(() => aui.subscribe(() => notifications++), [aui]);
+      return null;
+    };
+
+    const App = () => {
+      const [, setTick] = useState(0);
+      advance = () => setTick((tick) => tick + 1);
+      const runtime = useExternalStoreRuntime<ThreadMessage>({
+        messages: makeMessages(),
+        onNew: async () => {},
+      });
+
+      return (
+        <AssistantRuntimeProvider runtime={runtime} config={hoistedConfig}>
+          <Subscriber />
+        </AssistantRuntimeProvider>
+      );
+    };
+
+    render(<App />);
+    notifications = 0;
+
+    await act(async () => advance());
+    await act(async () => advance());
+    await act(async () => advance());
+
+    return notifications;
+  };
+
+  it("publishes once when the external-store runtime owner rerenders", async () => {
+    expect(await countPublicationsOverThreeRerenders(() => [])).toBe(3);
+  });
+
+  it("publishes once per rerender when the thread already has messages", async () => {
+    const messages = [message("m1", "user"), message("m2", "assistant")];
+    expect(await countPublicationsOverThreeRerenders(() => [...messages])).toBe(
+      3,
+    );
+  });
 });