@assistant-ui/react/src/unstable/webmcp/useWebMcpProvider.ts

Open-source TypeScript/React library for building production-grade AI chat experiences

"use client";

import {
  useCallback,
  useEffect,
  useInsertionEffect,
  useRef,
  useState,
  useSyncExternalStore,
} from "react";
import { useResources, useTapRoot, withKey } from "@assistant-ui/tap";
import { useAui } from "@assistant-ui/store";
import type { Tool } from "assistant-stream";
import { getDefaultWebMcpHost, type WebMcpHost } from "./webmcp-host";
import { defaultWebMcpFilter, toWebMcpInputSchema } from "./convertTools";
import { WebMcpRegistrationResource } from "./WebMcpRegistrationResource";
import { shallowEqual } from "@assistant-ui/store/internal";
import {
  useModelContextSnapshot,
  type ModelContextSnapshotSource,
} from "../useModelContextSnapshot";

/** @deprecated Experimental since 2026-08-29. Not scheduled for removal; the API may change in any release. */
export type Unstable_WebMcpProviderOptions = {
  filter?: (name: string, tool: Tool<any, any>) => boolean;
};

/** @deprecated Experimental since 2026-08-29. Not scheduled for removal; the API may change in any release. */
export type Unstable_WebMcpProviderResult = {
  status: "unsupported" | "active";
  registeredToolNames: readonly string[];
};

const EMPTY_NAMES: readonly string[] = Object.freeze([]);
const EMPTY_TOOLS: Record<string, Tool<any, any>> = Object.freeze({});

// The description is re-read on every sync so mutating it in place is
// observed; the schema is converted only when the tool object itself changes.
// Keyed on the tool so a schema or filter that throws warns once instead of
// on every model-context notify.
const warned = new WeakSet<Tool<any, any>>();

const signatures = new WeakMap<
  Tool<any, any>,
  { description: string | undefined; signature: string }
>();

const signatureOf = (tool: Tool<any, any>) => {
  const cached = signatures.get(tool);
  if (cached && cached.description === tool.description)
    return cached.signature;
  const signature = `${tool.description ?? ""}\u0000${JSON.stringify(
    toWebMcpInputSchema(tool),
  )}`;
  signatures.set(tool, { description: tool.description, signature });
  return signature;
};

const NO_SUBSCRIPTION = () => {};

// No `isEqual`: a caller republishing an unchanged tool set is how it asks for
// a re-sync, which is what re-attempts a tool whose filter has stopped throwing
// and what picks up a description edited in place on a stable tool object.
const modelContextToolSource: ModelContextSnapshotSource<
  Record<string, Tool<any, any>>
> = {
  empty: EMPTY_TOOLS,
  read: (aui) => aui.modelContext.getModelContext().tools ?? EMPTY_TOOLS,
  subscribe: (aui, onChange) =>
    aui.modelContext.subscribe?.(onChange) ?? NO_SUBSCRIPTION,
};

const useStableNames = (names: readonly (string | null)[]) => {
  const [cell] = useState(() => ({ names: EMPTY_NAMES }));
  const next = names.filter((name): name is string => name !== null).sort();
  const previous = cell.names;
  if (shallowEqual(previous, next)) {
    return previous;
  }
  cell.names = next;
  return next;
};

const useWebMcpRegistry = ({
  host,
  filter,
  tools,
  getCurrentTool,
}: {
  host: WebMcpHost;
  filter: (name: string, tool: Tool<any, any>) => boolean;
  tools: Record<string, Tool<any, any>>;
  getCurrentTool: (name: string) => Tool<any, any> | undefined;
}) => {
  const elements = [];
  for (const [name, tool] of Object.entries(tools)) {
    try {
      if (!filter(name, tool)) continue;
      const signature = signatureOf(tool);
      elements.push(
        withKey(
          name,
          WebMcpRegistrationResource({
            host,
            name,
            signature,
            tool,
            getCurrentTool,
          }),
          [host, name, signature, tool, getCurrentTool],
        ),
      );
    } catch (error) {
      if (warned.has(tool)) continue;
      warned.add(tool);
      console.warn(
        `[assistant-ui] Skipping WebMCP registration for tool "${name}": filter or schema conversion failed.`,
        error,
      );
    }
  }

  return useStableNames(useResources(elements));
};

/**
 * Publishes the frontend tools in the model context to a WebMCP-capable
 * browser, so the user's own browser agent can call them.
 *
 * Returns `status: "unsupported"` when the page exposes no
 * `document.modelContext` (or `navigator.modelContext`), and the sorted names
 * of the tools currently registered with the host.
 *
 * @deprecated Experimental since 2026-08-29. Not scheduled for removal; the API may change in any release.
 */
export const unstable_useWebMcpProvider = (
  options: Unstable_WebMcpProviderOptions = {},
): Unstable_WebMcpProviderResult => {
  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({
      host,
      filter,
      tools,
      getCurrentTool,
    });
  });
  const registeredToolNames = useSyncExternalStore(
    root.subscribe,
    root.getValue,
    () => EMPTY_NAMES,
  );

  const [published, setPublished] = useState(false);
  useEffect(() => {
    if (host.available) setPublished(true);
  }, [host]);

  return {
    status: published ? "active" : "unsupported",
    registeredToolNames,
  };
};