@assistant-ui/core 0.3.3 → 0.3.4

package/src/react/runtimes/RemoteThreadListHookInstanceManager.test.tsx

@@ -16,3 +16,163 @@ describe("RemoteThreadListHookInstanceManager", () => {
     );
   });
 });
+
+describe("RemoteThreadListHookInstanceManager.__internal_restartThreadRuntime", () => {
+  const makeManager = () =>
+    new RemoteThreadListHookInstanceManager(
+      () => ({}) as never,
+      {} as ThreadListRuntimeCore,
+    );
+
+  // no React binder attaches a runtime in these tests, so the returned promises
+  // stay pending or reject on stop; neither is what is under test here
+  const start = (manager: RemoteThreadListHookInstanceManager, id: string) => {
+    manager.startThreadRuntime(id).catch(() => {});
+  };
+  const restart = (
+    manager: RemoteThreadListHookInstanceManager,
+    id: string,
+  ) => {
+    manager.__internal_restartThreadRuntime(id).catch(() => {});
+  };
+
+  type InstanceInternals = {
+    instances: Map<
+      string,
+      {
+        runtime?: unknown;
+        publishedGeneration?: number;
+        generation: number;
+      }
+    >;
+    _notifySubscribers: () => void;
+  };
+  const internalsOf = (manager: RemoteThreadListHookInstanceManager) =>
+    manager as unknown as InstanceInternals;
+
+  const renderedKeys = (manager: RemoteThreadListHookInstanceManager) =>
+    Array.from(internalsOf(manager).instances.entries()).map(
+      ([id, { generation }]) => `${id}:${generation}`,
+    );
+
+  it("changes the binder key so React remounts the runtime hook", () => {
+    const manager = makeManager();
+    start(manager, "thread-1");
+    const before = renderedKeys(manager);
+
+    restart(manager, "thread-1");
+
+    expect(renderedKeys(manager)).not.toEqual(before);
+    expect(renderedKeys(manager)).toEqual(["thread-1:1"]);
+  });
+
+  it("keeps the thread rendered across the restart, unlike stop", () => {
+    const restarted = makeManager();
+    start(restarted, "thread-1");
+    restart(restarted, "thread-1");
+
+    const stopped = makeManager();
+    start(stopped, "thread-1");
+    stopped.stopThreadRuntime("thread-1");
+
+    expect(renderedKeys(restarted)).toHaveLength(1);
+    expect(renderedKeys(stopped)).toHaveLength(0);
+  });
+
+  it("starts the runtime when the thread is not alive yet", () => {
+    const manager = makeManager();
+
+    restart(manager, "thread-1");
+
+    expect(renderedKeys(manager)).toEqual(["thread-1:0"]);
+  });
+
+  it("stop then start in one tick yields a fresh key, so the old still-mounted binder cannot satisfy the new start", () => {
+    const manager = makeManager();
+    start(manager, "thread-1");
+    const before = renderedKeys(manager);
+
+    manager.stopThreadRuntime("thread-1");
+    start(manager, "thread-1");
+
+    expect(renderedKeys(manager)).not.toEqual(before);
+    expect(renderedKeys(manager)).toEqual(["thread-1:1"]);
+  });
+
+  // simulates the binder's updateRuntime: publish a runtime for the
+  // instance's current generation (or an explicit stale one)
+  const publish = (
+    manager: RemoteThreadListHookInstanceManager,
+    id: string,
+    runtime: unknown,
+    options?: { generation?: number },
+  ) => {
+    const instance = internalsOf(manager).instances.get(id)!;
+    instance.runtime = runtime;
+    instance.publishedGeneration = options?.generation ?? instance.generation;
+    internalsOf(manager)._notifySubscribers();
+  };
+
+  it("does not settle with the pre-restart runtime; only the incoming binder's publication resolves it", async () => {
+    const manager = makeManager();
+    start(manager, "thread-1");
+    publish(manager, "thread-1", { tag: "pre-reload-runtime" });
+
+    let settledWith = "NOT_SETTLED";
+    manager.__internal_restartThreadRuntime("thread-1").then((r) => {
+      settledWith = (r as { tag: string }).tag;
+    });
+
+    // the outgoing runtime rides across the restart (stays readable)…
+    expect(manager.getThreadRuntimeCore("thread-1")).toEqual({
+      tag: "pre-reload-runtime",
+    });
+    // …but must not count as attached
+    await Promise.resolve();
+    await Promise.resolve();
+    await Promise.resolve();
+    expect(settledWith).toBe("NOT_SETTLED");
+
+    publish(manager, "thread-1", { tag: "post-reload-runtime" });
+    await Promise.resolve();
+    expect(settledWith).toBe("post-reload-runtime");
+  });
+
+  it("keeps the outgoing runtime readable while the restart promise is pending", () => {
+    const manager = makeManager();
+    start(manager, "thread-1");
+    publish(manager, "thread-1", { tag: "pre-reload-runtime" });
+
+    restart(manager, "thread-1");
+
+    expect(manager.getThreadRuntimeCore("thread-1")).toEqual({
+      tag: "pre-reload-runtime",
+    });
+  });
+
+  it("a stale-generation publication does not resolve the restart promise", async () => {
+    const manager = makeManager();
+    start(manager, "thread-1");
+    publish(manager, "thread-1", { tag: "pre-reload-runtime" });
+
+    let settled = false;
+    manager.__internal_restartThreadRuntime("thread-1").then(() => {
+      settled = true;
+    });
+
+    // an outgoing binder re-publishing for its old generation (e.g. a late
+    // outerSubscribe callback) must not count as the new attachment
+    const staleGeneration =
+      internalsOf(manager).instances.get("thread-1")!.generation - 1;
+    publish(
+      manager,
+      "thread-1",
+      { tag: "stale-publication" },
+      { generation: staleGeneration },
+    );
+
+    await Promise.resolve();
+    await Promise.resolve();
+    expect(settled).toBe(false);
+  });
+});

package/src/react/runtimes/RemoteThreadListHookInstanceManager.tsx

@@ -18,14 +18,25 @@ import { ThreadListItemRuntimeProvider } from "../providers/ThreadListItemRuntim
 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 type { ThreadRuntimeImpl } from "../../runtime/api/thread-runtime";
+import {
+  getThreadRuntimeCoreIsRunning,
+  type ThreadRuntimeImpl,
+} from "../../runtime/api/thread-runtime";
 import { ThreadListRuntimeImpl } from "../../runtime/api/thread-list-runtime";
 
 type RemoteThreadListHook = () => AssistantRuntime;
 
 type RemoteThreadListHookInstance = {
-  runtime?: ThreadRuntimeCore;
+  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;
 };
 
 const ProviderRenderDetector: FC<{
@@ -41,6 +52,9 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
     StoreApi<{ useRuntime: RemoteThreadListHook }>
   >;
   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;
 
@@ -53,20 +67,18 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
     this.useRuntimeHook = create(() => ({ useRuntime: runtimeHook }));
   }
 
-  public startThreadRuntime(threadId: string) {
-    if (!this.instances.has(threadId)) {
-      this.instances.set(threadId, {});
-      this.useAliveThreadsKeysChanged.setState({}, true);
-    }
-
+  private _whenRuntimeAttached(threadId: string) {
     return new Promise<ThreadRuntimeCore>((resolve, reject) => {
       const callback = () => {
         const instance = this.instances.get(threadId);
         if (!instance) {
           dispose();
           reject(new Error("Thread was deleted before runtime was started"));
-        } else if (!instance.runtime) {
-          return; // misc update
+        } else if (
+          !instance.runtime ||
+          instance.publishedGeneration !== instance.generation
+        ) {
+          return; // not yet published by the current generation's binder
         } else {
           dispose();
           resolve(instance.runtime);
@@ -77,13 +89,104 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
     });
   }
 
+  public startThreadRuntime(threadId: string) {
+    if (!this.instances.has(threadId)) {
+      this.instances.set(threadId, {
+        generation: this.nextGeneration++,
+        isRunning: false,
+      });
+      this.useAliveThreadsKeysChanged.setState({}, true);
+    }
+
+    return this._whenRuntimeAttached(threadId);
+  }
+
+  public __internal_restartThreadRuntime(threadId: string) {
+    const instance = this.instances.get(threadId);
+    if (!instance) return this.startThreadRuntime(threadId);
+
+    instance.generation = this.nextGeneration++;
+    this.useAliveThreadsKeysChanged.setState({}, true);
+    this._notifySubscribers();
+
+    return this._whenRuntimeAttached(threadId);
+  }
+
   public getThreadRuntimeCore(threadId: string) {
     const instance = this.instances.get(threadId);
     if (!instance) return undefined;
     return instance.runtime;
   }
 
+  public __internal_isThreadRunning(threadId: string) {
+    return this.instances.get(threadId)?.isRunning ?? false;
+  }
+
+  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 _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.`,
+      );
+
+    // 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;
+    instance.runtime = runtime;
+    instance.publishedGeneration = generation;
+    if (previousRuntime !== runtime) {
+      this._trackRunning(instance);
+    }
+    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?.();
+
+    const runtime = instance.runtime;
+    if (!runtime) {
+      instance.unsubscribeRunning = undefined;
+      this._setRunning(instance, false);
+      return;
+    }
+
+    this._setRunning(instance, getThreadRuntimeCoreIsRunning(runtime));
+    instance.unsubscribeRunning = runtime.subscribe(() => {
+      this._setRunning(instance, getThreadRuntimeCoreIsRunning(runtime));
+    });
+  }
+
+  private _setRunning(
+    instance: RemoteThreadListHookInstance,
+    isRunning: boolean,
+  ) {
+    if (instance.isRunning === isRunning) return;
+    instance.isRunning = isRunning;
+    for (const callback of this.runningSubscribers) callback();
+  }
+
   public stopThreadRuntime(threadId: string) {
+    this.instances.get(threadId)?.unsubscribeRunning?.();
     this.instances.delete(threadId);
     this.useAliveThreadsKeysChanged.setState({}, true);
     this._notifySubscribers();
@@ -98,10 +201,9 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
 
   // 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 }>> = ({
-    threadId,
-    children,
-  }) => {
+  private _RuntimeBinder: FC<
+    PropsWithChildren<{ threadId: string; generation: number }>
+  > = ({ threadId, generation, children }) => {
     const { useRuntime } = this.useRuntimeHook();
     const runtime = useRuntime();
 
@@ -109,15 +211,12 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
       .__internal_threadBinding;
 
     const updateRuntime = useCallback(() => {
-      const aliveThread = this.instances.get(threadId);
-      if (!aliveThread)
-        throw new Error(
-          `Thread "${threadId}" runtime binding not found. This is a bug in assistant-ui.`,
-        );
-
-      aliveThread.runtime = threadBinding.getState();
-      this._notifySubscribers();
-    }, [threadId, threadBinding]);
+      this._publishThreadRuntime(
+        threadId,
+        threadBinding.getState(),
+        generation,
+      );
+    }, [threadId, generation, threadBinding]);
 
     const isMounted = useRef(false);
     if (!isMounted.current) {
@@ -168,8 +267,9 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
 
   private _OuterActiveThreadProvider: FC<{
     threadId: string;
+    generation: number;
     provider: ComponentType<PropsWithChildren>;
-  }> = memo(({ threadId, provider: Provider }) => {
+  }> = memo(({ threadId, generation, provider: Provider }) => {
     const runtime = useMemo(
       () => new ThreadListRuntimeImpl(this.parent).getItemById(threadId),
       [threadId],
@@ -195,7 +295,7 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
     return (
       <ThreadListItemRuntimeProvider runtime={runtime}>
         <Provider>
-          <this._RuntimeBinder threadId={threadId}>
+          <this._RuntimeBinder threadId={threadId} generation={generation}>
             <ProviderRenderDetector detectorRef={detectorRef} />
           </this._RuntimeBinder>
         </Provider>
@@ -208,12 +308,15 @@ export class RemoteThreadListHookInstanceManager extends BaseSubscribable {
   }> = ({ provider }) => {
     this.useAliveThreadsKeysChanged(); // trigger re-render on alive threads change
 
-    return Array.from(this.instances.keys()).map((threadId) => (
-      <this._OuterActiveThreadProvider
-        key={threadId}
-        threadId={threadId}
-        provider={provider}
-      />
-    ));
+    return Array.from(this.instances.entries()).map(
+      ([threadId, { generation }]) => (
+        <this._OuterActiveThreadProvider
+          key={`${threadId}:${generation}`}
+          threadId={threadId}
+          generation={generation}
+          provider={provider}
+        />
+      ),
+    );
   };
 }

package/src/react/runtimes/RemoteThreadListThreadListRuntimeCore.tsx

@@ -185,6 +185,9 @@ export class RemoteThreadListThreadListRuntimeCore
       options.runtimeHook,
       this,
     );
+    this._hookManager.__internal_subscribeRunningChanged(() =>
+      this._notifySubscribers(),
+    );
     this.useProvider = create(() => ({
       Provider: options.adapter.unstable_Provider ?? Fragment,
     }));
@@ -245,6 +248,34 @@ export class RemoteThreadListThreadListRuntimeCore
     }
   }
 
+  public async reloadMainThread(): Promise<void> {
+    const threadId = this._mainThreadId;
+    if (threadId === undefined) return;
+
+    // An unsent thread holds no remote state, so a refetch would only discard
+    // what the user has typed.
+    if (this.getItemById(threadId)?.status === "new") return;
+
+    const runtimeCore = this._hookManager.getThreadRuntimeCore(threadId);
+
+    try {
+      if (runtimeCore?.unstable_refetchThread) {
+        // Called on the core so class-method implementations keep `this`.
+        await runtimeCore.unstable_refetchThread();
+      } else {
+        await this._hookManager.__internal_restartThreadRuntime(threadId);
+      }
+    } catch (error) {
+      // delete and detach switch the main thread away before stopping the
+      // runtime, so a rejection once that has happened belongs to them.
+      if (threadId !== this._mainThreadId) return;
+      throw error;
+    }
+
+    if (threadId !== this._mainThreadId) return;
+    this._notifySubscribers();
+  }
+
   public reload() {
     this._loadGeneration++;
     this._loadThreadsPromise = undefined;
@@ -321,6 +352,12 @@ export class RemoteThreadListThreadListRuntimeCore
     return result;
   }
 
+  public unstable_isThreadRunning(threadIdOrRemoteId: string) {
+    const data = this.getItemById(threadIdOrRemoteId);
+    if (!data) return false;
+    return this._hookManager.__internal_isThreadRunning(data.id);
+  }
+
   public getItemById(threadIdOrRemoteId: string) {
     return getThreadData(this._state.value, threadIdOrRemoteId);
   }

package/src/react/runtimes/cloud/auiV0.ts

@@ -77,6 +77,7 @@ type AuiV0MessagePart =
       readonly data: string;
       readonly mimeType: string;
       readonly filename?: string;
+      readonly sourceType?: "url" | "id";
     };
 
 type AuiV0AttachmentPart =
@@ -94,6 +95,7 @@ type AuiV0AttachmentPart =
       readonly data: string;
       readonly mimeType: string;
       readonly filename?: string;
+      readonly sourceType?: "url" | "id";
     }
   | {
       readonly type: "audio";
@@ -157,6 +159,9 @@ const encodeAttachmentPart = (
         data: part.data,
         mimeType: part.mimeType,
         ...(part.filename != null ? { filename: part.filename } : undefined),
+        ...(part.sourceType != null
+          ? { sourceType: part.sourceType }
+          : undefined),
       };
 
     case "audio":
@@ -252,7 +257,7 @@ export function auiV0Encode(message: ThreadMessage): AuiV0Message {
           };
 
         case "tool-call": {
-          if (!isJSONValue(part.result)) {
+          if (part.result !== undefined && !isJSONValue(part.result)) {
             console.warn(
               `tool-call result is not JSON! ${JSON.stringify(part)}`,
             );
@@ -264,7 +269,7 @@ export function auiV0Encode(message: ThreadMessage): AuiV0Message {
             ...(JSON.stringify(part.args) === part.argsText
               ? { args: part.args }
               : { argsText: part.argsText }),
-            ...(part.result
+            ...(part.result !== undefined
               ? { result: part.result as ReadonlyJSONValue }
               : undefined),
             ...(part.isError ? { isError: true } : undefined),
@@ -281,6 +286,7 @@ export function auiV0Encode(message: ThreadMessage): AuiV0Message {
             data: part.data,
             mimeType: part.mimeType,
             ...(part.filename ? { filename: part.filename } : undefined),
+            ...(part.sourceType ? { sourceType: part.sourceType } : undefined),
           };
 
         default: {

package/src/react/types/MessagePartComponentTypes.ts

@@ -44,8 +44,16 @@ export type ImageMessagePartComponent = ComponentType<ImageMessagePartProps>;
 export type FileMessagePartProps = MessagePartState & FileMessagePart;
 export type FileMessagePartComponent = ComponentType<FileMessagePartProps>;
 
+/**
+ * @deprecated Use {@link FileMessagePartProps} and render `audio/*` mime types
+ * from the `File` slot.
+ */
 export type Unstable_AudioMessagePartProps = MessagePartState &
   Unstable_AudioMessagePart;
+/**
+ * @deprecated Use {@link FileMessagePartComponent} and render `audio/*` mime
+ * types from the `File` slot.
+ */
 export type Unstable_AudioMessagePartComponent =
   ComponentType<Unstable_AudioMessagePartProps>;
 

package/src/runtime/api/bindings.ts

@@ -37,6 +37,11 @@ export type MessageStateBinding = SubscribableWithState<
 
 export type ThreadListItemState = {
   readonly isMain: boolean;
+  /**
+   * Whether this thread has a run in progress, including a run that continues
+   * after the user switches to another thread.
+   */
+  readonly isRunning: boolean;
   readonly id: string;
   readonly remoteId: string | undefined;
   readonly externalId: string | undefined;

package/src/runtime/api/thread-list-runtime.ts

@@ -31,7 +31,10 @@ export type ThreadListState = {
   readonly isLoadingMore: boolean;
   readonly hasMore: boolean;
   readonly threadItems: Readonly<
-    Record<string, Omit<ThreadListItemState, "isMain" | "threadId">>
+    Record<
+      string,
+      Omit<ThreadListItemState, "isMain" | "threadId" | "isRunning">
+    >
   >;
 };
 
@@ -56,6 +59,18 @@ export type ThreadListRuntime = {
 
   getLoadThreadsPromise(): Promise<void>;
   reload(): Promise<void>;
+  /**
+   * Refetches the open thread's remote state, for state that changed out of
+   * band and so never reached the stream. When the runtime declares the
+   * in-place capability (`unstable_refetchThread`), composer drafts survive,
+   * existing messages stay rendered during the refetch, and the promise
+   * settles with the refetch, rejecting if it fails; that runtime also owns
+   * what happens to a run in progress, since this does not stop one. Runtimes
+   * without the capability have their hook remounted instead, which discards
+   * unsent composer input and ends any run, and the promise resolves once the
+   * new runtime attaches. A thread that has not been sent yet is left alone.
+   */
+  reloadMainThread(): Promise<void>;
   loadMore(): Promise<void>;
 };
 
@@ -91,6 +106,7 @@ const getThreadListItemState = (
     lastMessageAt: threadData.lastMessageAt,
     custom: threadData.custom,
     isMain: threadData.id === threadList.mainThreadId,
+    isRunning: threadList.unstable_isThreadRunning?.(threadData.id) ?? false,
   };
 };
 
@@ -155,6 +171,7 @@ export class ThreadListRuntimeImpl implements ThreadListRuntime {
     this.switchToNewThread = this.switchToNewThread.bind(this);
     this.getLoadThreadsPromise = this.getLoadThreadsPromise.bind(this);
     this.reload = this.reload.bind(this);
+    this.reloadMainThread = this.reloadMainThread.bind(this);
     this.loadMore = this.loadMore.bind(this);
     this.getState = this.getState.bind(this);
     this.subscribe = this.subscribe.bind(this);
@@ -183,6 +200,10 @@ export class ThreadListRuntimeImpl implements ThreadListRuntime {
     return this._core.reload?.() ?? RESOLVED_PROMISE;
   }
 
+  public reloadMainThread(): Promise<void> {
+    return this._core.reloadMainThread?.() ?? RESOLVED_PROMISE;
+  }
+
   public loadMore(): Promise<void> {
     return this._core.loadMore?.() ?? RESOLVED_PROMISE;
   }

package/src/runtime/api/thread-runtime.ts

@@ -193,22 +193,31 @@ export type ThreadState = {
   readonly voice: VoiceSessionState | undefined;
 };
 
+/**
+ * The canonical `isRunning` derivation. A runtime that tracks run state itself
+ * reports it directly; the rest fall back to the trailing assistant message.
+ */
+export const getThreadRuntimeCoreIsRunning = (
+  runtime: ThreadRuntimeCore,
+): boolean => {
+  if (runtime.isRunning !== undefined) return runtime.isRunning;
+  const lastMessage = runtime.messages.at(-1);
+  return (
+    lastMessage?.role === "assistant" && lastMessage.status.type === "running"
+  );
+};
+
 export const getThreadState = (
   runtime: ThreadRuntimeCore,
   threadListItemState: ThreadListItemState,
 ): ThreadState => {
-  const lastMessage = runtime.messages.at(-1);
   return Object.freeze({
     threadId: threadListItemState.id,
     metadata: threadListItemState,
     capabilities: runtime.capabilities,
     isDisabled: runtime.isDisabled,
     isLoading: runtime.isLoading,
-    isRunning:
-      runtime.isRunning ??
-      (lastMessage?.role !== "assistant"
-        ? false
-        : lastMessage.status.type === "running"),
+    isRunning: getThreadRuntimeCoreIsRunning(runtime),
     messages: runtime.messages,
     state: runtime.state,
     suggestions: runtime.suggestions,