@assistant-ui/react 0.15.22 → 0.15.23

package/src/unstable/webmcp/WebMcpRegistrationResource.ts

@@ -10,6 +10,7 @@ export type WebMcpRegistrationProps = {
   /** Re-registers when it changes; an execute-only edit reads through instead. */
   signature: string;
   tool: Tool<any, any>;
+  getCurrentTool: (name: string) => Tool<any, any> | undefined;
 };
 
 // The explainer specifies NotAllowedError for a page whose tools permission is
@@ -26,16 +27,12 @@ const useWebMcpRegistration = ({
   name,
   signature,
   tool,
+  getCurrentTool,
 }: WebMcpRegistrationProps): string | null => {
   // A refused name is remembered for as long as the tool stays in the model
   // context, so a permanent collision warns once rather than on every sync.
   const [refused, setRefused] = useState(false);
-  const toolRef = useRef(tool);
-  // Commit phase, so an abandoned render cannot hand the host a tool the tree
-  // never committed.
-  useEffect(() => {
-    toolRef.current = tool;
-  }, [tool]);
+  const fallbackToolRef = useRef(tool);
 
   useEffect(() => {
     if (refused) return undefined;
@@ -53,7 +50,11 @@ const useWebMcpRegistration = ({
 
     try {
       dispose = host.registerTool(
-        toWebMcpTool(name, () => toolRef.current, lifecycle.signal),
+        toWebMcpTool(
+          name,
+          () => getCurrentTool(name) ?? fallbackToolRef.current,
+          lifecycle.signal,
+        ),
         (error) =>
           refuse(
             notPermitted(error)
@@ -84,7 +85,7 @@ const useWebMcpRegistration = ({
         );
       }
     };
-  }, [host, name, signature, refused]);
+  }, [getCurrentTool, host, name, signature, refused]);
 
   return refused ? null : name;
 };

package/src/unstable/webmcp/useWebMcpProvider.test.tsx

@@ -1,7 +1,7 @@
 // @vitest-environment jsdom
 
-import { StrictMode, type ReactNode } from "react";
-import { cleanup } from "@testing-library/react";
+import { StrictMode, type ReactNode, useLayoutEffect } from "react";
+import { act, cleanup } from "@testing-library/react";
 import { afterEach, describe, expect, it, vi } from "vitest";
 import type { Tool } from "assistant-stream";
 import type {
@@ -195,6 +195,40 @@ describe("unstable_useWebMcpProvider", () => {
     expect(host.unregisterCalls).toEqual([]);
   });
 
+  it("calls through to the committed implementation from a descendant layout effect", async () => {
+    const host = useHost(createFakeWebMcpHost());
+    const provider = createProvider({
+      search: frontendTool({ execute: async () => "first" }),
+    });
+    let invoke = false;
+    let execution: ReturnType<WebMcpToolDescriptor["execute"]> | undefined;
+
+    const InvokeOnCommit = ({ children }: { children: ReactNode }) => {
+      useLayoutEffect(() => {
+        if (invoke) execution = host.registry.get("search")?.execute({});
+      });
+      return children;
+    };
+    const { rerender } = mountProvider(provider, {}, (children) => (
+      <InvokeOnCommit>{children}</InvokeOnCommit>
+    ));
+    await waitForNames(["search"]);
+
+    act(() => {
+      invoke = true;
+      provider.setTools({
+        search: frontendTool({ execute: async () => "second" }),
+      });
+      rerender({});
+    });
+
+    await expect(execution).resolves.toEqual({
+      content: [{ type: "text", text: "second" }],
+    });
+    expect(host.registerCalls).toEqual(["search"]);
+    expect(host.unregisterCalls).toEqual([]);
+  });
+
   it("re-registers when the description or the schema changes", async () => {
     const host = useHost(createFakeWebMcpHost());
     const provider = createProvider({ search: frontendTool() });

package/src/unstable/webmcp/useWebMcpProvider.ts

@@ -1,8 +1,15 @@
 "use client";
 
-import { useEffect, useState, useSyncExternalStore } from "react";
+import {
+  useCallback,
+  useEffect,
+  useInsertionEffect,
+  useRef,
+  useState,
+  useSyncExternalStore,
+} from "react";
 import { useResources, useTapRoot, withKey } from "@assistant-ui/tap";
-import { useAui, type AssistantClient } from "@assistant-ui/store";
+import { useAui } from "@assistant-ui/store";
 import type { Tool } from "assistant-stream";
 import { getDefaultWebMcpHost, type WebMcpHost } from "./webmcp-host";
 import { defaultWebMcpFilter, toWebMcpInputSchema } from "./convertTools";
@@ -73,20 +80,16 @@ const useStableNames = (names: readonly (string | null)[]) => {
 };
 
 const useWebMcpRegistry = ({
-  aui,
   host,
   filter,
+  tools,
+  getCurrentTool,
 }: {
-  aui: AssistantClient;
   host: WebMcpHost;
   filter: (name: string, tool: Tool<any, any>) => boolean;
+  tools: Record<string, Tool<any, any>>;
+  getCurrentTool: (name: string) => Tool<any, any> | undefined;
 }) => {
-  const tools = useModelContextSnapshot(
-    aui,
-    host.available,
-    modelContextToolSource,
-  );
-
   const elements = [];
   for (const [name, tool] of Object.entries(tools)) {
     try {
@@ -95,8 +98,14 @@ const useWebMcpRegistry = ({
       elements.push(
         withKey(
           name,
-          WebMcpRegistrationResource({ host, name, signature, tool }),
-          [host, name, signature, tool],
+          WebMcpRegistrationResource({
+            host,
+            name,
+            signature,
+            tool,
+            getCurrentTool,
+          }),
+          [host, name, signature, tool, getCurrentTool],
         ),
       );
     } catch (error) {
@@ -126,9 +135,27 @@ export const unstable_useWebMcpProvider = (
   const aui = useAui();
   const [host] = useState(getDefaultWebMcpHost);
   const filter = options.filter ?? defaultWebMcpFilter;
+  const tools = useModelContextSnapshot(
+    aui,
+    host.available,
+    modelContextToolSource,
+  );
+  const currentToolsRef = useRef(tools);
+  useInsertionEffect(() => {
+    currentToolsRef.current = tools;
+  }, [tools]);
+  const getCurrentTool = useCallback(
+    (name: string) => currentToolsRef.current[name],
+    [],
+  );
 
   const root = useTapRoot(function WebMcpProviderRoot() {
-    return useWebMcpRegistry({ aui, host, filter });
+    return useWebMcpRegistry({
+      host,
+      filter,
+      tools,
+      getCurrentTool,
+    });
   });
   const registeredToolNames = useSyncExternalStore(
     root.subscribe,