@assistant-ui/core 0.2.20 → 0.2.22

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

@@ -0,0 +1,81 @@
+import { describe, expect, it } from "vitest";
+import type { CompleteAttachment } from "../../types/attachment";
+import type { ThreadRuntimeCoreBinding } from "./thread-runtime";
+import {
+  MessageRuntimeImpl,
+  type MessageState,
+  type MessageStateBinding,
+} from "./message-runtime";
+
+const messagePath = {
+  ref: "threads.main.messages[0]",
+  threadSelector: { type: "main" as const },
+  messageSelector: { type: "index" as const, index: 0 },
+};
+
+const attachment: CompleteAttachment = {
+  id: "attachment-1",
+  type: "file",
+  name: "notes.txt",
+  contentType: "text/plain",
+  content: [{ type: "text", text: "notes" }],
+  status: { type: "complete" },
+};
+
+const message: MessageState = {
+  id: "message-1",
+  role: "assistant",
+  createdAt: new Date(0),
+  content: [
+    { type: "text", text: "Hello" },
+    {
+      type: "tool-call",
+      toolCallId: "call-1",
+      toolName: "weather",
+      args: {},
+      argsText: "{}",
+    },
+  ],
+  attachments: [attachment],
+  status: { type: "complete", reason: "stop" },
+  metadata: {
+    unstable_state: null,
+    unstable_annotations: [],
+    unstable_data: [],
+    steps: [],
+    custom: {},
+  },
+  parentId: null,
+  index: 0,
+  isLast: true,
+  branchNumber: 1,
+  branchCount: 1,
+  speech: undefined,
+};
+
+const messageBinding: MessageStateBinding = {
+  path: messagePath,
+  getState: () => message,
+  subscribe: () => () => {},
+};
+
+const threadBinding = {
+  subscribe: () => () => {},
+} as unknown as ThreadRuntimeCoreBinding;
+
+describe("MessageRuntimeImpl paths", () => {
+  it("appends nested selectors to the message path", () => {
+    const runtime = new MessageRuntimeImpl(messageBinding, threadBinding);
+
+    expect(runtime.composer.path.ref).toBe("threads.main.messages[0].composer");
+    expect(runtime.getMessagePartByIndex(0).path.ref).toBe(
+      "threads.main.messages[0].content[0]",
+    );
+    expect(runtime.getMessagePartByToolCallId("call-1").path.ref).toBe(
+      'threads.main.messages[0].content[toolCallId="call-1"]',
+    );
+    expect(runtime.getAttachmentByIndex(0).path.ref).toBe(
+      "threads.main.messages[0].attachments[0]",
+    );
+  });
+});

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

@@ -142,7 +142,7 @@ export class MessageRuntimeImpl implements MessageRuntime {
       new NestedSubscriptionSubject({
         path: {
           ...this.path,
-          ref: `${this.path.ref}${this.path.ref}.composer`,
+          ref: `${this.path.ref}.composer`,
           composerSource: "edit",
         },
         getState: this._getEditComposerRuntimeCore,
@@ -269,7 +269,7 @@ export class MessageRuntimeImpl implements MessageRuntime {
       new ShallowMemoizeSubject({
         path: {
           ...this.path,
-          ref: `${this.path.ref}${this.path.ref}.content[${idx}]`,
+          ref: `${this.path.ref}.content[${idx}]`,
           messagePartSelector: { type: "index", index: idx },
         },
         getState: () => {
@@ -287,9 +287,7 @@ export class MessageRuntimeImpl implements MessageRuntime {
       new ShallowMemoizeSubject({
         path: {
           ...this.path,
-          ref:
-            this.path.ref +
-            `${this.path.ref}.content[toolCallId=${JSON.stringify(toolCallId)}]`,
+          ref: `${this.path.ref}.content[toolCallId=${JSON.stringify(toolCallId)}]`,
           messagePartSelector: { type: "toolCallId", toolCallId },
         },
         getState: () => {
@@ -313,7 +311,7 @@ export class MessageRuntimeImpl implements MessageRuntime {
       new ShallowMemoizeSubject({
         path: {
           ...this.path,
-          ref: `${this.path.ref}${this.path.ref}.attachments[${idx}]`,
+          ref: `${this.path.ref}.attachments[${idx}]`,
           attachmentSource: "message",
           attachmentSelector: { type: "index", index: idx },
         },

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

@@ -2,6 +2,7 @@ import type { Unsubscribe } from "../../types/unsubscribe";
 import type { SubscribableWithState } from "../../subscribable/subscribable";
 import type { ThreadListItemRuntimePath } from "./paths";
 import type { ThreadListRuntimeCoreBinding } from "./thread-list-runtime";
+import { notifyEventListeners } from "../../utils/notify-event-listeners";
 
 export type ThreadListItemEventPayload = {
   /**
@@ -162,7 +163,11 @@ export class ThreadListItemRuntimeImpl implements ThreadListItemRuntime {
 
       if (event === "switchedTo" && !newIsMain) return;
       if (event === "switchedAway" && newIsMain) return;
-      (callback as (payload?: unknown) => void)({});
+      notifyEventListeners(
+        [callback as (payload?: unknown) => void],
+        {},
+        `Thread list item "${event}"`,
+      );
     });
   }
 

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

@@ -1,8 +1,9 @@
-import type {
-  Attachment,
-  CompleteAttachment,
-  CreateAttachment,
-  PendingAttachment,
+import {
+  isCreateAttachment,
+  type Attachment,
+  type CompleteAttachment,
+  type CreateAttachment,
+  type PendingAttachment,
 } from "../../types/attachment";
 import type { MessageRole, AppendMessage } from "../../types/message";
 import type { QuoteInfo } from "../../types/quote";
@@ -28,6 +29,7 @@ import {
   type QueueItemState,
 } from "../../store/scopes/queue-item";
 import { generateId } from "../../utils/id";
+import { notifyEventListeners } from "../../utils/notify-event-listeners";
 
 const isAttachmentComplete = (a: Attachment): a is CompleteAttachment =>
   a.status.type === "complete";
@@ -133,6 +135,10 @@ export abstract class BaseComposerRuntimeCore
     this._notifySubscribers();
   }
 
+  protected _isSending = false;
+  private _removedDuringSend = new Set<string>();
+  private _sendGeneration = 0;
+
   private _emptyTextAndAttachments() {
     this._attachments = [];
     this._text = "";
@@ -148,6 +154,14 @@ export abstract class BaseComposerRuntimeCore
   }
 
   public async reset() {
+    // A send whose adapter never settles must not brick the composer; reset is
+    // the escape hatch that releases the in-flight lock. Bumping the generation
+    // invalidates that send entirely so a late-settling upload can neither
+    // append the discarded draft nor touch a newer send's lock.
+    this._sendGeneration++;
+    this._isSending = false;
+    this._removedDuringSend.clear();
+
     if (
       this._attachments.length === 0 &&
       this._text === "" &&
@@ -175,7 +189,7 @@ export abstract class BaseComposerRuntimeCore
   }
 
   public async send(options?: SendOptions) {
-    if (!this.canSend) return;
+    if (!this.canSend || this._isSending) return;
 
     if (this._dictationSession) {
       this._dictationSession.cancel();
@@ -183,33 +197,75 @@ export abstract class BaseComposerRuntimeCore
     }
 
     const adapter = this.getAttachmentAdapter();
-    const attachments =
-      this.attachments.length > 0
-        ? Promise.all(
-            this.attachments.map(async (a) => {
-              if (isAttachmentComplete(a)) return a;
-              if (!adapter) throw new Error("Attachments are not supported");
-              const result = await adapter.send(a);
-              return result as CompleteAttachment;
-            }),
-          )
-        : [];
+    const attachmentTasks = this.attachments.map(async (a) => {
+      if (isAttachmentComplete(a)) return a;
+      if (!adapter) throw new Error("Attachments are not supported");
+      const result = await adapter.send(a);
+      return result as CompleteAttachment;
+    });
 
+    const originalAttachments = this.attachments;
     const text = this.text;
     const quote = this._quote;
     this._quote = undefined;
-    this._emptyTextAndAttachments();
+    this._text = "";
+    this._isSending = true;
+    const generation = ++this._sendGeneration;
+    this._notifySubscribers();
+
+    let resolvedAttachments: CompleteAttachment[];
+    try {
+      resolvedAttachments = await Promise.all(attachmentTasks);
+    } catch (e) {
+      if (generation === this._sendGeneration) {
+        if (!this.text.trim() && this._quote === undefined) {
+          this._text = text;
+          this._quote = quote;
+          this._notifySubscribers();
+        }
+        // Promise.all rejects on the first failure, but sibling uploads from
+        // this batch keep running; the send rejects immediately while the
+        // retry lock is held until they settle, or a retry could re-send
+        // attachments that are still in flight.
+        void Promise.allSettled(attachmentTasks).then(() => {
+          if (generation !== this._sendGeneration) return;
+          this._removedDuringSend.clear();
+          this._isSending = false;
+          this._notifySubscribers();
+        });
+      }
+      throw e;
+    }
+
+    // A reset during the upload discarded this send's draft; the settled
+    // uploads must not append it or touch the lock a newer send may own.
+    if (generation !== this._sendGeneration) return;
+
+    // Drop by id, not wholesale: the user may have added or removed chips while
+    // the upload was in flight.
+    const sentIds = new Set(originalAttachments.map((a) => a.id));
+    this._attachments = this._attachments.filter((a) => !sentIds.has(a.id));
+    this._isSending = false;
+    this._notifySubscribers();
+
+    // An attachment removed mid-upload can't be cancelled, but it can still be
+    // dropped from the outgoing message instead of silently being sent anyway.
+    const finalAttachments = resolvedAttachments.filter(
+      (a) => !this._removedDuringSend.has(a.id),
+    );
+    this._removedDuringSend.clear();
 
     const message: Omit<AppendMessage, "parentId" | "sourceId"> = {
       createdAt: new Date(),
       role: this.role,
       content: text ? [{ type: "text", text }] : [],
-      attachments: await attachments,
+      attachments: finalAttachments,
       runConfig: this.runConfig,
       metadata: { custom: { ...(quote ? { quote } : {}) } },
     };
 
-    this.handleSend(message, options);
+    const sendTask = this.handleSend(message, options);
+    if (sendTask) void sendTask.catch(() => {});
     this._notifyEventSubscribers("send", {});
   }
 
@@ -227,11 +283,11 @@ export abstract class BaseComposerRuntimeCore
   protected abstract handleSend(
     message: Omit<AppendMessage, "parentId" | "sourceId">,
     options?: SendOptions,
-  ): void;
+  ): void | Promise<void>;
   protected abstract handleCancel(): void;
 
   async addAttachment(fileOrAttachment: File | CreateAttachment) {
-    if (!(fileOrAttachment instanceof File)) {
+    if (isCreateAttachment(fileOrAttachment)) {
       const adapter = this.getAttachmentAdapter();
       if (
         adapter &&
@@ -321,7 +377,11 @@ export abstract class BaseComposerRuntimeCore
       if (lastAttachment) {
         upsertAttachment({
           ...lastAttachment,
-          status: { type: "incomplete", reason: "error" },
+          status: {
+            type: "incomplete",
+            reason: "error",
+            message: e instanceof Error ? e.message : String(e),
+          },
         });
       }
       this._safeEmitAttachmentAddError(
@@ -333,14 +393,15 @@ export abstract class BaseComposerRuntimeCore
       throw e;
     }
 
-    const hasError =
+    if (
       lastAttachment?.status.type === "incomplete" &&
-      lastAttachment.status.reason === "error";
-    if (hasError) {
+      lastAttachment.status.reason === "error"
+    ) {
       this._safeEmitAttachmentAddError(
         "adapter-error",
-        "Attachment upload did not complete successfully.",
-        lastAttachment?.id,
+        lastAttachment.status.message ??
+          "Attachment upload did not complete successfully.",
+        lastAttachment.id,
       );
     } else {
       this._notifyEventSubscribers("attachmentAdd", {});
@@ -373,12 +434,16 @@ export abstract class BaseComposerRuntimeCore
     if (index === -1) throw new Error("Attachment not found");
     const attachment = this._attachments[index]!;
 
+    // A send in flight may already be uploading this attachment; the upload
+    // can't be cancelled, so mark it to be dropped from the outgoing message
+    // before any await gives the upload a chance to settle first.
+    if (this._isSending) this._removedDuringSend.add(attachmentId);
+
     if (!isAttachmentComplete(attachment)) {
       const adapter = this.getAttachmentAdapter();
       if (!adapter) throw new Error("Attachments are not supported");
       await adapter.remove(attachment);
     }
-
     this._attachments = this._attachments.filter((a) => a.id !== attachmentId);
     this._notifySubscribers();
   }
@@ -544,7 +609,7 @@ export abstract class BaseComposerRuntimeCore
     const subscribers = this._eventSubscribers.get(event);
     if (!subscribers) return;
 
-    for (const callback of subscribers) callback(payload);
+    notifyEventListeners(subscribers, payload, `Composer runtime "${event}"`);
   }
 
   public unstable_on<E extends ComposerRuntimeEventType>(

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

@@ -34,6 +34,7 @@ import type { FeedbackAdapter } from "../../adapters/feedback";
 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";
 
 type BaseThreadAdapters = {
   speech?: SpeechSynthesisAdapter | undefined;
@@ -182,7 +183,7 @@ export abstract class BaseThreadRuntimeCore implements ThreadRuntimeCore {
     const subscribers = this._eventSubscribers.get(event);
     if (!subscribers) return;
 
-    for (const callback of subscribers) callback(payload);
+    notifyEventListeners(subscribers, payload, `Thread runtime "${event}"`);
   }
 
   public subscribe(callback: () => void): Unsubscribe {

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

@@ -15,7 +15,7 @@ export class DefaultEditComposerRuntimeCore extends BaseComposerRuntimeCore {
   }
 
   public get canSend() {
-    return !this.isEmpty;
+    return !this.isEmpty && !this._isSending;
   }
 
   protected getAttachmentAdapter() {
@@ -80,6 +80,7 @@ export class DefaultEditComposerRuntimeCore extends BaseComposerRuntimeCore {
     message: Omit<AppendMessage, "parentId" | "sourceId">,
     options?: SendOptions,
   ) {
+    let appendTask: void | Promise<void> = undefined;
     const text = getThreadMessageText(message as AppendMessage);
     const attachmentsChanged = !attachmentsEqual(
       message.attachments ?? [],
@@ -114,7 +115,7 @@ export class DefaultEditComposerRuntimeCore extends BaseComposerRuntimeCore {
         message,
         composerMetadata,
       );
-      this.runtime.append({
+      appendTask = this.runtime.append({
         ...enriched,
         content,
         parentId: this._parentId,
@@ -124,6 +125,7 @@ export class DefaultEditComposerRuntimeCore extends BaseComposerRuntimeCore {
     }
 
     this.handleCancel();
+    return appendTask;
   }
 
   public handleCancel() {

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

@@ -23,7 +23,7 @@ export class DefaultThreadComposerRuntimeCore
   }
 
   public get canSend() {
-    return !this.isEmpty && !this.runtime.isSendDisabled;
+    return !this.isEmpty && !this.runtime.isSendDisabled && !this._isSending;
   }
 
   public override get queue(): readonly QueueItemState[] {
@@ -94,7 +94,7 @@ export class DefaultThreadComposerRuntimeCore
     );
     const enriched = this.enrichWithComposerMetadata(message, composerMetadata);
 
-    this.runtime.append({
+    return this.runtime.append({
       ...(enriched as AppendMessage),
       parentId: this.runtime.messages.at(-1)?.id ?? null,
       sourceId: null,

package/src/runtime/utils/message-repository.ts

@@ -436,13 +436,18 @@ export class MessageRepository {
   export(): ExportedMessageRepository {
     const exportItems: ExportedMessageRepository["messages"] = [];
 
-    // Optimistic messages are ephemeral and never persisted. They're always
-    // leaf nodes, so skipping them can't orphan a persisted child.
+    // Optimistic messages are ephemeral and never persisted. A persisted child
+    // of an optimistic node is re-parented onto its nearest persisted ancestor
+    // so the exported tree never references a skipped id.
     for (const [, message] of this.messages) {
       if (message.current.metadata?.isOptimistic) continue;
+      let prev = message.prev;
+      while (prev && prev.current.metadata?.isOptimistic) {
+        prev = prev.prev;
+      }
       exportItems.push({
         message: message.current,
-        parentId: message.prev?.current.id ?? null,
+        parentId: prev?.current.id ?? null,
       });
     }
 

package/src/runtime/utils/thread-message-like.ts

@@ -20,6 +20,7 @@ import type {
 import type { CompleteAttachment } from "../../types/attachment";
 import type {
   MessageTiming,
+  PartProviderMetadata,
   TextMessagePart,
   ToolApprovalOption,
   ToolCallTiming,
@@ -61,6 +62,7 @@ export type ThreadMessageLike = {
             readonly messages?: readonly ThreadMessage[] | undefined;
             readonly interrupt?: { type: "human"; payload: unknown };
             readonly timing?: ToolCallTiming;
+            readonly providerMetadata?: PartProviderMetadata;
             readonly approval?: {
               readonly id: string;
               readonly approved?: boolean;
@@ -264,6 +266,7 @@ export const fromThreadMessageLike = (
         })),
         metadata: {
           custom: metadata?.custom ?? {},
+          ...(metadata?.isOptimistic && { isOptimistic: true }),
         },
       } satisfies ThreadUserMessage;