@assistant-ui/core 0.3.12 → 0.3.14

package/src/runtime/utils/thread-runtime-lifecycle.ts

@@ -0,0 +1,18 @@
+import type { ThreadRuntimeCore } from "../interfaces/thread-runtime-core";
+
+// Invalidation must stay re-entrant: StrictMode's simulated unmount runs the
+// effect cleanup while the runtime object survives into the next mount, so a
+// permanent disposed mark would swallow every later append.
+const generations = new WeakMap<ThreadRuntimeCore, number>();
+
+export const captureThreadRuntimeGeneration = (runtime: ThreadRuntimeCore) =>
+  generations.get(runtime) ?? 0;
+
+export const isThreadRuntimeGenerationCurrent = (
+  runtime: ThreadRuntimeCore,
+  generation: number,
+) => captureThreadRuntimeGeneration(runtime) === generation;
+
+export const invalidateThreadRuntime = (runtime: ThreadRuntimeCore) => {
+  generations.set(runtime, captureThreadRuntimeGeneration(runtime) + 1);
+};

package/src/runtimes/external-store/external-store-adapter.ts

@@ -109,6 +109,13 @@ type ExternalStoreAdapterBase<T> = {
   state?: ReadonlyJSONValue | undefined;
   extras?: unknown;
 
+  /**
+   * Applies a message list the runtime rewrote, and is what tells the runtime
+   * a removal it makes will survive the next snapshot. Without it, cancelling a
+   * run leaves a trailing user message in the thread and the composer
+   * untouched; an adapter that removes that message itself owns handing it
+   * back, because the runtime cannot see a removal it did not make.
+   */
   setMessages?: ((messages: readonly T[]) => void) | undefined;
   /**
    * Fires when the user explicitly switches branches via the runtime's

package/src/runtimes/external-store/external-store-shared-options.ts

@@ -14,5 +14,7 @@ export const pickExternalStoreSharedOptions = (
     unstable_capabilities: options.unstable_capabilities,
     suggestions: options.suggestions,
   }) satisfies {
-    [K in keyof Required<ExternalStoreSharedOptions>]: ExternalStoreSharedOptions[K];
+    [
+      K in keyof Required<ExternalStoreSharedOptions>
+    ]: ExternalStoreSharedOptions[K];
   };

package/src/runtimes/external-store/external-store-thread-list-runtime-core.ts

@@ -5,6 +5,8 @@ import type {
   ThreadListRuntimeCore,
 } from "../../runtime/interfaces/thread-list-runtime-core";
 import type { ExternalStoreThreadListAdapter } from "./external-store-adapter";
+import { invalidateThreadRuntime } from "../../runtime/utils/thread-runtime-lifecycle";
+import { notifySubscribers } from "../../subscribable/subscribable";
 
 export type ExternalStoreThreadFactory = () => ExternalStoreThreadRuntimeCore;
 
@@ -150,6 +152,7 @@ export class ExternalStoreThreadListRuntimeCore implements ThreadListRuntimeCore
 
     // `initialLoad ||`: `_mainThread!` must be assigned on construction.
     if (initialLoad || previousThreadId !== newThreadId) {
+      if (!initialLoad) invalidateThreadRuntime(this._mainThread);
       this._mainThreadId = newThreadId;
       this._mainThread = this.threadFactory();
     }
@@ -268,6 +271,6 @@ export class ExternalStoreThreadListRuntimeCore implements ThreadListRuntimeCore
   }
 
   private _notifySubscribers() {
-    for (const callback of this._subscriptions) callback();
+    notifySubscribers(this._subscriptions);
   }
 }

package/src/runtimes/external-store/external-store-thread-runtime-core.ts

@@ -1,4 +1,5 @@
 import type { AppendMessage, ThreadMessage } from "../../types/message";
+import type { Attachment } from "../../types/attachment";
 import type {
   AddToolResultOptions,
   ResumeRunConfig,
@@ -27,7 +28,10 @@ import type {
   RuntimeCapabilities,
   ThreadRuntimeCore,
 } from "../../runtime/interfaces/thread-runtime-core";
-import type { QueuePlacement } from "../../runtime/queue/external-thread-queue-adapter";
+import type {
+  ExternalThreadQueueAdapter,
+  QueuePlacement,
+} from "../../runtime/queue/external-thread-queue-adapter";
 import { BaseThreadRuntimeCore } from "../../runtime/base/base-thread-runtime-core";
 import type { ModelContextProvider } from "../../model-context/types";
 import {
@@ -37,9 +41,26 @@ import {
 import { generateId } from "../../utils/id";
 import { ToolInvocationTracker } from "../tool-invocations/ToolInvocationTracker";
 import { EMPTY_QUEUE_ITEMS } from "../../store/scopes/queue-item";
+import type { QuoteInfo } from "../../types/quote";
+import {
+  captureThreadRuntimeGeneration,
+  isThreadRuntimeGenerationCurrent,
+} from "../../runtime/utils/thread-runtime-lifecycle";
 
 const EMPTY_ARRAY: readonly ThreadSuggestion[] = Object.freeze([]);
 
+const observeAdapterCallback = (
+  name: "onAddToolResult" | "onRespondToToolApproval" | "onCancel",
+  result: Promise<void> | void,
+) => {
+  void Promise.resolve(result).catch((error) => {
+    console.error(
+      `[ExternalStoreThreadRuntimeCore] ${name} callback rejected`,
+      error,
+    );
+  });
+};
+
 const shallowEqual = (a: object, b: object): boolean => {
   const aKeys = Object.keys(a);
   if (aKeys.length !== Object.keys(b).length) return false;
@@ -117,6 +138,16 @@ export class ExternalStoreThreadRuntimeCore
 
   private _store!: ExternalStoreAdapter<any>;
 
+  private _getInitializePromise?: () => Promise<unknown> | undefined;
+
+  public __internal_setGetInitializePromise(
+    getPromise: () => Promise<unknown> | undefined,
+  ) {
+    this._getInitializePromise = getPromise;
+  }
+
+  private _transformedQueue: ExternalThreadQueueAdapter | undefined;
+
   /**
    * Client-side tool-invocations pipeline. Constructed lazily on first
    * snapshot — only when `adapter.unstable_enableToolInvocations === true`.
@@ -147,6 +178,19 @@ export class ExternalStoreThreadRuntimeCore
 
     const oldStore = this._store as ExternalStoreAdapter<any> | undefined;
     this._store = store;
+    if (oldStore?.queue !== store.queue) {
+      this._transformedQueue = undefined;
+      store.queue?.__internal_setDispatchTransform?.((message) => {
+        // Re-point at the tail, as LocalThreadRuntimeCore's driver does, so
+        // the prefix gated against is the one the message lands on whatever
+        // the host routes by. Queuing only ever accepts a tail append, so a
+        // later tail is the same intent.
+        const parentId = this.messages.at(-1)?.id ?? null;
+        return this.enrichAppendMetadata({ ...message, parentId }, parentId);
+      });
+      if (store.queue?.__internal_setDispatchTransform)
+        this._transformedQueue = store.queue;
+    }
     if (this.extras !== store.extras) {
       this.extras = store.extras;
     }
@@ -355,19 +399,22 @@ export class ExternalStoreThreadRuntimeCore
                 // rolled back). Drop the result.
                 return;
               }
-              this._store.onAddToolResult?.({
-                messageId,
-                toolCallId: command.toolCallId,
-                toolName: command.toolName,
-                result: command.result,
-                isError: command.isError,
-                ...(command.artifact !== undefined && {
-                  artifact: command.artifact,
-                }),
-                ...(command.modelContent !== undefined && {
-                  modelContent: command.modelContent,
+              observeAdapterCallback(
+                "onAddToolResult",
+                this._store.onAddToolResult?.({
+                  messageId,
+                  toolCallId: command.toolCallId,
+                  toolName: command.toolName,
+                  result: command.result,
+                  isError: command.isError,
+                  ...(command.artifact !== undefined && {
+                    artifact: command.artifact,
+                  }),
+                  ...(command.modelContent !== undefined && {
+                    modelContent: command.modelContent,
+                  }),
                 }),
-              });
+              );
             } catch (err) {
               console.error(
                 "[ExternalStoreThreadRuntimeCore] onAddToolResult dispatch failed",
@@ -465,22 +512,54 @@ export class ExternalStoreThreadRuntimeCore
     });
   }
 
-  public async append(message: AppendMessage): Promise<void> {
+  public async append(rawMessage: AppendMessage): Promise<void> {
     // sourceId marks an edit send; the parent may coincide with the head
     // after a resync (e.g. cancelRun dropped the edited message).
     const isEdit =
-      message.sourceId != null ||
-      message.parentId !== (this.messages.at(-1)?.id ?? null);
-
-    // Buffering does not start a run, so the tool-abort below must wait until
-    // the queue flushes. By then the prior run (and its tools) has settled.
+      rawMessage.sourceId != null ||
+      rawMessage.parentId !== (this.messages.at(-1)?.id ?? null);
+
+    // A transformed-queue send is stamped at flush; any other queue's
+    // transform would gate against its own thread's messages, so those stamp
+    // at send.
+    const message =
+      !isEdit &&
+      this._store.queue &&
+      this._store.queue === this._transformedQueue
+        ? rawMessage
+        : this.enrichAppendMetadata(rawMessage);
+
+    const generation = captureThreadRuntimeGeneration(this);
+    this.ensureInitialized();
+
+    // The getter call is what starts thread initialization.
+    const initPromise = this._getInitializePromise?.();
+
+    // The queue driver dispatches through the host adapter, outside this
+    // core, so the initialization barrier must run before a message can
+    // enter the queue.
     if (!isEdit && this._store.queue) {
+      if (initPromise) {
+        await initPromise;
+      }
+      if (!isThreadRuntimeGenerationCurrent(this, generation)) return;
+
+      // Buffering does not start a run, so the tool-abort below must wait
+      // until the queue flushes. By then the prior run (and its tools) has
+      // settled.
       if (message.steer ?? this._store.isRunning ?? false)
         this._store.queue.steer(message);
       else this._store.queue.enqueue(message);
       return;
     }
 
+    // The optimistic insert lives inside the adapter's dispatch, so holding
+    // `onNew` on initialization would keep the message off screen for the
+    // whole roundtrip. Seams that need the remote identity await
+    // `threadListItem.initialize()` themselves, and a rejection surfaces
+    // there.
+    void initPromise?.catch(() => {});
+
     // Auto-abort in-flight client-side tool executions when a new run is
     // about to start. Without this, a tool that finishes after the new turn
     // begins would feed a stale result into `onAddToolResult`, racing with
@@ -490,6 +569,7 @@ export class ExternalStoreThreadRuntimeCore
     if (message.startRun ?? message.role === "user") {
       await this._toolInvocations?.abort();
     }
+    if (!isThreadRuntimeGenerationCurrent(this, generation)) return;
 
     if (isEdit) {
       if (!this._store.onEdit)
@@ -572,17 +652,24 @@ export class ExternalStoreThreadRuntimeCore
     // imported state) is treated as historical — no streamCall/execute
     // fires for the loaded tool calls. The adapter is expected to update
     // its messages in response to onLoadExternalState; that update flows
-    // back here via __internal_setAdapter. We only clear adapter-side
-    // tool statuses when the tracker is the source of truth — otherwise
-    // we'd wipe statuses the adapter is managing on its own.
-    if (this._toolInvocations) {
-      this._toolInvocations.reset();
-      this._store.setToolStatuses?.({});
-    }
+    // back here via __internal_setAdapter. The tracker publishes the
+    // cleared status map itself, so adapter-side statuses reset only when
+    // the tracker is the source of truth.
+    this._toolInvocations?.reset();
 
     this._store.onLoadExternalState(state);
   }
 
+  /**
+   * Adapter-facing notification that the backing session was discarded.
+   * Clears session-scoped tool-invocation state and parks queued work,
+   * without run-cancel semantics (`onCancel`, composer draft restoration).
+   */
+  public unstable_notifySessionReset(): void {
+    this._toolInvocations?.reset();
+    this._store.queue?.__internal_notifyCancelled?.();
+  }
+
   public cancelRun(): void {
     if (!this._store.onCancel)
       throw new Error("Runtime does not support cancelling runs.");
@@ -592,41 +679,90 @@ export class ExternalStoreThreadRuntimeCore
     // cancel on it.
     void this._toolInvocations?.abort();
 
-    this._store.onCancel();
+    // Before the run is aborted, so the settle it produces keeps the pending
+    // items instead of dispatching the next one at the moment the user
+    // stopped.
+    this._store.queue?.__internal_notifyCancelled?.();
 
-    // Drop an empty optimistic head (placeholder or pre-stream message); a
-    // partially-streamed one is kept and re-supplied by the store on resync.
-    const head = this.repository.getMessages().at(-1);
-    if (head && head.metadata.isOptimistic && head.content.length === 0) {
-      this.repository.deleteMessage(head.id);
-    }
+    observeAdapterCallback("onCancel", this._store.onCancel());
+
+    this.dropEmptyOptimisticHead();
 
-    let messages = this.repository.getMessages();
+    const messages = this.repository.getMessages();
     const previousMessage = messages[messages.length - 1];
-    if (
+    const trailingUserLeaf =
+      this._store.setMessages !== undefined &&
       previousMessage?.role === "user" &&
-      previousMessage.id === messages.at(-1)?.id // ensure the previous message is a leaf node
-    ) {
-      this.repository.deleteMessage(previousMessage.id);
-      if (!this.composer.text.trim()) {
-        this.composer.setText(getThreadMessageText(previousMessage));
+      previousMessage.id === messages.at(-1)?.id && // ensure the previous message is a leaf node
+      previousMessage.content.every((part) => part.type === "text")
+        ? previousMessage
+        : undefined;
+
+    // Handing the message to the composer and taking it out of the thread are
+    // one move: the composer refuses while the user is writing, and removing
+    // the message then would leave it nowhere. A message the composer cannot
+    // hold whole, carrying content parts it has no home for, is not moved.
+    let movedLeaf:
+      | {
+          id: string;
+          draft: {
+            text: string;
+            attachments: readonly Attachment[];
+            quote: QuoteInfo | undefined;
+          };
+        }
+      | undefined;
+    if (trailingUserLeaf) {
+      const draft = {
+        text: getThreadMessageText(trailingUserLeaf),
+        attachments: trailingUserLeaf.attachments,
+        quote: trailingUserLeaf.metadata.custom.quote as QuoteInfo | undefined,
+      };
+      if (this.composer.restoreDraft(draft)) {
+        this.repository.deleteMessage(trailingUserLeaf.id);
+        movedLeaf = { id: trailingUserLeaf.id, draft };
       }
-
-      messages = this.repository.getMessages();
-    } else {
-      this._notifySubscribers();
     }
+    if (!movedLeaf) this._notifySubscribers();
 
-    // resync messages (for reloading, to restore the previous branch)
+    // The resync commits what the cancel left (a kept optimistic message, the
+    // restored branch) back to the store a macrotask later. The store may move
+    // in that gap; a server settling the cancelled turn lands in the same
+    // tick. Read the repository at flush time and re-apply the rollbacks to
+    // it, instead of stamping a snapshot captured above over the newer state.
     setTimeout(() => {
-      this.updateMessages(messages);
+      this.dropEmptyOptimisticHead();
+      if (movedLeaf) {
+        const current = this.repository.getMessages();
+        if (current.at(-1)?.id === movedLeaf.id) {
+          // Unanswered tail: the removal has not reached the store yet.
+          this.repository.deleteMessage(movedLeaf.id);
+        } else if (current.some((m) => m.id === movedLeaf.id)) {
+          // The store kept the turn in the thread; take the untouched draft
+          // back so the same content does not sit in both places.
+          this.composer.retractDraft(movedLeaf.draft);
+        }
+      }
+      this.updateMessages(this.repository.getMessages());
     }, 0);
   }
 
+  // Placeholder or pre-stream message; a partially-streamed one is kept and
+  // committed to the store by the cancel resync.
+  private dropEmptyOptimisticHead(): void {
+    const head = this.repository.getMessages().at(-1);
+    if (head && head.metadata.isOptimistic && head.content.length === 0) {
+      this.repository.deleteMessage(head.id);
+    }
+  }
+
   public addToolResult(options: AddToolResultOptions) {
     if (!this._store.onAddToolResult)
       throw new Error("Runtime does not support tool results.");
-    this._store.onAddToolResult?.(options);
+    observeAdapterCallback(
+      "onAddToolResult",
+      this._store.onAddToolResult(options),
+    );
   }
 
   public resumeToolCall(options: ResumeToolCallOptions) {
@@ -652,7 +788,10 @@ export class ExternalStoreThreadRuntimeCore
   public respondToToolApproval(options: RespondToToolApprovalOptions) {
     if (!this._store.onRespondToToolApproval)
       throw new Error("Runtime does not support tool approvals.");
-    this._store.onRespondToToolApproval(options);
+    observeAdapterCallback(
+      "onRespondToToolApproval",
+      this._store.onRespondToToolApproval(options),
+    );
   }
 
   public override reset(initialMessages?: readonly ThreadMessageLike[]) {