@assistant-ui/core 0.3.13 → 0.3.14

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;
@@ -158,8 +175,6 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
     this._notifySubscribers();
   }
 
-  // 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 +201,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 +319,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 +331,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.tsx

@@ -15,10 +15,11 @@ import {
   updateStatusReducer,
 } from "../../runtimes/remote-thread-list/remote-thread-state";
 import type {
+  RemoteThreadListAdapter,
   RemoteThreadListOptions,
-  RemoteThreadListProviderComponent,
 } from "../../runtimes/remote-thread-list/types";
 import { RemoteThreadListHookInstanceManager } from "./RemoteThreadListHookInstanceManager";
+import { isTitleSourceMessage } from "./RemoteThreadResource";
 import {
   type ComponentType,
   type FC,
@@ -27,15 +28,12 @@ import {
   useEffect,
   useId,
 } from "react";
+import { useAui } from "@assistant-ui/store";
 import { create } from "zustand";
 import { AssistantMessageStream } from "assistant-stream";
 import type { ModelContextProvider } from "../../model-context/types";
 import { RuntimeAdapterProvider } from "./RuntimeAdapterProvider";
-
-const asProviderComponent = (
-  provider: RemoteThreadListProviderComponent | undefined,
-): ComponentType<PropsWithChildren> =>
-  (provider ?? Fragment) as ComponentType<PropsWithChildren>;
+import { useStableRuntimeAdapters } from "./useRuntimeAdapters";
 
 const threadNotFoundError = (threadIdOrRemoteId: string, action: string) =>
   new Error(`Thread "${threadIdOrRemoteId}" not found while ${action}.`);
@@ -55,6 +53,7 @@ export class RemoteThreadListThreadListRuntimeCore
 {
   private _options!: RemoteThreadListOptions;
   private readonly _hookManager: RemoteThreadListHookInstanceManager;
+  private readonly _runtimeAdapters: { modelContext: ModelContextProvider };
 
   private _loadThreadsPromise: Promise<void> | undefined;
   private _loadMorePromise: Promise<void> | undefined;
@@ -74,6 +73,42 @@ export class RemoteThreadListThreadListRuntimeCore
     threadData: {},
   });
 
+  private readonly _useAdaptersProvider: FC<PropsWithChildren> = ({
+    children,
+  }) => {
+    const useAdapters = this._options.adapter.unstable_useAdapters;
+    if (useAdapters === undefined) return children;
+    return (
+      <this._SynthesizedAdapters useAdapters={useAdapters}>
+        {children}
+      </this._SynthesizedAdapters>
+    );
+  };
+
+  private readonly _SynthesizedAdapters: FC<
+    PropsWithChildren<{
+      useAdapters: NonNullable<RemoteThreadListAdapter["unstable_useAdapters"]>;
+    }>
+  > = ({ useAdapters, children }) => {
+    const adapters = useStableRuntimeAdapters(useAdapters());
+    if (adapters == null) return children;
+    return (
+      <RuntimeAdapterProvider adapters={adapters}>
+        {children}
+      </RuntimeAdapterProvider>
+    );
+  };
+
+  private resolveProvider(
+    adapter: RemoteThreadListAdapter,
+  ): ComponentType<PropsWithChildren> {
+    if (adapter.unstable_Provider !== undefined) {
+      return adapter.unstable_Provider as ComponentType<PropsWithChildren>;
+    }
+    if (adapter.unstable_useAdapters === undefined) return Fragment;
+    return this._useAdaptersProvider;
+  }
+
   public get threadItems() {
     return this._state.value.threadData;
   }
@@ -183,28 +218,27 @@ export class RemoteThreadListThreadListRuntimeCore
     return dedup;
   }
 
-  private readonly contextProvider: ModelContextProvider;
-
   constructor(
     options: RemoteThreadListOptions,
     contextProvider: ModelContextProvider,
   ) {
     super();
-    this.contextProvider = contextProvider;
 
     this._state.subscribe(() => {
       this._notifySubscribers();
       this._notifyThreadIdChange();
     });
+    this._runtimeAdapters = { modelContext: contextProvider };
     this._hookManager = new RemoteThreadListHookInstanceManager(
       options.runtimeHook,
       this,
     );
+    this._hookManager.__internal_setDefaultAdapters(this._runtimeAdapters);
     this._hookManager.__internal_subscribeRunningChanged(() =>
       this._notifySubscribers(),
     );
     this.useProvider = create(() => ({
-      Provider: asProviderComponent(options.adapter.unstable_Provider),
+      Provider: this.resolveProvider(options.adapter),
     }));
     this.__internal_setOptions(options);
     this.switchToNewThread();
@@ -225,7 +259,7 @@ export class RemoteThreadListThreadListRuntimeCore
 
     this._options = options;
 
-    const Provider = asProviderComponent(options.adapter.unstable_Provider);
+    const Provider = this.resolveProvider(options.adapter);
     if (Provider !== this.useProvider.getState().Provider) {
       this.useProvider.setState({ Provider }, true);
     }
@@ -536,7 +570,7 @@ export class RemoteThreadListThreadListRuntimeCore
       if (generation !== this._switchGeneration) return;
     }
 
-    const state = this._state.value;
+    const state = this._state.baseValue;
     let id: string | undefined = this._state.value.newThreadId;
     if (id === undefined) {
       do {
@@ -635,7 +669,10 @@ export class RemoteThreadListThreadListRuntimeCore
     const runtimeCore = this._hookManager.getThreadRuntimeCore(data.id);
     if (!runtimeCore) return; // thread is no longer running
 
-    const messages = runtimeCore.messages;
+    // Incomplete assistant turns (running status, possibly empty content)
+    // would make the payload race-dependent; the title reads settled
+    // messages only, matching the trigger's readiness gate.
+    const messages = runtimeCore.messages.filter(isTitleSourceMessage);
     const stream = await this._options.adapter.generateTitle(
       remoteId,
       messages,
@@ -823,6 +860,10 @@ export class RemoteThreadListThreadListRuntimeCore
     });
   }
 
+  public __internal_dispose() {
+    this._hookManager.__internal_dispose();
+  }
+
   public async detach(threadIdOrRemoteId: string): Promise<void> {
     const data = this.getItemById(threadIdOrRemoteId);
     if (!data) throw threadNotFoundError(threadIdOrRemoteId, "detaching it");
@@ -846,18 +887,16 @@ export class RemoteThreadListThreadListRuntimeCore
 
     const boundIds = this.useBoundIds();
     const { Provider } = this.useProvider();
-
-    const adapters = {
-      modelContext: this.contextProvider,
-    };
+    const aui = useAui();
+    const enabled = boundIds.length === 0 || boundIds[0] === id;
 
     return (
-      (boundIds.length === 0 || boundIds[0] === id) && (
-        // only render if the component is the first one mounted
-        <RuntimeAdapterProvider adapters={adapters}>
+      enabled && (
+        <RuntimeAdapterProvider adapters={this._runtimeAdapters}>
           <this._hookManager.__internal_RenderThreadRuntimes
             provider={Provider}
           />
+          <this._hookManager.__internal_Host parentClient={aui} />
         </RuntimeAdapterProvider>
       )
     );