@assistant-ui/core 0.3.13 → 0.3.15

package/src/react/runtimes/RemoteThreadListHookInstanceManager.tsx

@@ -1,44 +1,53 @@
 import {
   type FC,
   type RefObject,
-  useCallback,
-  useRef,
   useEffect,
-  useEffectEvent,
   useLayoutEffect,
+  useMemo,
+  useRef,
   memo,
   type PropsWithChildren,
   type ComponentType,
-  useMemo,
   Fragment,
 } from "react";
-import { type UseBoundStore, type StoreApi, create } from "zustand";
-import { useAui } from "@assistant-ui/store";
+import { create } from "zustand";
+import { useResources, withKey } from "@assistant-ui/tap";
+import type { AssistantClient } from "@assistant-ui/store";
 import { ThreadListItemRuntimeProvider } from "../providers/ThreadListItemRuntimeProvider";
 import type { ThreadRuntimeCore } from "../../runtime/interfaces/thread-runtime-core";
 import type { ThreadListRuntimeCore } from "../../runtime/interfaces/thread-list-runtime-core";
-import type { AssistantRuntime } from "../../runtime/api/assistant-runtime";
 import type { Unsubscribe } from "../../types/unsubscribe";
 import { BaseSubscribable } from "../../subscribable/subscribable";
-import {
-  getThreadRuntimeCoreIsRunning,
-  type ThreadRuntimeImpl,
-} from "../../runtime/api/thread-runtime";
+import { getThreadRuntimeCoreIsRunning } from "../../runtime/api/thread-runtime";
 import { ThreadListRuntimeImpl } from "../../runtime/api/thread-list-runtime";
-
-type RemoteThreadListHook = () => AssistantRuntime;
+import { invalidateThreadRuntime } from "../../runtime/utils/thread-runtime-lifecycle";
+import {
+  useRuntimeAdapters,
+  type RuntimeAdapters,
+} from "./RuntimeAdapterProvider";
+import {
+  RemoteThreadResource,
+  type RemoteThreadListHook,
+} from "./RemoteThreadResource";
 
 type RemoteThreadListHookInstance = {
   runtime?: ThreadRuntimeCore | undefined;
-  // A runtime riding across a restart stays readable, but only counts as
-  // attached once a binder of the current generation re-publishes it.
   publishedGeneration?: number | undefined;
-  // Part of the binder's React key, so only a bump remounts the hook.
   generation: number;
   isRunning: boolean;
   unsubscribeRunning?: Unsubscribe | undefined;
 };
 
+type AdapterSnapshot = {
+  defaultAdapters: RuntimeAdapters | null;
+  threadAdapters: ReadonlyMap<string, RuntimeAdapters | null>;
+};
+
+type HostSnapshot = {
+  threads: readonly { id: string; generation: number }[];
+  hookEpoch: number;
+};
+
 const ProviderRenderDetector: FC<{
   detectorRef: RefObject<boolean>;
 }> = ({ detectorRef }) => {
@@ -47,15 +56,23 @@ const ProviderRenderDetector: FC<{
   }, [detectorRef]);
   return null;
 };
+
 export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
-  private useRuntimeHook: UseBoundStore<
-    StoreApi<{ useRuntime: RemoteThreadListHook }>
-  >;
+  private runtimeHook: RemoteThreadListHook;
+  private readonly adapterStore = create<AdapterSnapshot>(() => ({
+    defaultAdapters: null,
+    threadAdapters: new Map(),
+  }));
+  private readonly hostStore = create<HostSnapshot>(() => ({
+    threads: [],
+    hookEpoch: 0,
+  }));
+  private readonly pendingThreadAdapters = new Map<
+    string,
+    RuntimeAdapters | null
+  >();
   private instances = new Map<string, RemoteThreadListHookInstance>();
-  // Manager-wide so it survives instance deletion: a stop and start within one
-  // React commit must not reuse a binder key.
   private nextGeneration = 0;
-  private useAliveThreadsKeysChanged = create(() => ({}));
   private parent: ThreadListRuntimeCore;
 
   constructor(
@@ -64,7 +81,7 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
   ) {
     super();
     this.parent = parent;
-    this.useRuntimeHook = create(() => ({ useRuntime: runtimeHook }));
+    this.runtimeHook = runtimeHook;
   }
 
   private _whenRuntimeAttached(threadId: string) {
@@ -78,7 +95,7 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
           !instance.runtime ||
           instance.publishedGeneration !== instance.generation
         ) {
-          return; // not yet published by the current generation's binder
+          return;
         } else {
           dispose();
           resolve(instance.runtime);
@@ -91,11 +108,12 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
 
   public startThreadRuntime(threadId: string) {
     if (!this.instances.has(threadId)) {
+      const generation = this.nextGeneration++;
       this.instances.set(threadId, {
-        generation: this.nextGeneration++,
+        generation,
         isRunning: false,
       });
-      this.useAliveThreadsKeysChanged.setState({}, true);
+      this._syncHostThreads();
     }
 
     return this._whenRuntimeAttached(threadId);
@@ -105,8 +123,9 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
     const instance = this.instances.get(threadId);
     if (!instance) return this.startThreadRuntime(threadId);
 
+    if (instance.runtime) invalidateThreadRuntime(instance.runtime);
     instance.generation = this.nextGeneration++;
-    this.useAliveThreadsKeysChanged.setState({}, true);
+    this._syncHostThreads();
     this._notifySubscribers();
 
     return this._whenRuntimeAttached(threadId);
@@ -124,29 +143,27 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
 
   private runningSubscribers = new Set<() => void>();
 
-  /**
-   * Fires when any thread crosses the running boundary. Separate from the
-   * general subscription so a run does not push the thread list through the
-   * channel that resolves pending runtime attachments.
-   */
   public __internal_subscribeRunningChanged(callback: () => void): Unsubscribe {
     this.runningSubscribers.add(callback);
     return () => this.runningSubscribers.delete(callback);
   }
 
+  private _publish = (
+    threadId: string,
+    runtime: ThreadRuntimeCore,
+    generation: number,
+  ) => {
+    this._publishThreadRuntime(threadId, runtime, generation);
+  };
+
   private _publishThreadRuntime(
     threadId: string,
     runtime: ThreadRuntimeCore,
     generation: number,
   ) {
     const instance = this.instances.get(threadId);
-    if (!instance)
-      throw new Error(
-        `Thread "${threadId}" runtime binding not found. This is a bug in assistant-ui.`,
-      );
+    if (!instance) return;
 
-    // An outgoing binder outlives its generation until React commits the key
-    // change, and must not publish over the incoming one.
     if (instance.generation !== generation) return;
 
     const previousRuntime = instance.runtime;
@@ -156,10 +173,20 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
       this._trackRunning(instance);
     }
     this._notifySubscribers();
+    if (previousRuntime !== undefined && previousRuntime !== runtime) {
+      for (const callback of this.replacedSubscribers) callback();
+    }
+  }
+
+  private replacedSubscribers = new Set<() => void>();
+
+  public __internal_subscribeRuntimeReplaced(
+    callback: () => void,
+  ): Unsubscribe {
+    this.replacedSubscribers.add(callback);
+    return () => this.replacedSubscribers.delete(callback);
   }
 
-  // Run state changes far more often than the thread list does, so the list is
-  // only notified when a thread crosses the running boundary.
   private _trackRunning(instance: RemoteThreadListHookInstance) {
     instance.unsubscribeRunning?.();
 
@@ -186,90 +213,111 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
   }
 
   public stopThreadRuntime(threadId: string) {
-    this.instances.get(threadId)?.unsubscribeRunning?.();
+    const instance = this.instances.get(threadId);
+    if (instance?.runtime) invalidateThreadRuntime(instance.runtime);
+    instance?.unsubscribeRunning?.();
     this.instances.delete(threadId);
-    this.useAliveThreadsKeysChanged.setState({}, true);
+    this.pendingThreadAdapters.delete(threadId);
+    this._syncHostThreads();
     this._notifySubscribers();
   }
 
   public setRuntimeHook(newRuntimeHook: RemoteThreadListHook) {
-    const prevRuntimeHook = this.useRuntimeHook.getState().useRuntime;
-    if (prevRuntimeHook !== newRuntimeHook) {
-      this.useRuntimeHook.setState({ useRuntime: newRuntimeHook }, true);
-    }
+    if (this.runtimeHook === newRuntimeHook) return;
+    this.runtimeHook = newRuntimeHook;
+    this.hostStore.setState(
+      (state) => ({ ...state, hookEpoch: state.hookEpoch + 1 }),
+      true,
+    );
   }
 
-  // Rendered as a child of the user's Provider so the runtime hook can
-  // read context the Provider injects (e.g. RuntimeAdapterProvider).
-  private _RuntimeBinder: FC<
-    PropsWithChildren<{ threadId: string; generation: number }>
-  > = ({ threadId, generation, children }) => {
-    const { useRuntime } = this.useRuntimeHook();
-    const runtime = useRuntime();
-
-    const threadBinding = (runtime.thread as ThreadRuntimeImpl)
-      .__internal_threadBinding;
-
-    const updateRuntime = useCallback(() => {
-      this._publishThreadRuntime(
-        threadId,
-        threadBinding.getState(),
-        generation,
-      );
-    }, [threadId, generation, threadBinding]);
-
-    const isMounted = useRef(false);
-    if (!isMounted.current) {
-      updateRuntime();
-    }
-
-    useEffect(() => {
-      isMounted.current = true;
-      updateRuntime();
-      return threadBinding.outerSubscribe(updateRuntime);
-    }, [threadBinding, updateRuntime]);
-
-    const aui = useAui();
-    const initPromiseRef = useRef<Promise<unknown> | undefined>(undefined);
-    const hasInitializedRef = useRef(false);
+  public __internal_setDefaultAdapters(adapters: RuntimeAdapters | null) {
+    const current = this.adapterStore.getState();
+    if (current.defaultAdapters === adapters) return;
+    this.adapterStore.setState({ ...current, defaultAdapters: adapters }, true);
+  }
 
-    useEffect(() => {
-      const runtimeCore = threadBinding.getState();
-      const setGetInitializePromise = (runtimeCore as Record<string, unknown>)
-        .__internal_setGetInitializePromise;
-      if (typeof setGetInitializePromise === "function") {
-        setGetInitializePromise.call(runtimeCore, () => initPromiseRef.current);
-      }
-    }, [threadBinding]);
+  public __internal_setThreadAdapters(
+    threadId: string,
+    adapters: RuntimeAdapters | null,
+  ) {
+    const current = this.adapterStore.getState();
+    const next = new Map(current.threadAdapters);
+    if (adapters === null) next.delete(threadId);
+    else next.set(threadId, adapters);
+    this.adapterStore.setState({ ...current, threadAdapters: next }, true);
+  }
 
-    const handleInitialize = useEffectEvent(() => {
-      if (hasInitializedRef.current) return;
+  public __internal_dispose() {
+    for (const threadId of [...this.instances.keys()]) {
+      this.stopThreadRuntime(threadId);
+    }
+  }
 
-      const state = aui.threadListItem.getState();
-      if (state.status !== "new") return;
-      hasInitializedRef.current = true;
+  private _syncHostThreads() {
+    this.hostStore.setState(
+      (state) => ({
+        ...state,
+        threads: Array.from(this.instances.entries()).map(
+          ([id, { generation }]) => ({ id, generation }),
+        ),
+      }),
+      true,
+    );
+  }
 
-      initPromiseRef.current = aui.threadListItem.initialize();
+  public __internal_useHost(parentClient: AssistantClient) {
+    const { threads, hookEpoch } = this.hostStore();
+    const adapters = this.adapterStore();
+    const runtimeHook = this.runtimeHook;
+    return useResources(
+      threads.map(({ id, generation }) => {
+        const threadAdapters = this.pendingThreadAdapters.has(id)
+          ? this.pendingThreadAdapters.get(id)!
+          : (adapters.threadAdapters.get(id) ?? adapters.defaultAdapters);
+        return withKey(
+          `${id}:${generation}:${hookEpoch}`,
+          RemoteThreadResource({
+            threadId: id,
+            generation,
+            parentList: this.parent,
+            runtimeHook,
+            parentClient,
+            adapters: threadAdapters,
+            publish: this._publish,
+          }),
+        );
+      }),
+    );
+  }
 
-      const dispose = runtime.thread.unstable_on("runEnd", () => {
-        dispose();
-        aui.threadListItem.generateTitle();
-      });
-    });
+  public __internal_RenderThreadRuntimes: FC<{
+    provider: ComponentType<PropsWithChildren>;
+  }> = ({ provider }) => {
+    this.hostStore();
 
-    useEffect(() => {
-      hasInitializedRef.current = false;
-      return runtime.threads.main.unstable_on("initialize", handleInitialize);
-    }, [runtime]);
+    return Array.from(this.instances.entries()).map(
+      ([threadId, { generation }]) => (
+        <this._OuterActiveThreadProvider
+          key={`${threadId}:${generation}`}
+          threadId={threadId}
+          provider={provider}
+        />
+      ),
+    );
+  };
 
-    return <>{children}</>;
+  public __internal_Host: FC<{ parentClient: AssistantClient }> = ({
+    parentClient,
+  }) => {
+    this.__internal_useHost(parentClient);
+    return null;
   };
 
   private _OuterActiveThreadProvider: FC<{
     threadId: string;
-    generation: number;
     provider: ComponentType<PropsWithChildren>;
-  }> = memo(({ threadId, generation, provider: Provider }) => {
+  }> = memo(({ threadId, provider: Provider }) => {
     const runtime = useMemo(
       () => new ThreadListRuntimeImpl(this.parent).getItemById(threadId),
       [threadId],
@@ -283,7 +331,7 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
             console.warn(
               "RemoteThreadListAdapter.unstable_Provider did not render its `children` synchronously. " +
                 "Render `children` on first commit; deferring them behind a loading state, Suspense boundary, " +
-                "or `useEffect` gate strands the runtime binder and leaves the thread without context.",
+                "or `useEffect` gate strands the runtime host and leaves the thread without context.",
             );
           }
         }, 100);
@@ -295,28 +343,27 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
     return (
       <ThreadListItemRuntimeProvider runtime={runtime}>
         <Provider>
-          <this._RuntimeBinder threadId={threadId} generation={generation}>
-            <ProviderRenderDetector detectorRef={detectorRef} />
-          </this._RuntimeBinder>
+          <this._AdapterSink threadId={threadId} detectorRef={detectorRef} />
         </Provider>
       </ThreadListItemRuntimeProvider>
     );
   });
 
-  public __internal_RenderThreadRuntimes: FC<{
-    provider: ComponentType<PropsWithChildren>;
-  }> = ({ provider }) => {
-    this.useAliveThreadsKeysChanged(); // trigger re-render on alive threads change
-
-    return Array.from(this.instances.entries()).map(
-      ([threadId, { generation }]) => (
-        <this._OuterActiveThreadProvider
-          key={`${threadId}:${generation}`}
-          threadId={threadId}
-          generation={generation}
-          provider={provider}
-        />
-      ),
-    );
+  private _AdapterSink: FC<{
+    threadId: string;
+    detectorRef: RefObject<boolean>;
+  }> = ({ threadId, detectorRef }) => {
+    const adapters = useRuntimeAdapters();
+    this.pendingThreadAdapters.set(threadId, adapters);
+    useLayoutEffect(() => {
+      this.__internal_setThreadAdapters(threadId, adapters);
+      return () => {
+        if (this.pendingThreadAdapters.get(threadId) === adapters) {
+          this.pendingThreadAdapters.delete(threadId);
+        }
+        this.__internal_setThreadAdapters(threadId, null);
+      };
+    }, [threadId, adapters]);
+    return <ProviderRenderDetector detectorRef={detectorRef} />;
   };
 }

package/src/react/runtimes/RemoteThreadListThreadListRuntimeCore.republish.test.ts

@@ -0,0 +1,87 @@
+import { describe, expect, it } from "vitest";
+import { RemoteThreadListThreadListRuntimeCore } from "./RemoteThreadListThreadListRuntimeCore";
+import { ThreadListRuntimeImpl } from "../../runtime/api/thread-list-runtime";
+import { ReadonlyThreadRuntimeCore } from "../../runtimes/readonly/ReadonlyThreadRuntimeCore";
+import {
+  contextProvider,
+  makeAdapter,
+} from "../../tests/remote-thread-list-test-helpers";
+import type { ThreadMessage } from "../../types/message";
+
+const message = (id: string): ThreadMessage =>
+  ({
+    id,
+    role: "user",
+    content: [{ type: "text", text: id }],
+    createdAt: new Date(0),
+    attachments: [],
+    metadata: { custom: {} },
+  }) as unknown as ThreadMessage;
+
+const setup = () => {
+  const core = new RemoteThreadListThreadListRuntimeCore(
+    { adapter: makeAdapter(), runtimeHook: () => ({}) as never },
+    contextProvider,
+  );
+  const hookManager = (
+    core as unknown as {
+      _hookManager: {
+        _publishThreadRuntime: (
+          id: string,
+          runtime: ReadonlyThreadRuntimeCore,
+          generation: number,
+        ) => void;
+        instances: Map<string, { generation: number }>;
+      };
+    }
+  )._hookManager;
+  const threadId = core.mainThreadId!;
+  const publish = (runtime: ReadonlyThreadRuntimeCore) =>
+    hookManager._publishThreadRuntime(
+      threadId,
+      runtime,
+      hookManager.instances.get(threadId)!.generation,
+    );
+  return { core, publish };
+};
+
+describe("RemoteThreadListThreadListRuntimeCore republish", () => {
+  it("keeps the subscribed main thread state on the republished core", async () => {
+    const { core, publish } = setup();
+    const first = new ReadonlyThreadRuntimeCore();
+    first.setMessages([message("m1")]);
+    publish(first);
+
+    const thread = new ThreadListRuntimeImpl(core).main;
+    thread.subscribe(() => {});
+    expect(thread.getState().messages.map((m) => m.id)).toEqual(["m1"]);
+
+    const second = new ReadonlyThreadRuntimeCore();
+    second.setMessages([message("m2")]);
+    publish(second);
+    await Promise.resolve();
+
+    const ids = thread.getState().messages.map((m) => m.id);
+    expect(ids).toEqual(["m2"]);
+    for (const id of ids) {
+      expect(() => thread.getMessageById(id).getState()).not.toThrow();
+    }
+  });
+
+  it("notifies thread subscribers when the core is republished", async () => {
+    const { core, publish } = setup();
+    publish(new ReadonlyThreadRuntimeCore());
+
+    const thread = new ThreadListRuntimeImpl(core).main;
+    let notified = 0;
+    thread.subscribe(() => notified++);
+
+    const next = new ReadonlyThreadRuntimeCore();
+    next.setMessages([message("m1")]);
+    publish(next);
+    await Promise.resolve();
+
+    expect(notified).toBeGreaterThan(0);
+    expect(thread.getState().messages.map((m) => m.id)).toEqual(["m1"]);
+  });
+});