reverse-engineering-anything/src/browser/CdpWebMcpDiscovery.ts

import { createHash } from "node:crypto";

import type { ProgressReporter } from "../application/ProgressReporter.js";
import type {
  DiscoverWebMcpToolsInput,
  WebMcpDiscovery,
} from "../domain/webMcpDiscovery.js";
import { inferJsonShape } from "../domain/jsonShape.js";
import type { CdpEndpointDiscovery, CdpEndpointTarget } from "./CdpEndpoint.js";
import type { CdpConnection, CdpEvent } from "./CdpConnection.js";
import { CdpCaptureCompleteness } from "./CdpCaptureCompleteness.js";
import { BrowserObservationError } from "../domain/browserObservationError.js";
import {
  allowedSanitizedUrl,
  delayWithCancellation,
  numberValue,
  recordValue,
  recordsValue,
  stringValue,
  type UnknownRecord,
} from "./CdpCaptureValues.js";
import {
  captureFrames,
  mainFrameUrl,
  walkFrameTrees,
} from "./CdpCaptureDocuments.js";
import { optionalCdpCommand } from "./CdpOptionalCommand.js";

interface DiscoveryContext {
  readonly connection: CdpConnection;
  readonly sessionId: string | undefined;
  readonly discovery: CdpEndpointDiscovery;
  readonly target: CdpEndpointTarget;
  readonly input: DiscoverWebMcpToolsInput;
  readonly signal?: AbortSignal;
  readonly progress?: ProgressReporter;
}

/** Discover WebMCP registrations without evaluating or invoking page code. */
export const discoverWebMcp = async (
  context: DiscoveryContext,
): Promise<WebMcpDiscovery> => {
  const origins = new Set(context.input.allowed_origins);
  const limitations = [
    "Page-declared tool metadata is untrusted and is never registered or invoked by REA.",
    "Discovery uses only WebMCP.enable and toolsAdded/toolsRemoved; WebMCP.invokeTool is never called.",
  ];
  const completeness = new CdpCaptureCompleteness();
  await context.connection.send(
    "Page.enable",
    {},
    context.sessionId,
    context.signal,
  );
  const frameTree = await context.connection.send(
    "Page.getFrameTree",
    {},
    context.sessionId,
    context.signal,
  );
  const initialUrl = mainFrameUrl(frameTree);
  if (allowedSanitizedUrl(initialUrl, origins) === undefined)
    throw new BrowserObservationError("inspect_web_page", "target_not_allowed");
  const frames = captureFrames(
    frameTree,
    origins,
    undefined,
    completeness,
  ).items;
  const frameUrls = initialFrameScope(frameTree, frames, completeness);
  const tools = new Map<string, WebMcpDiscovery["tools"]["items"][number]>();
  const frameScope = {
    origins,
    frames: frameUrls,
    tools,
    completeness,
    transientFrames: new Set<string>(),
    resolvedTransientFrames: new Set<string>(),
  };
  const removeListener = context.connection.onEvent((event) => {
    if (event.sessionId !== context.sessionId) return;
    if (ingestFrameScopeEvent(event, frameScope)) return;
    ingestWebMcpEvent({
      event,
      input: context.input,
      frames: frameUrls,
      tools,
      completeness,
      transientFrames: frameScope.transientFrames,
      resolvedTransientFrames: frameScope.resolvedTransientFrames,
    });
  });
  await context.progress?.report({
    phase: "browser_observation",
    completed: 1,
    total: 2,
    message: "Enabling passive WebMCP discovery",
  });
  try {
    const enabled = await optionalCdpCommand(
      context,
      "WebMCP.enable",
      {},
      limitations,
    );
    const available = enabled !== undefined;
    if (!available) {
      completeness.unavailable("webmcp_tools");
    } else
      // CDP sends the enable response before the required toolsAdded replay.
      await delayWithCancellation(
        context.input.observation_ms,
        "discover_webmcp_tools",
        context.signal,
      );
    await assertStableAuthorizedFrame(context, initialUrl, origins);
    if (frameScope.transientFrames.size > 0)
      completeness.attachLimited("webmcp_tools");
    return buildWebMcpResult({
      context,
      frameTree,
      tools,
      completeness,
      limitations,
      available,
    });
  } finally {
    removeListener();
    await context.progress?.report({
      phase: "browser_observation",
      completed: 2,
      total: 2,
      message: "WebMCP discovery complete",
      terminal: true,
    });
  }
};

interface FrameScopeState {
  readonly origins: ReadonlySet<string>;
  readonly frames: Map<string, FrameScopeFrame>;
  readonly tools: Map<string, WebMcpDiscovery["tools"]["items"][number]>;
  readonly completeness: CdpCaptureCompleteness;
  readonly transientFrames: Set<string>;
  readonly resolvedTransientFrames: Set<string>;
}

interface FrameScopeFrame {
  readonly url: string;
  readonly origin: string | null;
  readonly parentFrameId: string | null;
  readonly loaderId?: string | undefined;
}

const initialFrameScope = (
  frameTree: unknown,
  frames: ReturnType<typeof captureFrames>["items"],
  completeness: CdpCaptureCompleteness,
): Map<string, FrameScopeFrame> => {
  const scoped = new Map<string, FrameScopeFrame>(
    frames.map((frame) => [
      frame.frame_id,
      {
        url: frame.url,
        origin: frame.origin,
        parentFrameId: frame.parent_frame_id,
      },
    ]),
  );
  const root = recordValue(recordValue(frameTree)?.frameTree);
  if (root === undefined) return scoped;
  for (const tree of walkFrameTrees(root, () =>
    completeness.exclude("webmcp_tools", "invalid_protocol_value"),
  )) {
    const frame = recordValue(tree.frame);
    const frameId = stringValue(frame?.id);
    const prior = frameId === undefined ? undefined : scoped.get(frameId);
    if (frameId !== undefined && prior !== undefined)
      scoped.set(frameId, { ...prior, loaderId: stringValue(frame?.loaderId) });
  }
  return scoped;
};

const ingestFrameScopeEvent = (
  event: CdpEvent,
  state: FrameScopeState,
): boolean => {
  const params = recordValue(event.params);
  if (params === undefined) return false;
  switch (event.method) {
    case "Page.frameDetached":
      return handleFrameDetached(params, state);
    case "Page.frameNavigated":
      return handleFrameUpdate(recordValue(params.frame), state, true);
    case "Page.navigatedWithinDocument":
      return handleFrameUpdate(params, state, false);
    default:
      return false;
  }
};

const handleFrameDetached = (
  params: UnknownRecord,
  state: FrameScopeState,
): boolean => {
  const frameId = stringValue(params.frameId);
  if (frameId === undefined) {
    state.completeness.exclude("webmcp_tools", "invalid_protocol_value");
    return true;
  }
  removeFrame(frameId, state.frames, state.tools);
  state.transientFrames.delete(frameId);
  state.resolvedTransientFrames.delete(frameId);
  return true;
};

const handleFrameUpdate = (
  frame: UnknownRecord | undefined,
  state: FrameScopeState,
  documentCommitted: boolean,
): boolean => {
  if (frame === undefined) {
    state.completeness.exclude("webmcp_tools", "invalid_protocol_value");
    return true;
  }
  const frameId = stringValue(frame?.id) ?? stringValue(frame?.frameId);
  if (frameId === undefined) {
    state.completeness.exclude("webmcp_tools", "invalid_protocol_value");
    return true;
  }
  const rawUrl = stringValue(frame?.url);
  // Chromium can briefly commit an inherited, empty document while replacing
  // a subframe. That intermediate URL says nothing about the frame's eventual
  // scope or its registered tools; wait for the next committed URL (or the
  // authoritative frame-tree re-probe at the end of discovery).
  if (isTransientFrameUrl(rawUrl)) {
    if (documentCommitted) {
      removeFrameTools(frameId, state.tools);
      state.transientFrames.add(frameId);
      state.resolvedTransientFrames.delete(frameId);
    }
    return true;
  }
  updateFrameScope(frame, frameId, state, documentCommitted);
  return true;
};

const updateFrameScope = (
  frame: UnknownRecord,
  frameId: string,
  state: FrameScopeState,
  documentCommitted: boolean,
): void => {
  const url = allowedSanitizedUrl(frame?.url, state.origins);
  if (url === undefined || url.origin === null) {
    removeFrameTools(frameId, state.tools);
    state.frames.delete(frameId);
    state.completeness.exclude("webmcp_tools", "out_of_target_scope");
    return;
  }
  const previous = state.frames.get(frameId);
  const loaderId = stringValue(frame?.loaderId);
  if (hasNewDocument(documentCommitted, previous, url.url, loaderId))
    removeFrameTools(frameId, state.tools);
  if (documentCommitted && state.transientFrames.has(frameId))
    state.resolvedTransientFrames.add(frameId);
  state.frames.set(frameId, {
    url: url.url,
    origin: url.origin,
    ...(loaderId === undefined ? {} : { loaderId }),
    parentFrameId:
      (stringValue(frame?.parentId) ??
        state.frames.get(frameId)?.parentFrameId ??
        "") ||
      null,
  });
};

const hasNewDocument = (
  documentCommitted: boolean,
  previous: FrameScopeFrame | undefined,
  url: string,
  loaderId: string | undefined,
): boolean =>
  documentCommitted &&
  previous !== undefined &&
  (previous.url !== url ||
    (loaderId !== undefined &&
      previous.loaderId !== undefined &&
      previous.loaderId !== loaderId));

const removeFrameTools = (
  frameId: string,
  tools: Map<string, WebMcpDiscovery["tools"]["items"][number]>,
): void => {
  for (const [key, tool] of tools)
    if (tool.frame_id === frameId) tools.delete(key);
};

const isTransientFrameUrl = (url: string | undefined): boolean =>
  url === "" || url === "about:blank" || url === "about:srcdoc";

const removeFrame = (
  frameId: string,
  frames: Map<string, FrameScopeFrame>,
  tools: Map<string, WebMcpDiscovery["tools"]["items"][number]>,
): void => {
  const pending = [frameId];
  const removed = new Set<string>();
  while (pending.length > 0) {
    const current = pending.pop();
    if (current === undefined || removed.has(current)) continue;
    removed.add(current);
    for (const [candidate, frame] of frames)
      if (frame.parentFrameId === current) pending.push(candidate);
    frames.delete(current);
  }
  for (const [key, tool] of tools)
    if (removed.has(tool.frame_id)) tools.delete(key);
};

const assertStableAuthorizedFrame = async (
  context: DiscoveryContext,
  initialUrl: string | undefined,
  origins: ReadonlySet<string>,
): Promise<void> => {
  const finalTree = await context.connection.send(
    "Page.getFrameTree",
    {},
    context.sessionId,
    context.signal,
  );
  const finalUrl = mainFrameUrl(finalTree);
  if (allowedSanitizedUrl(finalUrl, origins) === undefined)
    throw new BrowserObservationError("inspect_web_page", "target_not_allowed");
  if (finalUrl !== initialUrl)
    throw new BrowserObservationError("inspect_web_page", "target_changed");
};

interface WebMcpIngestOptions {
  readonly event: CdpEvent;
  readonly input: DiscoverWebMcpToolsInput;
  readonly frames: ReadonlyMap<
    string,
    { readonly url: string; readonly origin: string | null }
  >;
  readonly tools: Map<string, WebMcpDiscovery["tools"]["items"][number]>;
  readonly completeness: CdpCaptureCompleteness;
  readonly transientFrames: Set<string>;
  readonly resolvedTransientFrames: Set<string>;
}

const ingestWebMcpEvent = (options: WebMcpIngestOptions): void => {
  const {
    event,
    input,
    frames,
    tools,
    completeness,
    transientFrames,
    resolvedTransientFrames,
  } = options;
  const params = recordValue(event.params);
  if (params === undefined) return;
  if (event.method === "WebMCP.toolsRemoved") {
    for (const removed of recordsValue(params.tools)) {
      const name = stringValue(removed.name);
      const frameId = stringValue(removed.frameId);
      if (name === undefined || frameId === undefined) continue;
      for (const [key, tool] of tools)
        if (tool.name === name && tool.frame_id === frameId) tools.delete(key);
    }
    return;
  }
  if (event.method !== "WebMCP.toolsAdded") return;
  for (const declared of recordsValue(params.tools)) {
    const frameId = stringValue(declared.frameId);
    if (frameId !== undefined && transientFrames.has(frameId)) {
      if (!resolvedTransientFrames.has(frameId)) continue;
      transientFrames.delete(frameId);
      resolvedTransientFrames.delete(frameId);
    }
    const normalized = normalizeTool(declared, input, frames, completeness);
    if (normalized === undefined) continue;
    tools.set(normalized.tool_key, normalized);
  }
};

const normalizeTool = (
  value: UnknownRecord,
  input: DiscoverWebMcpToolsInput,
  frames: ReadonlyMap<
    string,
    { readonly url: string; readonly origin: string | null }
  >,
  completeness: CdpCaptureCompleteness,
): WebMcpDiscovery["tools"]["items"][number] | undefined => {
  const frameId = stringValue(value.frameId);
  const name = stringValue(value.name);
  const frame = frameId === undefined ? undefined : frames.get(frameId);
  if (
    frameId === undefined ||
    name === undefined ||
    frame?.origin === null ||
    frame === undefined
  ) {
    completeness.exclude("webmcp_tools", "out_of_target_scope");
    return undefined;
  }
  const description = stringValue(value.description) ?? "";
  const annotations = recordValue(value.annotations);
  return {
    tool_key: toolKey(frame.url, name),
    name,
    description,
    frame_id: frameId,
    frame_url: frame.url,
    owner_origin: frame.origin,
    declaration_kind:
      numberValue(value.backendNodeId) === undefined
        ? "imperative"
        : "declarative",
    input_schema_shape: schemaShape(value.inputSchema),
    annotations: {
      read_only: booleanOrNull(annotations?.readOnly),
      untrusted_content: booleanOrNull(annotations?.untrustedContent),
      autosubmit: booleanOrNull(annotations?.autosubmit),
    },
    registration_source: registrationSource(
      recordValue(value.stackTrace),
      new Set(input.allowed_origins),
    ),
    trust: "page-declared-untrusted",
  };
};

const schemaShape = (value: unknown) => {
  if (recordValue(value) === undefined) return null;
  const encoded = JSON.stringify(value);
  const shape = inferJsonShape(encoded);
  if (shape === null)
    throw new BrowserObservationError("inspect_web_page", "protocol_error");
  return shape;
};

const registrationSource = (
  stack: UnknownRecord | undefined,
  allowedOrigins: ReadonlySet<string>,
): WebMcpDiscovery["tools"]["items"][number]["registration_source"] => {
  const frame = recordsValue(stack?.callFrames)[0];
  const url = allowedSanitizedUrl(frame?.url, allowedOrigins);
  return frame === undefined || url === undefined
    ? null
    : {
        url: url.url,
        line: integerOrNull(frame.lineNumber),
        column: integerOrNull(frame.columnNumber),
      };
};

interface WebMcpResultOptions {
  readonly context: DiscoveryContext;
  readonly frameTree: unknown;
  readonly tools: ReadonlyMap<
    string,
    WebMcpDiscovery["tools"]["items"][number]
  >;
  readonly completeness: CdpCaptureCompleteness;
  readonly limitations: string[];
  readonly available: boolean;
}

const buildWebMcpResult = (options: WebMcpResultOptions): WebMcpDiscovery => {
  const { context, frameTree, tools, completeness, limitations, available } =
    options;
  const targetUrl = mainFrameUrl(frameTree) ?? context.target.url;
  const sanitized = allowedSanitizedUrl(
    targetUrl,
    new Set(context.input.allowed_origins),
  );
  return {
    browser: context.discovery.version,
    target: {
      target_id: context.target.id,
      url: sanitized?.url ?? "[unsupported-url]",
      origin: sanitized?.origin ?? "",
    },
    status: available ? "available" : "unavailable",
    tools: {
      total: tools.size + completeness.droppedTotal,
      items: [...tools.values()].sort((left, right) =>
        left.tool_key.localeCompare(right.tool_key),
      ),
    },
    completeness: completeness.snapshot(),
    limitations,
  };
};

const toolKey = (frameUrl: string, name: string): string =>
  `webmcp_${createHash("sha256").update(`${frameUrl}\0${name}`).digest("hex")}`;

const booleanOrNull = (value: unknown): boolean | null =>
  typeof value === "boolean" ? value : null;

const integerOrNull = (value: unknown): number | null => {
  const number = numberValue(value);
  return number === undefined ? null : Math.max(0, Math.trunc(number));
};