@assistant-ui/core 0.3.16 → 0.3.18

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

@@ -83,6 +83,74 @@ describe("chunkExternalMessages", () => {
 });
 
 describe("convertExternalMessageChunk", () => {
+  it("keeps separate tool calls without IDs", () => {
+    const result = convertExternalMessageChunk(
+      {
+        inputs: [{}],
+        outputs: [
+          {
+            role: "assistant",
+            content: [
+              { type: "tool-call", toolName: "weather", args: {} },
+              { type: "tool-call", toolName: "search", args: {} },
+            ],
+          },
+        ],
+      },
+      0,
+      1,
+      false,
+      undefined,
+    );
+
+    expect(result.content).toMatchObject([
+      { type: "tool-call", toolName: "weather" },
+      { type: "tool-call", toolName: "search" },
+    ]);
+    const calls = result.content.filter((part) => part.type === "tool-call");
+    expect(calls[0]!.toolCallId).not.toBe(calls[1]!.toolCallId);
+    expect(calls.map((call) => call.toolCallId)).not.toContain("");
+  });
+
+  it("keeps separate tool calls with empty IDs", () => {
+    const result = convertExternalMessageChunk(
+      {
+        inputs: [{}],
+        outputs: [
+          {
+            role: "assistant",
+            content: [
+              {
+                type: "tool-call",
+                toolCallId: "",
+                toolName: "weather",
+                args: {},
+              },
+              {
+                type: "tool-call",
+                toolCallId: "",
+                toolName: "search",
+                args: {},
+              },
+            ],
+          },
+        ],
+      },
+      0,
+      1,
+      false,
+      undefined,
+    );
+
+    expect(result.content).toMatchObject([
+      { type: "tool-call", toolName: "weather" },
+      { type: "tool-call", toolName: "search" },
+    ]);
+    const calls = result.content.filter((part) => part.type === "tool-call");
+    expect(calls[0]?.toolCallId).not.toBe(calls[1]?.toolCallId);
+    expect(calls.map((call) => call.toolCallId)).not.toContain("");
+  });
+
   it("reuses a cached message when the error status is unchanged", () => {
     const input = {};
     const chunk = {
@@ -110,4 +178,101 @@ describe("convertExternalMessageChunk", () => {
 
     expect(second).toBe(first);
   });
+
+  it("replaces a cached complete status after cancellation", () => {
+    const input = {};
+    const chunk = {
+      inputs: [input],
+      outputs: [{ id: "m1", role: "assistant" as const, content: "partial" }],
+    };
+    const cancelled = new Set(["m1"]);
+    const generatedFallbackMessages = new WeakSet<object>();
+    const first = convertExternalMessageChunk(chunk, 0, 1, false, undefined, {
+      message: undefined,
+      generatedFallbackMessages,
+    });
+
+    const second = convertExternalMessageChunk(
+      chunk,
+      0,
+      1,
+      false,
+      undefined,
+      { message: first, generatedFallbackMessages },
+      cancelled,
+    );
+
+    expect(first.status).toMatchObject({ type: "complete", reason: "unknown" });
+    expect(second).not.toBe(first);
+    expect(second.status).toMatchObject({
+      type: "incomplete",
+      reason: "cancelled",
+    });
+
+    const failed = convertExternalMessageChunk(
+      chunk,
+      0,
+      1,
+      false,
+      "failed",
+      { message: second, generatedFallbackMessages },
+      cancelled,
+    );
+    expect(failed.status).toMatchObject({
+      type: "incomplete",
+      reason: "error",
+      error: "failed",
+    });
+  });
+
+  it("cancels a joined chunk when any joined message was stopped", () => {
+    const generatedFallbackMessages = new WeakSet<object>();
+    const result = convertExternalMessageChunk(
+      {
+        inputs: [{}, {}],
+        outputs: [
+          { id: "m1", role: "assistant" as const, content: "first" },
+          { id: "m2", role: "assistant" as const, content: "second" },
+        ],
+      },
+      0,
+      1,
+      false,
+      undefined,
+      { message: undefined, generatedFallbackMessages },
+      new Set(["m2"]),
+    );
+
+    expect(result.status).toMatchObject({
+      type: "incomplete",
+      reason: "cancelled",
+    });
+  });
+
+  it("cancels an earlier message that is no longer the last one", () => {
+    const cancelled = new Set(["m1"]);
+    const generatedFallbackMessages = new WeakSet<object>();
+    const build = (id: string, idx: number) =>
+      convertExternalMessageChunk(
+        {
+          inputs: [{}],
+          outputs: [{ id, role: "assistant" as const, content: "text" }],
+        },
+        idx,
+        2,
+        false,
+        undefined,
+        { message: undefined, generatedFallbackMessages },
+        cancelled,
+      );
+
+    expect(build("m1", 0).status).toMatchObject({
+      type: "incomplete",
+      reason: "cancelled",
+    });
+    expect(build("m2", 1).status).toMatchObject({
+      type: "complete",
+      reason: "unknown",
+    });
+  });
 });

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

@@ -46,6 +46,7 @@ export type ExternalMessageConverterMessage =
 export type ExternalMessageConverterMetadata = {
   readonly toolStatuses?: Record<string, ToolExecutionStatus>;
   readonly error?: ReadonlyJSONValue;
+  readonly cancelledMessageIds?: ReadonlySet<string>;
   readonly messageTiming?: Record<string, MessageTiming>;
 };
 
@@ -251,7 +252,7 @@ export const joinExternalMessages = (
 
           // Add content parts, merging reasoning parts with same parentId
           for (const part of content) {
-            if (part.type === "tool-call") {
+            if (part.type === "tool-call" && part.toolCallId) {
               const existingIdx = assistantMessage.content.findIndex(
                 (c) =>
                   c.type === "tool-call" && c.toolCallId === part.toolCallId,
@@ -358,7 +359,10 @@ export const chunkExternalMessages = <T>(
   return results;
 };
 
-export const shallowArrayEqual = (a: unknown[], b: unknown[]) => {
+export const shallowArrayEqual = (
+  a: readonly unknown[],
+  b: readonly unknown[],
+) => {
   if (a.length !== b.length) return false;
   for (let i = 0; i < a.length; i++) {
     if (a[i] !== b[i]) return false;
@@ -378,9 +382,18 @@ export const convertExternalMessageChunk = <T>(
   isRunning: boolean,
   error: ReadonlyJSONValue | undefined,
   cache?: ExternalMessageConversionCache,
+  cancelledMessageIds?: ReadonlySet<string>,
 ) => {
   const isLast = idx === chunkCount - 1;
   const joined = joinExternalMessages(message.outputs);
+  const isCancelled =
+    cancelledMessageIds !== undefined &&
+    message.outputs.some(
+      (output) =>
+        output.role !== "tool" &&
+        output.id != null &&
+        cancelledMessageIds.has(output.id),
+    );
   const hasInterruptedToolCalls =
     typeof joined.content === "object" &&
     joined.content.some(isInterruptedToolCall);
@@ -393,6 +406,7 @@ export const convertExternalMessageChunk = <T>(
     hasInterruptedToolCalls,
     hasPendingToolCalls,
     isLast ? error : undefined,
+    isCancelled,
   );
   const fallbackId = `${FALLBACK_ID_PREFIX}${idx}`;
 
@@ -488,6 +502,8 @@ export const convertExternalMessages = <T extends WeakKey>(
       chunks.length,
       isRunning,
       metadata.error,
+      undefined,
+      metadata.cancelledMessageIds,
     ),
   );
   return completeExternalMessageConversion(result, metadata.error);

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

@@ -134,6 +134,16 @@ export class MessageRepository {
     }
   }
 
+  private selectPathTo(message: RepositoryMessage) {
+    for (
+      let current: RepositoryMessage | null = message;
+      current;
+      current = current.prev
+    ) {
+      (current.prev ?? this.root).next = current;
+    }
+  }
+
   private performOp(
     newParent: RepositoryMessage | null,
     child: RepositoryMessage,
@@ -182,12 +192,17 @@ export class MessageRepository {
         child.current.id,
       ];
 
-      if (findHead(child) === this.head || newParentOrRoot.next === null) {
+      child.prev = newParent;
+
+      if (findHead(child) === this.head) {
+        this.selectPathTo(child);
+      } else if (newParentOrRoot.next === null) {
         newParentOrRoot.next = child;
+        if (this.head === newParentOrRoot) {
+          this.head = findHead(child);
+        }
       }
 
-      child.prev = newParent;
-
       const newLevel = newParent ? newParent.level + 1 : 0;
       this.updateLevels(child, newLevel);
     }
@@ -381,8 +396,7 @@ export class MessageRepository {
       );
 
     const previousHead = this.head;
-    const prevOrRoot = message.prev ?? this.root;
-    prevOrRoot.next = message;
+    this.selectPathTo(message);
 
     this.head = findHead(message);
 
@@ -423,17 +437,7 @@ export class MessageRepository {
     }
 
     this.head = message;
-    for (
-      let current: RepositoryMessage | null = message;
-      current;
-      current = current.prev
-    ) {
-      if (current.prev) {
-        current.prev.next = current;
-      } else {
-        this.root.next = current;
-      }
-    }
+    this.selectPathTo(message);
 
     this.evictOffBranchOptimisticMessages(previousHead, this.head);
 
@@ -456,7 +460,15 @@ export class MessageRepository {
     // 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) {
+    // Import and external-state conversion require parents before children, so
+    // the tree is walked in pre-order rather than iterated in insertion order.
+    const pending = [...this.root.children].reverse();
+    while (pending.length > 0) {
+      const message = this.messages.get(pending.pop()!);
+      if (!message) continue;
+      for (let i = message.children.length - 1; i >= 0; i--) {
+        pending.push(message.children[i]!);
+      }
       if (message.current.metadata?.isOptimistic) continue;
       let prev = message.prev;
       while (prev && prev.current.metadata?.isOptimistic) {

package/src/runtime/utils/resolveToolApprovalResponse.test.ts

@@ -88,4 +88,77 @@ describe("resolveToolApprovalResponse", () => {
       resolveToolApprovalResponse({ id: "a2" }, { optionId: "once" }),
     ).toThrow('no option with id "once"');
   });
+
+  it("resolves a free-form answer as approved and carries the text", () => {
+    expect(
+      resolveToolApprovalResponse(
+        { id: "q1", display: "text" },
+        { text: "staging" },
+      ),
+    ).toEqual({ approvalId: "q1", approved: true, text: "staging" });
+  });
+
+  it("carries the text alongside a chosen option", () => {
+    expect(
+      resolveToolApprovalResponse(
+        {
+          id: "q2",
+          display: "select",
+          allowFreeform: true,
+          options: [{ id: "other", kind: "_other" }],
+        },
+        { optionId: "other", approved: true, text: "somewhere else" },
+      ),
+    ).toEqual({
+      approvalId: "q2",
+      approved: true,
+      optionId: "other",
+      text: "somewhere else",
+    });
+  });
+
+  it("throws when the request does not accept a free-form answer", () => {
+    expect(() =>
+      resolveToolApprovalResponse(approval, { text: "anything" }),
+    ).toThrow("does not accept a free-form answer");
+    expect(() =>
+      resolveToolApprovalResponse(
+        { id: "a4", display: "select" },
+        { approved: true, text: "anything" },
+      ),
+    ).toThrow("does not accept a free-form answer");
+  });
+
+  it("refuses to infer approval from a bare answer on a decision", () => {
+    expect(() =>
+      resolveToolApprovalResponse(
+        { id: "a5", display: "decision", allowFreeform: true },
+        { text: "not this one" },
+      ),
+    ).toThrow("is a decision, not a question");
+    expect(() =>
+      resolveToolApprovalResponse(
+        { id: "a6", allowFreeform: true },
+        { text: "not this one" },
+      ),
+    ).toThrow("is a decision, not a question");
+  });
+
+  it("records the answer alongside an explicit decision on a gate", () => {
+    expect(
+      resolveToolApprovalResponse(
+        { id: "a7", display: "decision", allowFreeform: true },
+        { approved: false, text: "not this one" },
+      ),
+    ).toEqual({ approvalId: "a7", approved: false, text: "not this one" });
+  });
+
+  it("keeps a refusal a refusal when the request accepts text", () => {
+    expect(
+      resolveToolApprovalResponse(
+        { id: "q3", allowFreeform: true },
+        { approved: false, reason: "not answering" },
+      ),
+    ).toEqual({ approvalId: "q3", approved: false, reason: "not answering" });
+  });
 });

package/src/runtime/utils/resolveToolApprovalResponse.ts

@@ -1,6 +1,8 @@
-import type {
-  ToolApprovalOption,
-  ToolApprovalResponse,
+import {
+  toolApprovalAcceptsText,
+  type ToolApprovalDisplay,
+  type ToolApprovalOption,
+  type ToolApprovalResponse,
 } from "../../types/message";
 import type { RespondToToolApprovalOptions } from "../interfaces/thread-runtime-core";
 
@@ -12,16 +14,25 @@ const APPROVED_BY_KIND: Record<string, boolean> = {
 };
 
 /**
- * Resolves a renderer-facing approval response (boolean or optionId form)
- * against the approval's option list into the runtime-facing decision shape.
+ * Resolves a renderer-facing approval response (boolean, optionId, or
+ * free-form answer) against the approval's request shape into the
+ * runtime-facing decision shape.
  */
 export const resolveToolApprovalResponse = (
   approval: {
     readonly id: string;
+    readonly display?: ToolApprovalDisplay;
+    readonly allowFreeform?: boolean;
     readonly options?: readonly ToolApprovalOption[];
   },
   response: ToolApprovalResponse,
 ): RespondToToolApprovalOptions => {
+  const text = response.text;
+  if (text !== undefined && !toolApprovalAcceptsText(approval))
+    throw new Error(
+      `Tool approval "${approval.id}" does not accept a free-form answer; the request must declare display "text" or allowFreeform`,
+    );
+
   let approved: boolean;
   let optionId: string | undefined;
 
@@ -42,14 +53,24 @@ export const resolveToolApprovalResponse = (
       approved = APPROVED_BY_KIND[option.kind]!;
     }
     optionId = option.id;
-  } else {
+  } else if ("approved" in response) {
     approved = response.approved;
+  } else {
+    // A bare answer resolves a question, where answering is not refusing. On a
+    // decision the approval itself is the authorization, so inferring one from
+    // a typed note would let the note authorize the call.
+    if (approval.display !== "text" && approval.display !== "select")
+      throw new Error(
+        `Tool approval "${approval.id}" is a decision, not a question; respond with an explicit approved value, optionally alongside the answer`,
+      );
+    approved = true;
   }
 
   return {
     approvalId: approval.id,
     approved,
     ...(optionId !== undefined && { optionId }),
+    ...(text !== undefined && { text }),
     ...(response.reason != null && { reason: response.reason }),
   };
 };

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

@@ -23,8 +23,8 @@ import type {
   MessageTiming,
   PartProviderMetadata,
   TextMessagePart,
-  ToolApprovalOption,
   ToolCallTiming,
+  ToolCallMessagePart,
   ToolCallMessagePartMcpMetadata,
 } from "../../types/message";
 import type {
@@ -66,15 +66,7 @@ export type ThreadMessageLike = {
             readonly timing?: ToolCallTiming;
             readonly mcp?: ToolCallMessagePartMcpMetadata;
             readonly providerMetadata?: PartProviderMetadata;
-            readonly approval?: {
-              readonly id: string;
-              readonly approved?: boolean;
-              readonly reason?: string;
-              readonly isAutomatic?: boolean;
-              readonly options?: readonly ToolApprovalOption[];
-              readonly optionId?: string;
-              readonly resolution?: "cancelled" | "expired";
-            };
+            readonly approval?: NonNullable<ToolCallMessagePart["approval"]>;
           }
       )[];
   readonly id?: string | undefined;
@@ -187,7 +179,7 @@ export const fromThreadMessageLike = (
                 const { parentId, messages, ...basePart } = part;
                 const commonProps = {
                   ...basePart,
-                  toolCallId: part.toolCallId ?? `tool-${generateId()}`,
+                  toolCallId: part.toolCallId || `tool-${generateId()}`,
                   ...(parentId !== undefined && { parentId }),
                   ...(messages !== undefined && { messages }),
                 };

package/src/runtime/utils/tool-call-tree.test.ts

@@ -0,0 +1,191 @@
+import { describe, expect, it } from "vitest";
+import type {
+  ThreadAssistantMessage,
+  ThreadAssistantMessagePart,
+  ThreadMessage,
+  ToolCallMessagePart,
+} from "../../types/message";
+import {
+  iterateToolCallParts,
+  mapToolCallPartsDeep,
+  walkToolCallTree,
+} from "./tool-call-tree";
+
+const toolCall = (
+  toolCallId: string,
+  messages?: readonly ThreadMessage[],
+): ToolCallMessagePart => ({
+  type: "tool-call",
+  toolCallId,
+  toolName: "t",
+  argsText: "{}",
+  args: {},
+  ...(messages !== undefined && { messages }),
+});
+
+const assistant = (
+  id: string,
+  content: readonly ThreadAssistantMessagePart[],
+): ThreadAssistantMessage =>
+  ({
+    id,
+    role: "assistant",
+    content,
+    status: { type: "complete", reason: "unknown" },
+    metadata: {
+      unstable_state: null,
+      unstable_annotations: [],
+      unstable_data: [],
+      steps: [],
+      custom: {},
+    },
+    createdAt: new Date(0),
+    attachments: [],
+  }) as unknown as ThreadAssistantMessage;
+
+const user = (id: string): ThreadMessage =>
+  ({
+    id,
+    role: "user",
+    content: [{ type: "text", text: "hi" }],
+    metadata: { custom: {} },
+    createdAt: new Date(0),
+    attachments: [],
+  }) as unknown as ThreadMessage;
+
+// root
+// └─ a ─ nested message "m2"
+//         ├─ b ─ nested message "m3"
+//         │       └─ c
+//         └─ d
+const tree = () => {
+  const c = toolCall("c");
+  const b = toolCall("b", [assistant("m3", [c])]);
+  const d = toolCall("d");
+  const a = toolCall("a", [assistant("m2", [b, d])]);
+  return { a, b, c, d };
+};
+
+describe("walkToolCallTree", () => {
+  it("yields every part in document order, each ahead of its descendants", () => {
+    const { a } = tree();
+    const entries = [
+      ...walkToolCallTree([user("m0"), assistant("m1", [a, toolCall("e")])]),
+    ];
+
+    expect(entries.map((entry) => entry.part.toolCallId)).toEqual([
+      "a",
+      "b",
+      "c",
+      "d",
+      "e",
+    ]);
+  });
+
+  it("reports the message that directly holds each part", () => {
+    const { a } = tree();
+    const entries = [...walkToolCallTree([assistant("m1", [a])])];
+
+    expect(
+      Object.fromEntries(
+        entries.map((entry) => [entry.part.toolCallId, entry.messageId]),
+      ),
+    ).toEqual({ a: "m1", b: "m2", c: "m3", d: "m2" });
+  });
+
+  it("skips messages whose content is not an array", () => {
+    const malformed = {
+      id: "m1",
+      role: "assistant",
+    } as unknown as ThreadMessage;
+    expect([...walkToolCallTree([malformed])]).toEqual([]);
+  });
+
+  it("skips a non-assistant message carrying a tool-call part", () => {
+    const malformed = {
+      ...user("m2"),
+      content: [toolCall("b")],
+    } as unknown as ThreadMessage;
+    const a = toolCall("a", [malformed]);
+
+    expect(
+      [...walkToolCallTree([assistant("m1", [a])])].map(
+        (entry) => entry.part.toolCallId,
+      ),
+    ).toEqual(["a"]);
+  });
+});
+
+describe("iterateToolCallParts", () => {
+  it("matches walkToolCallTree over the same content", () => {
+    const { a } = tree();
+    const content = [a, toolCall("e")];
+
+    expect([...iterateToolCallParts(content)].map((p) => p.toolCallId)).toEqual(
+      [...walkToolCallTree([assistant("m1", content)])].map(
+        (entry) => entry.part.toolCallId,
+      ),
+    );
+  });
+
+  it("ignores non tool-call parts", () => {
+    expect([
+      ...iterateToolCallParts([
+        { type: "text", text: "hi" } as ThreadAssistantMessagePart,
+        toolCall("a"),
+      ]),
+    ]).toHaveLength(1);
+  });
+});
+
+describe("mapToolCallPartsDeep", () => {
+  it("rewrites nested parts and reports the change", () => {
+    const { a } = tree();
+    const content = [a];
+    const { content: next, changed } = mapToolCallPartsDeep(content, (part) =>
+      part.toolCallId === "c" ? { ...part, isError: true } : part,
+    );
+
+    expect(changed).toBe(true);
+    const parts = [...iterateToolCallParts(next)];
+    expect(parts.find((part) => part.toolCallId === "c")?.isError).toBe(true);
+    expect(
+      parts.find((part) => part.toolCallId === "d")?.isError,
+    ).toBeUndefined();
+  });
+
+  it("skips a nested message whose content is not an array", () => {
+    const malformed = {
+      id: "m2",
+      role: "assistant",
+    } as unknown as ThreadMessage;
+    const a = toolCall("a", [malformed]);
+
+    expect(() =>
+      mapToolCallPartsDeep([a], (part) => ({ ...part, isError: true })),
+    ).not.toThrow();
+  });
+
+  it("keeps the original content identity when nothing changed", () => {
+    const { a } = tree();
+    const content = [a, toolCall("e")];
+    const result = mapToolCallPartsDeep(content, (part) => part);
+
+    expect(result.changed).toBe(false);
+    expect(result.content).toBe(content);
+  });
+
+  it("preserves untouched sibling messages by identity", () => {
+    const nested = assistant("m2", [toolCall("b")]);
+    const sibling = user("m3");
+    const a = toolCall("a", [nested, sibling]);
+    const { content: next } = mapToolCallPartsDeep([a], (part) =>
+      part.toolCallId === "b" ? { ...part, isError: true } : part,
+    );
+
+    const mapped = next[0] as ToolCallMessagePart;
+    expect(mapped).not.toBe(a);
+    expect(mapped.messages?.[1]).toBe(sibling);
+    expect(mapped.messages?.[0]).not.toBe(nested);
+  });
+});