@assistant-ui/core 0.3.12 → 0.3.14

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

@@ -0,0 +1,131 @@
+import { describe, expect, it, vi } from "vitest";
+import { ModelContextRegistry } from "./registry";
+
+describe("ModelContextRegistry", () => {
+  it("notifies every subscriber and rethrows when a registration subscriber throws", () => {
+    const registry = new ModelContextRegistry();
+    const error = new Error("subscriber failed");
+    const laterSubscriber = vi.fn();
+
+    registry.subscribe(() => {
+      throw error;
+    });
+    registry.subscribe(laterSubscriber);
+
+    expect(() => registry.addTool({ toolName: "search" })).toThrow(error);
+
+    expect(registry.getModelContext().tools).toHaveProperty("search");
+    expect(laterSubscriber).toHaveBeenCalledTimes(1);
+  });
+
+  it("keeps a provider registered when a registration subscriber throws", () => {
+    const registry = new ModelContextRegistry();
+    const error = new Error("subscriber failed");
+    const laterSubscriber = vi.fn();
+
+    registry.subscribe(() => {
+      throw error;
+    });
+    registry.subscribe(laterSubscriber);
+
+    expect(() =>
+      registry.addProvider({
+        getModelContext: () => ({ system: "provider instructions" }),
+      }),
+    ).toThrow(error);
+
+    expect(registry.getModelContext().system).toBe("provider instructions");
+    expect(laterSubscriber).toHaveBeenCalledTimes(1);
+  });
+
+  it("notifies every subscriber and rethrows when registering instructions", () => {
+    const registry = new ModelContextRegistry();
+    const error = new Error("subscriber failed");
+    const laterSubscriber = vi.fn();
+
+    registry.subscribe(() => {
+      throw error;
+    });
+    registry.subscribe(laterSubscriber);
+
+    expect(() => registry.addInstruction("be concise")).toThrow(error);
+
+    expect(registry.getModelContext().system).toContain("be concise");
+    expect(laterSubscriber).toHaveBeenCalledTimes(1);
+  });
+
+  it("rethrows the provider subscription failure when a rollback subscriber throws", () => {
+    const registry = new ModelContextRegistry();
+    const subscriptionError = new Error("subscription failed");
+    const subscriberError = new Error("subscriber failed");
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => {});
+
+    registry.subscribe(() => {
+      throw subscriberError;
+    });
+
+    try {
+      expect(() =>
+        registry.addProvider({
+          getModelContext: () => ({ system: "provider instructions" }),
+          subscribe: () => {
+            throw subscriptionError;
+          },
+        }),
+      ).toThrow(subscriptionError);
+
+      expect(registry.getModelContext().system).toBeUndefined();
+      expect(consoleError).toHaveBeenCalledWith(subscriberError);
+    } finally {
+      consoleError.mockRestore();
+    }
+  });
+
+  it("rolls back providers that fail to subscribe", () => {
+    const registry = new ModelContextRegistry();
+    const error = new Error("subscription failed");
+    const observedSystems: Array<string | undefined> = [];
+
+    registry.subscribe(() => {
+      observedSystems.push(registry.getModelContext().system);
+    });
+
+    expect(() =>
+      registry.addProvider({
+        getModelContext: () => ({ system: "provider instructions" }),
+        subscribe: (callback) => {
+          callback();
+          throw error;
+        },
+      }),
+    ).toThrow(error);
+
+    expect(registry.getModelContext().system).toBeUndefined();
+    expect(observedSystems).toEqual(["provider instructions", undefined]);
+  });
+
+  it("notifies every subscriber and rethrows on provider updates", () => {
+    const registry = new ModelContextRegistry();
+    const error = new Error("subscriber failed");
+    const laterSubscriber = vi.fn();
+    let publishUpdate = () => {};
+
+    registry.addProvider({
+      getModelContext: () => ({ system: "provider instructions" }),
+      subscribe: (callback) => {
+        publishUpdate = callback;
+        return () => {};
+      },
+    });
+
+    registry.subscribe(() => {
+      throw error;
+    });
+    registry.subscribe(laterSubscriber);
+
+    expect(() => publishUpdate()).toThrow(error);
+    expect(laterSubscriber).toHaveBeenCalledTimes(1);
+  });
+});

package/src/model-context/registry.ts

@@ -7,6 +7,7 @@ import {
   type AssistantInstructionsConfig,
 } from "./types";
 import type { Unsubscribe } from "../types/unsubscribe";
+import { notifySubscribers as notifyStateSubscribers } from "../subscribable/subscribable";
 import type {
   ModelContextRegistryToolHandle,
   ModelContextRegistryInstructionHandle,
@@ -75,9 +76,7 @@ export class ModelContextRegistry implements ModelContextProvider {
   }
 
   private notifySubscribers(): void {
-    for (const callback of this._subscribers) {
-      callback();
-    }
+    notifyStateSubscribers(this._subscribers);
   }
 
   addTool<TArgs extends Record<string, unknown>, TResult>(
@@ -144,9 +143,20 @@ export class ModelContextRegistry implements ModelContextProvider {
 
     this._providers.set(id, provider);
 
-    const unsubscribe = provider.subscribe?.(() => {
-      this.notifySubscribers();
-    });
+    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();

package/src/react/RuntimeAdapter.test.ts

@@ -54,6 +54,7 @@ describe("RuntimeAdapter via the neutral store entry", () => {
     const handle = createAssistantClient(
       AuiConfig({ threads: RuntimeAdapter(runtime) }),
     );
+    handle.subscribe(() => {});
     const aui = handle.getClient();
 
     flushTapSync(() => aui.composer.setText("hello"));
@@ -75,6 +76,7 @@ describe("RuntimeAdapter via the neutral store entry", () => {
     const handle = createAssistantClient(
       AuiConfig({ threads: RuntimeAdapter(runtime) }),
     );
+    handle.subscribe(() => {});
 
     // Dev mode double-mounts the root (StrictMode emulation), so assert
     // presence rather than an exact count

package/src/react/RuntimeAdapter.ts

@@ -1,5 +1,5 @@
-import { useEffect } from "react";
 import { useResource, resource } from "@assistant-ui/tap";
+import type { AssistantClient, ScopesConfig } from "@assistant-ui/store";
 import type { AssistantRuntime } from "..";
 import {
   baseRuntimeAdapterTransformScopes,
@@ -8,6 +8,7 @@ import {
 import {
   attachTransformScopes,
   useAssistantClientRef,
+  useAssistantScopeEffect,
 } from "@assistant-ui/store/client";
 import { DataRenderers } from "./client/DataRenderers";
 import { Tools } from "./client/Tools";
@@ -15,11 +16,12 @@ import { Tools } from "./client/Tools";
 const useRuntimeAdapter = (runtime: AssistantRuntime) => {
   const clientRef = useAssistantClientRef();
 
-  useEffect(() => {
-    return runtime.registerModelContextProvider(
-      clientRef.current!.modelContext(),
-    );
-  }, [runtime, clientRef]);
+  useAssistantScopeEffect(
+    "modelContext",
+    () =>
+      runtime.registerModelContextProvider(clientRef.current!.modelContext()),
+    [runtime],
+  );
 
   return useResource(
     ThreadListClient({
@@ -31,7 +33,15 @@ const useRuntimeAdapter = (runtime: AssistantRuntime) => {
 
 export const RuntimeAdapter = resource(useRuntimeAdapter);
 
-attachTransformScopes(useRuntimeAdapter, (scopes, parent) => {
+/**
+ * The scope defaults `RuntimeAdapter` installs when it is used as the `threads`
+ * config entry. Adapter packages that wrap a runtime in their own config entry
+ * attach this to the wrapping resource for scope parity with `RuntimeAdapter`.
+ */
+export const runtimeAdapterTransformScopes = (
+  scopes: ScopesConfig,
+  parent: AssistantClient,
+): void => {
   baseRuntimeAdapterTransformScopes(scopes, parent);
 
   if (!scopes.tools && parent.tools.source === null) {
@@ -41,4 +51,6 @@ attachTransformScopes(useRuntimeAdapter, (scopes, parent) => {
   if (!scopes.dataRenderers && parent.dataRenderers.source === null) {
     scopes.dataRenderers = DataRenderers();
   }
-});
+};
+
+attachTransformScopes(useRuntimeAdapter, runtimeAdapterTransformScopes);

package/src/react/adapters/LocalStorageThreadListAdapter.tsx

@@ -21,7 +21,10 @@ import type {
   ExportedMessageRepository,
   ExportedMessageRepositoryItem,
 } from "../../internal";
-import { RuntimeAdapterProvider } from "../runtimes/RuntimeAdapterProvider";
+import {
+  RuntimeAdapterProvider,
+  type RuntimeAdapters,
+} from "../runtimes/RuntimeAdapterProvider";
 import type { TitleGenerationAdapter } from "./TitleGenerationAdapter";
 
 export type AsyncStorageLike = {
@@ -338,29 +341,40 @@ class AsyncStorageHistoryAdapter implements ThreadHistoryAdapter {
   }
 }
 
+const useLocalStorageThreadAdapters = (
+  storage: AsyncStorageLike,
+  prefix: string,
+  mutationQueue: KeyedMutationQueue,
+): RuntimeAdapters => {
+  const aui = useAui();
+  // Not useEffectEvent: history adapter methods run during render (SSR load).
+  const auiRef = useRef(aui);
+  useEffect(() => {
+    auiRef.current = aui;
+  });
+  const [history] = useState(
+    () =>
+      new AsyncStorageHistoryAdapter(
+        storage,
+        () => auiRef.current,
+        prefix,
+        mutationQueue,
+      ),
+  );
+  return useMemo(() => ({ history }), [history]);
+};
+
 const createHistoryProvider = (
   storage: AsyncStorageLike,
   prefix: string,
   mutationQueue: KeyedMutationQueue,
 ): FC<PropsWithChildren> => {
   const Provider: FC<PropsWithChildren> = ({ children }) => {
-    const aui = useAui();
-    // Not useEffectEvent: history adapter methods run during render (SSR load).
-    const auiRef = useRef(aui);
-    useEffect(() => {
-      auiRef.current = aui;
-    });
-    const [history] = useState(
-      () =>
-        new AsyncStorageHistoryAdapter(
-          storage,
-          () => auiRef.current,
-          prefix,
-          mutationQueue,
-        ),
+    const adapters = useLocalStorageThreadAdapters(
+      storage,
+      prefix,
+      mutationQueue,
     );
-    const adapters = useMemo(() => ({ history }), [history]);
-
     return (
       <RuntimeAdapterProvider adapters={adapters}>
         {children}
@@ -406,6 +420,9 @@ export const createLocalStorageAdapter = (
 
   const adapter: RemoteThreadListAdapter = {
     unstable_Provider: createHistoryProvider(storage, prefix, mutationQueue),
+    unstable_useAdapters: function useLocalStorageAdapters() {
+      return useLocalStorageThreadAdapters(storage, prefix, mutationQueue);
+    },
 
     async list(): Promise<RemoteThreadListResponse> {
       const threads = await loadThreadMetadata();

package/src/react/client/AdaptedRemoteThread.ts

@@ -0,0 +1,41 @@
+import { useEffect, useMemo } from "react";
+import { resource, type ResourceElement } from "@assistant-ui/tap";
+import { useClientResource } from "@assistant-ui/store/client";
+import type { ClientOutput } from "@assistant-ui/store";
+import type { RuntimeAdapters } from "../../runtimes/remote-thread-list/types";
+import { isDevelopment } from "../../store/internal";
+import {
+  useRuntimeAdapters,
+  useRuntimeAdaptersProvider,
+  useStableRuntimeAdapters,
+} from "../runtimes/useRuntimeAdapters";
+
+const useAdaptedRemoteThread = ({
+  useAdapters,
+  thread,
+}: {
+  useAdapters: () => RuntimeAdapters | null | undefined;
+  thread: ResourceElement<ClientOutput<"thread">>;
+}): ClientOutput<"thread"> => {
+  const parent = useRuntimeAdapters();
+  const stableAdapters = useStableRuntimeAdapters(useAdapters());
+
+  const unkeyedHistory =
+    stableAdapters?.history != null && thread.key === undefined;
+  useEffect(() => {
+    if (!isDevelopment || !unkeyedHistory) return;
+    console.warn(
+      "[assistant-ui] RemoteThreadList received a history adapter, but the thread factory is unkeyed. Key the thread resource by thread id (withKey) so switching threads reloads history.",
+    );
+  }, [unkeyedHistory]);
+
+  const merged = useMemo(
+    () => (stableAdapters == null ? parent : { ...parent, ...stableAdapters }),
+    [parent, stableAdapters],
+  );
+  return useRuntimeAdaptersProvider(merged, function useBoundRemoteThread() {
+    return useClientResource(thread).methods;
+  });
+};
+
+export const AdaptedRemoteThread = resource(useAdaptedRemoteThread);

package/src/react/client/InMemoryThreadList.test.tsx

@@ -0,0 +1,148 @@
+// @vitest-environment jsdom
+
+import { act, render } from "@testing-library/react";
+import { describe, expect, it, vi } from "vitest";
+import { resource } from "@assistant-ui/tap";
+import { AuiProvider, useAui, useAuiEvent } from "@assistant-ui/store";
+import { InMemoryThreadList } from "./InMemoryThreadList";
+
+const stubComposer = { getState: () => ({}) };
+const stubSuggestions = { getState: () => ({ suggestions: [] }) };
+const useStubThread = (_props: { threadId: string }) => ({
+  getState: () => ({ isRunning: false }),
+  composer: () => stubComposer,
+  suggestions: () => stubSuggestions,
+});
+const StubThread = resource(useStubThread);
+
+const setup = () => {
+  const selectionChanged = vi.fn();
+  let aui!: ReturnType<typeof useAui>;
+  const Consumer = () => {
+    useAuiEvent("threads.selectionChanged" as never, selectionChanged as never);
+    return null;
+  };
+  const Harness = () => {
+    aui = useAui({
+      threads: InMemoryThreadList({
+        thread: (threadId) => StubThread({ threadId }) as never,
+      }),
+    } as never);
+    return (
+      <AuiProvider value={aui}>
+        <Consumer />
+      </AuiProvider>
+    );
+  };
+  render(<Harness />);
+  return { getAui: () => aui, selectionChanged };
+};
+
+describe("InMemoryThreadList selection events", () => {
+  it("does not emit for the initially selected thread on mount", async () => {
+    const { selectionChanged } = setup();
+    await act(async () => {});
+    expect(selectionChanged).not.toHaveBeenCalled();
+  });
+
+  it("emits on switchToNewThread and on switching back", async () => {
+    const { getAui, selectionChanged } = setup();
+    await act(async () => {});
+
+    await act(async () => {
+      getAui().threads.switchToNewThread();
+    });
+    await act(async () => {});
+
+    const newThreadId = getAui().threads.getState().mainThreadId;
+    expect(newThreadId).not.toBe("main");
+    expect(selectionChanged).toHaveBeenCalledExactlyOnceWith({
+      threadId: newThreadId,
+      previousThreadId: "main",
+    });
+
+    await act(async () => {
+      getAui().threads.switchToThread("main");
+    });
+    await act(async () => {});
+
+    expect(selectionChanged).toHaveBeenCalledTimes(2);
+    expect(selectionChanged).toHaveBeenLastCalledWith({
+      threadId: "main",
+      previousThreadId: newThreadId,
+    });
+  });
+
+  it("emits when deleting the selected thread falls back to another", async () => {
+    const { getAui, selectionChanged } = setup();
+    await act(async () => {});
+
+    await act(async () => {
+      getAui().threads.switchToNewThread();
+    });
+    await act(async () => {});
+
+    const newThreadId = getAui().threads.getState().mainThreadId;
+    selectionChanged.mockClear();
+
+    await act(async () => {
+      getAui().threads.item({ id: newThreadId }).delete();
+    });
+    await act(async () => {});
+
+    expect(getAui().threads.getState().mainThreadId).toBe("main");
+    expect(selectionChanged).toHaveBeenCalledExactlyOnceWith({
+      threadId: "main",
+      previousThreadId: newThreadId,
+    });
+  });
+});
+
+describe("InMemoryThreadList delete", () => {
+  it("notifies onDelete with the removed thread id", async () => {
+    const onDelete = vi.fn();
+    let aui!: ReturnType<typeof useAui>;
+    const Harness = () => {
+      aui = useAui({
+        threads: InMemoryThreadList({
+          thread: (threadId) => StubThread({ threadId }) as never,
+          onDelete,
+        }),
+      } as never);
+      return <AuiProvider value={aui}>{null}</AuiProvider>;
+    };
+    render(<Harness />);
+    await act(async () => {});
+
+    await act(async () => {
+      aui.threads.switchToNewThread();
+    });
+    const doomed = aui.threads.getState().mainThreadId;
+    await act(async () => {
+      aui.threads.item({ id: doomed }).delete();
+    });
+    expect(onDelete).toHaveBeenCalledExactlyOnceWith(doomed);
+  });
+
+  it("starts a fresh thread when the last one is deleted", async () => {
+    let aui!: ReturnType<typeof useAui>;
+    const Harness = () => {
+      aui = useAui({
+        threads: InMemoryThreadList({
+          thread: (threadId) => StubThread({ threadId }) as never,
+        }),
+      } as never);
+      return <AuiProvider value={aui}>{null}</AuiProvider>;
+    };
+    render(<Harness />);
+    await act(async () => {});
+
+    await act(async () => {
+      aui.threads.item({ id: "main" }).delete();
+    });
+    const state = aui.threads.getState();
+    expect(state.threadIds).toHaveLength(1);
+    expect(state.mainThreadId).toBe(state.threadIds[0]);
+    expect(state.mainThreadId).not.toBe("main");
+  });
+});