@assistant-ui/core 0.3.19 → 0.3.21

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

@@ -25,6 +25,8 @@ export type MessageQueueDriver = {
 
 export type MessageQueueController = {
   readonly adapter: ExternalThreadQueueAdapter;
+  hold: () => void;
+  release: () => void;
   /** Mark a run as in flight so concurrent sends buffer; call on the rising edge. */
   notifyBusy: () => void;
   /** Advances to the next pending message; call on the run's falling edge. */
@@ -88,6 +90,7 @@ export const createMessageQueue = (
 
   let running = false;
   let paused = false;
+  let held = 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;
@@ -133,7 +136,7 @@ export const createMessageQueue = (
   };
 
   const advance = () => {
-    if (running || paused) return;
+    if (running || paused || held) return;
     const lane: Lane = lanes.steer.length > 0 ? "steer" : "queue";
     const head = lanes[lane][0];
     if (!head) return;
@@ -334,6 +337,13 @@ export const createMessageQueue = (
 
   return {
     adapter,
+    hold: () => {
+      held = true;
+    },
+    release: () => {
+      held = false;
+      advance();
+    },
     notifyBusy: () => {
       paused = false;
       // a cancelled run's settle that is still outstanding belongs to a run

package/src/runtime/utils/attachment-add-operations.ts

@@ -1,3 +1,5 @@
+import type { Attachment } from "../../types/attachment";
+
 export type AttachmentAddOperation = {
   cancelled: boolean;
   attachmentIds: Set<string>;
@@ -5,6 +7,10 @@ export type AttachmentAddOperation = {
 
 export class AttachmentAddOperations {
   private readonly operations = new Set<AttachmentAddOperation>();
+  private readonly uploading = new Map<
+    string,
+    { operation: AttachmentAddOperation; waiters: Set<() => void> }
+  >();
 
   start() {
     const operation: AttachmentAddOperation = {
@@ -15,14 +21,32 @@ export class AttachmentAddOperations {
     return operation;
   }
 
-  accept(operation: AttachmentAddOperation, attachmentId: string) {
+  accept(
+    operation: AttachmentAddOperation,
+    attachment: Pick<Attachment, "id" | "status">,
+  ) {
     if (operation.cancelled) return false;
-    operation.attachmentIds.add(attachmentId);
+    operation.attachmentIds.add(attachment.id);
+    // The composer shows one attachment per id, so the add that updated it
+    // last owns its upload; an add that ended earlier no longer speaks for it.
+    const entry = this.uploading.get(attachment.id);
+    if (
+      entry?.operation !== operation &&
+      (entry || attachment.status.type === "running")
+    )
+      this.uploading.set(attachment.id, {
+        operation,
+        waiters: entry?.waiters ?? new Set(),
+      });
+    if (attachment.status.type !== "running")
+      this.settle(attachment.id, operation);
     return true;
   }
 
   finish(operation: AttachmentAddOperation) {
     this.operations.delete(operation);
+    for (const attachmentId of operation.attachmentIds)
+      this.settle(attachmentId, operation);
   }
 
   isCancelled(operation: AttachmentAddOperation) {
@@ -35,6 +59,7 @@ export class AttachmentAddOperations {
       operation.cancelled = true;
       this.operations.delete(operation);
     }
+    this.settle(attachmentId);
   }
 
   cancelAll() {
@@ -42,6 +67,24 @@ export class AttachmentAddOperations {
       operation.cancelled = true;
     }
     this.operations.clear();
+    for (const attachmentId of [...this.uploading.keys()])
+      this.settle(attachmentId);
+  }
+
+  whenSendable(attachmentId: string): Promise<void> | undefined {
+    const entry = this.uploading.get(attachmentId);
+    if (!entry) return undefined;
+    return new Promise((resolve) => entry.waiters.add(resolve));
+  }
+
+  // An attachment id outlives the add that produced it, so an add that ends
+  // after a newer one took the id over must not release its upload.
+  private settle(attachmentId: string, operation?: AttachmentAddOperation) {
+    const entry = this.uploading.get(attachmentId);
+    if (!entry) return;
+    if (operation && entry.operation !== operation) return;
+    this.uploading.delete(attachmentId);
+    for (const resolve of entry.waiters) resolve();
   }
 }
 

package/src/runtime/utils/attachment-send-operations.ts

@@ -0,0 +1,54 @@
+import type { AttachmentAdapter } from "../../adapters/attachment";
+import {
+  isAttachmentComplete,
+  type Attachment,
+  type CompleteAttachment,
+} from "../../types/attachment";
+
+export class AttachmentSendOperations {
+  // Draft attachments stay pending so discarding an unsent upload still calls remove.
+  private readonly entries = new WeakMap<
+    Attachment,
+    { result?: CompleteAttachment }
+  >();
+  // Removal marks are per-object and never bulk-cleared: a removed attachment
+  // either leaves the draft or is replaced via transfer with a fresh unmarked
+  // object, so a mark cannot leak into a later send's batch.
+  private readonly removed = new WeakSet<Attachment>();
+
+  markRemoved(attachment: Attachment) {
+    this.removed.add(attachment);
+  }
+
+  unmarkRemoved(attachment: Attachment) {
+    this.removed.delete(attachment);
+  }
+
+  isRemoved(attachment: Attachment) {
+    return this.removed.has(attachment);
+  }
+
+  async send(
+    attachment: Attachment,
+    adapter: AttachmentAdapter | undefined,
+    signal?: AbortSignal,
+  ): Promise<CompleteAttachment> {
+    if (isAttachmentComplete(attachment)) return attachment;
+    const entry = this.entries.get(attachment) ?? {};
+    if (entry.result) return entry.result;
+    if (!adapter) throw new Error("Attachments are not supported");
+    this.entries.set(attachment, entry);
+    const result = await adapter.send(
+      attachment,
+      signal ? { signal } : undefined,
+    );
+    entry.result = result;
+    return result;
+  }
+
+  transfer(original: Attachment, replacement: Attachment) {
+    const entry = this.entries.get(original);
+    if (entry) this.entries.set(replacement, entry);
+    return replacement;
+  }
+}

package/src/runtime/utils/auto-status.test.ts

@@ -48,6 +48,21 @@ describe("getAutoStatus", () => {
     ).toMatchObject({ type: "complete", reason: "unknown" });
   });
 
+  it.each([false, 0, ""])("preserves the falsy error payload %j", (error) => {
+    expect(getAutoStatus(true, false, false, false, error)).toMatchObject({
+      type: "incomplete",
+      reason: "error",
+      error,
+    });
+  });
+
+  it("treats a null error as no error", () => {
+    expect(getAutoStatus(true, false, false, false, null)).toMatchObject({
+      type: "complete",
+      reason: "unknown",
+    });
+  });
+
   it.each([
     ["running", { type: "running" }, true, false, false, undefined],
     [
@@ -143,6 +158,45 @@ describe("getAutoStatus", () => {
     ).toMatchObject({ type: "requires-action", reason: "interrupt" });
   });
 
+  it.each([
+    ["approval", { approval: { id: "approval-1" } }],
+    ["interrupt", { interrupt: { type: "human" as const, payload: {} } }],
+  ] as const)(
+    "keeps a tool call with a pending %s interrupted beside its result",
+    (_label, action) => {
+      expect(
+        getContentAutoStatus(
+          [{ ...pendingToolCall(), result: "partial output", ...action }],
+          true,
+          false,
+        ),
+      ).toMatchObject({ type: "requires-action", reason: "interrupt" });
+    },
+  );
+
+  it.each([
+    ["a decision", { approved: true }],
+    ["a rejection", { approved: false }],
+    ["a resolution", { resolution: "cancelled" as const }],
+  ])(
+    "settles a tool call whose approval carries %s beside its result",
+    (_label, settled) => {
+      expect(
+        getContentAutoStatus(
+          [
+            {
+              ...pendingToolCall(),
+              result: "sunny",
+              approval: { id: "approval-1", ...settled },
+            },
+          ],
+          true,
+          false,
+        ),
+      ).toMatchObject({ type: "complete" });
+    },
+  );
+
   it.each([
     ["a result", { ...pendingToolCall(), result: {} }],
     ["no nested messages", pendingToolCall()],

package/src/runtime/utils/auto-status.ts

@@ -1,6 +1,7 @@
 import type { MessageStatus } from "../../types/message";
 import type { ReadonlyJSONValue } from "assistant-stream/utils";
 import type { ThreadMessageLike } from "./thread-message-like";
+import { hasPendingToolAction } from "../../utils/normalizePartStatus";
 
 type ThreadMessageLikeContentItem = Exclude<
   ThreadMessageLike["content"],
@@ -21,15 +22,7 @@ export const isBackgroundToolCall = (
 
 export const isInterruptedToolCall = (
   c: ThreadMessageLikeContentItem,
-): boolean => {
-  if (c.type !== "tool-call" || c.result !== undefined) return false;
-  return (
-    c.interrupt != null ||
-    (c.approval != null &&
-      c.approval.approved === undefined &&
-      c.approval.resolution === undefined)
-  );
-};
+): boolean => c.type === "tool-call" && hasPendingToolAction(c);
 
 const symbolAutoStatus = Symbol("autoStatus");
 
@@ -87,7 +80,7 @@ export const getAutoStatus = (
   isCancelled?: boolean,
   hasBackgroundToolCalls?: boolean,
 ): MessageStatus => {
-  if (isLast && error) {
+  if (isLast && error !== undefined && error !== null) {
     return Object.assign(
       {
         type: "incomplete" as const,

package/src/runtime/utils/external-message-conversion.test.ts

@@ -2,6 +2,7 @@ import { describe, expect, it, vi } from "vitest";
 import type { ThreadMessage } from "../../types/message";
 import {
   chunkExternalMessages,
+  completeExternalMessageConversion,
   convertExternalMessageCallback,
   convertExternalMessageChunk,
   joinExternalMessages,
@@ -10,6 +11,25 @@ import {
   type ExternalMessageConverterMessage,
 } from "./external-message-conversion";
 
+describe("completeExternalMessageConversion", () => {
+  it.each([false, 0, ""])(
+    "adds an assistant message for the falsy error payload %j",
+    (error) => {
+      const result = completeExternalMessageConversion([], error);
+
+      expect(result).toHaveLength(1);
+      expect(result[0]).toMatchObject({
+        role: "assistant",
+        status: { type: "incomplete", reason: "error", error },
+      });
+    },
+  );
+
+  it("adds no assistant message for a null error", () => {
+    expect(completeExternalMessageConversion([], null)).toHaveLength(0);
+  });
+});
+
 describe("convertExternalMessageCallback", () => {
   it.each([
     ["a missing toolCallId", { role: "tool", result: "ok" }],
@@ -69,6 +89,34 @@ describe("joinExternalMessages", () => {
     ]);
   });
 
+  it("settles a preliminary tool call when its tool message lands", () => {
+    const messages = [
+      {
+        role: "assistant",
+        content: [
+          {
+            type: "tool-call",
+            toolCallId: "call-1",
+            toolName: "bash",
+            args: {},
+            result: "partial output",
+            isPreliminary: true,
+          },
+        ],
+      },
+      {
+        role: "tool",
+        toolCallId: "call-1",
+        toolName: "bash",
+        result: "final output",
+      },
+    ] as unknown as ExternalMessageConverterMessage[];
+
+    const [part] = joinExternalMessages(messages).content;
+    expect(part).toMatchObject({ type: "tool-call", result: "final output" });
+    expect(part).not.toHaveProperty("isPreliminary");
+  });
+
   it("does not merge malformed NaN tool-call IDs", () => {
     const messages = [
       {
@@ -152,6 +200,49 @@ describe("chunkExternalMessages", () => {
     ]);
     expect(chunks[1]?.inputs).toEqual([answer]);
   });
+
+  it("keeps a voice assistant transcript out of the neighbouring assistant chunks", () => {
+    const typed = {};
+    const spoken = {};
+    const spokenAgain = {};
+    const typedAgain = {};
+    const callbackResults: ExternalMessageConverterCallbackResult<object>[] = [
+      { input: typed, outputs: [{ role: "assistant", content: "typed" }] },
+      {
+        input: spoken,
+        outputs: [
+          {
+            role: "assistant",
+            content: "spoken",
+            metadata: { modality: "voice" },
+          },
+        ],
+      },
+      {
+        input: spokenAgain,
+        outputs: [
+          {
+            role: "assistant",
+            content: "spoken again",
+            metadata: { modality: "voice" },
+          },
+        ],
+      },
+      {
+        input: typedAgain,
+        outputs: [{ role: "assistant", content: "typed again" }],
+      },
+    ];
+
+    const chunks = chunkExternalMessages(callbackResults);
+
+    expect(chunks.map((chunk) => chunk.inputs)).toEqual([
+      [typed],
+      [spoken],
+      [spokenAgain],
+      [typedAgain],
+    ]);
+  });
 });
 
 describe("convertExternalMessageChunk", () => {

package/src/runtime/utils/external-message-conversion.ts

@@ -157,9 +157,8 @@ export const joinExternalMessages = (
       // Ignore orphaned tool results so one bad tool message does not
       // prevent rendering the rest of the conversation.
       if (toolCallIdx !== -1) {
-        const toolCall = assistantMessage.content[
-          toolCallIdx
-        ]! as ToolCallMessagePart;
+        const { isPreliminary: _isPreliminary, ...toolCall } = assistantMessage
+          .content[toolCallIdx]! as ToolCallMessagePart;
         if (output.toolName != null) {
           if (toolCall.toolName !== output.toolName)
             throw new Error(
@@ -348,11 +347,14 @@ export const chunkExternalMessages = <T>(
 
   for (const callbackResult of callbackResults) {
     for (const output of callbackResult.outputs) {
+      const isVoice =
+        output.role === "assistant" && output.metadata?.modality === "voice";
       if (
         (pendingNone && output.role !== "tool") ||
         !isAssistant ||
         output.role === "user" ||
-        output.role === "system"
+        output.role === "system" ||
+        isVoice
       ) {
         flush();
       }
@@ -365,7 +367,8 @@ export const chunkExternalMessages = <T>(
 
       if (
         output.role === "assistant" &&
-        (output.convertConfig?.joinStrategy === "none" ||
+        (isVoice ||
+          output.convertConfig?.joinStrategy === "none" ||
           joinStrategy === "none")
       ) {
         pendingNone = true;
@@ -525,7 +528,7 @@ export const completeExternalMessageConversion = (
   messages: ThreadMessage[],
   error: ReadonlyJSONValue | undefined,
 ) => {
-  if (error) {
+  if (error !== undefined && error !== null) {
     const lastMessage = messages.at(-1);
     if (!lastMessage || lastMessage.role !== "assistant") {
       messages.push(createErrorAssistantMessage(error));

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

@@ -1,6 +1,9 @@
 import { describe, expect, it } from "vitest";
 import type { ThreadMessage } from "../../types/message";
-import { ExportedMessageRepository } from "./message-repository";
+import {
+  ExportedMessageRepository,
+  MessageRepository,
+} from "./message-repository";
 
 const nestedRunningAssistant: ThreadMessage = {
   id: "nested-assistant",
@@ -47,3 +50,68 @@ describe("ExportedMessageRepository", () => {
     }
   });
 });
+
+const message = (id: string, text = id): ThreadMessage => ({
+  id,
+  createdAt: new Date(0),
+  role: "user",
+  content: [{ type: "text", text }],
+  attachments: [],
+  metadata: { custom: {} },
+});
+
+const snapshot = (repository: MessageRepository) => ({
+  headId: repository.headId,
+  visible: repository.getMessages(),
+  exported: repository.export(),
+});
+
+describe("MessageRepository rejected operations", () => {
+  it("keeps the old content when an update is rejected as a cycle", () => {
+    const repository = new MessageRepository();
+    const original = message("a");
+    repository.addOrUpdateMessage(null, original);
+    const before = snapshot(repository);
+
+    expect(() =>
+      repository.addOrUpdateMessage("a", message("a", "edited")),
+    ).toThrow(/same id already exists in the parent tree/);
+
+    expect(snapshot(repository)).toEqual(before);
+    expect(repository.getMessage("a").message).toBe(original);
+  });
+
+  it("moves no child when the replacement is a descendant of the deleted message", () => {
+    const repository = new MessageRepository();
+    repository.addOrUpdateMessage(null, message("a"));
+    repository.addOrUpdateMessage("a", message("b"));
+    repository.addOrUpdateMessage("a", message("c"));
+    repository.addOrUpdateMessage("c", message("d"));
+    const before = snapshot(repository);
+
+    expect(() => repository.deleteMessage("a", "c")).toThrow(
+      /Replacement is the deleted message or one of its descendants/,
+    );
+    expect(snapshot(repository)).toEqual(before);
+
+    expect(() => repository.deleteMessage("a", "d")).toThrow(
+      /Replacement is the deleted message or one of its descendants/,
+    );
+    expect(snapshot(repository)).toEqual(before);
+    expect(repository.getMessage("b").parentId).toBe("a");
+    expect(repository.getBranches("b")).toEqual(["b", "c"]);
+  });
+
+  it("rejects a message as its own replacement instead of leaving the head on it", () => {
+    const repository = new MessageRepository();
+    repository.addOrUpdateMessage(null, message("a"));
+    const before = snapshot(repository);
+
+    expect(() => repository.deleteMessage("a", "a")).toThrow(
+      /Replacement is the deleted message or one of its descendants/,
+    );
+
+    expect(snapshot(repository)).toEqual(before);
+    expect(repository.getMessage("a").message.id).toBe("a");
+  });
+});

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

@@ -262,8 +262,8 @@ export class MessageRepository {
       );
 
     if (existingItem) {
-      existingItem.current = message;
       this.performOp(prev, existingItem, "relink");
+      existingItem.current = message;
       this._messages.dirty();
       return;
     }
@@ -319,6 +319,13 @@ export class MessageRepository {
         "MessageRepository(deleteMessage): Replacement not found. This is likely an internal bug in assistant-ui.",
       );
 
+    for (let current = replacement; current; current = current.prev) {
+      if (current === message)
+        throw new Error(
+          "MessageRepository(deleteMessage): Replacement is the deleted message or one of its descendants. This is likely an internal bug in assistant-ui.",
+        );
+    }
+
     for (const child of message.children) {
       const childMessage = this.messages.get(child);
       if (!childMessage)
@@ -350,12 +357,12 @@ export class MessageRepository {
   }
 
   /**
-   * Evicts optimistic messages (`metadata.isOptimistic`) the head just moved
-   * away from. Since eviction runs on every head move, the only optimistic
-   * messages in the repository live on the branch the head previously pointed
-   * at — so we walk just that branch rather than the whole repository. Keeps a
-   * client→server id swap from leaving a phantom sibling, and drops off-branch
-   * placeholders.
+   * Evicts optimistic messages (`metadata.isOptimistic`) on the branch the head
+   * just moved away from. Only that branch is walked, so an optimistic message
+   * added off the head branch (such as the server-id copy that replaces a
+   * client-id placeholder before the head moves to it) is kept until a
+   * `switchToBranch` or `resetHead` moves the head off the branch it is on.
+   * Keeps a client→server id swap from leaving a phantom sibling.
    */
   private evictOffBranchOptimisticMessages(
     previousHead: RepositoryMessage | null,