@assistant-ui/core 0.3.15 → 0.3.17

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

@@ -0,0 +1,176 @@
+import type { ThreadMessage } from "../../types/message";
+import {
+  MessageRepository,
+  type ExportedMessageRepository,
+} from "./message-repository";
+
+export type MessageRepositorySessionOptions = {
+  decorateExport?: (
+    exported: ExportedMessageRepository,
+    repository: MessageRepository,
+  ) => ExportedMessageRepository;
+};
+
+export const createMessageRepositorySession = (
+  options: MessageRepositorySessionOptions = {},
+) => {
+  const repository = new MessageRepository();
+  let exportedRepository: ExportedMessageRepository | undefined;
+
+  const getMessages = (): readonly ThreadMessage[] => {
+    return repository.getMessages();
+  };
+
+  const tryGetMessage = (messageId: string) => {
+    try {
+      return repository.getMessage(messageId);
+    } catch {
+      return undefined;
+    }
+  };
+
+  const tryGetMessages = (
+    messageId: string,
+  ): readonly ThreadMessage[] | undefined => {
+    try {
+      return repository.getMessages(messageId);
+    } catch {
+      return undefined;
+    }
+  };
+
+  const hasMessage = (messageId: string): boolean => {
+    return tryGetMessage(messageId) !== undefined;
+  };
+
+  const addOrUpdateMessage = (
+    parentId: string | null,
+    message: ThreadMessage,
+  ): void => {
+    repository.addOrUpdateMessage(parentId, message);
+    exportedRepository = undefined;
+  };
+
+  const deleteMessage = (
+    messageId: string,
+    replacementId?: string | null,
+  ): void => {
+    repository.deleteMessage(messageId, replacementId);
+    exportedRepository = undefined;
+  };
+
+  const tryDeleteMessage = (messageId: string): boolean => {
+    if (!hasMessage(messageId)) return false;
+    deleteMessage(messageId);
+    return true;
+  };
+
+  const switchToBranch = (messageId: string): void => {
+    repository.switchToBranch(messageId);
+    exportedRepository = undefined;
+  };
+
+  const resetHead = (messageId: string | null): void => {
+    repository.resetHead(messageId);
+    exportedRepository = undefined;
+  };
+
+  const clear = (): void => {
+    repository.clear();
+    exportedRepository = undefined;
+  };
+
+  const updateMessage = (
+    messageId: string,
+    updater: (message: ThreadMessage) => ThreadMessage,
+  ): boolean => {
+    const item = tryGetMessage(messageId);
+    if (!item) return false;
+    const message = updater(item.message);
+    if (message === item.message) return false;
+    addOrUpdateMessage(item.parentId, message);
+    return true;
+  };
+
+  const applyExternalMessageRepository = (
+    loaded: ExportedMessageRepository,
+  ): void => {
+    const headId = loaded.headId ?? loaded.messages.at(-1)?.message.id ?? null;
+    const ids = new Set<string>();
+    let degenerate = false;
+    for (const { message } of loaded.messages) {
+      if (ids.has(message.id)) {
+        degenerate = true;
+        break;
+      }
+      ids.add(message.id);
+    }
+    if (headId !== null && !ids.has(headId)) degenerate = true;
+
+    if (!degenerate) {
+      clear();
+      let pending = [...loaded.messages];
+      const importedIds = new Set<string>();
+
+      while (pending.length > 0) {
+        const unresolved: typeof pending = [];
+        let progressed = false;
+        for (const item of pending) {
+          if (item.parentId !== null && !importedIds.has(item.parentId)) {
+            unresolved.push(item);
+            continue;
+          }
+          addOrUpdateMessage(item.parentId, item.message);
+          importedIds.add(item.message.id);
+          progressed = true;
+        }
+        if (!progressed) {
+          degenerate = true;
+          break;
+        }
+        pending = unresolved;
+      }
+    }
+
+    if (degenerate) {
+      clear();
+      let previousId: string | null = null;
+      for (const { message } of loaded.messages) {
+        const existing = tryGetMessage(message.id);
+        addOrUpdateMessage(existing ? existing.parentId : previousId, message);
+        previousId = message.id;
+      }
+      resetHead(previousId);
+    } else {
+      resetHead(headId);
+    }
+  };
+
+  return {
+    getMessages,
+    get headId() {
+      return repository.headId;
+    },
+    export: (): ExportedMessageRepository => {
+      exportedRepository ??= options.decorateExport
+        ? options.decorateExport(repository.export(), repository)
+        : repository.export();
+      return exportedRepository;
+    },
+    tryGetMessage,
+    tryGetMessages,
+    hasMessage,
+    addOrUpdateMessage,
+    deleteMessage,
+    tryDeleteMessage,
+    switchToBranch,
+    resetHead,
+    clear,
+    updateMessage,
+    applyExternalMessageRepository,
+  };
+};
+
+export type MessageRepositorySession = ReturnType<
+  typeof createMessageRepositorySession
+>;

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

@@ -2,7 +2,7 @@ import type { ThreadMessage } from "../../types/message";
 import type { RunConfig } from "../../types/message";
 import { generateId } from "../../utils/id";
 import type { ThreadMessageLike } from "./thread-message-like";
-import { getAutoStatus } from "./auto-status";
+import { getContentAutoStatus } from "./auto-status";
 import { fromThreadMessageLike } from "./thread-message-like";
 
 export type ExportedMessageRepositoryItem = {
@@ -28,7 +28,7 @@ export const ExportedMessageRepository = {
       fromThreadMessageLike(
         m,
         generateId(),
-        getAutoStatus(false, false, false, false, undefined),
+        getContentAutoStatus(m.content, false, false),
       ),
     );
 
@@ -47,7 +47,6 @@ export const ExportedMessageRepository = {
     }[],
     options?: { headId?: string | null },
   ): ExportedMessageRepository => {
-    const fallbackStatus = getAutoStatus(false, false, false, false, undefined);
     return {
       ...(options?.headId !== undefined
         ? { headId: options.headId }
@@ -60,7 +59,11 @@ export const ExportedMessageRepository = {
         }
         return {
           parentId,
-          message: fromThreadMessageLike(message, message.id, fallbackStatus),
+          message: fromThreadMessageLike(
+            message,
+            message.id,
+            getContentAutoStatus(message.content, false, false),
+          ),
         };
       }),
     };

package/src/runtime/utils/pending-tool-calls.test.ts

@@ -0,0 +1,57 @@
+import { describe, expect, it } from "vitest";
+import {
+  createToolCallCancellationStub,
+  scanPendingToolCalls,
+} from "./pending-tool-calls";
+
+type Message =
+  | { type: "assistant"; toolCalls: { id: string; name: string }[] }
+  | { type: "tool"; toolCallId: string };
+
+const getPendingToolCalls = (messages: readonly Message[]) =>
+  scanPendingToolCalls(
+    messages,
+    (message) =>
+      message.type === "assistant"
+        ? { toolCalls: message.toolCalls }
+        : { toolCallId: message.toolCallId },
+    (toolCall) => toolCall.id,
+  );
+
+describe("scanPendingToolCalls", () => {
+  it("returns tool calls without matching results in history order", () => {
+    expect(
+      getPendingToolCalls([
+        {
+          type: "assistant",
+          toolCalls: [
+            { id: "call-1", name: "first" },
+            { id: "call-2", name: "second" },
+          ],
+        },
+        { type: "tool", toolCallId: "call-1" },
+        {
+          type: "assistant",
+          toolCalls: [{ id: "call-3", name: "third" }],
+        },
+      ]),
+    ).toEqual([
+      { id: "call-2", name: "second" },
+      { id: "call-3", name: "third" },
+    ]);
+  });
+});
+
+describe("createToolCallCancellationStub", () => {
+  it("creates an error tool result with the cancellation payload", () => {
+    expect(
+      createToolCallCancellationStub({ id: "call-1", name: "lookup" }),
+    ).toEqual({
+      type: "tool",
+      name: "lookup",
+      tool_call_id: "call-1",
+      content: JSON.stringify({ cancelled: true }),
+      status: "error",
+    });
+  });
+});

package/src/runtime/utils/pending-tool-calls.ts

@@ -0,0 +1,35 @@
+type PendingToolCallMessage<TToolCall> =
+  | { toolCalls: readonly TToolCall[] }
+  | { toolCallId: string }
+  | undefined;
+
+export const scanPendingToolCalls = <TMessage, TToolCall>(
+  messages: readonly TMessage[],
+  getMessage: (message: TMessage) => PendingToolCallMessage<TToolCall>,
+  getToolCallId: (toolCall: TToolCall) => string,
+): TToolCall[] => {
+  const pending = new Map<string, TToolCall>();
+  for (const message of messages) {
+    const value = getMessage(message);
+    if (!value) continue;
+    if ("toolCalls" in value) {
+      for (const toolCall of value.toolCalls) {
+        pending.set(getToolCallId(toolCall), toolCall);
+      }
+    } else {
+      pending.delete(value.toolCallId);
+    }
+  }
+  return [...pending.values()];
+};
+
+export const createToolCallCancellationStub = (toolCall: {
+  readonly id: string;
+  readonly name: string;
+}) => ({
+  type: "tool" as const,
+  name: toolCall.name,
+  tool_call_id: toolCall.id,
+  content: JSON.stringify({ cancelled: true }),
+  status: "error" as const,
+});

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,9 @@ import type {
   MessageTiming,
   PartProviderMetadata,
   TextMessagePart,
-  ToolApprovalOption,
   ToolCallTiming,
+  ToolCallMessagePart,
+  ToolCallMessagePartMcpMetadata,
 } from "../../types/message";
 import type {
   ReadonlyJSONObject,
@@ -63,16 +64,9 @@ export type ThreadMessageLike = {
             readonly messages?: readonly ThreadMessage[] | undefined;
             readonly interrupt?: { type: "human"; payload: unknown };
             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;

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);
+  });
+});