@assistant-ui/core 0.3.2 → 0.3.4

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

@@ -1,13 +1,8 @@
 import type {
-  DataMessagePart,
-  FileMessagePart,
-  ImageMessagePart,
   MessageStatus,
   SourceProviderMetadata,
   ThreadMessage,
-  TextMessagePart,
   ToolApprovalOption,
-  Unstable_AudioMessagePart,
 } from "../../../types/message";
 import type { CompleteAttachment } from "../../../types/attachment";
 import { fromThreadMessageLike } from "../../../runtime/utils/thread-message-like";
@@ -82,14 +77,38 @@ type AuiV0MessagePart =
       readonly data: string;
       readonly mimeType: string;
       readonly filename?: string;
+      readonly sourceType?: "url" | "id";
     };
 
 type AuiV0AttachmentPart =
-  | TextMessagePart
-  | ImageMessagePart
-  | FileMessagePart
-  | Unstable_AudioMessagePart
-  | DataMessagePart<ReadonlyJSONValue>;
+  | {
+      readonly type: "text";
+      readonly text: string;
+    }
+  | {
+      readonly type: "image";
+      readonly image: string;
+      readonly filename?: string;
+    }
+  | {
+      readonly type: "file";
+      readonly data: string;
+      readonly mimeType: string;
+      readonly filename?: string;
+      readonly sourceType?: "url" | "id";
+    }
+  | {
+      readonly type: "audio";
+      readonly audio: {
+        readonly data: string;
+        readonly format: "mp3" | "wav";
+      };
+    }
+  | {
+      readonly type: "data";
+      readonly name: string;
+      readonly data: ReadonlyJSONValue;
+    };
 
 type AuiV0Attachment = {
   readonly id: string;
@@ -125,16 +144,41 @@ const encodeAttachmentPart = (
   const type = part.type;
   switch (type) {
     case "text":
+      return { type: "text", text: part.text };
+
     case "image":
+      return {
+        type: "image",
+        image: part.image,
+        ...(part.filename != null ? { filename: part.filename } : undefined),
+      };
+
     case "file":
+      return {
+        type: "file",
+        data: part.data,
+        mimeType: part.mimeType,
+        ...(part.filename != null ? { filename: part.filename } : undefined),
+        ...(part.sourceType != null
+          ? { sourceType: part.sourceType }
+          : undefined),
+      };
+
     case "audio":
-      return part;
+      return {
+        type: "audio",
+        audio: { data: part.audio.data, format: part.audio.format },
+      };
 
     case "data": {
       if (!isJSONValue(part.data)) {
         console.warn(`attachment data is not JSON! ${JSON.stringify(part)}`);
       }
-      return { ...part, data: part.data as ReadonlyJSONValue };
+      return {
+        type: "data",
+        name: part.name,
+        data: part.data as ReadonlyJSONValue,
+      };
     }
 
     default: {
@@ -213,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)}`,
             );
@@ -225,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),
@@ -242,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/message-runtime.test.ts

@@ -1,8 +1,10 @@
 import { describe, expect, it } from "vitest";
 import type { CompleteAttachment } from "../../types/attachment";
+import type { ThreadAssistantMessage } from "../../types/message";
 import type { ThreadRuntimeCoreBinding } from "./thread-runtime";
 import {
   MessageRuntimeImpl,
+  toMessagePartStatus,
   type MessageState,
   type MessageStateBinding,
 } from "./message-runtime";
@@ -63,6 +65,144 @@ const threadBinding = {
   subscribe: () => () => {},
 } as unknown as ThreadRuntimeCoreBinding;
 
+const createAssistantMessage = (
+  content: ThreadAssistantMessage["content"],
+  status: ThreadAssistantMessage["status"],
+): ThreadAssistantMessage => ({
+  id: "message-1",
+  role: "assistant",
+  createdAt: new Date(0),
+  content,
+  status,
+  metadata: {
+    unstable_state: null,
+    unstable_annotations: [],
+    unstable_data: [],
+    steps: [],
+    custom: {},
+  },
+});
+
+describe("toMessagePartStatus", () => {
+  it("honours a supplied running status on a non-last part", () => {
+    const message = createAssistantMessage(
+      [
+        { type: "text", text: "first", status: { type: "running" } },
+        { type: "text", text: "last" },
+      ],
+      { type: "running" },
+    );
+
+    expect(toMessagePartStatus(message, 0, message.content[0]!)).toEqual({
+      type: "running",
+    });
+  });
+
+  it("honours a supplied complete status on the last part", () => {
+    const message = createAssistantMessage(
+      [{ type: "reasoning", text: "done", status: { type: "complete" } }],
+      { type: "running" },
+    );
+
+    expect(toMessagePartStatus(message, 0, message.content[0]!)).toEqual({
+      type: "complete",
+    });
+  });
+
+  it("ignores supplied statuses after the message completes", () => {
+    const message = createAssistantMessage(
+      [{ type: "text", text: "truncated", status: { type: "running" } }],
+      { type: "complete", reason: "stop" },
+    );
+
+    expect(toMessagePartStatus(message, 0, message.content[0]!)).toEqual({
+      type: "complete",
+      reason: "stop",
+    });
+  });
+
+  it("falls back to positional statuses for statusless running parts", () => {
+    const message = createAssistantMessage(
+      [
+        { type: "text", text: "first" },
+        { type: "reasoning", text: "last" },
+      ],
+      { type: "running" },
+    );
+
+    expect(toMessagePartStatus(message, 0, message.content[0]!)).toEqual({
+      type: "complete",
+    });
+    expect(toMessagePartStatus(message, 1, message.content[1]!)).toEqual({
+      type: "running",
+    });
+  });
+
+  it("preserves tool-call status derivation", () => {
+    const unresolved = createAssistantMessage(
+      [
+        {
+          type: "tool-call",
+          toolCallId: "call-1",
+          toolName: "weather",
+          args: {},
+          argsText: "{}",
+        },
+      ],
+      { type: "running" },
+    );
+    const resolved = createAssistantMessage(
+      [
+        {
+          type: "tool-call",
+          toolCallId: "call-1",
+          toolName: "weather",
+          args: {},
+          argsText: "{}",
+          result: "sunny",
+        },
+      ],
+      { type: "running" },
+    );
+
+    expect(toMessagePartStatus(unresolved, 0, unresolved.content[0]!)).toEqual({
+      type: "running",
+    });
+    expect(toMessagePartStatus(resolved, 0, resolved.content[0]!)).toEqual({
+      type: "complete",
+    });
+  });
+
+  it("normalizes supplied upstream statuses", () => {
+    const upstreamComplete = {
+      type: "text",
+      text: "done",
+      status: { type: "complete", reason: "unknown" },
+    } as unknown as ThreadAssistantMessage["content"][number];
+    const upstreamIncomplete = {
+      type: "reasoning",
+      text: "interrupted",
+      status: {
+        type: "incomplete",
+        reason: "unknown",
+        error: "upstream error",
+      },
+    } as unknown as ThreadAssistantMessage["content"][number];
+    const message = createAssistantMessage(
+      [upstreamComplete, upstreamIncomplete],
+      { type: "running" },
+    );
+
+    expect(toMessagePartStatus(message, 0, message.content[0]!)).toEqual({
+      type: "complete",
+    });
+    expect(toMessagePartStatus(message, 1, message.content[1]!)).toEqual({
+      type: "incomplete",
+      reason: "other",
+    });
+  });
+});
+
 describe("MessageRuntimeImpl paths", () => {
   it("appends nested selectors to the message path", () => {
     const runtime = new MessageRuntimeImpl(messageBinding, threadBinding);

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

@@ -8,6 +8,10 @@ import type {
 } from "../../types/message";
 import type { Unsubscribe } from "../../types/unsubscribe";
 import type { MessagePartStatus, RunConfig } from "../../types/message";
+import {
+  COMPLETE_STATUS,
+  normalizePartStatus,
+} from "../../utils/normalizePartStatus";
 import { getThreadMessageText } from "../../utils/text";
 import { NestedSubscriptionSubject } from "../../subscribable/subscribable";
 import {
@@ -32,10 +36,6 @@ import type { MessageRuntimePath } from "./paths";
 import type { ThreadRuntimeCoreBinding } from "./thread-runtime";
 import type { MessageStateBinding } from "./bindings";
 
-const COMPLETE_STATUS: MessagePartStatus = Object.freeze({
-  type: "complete",
-});
-
 export const toMessagePartStatus = (
   message: ThreadMessage,
   partIndex: number,
@@ -51,6 +51,11 @@ export const toMessagePartStatus = (
     }
   }
 
+  if (message.status.type === "running") {
+    const status = normalizePartStatus(part);
+    if (status) return status;
+  }
+
   const isLastPart = partIndex === Math.max(0, message.content.length - 1);
   if (message.status.type === "requires-action") return COMPLETE_STATUS;
   return isLastPart ? (message.status as MessagePartStatus) : COMPLETE_STATUS;
@@ -70,7 +75,7 @@ const getMessagePartState = (
   return Object.freeze({
     ...part,
     ...{ [symbolInnerMessage]: (part as any)[symbolInnerMessage] },
-    status,
+    status: status as MessagePartStatus,
   });
 };