@assistant-ui/core 0.3.15 → 0.3.17

package/src/react/runtimes/RemoteThreadListHookInstanceManager.tsx

@@ -10,17 +10,27 @@ import {
   type ComponentType,
   Fragment,
 } from "react";
-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 {
+  ThreadRuntimeCore,
+  ThreadRuntimeEventType,
+} from "../../runtime/interfaces/thread-runtime-core";
+import type {
+  ThreadListRuntimeCore,
+  ThreadListRuntimeEvent,
+} from "../../runtime/interfaces/thread-list-runtime-core";
 import type { Unsubscribe } from "../../types/unsubscribe";
-import { BaseSubscribable } from "../../subscribable/subscribable";
+import {
+  BaseSubscribable,
+  WritableSubscribable,
+} from "../../subscribable/subscribable";
+import { useSubscribable } from "../../store/runtime-clients/useSubscribable";
 import { getThreadRuntimeCoreIsRunning } from "../../runtime/api/thread-runtime";
 import { ThreadListRuntimeImpl } from "../../runtime/api/thread-list-runtime";
 import { invalidateThreadRuntime } from "../../runtime/utils/thread-runtime-lifecycle";
+import { notifyEventListeners } from "../../utils/notify-event-listeners";
 import {
   useRuntimeAdapters,
   type RuntimeAdapters,
@@ -30,6 +40,13 @@ import {
   type RemoteThreadListHook,
 } from "./RemoteThreadResource";
 
+const THREAD_EVENTS = [
+  "runStart",
+  "runEnd",
+  "initialize",
+  "modelContextUpdate",
+] as const satisfies readonly ThreadRuntimeEventType[];
+
 type RemoteThreadListHookInstance = {
   runtime?: ThreadRuntimeCore | undefined;
   publishedGeneration?: number | undefined;
@@ -59,14 +76,14 @@ const ProviderRenderDetector: FC<{
 
 export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
   private runtimeHook: RemoteThreadListHook;
-  private readonly adapterStore = create<AdapterSnapshot>(() => ({
+  private readonly adapterStore = new WritableSubscribable<AdapterSnapshot>({
     defaultAdapters: null,
     threadAdapters: new Map(),
-  }));
-  private readonly hostStore = create<HostSnapshot>(() => ({
+  });
+  private readonly hostStore = new WritableSubscribable<HostSnapshot>({
     threads: [],
     hookEpoch: 0,
-  }));
+  });
   private readonly pendingThreadAdapters = new Map<
     string,
     RuntimeAdapters | null
@@ -148,6 +165,17 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
     return () => this.runningSubscribers.delete(callback);
   }
 
+  private threadEventSubscribers = new Set<
+    (event: ThreadListRuntimeEvent) => void
+  >();
+
+  public __internal_subscribeThreadEvents(
+    callback: (event: ThreadListRuntimeEvent) => void,
+  ): Unsubscribe {
+    this.threadEventSubscribers.add(callback);
+    return () => this.threadEventSubscribers.delete(callback);
+  }
+
   private _publish = (
     threadId: string,
     runtime: ThreadRuntimeCore,
@@ -170,7 +198,7 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
     instance.runtime = runtime;
     instance.publishedGeneration = generation;
     if (previousRuntime !== runtime) {
-      this._trackRunning(instance);
+      this._trackRunning(threadId, instance);
     }
     this._notifySubscribers();
     if (previousRuntime !== undefined && previousRuntime !== runtime) {
@@ -187,7 +215,10 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
     return () => this.replacedSubscribers.delete(callback);
   }
 
-  private _trackRunning(instance: RemoteThreadListHookInstance) {
+  private _trackRunning(
+    threadId: string,
+    instance: RemoteThreadListHookInstance,
+  ) {
     instance.unsubscribeRunning?.();
 
     const runtime = instance.runtime;
@@ -198,9 +229,23 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
     }
 
     this._setRunning(instance, getThreadRuntimeCoreIsRunning(runtime));
-    instance.unsubscribeRunning = runtime.subscribe(() => {
-      this._setRunning(instance, getThreadRuntimeCoreIsRunning(runtime));
-    });
+    const unsubscribers = [
+      runtime.subscribe(() => {
+        this._setRunning(instance, getThreadRuntimeCoreIsRunning(runtime));
+      }),
+      ...THREAD_EVENTS.map((type) =>
+        runtime.unstable_on(type, () => {
+          notifyEventListeners(
+            this.threadEventSubscribers,
+            { threadId, type },
+            `Thread event "${type}"`,
+          );
+        }),
+      ),
+    ];
+    instance.unsubscribeRunning = () => {
+      for (const unsubscribe of unsubscribers) unsubscribe();
+    };
   }
 
   private _setRunning(
@@ -225,16 +270,14 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
   public setRuntimeHook(newRuntimeHook: RemoteThreadListHook) {
     if (this.runtimeHook === newRuntimeHook) return;
     this.runtimeHook = newRuntimeHook;
-    this.hostStore.setState(
-      (state) => ({ ...state, hookEpoch: state.hookEpoch + 1 }),
-      true,
-    );
+    const host = this.hostStore.getState();
+    this.hostStore.setState({ ...host, hookEpoch: host.hookEpoch + 1 });
   }
 
   public __internal_setDefaultAdapters(adapters: RuntimeAdapters | null) {
     const current = this.adapterStore.getState();
     if (current.defaultAdapters === adapters) return;
-    this.adapterStore.setState({ ...current, defaultAdapters: adapters }, true);
+    this.adapterStore.setState({ ...current, defaultAdapters: adapters });
   }
 
   public __internal_setThreadAdapters(
@@ -245,7 +288,7 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
     const next = new Map(current.threadAdapters);
     if (adapters === null) next.delete(threadId);
     else next.set(threadId, adapters);
-    this.adapterStore.setState({ ...current, threadAdapters: next }, true);
+    this.adapterStore.setState({ ...current, threadAdapters: next });
   }
 
   public __internal_dispose() {
@@ -255,20 +298,18 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
   }
 
   private _syncHostThreads() {
-    this.hostStore.setState(
-      (state) => ({
-        ...state,
-        threads: Array.from(this.instances.entries()).map(
-          ([id, { generation }]) => ({ id, generation }),
-        ),
-      }),
-      true,
-    );
+    const host = this.hostStore.getState();
+    this.hostStore.setState({
+      ...host,
+      threads: Array.from(this.instances.entries()).map(
+        ([id, { generation }]) => ({ id, generation }),
+      ),
+    });
   }
 
   public __internal_useHost(parentClient: AssistantClient) {
-    const { threads, hookEpoch } = this.hostStore();
-    const adapters = this.adapterStore();
+    const { threads, hookEpoch } = useSubscribable(this.hostStore);
+    const adapters = useSubscribable(this.adapterStore);
     const runtimeHook = this.runtimeHook;
     return useResources(
       threads.map(({ id, generation }) => {
@@ -294,7 +335,7 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
   public __internal_RenderThreadRuntimes: FC<{
     provider: ComponentType<PropsWithChildren>;
   }> = ({ provider }) => {
-    this.hostStore();
+    useSubscribable(this.hostStore);
 
     return Array.from(this.instances.entries()).map(
       ([threadId, { generation }]) => (

package/src/react/runtimes/RemoteThreadListThreadListRuntimeCore.switch-delete.test.ts

@@ -0,0 +1,169 @@
+import { describe, expect, it, vi } from "vitest";
+import { RemoteThreadListThreadListRuntimeCore } from "./RemoteThreadListThreadListRuntimeCore";
+import {
+  contextProvider,
+  createCore,
+  deferred,
+  makeAdapter,
+} from "../../tests/remote-thread-list-test-helpers";
+
+describe("RemoteThreadListThreadListRuntimeCore switch/delete ordering", () => {
+  it("rejects when a thread is deleted before its runtime attaches", async () => {
+    const adapter = makeAdapter({
+      list: vi.fn(async () => ({
+        threads: [
+          {
+            status: "archived" as const,
+            remoteId: "thread-b",
+            externalId: "thread-b",
+            title: "Thread B",
+          },
+        ],
+      })),
+    });
+    const core = new RemoteThreadListThreadListRuntimeCore(
+      { adapter, runtimeHook: () => ({}) as never },
+      contextProvider,
+    );
+    await core.getLoadThreadsPromise();
+    const initialMainThreadId = core.mainThreadId;
+
+    const switchToB = core.switchToThread("thread-b");
+    const rejection = expect(switchToB).rejects.toThrow(
+      "Thread was deleted before runtime was started",
+    );
+    await core.delete("thread-b");
+    await rejection;
+
+    expect(adapter.unarchive).not.toHaveBeenCalled();
+    expect(core.getItemById("thread-b")).toBeUndefined();
+    expect(core.mainThreadId).toBe(initialMainThreadId);
+  });
+
+  it("does not select a thread deleted while its switch is unarchiving", async () => {
+    const unarchive = deferred<void>();
+    const adapter = makeAdapter({
+      list: vi.fn(async () => ({
+        threads: [
+          {
+            status: "archived" as const,
+            remoteId: "thread-b",
+            externalId: "thread-b",
+            title: "Thread B",
+          },
+        ],
+      })),
+      unarchive: vi.fn(() => unarchive.promise),
+    });
+    const core = createCore(adapter);
+    await core.getLoadThreadsPromise();
+    const initialMainThreadId = core.mainThreadId;
+
+    const switchToB = core.switchToThread("thread-b");
+    await vi.waitFor(() => {
+      expect(adapter.unarchive).toHaveBeenCalledWith("thread-b");
+    });
+    await core.delete("thread-b");
+
+    unarchive.resolve();
+    await switchToB;
+
+    expect(core.getItemById("thread-b")).toBeUndefined();
+    expect(core.mainThreadId).toBe(initialMainThreadId);
+    expect(core.getItemById(core.mainThreadId!)).toBeDefined();
+  });
+
+  it("does not start unarchive after initialization when the target was deleted", async () => {
+    const initialization = deferred<{
+      remoteId: string;
+      externalId: string | undefined;
+    }>();
+    const adapter = makeAdapter({
+      list: vi.fn(async () => ({
+        threads: [
+          {
+            status: "archived" as const,
+            remoteId: "thread-b",
+            externalId: "thread-b",
+            title: "Thread B",
+          },
+        ],
+      })),
+    });
+    const core = createCore(adapter);
+    await core.getLoadThreadsPromise();
+    const target = core.getItemById("thread-b")!;
+    (
+      target as {
+        initializeTask: Promise<{
+          remoteId: string;
+          externalId: string | undefined;
+        }>;
+      }
+    ).initializeTask = initialization.promise;
+    const initialMainThreadId = core.mainThreadId;
+
+    const switchToB = core.switchToThread("thread-b");
+    const deleteB = core.delete("thread-b");
+    await vi.waitFor(() => {
+      expect(core.getItemById("thread-b")).toBeUndefined();
+    });
+
+    initialization.resolve({
+      remoteId: "thread-b",
+      externalId: "thread-b",
+    });
+    await Promise.all([switchToB, deleteB]);
+
+    expect(adapter.unarchive).not.toHaveBeenCalled();
+    expect(core.mainThreadId).toBe(initialMainThreadId);
+    expect(core.getItemById(core.mainThreadId!)).toBeDefined();
+  });
+
+  it("does not unarchive again when the target became regular during initialization", async () => {
+    const initialization = deferred<{
+      remoteId: string;
+      externalId: string | undefined;
+    }>();
+    const adapter = makeAdapter({
+      list: vi.fn(async () => ({
+        threads: [
+          {
+            status: "archived" as const,
+            remoteId: "thread-b",
+            externalId: "thread-b",
+            title: "Thread B",
+          },
+        ],
+      })),
+    });
+    const core = createCore(adapter);
+    await core.getLoadThreadsPromise();
+    const target = core.getItemById("thread-b")!;
+    (
+      target as {
+        initializeTask: Promise<{
+          remoteId: string;
+          externalId: string | undefined;
+        }>;
+      }
+    ).initializeTask = initialization.promise;
+
+    const switchToB = core.switchToThread("thread-b");
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    const unarchiveB = core.unarchive("thread-b");
+    await vi.waitFor(() => {
+      expect(core.getItemById("thread-b")?.status).toBe("regular");
+    });
+
+    initialization.resolve({
+      remoteId: "thread-b",
+      externalId: "thread-b",
+    });
+    await Promise.all([switchToB, unarchiveB]);
+
+    expect(adapter.unarchive).toHaveBeenCalledTimes(1);
+    expect(adapter.unarchive).toHaveBeenCalledWith("thread-b");
+    expect(core.mainThreadId).toBe("thread-b");
+  });
+});