@assistant-ui/core 0.3.19 → 0.3.21

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

@@ -2,9 +2,24 @@
  * @vitest-environment jsdom
  */
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import type { Tool } from "assistant-stream";
 import { AssistantFrameProvider } from "./provider";
 import { FRAME_MESSAGE_CHANNEL } from "./types";
 
+const createTool = <TResult>(
+  execute: NonNullable<
+    Extract<
+      Tool<Record<string, unknown>, TResult>,
+      { type: "frontend" }
+    >["execute"]
+  >,
+) =>
+  ({
+    type: "frontend",
+    parameters: { type: "object", properties: {} },
+    execute,
+  }) satisfies Tool<Record<string, unknown>, TResult>;
+
 describe("AssistantFrameProvider", () => {
   let messageHandler: ((event: MessageEvent) => void) | undefined;
   let parentWindow: Window;
@@ -69,6 +84,45 @@ describe("AssistantFrameProvider", () => {
     vi.spyOn(window, "removeEventListener").mockImplementation(() => {});
   });
 
+  it.each([
+    null,
+    {},
+    { type: "tool-call", id: null, toolName: "sensitiveTool", args: {} },
+    { type: "tool-call", id: "tool-call-1", toolName: null, args: {} },
+    { type: "tool-call", id: "tool-call-1", toolName: "sensitiveTool" },
+    { type: "tool-cancel", id: null },
+  ])("ignores malformed frame messages", async (message) => {
+    const execute = vi.fn(async () => "result");
+    AssistantFrameProvider.addModelContextProvider(
+      {
+        getModelContext: () => ({
+          tools: { sensitiveTool: createTool(execute) },
+        }),
+      },
+      "https://parent.example",
+    );
+
+    expect(() =>
+      messageHandler?.(
+        new MessageEvent("message", {
+          data: { channel: FRAME_MESSAGE_CHANNEL, message },
+          origin: "https://parent.example",
+          source: parentWindow,
+        }),
+      ),
+    ).not.toThrow();
+
+    await Promise.resolve();
+
+    expect(execute).not.toHaveBeenCalled();
+    expect(parentWindow.postMessage).not.toHaveBeenCalledWith(
+      expect.objectContaining({
+        message: expect.objectContaining({ type: "tool-result" }),
+      }),
+      expect.anything(),
+    );
+  });
+
   afterEach(() => {
     AssistantFrameProvider.dispose();
     vi.restoreAllMocks();
@@ -80,7 +134,7 @@ describe("AssistantFrameProvider", () => {
       {
         getModelContext: () => ({
           tools: {
-            sensitiveTool: { execute },
+            sensitiveTool: createTool(execute),
           },
         }),
       },
@@ -104,7 +158,7 @@ describe("AssistantFrameProvider", () => {
     const execute = vi.fn(async () => "result");
     AssistantFrameProvider.addModelContextProvider({
       getModelContext: () => ({
-        tools: { sensitiveTool: { execute } },
+        tools: { sensitiveTool: createTool(execute) },
       }),
     });
 
@@ -163,7 +217,7 @@ describe("AssistantFrameProvider", () => {
     AssistantFrameProvider.addModelContextProvider({
       getModelContext: () => ({
         tools: {
-          sensitiveTool: { execute },
+          sensitiveTool: createTool(execute),
         },
       }),
     });
@@ -192,7 +246,7 @@ describe("AssistantFrameProvider", () => {
     });
     AssistantFrameProvider.addModelContextProvider({
       getModelContext: () => ({
-        tools: { sensitiveTool: { execute } },
+        tools: { sensitiveTool: createTool(execute) },
       }),
     });
 
@@ -233,7 +287,7 @@ describe("AssistantFrameProvider", () => {
     );
     AssistantFrameProvider.addModelContextProvider({
       getModelContext: () => ({
-        tools: { sensitiveTool: { execute } },
+        tools: { sensitiveTool: createTool(execute) },
       }),
     });
 
@@ -270,7 +324,9 @@ describe("AssistantFrameProvider", () => {
       },
     );
     AssistantFrameProvider.addModelContextProvider({
-      getModelContext: () => ({ tools: { sensitiveTool: { execute } } }),
+      getModelContext: () => ({
+        tools: { sensitiveTool: createTool(execute) },
+      }),
     });
 
     dispatchToolCall(window.location.origin, parentWindow, "tool-a");
@@ -298,7 +354,9 @@ describe("AssistantFrameProvider", () => {
       },
     );
     AssistantFrameProvider.addModelContextProvider({
-      getModelContext: () => ({ tools: { sensitiveTool: { execute } } }),
+      getModelContext: () => ({
+        tools: { sensitiveTool: createTool(execute) },
+      }),
     });
 
     dispatchToolCall(window.location.origin, parentWindow, "duplicate");
@@ -310,6 +368,28 @@ describe("AssistantFrameProvider", () => {
     expect(signals[1]?.aborted).toBe(false);
   });
 
+  it("withdraws its context from the parent when disposed", () => {
+    AssistantFrameProvider.addModelContextProvider(
+      {
+        getModelContext: () => ({
+          tools: { sensitiveTool: createTool(vi.fn()) },
+        }),
+      },
+      "https://parent.example",
+    );
+    vi.mocked(parentWindow.postMessage).mockClear();
+
+    AssistantFrameProvider.dispose();
+
+    expect(parentWindow.postMessage).toHaveBeenCalledExactlyOnceWith(
+      {
+        channel: FRAME_MESSAGE_CHANNEL,
+        message: { type: "model-context-update", context: {} },
+      },
+      "https://parent.example",
+    );
+  });
+
   it("aborts in-flight tool calls when the provider is disposed", async () => {
     let toolSignal: AbortSignal | undefined;
     const execute = vi.fn(
@@ -325,7 +405,9 @@ describe("AssistantFrameProvider", () => {
       },
     );
     AssistantFrameProvider.addModelContextProvider({
-      getModelContext: () => ({ tools: { sensitiveTool: { execute } } }),
+      getModelContext: () => ({
+        tools: { sensitiveTool: createTool(execute) },
+      }),
     });
 
     dispatchToolCall(window.location.origin);
@@ -371,7 +453,9 @@ describe("AssistantFrameProvider", () => {
       },
     );
     const removeProvider = AssistantFrameProvider.addModelContextProvider({
-      getModelContext: () => ({ tools: { sensitiveTool: { execute } } }),
+      getModelContext: () => ({
+        tools: { sensitiveTool: createTool(execute) },
+      }),
     });
 
     dispatchToolCall(window.location.origin);
@@ -417,7 +501,9 @@ describe("AssistantFrameProvider", () => {
       },
     );
     const provider = {
-      getModelContext: () => ({ tools: { sensitiveTool: { execute } } }),
+      getModelContext: () => ({
+        tools: { sensitiveTool: createTool(execute) },
+      }),
     };
     const removeFirst =
       AssistantFrameProvider.addModelContextProvider(provider);
@@ -438,7 +524,7 @@ describe("AssistantFrameProvider", () => {
     const shadowedExecute = vi.fn(async () => "shadowed");
     const removeShadowed = AssistantFrameProvider.addModelContextProvider({
       getModelContext: () => ({
-        tools: { sensitiveTool: { execute: shadowedExecute } },
+        tools: { sensitiveTool: createTool(shadowedExecute) },
       }),
     });
 
@@ -457,7 +543,7 @@ describe("AssistantFrameProvider", () => {
     );
     const removeOwner = AssistantFrameProvider.addModelContextProvider({
       getModelContext: () => ({
-        tools: { sensitiveTool: { execute } },
+        tools: { sensitiveTool: createTool(execute) },
       }),
     });
 
@@ -490,10 +576,10 @@ describe("AssistantFrameProvider", () => {
       },
     );
     const removeFirst = AssistantFrameProvider.addModelContextProvider({
-      getModelContext: () => ({ tools: { firstTool: { execute } } }),
+      getModelContext: () => ({ tools: { firstTool: createTool(execute) } }),
     });
     const removeSecond = AssistantFrameProvider.addModelContextProvider({
-      getModelContext: () => ({ tools: { secondTool: { execute } } }),
+      getModelContext: () => ({ tools: { secondTool: createTool(execute) } }),
     });
 
     dispatchToolCall(
@@ -540,7 +626,7 @@ describe("AssistantFrameProvider", () => {
       {
         getModelContext: () => ({
           tools: {
-            sensitiveTool: { execute },
+            sensitiveTool: createTool(execute),
           },
         }),
       },
@@ -567,7 +653,7 @@ describe("AssistantFrameProvider", () => {
       {
         getModelContext: () => ({
           tools: {
-            sensitiveTool: { execute },
+            sensitiveTool: createTool(execute),
           },
         }),
       },
@@ -609,7 +695,7 @@ describe("AssistantFrameProvider", () => {
       AssistantFrameProvider.addModelContextProvider(
         {
           getModelContext: () => ({
-            tools: { sensitiveTool: { execute } },
+            tools: { sensitiveTool: createTool(execute) },
           }),
           subscribe: () => {
             throw new Error("subscribe failed");
@@ -648,7 +734,7 @@ describe("AssistantFrameProvider", () => {
     expect(() =>
       AssistantFrameProvider.addModelContextProvider({
         getModelContext: () => ({
-          tools: { sensitiveTool: { execute } },
+          tools: { sensitiveTool: createTool(execute) },
         }),
         subscribe: () => {
           dispatchToolCall(window.location.origin);
@@ -678,7 +764,7 @@ describe("AssistantFrameProvider", () => {
     let subscriptionCount = 0;
     const provider = {
       getModelContext: () => ({
-        tools: { sensitiveTool: { execute } },
+        tools: { sensitiveTool: createTool(execute) },
       }),
       subscribe: () => {
         subscriptionCount += 1;
@@ -934,7 +1020,7 @@ describe("AssistantFrameProvider", () => {
       {
         getModelContext: () => ({
           tools: {
-            sensitiveTool: { execute },
+            sensitiveTool: createTool(execute),
           },
         }),
       },

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

@@ -7,6 +7,7 @@ import {
   type SerializedModelContext,
   type SerializedTool,
 } from "./types";
+import { isFrameMessage } from "./validate";
 
 const serializeTool = (tool: Tool<any, any>): SerializedTool => ({
   ...(tool.description && { description: tool.description }),
@@ -100,7 +101,8 @@ export class AssistantFrameProvider {
     if (event.source !== window.parent) return;
     if (event.data?.channel !== FRAME_MESSAGE_CHANNEL) return;
 
-    const message = event.data.message as FrameMessage;
+    const message = event.data.message;
+    if (!isFrameMessage(message)) return;
 
     switch (message.type) {
       case "model-context-request":
@@ -269,6 +271,10 @@ export class AssistantFrameProvider {
 
   private broadcastUpdate() {
     if (this._disposed) return;
+    this.postModelContext();
+  }
+
+  private postModelContext() {
     if (window.parent && window.parent !== window) {
       const updateMessage: FrameMessage = {
         type: "model-context-update",
@@ -427,6 +433,7 @@ export class AssistantFrameProvider {
       });
       instance._providerUnsubscribes.clear();
       instance._providers.clear();
+      runCleanup(() => instance.postModelContext());
       instance._activeToolCalls.forEach(({ abortController, event }, id) => {
         runCleanup(() => {
           abortController.abort();

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

@@ -0,0 +1,57 @@
+import type {
+  FrameMessage,
+  SerializedModelContext,
+  SerializedTool,
+} from "./types";
+
+const isRecord = (value: unknown): value is Record<string, unknown> =>
+  typeof value === "object" && value !== null && !Array.isArray(value);
+
+const hasOptionalPropertyType = (
+  value: Record<string, unknown>,
+  property: string,
+  type: "boolean" | "string",
+) => value[property] === undefined || typeof value[property] === type;
+
+const isSerializedTool = (value: unknown): value is SerializedTool =>
+  isRecord(value) &&
+  Object.hasOwn(value, "parameters") &&
+  hasOptionalPropertyType(value, "description", "string") &&
+  hasOptionalPropertyType(value, "disabled", "boolean") &&
+  hasOptionalPropertyType(value, "type", "string");
+
+const isSerializedModelContext = (
+  value: unknown,
+): value is SerializedModelContext => {
+  if (!isRecord(value)) return false;
+  if (!hasOptionalPropertyType(value, "system", "string")) return false;
+  if (value.tools === undefined) return true;
+  if (!isRecord(value.tools)) return false;
+  return Object.values(value.tools).every(isSerializedTool);
+};
+
+export const isFrameMessage = (value: unknown): value is FrameMessage => {
+  if (!isRecord(value) || typeof value.type !== "string") return false;
+
+  switch (value.type) {
+    case "model-context-request":
+      return true;
+    case "model-context-update":
+      return isSerializedModelContext(value.context);
+    case "tool-call":
+      return (
+        typeof value.id === "string" &&
+        typeof value.toolName === "string" &&
+        Object.hasOwn(value, "args")
+      );
+    case "tool-cancel":
+      return typeof value.id === "string";
+    case "tool-result":
+      return (
+        typeof value.id === "string" &&
+        (value.error == null || typeof value.error === "string")
+      );
+    default:
+      return false;
+  }
+};

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

@@ -195,6 +195,88 @@ describe("shallowMergeInteractableState", () => {
       title: "final",
     });
   });
+
+  it("keeps the unwritten fields of the record still streaming", () => {
+    const prev = { title: "old", settings: { name: "n", size: 1 } };
+
+    expect(
+      shallowMergeInteractableState(
+        prev,
+        { settings: {} },
+        { partialPath: ["settings"] },
+      ),
+    ).toEqual({ title: "old", settings: { name: "n", size: 1 } });
+    expect(
+      shallowMergeInteractableState(
+        prev,
+        { settings: { name: "m" } },
+        { partialPath: ["settings", "name"] },
+      ),
+    ).toEqual({ title: "old", settings: { name: "m", size: 1 } });
+  });
+
+  it("replaces a nested record once the parser has closed it", () => {
+    expect(
+      shallowMergeInteractableState(
+        { title: "old", settings: { name: "n", size: 1, stale: true } },
+        { settings: { name: "m", size: 2 }, title: "ne" },
+        { partialPath: ["title"] },
+      ),
+    ).toEqual({ title: "ne", settings: { name: "m", size: 2 } });
+    expect(
+      shallowMergeInteractableState(
+        { settings: { name: "n", size: 1, stale: true } },
+        { settings: { name: "m", size: 2 } },
+      ),
+    ).toEqual({ settings: { name: "m", size: 2 } });
+  });
+
+  it("overlays along the whole partial path", () => {
+    expect(
+      shallowMergeInteractableState(
+        {
+          settings: {
+            name: "n",
+            theme: { fg: "black", bg: "white" },
+            tags: ["a"],
+          },
+        },
+        { settings: { theme: { fg: "r" }, tags: ["b"] } },
+        { partialPath: ["settings", "theme", "fg"] },
+      ),
+    ).toEqual({
+      settings: { name: "n", theme: { fg: "r", bg: "white" }, tags: ["b"] },
+    });
+  });
+
+  it("overlays a streaming array item patch onto the item", () => {
+    const prev = {
+      tasks: [
+        { id: "a", title: "A", meta: { owner: "me", due: "soon" } },
+        { id: "b", title: "B", meta: { owner: "you", due: "later" } },
+      ],
+    };
+
+    expect(
+      shallowMergeInteractableState(
+        prev,
+        {
+          tasks: {
+            update: [
+              { id: "b", meta: { owner: "them" } },
+              { id: "a", meta: { owner: "u" } },
+            ],
+          },
+        },
+        { partialPath: ["tasks", "update", "1", "meta", "owner"] },
+      ),
+    ).toEqual({
+      tasks: [
+        { id: "a", title: "A", meta: { owner: "u", due: "soon" } },
+        { id: "b", title: "B", meta: { owner: "them" } },
+      ],
+    });
+  });
 });
 
 describe("findModelKnownState", () => {

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

@@ -75,10 +75,31 @@ const getArrayItemId = (item: unknown): string | number | undefined => {
   return typeof id === "string" || typeof id === "number" ? id : undefined;
 };
 
+/**
+ * Overlays a streamed `value` onto `prev` along the parser's `partialPath`:
+ * the record still being written keeps the fields of `prev` it has not
+ * reached yet, while every field the parser has closed replaces its previous
+ * value, as the final merge will.
+ */
+export function overlayPartialPath(
+  prev: unknown,
+  value: unknown,
+  partialPath: readonly string[] = [],
+): unknown {
+  if (!isRecord(prev) || !isRecord(value)) return value;
+  const next: Record<string, unknown> = { ...prev, ...value };
+  const [key, ...rest] = partialPath;
+  if (key !== undefined && Object.hasOwn(value, key)) {
+    next[key] = overlayPartialPath(prev[key], value[key], rest);
+  }
+  return next;
+}
+
 function applyArrayUpdate(
   prev: unknown[],
   update: Record<string, unknown>,
   mintId?: () => string | undefined,
+  partialPath?: readonly string[],
 ) {
   let next = Array.isArray(update.set) ? [...update.set] : [...prev];
 
@@ -102,11 +123,18 @@ function applyArrayUpdate(
       }
     }
 
+    const streaming =
+      partialPath?.[0] === "update"
+        ? { patch: patches[Number(partialPath[1])], rest: partialPath.slice(2) }
+        : undefined;
     next = next.map((item) => {
       const id = getArrayItemId(item);
       if (id === undefined || !isRecord(item)) return item;
       const patch = patchesById.get(id);
-      return patch ? { ...item, ...patch } : item;
+      if (!patch) return item;
+      return streaming && patch === streaming.patch
+        ? overlayPartialPath(item, patch, streaming.rest)
+        : { ...item, ...patch };
     });
   }
 
@@ -142,6 +170,13 @@ export function shallowMergeInteractableState(
          */
         idFactory?: (field: string) => string | undefined;
         idKeyedFields?: ReadonlySet<string>;
+        /**
+         * The path of the value the argument parser reports as still
+         * streaming. The record there is overlaid onto its previous value
+         * so its unwritten fields hold until the model writes them; omitted
+         * once the arguments are complete.
+         */
+        partialPath?: readonly string[] | undefined;
       }
     | undefined,
 ): unknown {
@@ -149,16 +184,24 @@ export function shallowMergeInteractableState(
   const baseline = isRecord(options?.arrayBaseline)
     ? options.arrayBaseline
     : prev;
+  const partialPath = options?.partialPath;
   const next = Object.entries(prev);
   for (const [key, value] of Object.entries(partial)) {
     const baseValue = baseline[key];
+    const streamingPath =
+      partialPath?.[0] === key ? partialPath.slice(1) : undefined;
     if (Array.isArray(baseValue) && isRecord(value)) {
       const mintId =
         options?.idFactory &&
         (options.idKeyedFields === undefined || options.idKeyedFields.has(key))
           ? () => options.idFactory?.(key)
           : undefined;
-      next.push([key, applyArrayUpdate(baseValue, value, mintId)]);
+      next.push([
+        key,
+        applyArrayUpdate(baseValue, value, mintId, streamingPath),
+      ]);
+    } else if (streamingPath) {
+      next.push([key, overlayPartialPath(prev[key], value, streamingPath)]);
     } else {
       next.push([key, value]);
     }

package/src/react/RuntimeAdapter.test.ts

@@ -5,13 +5,16 @@ import { RuntimeAdapter } from "../store";
 import { ExternalStoreRuntimeCore } from "../runtimes/internal";
 import { AssistantRuntimeImpl } from "../runtime/internal";
 import type { ExternalStoreAdapter } from "../runtimes/external-store/external-store-adapter";
+import type { AppendMessage } from "../types/message";
 
 type DemoMessage = { role: "user" | "assistant"; text: string };
 
 const createRuntime = (
   overrides?: Partial<ExternalStoreAdapter<DemoMessage>>,
 ) => {
-  const onNew = vi.fn(async () => {});
+  const onNew = vi.fn<(message: AppendMessage) => Promise<void>>(
+    async () => {},
+  );
   const adapter: ExternalStoreAdapter<DemoMessage> = {
     messages: [{ role: "user", text: "hi" }],
     convertMessage: (message) => ({