@assistant-ui/core 0.3.17 → 0.3.19

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

@@ -120,6 +120,42 @@ describe("AssistantFrameProvider", () => {
     await vi.waitFor(() => expect(execute).toHaveBeenCalledOnce());
   });
 
+  it("does not broadcast from a disposed provider", async () => {
+    vi.useFakeTimers();
+
+    try {
+      AssistantFrameProvider.addModelContextProvider({
+        getModelContext: () => ({ system: "disposed context" }),
+      });
+      AssistantFrameProvider.dispose();
+
+      AssistantFrameProvider.addModelContextProvider({
+        getModelContext: () => ({ system: "current context" }),
+      });
+      vi.mocked(parentWindow.postMessage).mockClear();
+
+      await vi.runAllTimersAsync();
+
+      expect(parentWindow.postMessage).toHaveBeenCalledOnce();
+      expect(parentWindow.postMessage).toHaveBeenCalledWith(
+        {
+          channel: FRAME_MESSAGE_CHANNEL,
+          message: {
+            type: "model-context-update",
+            context: {
+              system: "current context",
+              tools: {},
+            },
+          },
+        },
+        window.location.origin,
+      );
+    } finally {
+      AssistantFrameProvider.dispose();
+      vi.useRealTimers();
+    }
+  });
+
   it("reports a failure even when the thrown error has an empty message", async () => {
     const execute = vi.fn(async () => {
       throw new Error();
@@ -146,6 +182,41 @@ describe("AssistantFrameProvider", () => {
     });
   });
 
+  it("reports tool results that cannot cross the frame boundary", async () => {
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => undefined);
+    const execute = vi.fn(async () => () => undefined);
+    vi.mocked(parentWindow.postMessage).mockImplementation((data) => {
+      structuredClone(data);
+    });
+    AssistantFrameProvider.addModelContextProvider({
+      getModelContext: () => ({
+        tools: { sensitiveTool: { execute } },
+      }),
+    });
+
+    dispatchToolCall(window.location.origin);
+
+    await vi.waitFor(() => {
+      expect(consoleError).toHaveBeenCalledWith(
+        "[assistant-ui] AssistantFrame tool result could not be sent.",
+        expect.objectContaining({ name: "DataCloneError" }),
+      );
+      expect(parentWindow.postMessage).toHaveBeenCalledWith(
+        {
+          channel: FRAME_MESSAGE_CHANNEL,
+          message: {
+            type: "tool-result",
+            id: "tool-call-1",
+            error: "Tool result could not be sent across the frame boundary",
+          },
+        },
+        { targetOrigin: window.location.origin },
+      );
+    });
+  });
+
   it("aborts in-flight tool calls when the parent cancels them", async () => {
     let toolSignal: AbortSignal | undefined;
     const execute = vi.fn(

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

@@ -47,13 +47,18 @@ export class AssistantFrameProvider {
   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 {
@@ -106,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":
@@ -156,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) {
@@ -244,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",
@@ -375,6 +400,11 @@ export class AssistantFrameProvider {
   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

@@ -104,6 +104,18 @@ describe("interactableToolName", () => {
 });
 
 describe("shallowMergeInteractableState", () => {
+  it("adds prototype-named fields as own properties", () => {
+    const value = { enabled: true };
+    const result = shallowMergeInteractableState(
+      { title: "Example" },
+      Object.fromEntries([["__proto__", value]]),
+    ) as Record<string, unknown>;
+
+    expect(Object.getPrototypeOf(result)).toBe(Object.prototype);
+    expect(Object.hasOwn(result, "__proto__")).toBe(true);
+    expect(result["__proto__"]).toBe(value);
+  });
+
   it("applies array operations from a baseline", () => {
     const prev = {
       tasks: [
@@ -131,6 +143,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(
@@ -425,6 +465,24 @@ describe("gateInteractableComposerMetadata", () => {
     expect(gated?.interactables).toEqual([entry("a", { v: 1 })]);
   });
 
+  it.each(["__proto__", "toString"])(
+    "treats a removed prototype-named field %s as a full snapshot",
+    (field) => {
+      const known = Object.fromEntries([
+        [field, { enabled: true }],
+        ["title", "draft"],
+        ["stable", true],
+      ]);
+      const current = { title: "edited", stable: true };
+      const meta = { interactables: [entry("a", current)] };
+      const history = [userMsg([entry("a", known)])];
+
+      const gated = gateInteractableComposerMetadata(meta, history);
+
+      expect(gated?.interactables).toEqual([entry("a", current)]);
+    },
+  );
+
   it("omits an interactable the model already knows via its own update_* call", () => {
     const meta = { interactables: [entry("a", { v: 2 }, "note")] };
     const history = [

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;
     });
   }
@@ -143,7 +149,7 @@ export function shallowMergeInteractableState(
   const baseline = isRecord(options?.arrayBaseline)
     ? options.arrayBaseline
     : prev;
-  const next = { ...prev };
+  const next = Object.entries(prev);
   for (const [key, value] of Object.entries(partial)) {
     const baseValue = baseline[key];
     if (Array.isArray(baseValue) && isRecord(value)) {
@@ -152,12 +158,12 @@ export function shallowMergeInteractableState(
         (options.idKeyedFields === undefined || options.idKeyedFields.has(key))
           ? () => options.idFactory?.(key)
           : undefined;
-      next[key] = applyArrayUpdate(baseValue, value, mintId);
+      next.push([key, applyArrayUpdate(baseValue, value, mintId)]);
     } else {
-      next[key] = value;
+      next.push([key, value]);
     }
   }
-  return next;
+  return Object.fromEntries(next);
 }
 
 /**
@@ -173,17 +179,17 @@ function shallowDiffInteractableState(
 ): Record<string, unknown> | undefined {
   if (!isRecord(known) || !isRecord(next)) return undefined;
   for (const key of Object.keys(known)) {
-    if (!(key in next)) return undefined;
+    if (!Object.hasOwn(next, key)) return undefined;
   }
-  const diff: Record<string, unknown> = {};
+  const diff: [string, unknown][] = [];
   for (const [key, value] of Object.entries(next)) {
-    if (!(key in known) || !isJSONValueEqual(known[key], value)) {
-      diff[key] = value;
+    if (!Object.hasOwn(known, key) || !isJSONValueEqual(known[key], value)) {
+      diff.push([key, value]);
     }
   }
-  const changed = Object.keys(diff).length;
+  const changed = diff.length;
   if (changed === 0 || changed === Object.keys(next).length) return undefined;
-  return diff;
+  return Object.fromEntries(diff);
 }
 
 type ToolCallLikePart = {

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/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/RuntimeAdapter.test.ts

@@ -69,6 +69,40 @@ describe("RuntimeAdapter via the neutral store entry", () => {
     handle.destroy();
   });
 
+  it("emits composer.send for a thread.append, flagging a suggestion", async () => {
+    const { runtime, onNew } = createRuntime({
+      suggestions: [{ prompt: "hi there" }],
+    });
+    const handle = createAssistantClient(
+      AuiConfig({ threads: RuntimeAdapter(runtime) }),
+    );
+    handle.subscribe(() => {});
+    const aui = handle.getClient();
+    const sent = vi.fn();
+    aui.on({ scope: "thread", event: "composer.send" }, sent);
+
+    flushTapSync(() =>
+      aui.thread.append({ content: [{ type: "text", text: "hi there" }] }),
+    );
+    flushTapSync(() => aui.thread.append("plain text"));
+    await new Promise((resolve) => setTimeout(resolve));
+
+    expect(onNew).toHaveBeenCalledTimes(2);
+    expect(sent).toHaveBeenCalledTimes(2);
+    expect(sent.mock.calls[0]![0]).toMatchObject({
+      chars: 8,
+      attachments: 0,
+      suggestion: true,
+    });
+    expect(sent.mock.calls[1]![0]).toMatchObject({
+      chars: 10,
+      attachments: 0,
+    });
+    expect(sent.mock.calls[1]![0]).not.toHaveProperty("suggestion");
+
+    handle.destroy();
+  });
+
   it("registers the client's model context on the runtime", () => {
     const { runtime } = createRuntime();
     const register = vi.spyOn(runtime, "registerModelContextProvider");