@assistant-ui/core 0.3.15 → 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),
+          ),
         };
       }),
     };

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