@assistant-ui/core 0.3.15 → 0.3.17

package/src/react/runtimes/external-message-converter.ts

@@ -1,432 +1,46 @@
 "use client";
 
 import { useMemo, useState } from "react";
-import { ThreadMessageConverter } from "../../runtimes/external-store/thread-message-converter";
-import {
-  getExternalStoreMessages,
-  symbolInnerMessage,
-  bindExternalStoreMessage,
-} from "../../runtime/utils/external-store-message";
 import {
-  fromThreadMessageLike,
-  type ThreadMessageLike,
-} from "../../runtime/utils/thread-message-like";
-import { getAutoStatus, isAutoStatus } from "../../runtime/utils/auto-status";
-import type { ToolExecutionStatus } from "../../runtimes/tool-invocations/ToolInvocationTracker";
-import type { ReadonlyJSONValue } from "assistant-stream/utils";
-import { generateErrorMessageId } from "../../utils/id";
-import type {
-  ThreadAssistantMessage,
-  ThreadMessage,
-  ToolCallMessagePart,
-} from "../../types/message";
-import type { MessageTiming } from "../../types/message";
-
-export type JoinStrategy = "concat-content" | "none";
-
-type ThreadMessageLikeContentItem = Exclude<
-  ThreadMessageLike["content"],
-  string
->[number];
+  chunkExternalMessages,
+  completeExternalMessageConversion,
+  convertExternalMessageCallback,
+  convertExternalMessageChunk,
+  convertExternalMessages as convertExternalMessagesInternal,
+  shallowArrayEqual,
+  type ExternalMessageConverterCallback,
+  type ExternalMessageConverterCallbackResult,
+  type ExternalMessageConverterChunk,
+  type ExternalMessageConverterMessage,
+  type ExternalMessageConverterMetadata,
+  type JoinStrategy,
+} from "../../runtime/utils/external-message-conversion";
+import { bindExternalStoreMessage } from "../../runtime/utils/external-store-message";
+import { ThreadMessageConverter } from "../../runtimes/external-store/thread-message-converter";
+import type { ThreadMessage } from "../../types/message";
 
-const isPendingToolCall = (c: ThreadMessageLikeContentItem): boolean =>
-  c.type === "tool-call" && c.result === undefined;
+// Generatedness is tracked by identity, not by id shape: a caller-supplied id
+// that happens to match the generated pattern must never be rewritten.
+const generatedFallbackMessages = new WeakSet<object>();
 
-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)
-  );
-};
+export type { JoinStrategy };
 
 export namespace useExternalMessageConverter {
-  export type Message =
-    | (ThreadMessageLike & {
-        readonly convertConfig?: {
-          readonly joinStrategy?: JoinStrategy;
-        };
-      })
-    | {
-        role: "tool";
-        toolCallId: string;
-        toolName?: string | undefined;
-        result: any;
-        artifact?: any;
-        isError?: boolean;
-        messages?: readonly ThreadMessage[];
-      };
-
-  export type Metadata = {
-    readonly toolStatuses?: Record<string, ToolExecutionStatus>;
-    readonly error?: ReadonlyJSONValue;
-    readonly messageTiming?: Record<string, MessageTiming>;
-  };
-
-  export type Callback<T> = (
-    message: T,
-    metadata: Metadata,
-  ) => Message | Message[];
-}
-
-type CallbackResult<T> = {
-  input: T;
-  outputs: useExternalMessageConverter.Message[];
-};
-
-const stringifyForError = (value: unknown) => {
-  let text;
-  try {
-    text = value instanceof Error ? String(value) : JSON.stringify(value);
-  } catch {
-    /* circular */
-  }
-  text ??= String(value);
-  return text.length > 200 ? `${text.slice(0, 200)}…` : text;
-};
-
-const toCallbackOutputs = (
-  output:
-    | useExternalMessageConverter.Message
-    | useExternalMessageConverter.Message[],
-  input: unknown,
-): useExternalMessageConverter.Message[] => {
-  const outputs = Array.isArray(output) ? output : [output];
-  for (const o of outputs) {
-    const valid =
-      typeof o === "object" &&
-      o !== null &&
-      (o.role === "tool" ||
-        ((o.role === "assistant" || o.role === "user" || o.role === "system") &&
-          (typeof o.content === "string" || Array.isArray(o.content))));
-    if (!valid)
-      throw new Error(
-        `External message converter: the converter callback returned an invalid message (${stringifyForError(o)}) for input ${stringifyForError(input)}. Return an empty array to skip a message.`,
-      );
-  }
-  return outputs;
-};
-
-type CallbackCacheEntry<T> = CallbackResult<T> & {
-  metadata: useExternalMessageConverter.Metadata;
-  callback: useExternalMessageConverter.Callback<T>;
-};
-
-type ChunkResult<T> = {
-  inputs: T[];
-  outputs: useExternalMessageConverter.Message[];
-};
-
-type Mutable<T> = {
-  -readonly [P in keyof T]: T[P];
-};
-
-const mergeInnerMessages = (existing: object, incoming: object) => ({
-  [symbolInnerMessage]: [
-    ...((existing as any)[symbolInnerMessage] ?? []),
-    ...((incoming as any)[symbolInnerMessage] ?? []),
-  ],
-});
-
-const joinExternalMessages = (
-  messages: readonly useExternalMessageConverter.Message[],
-): ThreadMessageLike => {
-  const assistantMessage: Mutable<Omit<ThreadMessageLike, "metadata">> & {
-    content: Exclude<ThreadMessageLike["content"][0], string>[];
-    metadata?: Mutable<ThreadMessageLike["metadata"]>;
-  } = {
-    role: "assistant",
-    content: [],
-  };
-  for (const output of messages) {
-    if (output.role === "tool") {
-      const toolCallIdx = assistantMessage.content.findIndex(
-        (c) => c.type === "tool-call" && c.toolCallId === output.toolCallId,
-      );
-      // 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;
-        if (output.toolName != null) {
-          if (toolCall.toolName !== output.toolName)
-            throw new Error(
-              `Tool call name ${output.toolCallId} ${output.toolName} does not match existing tool call ${toolCall.toolName}`,
-            );
-        }
-        assistantMessage.content[toolCallIdx] = {
-          ...toolCall,
-          ...{
-            [symbolInnerMessage]: [
-              ...((toolCall as any)[symbolInnerMessage] ?? []),
-              output,
-            ],
-          },
-          result: output.result,
-          artifact: output.artifact,
-          isError: output.isError,
-          messages: output.messages,
-        };
-      }
-    } else {
-      const role = output.role;
-      const content = (
-        typeof output.content === "string"
-          ? [{ type: "text" as const, text: output.content }]
-          : output.content
-      ).map((c) => ({
-        ...c,
-        ...{ [symbolInnerMessage]: [output] },
-      }));
-      switch (role) {
-        case "system":
-        case "user":
-          return {
-            ...output,
-            content,
-          };
-        case "assistant":
-          if (assistantMessage.content.length === 0) {
-            assistantMessage.id = output.id;
-            assistantMessage.createdAt ??= output.createdAt;
-            assistantMessage.status ??= output.status;
-
-            if (output.attachments) {
-              assistantMessage.attachments = [
-                ...(assistantMessage.attachments ?? []),
-                ...output.attachments,
-              ];
-            }
-          }
-
-          if (output.metadata) {
-            assistantMessage.metadata ??= {};
-            if (output.metadata.unstable_state !== undefined) {
-              assistantMessage.metadata.unstable_state =
-                output.metadata.unstable_state;
-            }
-            if (output.metadata.unstable_annotations) {
-              assistantMessage.metadata.unstable_annotations = [
-                ...(assistantMessage.metadata.unstable_annotations ?? []),
-                ...output.metadata.unstable_annotations,
-              ];
-            }
-            if (output.metadata.unstable_data) {
-              assistantMessage.metadata.unstable_data = [
-                ...(assistantMessage.metadata.unstable_data ?? []),
-                ...output.metadata.unstable_data,
-              ];
-            }
-            if (output.metadata.steps) {
-              assistantMessage.metadata.steps = [
-                ...(assistantMessage.metadata.steps ?? []),
-                ...output.metadata.steps,
-              ];
-            }
-            if (output.metadata.custom) {
-              assistantMessage.metadata.custom = {
-                ...(assistantMessage.metadata.custom ?? {}),
-                ...output.metadata.custom,
-              };
-            }
-
-            if (output.metadata.timing) {
-              assistantMessage.metadata.timing = output.metadata.timing;
-            }
-
-            if (output.metadata.submittedFeedback) {
-              assistantMessage.metadata.submittedFeedback =
-                output.metadata.submittedFeedback;
-            }
-
-            if (output.metadata.isOptimistic) {
-              assistantMessage.metadata.isOptimistic = true;
-            }
-            // TODO keep this in sync with ThreadMessageLike["metadata"] / fromThreadMessageLike
-          }
-
-          // Add content parts, merging reasoning parts with same parentId
-          for (const part of content) {
-            if (part.type === "tool-call") {
-              const existingIdx = assistantMessage.content.findIndex(
-                (c) =>
-                  c.type === "tool-call" && c.toolCallId === part.toolCallId,
-              );
-              if (existingIdx !== -1) {
-                const existing = assistantMessage.content[
-                  existingIdx
-                ] as typeof part;
-                assistantMessage.content[existingIdx] = {
-                  ...existing,
-                  ...part,
-                  ...mergeInnerMessages(existing, part),
-                };
-                continue;
-              }
-            }
-
-            if (
-              part.type === "reasoning" &&
-              "parentId" in part &&
-              part.parentId
-            ) {
-              const existingIdx = assistantMessage.content.findIndex(
-                (c) =>
-                  c.type === "reasoning" &&
-                  "parentId" in c &&
-                  c.parentId === part.parentId,
-              );
-              if (existingIdx !== -1) {
-                const existing = assistantMessage.content[
-                  existingIdx
-                ] as typeof part;
-                assistantMessage.content[existingIdx] = {
-                  ...existing,
-                  text: `${existing.text}\n\n${part.text}`,
-                  ...mergeInnerMessages(existing, part),
-                };
-                continue;
-              }
-            }
-            assistantMessage.content.push(part);
-          }
-          break;
-        default: {
-          const unsupportedRole: never = role;
-          throw new Error(`Unknown message role: ${unsupportedRole}`);
-        }
-      }
-    }
-  }
-  return assistantMessage;
-};
-
-const chunkExternalMessages = <T>(
-  callbackResults: CallbackResult<T>[],
-  joinStrategy?: JoinStrategy,
-) => {
-  const results: ChunkResult<T>[] = [];
-  let isAssistant = false;
-  let pendingNone = false; // true if the previous assistant message had joinStrategy "none"
-  let inputs: T[] = [];
-  let outputs: useExternalMessageConverter.Message[] = [];
-
-  const flush = () => {
-    if (outputs.length) {
-      results.push({
-        inputs,
-        outputs,
-      });
-    }
-    inputs = [];
-    outputs = [];
-    isAssistant = false;
-    pendingNone = false;
-  };
-
-  for (const callbackResult of callbackResults) {
-    for (const output of callbackResult.outputs) {
-      if (
-        (pendingNone && output.role !== "tool") ||
-        !isAssistant ||
-        output.role === "user" ||
-        output.role === "system"
-      ) {
-        flush();
-      }
-      isAssistant = output.role === "assistant" || output.role === "tool";
-
-      if (inputs.at(-1) !== callbackResult.input) {
-        inputs.push(callbackResult.input);
-      }
-      outputs.push(output);
-
-      if (
-        output.role === "assistant" &&
-        (output.convertConfig?.joinStrategy === "none" ||
-          joinStrategy === "none")
-      ) {
-        pendingNone = true;
-      }
-    }
-  }
-  flush();
-  return results;
-};
-
-function createErrorAssistantMessage(
-  error: ReadonlyJSONValue,
-): ThreadAssistantMessage {
-  const msg: ThreadAssistantMessage = {
-    id: generateErrorMessageId(),
-    role: "assistant",
-    content: [],
-    status: { type: "incomplete", reason: "error", error },
-    createdAt: new Date(),
-    metadata: {
-      unstable_state: null,
-      unstable_annotations: [],
-      unstable_data: [],
-      custom: {},
-      steps: [],
-    },
-  };
-  bindExternalStoreMessage(msg, []);
-  return msg;
+  export type Message = ExternalMessageConverterMessage;
+  export type Metadata = ExternalMessageConverterMetadata;
+  export type Callback<T> = ExternalMessageConverterCallback<T>;
 }
 
-export const convertExternalMessages = <T extends WeakKey>(
+export const convertExternalMessages: <T extends WeakKey>(
   messages: T[],
   callback: useExternalMessageConverter.Callback<T>,
   isRunning: boolean,
   metadata: useExternalMessageConverter.Metadata,
-) => {
-  const callbackResults: CallbackResult<T>[] = [];
-  for (const message of messages) {
-    const output = callback(message, metadata);
-    const outputs = toCallbackOutputs(output, message);
-    const result = { input: message, outputs };
-    callbackResults.push(result);
-  }
-
-  const chunks = chunkExternalMessages(callbackResults);
-
-  const result = chunks.map((message, idx) => {
-    const isLast = idx === chunks.length - 1;
-    const joined = joinExternalMessages(message.outputs);
-    const hasInterruptedToolCalls =
-      typeof joined.content === "object" &&
-      joined.content.some(isInterruptedToolCall);
-    const hasPendingToolCalls =
-      typeof joined.content === "object" &&
-      joined.content.some(isPendingToolCall);
-    const autoStatus = getAutoStatus(
-      isLast,
-      isRunning,
-      hasInterruptedToolCalls,
-      hasPendingToolCalls,
-      isLast ? metadata.error : undefined,
-    );
-    const newMessage = fromThreadMessageLike(
-      joined,
-      idx.toString(),
-      autoStatus,
-    );
-    bindExternalStoreMessage(newMessage, message.inputs);
-    return newMessage;
-  });
+) => ThreadMessage[] = convertExternalMessagesInternal;
 
-  if (metadata.error) {
-    const lastMessage = result.at(-1);
-    if (!lastMessage || lastMessage.role !== "assistant") {
-      result.push(createErrorAssistantMessage(metadata.error));
-    }
-  }
-
-  return result;
+type CallbackCacheEntry<T> = ExternalMessageConverterCallbackResult<T> & {
+  metadata: useExternalMessageConverter.Metadata;
+  callback: useExternalMessageConverter.Callback<T>;
 };
 
 export const useExternalMessageConverter = <T extends WeakKey>({
@@ -452,8 +66,8 @@ export const useExternalMessageConverter = <T extends WeakKey>({
   const [caches] = useState(() => ({
     callbackCache: new WeakMap<T, CallbackCacheEntry<T>>(),
     chunkCache: new WeakMap<
-      useExternalMessageConverter.Message,
-      ChunkResult<T>
+      ExternalMessageConverterMessage,
+      ExternalMessageConverterChunk<T>
     >(),
     converterCache: new ThreadMessageConverter(),
   }));
@@ -468,7 +82,7 @@ export const useExternalMessageConverter = <T extends WeakKey>({
   );
 
   return useMemo(() => {
-    const callbackResults: CallbackResult<T>[] = [];
+    const callbackResults: ExternalMessageConverterCallbackResult<T>[] = [];
     for (const message of messages) {
       let result = state.callbackCache.get(message);
       if (
@@ -476,11 +90,12 @@ export const useExternalMessageConverter = <T extends WeakKey>({
         result.metadata !== state.metadata ||
         result.callback !== state.callback
       ) {
-        const output = state.callback(message, state.metadata);
-        const outputs = toCallbackOutputs(output, message);
         result = {
-          input: message,
-          outputs,
+          ...convertExternalMessageCallback(
+            message,
+            state.callback,
+            state.metadata,
+          ),
           metadata: state.metadata,
           callback: state.callback,
         };
@@ -490,77 +105,40 @@ export const useExternalMessageConverter = <T extends WeakKey>({
     }
 
     const chunks = chunkExternalMessages(callbackResults, joinStrategy).map(
-      (m) => {
-        const key = m.outputs[0];
-        if (!key) return m;
+      (message) => {
+        const key = message.outputs[0];
+        if (!key) return message;
 
         const cached = state.chunkCache.get(key);
-        if (cached && shallowArrayEqual(cached.outputs, m.outputs))
+        if (cached && shallowArrayEqual(cached.outputs, message.outputs)) {
           return cached;
-        state.chunkCache.set(key, m);
-        return m;
+        }
+        state.chunkCache.set(key, message);
+        return message;
       },
     );
 
     const threadMessages = state.converterCache.convertMessages(
       chunks,
-      (cache, message, idx) => {
-        const isLast = idx === chunks.length - 1;
-
-        const joined = joinExternalMessages(message.outputs);
-        const hasInterruptedToolCalls =
-          typeof joined.content === "object" &&
-          joined.content.some(isInterruptedToolCall);
-        const hasPendingToolCalls =
-          typeof joined.content === "object" &&
-          joined.content.some(isPendingToolCall);
-        const autoStatus = getAutoStatus(
-          isLast,
+      (cache, message, idx) =>
+        convertExternalMessageChunk(
+          message,
+          idx,
+          chunks.length,
           isRunning,
-          hasInterruptedToolCalls,
-          hasPendingToolCalls,
-          isLast ? state.metadata.error : undefined,
-        );
-
-        if (
-          cache &&
-          (cache.role !== "assistant" ||
-            !isAutoStatus(cache.status) ||
-            cache.status === autoStatus)
-        ) {
-          const inputs = getExternalStoreMessages<T>(cache);
-          if (shallowArrayEqual(inputs, message.inputs)) {
-            return cache;
-          }
-        }
-
-        const newMessage = fromThreadMessageLike(
-          joined,
-          idx.toString(),
-          autoStatus,
-        );
-        bindExternalStoreMessage(newMessage, message.inputs);
-        return newMessage;
-      },
+          state.metadata.error,
+          {
+            message: cache,
+            generatedFallbackMessages,
+          },
+          state.metadata.cancelledMessageIds,
+        ),
     );
 
     bindExternalStoreMessage(threadMessages, messages);
-
-    if (state.metadata.error) {
-      const lastMessage = threadMessages.at(-1);
-      if (!lastMessage || lastMessage.role !== "assistant") {
-        threadMessages.push(createErrorAssistantMessage(state.metadata.error));
-      }
-    }
-
-    return threadMessages;
+    return completeExternalMessageConversion(
+      threadMessages,
+      state.metadata.error,
+    );
   }, [state, messages, isRunning, joinStrategy]);
 };
-
-const shallowArrayEqual = (a: unknown[], b: unknown[]) => {
-  if (a.length !== b.length) return false;
-  for (let i = 0; i < a.length; i++) {
-    if (a[i] !== b[i]) return false;
-  }
-  return true;
-};