@assistant-ui/core 0.3.13 → 0.3.14

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

@@ -1,8 +1,12 @@
 // @vitest-environment jsdom
 
-import { act, render } from "@testing-library/react";
+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";
@@ -10,6 +14,15 @@ import type { ExternalStoreAdapter } from "../runtimes/external-store/external-s
 
 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 = [
     {
@@ -110,6 +123,57 @@ describe("thread switch events", () => {
     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();
@@ -118,6 +182,10 @@ describe("thread switch events", () => {
         { scope: "*", event: "threadListItem.switchedTo" } as never,
         anySwitch as never,
       );
+      useAuiEvent(
+        { scope: "*", event: "threads.selectionChanged" } as never,
+        anySwitch as never,
+      );
       return null;
     };
     const Harness = () => {
@@ -133,4 +201,34 @@ describe("thread switch events", () => {
 
     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/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);
+  });
+});

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

@@ -2,23 +2,41 @@ import {
   type ModelContextProvider,
   mergeModelContexts,
 } from "../model-context/types";
+import { notifySubscribers as notifyStateSubscribers } from "../subscribable/subscribable";
+import type { Unsubscribe } from "../types/unsubscribe";
 
 export class CompositeContextProvider implements ModelContextProvider {
-  private _providers = new Set<ModelContextProvider>();
+  private _providers = new Map<symbol, ModelContextProvider>();
+  private _providerUnsubscribes = new Map<symbol, Unsubscribe | undefined>();
 
   getModelContext() {
-    return mergeModelContexts(this._providers);
+    return mergeModelContexts(new Set(this._providers.values()));
   }
 
   registerModelContextProvider(provider: ModelContextProvider) {
-    this._providers.add(provider);
-    const unsubscribe = provider.subscribe?.(() => {
-      this.notifySubscribers();
-    });
+    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();
     return () => {
-      this._providers.delete(provider);
+      this._providers.delete(id);
+      const unsubscribe = this._providerUnsubscribes.get(id);
       unsubscribe?.();
+      this._providerUnsubscribes.delete(id);
       this.notifySubscribers();
     };
   }
@@ -26,7 +44,7 @@ export class CompositeContextProvider implements ModelContextProvider {
   private _subscribers = new Set<() => void>();
 
   notifySubscribers() {
-    for (const callback of this._subscribers) callback();
+    notifyStateSubscribers(this._subscribers);
   }
 
   subscribe(callback: () => void) {

package/src/utils/notify-event-listeners.ts

@@ -1,8 +1,9 @@
-type EventListener = (payload?: unknown) => void;
+type EventListener<T> = (payload: T) => unknown;
 
-export const notifyEventListeners = (
-  listeners: Iterable<EventListener>,
-  payload: unknown,
+// A function payload is indistinguishable from a factory and gets invoked; wrap a function-typed payload in a factory.
+export const notifyEventListeners = <T>(
+  listeners: Iterable<EventListener<T>>,
+  payloadOrFactory: T | (() => T),
   errorContext: string,
 ) => {
   const reportError = (error: unknown) => {
@@ -14,7 +15,11 @@ export const notifyEventListeners = (
 
   for (const listener of listeners) {
     try {
-      const result = listener(payload) as unknown;
+      const result = listener(
+        typeof payloadOrFactory === "function"
+          ? (payloadOrFactory as () => T)()
+          : payloadOrFactory,
+      );
       if (
         result !== null &&
         (typeof result === "object" || typeof result === "function") &&