React Embedded AnalyticsWebMCP

"use client";

import React, {
  useCallback,
  useMemo,
  useRef,
  useState,
  StrictMode,
} from "react";
import { createRoot } from "react-dom/client";
import { AgStudio, AgStudioProvider, AgStudioRef } from "ag-studio-react";
import {
  AgDataEngine,
  AgDataSourcesDefinition,
  AgReportState,
  AgStudioAiModule,
  AgStudioApi,
  AgStudioMode,
  AgStudioProperties,
  enableStudioDevValidations,
} from "ag-studio";
import { getGhcnCitiesData } from "./shared/ghcnCities/data.tsx";
import { ghcnCitiesReportState } from "./shared/ghcnCities/state.tsx";
import { PageUpdater } from "./interfaces.tsx";
import { StudioWebMcpBridge, createWebMcpBridge } from "./webmcpBridge.tsx";

if (process.env.NODE_ENV !== "production") {
  // Enable extended validations only for development
  enableStudioDevValidations();
}

// `view_widget` is the only one of the four tools that can go uncallable: its widget-id enum comes
// from the *selected* page's widgets, so it is withheld whenever that page has none. Two controls
// show this - clearing the current page, and the page toolbar's blank page, which starts empty.
const populatedPages = ghcnCitiesReportState.pages;

let bridge: StudioWebMcpBridge | undefined;

let reconcileCount = 0;

const NOTICE_TEXT = [
  "This browser does not expose document.modelContext, so nothing is registered with the browser.",
  "WebMCP is experimental: Chrome 149+ behind an origin trial or chrome://flags/#enable-webmcp-testing,",
  "and it requires a secure context. The list below is what would be advertised.",
].join(" ");

const updateModeButton: (mode: string) => void = (mode: string) => {
  document.querySelector<HTMLElement>("#modeButton")!.textContent =
    mode === "edit" ? "Switch to view mode" : "Switch to edit mode";
};

const reconcileAndRender: () => Promise<void> = async () => {
  const errorEl = document.querySelector<HTMLElement>("#webmcpError")!;
  try {
    await bridge?.reconcile();
    renderPanel();
    errorEl.textContent = bridge?.lastError() ?? "";
  } catch (err) {
    errorEl.textContent = `Reconcile failed: ${err instanceof Error ? err.message : String(err)}`;
  }
};

/**
 * The notice and the call button belong to the panel, not to start-up: a framework renders its
 * template after mount, so anything written at mount time has nothing to write into yet.
 */
const renderSupportNotice: () => void = () => {
  if (document.modelContext != null) {
    return;
  }
  document.querySelector<HTMLElement>("#webmcpNotice")!.textContent =
    NOTICE_TEXT;
  document.querySelector<HTMLButtonElement>("#callToolButton")!.disabled = true;
};

const renderPanel: () => void = () => {
  renderSupportNotice();
  const container = document.querySelector<HTMLElement>("#advertisedTools")!;
  container.dataset.reconciles = String(++reconcileCount);
  container.replaceChildren();
  for (const { name, registrations } of bridge?.getAdvertisedTools() ?? []) {
    const item = document.createElement("div");
    item.className = "advertised-tool";
    item.dataset.tool = name;
    item.dataset.registrations = String(registrations);
    item.textContent = `${name} (registrations: ${registrations})`;
    container.appendChild(item);
  }
};

const renderLiveTools: () => Promise<void> = async () => {
  const live = document.querySelector<HTMLElement>("#liveTools")!;
  const { modelContext } = document;
  if (modelContext == null) return;
  try {
    const advertised = await modelContext.getTools();
    live.textContent = `Registered with the browser: ${advertised.map((tool) => tool.name).join(", ")}`;
  } catch (err) {
    document.querySelector<HTMLElement>("#webmcpError")!.textContent =
      `Could not read the registered tools: ${err instanceof Error ? err.message : String(err)}`;
  }
};

const StudioExample = () => {
  const studioRef = useRef<AgStudioRef>(null);
  const containerStyle = useMemo(() => ({ width: "100%", height: "100%" }), []);
  const studioStyle = useMemo(() => ({ height: "100%", width: "100%" }), []);
  const [data, setData] = useState<AgDataSource>(
    getGhcnCitiesData("https://www.ag-grid.com/studio/example-assets"),
  );
  const initialState = useMemo<AgReportState>(() => {
    return ghcnCitiesReportState;
  }, []);

  const onStudioReady = useCallback((event) => {
    const studio = event.api.getAiTools();
    bridge = createWebMcpBridge([
      { tool: studio.viewSchema(), readOnly: true },
      { tool: studio.viewPage(), readOnly: true },
      { tool: studio.viewWidget(), readOnly: true },
      { tool: studio.executeQuery(), readOnly: true },
    ]);
    void reconcileAndRender();
  }, []);

  const onStateUpdated = useCallback(() => {
    void reconcileAndRender();
  }, []);

  const onRenderStateChanged = useCallback(() => {
    void reconcileAndRender();
  }, []);

  const onStudioPreDestroyed = useCallback(() => {
    bridge?.destroy();
  }, []);

  const toggleMode = useCallback((): void => {
    const mode =
      studioRef.current!.api.getProperty("mode") === "edit" ? "view" : "edit";
    studioRef.current!.api.setProperty("mode", mode);
    updateModeButton(mode);
  }, []);

  const callTool = useCallback(async (): Promise<void> => {
    const output = document.querySelector<HTMLElement>("#toolOutput")!;
    const { modelContext } = document;
    if (modelContext == null) return;
    try {
      const advertised = await modelContext.getTools();
      const tool = advertised.find(
        (candidate) => candidate.name === "view_schema",
      );
      if (tool == null) {
        output.textContent = "view_schema is not currently advertised.";
        return;
      }
      // `view_schema` takes no parameters at all, so an empty object is a complete call.
      const result = await modelContext.executeTool(tool, JSON.stringify({}));
      // The browser hands back the tool's `content` array; older builds returned a bare string.
      output.textContent =
        typeof result === "string"
          ? result
          : result.content.map((part) => part.text).join("\n");
    } catch (err) {
      output.textContent = `Tool call failed: ${err instanceof Error ? err.message : String(err)}`;
    }
  }, []);

  // `api.getState()` hands back Studio's live state object, and `api.setState()` ignores a call whose
  // argument is that same reference. Every control below therefore builds a new state object, and a
  // new object for the slice it changes - Studio also compares each state slice by reference.
  const applyState = useCallback((state: AgReportState): void => {
    studioRef.current!.api.setState(state);
    void reconcileAndRender();
  }, []);

  const setPage = useCallback(
    (pageId: string): void => {
      applyState({
        ...studioRef.current!.api.getState(),
        selectedPageId: pageId,
      });
    },
    [applyState],
  );

  const addPage = useCallback((): void => {
    const state = studioRef.current!.api.getState();
    if (state.pages.some((page) => page.id === "secondary")) return;
    applyState({ ...state, pages: [...state.pages, { id: "secondary" }] });
  }, [applyState]);

  const removePage = useCallback((): void => {
    // Never remove the last page: with no pages `view_widget` goes uncallable, which would collide
    // the page scenario with the widget scenario.
    const state = studioRef.current!.api.getState();
    applyState({
      ...state,
      pages: state.pages.filter((page) => page.id !== "secondary"),
      selectedPageId: "temperature",
    });
  }, [applyState]);

  const addCalculatedField = useCallback((): void => {
    applyState({
      ...studioRef.current!.api.getState(),
      schema: {
        expressions: [
          { isMeasure: false, id: "temp_range", tableId: "weather" },
        ],
        fields: {
          temp_range: {
            name: "Temp Range",
            expression: "[weather.tmax] - [weather.tmin]",
          },
        },
      },
    });
  }, [applyState]);

  const removeCalculatedField = useCallback((): void => {
    applyState({ ...studioRef.current!.api.getState(), schema: {} });
  }, [applyState]);

  const updateSelectedPage = useCallback(
    (update: PageUpdater): void => {
      const state = studioRef.current!.api.getState();
      applyState({
        ...state,
        pages: state.pages.map((page) =>
          page.id === state.selectedPageId ? update(page) : page,
        ),
      });
    },
    [applyState],
  );

  const removeWidgets = useCallback((): void => {
    updateSelectedPage((page) => ({ ...page, widgets: {}, widgetLayout: {} }));
  }, [updateSelectedPage]);

  const restoreWidgets = useCallback((): void => {
    updateSelectedPage((page) => {
      const original = populatedPages.find(
        (candidate) => candidate.id === page.id,
      );
      return original
        ? {
            ...page,
            widgets: original.widgets,
            widgetLayout: original.widgetLayout,
          }
        : page;
    });
  }, [updateSelectedPage, populatedPages]);

  return (
    <div style={containerStyle}>
      <div style={{ display: "flex", flexDirection: "column", height: "100%" }}>
        <div className="example-controls">
          <div className="controls-row">
            <button type="button" onClick={() => setPage("temperature")}>
              Temperature
            </button>
            <button type="button" onClick={() => setPage("precipitation")}>
              Precipitation
            </button>
            <button type="button" onClick={() => setPage("blank")}>
              Blank
            </button>
          </div>
          <div className="controls-row">
            <button type="button" id="modeButton" onClick={toggleMode}>
              Switch to view mode
            </button>
            <button type="button" onClick={removeWidgets}>
              Remove widgets
            </button>
            <button type="button" onClick={restoreWidgets}>
              Restore widgets
            </button>
            <button type="button" onClick={addPage}>
              Add page
            </button>
            <button type="button" onClick={removePage}>
              Remove page
            </button>
            <button type="button" onClick={addCalculatedField}>
              Add calculated field
            </button>
            <button type="button" onClick={removeCalculatedField}>
              Remove calculated field
            </button>
            <button type="button" id="callToolButton" onClick={callTool}>
              Call a tool
            </button>
          </div>
        </div>
        <div className="webmcp-panel">
          <div className="webmcp-notice" id="webmcpNotice"></div>
          <div className="webmcp-error" id="webmcpError"></div>
          <div id="advertisedTools"></div>
          <div id="liveTools"></div>
          <div className="webmcp-output" id="toolOutput"></div>
        </div>

        <AgStudio
          ref={studioRef}
          style={studioStyle}
          className="my-studio-container"
          data={data}
          mode={"edit"}
          initialState={initialState}
          onStudioReady={onStudioReady}
          onStateUpdated={onStateUpdated}
          onRenderStateChanged={onRenderStateChanged}
          onStudioPreDestroyed={onStudioPreDestroyed}
        />
      </div>
    </div>
  );
};

const root = createRoot(document.getElementById("root")!);
root.render(
  <StrictMode>
    <AgStudioProvider modules={[AgStudioAiModule]}>
      <StudioExample />
    </AgStudioProvider>
  </StrictMode>,
);