@assistant-ui/core 0.2.12 → 0.2.14

package/src/react/client/Interactables.ts

@@ -18,201 +18,183 @@ import { buildInteractableModelContext } from "./interactable-model-context";
 
 const PERSISTENCE_DEBOUNCE_MS = 500;
 
-export const Interactables = resource(
-  function Interactables(): ClientOutput<"interactables"> {
-    const [state, setState] = useState<InteractablesState>(() => ({
-      definitions: {},
-      persistence: {},
-    }));
+const useInteractables = (): ClientOutput<"interactables"> => {
+  const [state, setState] = useState<InteractablesState>(() => ({
+    definitions: {},
+    persistence: {},
+  }));
+
+  const clientRef = useAssistantClientRef();
+
+  const stateRef = useRef(state);
+  useEffect(() => {
+    stateRef.current = state;
+  }, [state]);
 
-    const clientRef = useAssistantClientRef();
+  const subscribersRef = useRef(new Set<() => void>());
+  const partialSchemaCacheRef = useRef(
+    new Map<string, InteractableStateSchema>(),
+  );
+  const detachedStateRef = useRef(new Map<string, unknown>());
 
-    const stateRef = useRef(state);
-    useEffect(() => {
-      stateRef.current = state;
-    }, [state]);
+  const adapterRef = useRef<InteractablePersistenceAdapter | undefined>(
+    undefined,
+  );
+  const debounceTimerRef = useRef<ReturnType<typeof setTimeout> | undefined>(
+    undefined,
+  );
+  const syncSeqRef = useRef(0);
+  const hasPendingLocalChangeRef = useRef(false);
+  const flushResolversRef = useRef<Array<() => void>>([]);
+  const dirtyIdsRef = useRef(new Set<string>());
 
-    const subscribersRef = useRef(new Set<() => void>());
-    const partialSchemaCacheRef = useRef(
-      new Map<string, InteractableStateSchema>(),
-    );
-    const detachedStateRef = useRef(new Map<string, unknown>());
+  const runPersistence = useCallback(async () => {
+    const adapter = adapterRef.current;
+    if (!adapter) {
+      for (const resolve of flushResolversRef.current) resolve();
+      flushResolversRef.current = [];
+      return;
+    }
 
-    const adapterRef = useRef<InteractablePersistenceAdapter | undefined>(
-      undefined,
-    );
-    const debounceTimerRef = useRef<ReturnType<typeof setTimeout> | undefined>(
-      undefined,
-    );
-    const syncSeqRef = useRef(0);
-    const hasPendingLocalChangeRef = useRef(false);
-    const flushResolversRef = useRef<Array<() => void>>([]);
-    const dirtyIdsRef = useRef(new Set<string>());
+    const seq = ++syncSeqRef.current;
+    const dirtyIds = new Set(dirtyIdsRef.current);
+    dirtyIdsRef.current.clear();
+    hasPendingLocalChangeRef.current = true;
 
-    const runPersistence = useCallback(async () => {
-      const adapter = adapterRef.current;
-      if (!adapter) {
+    // Snapshot before any await so unregistered definitions are still included.
+    const exported = stateRef.current.definitions;
+    const payload: InteractablePersistedState = {};
+    for (const [id, def] of Object.entries(exported)) {
+      payload[id] = { name: def.name, state: def.state };
+    }
+
+    setState((prev) => ({
+      ...prev,
+      persistence: {
+        ...prev.persistence,
+        ...Object.fromEntries(
+          [...dirtyIds].map((id) => [
+            id,
+            { isPending: true, error: undefined },
+          ]),
+        ),
+      },
+    }));
+
+    try {
+      await adapter.save(payload);
+      if (syncSeqRef.current === seq) {
+        hasPendingLocalChangeRef.current = false;
+        setState((prev) => {
+          const persistence = { ...prev.persistence };
+          for (const id of dirtyIds) delete persistence[id];
+          return { ...prev, persistence };
+        });
+      }
+    } catch (e) {
+      if (syncSeqRef.current === seq) {
+        hasPendingLocalChangeRef.current = false;
+        setState((prev) => ({
+          ...prev,
+          persistence: {
+            ...prev.persistence,
+            ...Object.fromEntries(
+              [...dirtyIds].map((id) => [id, { isPending: false, error: e }]),
+            ),
+          },
+        }));
+      }
+    } finally {
+      if (dirtyIdsRef.current.size > 0 && adapterRef.current) {
+        runPersistence();
+      } else {
         for (const resolve of flushResolversRef.current) resolve();
         flushResolversRef.current = [];
-        return;
       }
+    }
+  }, []);
 
-      const seq = ++syncSeqRef.current;
-      const dirtyIds = new Set(dirtyIdsRef.current);
-      dirtyIdsRef.current.clear();
-      hasPendingLocalChangeRef.current = true;
-
-      // Snapshot before any await so unregistered definitions are still included.
-      const exported = stateRef.current.definitions;
-      const payload: InteractablePersistedState = {};
-      for (const [id, def] of Object.entries(exported)) {
-        payload[id] = { name: def.name, state: def.state };
+  const schedulePersistence = useCallback(
+    (id: string) => {
+      if (!adapterRef.current) return;
+      dirtyIdsRef.current.add(id);
+      if (debounceTimerRef.current !== undefined) {
+        clearTimeout(debounceTimerRef.current);
       }
-
-      setState((prev) => ({
-        ...prev,
-        persistence: {
-          ...prev.persistence,
-          ...Object.fromEntries(
-            [...dirtyIds].map((id) => [
-              id,
-              { isPending: true, error: undefined },
-            ]),
-          ),
-        },
-      }));
-
-      try {
-        await adapter.save(payload);
-        if (syncSeqRef.current === seq) {
-          hasPendingLocalChangeRef.current = false;
-          setState((prev) => {
-            const persistence = { ...prev.persistence };
-            for (const id of dirtyIds) delete persistence[id];
-            return { ...prev, persistence };
-          });
-        }
-      } catch (e) {
-        if (syncSeqRef.current === seq) {
-          hasPendingLocalChangeRef.current = false;
-          setState((prev) => ({
-            ...prev,
-            persistence: {
-              ...prev.persistence,
-              ...Object.fromEntries(
-                [...dirtyIds].map((id) => [id, { isPending: false, error: e }]),
-              ),
-            },
-          }));
-        }
-      } finally {
-        if (dirtyIdsRef.current.size > 0 && adapterRef.current) {
+      debounceTimerRef.current = setTimeout(() => {
+        debounceTimerRef.current = undefined;
+        if (!hasPendingLocalChangeRef.current) {
           runPersistence();
         } else {
-          for (const resolve of flushResolversRef.current) resolve();
-          flushResolversRef.current = [];
-        }
-      }
-    }, []);
-
-    const schedulePersistence = useCallback(
-      (id: string) => {
-        if (!adapterRef.current) return;
-        dirtyIdsRef.current.add(id);
-        if (debounceTimerRef.current !== undefined) {
-          clearTimeout(debounceTimerRef.current);
-        }
-        debounceTimerRef.current = setTimeout(() => {
-          debounceTimerRef.current = undefined;
-          if (!hasPendingLocalChangeRef.current) {
+          debounceTimerRef.current = setTimeout(() => {
+            debounceTimerRef.current = undefined;
             runPersistence();
-          } else {
-            debounceTimerRef.current = setTimeout(() => {
-              debounceTimerRef.current = undefined;
-              runPersistence();
-            }, PERSISTENCE_DEBOUNCE_MS);
-          }
-        }, PERSISTENCE_DEBOUNCE_MS);
-      },
-      [runPersistence],
-    );
+          }, PERSISTENCE_DEBOUNCE_MS);
+        }
+      }, PERSISTENCE_DEBOUNCE_MS);
+    },
+    [runPersistence],
+  );
 
-    const exportState = useCallback((): InteractablePersistedState => {
-      const result: InteractablePersistedState = {};
-      for (const [id, def] of Object.entries(stateRef.current.definitions)) {
-        result[id] = { name: def.name, state: def.state };
-      }
-      return result;
-    }, []);
+  const exportState = useCallback((): InteractablePersistedState => {
+    const result: InteractablePersistedState = {};
+    for (const [id, def] of Object.entries(stateRef.current.definitions)) {
+      result[id] = { name: def.name, state: def.state };
+    }
+    return result;
+  }, []);
 
-    const importState = useCallback((saved: InteractablePersistedState) => {
+  const importState = useCallback((saved: InteractablePersistedState) => {
+    for (const [id, entry] of Object.entries(saved)) {
+      detachedStateRef.current.set(id, entry.state);
+    }
+    setState((prev) => {
+      let changed = false;
+      const definitions = { ...prev.definitions };
       for (const [id, entry] of Object.entries(saved)) {
-        detachedStateRef.current.set(id, entry.state);
-      }
-      setState((prev) => {
-        let changed = false;
-        const definitions = { ...prev.definitions };
-        for (const [id, entry] of Object.entries(saved)) {
-          if (definitions[id]) {
-            definitions[id] = { ...definitions[id], state: entry.state };
-            changed = true;
-          }
+        if (definitions[id]) {
+          definitions[id] = { ...definitions[id], state: entry.state };
+          changed = true;
         }
-        return changed ? { ...prev, definitions } : prev;
-      });
-    }, []);
-
-    const setPersistenceAdapter = useCallback(
-      (adapter: InteractablePersistenceAdapter | undefined) => {
-        adapterRef.current = adapter;
-      },
-      [],
-    );
-
-    const flush = useCallback(async () => {
-      if (debounceTimerRef.current !== undefined) {
-        clearTimeout(debounceTimerRef.current);
-        debounceTimerRef.current = undefined;
       }
-      if (!adapterRef.current) return;
-      if (!hasPendingLocalChangeRef.current && dirtyIdsRef.current.size === 0)
-        return;
-      const p = new Promise<void>((resolve) => {
-        flushResolversRef.current.push(resolve);
-      });
-      if (!hasPendingLocalChangeRef.current) {
-        runPersistence();
-      }
-      return p;
-    }, [runPersistence]);
+      return changed ? { ...prev, definitions } : prev;
+    });
+  }, []);
 
-    const flushIfPending = useCallback(() => {
-      if (adapterRef.current && debounceTimerRef.current !== undefined) {
-        clearTimeout(debounceTimerRef.current);
-        debounceTimerRef.current = undefined;
-        runPersistence();
-      }
-    }, [runPersistence]);
+  const setPersistenceAdapter = useCallback(
+    (adapter: InteractablePersistenceAdapter | undefined) => {
+      adapterRef.current = adapter;
+    },
+    [],
+  );
 
-    const setDefState = useCallback(
-      (id: string, updater: (prev: unknown) => unknown) => {
-        setState((prev) => {
-          const existing = prev.definitions[id];
-          if (!existing) return prev;
-          return {
-            ...prev,
-            definitions: {
-              ...prev.definitions,
-              [id]: { ...existing, state: updater(existing.state) },
-            },
-          };
-        });
-        if (stateRef.current.definitions[id]) schedulePersistence(id);
-      },
-      [schedulePersistence],
-    );
+  const flush = useCallback(async () => {
+    if (debounceTimerRef.current !== undefined) {
+      clearTimeout(debounceTimerRef.current);
+      debounceTimerRef.current = undefined;
+    }
+    if (!adapterRef.current) return;
+    if (!hasPendingLocalChangeRef.current && dirtyIdsRef.current.size === 0)
+      return;
+    const p = new Promise<void>((resolve) => {
+      flushResolversRef.current.push(resolve);
+    });
+    if (!hasPendingLocalChangeRef.current) {
+      runPersistence();
+    }
+    return p;
+  }, [runPersistence]);
+
+  const flushIfPending = useCallback(() => {
+    if (adapterRef.current && debounceTimerRef.current !== undefined) {
+      clearTimeout(debounceTimerRef.current);
+      debounceTimerRef.current = undefined;
+      runPersistence();
+    }
+  }, [runPersistence]);
 
-    const setDefSelected = useCallback((id: string, selected: boolean) => {
+  const setDefState = useCallback(
+    (id: string, updater: (prev: unknown) => unknown) => {
       setState((prev) => {
         const existing = prev.definitions[id];
         if (!existing) return prev;
@@ -220,107 +202,125 @@ export const Interactables = resource(
           ...prev,
           definitions: {
             ...prev.definitions,
-            [id]: { ...existing, selected },
+            [id]: { ...existing, state: updater(existing.state) },
           },
         };
       });
-    }, []);
+      if (stateRef.current.definitions[id]) schedulePersistence(id);
+    },
+    [schedulePersistence],
+  );
 
-    const provider = useMemo(
-      () => ({
-        getModelContext: () => {
-          const defs = stateRef.current.definitions;
-          return (
-            buildInteractableModelContext(
-              defs,
-              partialSchemaCacheRef.current,
-              setDefState,
-            ) ?? {}
-          );
-        },
-        subscribe: (callback: () => void) => {
-          subscribersRef.current.add(callback);
-          return () => {
-            subscribersRef.current.delete(callback);
-          };
+  const setDefSelected = useCallback((id: string, selected: boolean) => {
+    setState((prev) => {
+      const existing = prev.definitions[id];
+      if (!existing) return prev;
+      return {
+        ...prev,
+        definitions: {
+          ...prev.definitions,
+          [id]: { ...existing, selected },
         },
-      }),
-      [setDefState],
-    );
+      };
+    });
+  }, []);
+
+  const provider = useMemo(
+    () => ({
+      getModelContext: () => {
+        const defs = stateRef.current.definitions;
+        return (
+          buildInteractableModelContext(
+            defs,
+            partialSchemaCacheRef.current,
+            setDefState,
+          ) ?? {}
+        );
+      },
+      subscribe: (callback: () => void) => {
+        subscribersRef.current.add(callback);
+        return () => {
+          subscribersRef.current.delete(callback);
+        };
+      },
+    }),
+    [setDefState],
+  );
 
-    useEffect(() => {
-      for (const cb of subscribersRef.current) cb();
-    }, [state]);
+  useEffect(() => {
+    for (const cb of subscribersRef.current) cb();
+  }, [state]);
 
-    useEffect(() => {
-      return clientRef.current!.modelContext().register(provider);
-    }, [clientRef, provider]);
+  useEffect(() => {
+    return clientRef.current!.modelContext().register(provider);
+  }, [clientRef, provider]);
 
-    const register = useCallback(
-      (def: InteractableRegistration) => {
-        try {
-          const jsonSchema = toJSONSchema(def.stateSchema);
-          partialSchemaCacheRef.current.set(
-            def.id,
-            toPartialJSONSchema(jsonSchema),
-          );
-        } catch (e) {
-          console.warn(
-            `[Interactables] Failed to create partial schema for "${def.name}". The update tool will require all fields.`,
-            e,
-          );
-        }
+  const register = useCallback(
+    (def: InteractableRegistration) => {
+      try {
+        const jsonSchema = toJSONSchema(def.stateSchema);
+        partialSchemaCacheRef.current.set(
+          def.id,
+          toPartialJSONSchema(jsonSchema),
+        );
+      } catch (e) {
+        console.warn(
+          `[Interactables] Failed to create partial schema for "${def.name}". The update tool will require all fields.`,
+          e,
+        );
+      }
 
-        const detached = detachedStateRef.current.get(def.id);
-        detachedStateRef.current.delete(def.id);
+      const detached = detachedStateRef.current.get(def.id);
+      detachedStateRef.current.delete(def.id);
 
-        setState((prev) => ({
-          ...prev,
-          definitions: {
-            ...prev.definitions,
-            [def.id]: {
-              id: def.id,
-              name: def.name,
-              description: def.description,
-              stateSchema: def.stateSchema,
-              state:
-                prev.definitions[def.id]?.state ?? detached ?? def.initialState,
-              selected: def.selected,
-            },
+      setState((prev) => ({
+        ...prev,
+        definitions: {
+          ...prev.definitions,
+          [def.id]: {
+            id: def.id,
+            name: def.name,
+            description: def.description,
+            stateSchema: def.stateSchema,
+            state:
+              prev.definitions[def.id]?.state ?? detached ?? def.initialState,
+            selected: def.selected,
           },
-        }));
+        },
+      }));
 
-        return () => {
-          flushIfPending();
-          setState((prev) => {
-            const existing = prev.definitions[def.id];
-            if (existing) {
-              detachedStateRef.current.set(def.id, existing.state);
-            }
-            partialSchemaCacheRef.current.delete(def.id);
-            const { [def.id]: _, ...rest } = prev.definitions;
-            const { [def.id]: __, ...restPersistence } = prev.persistence;
-            return { ...prev, definitions: rest, persistence: restPersistence };
-          });
-        };
-      },
-      [flushIfPending],
-    );
+      return () => {
+        flushIfPending();
+        setState((prev) => {
+          const existing = prev.definitions[def.id];
+          if (existing) {
+            detachedStateRef.current.set(def.id, existing.state);
+          }
+          partialSchemaCacheRef.current.delete(def.id);
+          const { [def.id]: _, ...rest } = prev.definitions;
+          const { [def.id]: __, ...restPersistence } = prev.persistence;
+          return { ...prev, definitions: rest, persistence: restPersistence };
+        });
+      };
+    },
+    [flushIfPending],
+  );
+
+  return {
+    getState: () => state,
+    register,
+    setState: setDefState,
+    setSelected: setDefSelected,
+    exportState,
+    importState,
+    setPersistenceAdapter,
+    flush,
+  };
+};
 
-    return {
-      getState: () => state,
-      register,
-      setState: setDefState,
-      setSelected: setDefSelected,
-      exportState,
-      importState,
-      setPersistenceAdapter,
-      flush,
-    };
-  },
-);
+export const Interactables = resource(useInteractables);
 
-attachTransformScopes(Interactables, (scopes, parent) => {
+attachTransformScopes(useInteractables, (scopes, parent) => {
   if (!scopes.modelContext && parent.modelContext.source === null) {
     scopes.modelContext = ModelContext();
   }

package/src/react/client/Tools.ts

@@ -30,7 +30,7 @@ export type { McpAppResourceOutput };
  * context, while each tool renderer is registered with the tools scope for
  * message rendering.
  */
-export const Tools = resource(function Tools({
+const useTools = ({
   toolkit,
   mcpApp,
 }: {
@@ -38,7 +38,7 @@ export const Tools = resource(function Tools({
   toolkit?: Toolkit;
   /** Optional MCP app resource whose tools should be merged into context. */
   mcpApp?: ResourceElement<McpAppResourceOutput> | undefined;
-}): ClientOutput<"tools"> {
+}): ClientOutput<"tools"> => {
   const mcpAppOutputs = useResources(
     () => (mcpApp ? [withKey("mcpApp", mcpApp)] : []),
     [mcpApp],
@@ -155,9 +155,11 @@ export const Tools = resource(function Tools({
     getState: () => state,
     setToolUI,
   };
-});
+};
+
+export const Tools = resource(useTools);
 
-attachTransformScopes(Tools, (scopes, parent) => {
+attachTransformScopes(useTools, (scopes, parent) => {
   if (!scopes.modelContext && parent.modelContext.source === null) {
     scopes.modelContext = ModelContext();
   }

package/src/react/primitives/message/MessageGroupedParts.tsx

@@ -49,8 +49,9 @@ export namespace MessagePrimitiveGroupedParts {
    * which running states qualify:
    * - `"never"` — never.
    * - `"empty"` — only when the message has no parts yet.
-   * - `"no-text"` (default) — when the last part isn't `text`/`reasoning`
-   *   (e.g. it ended on a tool call, so the assistant likely isn't done).
+   * - `"no-text"` (default) — when the message has no parts yet or the last
+   *   part isn't `text`/`reasoning` (e.g. it ended on a tool call, so the
+   *   assistant likely isn't done).
    * - `"always"` — whenever the message is running, regardless of parts.
    */
   export type IndicatorMode = "never" | "empty" | "no-text" | "always";
@@ -153,7 +154,8 @@ const shouldShowIndicator = (
     case "no-text": {
       const last = parts[parts.length - 1];
       return (
-        last !== undefined && last.type !== "text" && last.type !== "reasoning"
+        last === undefined ||
+        (last.type !== "text" && last.type !== "reasoning")
       );
     }
   }
@@ -179,9 +181,14 @@ const renderNode = <TKey extends `group-${string}`>(
   render: (info: MessagePrimitiveGroupedParts.RenderInfo<TKey>) => ReactNode,
 ): ReactNode => {
   if (node.type === "part") {
-    // Key by absolute part index, not structural nodeKey — prevents zombie fiber subscriptions when parts reshape (#4051).
+    // Key by part identity when available, else absolute part index — never
+    // the structural nodeKey, which leaves zombie fiber subscriptions when
+    // parts reshape (#4051).
     return (
-      <MessagePartChildren key={`part-${node.index}`} index={node.index}>
+      <MessagePartChildren
+        key={node.idKey ? `part-${node.idKey}` : `part-${node.index}`}
+        index={node.index}
+      >
         {({ part }) => render({ part, children: <PartChildrenSentinel /> })}
       </MessagePartChildren>
     );
@@ -195,7 +202,7 @@ const renderNode = <TKey extends `group-${string}`>(
   };
 
   return (
-    <Fragment key={node.nodeKey}>
+    <Fragment key={node.idKey ?? node.nodeKey}>
       {render({
         part: groupPart,
         children: (
@@ -264,7 +271,12 @@ export const MessagePrimitiveGroupedParts = <TKey extends `group-${string}`>({
   const memoDep = memoKey ?? groupBy;
   const tree = useMemo(() => {
     const context: GroupByContext = { toolUIs };
-    return buildGroupTree(parts.map((part) => groupBy(part, context) ?? []));
+    return buildGroupTree(
+      parts.map((part) => groupBy(part, context) ?? []),
+      parts.map((part) =>
+        part.type === "tool-call" ? part.toolCallId : undefined,
+      ),
+    );
     // oxlint-disable-next-line react/exhaustive-deps -- groupBy is captured via memoDep (either its identity or the helper's memoKey fingerprint); listing it directly would defeat the helper-tagged memo path
   }, [parts, memoDep, toolUIs]);