@assistant-ui/core 0.3.16 → 0.3.18

package/src/react/client/Tools.ts

@@ -16,10 +16,12 @@ import type { Tool } from "assistant-stream";
 import {
   isStandaloneToolDisplay,
   makeToolCallTextComponent,
+  type ToolCallText,
   type Toolkit,
 } from "../model-context/toolbox";
 import type { ToolCallMessagePartComponent } from "../types/MessagePartComponentTypes";
 import { ModelContext } from "../../store/clients/model-context-client";
+import { nullProtoRecord } from "../../utils/record";
 
 export type { McpAppResourceOutput };
 
@@ -43,7 +45,9 @@ const useTools = ({
   const mcpAppOutputs = useResources(mcpApp ? [withKey("mcpApp", mcpApp)] : []);
   const mcpAppOutput = mcpAppOutputs[0];
 
-  const [toolUIs, setToolUIs] = useState<ToolsState["toolUIs"]>(() => ({}));
+  const [toolUIs, setToolUIs] = useState<ToolsState["toolUIs"]>(() =>
+    nullProtoRecord(),
+  );
 
   const state = useMemo(
     (): ToolsState => ({
@@ -59,29 +63,38 @@ const useTools = ({
     (
       toolName: string,
       render: ToolCallMessagePartComponent,
-      options?: { standalone?: boolean },
+      options?: {
+        standalone?: boolean;
+        renderText?: ToolCallText<any, any> | undefined;
+      },
     ) => {
       // One registration object per call; identity is the removal key, so
       // the per-name list stays correctly ref-counted across re-registers.
       const registration = {
         render,
+        renderText: options?.renderText,
         standalone: options?.standalone ?? false,
       };
 
-      setToolUIs((prev) => ({
-        ...prev,
-        [toolName]: [...(prev[toolName] ?? []), registration],
-      }));
+      setToolUIs((prev) => {
+        const next = nullProtoRecord(prev);
+        next[toolName] = [...(next[toolName] ?? []), registration];
+        return next;
+      });
 
       return () => {
         setToolUIs((prev) => {
-          const next = prev[toolName]?.filter((r) => r !== registration) ?? [];
-          if (next.length > 0) return { ...prev, [toolName]: next };
+          const registrations =
+            prev[toolName]?.filter((r) => r !== registration) ?? [];
+          const next = nullProtoRecord(prev);
+          if (registrations.length > 0) {
+            next[toolName] = registrations;
+            return next;
+          }
           // Drop the key entirely so repeatedly mounted/unmounted tools
           // don't leave empty arrays accumulating across a long session.
-          const rest = { ...prev };
-          delete rest[toolName];
-          return rest;
+          delete next[toolName];
+          return next;
         });
       };
     },
@@ -103,8 +116,12 @@ const useTools = ({
           : undefined);
       if (render) {
         unsubscribes.push(
+          // Registration has to be undone on unmount, so the registry write and
+          // its unsubscribe belong to the same effect.
+          // eslint-disable-next-line react-hooks/set-state-in-effect
           setToolUI(toolName, render, {
             standalone: isStandaloneToolDisplay(tool),
+            renderText: toolRenderText,
           }),
         );
       }
@@ -135,7 +152,7 @@ const useTools = ({
           acc[name] = rest as Tool<any, any>;
           return acc;
         },
-        {} as Record<string, Tool<any, any>>,
+        nullProtoRecord<Tool<any, any>>(),
       );
 
       const modelContextProvider = {

package/src/react/index.ts

@@ -235,6 +235,10 @@ export { RemoteThreadListHookInstanceManager } from "./runtimes/RemoteThreadList
 export { RemoteThreadListThreadListRuntimeCore } from "./runtimes/RemoteThreadListThreadListRuntimeCore";
 export { useRemoteThreadListRuntime } from "./runtimes/useRemoteThreadListRuntime";
 export { useCloudThreadListAdapter } from "./runtimes/cloud/useCloudThreadListAdapter";
+export {
+  createCloudThreadListAdapter,
+  type CloudThreadListAdapterOptions,
+} from "./runtimes/cloud/createCloudThreadListAdapter";
 export { useCloudThreadListRuntime } from "./runtimes/cloud/useCloudThreadListRuntime";
 export {
   useAssistantTransportRuntime,

package/src/react/interactables-legacy/Interactables.ts

@@ -18,11 +18,12 @@ import { ModelContext } from "../../store";
 import { buildInteractableModelContext } from "./interactable-model-context";
 import { notifySubscribers as notifyStateSubscribers } from "../../subscribable/subscribable";
 import { useInteractablePersistenceQueue } from "../interactables-shared/useInteractablePersistenceQueue";
+import { nullProtoRecord } from "../../utils/record";
 
 const useInteractables = (): ClientOutput<"interactables"> => {
   const [state, setState] = useState<InteractablesState>(() => ({
-    definitions: {},
-    persistence: {},
+    definitions: nullProtoRecord(),
+    persistence: nullProtoRecord(),
   }));
 
   const clientRef = useAssistantClientRef();
@@ -49,7 +50,7 @@ const useInteractables = (): ClientOutput<"interactables"> => {
   );
 
   const exportState = useCallback((): InteractablePersistedState => {
-    const result: InteractablePersistedState = {};
+    const result = nullProtoRecord<InteractablePersistedState[string]>();
     for (const [id, def] of Object.entries(stateRef.current.definitions)) {
       result[id] = { name: def.name, state: def.state };
     }
@@ -86,7 +87,7 @@ const useInteractables = (): ClientOutput<"interactables"> => {
       }
       setStateAndRef((prev) => {
         let changed = false;
-        const definitions = { ...prev.definitions };
+        const definitions = nullProtoRecord(prev.definitions);
         for (const [id, entry] of Object.entries(saved)) {
           if (definitions[id]) {
             definitions[id] = { ...definitions[id], state: entry.state };
@@ -114,10 +115,9 @@ const useInteractables = (): ClientOutput<"interactables"> => {
         if (!existing) return prev;
         return {
           ...prev,
-          definitions: {
-            ...prev.definitions,
+          definitions: nullProtoRecord(prev.definitions, {
             [id]: { ...existing, state: updater(existing.state) },
-          },
+          }),
         };
       });
       if (stateRef.current.definitions[id]) schedulePersistence(id);
@@ -132,10 +132,9 @@ const useInteractables = (): ClientOutput<"interactables"> => {
         if (!existing) return prev;
         return {
           ...prev,
-          definitions: {
-            ...prev.definitions,
+          definitions: nullProtoRecord(prev.definitions, {
             [id]: { ...existing, selected },
-          },
+          }),
         };
       });
     },
@@ -194,8 +193,7 @@ const useInteractables = (): ClientOutput<"interactables"> => {
 
       setStateAndRef((prev) => ({
         ...prev,
-        definitions: {
-          ...prev.definitions,
+        definitions: nullProtoRecord(prev.definitions, {
           [def.id]: {
             id: def.id,
             name: def.name,
@@ -205,7 +203,7 @@ const useInteractables = (): ClientOutput<"interactables"> => {
               prev.definitions[def.id]?.state ?? detached ?? def.initialState,
             selected: def.selected,
           },
-        },
+        }),
       }));
 
       return () => {
@@ -216,9 +214,11 @@ const useInteractables = (): ClientOutput<"interactables"> => {
             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 };
+          const definitions = nullProtoRecord(prev.definitions);
+          const persistence = nullProtoRecord(prev.persistence);
+          delete definitions[def.id];
+          delete persistence[def.id];
+          return { ...prev, definitions, persistence };
         });
       };
     },

package/src/react/interactables-shared/useInteractablePersistenceQueue.ts

@@ -1,4 +1,5 @@
 import { useCallback, useRef, type RefObject } from "react";
+import { nullProtoRecord } from "../../utils/record";
 
 const PERSISTENCE_DEBOUNCE_MS = 500;
 
@@ -90,15 +91,13 @@ export const useInteractablePersistenceQueue = <State>({
       const { adapter, payload, dirtyIds, seq } = resolved;
       inFlightPersistenceRef.current += 1;
 
-      updatePersistenceStatus((prev) => ({
-        ...prev,
-        ...Object.fromEntries(
-          [...dirtyIds].map((id) => [
-            id,
-            { isPending: true, error: undefined },
-          ]),
-        ),
-      }));
+      updatePersistenceStatus((prev) => {
+        const persistence = nullProtoRecord(prev);
+        for (const id of dirtyIds) {
+          persistence[id] = { isPending: true, error: undefined };
+        }
+        return persistence;
+      });
 
       const settleBatch = (status: PersistenceStatus | undefined) => {
         const settledIds: string[] = [];
@@ -114,7 +113,7 @@ export const useInteractablePersistenceQueue = <State>({
         if (settledIds.length === 0) return;
         updatePersistenceStatus((prev) => {
           let changed = false;
-          const persistence = { ...prev };
+          const persistence = nullProtoRecord(prev);
           for (const id of settledIds) {
             if (prev[id] === undefined) continue;
             if (status === undefined) delete persistence[id];

package/src/react/model-context/makeAssistantTool.ts

@@ -8,7 +8,7 @@ import { type AssistantToolProps, useAssistantTool } from "./useAssistantTool";
  * subtree.
  *
  * @deprecated Use a toolkit with `Tools({ toolkit })` and register it via
- * `useAui({ tools: Tools({ toolkit }) })` instead. See
+ * `AuiConfig({ tools: Tools({ toolkit }) })` on the provider's `config` prop instead. See
  * https://assistant-ui.com/docs/migrations/toolkit-tools.
  */
 export type AssistantTool = FC & {
@@ -25,7 +25,7 @@ export type AssistantTool = FC & {
  * @param tool - Tool definition and name to register.
  *
  * @deprecated Use a toolkit with `Tools({ toolkit })` and register it via
- * `useAui({ tools: Tools({ toolkit }) })` instead. See
+ * `AuiConfig({ tools: Tools({ toolkit }) })` on the provider's `config` prop instead. See
  * https://assistant-ui.com/docs/migrations/toolkit-tools.
  */
 export const makeAssistantTool = <

package/src/react/model-context/modelContextHooks.test.tsx

@@ -0,0 +1,131 @@
+// @vitest-environment jsdom
+
+import { act, cleanup, render, waitFor } from "@testing-library/react";
+import { startTransition, Suspense } from "react";
+import { afterEach, describe, expect, it, vi } from "vitest";
+
+const { aui, getProvider, resetProvider } = vi.hoisted(() => {
+  type TestProvider = {
+    getModelContext: () => {
+      system?: string | undefined;
+      tools?:
+        | Record<string, { execute?: (() => unknown) | undefined }>
+        | undefined;
+    };
+  };
+  let provider: TestProvider | undefined;
+
+  return {
+    aui: {
+      modelContext: {
+        register: vi.fn((nextProvider: TestProvider) => {
+          provider = nextProvider;
+          return () => {
+            if (provider === nextProvider) provider = undefined;
+          };
+        }),
+      },
+    },
+    getProvider: () => provider,
+    resetProvider: () => {
+      provider = undefined;
+    },
+  };
+});
+
+vi.mock("@assistant-ui/store", async (importOriginal) => ({
+  ...(await importOriginal()),
+  useAui: () => aui,
+}));
+
+import { useAssistantContext } from "./useAssistantContext";
+import { useAuiToolOverrides } from "./useAuiToolOverrides";
+
+afterEach(() => {
+  cleanup();
+  resetProvider();
+});
+
+describe("model context hooks", () => {
+  it("keeps assistant context scoped to the committed render", async () => {
+    const pending = new Promise<never>(() => {});
+    const renderB = vi.fn();
+    const Probe = ({ value, suspend }: { value: string; suspend: boolean }) => {
+      useAssistantContext({ getContext: () => value });
+      if (value === "workspace-b") renderB();
+      if (suspend) throw pending;
+      return null;
+    };
+
+    const view = render(
+      <Suspense fallback={null}>
+        <Probe value="workspace-a" suspend={false} />
+      </Suspense>,
+    );
+    await waitFor(() => expect(getProvider()).toBeDefined());
+
+    act(() => {
+      startTransition(() => {
+        view.rerender(
+          <Suspense fallback={null}>
+            <Probe value="workspace-b" suspend />
+          </Suspense>,
+        );
+      });
+    });
+
+    expect(renderB).toHaveBeenCalled();
+    expect(getProvider()?.getModelContext().system).toBe("workspace-a");
+
+    view.rerender(
+      <Suspense fallback={null}>
+        <Probe value="workspace-b" suspend={false} />
+      </Suspense>,
+    );
+
+    expect(getProvider()?.getModelContext().system).toBe("workspace-b");
+  });
+
+  it("keeps tool overrides scoped to the committed render", async () => {
+    const pending = new Promise<never>(() => {});
+    const renderB = vi.fn();
+    const Probe = ({ value, suspend }: { value: string; suspend: boolean }) => {
+      useAuiToolOverrides({ search: { execute: () => value } });
+      if (value === "workspace-b") renderB();
+      if (suspend) throw pending;
+      return null;
+    };
+
+    const view = render(
+      <Suspense fallback={null}>
+        <Probe value="workspace-a" suspend={false} />
+      </Suspense>,
+    );
+    await waitFor(() => expect(getProvider()).toBeDefined());
+
+    act(() => {
+      startTransition(() => {
+        view.rerender(
+          <Suspense fallback={null}>
+            <Probe value="workspace-b" suspend />
+          </Suspense>,
+        );
+      });
+    });
+
+    expect(renderB).toHaveBeenCalled();
+    expect(getProvider()?.getModelContext().tools?.search?.execute?.()).toBe(
+      "workspace-a",
+    );
+
+    view.rerender(
+      <Suspense fallback={null}>
+        <Probe value="workspace-b" suspend={false} />
+      </Suspense>,
+    );
+
+    expect(getProvider()?.getModelContext().tools?.search?.execute?.()).toBe(
+      "workspace-b",
+    );
+  });
+});

package/src/react/model-context/toolbox.test.ts

@@ -19,7 +19,7 @@ describe("makeToolCallTextComponent", () => {
         status: { type: "running" },
         addResult: () => {},
         resume: () => {},
-        respondToApproval: () => {},
+        respondToApproval: async () => {},
       }),
     ).toBe("Searching...");
 
@@ -34,7 +34,7 @@ describe("makeToolCallTextComponent", () => {
         status: { type: "complete" },
         addResult: () => {},
         resume: () => {},
-        respondToApproval: () => {},
+        respondToApproval: async () => {},
       }),
     ).toBe("Done searching");
   });
@@ -59,7 +59,7 @@ describe("makeToolCallTextComponent", () => {
         status: { type: "running" },
         addResult: () => {},
         resume: () => {},
-        respondToApproval: () => {},
+        respondToApproval: async () => {},
       }),
     ).toBe("Searching docs...");
 
@@ -74,7 +74,7 @@ describe("makeToolCallTextComponent", () => {
         status: { type: "complete" },
         addResult: () => {},
         resume: () => {},
-        respondToApproval: () => {},
+        respondToApproval: async () => {},
       }),
     ).toBe("Found 3 results for docs");
   });
@@ -94,7 +94,7 @@ describe("makeToolCallTextComponent", () => {
         argsText: "{}",
         addResult: () => {},
         resume: () => {},
-        respondToApproval: () => {},
+        respondToApproval: async () => {},
       } as unknown as ToolCallMessagePartProps),
     ).toBe("Complete");
   });
@@ -115,7 +115,7 @@ describe("makeToolCallTextComponent", () => {
         status: { type: "incomplete", reason: "error" },
         addResult: () => {},
         resume: () => {},
-        respondToApproval: () => {},
+        respondToApproval: async () => {},
       }),
     ).toBe("Finished");
   });
@@ -136,7 +136,7 @@ describe("makeToolCallTextComponent", () => {
         status: { type: "requires-action", reason: "interrupt" },
         addResult: () => {},
         resume: () => {},
-        respondToApproval: () => {},
+        respondToApproval: async () => {},
       }),
     ).toBe("Waiting for approval...");
   });
@@ -156,7 +156,7 @@ describe("makeToolCallTextComponent", () => {
         status: { type: "complete" },
         addResult: () => {},
         resume: () => {},
-        respondToApproval: () => {},
+        respondToApproval: async () => {},
       }),
     ).toBeNull();
   });
@@ -176,7 +176,7 @@ describe("makeToolCallTextComponent", () => {
         status: { type: "running" },
         addResult: () => {},
         resume: () => {},
-        respondToApproval: () => {},
+        respondToApproval: async () => {},
       }),
     ).toBeNull();
   });

package/src/react/model-context/toolbox.ts

@@ -9,6 +9,15 @@ import type {
   ToolCallMessagePartComponent,
   ToolCallMessagePartProps,
 } from "../types/MessagePartComponentTypes";
+import {
+  resolveToolCallText,
+  type ToolCallText as NeutralToolCallText,
+} from "../../model-context/tool-call-text";
+
+export type ToolCallText<
+  TArgs extends Record<string, unknown>,
+  TResult,
+> = NeutralToolCallText<TArgs, TResult, ReactNode>;
 
 /**
  * Resolves whether a tool's UI should be presented standalone (outside the
@@ -61,42 +70,6 @@ type ToolStreamCall<TArgs extends Record<string, unknown>, TResult> = (
   context: ToolExecuteContext,
 ) => void;
 
-type ToolCallRunningText<TArgs extends Record<string, unknown>> =
-  | ReactNode
-  | ((options: { args: TArgs }) => ReactNode);
-
-type ToolCallCompleteText<TArgs extends Record<string, unknown>, TResult> =
-  | ReactNode
-  | ((options: { args: TArgs; result: TResult | undefined }) => ReactNode);
-
-export type ToolCallText<TArgs extends Record<string, unknown>, TResult> =
-  | {
-      running: ToolCallRunningText<TArgs>;
-      complete?: ToolCallCompleteText<TArgs, TResult> | undefined;
-    }
-  | {
-      running?: ToolCallRunningText<TArgs> | undefined;
-      complete: ToolCallCompleteText<TArgs, TResult>;
-    };
-
-const resolveToolCallText = <TArgs extends Record<string, unknown>, TResult>(
-  text: ToolCallText<TArgs, TResult>,
-  part: ToolCallMessagePartProps<TArgs, TResult>,
-): ReactNode => {
-  const isRunning =
-    part.status?.type === "running" || part.status?.type === "requires-action";
-
-  if (!isRunning) {
-    const value = text.complete;
-    if (typeof value !== "function") return value ?? null;
-    return value({ args: part.args, result: part.result });
-  }
-
-  const value = text.running;
-  if (typeof value !== "function") return value ?? null;
-  return value({ args: part.args });
-};
-
 export const makeToolCallTextComponent = <
   TArgs extends Record<string, unknown>,
   TResult,

package/src/react/model-context/useAssistantContext.ts

@@ -1,4 +1,4 @@
-import { useEffect, useRef } from "react";
+import { useEffect, useInsertionEffect, useRef } from "react";
 import { useAui } from "@assistant-ui/store";
 import type { AssistantContextConfig } from "../..";
 
@@ -8,7 +8,9 @@ export const useAssistantContext = (config: AssistantContextConfig) => {
   const { getContext, disabled = false } = config;
   const aui = useAui();
   const getContextRef = useRef(getContext);
-  getContextRef.current = getContext;
+  useInsertionEffect(() => {
+    getContextRef.current = getContext;
+  }, [getContext]);
 
   useEffect(() => {
     if (disabled) return;

package/src/react/model-context/useAssistantTool.ts

@@ -12,7 +12,7 @@ import {
  * Props used to register a tool from React.
  *
  * @deprecated Use a toolkit with `Tools({ toolkit })` and register it via
- * `useAui({ tools: Tools({ toolkit }) })` instead. See
+ * `AuiConfig({ tools: Tools({ toolkit }) })` on the provider's `config` prop instead. See
  * https://assistant-ui.com/docs/migrations/toolkit-tools.
  */
 export type AssistantToolProps<
@@ -39,7 +39,7 @@ export type AssistantToolProps<
  * @param tool - Tool definition and name to register.
  *
  * @deprecated Use a toolkit with `Tools({ toolkit })` and register it via
- * `useAui({ tools: Tools({ toolkit }) })` instead. See
+ * `AuiConfig({ tools: Tools({ toolkit }) })` on the provider's `config` prop instead. See
  * https://assistant-ui.com/docs/migrations/toolkit-tools.
  *
  * @example

package/src/react/model-context/useAuiToolOverrides.ts

@@ -1,4 +1,4 @@
-import { useEffect, useRef } from "react";
+import { useEffect, useInsertionEffect, useRef } from "react";
 import { useAui } from "@assistant-ui/store";
 import type { Tool } from "assistant-stream";
 
@@ -25,7 +25,9 @@ type AuiToolOverrides = Record<string, AuiToolOverride<any, any>>;
 export function useAuiToolOverrides(overrides: AuiToolOverrides): void {
   const aui = useAui();
   const overridesRef = useRef(overrides);
-  overridesRef.current = overrides;
+  useInsertionEffect(() => {
+    overridesRef.current = overrides;
+  }, [overrides]);
 
   useEffect(() => {
     return aui.modelContext.register({

package/src/react/model-context/useInlineRender.test.tsx

@@ -0,0 +1,48 @@
+// @vitest-environment jsdom
+import { useState, type FC } from "react";
+import { act, cleanup, render, screen } from "@testing-library/react";
+import { afterEach, describe, expect, it } from "vitest";
+import { useInlineRender } from "./useInlineRender";
+import type { ToolCallMessagePartProps } from "../types/MessagePartComponentTypes";
+
+afterEach(() => cleanup());
+
+type Props = ToolCallMessagePartProps<any, any>;
+
+// The point of useInlineRender is a component whose identity survives a
+// changing toolUI, so inner state must not be lost when toolUI swaps.
+const makeToolUI = (label: string): FC<Props> =>
+  function Inner() {
+    const [n, setN] = useState(0);
+    return (
+      <button onClick={() => setN((v) => v + 1)}>
+        {label}:{n}
+      </button>
+    );
+  };
+
+describe("useInlineRender", () => {
+  it("follows toolUI changes without remounting the inner tree", async () => {
+    let swap!: (fc: FC<Props>) => void;
+    const identities = new Set<unknown>();
+
+    const Host = () => {
+      const [toolUI, setToolUI] = useState(() => makeToolUI("a"));
+      swap = (fc) => setToolUI(() => fc);
+      const ToolUI = useInlineRender(toolUI);
+      identities.add(ToolUI);
+      return <ToolUI {...({} as Props)} />;
+    };
+
+    render(<Host />);
+    expect(screen.getByRole("button").textContent).toBe("a:0");
+
+    await act(async () => screen.getByRole("button").click());
+    expect(screen.getByRole("button").textContent).toBe("a:1");
+
+    await act(async () => swap(makeToolUI("b")));
+    // label follows the new toolUI, counter survives => no remount
+    expect(screen.getByRole("button").textContent).toBe("b:1");
+    expect(identities.size).toBe(1);
+  });
+});

package/src/react/model-context/useInlineRender.ts

@@ -1,25 +1,29 @@
 import { type FC, useCallback, useEffect, useState } from "react";
 import type { ToolCallMessagePartProps } from "../types/MessagePartComponentTypes";
-import { create } from "zustand";
+import { WritableSubscribable } from "../../subscribable/subscribable";
+import { useSubscribable } from "../../store/runtime-clients/useSubscribable";
 
 export const useInlineRender = <TArgs, TResult>(
   toolUI: FC<ToolCallMessagePartProps<TArgs, TResult>>,
 ): FC<ToolCallMessagePartProps<TArgs, TResult>> => {
-  const [useToolUIStore] = useState(() =>
-    create(() => ({
-      toolUI,
-    })),
+  const [toolUIStore] = useState(
+    () =>
+      new WritableSubscribable<FC<ToolCallMessagePartProps<TArgs, TResult>>>(
+        toolUI,
+      ),
   );
 
   useEffect(() => {
-    useToolUIStore.setState({ toolUI });
-  }, [toolUI, useToolUIStore]);
+    toolUIStore.setState(toolUI);
+  }, [toolUI, toolUIStore]);
 
   return useCallback(
     function ToolUI(args) {
-      const store = useToolUIStore();
-      return store.toolUI(args);
+      // Invoked as a plain function so its hooks belong to this component,
+      // which keeps its identity while `toolUI` changes underneath.
+      const currentToolUI = useSubscribable(toolUIStore);
+      return currentToolUI(args);
     },
-    [useToolUIStore],
+    [toolUIStore],
   );
 };