@assistant-ui/core 0.3.19 → 0.3.21

package/src/react/runtimes/RemoteThreadListHookInstanceManager.tsx

@@ -10,7 +10,7 @@ import {
   type ComponentType,
   Fragment,
 } from "react";
-import { useResources, withKey } from "@assistant-ui/tap";
+import { useResources, useTapHost, withKey } from "@assistant-ui/tap";
 import type { AssistantClient } from "@assistant-ui/store";
 import { ThreadListItemRuntimeProvider } from "../providers/ThreadListItemRuntimeProvider";
 import type {
@@ -25,12 +25,16 @@ import type { Unsubscribe } from "../../types/unsubscribe";
 import {
   BaseSubscribable,
   notifySubscribers,
+  runCleanups,
   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 {
+  disposeThreadRuntime,
+  supersedeThreadRuntime,
+} from "../../runtime/utils/thread-runtime-lifecycle";
 import { notifyEventListeners } from "../../utils/notify-event-listeners";
 import {
   useRuntimeAdapters,
@@ -149,18 +153,21 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
     const instance = this.instances.get(threadId);
     if (!instance) return this.startThreadRuntime(threadId);
 
-    if (instance.runtime) invalidateThreadRuntime(instance.runtime);
-    // Detach before aborting, as stopThreadRuntime does: the abort runs the
-    // destroy listeners synchronously, and a listener that stops the outgoing
-    // runtime would otherwise emit that generation's terminal events through
-    // the subscription the next generation is about to reuse.
-    instance.unsubscribeRunning?.();
-    instance.unsubscribeRunning = undefined;
-    instance.destroy.abort();
-    instance.destroy = new AbortController();
-    instance.generation = this.nextGeneration++;
-    this._syncHostThreads();
-    this._notifySubscribers();
+    // Detach before superseding and aborting, as stopThreadRuntime does: both
+    // tear the outgoing runtime down synchronously, and its terminal events
+    // would otherwise reach the subscription the next generation is about to
+    // reuse.
+    try {
+      instance.unsubscribeRunning?.();
+    } finally {
+      instance.unsubscribeRunning = undefined;
+      if (instance.runtime) supersedeThreadRuntime(instance.runtime);
+      instance.destroy.abort();
+      instance.destroy = new AbortController();
+      instance.generation = this.nextGeneration++;
+      this._syncHostThreads();
+      this._notifySubscribers();
+    }
 
     return this._whenRuntimeAttached(threadId);
   }
@@ -214,12 +221,15 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
     const previousRuntime = instance.runtime;
     instance.runtime = runtime;
     instance.publishedGeneration = generation;
-    if (previousRuntime !== runtime) {
-      this._trackRunning(threadId, instance);
-    }
-    this._notifySubscribers();
-    if (previousRuntime !== undefined && previousRuntime !== runtime) {
-      notifySubscribers(this.replacedSubscribers);
+    try {
+      if (previousRuntime !== runtime) {
+        this._trackRunning(threadId, instance);
+      }
+    } finally {
+      this._notifySubscribers();
+      if (previousRuntime !== undefined && previousRuntime !== runtime) {
+        notifySubscribers(this.replacedSubscribers);
+      }
     }
   }
 
@@ -236,33 +246,33 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
     threadId: string,
     instance: RemoteThreadListHookInstance,
   ) {
-    instance.unsubscribeRunning?.();
-
-    const runtime = instance.runtime;
-    if (!runtime) {
-      instance.unsubscribeRunning = undefined;
-      this._setRunning(instance, false);
-      return;
-    }
-
-    this._setRunning(instance, getThreadRuntimeCoreIsRunning(runtime));
-    const unsubscribers = [
-      runtime.subscribe(() => {
+    const previousCleanup = instance.unsubscribeRunning;
+    instance.unsubscribeRunning = undefined;
+    try {
+      previousCleanup?.();
+    } finally {
+      const runtime = instance.runtime;
+      if (!runtime) {
+        this._setRunning(instance, false);
+      } else {
         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();
-    };
+        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 = () => runCleanups(unsubscribers);
+      }
+    }
   }
 
   private _setRunning(
@@ -276,13 +286,16 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
 
   public stopThreadRuntime(threadId: string) {
     const instance = this.instances.get(threadId);
-    if (instance?.runtime) invalidateThreadRuntime(instance.runtime);
-    instance?.unsubscribeRunning?.();
-    instance?.destroy.abort();
-    this.instances.delete(threadId);
-    this.pendingThreadAdapters.delete(threadId);
-    this._syncHostThreads();
-    this._notifySubscribers();
+    try {
+      instance?.unsubscribeRunning?.();
+    } finally {
+      if (instance?.runtime) disposeThreadRuntime(instance.runtime);
+      instance?.destroy.abort();
+      this.instances.delete(threadId);
+      this.pendingThreadAdapters.delete(threadId);
+      this._syncHostThreads();
+      this._notifySubscribers();
+    }
   }
 
   public setRuntimeHook(newRuntimeHook: RemoteThreadListHook) {
@@ -311,7 +324,14 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
 
   public __internal_dispose() {
     for (const threadId of [...this.instances.keys()]) {
-      this.stopThreadRuntime(threadId);
+      try {
+        this.stopThreadRuntime(threadId);
+      } catch (error) {
+        console.error(
+          "[assistant-ui] Thread runtime cleanup threw while stopping a thread",
+          error,
+        );
+      }
     }
   }
 
@@ -329,30 +349,37 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
     });
   }
 
+  private _threadElements(
+    parentClient: AssistantClient,
+    { threads, hookEpoch }: HostSnapshot,
+    adapters: AdapterSnapshot,
+  ) {
+    const runtimeHook = this.runtimeHook;
+    return threads.map(({ id, generation, destroySignal }) => {
+      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,
+          destroySignal,
+        }),
+      );
+    });
+  }
+
+  /** @deprecated Commits the hosted threads after descendant layout effects; render `__internal_Host` instead. */
   public __internal_useHost(parentClient: AssistantClient) {
-    const { threads, hookEpoch } = useSubscribable(this.hostStore);
+    const host = useSubscribable(this.hostStore);
     const adapters = useSubscribable(this.adapterStore);
-    const runtimeHook = this.runtimeHook;
-    return useResources(
-      threads.map(({ id, generation, destroySignal }) => {
-        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,
-            destroySignal,
-          }),
-        );
-      }),
-    );
+    return useResources(this._threadElements(parentClient, host, adapters));
   }
 
   public __internal_RenderThreadRuntimes: FC<{
@@ -374,7 +401,15 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
   public __internal_Host: FC<{ parentClient: AssistantClient }> = ({
     parentClient,
   }) => {
-    this.__internal_useHost(parentClient);
+    const host = useSubscribable(this.hostStore);
+    const adapters = useSubscribable(this.adapterStore);
+    const elements = this._threadElements(parentClient, host, adapters);
+    const { effects } = useTapHost(function RemoteThreadResources() {
+      return useResources(elements);
+    });
+    // Descendant layout effects may already dispatch to the thread resources,
+    // and tap commits a hosted resource only when its host effects run.
+    useLayoutEffect(effects);
     return null;
   };
 

package/src/react/runtimes/RemoteThreadListThreadListRuntimeCore.tsx

@@ -7,6 +7,7 @@ import {
   WritableSubscribable,
 } from "../../subscribable/subscribable";
 import { useSubscribable } from "../../store/runtime-clients/useSubscribable";
+import { handleThreadListAction } from "../../store/runtime-clients/handle-thread-list-action";
 import { nullProtoRecord } from "../../utils/record";
 import { OptimisticState } from "../../runtimes/remote-thread-list/optimistic-state";
 import { EMPTY_THREAD_CORE } from "../../runtimes/remote-thread-list/empty-thread-core";
@@ -461,8 +462,14 @@ export class RemoteThreadListThreadListRuntimeCore
       Object.values(nextState.threadData).map((item) => item.id),
     );
     for (const item of Object.values(state.threadData)) {
-      if (!nextIds.has(item.id)) {
+      if (nextIds.has(item.id)) continue;
+      try {
         this._hookManager.stopThreadRuntime(item.id);
+      } catch (error) {
+        console.error(
+          "[assistant-ui] Thread runtime cleanup threw while stopping a thread",
+          error,
+        );
       }
     }
     void this._hookManager.startThreadRuntime(this._mainThreadId).then(
@@ -769,8 +776,12 @@ export class RemoteThreadListThreadListRuntimeCore
 
   private _switchToThreadFromProp(threadId: string | undefined): Promise<void> {
     return threadId !== undefined
-      ? this._startSwitchToThread(threadId, undefined, false)
-      : this._startSwitchToNewThread(false);
+      ? handleThreadListAction("switch", () =>
+          this._startSwitchToThread(threadId, undefined, false),
+        )
+      : handleThreadListAction("create", () =>
+          this._startSwitchToNewThread(false),
+        );
   }
 
   private _startSwitchToNewThread(emitThreadIdChange: boolean): Promise<void> {
@@ -1012,6 +1023,11 @@ export class RemoteThreadListThreadListRuntimeCore
     let lastAwaitedTask: Promise<void> | undefined;
 
     while (threadId === this._mainThreadId) {
+      // Rechecked each pass: the draft can become the new thread again
+      // mid-loop when its failed first save rolls it back, and switching to a
+      // new thread then re-adopts it, so no switch can move main off it.
+      if (threadId === this.newThreadId)
+        throw new Error("Cannot ensure new thread is not main");
       let switchTask = this._switchTask;
       const startedFallback = !switchTask || switchTask === lastAwaitedTask;
       if (startedFallback) switchTask = this.switchToNewThread();

package/src/react/runtimes/RemoteThreadResource.ts

@@ -106,15 +106,18 @@ const useRemoteThreadBinder = ({
 
   const initPromiseRef = useRef<Promise<unknown> | undefined>(undefined);
   const hasInitializedRef = useRef(false);
+  // Any caller's initialize() moves the item off "new"; a thread born "new"
+  // here still joins that initialization so its title arms.
+  const bornNewRef = useRef(itemRuntime.getState().status === "new");
   const titleDisposeRef = useRef<(() => void) | undefined>(undefined);
   const titleAliveRef = useRef(false);
 
   const handleInitialize = useEffectEvent(() => {
     if (hasInitializedRef.current) return;
 
-    const state = itemRuntime.getState();
-    if (state.status !== "new") return;
+    if (itemRuntime.getState().status !== "new" && !bornNewRef.current) return;
     hasInitializedRef.current = true;
+    bornNewRef.current = false;
 
     const initPromise = itemRuntime.initialize();
     initPromiseRef.current = initPromise;

package/src/react/runtimes/assistant-transport/replayBoundaryStream.test.ts

@@ -322,6 +322,7 @@ describe("createReplayBoundaryStream", () => {
 
   it("releases the response body reader after a read failure", async () => {
     const { waitForRender, releaseNext } = createRenderWait();
+    const setReplaying = vi.fn();
     const error = new Error("stream failed");
     const body = new ReadableStream<Uint8Array>({
       start(controller) {
@@ -330,13 +331,14 @@ describe("createReplayBoundaryStream", () => {
     });
     const streamPromise = createReplayBoundaryStream(
       new Response(body, { headers: { [REPLAY_CONTENT_LENGTH_HEADER]: "10" } }),
-      { setReplaying: vi.fn(), waitForRender },
+      { setReplaying, waitForRender },
     );
 
     await releaseNext();
     const stream = await streamPromise;
 
     await expect(stream.getReader().read()).rejects.toBe(error);
+    expect(setReplaying).toHaveBeenLastCalledWith(false);
     expect(body.locked).toBe(false);
   });
 

package/src/react/runtimes/assistant-transport/replayBoundaryStream.ts

@@ -90,8 +90,15 @@ export const createReplayBoundaryStream = async (
   const reader = body.getReader();
   let bytesForwarded = 0;
   let replayFinished = false;
+  let replayCleared = false;
   let readerCleanup: Promise<void> | undefined;
 
+  const clearReplay = () => {
+    if (replayCleared) return;
+    replayCleared = true;
+    setReplaying(false);
+  };
+
   const releaseReader = () => {
     if (readerCleanup) return readerCleanup;
     reader.releaseLock();
@@ -117,7 +124,7 @@ export const createReplayBoundaryStream = async (
 
     // Let replay bytes drain before rendering live mode, then render live mode before releasing live bytes.
     await waitForReplayRender();
-    setReplaying(false);
+    clearReplay();
     await waitForReplayRender();
   };
 
@@ -158,6 +165,7 @@ export const createReplayBoundaryStream = async (
         await finishReplay();
         controller.enqueue(value.subarray(replayBytesInChunk));
       } catch (error) {
+        clearReplay();
         await cancelReader(error).catch(() => {});
         throw error;
       }
@@ -165,7 +173,7 @@ export const createReplayBoundaryStream = async (
     async cancel(reason) {
       const wasFinished = replayFinished;
       replayFinished = true;
-      if (!wasFinished) setReplaying(false);
+      if (!wasFinished) clearReplay();
       await cancelReader(reason);
     },
   });

package/src/react/runtimes/assistant-transport/runManager.ts

@@ -49,6 +49,9 @@ export function useRunManager(config: {
       try {
         if (!disposeAborted()) {
           await onRunRef.current(ac.signal);
+          // A fully received body is not errored by abort(), so a cancelled
+          // run can still resolve.
+          if (ac.signal.aborted) throw ac.signal.reason;
         }
       } catch (error) {
         if (!disposeAborted() && !stateRef.current.disposed) {

package/src/react/runtimes/assistant-transport/transport-parentid.test.tsx

@@ -55,6 +55,14 @@ const setupRuntime = () => {
   return { App, fetchMock, requestBodies, runtimeRef };
 };
 
+const toolResult = (toolCallId: string): AssistantTransportCommand => ({
+  type: "add-tool-result",
+  toolCallId,
+  toolName: "tool",
+  result: {},
+  isError: false,
+});
+
 describe("assistant transport parentId lifetime", () => {
   afterEach(() => {
     vi.unstubAllGlobals();
@@ -94,4 +102,53 @@ describe("assistant transport parentId lifetime", () => {
     await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2));
     expect(Object.hasOwn(requestBodies[1]!, "parentId")).toBe(false);
   });
+
+  it.each(["cancel", "error"] as const)(
+    "drops the parentId of a message discarded by a run's %s",
+    async (outcome) => {
+      const { App, fetchMock, requestBodies, runtimeRef } = setupRuntime();
+      let settleFirstRun!: () => void;
+      fetchMock.mockImplementationOnce(
+        (_url: RequestInfo | URL, init?: RequestInit) => {
+          requestBodies.push(JSON.parse(init!.body as string));
+          return new Promise<Response>((_resolve, reject) => {
+            settleFirstRun = () => reject(new Error("network down"));
+            init!.signal!.addEventListener(
+              "abort",
+              () => reject(init!.signal!.reason),
+              { once: true },
+            );
+          });
+        },
+      );
+
+      await act(async () => {
+        render(<App />);
+      });
+      await waitFor(() => expect(runtimeRef.current).not.toBeNull());
+      const extras = () =>
+        runtimeRef.current!.thread.getState().extras as {
+          sendCommand: (command: AssistantTransportCommand) => void;
+        };
+
+      act(() => extras().sendCommand(toolResult("t0")));
+      await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(1));
+      await act(async () => {
+        void runtimeRef.current!.thread.append("m2");
+      });
+      await act(async () => {
+        if (outcome === "cancel") runtimeRef.current!.thread.cancelRun();
+        else settleFirstRun();
+      });
+      await waitFor(() =>
+        expect(runtimeRef.current!.thread.getState().isRunning).toBe(false),
+      );
+
+      act(() => extras().sendCommand(toolResult("t1")));
+
+      await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2));
+      expect(requestBodies[1]!["commands"]).toEqual([toolResult("t1")]);
+      expect(Object.hasOwn(requestBodies[1]!, "parentId")).toBe(false);
+    },
+  );
 });

package/src/react/runtimes/assistant-transport/transport-scheduling.test.ts

@@ -1,8 +1,15 @@
 // @vitest-environment jsdom
 
-import { act, render, renderHook, waitFor } from "@testing-library/react";
+import { act, renderHook, waitFor } from "@testing-library/react";
 import { describe, expect, it, vi } from "vitest";
-import { createElement, StrictMode, useLayoutEffect, useRef } from "react";
+import {
+  createElement,
+  StrictMode,
+  useEffect,
+  useLayoutEffect,
+  useRef,
+} from "react";
+import { createRoot } from "react-dom/client";
 import { useCommandQueue } from "./commandQueue";
 import { useRunManager } from "./runManager";
 import type { AssistantTransportCommand } from "./types";
@@ -85,52 +92,70 @@ const useTransportSchedulingHarness = (
 };
 
 describe("assistant transport scheduling contracts", () => {
-  it("uses current callbacks when scheduled by a descendant layout effect", async () => {
-    const onRunA = vi.fn(async () => {});
-    const onRunB = vi.fn(async () => {});
-    // act flushes passive effects before real microtasks, so preserve the
-    // browser's microtask-before-passive-effect ordering at the layout boundary.
-    const queueMicrotaskSpy = vi
-      .spyOn(globalThis, "queueMicrotask")
-      .mockImplementation((callback) => callback());
-
-    const Scheduler = ({
-      schedule,
-      enabled,
-    }: {
-      schedule: () => void;
-      enabled: boolean;
-    }) => {
-      useLayoutEffect(() => {
-        if (enabled) schedule();
-      }, [enabled, schedule]);
-      return null;
-    };
-    const Probe = ({
-      enabled,
-      onRun,
-    }: {
-      enabled: boolean;
-      onRun: (signal: AbortSignal) => Promise<void>;
-    }) => {
-      const runManager = useRunManager({ onRun });
-      return createElement(Scheduler, {
-        enabled,
-        schedule: runManager.schedule,
-      });
-    };
-
+  it("reads the committed callbacks when a run settles inside a yielded commit", async () => {
+    // A render past the scheduler's 5 ms frame budget yields before passive
+    // effects, and act would drain them first, so this renders outside act.
+    vi.stubGlobal("IS_REACT_ACT_ENVIRONMENT", false);
     try {
-      const view = render(
-        createElement(Probe, { enabled: false, onRun: onRunA }),
+      const events: string[] = [];
+      const pendingRuns: (() => void)[] = [];
+      let schedule!: () => void;
+      const settleRun = () => {
+        pendingRuns.shift()?.();
+        schedule();
+      };
+
+      const Settle = ({ onLayout }: { onLayout: (() => void) | undefined }) => {
+        useLayoutEffect(() => {
+          onLayout?.();
+        }, [onLayout]);
+        return null;
+      };
+      const Probe = ({
+        label,
+        renderMs,
+        onLayout,
+      }: {
+        label: string;
+        renderMs: number;
+        onLayout?: () => void;
+      }) => {
+        const runManager = useRunManager({
+          onRun: () => {
+            events.push(`run:${label}`);
+            return new Promise<void>((resolve) => pendingRuns.push(resolve));
+          },
+          onFinish: () => events.push(`finish:${label}`),
+        });
+        schedule = runManager.schedule;
+        useEffect(() => {
+          events.push(`passive:${label}`);
+        }, [label]);
+        const renderEnd = performance.now() + renderMs;
+        while (performance.now() < renderEnd) {}
+        return createElement(Settle, { onLayout });
+      };
+
+      const root = createRoot(document.createElement("div"));
+      root.render(createElement(Probe, { label: "A", renderMs: 0 }));
+      await vi.waitFor(() => expect(events).toEqual(["passive:A"]));
+      schedule();
+      await vi.waitFor(() => expect(events).toEqual(["passive:A", "run:A"]));
+      root.render(
+        createElement(Probe, { label: "B", renderMs: 30, onLayout: settleRun }),
       );
-      view.rerender(createElement(Probe, { enabled: true, onRun: onRunB }));
-      await act(async () => {});
-
-      expect(onRunB).toHaveBeenCalledTimes(1);
-      expect(onRunA).not.toHaveBeenCalled();
+      await vi.waitFor(() => expect(events).toContain("passive:B"));
+      root.unmount();
+
+      expect(events).toEqual([
+        "passive:A",
+        "run:A",
+        "finish:B",
+        "run:B",
+        "passive:B",
+      ]);
     } finally {
-      queueMicrotaskSpy.mockRestore();
+      vi.unstubAllGlobals();
     }
   });
 
@@ -196,7 +221,9 @@ describe("assistant transport scheduling contracts", () => {
         onRun: () => {
           throw new Error("network error");
         },
-        onError: (commands) => seen.push(commands),
+        onError: (commands) => {
+          seen.push(commands);
+        },
       }),
     );