@assistant-ui/core 0.3.14 → 0.3.16

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),
+          ),
         };
       }),
     };
@@ -81,9 +84,9 @@ type RepositoryMessage = RepositoryParent & {
 const findHead = (
   message: RepositoryMessage | RepositoryParent,
 ): RepositoryMessage | null => {
-  if (message.next) return findHead(message.next);
-  if ("current" in message) return message;
-  return null;
+  let current = message;
+  while (current.next) current = current.next;
+  return "current" in current ? current : null;
 };
 
 class CachedValue<T> {
@@ -116,11 +119,17 @@ export class MessageRepository {
   };
 
   private updateLevels(message: RepositoryMessage, newLevel: number) {
-    message.level = newLevel;
-    for (const childId of message.children) {
-      const childMessage = this.messages.get(childId);
-      if (childMessage) {
-        this.updateLevels(childMessage, newLevel + 1);
+    const pending = [{ message, level: newLevel }];
+
+    while (pending.length > 0) {
+      const current = pending.pop()!;
+      current.message.level = current.level;
+
+      for (const childId of current.message.children) {
+        const childMessage = this.messages.get(childId);
+        if (childMessage) {
+          pending.push({ message: childMessage, level: current.level + 1 });
+        }
       }
     }
   }
@@ -135,7 +144,8 @@ export class MessageRepository {
 
     if (operation === "relink" && parentOrRoot === newParentOrRoot) return;
 
-    if (operation !== "cut") {
+    // `link` receives a fresh ID from `addOrUpdateMessage`; only `relink` can introduce a cycle.
+    if (operation === "relink") {
       for (
         let current: RepositoryMessage | null = newParent;
         current;
@@ -143,7 +153,7 @@ export class MessageRepository {
       ) {
         if (current.current.id === child.current.id) {
           throw new Error(
-            "MessageRepository(performOp/link): A message with the same id already exists in the parent tree. This error occurs if the same message id is found multiple times. This is likely an internal bug in assistant-ui.",
+            "MessageRepository(performOp/relink): A message with the same id already exists in the parent tree. This error occurs if the same message id is found multiple times. This is likely an internal bug in assistant-ui.",
           );
         }
       }
@@ -396,16 +406,17 @@ export class MessageRepository {
     const previousHead = this.head;
 
     if (message.children.length > 0) {
-      const deleteDescendants = (msg: RepositoryMessage) => {
-        for (const childId of msg.children) {
-          const childMessage = this.messages.get(childId);
-          if (childMessage) {
-            deleteDescendants(childMessage);
-            this.messages.delete(childId);
+      const pending = [...message.children];
+      while (pending.length > 0) {
+        const childId = pending.pop()!;
+        const childMessage = this.messages.get(childId);
+        if (childMessage) {
+          for (const descendantId of childMessage.children) {
+            pending.push(descendantId);
           }
+          this.messages.delete(childId);
         }
-      };
-      deleteDescendants(message);
+      }
 
       message.children = [];
       message.next = null;

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/thread-message-like.ts

@@ -25,6 +25,7 @@ import type {
   TextMessagePart,
   ToolApprovalOption,
   ToolCallTiming,
+  ToolCallMessagePartMcpMetadata,
 } from "../../types/message";
 import type {
   ReadonlyJSONObject,
@@ -63,6 +64,7 @@ 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;

package/src/runtimes/external-store/external-store-thread-list-runtime-core.ts

@@ -1,4 +1,3 @@
-import type { Unsubscribe } from "../../types/unsubscribe";
 import type { ExternalStoreThreadRuntimeCore } from "./external-store-thread-runtime-core";
 import type {
   ThreadListItemCoreState,
@@ -6,7 +5,7 @@ import type {
 } from "../../runtime/interfaces/thread-list-runtime-core";
 import type { ExternalStoreThreadListAdapter } from "./external-store-adapter";
 import { invalidateThreadRuntime } from "../../runtime/utils/thread-runtime-lifecycle";
-import { notifySubscribers } from "../../subscribable/subscribable";
+import { BaseSubscribable } from "../../subscribable/subscribable";
 
 export type ExternalStoreThreadFactory = () => ExternalStoreThreadRuntimeCore;
 
@@ -24,7 +23,10 @@ const DEFAULT_THREAD_DATA = Object.freeze({
   [DEFAULT_THREAD_ID]: DEFAULT_THREAD,
 });
 
-export class ExternalStoreThreadListRuntimeCore implements ThreadListRuntimeCore {
+export class ExternalStoreThreadListRuntimeCore
+  extends BaseSubscribable
+  implements ThreadListRuntimeCore
+{
   private _mainThreadId: string = DEFAULT_THREAD_ID;
   private _threads: readonly string[] = DEFAULT_THREADS;
   private _archivedThreads: readonly string[] = EMPTY_ARRAY;
@@ -68,6 +70,7 @@ export class ExternalStoreThreadListRuntimeCore implements ThreadListRuntimeCore
     adapter: ExternalStoreThreadListAdapter = {},
     threadFactory: ExternalStoreThreadFactory,
   ) {
+    super();
     this.threadFactory = threadFactory;
     this.__internal_setAdapter(adapter, true);
   }
@@ -262,15 +265,4 @@ export class ExternalStoreThreadListRuntimeCore implements ThreadListRuntimeCore
   public generateTitle(): never {
     throw new Error("Method not implemented.");
   }
-
-  private _subscriptions = new Set<() => void>();
-
-  public subscribe(callback: () => void): Unsubscribe {
-    this._subscriptions.add(callback);
-    return () => this._subscriptions.delete(callback);
-  }
-
-  private _notifySubscribers() {
-    notifySubscribers(this._subscriptions);
-  }
 }