@assistant-ui/react 0.12.14 → 0.12.15

package/src/legacy-runtime/runtime-cores/assistant-transport/useToolInvocations.ts

@@ -14,6 +14,7 @@ import {
   AssistantMetaTransformStream,
   type ReadonlyJSONValue,
 } from "assistant-stream/utils";
+import { isJSONValueEqual } from "../../../utils/json/is-json-equal";
 
 const isArgsTextComplete = (argsText: string) => {
   try {
@@ -24,6 +25,21 @@ const isArgsTextComplete = (argsText: string) => {
   }
 };
 
+const parseArgsText = (argsText: string) => {
+  try {
+    return JSON.parse(argsText);
+  } catch {
+    return undefined;
+  }
+};
+
+const isEquivalentCompleteArgsText = (previous: string, next: string) => {
+  const previousValue = parseArgsText(previous);
+  const nextValue = parseArgsText(next);
+  if (previousValue === undefined || nextValue === undefined) return false;
+  return isJSONValueEqual(previousValue, nextValue);
+};
+
 type UseToolInvocationsParams = {
   state: AssistantTransportState;
   getTools: () => Record<string, Tool> | undefined;
@@ -69,6 +85,7 @@ export function useToolInvocations({
 
   const acRef = useRef<AbortController>(new AbortController());
   const executingCountRef = useRef(0);
+  const startedExecutionToolCallIdsRef = useRef<Set<string>>(new Set());
   const settledResolversRef = useRef<Array<() => void>>([]);
   const toolCallIdAliasesRef = useRef<Map<string, string>>(new Map());
   const ignoredResultToolCallIdsRef = useRef<Set<string>>(new Set());
@@ -151,6 +168,7 @@ export function useToolInvocations({
           if (ignoredResultToolCallIdsRef.current.has(toolCallId)) {
             return;
           }
+          startedExecutionToolCallIdsRef.current.add(toolCallId);
           const logicalToolCallId = getLogicalToolCallId(toolCallId);
           executingCountRef.current++;
           setToolStatuses((prev) => ({
@@ -159,7 +177,19 @@ export function useToolInvocations({
           }));
         },
         onExecutionEnd: (toolCallId: string) => {
+          const wasStarted =
+            startedExecutionToolCallIdsRef.current.delete(toolCallId);
           if (ignoredResultToolCallIdsRef.current.has(toolCallId)) {
+            if (wasStarted) {
+              executingCountRef.current--;
+              if (executingCountRef.current === 0) {
+                settledResolversRef.current.forEach((resolve) => resolve());
+                settledResolversRef.current = [];
+              }
+            }
+            return;
+          }
+          if (!wasStarted) {
             return;
           }
           const logicalToolCallId = getLogicalToolCallId(toolCallId);
@@ -244,6 +274,64 @@ export function useToolInvocations({
       return setToolState(toolCallId, { ...state, ...patch });
     };
 
+    const hasExecutableTool = (toolName: string) => {
+      const tool = getTools()?.[toolName];
+      return tool?.execute !== undefined || tool?.streamCall !== undefined;
+    };
+
+    const shouldCloseArgsStream = ({
+      toolName,
+      argsText,
+      hasResult,
+    }: {
+      toolName: string;
+      argsText: string;
+      hasResult: boolean;
+    }) => {
+      if (hasResult) return true;
+      if (!hasExecutableTool(toolName)) {
+        // Non-executable tools can emit parseable snapshots mid-stream.
+        // Wait until the run settles before closing the args stream.
+        return !state.isRunning && isArgsTextComplete(argsText);
+      }
+      return isArgsTextComplete(argsText);
+    };
+
+    const restartToolArgsStream = ({
+      toolCallId,
+      toolName,
+      state,
+    }: {
+      toolCallId: string;
+      toolName: string;
+      state: ToolState;
+    }) => {
+      ignoredResultToolCallIdsRef.current.add(state.streamToolCallId);
+      state.controller.argsText.close();
+
+      const streamToolCallId = `${toolCallId}:rewrite:${rewriteCounterRef.current++}`;
+      toolCallIdAliasesRef.current.set(streamToolCallId, toolCallId);
+      const toolCallController = controller.addToolCallPart({
+        toolName,
+        toolCallId: streamToolCallId,
+      });
+
+      if (process.env.NODE_ENV !== "production") {
+        console.warn("started replacement stream tool call", {
+          toolCallId,
+          streamToolCallId,
+        });
+      }
+
+      return setToolState(toolCallId, {
+        ...createToolState({
+          controller: toolCallController,
+          streamToolCallId,
+        }),
+        hasResult: state.hasResult,
+      });
+    };
+
     const processMessages = (
       messages: readonly (typeof state.messages)[number][],
     ) => {
@@ -276,89 +364,141 @@ export function useToolInvocations({
               }
 
               if (content.argsText !== lastState.argsText) {
+                let shouldWriteArgsText = true;
+
                 if (lastState.argsComplete) {
-                  if (process.env.NODE_ENV !== "production") {
-                    console.warn(
-                      "argsText updated after controller was closed:",
-                      {
-                        previous: lastState.argsText,
-                        next: content.argsText,
-                      },
-                    );
+                  if (
+                    isEquivalentCompleteArgsText(
+                      lastState.argsText,
+                      content.argsText,
+                    )
+                  ) {
+                    lastState = patchToolState(content.toolCallId, lastState, {
+                      argsText: content.argsText,
+                    });
+                    shouldWriteArgsText = false;
                   }
-                } else {
-                  if (!content.argsText.startsWith(lastState.argsText)) {
-                    // Check if this is key reordering (both are complete JSON)
-                    // This happens when transitioning from streaming to complete state
-                    // and the provider returns keys in a different order
-                    if (
-                      isArgsTextComplete(lastState.argsText) &&
-                      isArgsTextComplete(content.argsText)
-                    ) {
-                      lastState.controller.argsText.close();
-                      patchToolState(content.toolCallId, lastState, {
-                        argsText: content.argsText,
-                        argsComplete: true,
-                      });
-                      return; // Continue to next content part
-                    }
+
+                  if (shouldWriteArgsText) {
+                    const canRestartClosedArgsStream =
+                      !lastState.hasResult &&
+                      !startedExecutionToolCallIdsRef.current.has(
+                        lastState.streamToolCallId,
+                      );
+
                     if (process.env.NODE_ENV !== "production") {
                       console.warn(
-                        "argsText rewrote previous snapshot, restarting tool args stream:",
+                        canRestartClosedArgsStream
+                          ? "argsText updated after controller was closed, restarting tool args stream:"
+                          : "argsText updated after controller was closed:",
                         {
                           previous: lastState.argsText,
                           next: content.argsText,
-                          toolCallId: content.toolCallId,
                         },
                       );
                     }
 
-                    ignoredResultToolCallIdsRef.current.add(
-                      lastState.streamToolCallId,
-                    );
-                    lastState.controller.argsText.close();
+                    if (!canRestartClosedArgsStream) {
+                      lastState = patchToolState(
+                        content.toolCallId,
+                        lastState,
+                        {
+                          argsText: content.argsText,
+                        },
+                      );
+                      shouldWriteArgsText = false;
+                    }
+                  }
 
-                    const streamToolCallId = `${content.toolCallId}:rewrite:${rewriteCounterRef.current++}`;
-                    toolCallIdAliasesRef.current.set(
-                      streamToolCallId,
-                      content.toolCallId,
-                    );
-                    const toolCallController = controller.addToolCallPart({
+                  if (shouldWriteArgsText) {
+                    lastState = restartToolArgsStream({
+                      toolCallId: content.toolCallId,
                       toolName: content.toolName,
-                      toolCallId: streamToolCallId,
+                      state: lastState,
                     });
+                  }
+                } else if (!content.argsText.startsWith(lastState.argsText)) {
+                  // Check if this is key reordering (both are complete JSON)
+                  // This happens when transitioning from streaming to complete state
+                  // and the provider returns keys in a different order
+                  if (
+                    isArgsTextComplete(lastState.argsText) &&
+                    isArgsTextComplete(content.argsText) &&
+                    isEquivalentCompleteArgsText(
+                      lastState.argsText,
+                      content.argsText,
+                    )
+                  ) {
+                    const shouldClose = shouldCloseArgsStream({
+                      toolName: content.toolName,
+                      argsText: content.argsText,
+                      hasResult: content.result !== undefined,
+                    });
+                    if (shouldClose) {
+                      lastState.controller.argsText.close();
+                    }
+                    lastState = patchToolState(content.toolCallId, lastState, {
+                      argsText: content.argsText,
+                      argsComplete: shouldClose,
+                    });
+                    shouldWriteArgsText = false;
+                  }
+                  if (shouldWriteArgsText) {
                     if (process.env.NODE_ENV !== "production") {
-                      console.warn("started replacement stream tool call", {
-                        toolCallId: content.toolCallId,
-                        streamToolCallId,
-                      });
+                      console.warn(
+                        "argsText rewrote previous snapshot, restarting tool args stream:",
+                        {
+                          previous: lastState.argsText,
+                          next: content.argsText,
+                          toolCallId: content.toolCallId,
+                        },
+                      );
                     }
-                    lastState = setToolState(content.toolCallId, {
-                      ...createToolState({
-                        controller: toolCallController,
-                        streamToolCallId,
-                      }),
-                      hasResult: lastState.hasResult,
+                    lastState = restartToolArgsStream({
+                      toolCallId: content.toolCallId,
+                      toolName: content.toolName,
+                      state: lastState,
                     });
                   }
+                }
 
+                if (shouldWriteArgsText) {
                   const argsTextDelta = content.argsText.slice(
                     lastState.argsText.length,
                   );
                   lastState.controller.argsText.append(argsTextDelta);
 
-                  const shouldClose = isArgsTextComplete(content.argsText);
+                  const shouldClose = shouldCloseArgsStream({
+                    toolName: content.toolName,
+                    argsText: content.argsText,
+                    hasResult: content.result !== undefined,
+                  });
                   if (shouldClose) {
                     lastState.controller.argsText.close();
                   }
 
-                  patchToolState(content.toolCallId, lastState, {
+                  lastState = patchToolState(content.toolCallId, lastState, {
                     argsText: content.argsText,
                     argsComplete: shouldClose,
                   });
                 }
               }
 
+              if (!lastState.argsComplete) {
+                const shouldClose = shouldCloseArgsStream({
+                  toolName: content.toolName,
+                  argsText: content.argsText,
+                  hasResult: content.result !== undefined,
+                });
+                if (shouldClose) {
+                  lastState.controller.argsText.close();
+                  lastState = patchToolState(content.toolCallId, lastState, {
+                    argsText: content.argsText,
+                    argsComplete: true,
+                  });
+                }
+              }
+
               if (content.result !== undefined && !lastState.hasResult) {
                 patchToolState(content.toolCallId, lastState, {
                   hasResult: true,
@@ -390,7 +530,7 @@ export function useToolInvocations({
     if (isInitialState.current) {
       isInitialState.current = false;
     }
-  }, [state, controller]);
+  }, [state, controller, getTools]);
 
   const abort = (): Promise<void> => {
     humanInputRef.current.forEach(({ reject }) => {
@@ -414,6 +554,7 @@ export function useToolInvocations({
     reset: () => {
       isInitialState.current = true;
       void abort().finally(() => {
+        startedExecutionToolCallIdsRef.current.clear();
         toolCallIdAliasesRef.current.clear();
         ignoredResultToolCallIdsRef.current.clear();
         rewriteCounterRef.current = 0;

package/src/legacy-runtime/runtime-cores/external-store/external-message-converter.ts

@@ -62,6 +62,13 @@ 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 => {
@@ -77,6 +84,8 @@ const joinExternalMessages = (
       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
@@ -100,10 +109,6 @@ const joinExternalMessages = (
           isError: output.isError,
           messages: output.messages,
         };
-      } else {
-        throw new Error(
-          `Tool call ${output.toolCallId} ${output.toolName} not found in assistant message`,
-        );
       }
     } else {
       const role = output.role;
@@ -180,6 +185,24 @@ 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) {
+                const existing = assistantMessage.content[
+                  existingIdx
+                ] as typeof part;
+                assistantMessage.content[existingIdx] = {
+                  ...existing,
+                  ...part,
+                  ...mergeInnerMessages(existing, part),
+                };
+                continue;
+              }
+            }
+
             if (
               part.type === "reasoning" &&
               "parentId" in part &&
@@ -198,12 +221,7 @@ const joinExternalMessages = (
                 assistantMessage.content[existingIdx] = {
                   ...existing,
                   text: `${existing.text}\n\n${part.text}`,
-                  ...{
-                    [symbolInnerMessage]: [
-                      ...((existing as any)[symbolInnerMessage] ?? []),
-                      ...((part as any)[symbolInnerMessage] ?? []),
-                    ],
-                  },
+                  ...mergeInnerMessages(existing, part),
                 };
                 continue;
               }

package/src/primitives/actionBar/ActionBarInteractionContext.ts

@@ -0,0 +1,13 @@
+"use client";
+
+import { createContext, useContext } from "react";
+
+export type ActionBarInteractionContextValue = {
+  acquireInteractionLock: () => () => void;
+};
+
+export const ActionBarInteractionContext =
+  createContext<ActionBarInteractionContextValue | null>(null);
+
+export const useActionBarInteractionContext = () =>
+  useContext(ActionBarInteractionContext);

package/src/primitives/actionBar/ActionBarRoot.tsx

@@ -1,11 +1,19 @@
 "use client";
 
 import { Primitive } from "@radix-ui/react-primitive";
-import { type ComponentRef, forwardRef, ComponentPropsWithoutRef } from "react";
+import {
+  type ComponentRef,
+  forwardRef,
+  ComponentPropsWithoutRef,
+  useCallback,
+  useMemo,
+  useState,
+} from "react";
 import {
   useActionBarFloatStatus,
   HideAndFloatStatus,
 } from "./useActionBarFloatStatus";
+import { ActionBarInteractionContext } from "./ActionBarInteractionContext";
 
 type PrimitiveDivProps = ComponentPropsWithoutRef<typeof Primitive.div>;
 
@@ -60,22 +68,44 @@ export const ActionBarPrimitiveRoot = forwardRef<
   ActionBarPrimitiveRoot.Element,
   ActionBarPrimitiveRoot.Props
 >(({ hideWhenRunning, autohide, autohideFloat, ...rest }, ref) => {
+  const [interactionCount, setInteractionCount] = useState(0);
+
+  const acquireInteractionLock = useCallback(() => {
+    let released = false;
+
+    setInteractionCount((count) => count + 1);
+
+    return () => {
+      if (released) return;
+      released = true;
+      setInteractionCount((count) => Math.max(0, count - 1));
+    };
+  }, []);
+
+  const interactionContext = useMemo(
+    () => ({ acquireInteractionLock }),
+    [acquireInteractionLock],
+  );
+
   const hideAndfloatStatus = useActionBarFloatStatus({
     hideWhenRunning,
     autohide,
     autohideFloat,
+    forceVisible: interactionCount > 0,
   });
 
   if (hideAndfloatStatus === HideAndFloatStatus.Hidden) return null;
 
   return (
-    <Primitive.div
-      {...(hideAndfloatStatus === HideAndFloatStatus.Floating
-        ? { "data-floating": "true" }
-        : null)}
-      {...rest}
-      ref={ref}
-    />
+    <ActionBarInteractionContext.Provider value={interactionContext}>
+      <Primitive.div
+        {...(hideAndfloatStatus === HideAndFloatStatus.Floating
+          ? { "data-floating": "true" }
+          : null)}
+        {...rest}
+        ref={ref}
+      />
+    </ActionBarInteractionContext.Provider>
   );
 });
 

package/src/primitives/actionBar/useActionBarFloatStatus.ts

@@ -12,24 +12,27 @@ export type UseActionBarFloatStatusProps = {
   hideWhenRunning?: boolean | undefined;
   autohide?: "always" | "not-last" | "never" | undefined;
   autohideFloat?: "always" | "single-branch" | "never" | undefined;
+  forceVisible?: boolean | undefined;
 };
 
 export const useActionBarFloatStatus = ({
   hideWhenRunning,
   autohide,
   autohideFloat,
+  forceVisible,
 }: UseActionBarFloatStatusProps) => {
   return useAuiState((s) => {
     if (hideWhenRunning && s.thread.isRunning) return HideAndFloatStatus.Hidden;
 
     const autohideEnabled =
       autohide === "always" || (autohide === "not-last" && !s.message.isLast);
+    const isVisibleByInteraction = forceVisible || s.message.isHovering;
 
     // normal status
     if (!autohideEnabled) return HideAndFloatStatus.Normal;
 
     // hidden status
-    if (!s.message.isHovering) return HideAndFloatStatus.Hidden;
+    if (!isVisibleByInteraction) return HideAndFloatStatus.Hidden;
 
     // floating status
     if (

package/src/primitives/actionBarMore/ActionBarMoreRoot.tsx

@@ -1,8 +1,9 @@
 "use client";
 
-import { FC } from "react";
+import { FC, useCallback, useEffect, useRef } from "react";
 import { DropdownMenu as DropdownMenuPrimitive } from "radix-ui";
 import { ScopedProps, useDropdownMenuScope } from "./scope";
+import { useActionBarInteractionContext } from "../actionBar/ActionBarInteractionContext";
 
 export namespace ActionBarMorePrimitiveRoot {
   export type Props = DropdownMenuPrimitive.DropdownMenuProps;
@@ -12,11 +13,60 @@ export const ActionBarMorePrimitiveRoot: FC<
   ActionBarMorePrimitiveRoot.Props
 > = ({
   __scopeActionBarMore,
+  open,
+  onOpenChange,
   ...rest
 }: ScopedProps<ActionBarMorePrimitiveRoot.Props>) => {
   const scope = useDropdownMenuScope(__scopeActionBarMore);
+  const actionBarInteraction = useActionBarInteractionContext();
+  const releaseInteractionLockRef = useRef<(() => void) | null>(null);
+  const isControlled = open !== undefined;
 
-  return <DropdownMenuPrimitive.Root {...scope} {...rest} />;
+  const setInteractionOpen = useCallback(
+    (nextOpen: boolean) => {
+      if (nextOpen) {
+        if (releaseInteractionLockRef.current) return;
+        releaseInteractionLockRef.current =
+          actionBarInteraction?.acquireInteractionLock() ?? null;
+        return;
+      }
+
+      releaseInteractionLockRef.current?.();
+      releaseInteractionLockRef.current = null;
+    },
+    [actionBarInteraction],
+  );
+
+  const handleOpenChange = useCallback(
+    (nextOpen: boolean) => {
+      if (!isControlled) {
+        setInteractionOpen(nextOpen);
+      }
+      onOpenChange?.(nextOpen);
+    },
+    [isControlled, setInteractionOpen, onOpenChange],
+  );
+
+  useEffect(() => {
+    if (!isControlled) return;
+    setInteractionOpen(Boolean(open));
+  }, [isControlled, open, setInteractionOpen]);
+
+  useEffect(() => {
+    return () => {
+      releaseInteractionLockRef.current?.();
+      releaseInteractionLockRef.current = null;
+    };
+  }, []);
+
+  return (
+    <DropdownMenuPrimitive.Root
+      {...scope}
+      {...rest}
+      {...(open !== undefined ? { open } : null)}
+      onOpenChange={handleOpenChange}
+    />
+  );
 };
 
 ActionBarMorePrimitiveRoot.displayName = "ActionBarMorePrimitive.Root";

package/src/tests/external-message-converter.test.ts

@@ -158,6 +158,86 @@ describe("convertExternalMessages", () => {
       expect(toolCallParts).toHaveLength(1);
       expect((toolCallParts[0] as any).result).toEqual({ data: "result" });
     });
+
+    it("should merge duplicate tool calls by toolCallId across assistant messages", () => {
+      const messages = [
+        {
+          id: "msg1",
+          role: "assistant" as const,
+          content: [
+            {
+              type: "tool-call" as const,
+              toolCallId: "tc1",
+              toolName: "search",
+              args: { query: "old" },
+              argsText: '{"query":"old"',
+            },
+          ],
+        },
+        {
+          id: "msg2",
+          role: "assistant" as const,
+          content: [
+            {
+              type: "tool-call" as const,
+              toolCallId: "tc1",
+              toolName: "search",
+              args: { query: "new" },
+              argsText: '{"query":"new"}',
+            },
+          ],
+        },
+      ];
+
+      const callback: useExternalMessageConverter.Callback<
+        (typeof messages)[number]
+      > = (msg) => msg;
+
+      const result = convertExternalMessages(messages, callback, false, {});
+
+      expect(result).toHaveLength(1);
+      expect(result[0]!.role).toBe("assistant");
+      const toolCallParts = result[0]!.content.filter(
+        (p) => p.type === "tool-call",
+      );
+      expect(toolCallParts).toHaveLength(1);
+      expect((toolCallParts[0] as any).args).toEqual({ query: "new" });
+      expect((toolCallParts[0] as any).argsText).toBe('{"query":"new"}');
+    });
+
+    it("should ignore orphaned tool results without throwing", () => {
+      const messages = [
+        {
+          id: "msg1",
+          role: "assistant" as const,
+          content: "First response",
+        },
+        {
+          role: "tool" as const,
+          toolCallId: "missing-tool-call",
+          toolName: "search",
+          result: { data: "orphan result" },
+        },
+        {
+          id: "msg2",
+          role: "assistant" as const,
+          content: "Second response",
+        },
+      ];
+
+      const callback: useExternalMessageConverter.Callback<
+        (typeof messages)[number]
+      > = (msg) => msg;
+
+      const result = convertExternalMessages(messages, callback, false, {});
+      expect(result).toHaveLength(1);
+      expect(result[0]!.role).toBe("assistant");
+
+      const textParts = result[0]!.content.filter((p) => p.type === "text");
+      expect(textParts).toHaveLength(2);
+      expect((textParts[0] as any).text).toBe("First response");
+      expect((textParts[1] as any).text).toBe("Second response");
+    });
   });
 
   describe("synthetic error message", () => {