@assistant-ui/core 0.3.12 → 0.3.14

package/src/tests/session-reset-notification.test.ts

@@ -0,0 +1,145 @@
+import { describe, expect, it, vi } from "vitest";
+import { ExternalStoreRuntimeCore } from "../runtimes/external-store/external-store-runtime-core";
+import { LocalRuntimeCore } from "../runtimes/local/local-runtime-core";
+import type { ChatModelAdapter } from "../runtime/utils/chat-model-adapter";
+import type { ThreadMessage } from "../types/message";
+
+async function waitFor(
+  predicate: () => unknown,
+  timeoutMs = 500,
+): Promise<void> {
+  const start = Date.now();
+  while (Date.now() - start < timeoutMs) {
+    try {
+      await predicate();
+      return;
+    } catch {
+      // retry
+    }
+    await new Promise((r) => setTimeout(r, 5));
+  }
+  await predicate();
+}
+
+const trailingUserMessage = {
+  id: "u1",
+  role: "user",
+  createdAt: new Date(),
+  content: [{ type: "text", text: "draft in flight" }],
+  attachments: [],
+  metadata: { custom: {} },
+} as ThreadMessage;
+
+const toolCallMessage = {
+  id: "a1",
+  role: "assistant",
+  createdAt: new Date(),
+  status: { type: "requires-action", reason: "tool-calls" },
+  metadata: {
+    unstable_state: null,
+    unstable_annotations: [],
+    unstable_data: [],
+    steps: [],
+    custom: {},
+  },
+  content: [
+    {
+      type: "tool-call",
+      toolCallId: "tool-1",
+      toolName: "send_email",
+      args: { to: "dev@example.com" },
+      argsText: '{"to":"dev@example.com"}',
+    },
+  ],
+} as unknown as ThreadMessage;
+
+const externalThread = () => {
+  const setToolStatuses = vi.fn();
+  const setMessages = vi.fn();
+  const onCancel = vi.fn();
+  const notifyCancelled = vi.fn();
+  const store = {
+    messages: [trailingUserMessage] as ThreadMessage[],
+    onNew: vi.fn(async () => {}),
+    onCancel,
+    setMessages,
+    unstable_enableToolInvocations: true,
+    setToolStatuses,
+    queue: {
+      __internal_notifyCancelled: notifyCancelled,
+    } as never,
+  };
+  const core = new ExternalStoreRuntimeCore(store);
+  core.registerModelContextProvider({
+    getModelContext: () => ({
+      tools: {
+        send_email: {
+          parameters: { type: "object", properties: {} },
+          execute: vi.fn(() => new Promise(() => {})),
+        },
+      },
+    }),
+  });
+  return {
+    core,
+    store,
+    thread: core.threads.getMainThreadRuntimeCore(),
+    setToolStatuses,
+    setMessages,
+    onCancel,
+    notifyCancelled,
+  };
+};
+
+const driveExecutingTool = async (
+  harness: ReturnType<typeof externalThread>,
+) => {
+  harness.core.setAdapter({
+    ...harness.store,
+    messages: [trailingUserMessage, toolCallMessage],
+  });
+  await waitFor(() => {
+    const statuses = harness.setToolStatuses.mock.lastCall?.[0] as
+      | Record<string, { type: string }>
+      | undefined;
+    expect(statuses?.["tool-1"]?.type).toBe("executing");
+  });
+};
+
+describe("unstable_notifySessionReset", () => {
+  it("clears session-scoped tool state once and parks queued work", async () => {
+    const harness = externalThread();
+    await driveExecutingTool(harness);
+
+    harness.setToolStatuses.mockClear();
+    harness.thread.unstable_notifySessionReset();
+
+    expect(harness.setToolStatuses).toHaveBeenCalledTimes(1);
+    expect(harness.setToolStatuses).toHaveBeenLastCalledWith({});
+    expect(harness.notifyCancelled).toHaveBeenCalledTimes(1);
+  });
+
+  it("carries no run-cancel semantics", async () => {
+    const harness = externalThread();
+    await driveExecutingTool(harness);
+
+    harness.thread.unstable_notifySessionReset();
+
+    expect(harness.onCancel).not.toHaveBeenCalled();
+    expect(harness.setMessages).not.toHaveBeenCalled();
+    expect(harness.thread.messages.at(-1)?.id).toBe("a1");
+  });
+
+  it("throws on runtimes without a backing session", () => {
+    const local = new LocalRuntimeCore(
+      {
+        adapters: { chatModel: { async *run() {} } satisfies ChatModelAdapter },
+      },
+      undefined,
+    );
+
+    expect(() =>
+      local.threads.getMainThreadRuntimeCore().unstable_notifySessionReset(),
+    ).toThrow("Runtime does not support resetting sessions.");
+  });
+});

package/src/tests/thread-message-like.test.ts

@@ -382,3 +382,42 @@ describe("fromThreadMessageLike", () => {
     });
   });
 });
+
+describe("fromThreadMessageLike provider metadata", () => {
+  it("keeps it on image and file parts", () => {
+    const result = fromThreadMessageLike(
+      {
+        role: "user",
+        content: [
+          {
+            type: "image",
+            image: "https://example.com/cat.png",
+            providerMetadata: { agui: { file_id: "f_1" } },
+          },
+          {
+            type: "file",
+            data: "https://example.com/spec.pdf",
+            mimeType: "application/pdf",
+            providerMetadata: { agui: { file_id: "f_2" } },
+          },
+        ],
+      },
+      fallbackId,
+      fallbackStatus,
+    );
+
+    expect(result.content).toEqual([
+      {
+        type: "image",
+        image: "https://example.com/cat.png",
+        providerMetadata: { agui: { file_id: "f_1" } },
+      },
+      {
+        type: "file",
+        data: "https://example.com/spec.pdf",
+        mimeType: "application/pdf",
+        providerMetadata: { agui: { file_id: "f_2" } },
+      },
+    ]);
+  });
+});

package/src/tests/thread-switch-events.test.tsx

@@ -0,0 +1,234 @@
+// @vitest-environment jsdom
+
+import { act, render, waitFor } from "@testing-library/react";
+import { describe, expect, it, vi } from "vitest";
+import { AuiProvider, useAui, useAuiEvent } from "@assistant-ui/store";
+import { AssistantRuntimeProvider } from "../react/AssistantRuntimeProvider";
+import { useExternalStoreRuntime } from "../react/runtimes/useExternalStoreRuntime";
+import { useRemoteThreadListRuntime } from "../react/runtimes/useRemoteThreadListRuntime";
+import { makeAdapter } from "./remote-thread-list-test-helpers";
+import { RuntimeAdapter } from "../react/RuntimeAdapter";
+import { AssistantRuntimeImpl } from "../runtime/api/assistant-runtime";
+import { ExternalStoreRuntimeCore } from "../runtimes/external-store/external-store-runtime-core";
+import type { ExternalStoreAdapter } from "../runtimes/external-store/external-store-adapter";
+
+type DemoMessage = { id: string; role: "user" | "assistant"; text: string };
+
+const EMPTY_MESSAGES: readonly never[] = [];
+
+const useTestThreadRuntime = () =>
+  useExternalStoreRuntime({
+    messages: EMPTY_MESSAGES,
+    isRunning: false,
+    onNew: async () => {},
+  });
+
+const createRuntime = () => {
+  const threads = [
+    {
+      id: "t1",
+      title: "one",
+      messages: [{ id: "m1", role: "user" as const, text: "a" }],
+    },
+    {
+      id: "t2",
+      title: "two",
+      messages: [{ id: "m2", role: "user" as const, text: "b" }],
+    },
+  ];
+  let currentId = "t1";
+  const makeAdapter = (): ExternalStoreAdapter<DemoMessage> => ({
+    messages: threads.find((t) => t.id === currentId)!.messages,
+    convertMessage: (m) => ({
+      id: m.id,
+      role: m.role,
+      content: [{ type: "text", text: m.text }],
+    }),
+    onNew: async () => {},
+    adapters: {
+      threadList: {
+        threadId: currentId,
+        threads: threads.map((t) => ({
+          status: "regular" as const,
+          id: t.id,
+          title: t.title,
+        })),
+        onSwitchToThread: (threadId: string) => {
+          currentId = threadId;
+          sync();
+        },
+        onSwitchToNewThread: () => {},
+      },
+    },
+  });
+  const core = new ExternalStoreRuntimeCore(makeAdapter());
+  const runtime = new AssistantRuntimeImpl(core);
+  const sync = () => core.setAdapter(makeAdapter());
+  return runtime;
+};
+
+describe("thread switch events", () => {
+  it("delivers switchedTo to default-scope, star-scope, and aui.on listeners", async () => {
+    const runtime = createRuntime();
+    const defaultScope = vi.fn();
+    const starScope = vi.fn();
+    const auiOn = vi.fn();
+    const switchedAwayStar = vi.fn();
+    let aui!: ReturnType<typeof useAui>;
+    const Consumer = () => {
+      useAuiEvent("threadListItem.switchedTo" as never, defaultScope as never);
+      useAuiEvent(
+        { scope: "*", event: "threadListItem.switchedTo" } as never,
+        starScope as never,
+      );
+      useAuiEvent(
+        { scope: "*", event: "threadListItem.switchedAway" } as never,
+        switchedAwayStar as never,
+      );
+      return null;
+    };
+    const Harness = () => {
+      aui = useAui({ threads: RuntimeAdapter(runtime) } as never);
+      return (
+        <AuiProvider value={aui}>
+          <Consumer />
+        </AuiProvider>
+      );
+    };
+    render(<Harness />);
+    await act(async () => {});
+
+    aui.on("threadListItem.switchedTo" as never, auiOn as never);
+
+    await act(async () => {
+      aui.threads.item({ index: 1 }).switchTo();
+    });
+    await act(async () => {});
+
+    expect(defaultScope).toHaveBeenCalledExactlyOnceWith({ threadId: "t2" });
+    expect(starScope).toHaveBeenCalledExactlyOnceWith({ threadId: "t2" });
+    expect(auiOn).toHaveBeenCalledExactlyOnceWith({ threadId: "t2" });
+    expect(switchedAwayStar).toHaveBeenCalledExactlyOnceWith({
+      threadId: "t1",
+    });
+
+    await act(async () => {
+      aui.threads.item({ index: 0 }).switchTo();
+    });
+    await act(async () => {});
+
+    expect(defaultScope).toHaveBeenCalledTimes(2);
+    expect(defaultScope).toHaveBeenLastCalledWith({ threadId: "t1" });
+    expect(switchedAwayStar).toHaveBeenCalledTimes(2);
+    expect(switchedAwayStar).toHaveBeenLastCalledWith({ threadId: "t2" });
+  });
+
+  it("delivers threads.selectionChanged with the new and previous thread ids", async () => {
+    const runtime = createRuntime();
+    const selectionChanged = vi.fn();
+    const auiOn = vi.fn();
+    let aui!: ReturnType<typeof useAui>;
+    const Consumer = () => {
+      useAuiEvent(
+        "threads.selectionChanged" as never,
+        selectionChanged as never,
+      );
+      return null;
+    };
+    const Harness = () => {
+      aui = useAui({ threads: RuntimeAdapter(runtime) } as never);
+      return (
+        <AuiProvider value={aui}>
+          <Consumer />
+        </AuiProvider>
+      );
+    };
+    render(<Harness />);
+    await act(async () => {});
+
+    aui.on("threads.selectionChanged" as never, auiOn as never);
+
+    await act(async () => {
+      aui.threads.item({ index: 1 }).switchTo();
+    });
+    await act(async () => {});
+
+    expect(selectionChanged).toHaveBeenCalledExactlyOnceWith({
+      threadId: "t2",
+      previousThreadId: "t1",
+    });
+    expect(auiOn).toHaveBeenCalledExactlyOnceWith({
+      threadId: "t2",
+      previousThreadId: "t1",
+    });
+
+    await act(async () => {
+      aui.threads.item({ index: 0 }).switchTo();
+    });
+    await act(async () => {});
+
+    expect(selectionChanged).toHaveBeenCalledTimes(2);
+    expect(selectionChanged).toHaveBeenLastCalledWith({
+      threadId: "t1",
+      previousThreadId: "t2",
+    });
+  });
+
+  it("does not emit for the initially selected thread on mount", async () => {
+    const runtime = createRuntime();
+    const anySwitch = vi.fn();
+    const Consumer = () => {
+      useAuiEvent(
+        { scope: "*", event: "threadListItem.switchedTo" } as never,
+        anySwitch as never,
+      );
+      useAuiEvent(
+        { scope: "*", event: "threads.selectionChanged" } as never,
+        anySwitch as never,
+      );
+      return null;
+    };
+    const Harness = () => {
+      const aui = useAui({ threads: RuntimeAdapter(runtime) } as never);
+      return (
+        <AuiProvider value={aui}>
+          <Consumer />
+        </AuiProvider>
+      );
+    };
+    render(<Harness />);
+    await act(async () => {});
+
+    expect(anySwitch).not.toHaveBeenCalled();
+  });
+
+  it("emits when a deep-linked initial thread resolves after mount", async () => {
+    const adapter = makeAdapter();
+    const selectionChanged = vi.fn();
+    const Listener = () => {
+      useAuiEvent("threads.selectionChanged", selectionChanged);
+      return null;
+    };
+    const Harness = () => {
+      const runtime = useRemoteThreadListRuntime({
+        adapter,
+        initialThreadId: "thread-a",
+        runtimeHook: useTestThreadRuntime,
+      });
+      return (
+        <AssistantRuntimeProvider runtime={runtime}>
+          <Listener />
+        </AssistantRuntimeProvider>
+      );
+    };
+    render(<Harness />);
+
+    await waitFor(() => expect(selectionChanged).toHaveBeenCalledTimes(1));
+    const payload = selectionChanged.mock.calls[0]![0] as {
+      threadId: string;
+      previousThreadId: string;
+    };
+    expect(payload.threadId).toBe("thread-a");
+    expect(payload.previousThreadId).toMatch(/^__LOCALID_/);
+  });
+});

package/src/tests/tools-scope-migration.test.tsx

@@ -0,0 +1,58 @@
+// @vitest-environment jsdom
+
+import { act, cleanup, render } from "@testing-library/react";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { withKey } from "@assistant-ui/tap";
+import { useAui } from "@assistant-ui/store";
+import { Tools } from "../react/client/Tools";
+import { ModelContext } from "../store/clients/model-context-client";
+import type { Toolkit } from "../react/model-context/toolbox";
+
+type AnyClient = Record<string, any>;
+
+const toolkit = {
+  lookup: {
+    description: "look something up",
+    parameters: { type: "object", properties: {} },
+    execute: async () => "ok",
+  },
+} as unknown as Toolkit;
+
+afterEach(() => {
+  cleanup();
+});
+
+describe("Tools model-context registration", () => {
+  it("migrates to a structurally replaced modelContext scope", async () => {
+    let aui!: AnyClient;
+    const Harness = ({ generation }: { generation: number }) => {
+      aui = useAui({
+        tools: Tools({ toolkit }),
+        modelContext: withKey(String(generation), ModelContext()),
+      } as never);
+      return null;
+    };
+    const view = render(<Harness generation={0} />);
+
+    expect(
+      Object.keys(aui.modelContext().getModelContext().tools ?? {}),
+    ).toContain("lookup");
+
+    // A rerender inside an outer act would not land until that act exits,
+    // so the swap happens first and only the scheduler wait is wrapped
+    view.rerender(<Harness generation={1} />);
+
+    // The registration lands one scheduler macrotask after the swap: the
+    // mount-time register marks the tap scheduler dirty, which defers the
+    // structural publish to the scheduler's drain. The remounted scope
+    // starts from fresh state, so the tools reappearing proves the
+    // registration migrated to the replacement instance.
+    await act(async () => {
+      await vi.waitFor(() =>
+        expect(
+          Object.keys(aui.modelContext().getModelContext().tools ?? {}),
+        ).toContain("lookup"),
+      );
+    });
+  });
+});

package/src/types/error.ts

@@ -52,6 +52,30 @@ export const isAssistantError = (value: unknown): value is AssistantError => {
   return true;
 };
 
+const MESSAGE_NOT_SENT = Symbol.for("assistant-ui.message-not-sent");
+
+/**
+ * Rejection reason for a send that never reached the backend, so nothing ran
+ * and nothing is recoverable from the thread. A runtime adapter throws it from
+ * `onNew` to hand the message back to the thread composer, which restores the
+ * draft it cleared at dispatch time when nothing has claimed the composer
+ * since. An edit composer closes at dispatch, so a rejected edit is not
+ * restored.
+ */
+export class MessageNotSentError extends Error {
+  readonly [MESSAGE_NOT_SENT] = true;
+
+  constructor(message = "The message was not sent.") {
+    super(message);
+    this.name = "MessageNotSentError";
+  }
+}
+
+export const isMessageNotSentError = (
+  error: unknown,
+): error is MessageNotSentError =>
+  typeof error === "object" && error !== null && MESSAGE_NOT_SENT in error;
+
 export const toAssistantError = (error: unknown): AssistantError => {
   if (isAssistantError(error)) return error;
   if (error instanceof Error) {

package/src/types/message.ts

@@ -56,6 +56,7 @@ export type ImageMessagePart = {
   readonly type: "image";
   readonly image: string;
   readonly filename?: string;
+  readonly providerMetadata?: PartProviderMetadata;
 };
 
 export type FileMessagePart = {
@@ -65,6 +66,7 @@ export type FileMessagePart = {
   readonly mimeType: string;
   /** How `data` goes on the wire: a url or id reference; omitted = inferred (http(s) → url, else base64). "url" is honored by the LangChain-family, A2A, AG-UI, and Google ADK runtimes; "id" by the LangChain family only. */
   readonly sourceType?: "url" | "id";
+  readonly providerMetadata?: PartProviderMetadata;
   readonly parentId?: string;
 };
 

package/src/utils/composite-context-provider.test.ts

@@ -0,0 +1,139 @@
+import { describe, expect, it, vi } from "vitest";
+import { CompositeContextProvider } from "./composite-context-provider";
+
+describe("CompositeContextProvider", () => {
+  it("notifies every subscriber and rethrows when a registration subscriber throws", () => {
+    const composite = new CompositeContextProvider();
+    const error = new Error("subscriber failed");
+    const laterSubscriber = vi.fn();
+
+    composite.subscribe(() => {
+      throw error;
+    });
+    composite.subscribe(laterSubscriber);
+
+    expect(() =>
+      composite.registerModelContextProvider({
+        getModelContext: () => ({ system: "provider instructions" }),
+      }),
+    ).toThrow(error);
+
+    expect(composite.getModelContext().system).toBe("provider instructions");
+    expect(laterSubscriber).toHaveBeenCalledTimes(1);
+  });
+
+  it("rolls back providers that fail to subscribe", () => {
+    const composite = new CompositeContextProvider();
+    const error = new Error("subscription failed");
+    const observedSystems: Array<string | undefined> = [];
+
+    composite.subscribe(() => {
+      observedSystems.push(composite.getModelContext().system);
+    });
+
+    expect(() =>
+      composite.registerModelContextProvider({
+        getModelContext: () => ({ system: "provider instructions" }),
+        subscribe: (callback) => {
+          callback();
+          throw error;
+        },
+      }),
+    ).toThrow(error);
+
+    expect(composite.getModelContext().system).toBeUndefined();
+    expect(observedSystems).toEqual(["provider instructions", undefined]);
+  });
+
+  it("rethrows the provider subscription failure when a rollback subscriber throws", () => {
+    const composite = new CompositeContextProvider();
+    const subscriptionError = new Error("subscription failed");
+    const subscriberError = new Error("subscriber failed");
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => {});
+
+    composite.subscribe(() => {
+      throw subscriberError;
+    });
+
+    try {
+      expect(() =>
+        composite.registerModelContextProvider({
+          getModelContext: () => ({ system: "provider instructions" }),
+          subscribe: () => {
+            throw subscriptionError;
+          },
+        }),
+      ).toThrow(subscriptionError);
+
+      expect(composite.getModelContext().system).toBeUndefined();
+      expect(consoleError).toHaveBeenCalledWith(subscriberError);
+    } finally {
+      consoleError.mockRestore();
+    }
+  });
+
+  it("preserves an earlier registration when a duplicate fails to subscribe", () => {
+    const composite = new CompositeContextProvider();
+    const error = new Error("subscription failed");
+    let subscriptionCount = 0;
+    const provider = {
+      getModelContext: () => ({ system: "provider instructions" }),
+      subscribe: () => {
+        subscriptionCount++;
+        if (subscriptionCount === 2) throw error;
+        return () => {};
+      },
+    };
+
+    const unregister = composite.registerModelContextProvider(provider);
+
+    expect(() => composite.registerModelContextProvider(provider)).toThrow(
+      error,
+    );
+    expect(composite.getModelContext().system).toBe("provider instructions");
+
+    unregister();
+    expect(composite.getModelContext().system).toBeUndefined();
+  });
+
+  it("keeps duplicate registrations independent", () => {
+    const composite = new CompositeContextProvider();
+    const provider = {
+      getModelContext: () => ({ system: "provider instructions" }),
+    };
+
+    const unregisterFirst = composite.registerModelContextProvider(provider);
+    const unregisterSecond = composite.registerModelContextProvider(provider);
+
+    unregisterFirst();
+    expect(composite.getModelContext().system).toBe("provider instructions");
+
+    unregisterSecond();
+    expect(composite.getModelContext().system).toBeUndefined();
+  });
+
+  it("notifies every subscriber and rethrows on provider updates", () => {
+    const composite = new CompositeContextProvider();
+    const error = new Error("subscriber failed");
+    const laterSubscriber = vi.fn();
+    let publishUpdate = () => {};
+
+    composite.registerModelContextProvider({
+      getModelContext: () => ({ system: "provider instructions" }),
+      subscribe: (callback) => {
+        publishUpdate = callback;
+        return () => {};
+      },
+    });
+
+    composite.subscribe(() => {
+      throw error;
+    });
+    composite.subscribe(laterSubscriber);
+
+    expect(() => publishUpdate()).toThrow(error);
+    expect(laterSubscriber).toHaveBeenCalledTimes(1);
+  });
+});