@assistant-ui/core 0.3.11 → 0.3.13

package/src/runtime/base/base-composer-runtime-core.ts

@@ -6,6 +6,7 @@ import {
   type PendingAttachment,
 } from "../../types/attachment";
 import type { MessageRole, AppendMessage } from "../../types/message";
+import { isMessageNotSentError } from "../../types/error";
 import type { QuoteInfo } from "../../types/quote";
 import type { Unsubscribe } from "../../types/unsubscribe";
 import type { RunConfig } from "../../types/message";
@@ -118,15 +119,22 @@ export abstract class BaseComposerRuntimeCore
     if (this._text === value) return;
 
     this._text = value;
-    if (this._dictation) {
-      this._dictationBaseText = value;
-      this._currentInterimText = "";
-      const { status, inputDisabled } = this._dictation;
-      this._dictation = inputDisabled ? { status, inputDisabled } : { status };
-    }
+    this._rebaseDictation(value);
     this._notifySubscribers();
   }
 
+  // A live dictation session appends to the text it last saw, so any write
+  // that bypasses `setText` has to move that baseline or the next transcript
+  // overwrites what was just written.
+  private _rebaseDictation(value: string) {
+    if (!this._dictation) return;
+
+    this._dictationBaseText = value;
+    this._currentInterimText = "";
+    const { status, inputDisabled } = this._dictation;
+    this._dictation = inputDisabled ? { status, inputDisabled } : { status };
+  }
+
   public setRole(role: MessageRole) {
     if (this._role === role) return;
 
@@ -245,6 +253,7 @@ export abstract class BaseComposerRuntimeCore
       if (generation === this._sendGeneration) {
         if (!this.text.trim() && this._quote === undefined) {
           this._text = text;
+          this._rebaseDictation(text);
           this._quote = quote;
           this._notifySubscribers();
         }
@@ -289,11 +298,94 @@ export abstract class BaseComposerRuntimeCore
       metadata: { custom: { ...(quote ? { quote } : {}) } },
     };
 
-    const sendTask = this.handleSend(message, options);
-    if (sendTask) void sendTask.catch(() => {});
+    const draft = { text, quote, attachments: finalAttachments };
+    let sendTask: void | Promise<void>;
+    try {
+      sendTask = this.handleSend(message, options);
+    } catch (error) {
+      this._restoreUnsentDraft(error, generation, draft);
+      throw error;
+    }
+    if (sendTask)
+      void sendTask.catch((error) => {
+        this._restoreUnsentDraft(error, generation, draft);
+      });
     this._notifyEventSubscribers("send", {});
   }
 
+  /**
+   * Take a message back into the composer when it has nowhere else to live:
+   * a send the runtime never dispatched, or a message a cancelled run is
+   * removing from the thread. Reports whether the composer accepted it, so a
+   * caller that is also removing the message can keep it instead of dropping
+   * it. Refused, and left untouched, while the composer holds anything of its
+   * own.
+   */
+  public restoreDraft(draft: {
+    text: string;
+    quote?: QuoteInfo | undefined;
+    attachments?: readonly Attachment[] | undefined;
+  }): boolean {
+    if (
+      this._text.trim() ||
+      this._quote !== undefined ||
+      this._attachments.length > 0
+    )
+      return false;
+
+    this._text = draft.text;
+    this._rebaseDictation(draft.text);
+    this._quote = draft.quote;
+    this._attachments = draft.attachments ?? [];
+    this._notifySubscribers();
+    return true;
+  }
+
+  /**
+   * Inverse of `restoreDraft`: clears the composer while it still holds
+   * exactly the given draft. A draft the user has edited since is left
+   * untouched.
+   */
+  public retractDraft(draft: {
+    text: string;
+    quote?: QuoteInfo | undefined;
+    attachments?: readonly Attachment[] | undefined;
+  }): void {
+    const attachmentsUntouched =
+      draft.attachments !== undefined
+        ? this._attachments === draft.attachments
+        : this._attachments.length === 0;
+    if (
+      this._text !== draft.text ||
+      this._quote !== draft.quote ||
+      !attachmentsUntouched
+    )
+      return;
+
+    this._text = "";
+    this._rebaseDictation("");
+    this._quote = undefined;
+    this._attachments = [];
+    this._notifySubscribers();
+  }
+
+  // The generation check is what a reset and a later send use to invalidate a
+  // draft, so of several queued drafts only the most recent one is still
+  // restorable.
+  private _restoreUnsentDraft(
+    error: unknown,
+    generation: number,
+    draft: {
+      text: string;
+      quote: QuoteInfo | undefined;
+      attachments: readonly CompleteAttachment[];
+    },
+  ) {
+    if (!isMessageNotSentError(error)) return;
+    if (generation !== this._sendGeneration) return;
+    this.restoreDraft(draft);
+  }
+
   public cancel() {
     this.handleCancel();
   }

package/src/runtime/base/base-thread-runtime-core.ts

@@ -35,6 +35,7 @@ import type { AttachmentAdapter } from "../../adapters/attachment";
 import type { RealtimeVoiceAdapter } from "../../adapters/voice";
 import type { ThreadMessageLike } from "../utils/thread-message-like";
 import { notifyEventListeners } from "../../utils/notify-event-listeners";
+import { gateInteractableComposerMetadata } from "../../model-context/interactable-composer-metadata";
 
 type BaseThreadAdapters = {
   speech?: SpeechSynthesisAdapter | undefined;
@@ -123,6 +124,41 @@ export abstract class BaseThreadRuntimeCore implements ThreadRuntimeCore {
     return this._contextProvider.getModelContext();
   }
 
+  /**
+   * Stamps provider-contributed composer metadata onto an outgoing message.
+   * Called at dispatch rather than in the composer, so programmatic sends are
+   * covered too, and exactly once per message: a queued send is stamped when
+   * it leaves the lane, never when it enters.
+   *
+   * Only user messages are stamped, matching the readers: both the version
+   * fold and the model injection skip every other role.
+   *
+   * @param anchorId Message the gated branch prefix ends at. A queued send
+   * passes the current tail, having waited through a run that grew the prefix
+   * past the parent it was created with.
+   */
+  protected enrichAppendMetadata(
+    message: AppendMessage,
+    anchorId: string | null = message.parentId,
+  ): AppendMessage {
+    if (message.role !== "user") return message;
+    const messages = this.messages;
+    const parentIndex =
+      anchorId === null ? -1 : messages.findIndex((m) => m.id === anchorId);
+    const composerMetadata = gateInteractableComposerMetadata(
+      this.getModelContext().unstable_composerMetadata,
+      messages.slice(0, parentIndex + 1),
+    );
+    if (!composerMetadata) return message;
+    return {
+      ...message,
+      metadata: {
+        ...message.metadata,
+        custom: { ...message.metadata?.custom, ...composerMetadata },
+      },
+    };
+  }
+
   private _editComposers = new Map<string, DefaultEditComposerRuntimeCore>();
   public getEditComposer(messageId: string) {
     return this._editComposers.get(messageId);

package/src/runtime/base/default-edit-composer-runtime-core.ts

@@ -7,7 +7,6 @@ import type { DictationAdapter } from "../../adapters/speech";
 import type { SendOptions } from "../interfaces/composer-runtime-core";
 import type { ThreadRuntimeCore } from "../interfaces/thread-runtime-core";
 import { BaseComposerRuntimeCore } from "./base-composer-runtime-core";
-import { gateInteractableComposerMetadata } from "../../model-context/interactable-composer-metadata";
 
 export class DefaultEditComposerRuntimeCore extends BaseComposerRuntimeCore {
   public get canCancel() {
@@ -97,21 +96,8 @@ export class DefaultEditComposerRuntimeCore extends BaseComposerRuntimeCore {
             ...this._nonTextPassthrough,
           ] as AppendMessage["content"])
         : message.content;
-    // Gate live state against the new branch's prefix (messages up to the
-    // parent): an unchanged interactable re-stamps the prior baseline, an
-    // interactable edited since the original message stamps its newest state.
-    const messages = this.runtime.messages;
-    const parentIndex =
-      this._parentId === null
-        ? -1
-        : messages.findIndex((m) => m.id === this._parentId);
-    const composerMetadata = gateInteractableComposerMetadata(
-      this.runtime.getModelContext().unstable_composerMetadata,
-      messages.slice(0, parentIndex + 1),
-    );
-    const enriched = this.enrichWithComposerMetadata(message, composerMetadata);
     const appendTask = this.runtime.append({
-      ...enriched,
+      ...message,
       content,
       parentId: this._parentId,
       sourceId: this._sourceId,

package/src/runtime/base/default-thread-composer-runtime-core.ts

@@ -12,7 +12,6 @@ import {
   type QueueItemState,
 } from "../../store/scopes/queue-item";
 import { BaseComposerRuntimeCore } from "./base-composer-runtime-core";
-import { gateInteractableComposerMetadata } from "../../model-context/interactable-composer-metadata";
 
 export class DefaultThreadComposerRuntimeCore
   extends BaseComposerRuntimeCore
@@ -119,17 +118,8 @@ export class DefaultThreadComposerRuntimeCore
     message: Omit<AppendMessage, "parentId" | "sourceId">,
     options?: SendOptions,
   ) {
-    // Merge provider-contributed metadata onto the outgoing user message
-    // (same metadata.custom append path quotes ride). The interactables gate
-    // runs here because it needs thread history, unavailable to the provider.
-    const composerMetadata = gateInteractableComposerMetadata(
-      this.runtime.getModelContext().unstable_composerMetadata,
-      this.runtime.messages,
-    );
-    const enriched = this.enrichWithComposerMetadata(message, composerMetadata);
-
     return this.runtime.append({
-      ...(enriched as AppendMessage),
+      ...(message as AppendMessage),
       parentId: this.runtime.messages.at(-1)?.id ?? null,
       sourceId: null,
       startRun: options?.startRun,

package/src/runtime/queue/external-thread-queue-adapter.ts

@@ -27,4 +27,21 @@ export type ExternalThreadQueueAdapter = {
   move: (queueItemId: string, placement: QueuePlacement) => void;
   edit: (queueItemId: string, message: AppendMessage) => void;
   remove: (queueItemId: string) => void;
+  /**
+   * Installed by the owning runtime, applied to a message as it leaves a lane.
+   * A message can wait through an entire run, so state the runtime attaches at
+   * send time is recomputed against the thread it actually lands on. Adapters
+   * from `createMessageQueue` provide this; a hand-rolled adapter that omits
+   * it dispatches whatever was enqueued.
+   */
+  __internal_setDispatchTransform?:
+    | ((transform: (message: AppendMessage) => AppendMessage) => void)
+    | undefined;
+  /**
+   * Called by the owning runtime when the user cancels the run, before the run
+   * is aborted. Pauses the queue so the cancelled run's settle keeps the
+   * pending items instead of dispatching the next one; the next send re-arms
+   * draining. A hand-rolled adapter that omits it drains on cancel.
+   */
+  __internal_notifyCancelled?: (() => void) | undefined;
 };

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

@@ -79,10 +79,12 @@ export const createMessageQueue = (
 
   let running = false;
   let paused = false;
+  let dispatchTransform: (message: AppendMessage) => AppendMessage = (m) => m;
   // swallow the cancelled run's settle when steering so it does not double-advance
   let suppressIdle = 0;
   // settles from cancelled runs that must drop `running` without advancing
   let cancelSettles = 0;
+  let interrupting = false;
 
   const notify = () => {
     for (const callback of subscribers) callback();
@@ -117,7 +119,7 @@ export const createMessageQueue = (
     setLanes({ ...lanes, [lane]: lanes[lane].slice(1) });
     if (!message) return;
     running = true;
-    driver.run(message, { steer: false });
+    driver.run(dispatchTransform(message), { steer: false });
   };
 
   const interrupt = (message: AppendMessage) => {
@@ -126,9 +128,17 @@ export const createMessageQueue = (
     // already cancel-notified
     suppressIdle += Math.max(cancelSettles, 1);
     cancelSettles = 0;
-    driver.cancel!();
+    // a driver whose cancel routes through the runtime notifies this queue
+    // back; the interrupt already accounted for that settle and is dispatching
+    // in its place
+    interrupting = true;
+    try {
+      driver.cancel!();
+    } finally {
+      interrupting = false;
+    }
     running = true;
-    driver.run(message, { steer: true });
+    driver.run(dispatchTransform(message), { steer: true });
   };
 
   const push = (lane: Lane, message: AppendMessage) => {
@@ -241,6 +251,14 @@ export const createMessageQueue = (
     });
   };
 
+  const notifyCancelled = () => {
+    if (interrupting) return;
+    if (running && cancelSettles === 0) {
+      paused = true;
+      cancelSettles = 1;
+    }
+  };
+
   const adapter: ExternalThreadQueueAdapter = {
     items: lanes.queue,
     steerItems: lanes.steer,
@@ -249,6 +267,10 @@ export const createMessageQueue = (
     move,
     edit,
     remove,
+    __internal_setDispatchTransform: (transform) => {
+      dispatchTransform = transform;
+    },
+    __internal_notifyCancelled: notifyCancelled,
   };
 
   return {
@@ -270,12 +292,7 @@ export const createMessageQueue = (
       running = false;
       advance();
     },
-    notifyCancelled: () => {
-      if (running && cancelSettles === 0) {
-        paused = true;
-        cancelSettles = 1;
-      }
-    },
+    notifyCancelled,
     clear: () => {
       messages.clear();
       setLanes({ queue: EMPTY_QUEUE_ITEMS, steer: EMPTY_QUEUE_ITEMS });

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-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,6 +41,7 @@ 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";
 
 const EMPTY_ARRAY: readonly ThreadSuggestion[] = Object.freeze([]);
 
@@ -117,6 +122,8 @@ export class ExternalStoreThreadRuntimeCore
 
   private _store!: ExternalStoreAdapter<any>;
 
+  private _transformedQueue: ExternalThreadQueueAdapter | undefined;
+
   /**
    * Client-side tool-invocations pipeline. Constructed lazily on first
    * snapshot — only when `adapter.unstable_enableToolInvocations === true`.
@@ -147,6 +154,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;
     }
@@ -465,22 +485,30 @@ 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);
+      rawMessage.sourceId != null ||
+      rawMessage.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.
     if (!isEdit && this._store.queue) {
-      if (message.steer ?? this._store.isRunning ?? false)
-        this._store.queue.steer(message);
-      else this._store.queue.enqueue(message);
+      // Skip only for the queue this core actually installed on: another
+      // core's transform would gate against its own thread's messages.
+      const queued =
+        this._store.queue === this._transformedQueue
+          ? rawMessage
+          : this.enrichAppendMetadata(rawMessage);
+      if (queued.steer ?? this._store.isRunning ?? false)
+        this._store.queue.steer(queued);
+      else this._store.queue.enqueue(queued);
       return;
     }
 
+    const message = this.enrichAppendMetadata(rawMessage);
+
     // 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
@@ -592,37 +620,83 @@ export class ExternalStoreThreadRuntimeCore
     // cancel on it.
     void this._toolInvocations?.abort();
 
+    // 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?.();
+
     this._store.onCancel();
 
-    // 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);
-    }
+    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.");

package/src/runtimes/local/local-thread-runtime-core.ts

@@ -312,7 +312,10 @@ export class LocalThreadRuntimeCore
     this._queue?.adapter.remove(queueItemId);
   }
 
-  private async _runAppend(message: AppendMessage): Promise<void> {
+  private async _runAppend(rawMessage: AppendMessage): Promise<void> {
+    // Stamped here rather than in `append` so a queued message is gated after
+    // the flush re-pointed its parentId at the current tail.
+    const message = this.enrichAppendMetadata(rawMessage);
     this.ensureInitialized();
 
     const initPromise = this._getInitializePromise?.();