@assistant-ui/core 0.3.17 → 0.3.19

package/src/runtime/utils/auto-status.test.ts

@@ -1,5 +1,33 @@
 import { describe, expect, it } from "vitest";
-import { getAutoStatus } from "./auto-status";
+import type { MessageStatus, ThreadMessage } from "../../types/message";
+import {
+  getAutoStatus,
+  getContentAutoStatus,
+  isBackgroundToolCall,
+} from "./auto-status";
+
+const nestedAssistantMessage = (status: MessageStatus): ThreadMessage => ({
+  id: "nested-assistant",
+  createdAt: new Date(0),
+  role: "assistant",
+  content: [],
+  status,
+  metadata: {
+    unstable_state: {},
+    unstable_annotations: [],
+    unstable_data: [],
+    steps: [],
+    custom: {},
+  },
+});
+
+const pendingToolCall = (messages?: readonly ThreadMessage[]) => ({
+  type: "tool-call" as const,
+  toolName: "delegate",
+  args: {},
+  argsText: "",
+  messages,
+});
 
 describe("getAutoStatus", () => {
   it("reports a cancelled message as incomplete", () => {
@@ -54,4 +82,87 @@ describe("getAutoStatus", () => {
       ).toMatchObject(expected);
     },
   );
+
+  it.each([true, false])(
+    "keeps a background tool call running when isLast is %s",
+    (isLast) => {
+      expect(
+        getContentAutoStatus(
+          [pendingToolCall([nestedAssistantMessage({ type: "running" })])],
+          isLast,
+          false,
+        ),
+      ).toMatchObject({ type: "running" });
+    },
+  );
+
+  it("reads the trailing nested message, not an earlier running one", () => {
+    expect(
+      isBackgroundToolCall(
+        pendingToolCall([
+          nestedAssistantMessage({ type: "running" }),
+          nestedAssistantMessage({ type: "complete", reason: "unknown" }),
+        ]),
+      ),
+    ).toBe(false);
+  });
+
+  it("lets cancellation clear a background tool call", () => {
+    expect(
+      getAutoStatus(true, false, false, true, undefined, true, true),
+    ).toMatchObject({ type: "requires-action", reason: "tool-calls" });
+  });
+
+  it("keeps a settled nested tool call pending", () => {
+    expect(
+      getContentAutoStatus(
+        [
+          pendingToolCall([
+            nestedAssistantMessage({ type: "complete", reason: "unknown" }),
+          ]),
+        ],
+        true,
+        false,
+      ),
+    ).toMatchObject({ type: "requires-action", reason: "tool-calls" });
+  });
+
+  it("keeps an interrupted tool call ahead of a background tool call", () => {
+    expect(
+      getContentAutoStatus(
+        [
+          {
+            ...pendingToolCall(),
+            interrupt: { type: "human", payload: {} },
+          },
+          pendingToolCall([nestedAssistantMessage({ type: "running" })]),
+        ],
+        false,
+        false,
+      ),
+    ).toMatchObject({ type: "requires-action", reason: "interrupt" });
+  });
+
+  it.each([
+    ["a result", { ...pendingToolCall(), result: {} }],
+    ["no nested messages", pendingToolCall()],
+    [
+      "only a user nested message",
+      pendingToolCall([
+        {
+          id: "nested-user",
+          createdAt: new Date(0),
+          role: "user",
+          content: [],
+          attachments: [],
+          metadata: { custom: {} },
+        },
+      ]),
+    ],
+  ])(
+    "does not treat a tool call with %s as a background call",
+    (_label, call) => {
+      expect(isBackgroundToolCall(call)).toBe(false);
+    },
+  );
 });

package/src/runtime/utils/auto-status.ts

@@ -10,6 +10,15 @@ type ThreadMessageLikeContentItem = Exclude<
 export const isPendingToolCall = (c: ThreadMessageLikeContentItem): boolean =>
   c.type === "tool-call" && c.result === undefined;
 
+/** A pending tool call whose nested conversation is still streaming; the outer message stays running until the nested run settles. */
+export const isBackgroundToolCall = (
+  c: ThreadMessageLikeContentItem,
+): boolean => {
+  if (c.type !== "tool-call" || c.result !== undefined) return false;
+  const last = c.messages?.at(-1);
+  return last?.role === "assistant" && last.status.type === "running";
+};
+
 export const isInterruptedToolCall = (
   c: ThreadMessageLikeContentItem,
 ): boolean => {
@@ -76,6 +85,7 @@ export const getAutoStatus = (
   hasPendingToolCalls: boolean,
   error?: ReadonlyJSONValue,
   isCancelled?: boolean,
+  hasBackgroundToolCalls?: boolean,
 ): MessageStatus => {
   if (isLast && error) {
     return Object.assign(
@@ -92,13 +102,26 @@ export const getAutoStatus = (
     ? AUTO_STATUS_RUNNING
     : hasInterruptedToolCalls
       ? AUTO_STATUS_INTERRUPT
-      : hasPendingToolCalls
-        ? AUTO_STATUS_PENDING
-        : isCancelled
-          ? AUTO_STATUS_CANCELLED
-          : AUTO_STATUS_COMPLETE;
+      : hasBackgroundToolCalls && !isCancelled
+        ? AUTO_STATUS_RUNNING
+        : hasPendingToolCalls
+          ? AUTO_STATUS_PENDING
+          : isCancelled
+            ? AUTO_STATUS_CANCELLED
+            : AUTO_STATUS_COMPLETE;
 };
 
+/** Status for a message entering the repository: the nested conversation is not being fed here, so it never keeps the message running. */
+export const getRepositoryContentAutoStatus = (
+  content: ThreadMessageLike["content"],
+): MessageStatus =>
+  getAutoStatus(
+    false,
+    false,
+    typeof content !== "string" && content.some(isInterruptedToolCall),
+    typeof content !== "string" && content.some(isPendingToolCall),
+  );
+
 export const getContentAutoStatus = (
   content: ThreadMessageLike["content"],
   isLast: boolean,
@@ -109,4 +132,7 @@ export const getContentAutoStatus = (
     isRunning,
     typeof content !== "string" && content.some(isInterruptedToolCall),
     typeof content !== "string" && content.some(isPendingToolCall),
+    undefined,
+    undefined,
+    typeof content !== "string" && content.some(isBackgroundToolCall),
   );

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

@@ -1,10 +1,13 @@
 import { describe, expect, it, vi } from "vitest";
+import type { ThreadMessage } from "../../types/message";
 import {
   chunkExternalMessages,
   convertExternalMessageCallback,
   convertExternalMessageChunk,
+  joinExternalMessages,
   type ExternalMessageConverterCallback,
   type ExternalMessageConverterCallbackResult,
+  type ExternalMessageConverterMessage,
 } from "./external-message-conversion";
 
 describe("convertExternalMessageCallback", () => {
@@ -30,6 +33,75 @@ describe("convertExternalMessageCallback", () => {
   });
 });
 
+describe("joinExternalMessages", () => {
+  it("preserves strict equality for malformed numeric tool-call IDs", () => {
+    const messages = [
+      {
+        role: "assistant",
+        content: [
+          {
+            type: "tool-call",
+            toolCallId: 1,
+            toolName: "search",
+            args: { query: "old" },
+          },
+        ],
+      },
+      {
+        role: "assistant",
+        content: [
+          {
+            type: "tool-call",
+            toolCallId: 1,
+            toolName: "search",
+            args: { query: "new" },
+          },
+        ],
+      },
+    ] as unknown as ExternalMessageConverterMessage[];
+
+    expect(joinExternalMessages(messages).content).toMatchObject([
+      {
+        type: "tool-call",
+        toolCallId: 1,
+        args: { query: "new" },
+      },
+    ]);
+  });
+
+  it("does not merge malformed NaN tool-call IDs", () => {
+    const messages = [
+      {
+        role: "assistant",
+        content: [
+          {
+            type: "tool-call",
+            toolCallId: Number.NaN,
+            toolName: "search",
+            args: { query: "first" },
+          },
+          {
+            type: "tool-call",
+            toolCallId: Number.NaN,
+            toolName: "search",
+            args: { query: "second" },
+          },
+        ],
+      },
+      { role: "tool", toolCallId: Number.NaN, result: "found" },
+    ] as unknown as ExternalMessageConverterMessage[];
+
+    const result = joinExternalMessages(messages);
+
+    expect(result.content).toMatchObject([
+      { type: "tool-call", args: { query: "first" } },
+      { type: "tool-call", args: { query: "second" } },
+    ]);
+    expect(result.content[0]).not.toHaveProperty("result");
+    expect(result.content[1]).not.toHaveProperty("result");
+  });
+});
+
 describe("chunkExternalMessages", () => {
   it("keeps a tool result with a non-joining assistant before starting the next chunk", () => {
     const toolCall = {};
@@ -83,6 +155,136 @@ describe("chunkExternalMessages", () => {
 });
 
 describe("convertExternalMessageChunk", () => {
+  it("keeps an ended outer run running for a background tool call", () => {
+    const nestedAssistant: ThreadMessage = {
+      id: "nested-assistant",
+      createdAt: new Date(0),
+      role: "assistant",
+      content: [],
+      status: { type: "running" },
+      metadata: {
+        unstable_state: {},
+        unstable_annotations: [],
+        unstable_data: [],
+        steps: [],
+        custom: {},
+      },
+    };
+    const result = convertExternalMessageChunk(
+      {
+        inputs: [{}],
+        outputs: [
+          {
+            role: "assistant",
+            content: [
+              {
+                type: "tool-call",
+                toolName: "delegate",
+                args: {},
+                messages: [nestedAssistant],
+              },
+            ],
+          },
+        ],
+      },
+      0,
+      1,
+      false,
+      undefined,
+    );
+
+    expect(result.status).toMatchObject({ type: "running" });
+  });
+
+  it("preserves assistant message modality", () => {
+    const result = convertExternalMessageChunk(
+      {
+        inputs: [{}],
+        outputs: [
+          {
+            role: "assistant",
+            content: "Spoken reply",
+            metadata: { modality: "voice" },
+          },
+        ],
+      },
+      0,
+      1,
+      false,
+      undefined,
+    );
+
+    expect(result.metadata.modality).toBe("voice");
+  });
+
+  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 = {

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

@@ -1,5 +1,6 @@
 import type { ReadonlyJSONValue } from "assistant-stream/utils";
 import type { ToolExecutionStatus } from "../../runtimes/tool-invocations/ToolInvocationTracker";
+import { ThreadMessageConverter } from "../../runtimes/external-store/thread-message-converter";
 import type {
   ThreadAssistantMessage,
   ThreadMessage,
@@ -11,6 +12,7 @@ import { isJSONValueEqual } from "../../utils/json/is-json-equal";
 import {
   getAutoStatus,
   isAutoStatus,
+  isBackgroundToolCall,
   isInterruptedToolCall,
   isPendingToolCall,
 } from "./auto-status";
@@ -132,6 +134,9 @@ const mergeInnerMessages = (existing: object, incoming: object) => ({
   ],
 });
 
+const isNaNToolCallId = (toolCallId: unknown) =>
+  typeof toolCallId === "number" && Number.isNaN(toolCallId);
+
 export const joinExternalMessages = (
   messages: readonly ExternalMessageConverterMessage[],
 ): ThreadMessageLike => {
@@ -142,11 +147,13 @@ export const joinExternalMessages = (
     role: "assistant",
     content: [],
   };
+  const toolCallIndices = new Map<unknown, number>();
+  const reasoningIndices = new Map<string, number>();
   for (const output of messages) {
     if (output.role === "tool") {
-      const toolCallIdx = assistantMessage.content.findIndex(
-        (c) => c.type === "tool-call" && c.toolCallId === output.toolCallId,
-      );
+      const toolCallIdx = !isNaNToolCallId(output.toolCallId)
+        ? (toolCallIndices.get(output.toolCallId) ?? -1)
+        : -1;
       // Ignore orphaned tool results so one bad tool message does not
       // prevent rendering the rest of the conversation.
       if (toolCallIdx !== -1) {
@@ -239,6 +246,10 @@ export const joinExternalMessages = (
               assistantMessage.metadata.timing = output.metadata.timing;
             }
 
+            if (output.metadata.modality) {
+              assistantMessage.metadata.modality = output.metadata.modality;
+            }
+
             if (output.metadata.submittedFeedback) {
               assistantMessage.metadata.submittedFeedback =
                 output.metadata.submittedFeedback;
@@ -252,12 +263,9 @@ export const joinExternalMessages = (
 
           // 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) {
+            if (part.type === "tool-call" && part.toolCallId) {
+              const existingIdx = toolCallIndices.get(part.toolCallId);
+              if (existingIdx !== undefined) {
                 const existing = assistantMessage.content[
                   existingIdx
                 ] as typeof part;
@@ -275,13 +283,8 @@ export const joinExternalMessages = (
               "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 existingIdx = reasoningIndices.get(part.parentId);
+              if (existingIdx !== undefined) {
                 const existing = assistantMessage.content[
                   existingIdx
                 ] as typeof part;
@@ -293,7 +296,21 @@ export const joinExternalMessages = (
                 continue;
               }
             }
+            const partIdx = assistantMessage.content.length;
             assistantMessage.content.push(part);
+            if (
+              part.type === "tool-call" &&
+              !isNaNToolCallId(part.toolCallId) &&
+              !toolCallIndices.has(part.toolCallId)
+            ) {
+              toolCallIndices.set(part.toolCallId, partIdx);
+            } else if (
+              part.type === "reasoning" &&
+              "parentId" in part &&
+              part.parentId
+            ) {
+              reasoningIndices.set(part.parentId, partIdx);
+            }
           }
           break;
         default: {
@@ -370,18 +387,46 @@ export const shallowArrayEqual = (
   return true;
 };
 
-type ExternalMessageConversionCache = {
+type ExternalMessageChunkConversionCache = {
   message: ThreadMessage | undefined;
   generatedFallbackMessages: WeakSet<object>;
 };
 
+type ExternalMessageConversionCallbackCacheEntry<T> =
+  ExternalMessageConverterCallbackResult<T> & {
+    metadata: ExternalMessageConverterMetadata;
+    callback: ExternalMessageConverterCallback<T>;
+  };
+
+export type InternalExternalMessageConversionCache<
+  T extends WeakKey = WeakKey,
+> = {
+  callbackCache: WeakMap<T, ExternalMessageConversionCallbackCacheEntry<T>>;
+  chunkCache: WeakMap<
+    ExternalMessageConverterMessage,
+    ExternalMessageConverterChunk<T>
+  >;
+  converterCache: ThreadMessageConverter;
+  // Generatedness is tracked by identity, not by id shape: a caller-supplied id that happens to match the generated pattern must never be rewritten.
+  generatedFallbackMessages: WeakSet<object>;
+};
+
+export const createExternalMessageConversionCache = <
+  T extends WeakKey = WeakKey,
+>(): InternalExternalMessageConversionCache<T> => ({
+  callbackCache: new WeakMap(),
+  chunkCache: new WeakMap(),
+  converterCache: new ThreadMessageConverter(),
+  generatedFallbackMessages: new WeakSet(),
+});
+
 export const convertExternalMessageChunk = <T>(
   message: ExternalMessageConverterChunk<T>,
   idx: number,
   chunkCount: number,
   isRunning: boolean,
   error: ReadonlyJSONValue | undefined,
-  cache?: ExternalMessageConversionCache,
+  cache?: ExternalMessageChunkConversionCache,
   cancelledMessageIds?: ReadonlySet<string>,
 ) => {
   const isLast = idx === chunkCount - 1;
@@ -400,6 +445,9 @@ export const convertExternalMessageChunk = <T>(
   const hasPendingToolCalls =
     typeof joined.content === "object" &&
     joined.content.some(isPendingToolCall);
+  const hasBackgroundToolCalls =
+    typeof joined.content === "object" &&
+    joined.content.some(isBackgroundToolCall);
   const autoStatus = getAutoStatus(
     isLast,
     isRunning,
@@ -407,6 +455,7 @@ export const convertExternalMessageChunk = <T>(
     hasPendingToolCalls,
     isLast ? error : undefined,
     isCancelled,
+    hasBackgroundToolCalls,
   );
   const fallbackId = `${FALLBACK_ID_PREFIX}${idx}`;
 
@@ -490,21 +539,63 @@ export const convertExternalMessages = <T extends WeakKey>(
   callback: ExternalMessageConverterCallback<T>,
   isRunning: boolean,
   metadata: ExternalMessageConverterMetadata,
+  joinStrategy?: JoinStrategy,
+  cache?: InternalExternalMessageConversionCache<T>,
 ) => {
-  const callbackResults = messages.map((message) =>
-    convertExternalMessageCallback(message, callback, metadata),
-  );
-  const chunks = chunkExternalMessages(callbackResults);
-  const result = chunks.map((message, idx) =>
-    convertExternalMessageChunk(
-      message,
-      idx,
-      chunks.length,
-      isRunning,
-      metadata.error,
-      undefined,
-      metadata.cancelledMessageIds,
-    ),
+  const callbackResults = messages.map((message) => {
+    let result = cache?.callbackCache.get(message);
+    if (
+      !result ||
+      result.metadata !== metadata ||
+      result.callback !== callback
+    ) {
+      result = {
+        ...convertExternalMessageCallback(message, callback, metadata),
+        metadata,
+        callback,
+      };
+      cache?.callbackCache.set(message, result);
+    }
+    return result;
+  });
+  const chunks = chunkExternalMessages(callbackResults, joinStrategy).map(
+    (message) => {
+      const key = message.outputs[0];
+      if (!key || !cache) return message;
+
+      const cached = cache.chunkCache.get(key);
+      if (cached && shallowArrayEqual(cached.outputs, message.outputs)) {
+        return cached;
+      }
+      cache.chunkCache.set(key, message);
+      return message;
+    },
   );
+  const result = cache
+    ? cache.converterCache.convertMessages(chunks, (cached, message, idx) =>
+        convertExternalMessageChunk(
+          message,
+          idx,
+          chunks.length,
+          isRunning,
+          metadata.error,
+          {
+            message: cached,
+            generatedFallbackMessages: cache.generatedFallbackMessages,
+          },
+          metadata.cancelledMessageIds,
+        ),
+      )
+    : chunks.map((message, idx) =>
+        convertExternalMessageChunk(
+          message,
+          idx,
+          chunks.length,
+          isRunning,
+          metadata.error,
+          undefined,
+          metadata.cancelledMessageIds,
+        ),
+      );
   return completeExternalMessageConversion(result, metadata.error);
 };