ai agent

src/agent/mcpServer.ts

@@ -0,0 +1,499 @@
+/**
+ * Minimal MCP (Model Context Protocol) stdio server exposing the RA3 Mod XML
+ * agent snapshot query API.
+ *
+ * This is intentionally dependency-free. It speaks the JSON-RPC-over-stdio
+ * subset used by MCP clients:
+ *
+ *   initialize
+ *   notifications/initialized
+ *   ping
+ *   tools/list
+ *   tools/call
+ *
+ * Usage:
+ *   node out/agent/mcpServer.js --project D:/Mods/Example
+ *   node out/agent/mcpServer.js --snapshot /path/to/snapshot.json.gz
+ */
+
+import { createInterface } from "node:readline";
+import {
+  isProcessAlive,
+  readEndpoint,
+  readEndpointForProject,
+} from "./endpoint";
+import { readSnapshotFile, snapshotPathForProject } from "./snapshot";
+import {
+  findAssets,
+  findDefine,
+  findReferenceGroups,
+  isFileActive,
+  listAssetsByType,
+  resolveIncludeSource,
+  statusFromSnapshot,
+} from "./query";
+import type { AgentIndexSnapshot } from "./types";
+
+interface McpTool {
+  name: string;
+  description: string;
+  inputSchema: Record<string, unknown>;
+  handler: (args: Record<string, unknown>, snapshot: AgentIndexSnapshot | null) => unknown;
+}
+
+const USAGE_GUIDE = `RA3 Mod XML index query tools
+
+This MCP server exposes the semantic index built by the RA3 Mod XML VS Code extension.
+
+Use these tools instead of full-text grepping the XML tree when you need exact facts:
+- find_asset(id, type?) -> definition sites (file/line/origin/stream)
+- find_references(id, type?) -> semantic reference sites
+- get_asset_references(id, type?, depth?, targetTypes?) -> outgoing reference EDGES with element context
+- list_assets_by_type(type, prefix?, limit?) -> assets of a type
+- is_file_active(path) -> whether a file is part of an indexed include stream
+- find_define(name) -> $DEFINE definitions
+- resolve_include(source) -> candidate source file
+- get_status() -> current index state
+
+Tips:
+- Asset ids are case-insensitive.
+- Prefer passing type when the same id exists for multiple asset types.
+- Always check the returned index state; if it is stale/incomplete, treat results as provisional.
+- Use get_asset_references to follow "which weapon/model/upgrade does this asset use" chains.
+  It returns edges annotated with the element name, parent element and attribute that produced
+  them, plus the exact XML file/line, so you do not have to read source to find the link.
+  Start with depth 1 (the default) and pass targetTypes (e.g. ["WeaponTemplate"]) to cut noise.
+  Edges with a "definedIn" field come from an inheritFrom ancestor's XML.
+  When "truncated" is true, read "omittedByTargetType" and narrow the query instead of retrying.
+- Do not attempt to read the entire snapshot file; query narrowly.`;
+
+const TOOLS: McpTool[] = [
+  {
+    name: "get_status",
+    description: "Returns the current index state and basic statistics.",
+    inputSchema: { type: "object", properties: {} },
+    handler: (_args, snapshot) => statusFromSnapshot(snapshot),
+  },
+  {
+    name: "find_asset",
+    description: "Finds asset definitions by id, optionally filtered by asset type.",
+    inputSchema: {
+      type: "object",
+      properties: {
+        id: { type: "string", description: "Asset id to find" },
+        type: { type: "string", description: "Optional asset type filter" },
+      },
+      required: ["id"],
+    },
+    handler: (args, snapshot) => {
+      const id = String(args.id ?? "");
+      if (!snapshot) return statusFromSnapshot(snapshot);
+      return {
+        index: statusFromSnapshot(snapshot),
+        data: findAssets(snapshot, id, args.type ? String(args.type) : null),
+      };
+    },
+  },
+  {
+    name: "find_references",
+    description: "Finds semantic reference sites pointing to an asset id, optionally filtered by type.",
+    inputSchema: {
+      type: "object",
+      properties: {
+        id: { type: "string", description: "Asset id whose references to find" },
+        type: { type: "string", description: "Optional asset type filter" },
+      },
+      required: ["id"],
+    },
+    handler: (args, snapshot) => {
+      const id = String(args.id ?? "");
+      if (!snapshot) return statusFromSnapshot(snapshot);
+      return {
+        index: statusFromSnapshot(snapshot),
+        data: findReferenceGroups(snapshot, id, args.type ? String(args.type) : null),
+      };
+    },
+  },
+  {
+    name: "list_assets_by_type",
+    description: "Lists asset definitions of one type, optionally filtered by id prefix.",
+    inputSchema: {
+      type: "object",
+      properties: {
+        type: { type: "string", description: "Asset type" },
+        prefix: { type: "string", description: "Optional id prefix" },
+        limit: { type: "number", description: "Maximum number of results" },
+      },
+      required: ["type"],
+    },
+    handler: (args, snapshot) => {
+      const type = String(args.type ?? "");
+      if (!snapshot) return statusFromSnapshot(snapshot);
+      const limit = typeof args.limit === "number" ? args.limit : undefined;
+      return {
+        index: statusFromSnapshot(snapshot),
+        data: listAssetsByType(snapshot, type, args.prefix ? String(args.prefix) : "", limit),
+      };
+    },
+  },
+  {
+    name: "is_file_active",
+    description: "Returns whether a file belongs to an indexed include stream (i.e. is not a dead file).",
+    inputSchema: {
+      type: "object",
+      properties: {
+        path: { type: "string", description: "Absolute file path" },
+      },
+      required: ["path"],
+    },
+    handler: (args, snapshot) => {
+      const path = String(args.path ?? "");
+      if (!snapshot) return statusFromSnapshot(snapshot);
+      return {
+        index: statusFromSnapshot(snapshot),
+        data: { active: isFileActive(snapshot, path) },
+      };
+    },
+  },
+  {
+    name: "find_define",
+    description: "Finds $DEFINE constants by name.",
+    inputSchema: {
+      type: "object",
+      properties: {
+        name: { type: "string", description: "Define name (with or without leading $)" },
+      },
+      required: ["name"],
+    },
+    handler: (args, snapshot) => {
+      const name = String(args.name ?? "").replace(/^\$/, "");
+      if (!snapshot) return statusFromSnapshot(snapshot);
+      return {
+        index: statusFromSnapshot(snapshot),
+        data: findDefine(snapshot, name),
+      };
+    },
+  },
+  {
+    name: "resolve_include",
+    description: "Resolves an Include source string from the snapshot's candidate list.",
+    inputSchema: {
+      type: "object",
+      properties: {
+        source: { type: "string", description: "Include source, e.g. DATA:Units/Example.xml" },
+      },
+      required: ["source"],
+    },
+    handler: (args, snapshot) => {
+      const source = String(args.source ?? "");
+      if (!snapshot) return statusFromSnapshot(snapshot);
+      return {
+        index: statusFromSnapshot(snapshot),
+        data: resolveIncludeSource(snapshot, source),
+      };
+    },
+  },
+  {
+    name: "get_asset_references",
+    description:
+      "Returns the outgoing references (edges) of an asset: which assets it references, through which element/attribute, and at which file/line. Also follows inheritFrom ancestors (marked with definedIn). Live index required.",
+    inputSchema: {
+      type: "object",
+      properties: {
+        id: { type: "string", description: "Asset id whose outgoing references to return" },
+        type: { type: "string", description: "Optional asset type filter" },
+        depth: {
+          type: "number",
+          description:
+            "Levels of assets to expand: 1 (default) = the asset itself, including inherited XML; max 3.",
+        },
+        targetTypes: {
+          type: "array",
+          items: { type: "string" },
+          description:
+            "Only keep edges whose target is assignable to one of these types, e.g. [\"WeaponTemplate\"].",
+        },
+        maxEdges: { type: "number", description: "Hard cap on returned edges (default 200)." },
+        includeUnresolved: {
+          type: "boolean",
+          description: "Also return edges whose reference value could not be resolved.",
+        },
+      },
+      required: ["id"],
+    },
+    // Live-only: element context is not stored in the on-disk snapshot.
+    handler: () => ({
+      index: { state: "no_index" },
+      error:
+        "get_asset_references requires a live index. Open the project in VS Code (with AI Agent access enabled) and retry.",
+    }),
+  },
+  {
+    name: "get_usage_guide",
+    description: "Returns guidance for using the RA3 Mod XML index tools.",
+    inputSchema: { type: "object", properties: {} },
+    handler: () => ({ text: USAGE_GUIDE }),
+  },
+];
+
+/** Tools that can only be answered by the live extension server. */
+const LIVE_ONLY_TOOLS = new Set(["get_asset_references"]);
+
+export function liveUrlForTool(
+  endpointUrl: string,
+  projectDir: string | null,
+  toolName: string,
+  args: Record<string, unknown>,
+): string | null {
+  const base = endpointUrl.replace(/\/$/, "");
+  const q = new URLSearchParams();
+  // Always pin the requested project. Without this the server would silently
+  // answer from whatever project its active editor points at.
+  if (projectDir) q.set("project", projectDir);
+  switch (toolName) {    case "get_status":
+      return q.toString() ? `${base}/status?${q}` : `${base}/status`;
+    case "find_asset":
+      q.set("id", String(args.id ?? ""));
+      if (args.type != null) q.set("type", String(args.type));
+      return `${base}/find_asset?${q}`;
+    case "find_references":
+      q.set("id", String(args.id ?? ""));
+      if (args.type != null) q.set("type", String(args.type));
+      return `${base}/find_references?${q}`;
+    case "get_asset_references":
+      q.set("id", String(args.id ?? ""));
+      if (args.type != null) q.set("type", String(args.type));
+      if (args.depth != null) q.set("depth", String(args.depth));
+      if (Array.isArray(args.targetTypes)) {
+        q.set("targetTypes", (args.targetTypes as unknown[]).map(String).join(","));
+      }
+      if (args.maxEdges != null) q.set("maxEdges", String(args.maxEdges));
+      if (args.includeUnresolved != null) {
+        q.set("includeUnresolved", String(args.includeUnresolved));
+      }
+      return `${base}/get_asset_references?${q}`;
+    case "list_assets_by_type":
+      q.set("type", String(args.type ?? ""));
+      if (args.prefix != null) q.set("prefix", String(args.prefix));
+      if (args.limit != null) q.set("limit", String(args.limit));
+      return `${base}/list_assets?${q}`;
+    case "is_file_active":
+      q.set("path", String(args.path ?? ""));
+      return `${base}/is_file_active?${q}`;
+    case "find_define":
+      q.set("name", String(args.name ?? ""));
+      return `${base}/find_define?${q}`;
+    case "resolve_include":
+      q.set("source", String(args.source ?? ""));
+      return `${base}/resolve_include?${q}`;
+    default:
+      return null;
+  }
+}
+
+export function normalizePath(p: string): string {
+  return p.replace(/\\/g, "/").replace(/\/+$/, "").toLowerCase();
+}
+
+/**
+ * Rejects a live response that belongs to a different project than the one
+ * this MCP client was started for. This is the guard against the
+ * multi-window cross-talk described in docs/ai-agent-integration-plan.md §10:
+ * returning another project's data silently would be worse than returning
+ * nothing, because the agent cannot tell the difference.
+ */
+export function responseProjectMismatch(
+  payload: unknown,
+  projectDir: string | null,
+): boolean {
+  if (!projectDir) return false;
+  const index = (payload as { index?: { projectDir?: string } } | null)?.index;
+  const reported = index?.projectDir;
+  if (!reported) return false;
+  return normalizePath(reported) !== normalizePath(projectDir);
+}
+
+/** Cooldown after a failed live attempt, to avoid a probe per tool call. */
+const LIVE_RETRY_COOLDOWN_MS = 5000;
+let liveUnavailableUntil = 0;
+
+interface LiveResult {
+  payload: unknown;
+  /** True when the live server answered but for a different project. */
+  mismatched: boolean;
+}
+
+/** Tries the live extension server; returns null when unavailable. */
+async function tryLiveQuery(
+  toolName: string,
+  args: Record<string, unknown>,
+  projectDir: string | null,
+  agentHome?: string,
+): Promise<LiveResult | null> {
+  if (Date.now() < liveUnavailableUntil) return null;
+  // Prefer the per-project endpoint so two open windows cannot shadow each
+  // other; fall back to the legacy global file only when it matches.
+  let endpoint = projectDir
+    ? await readEndpointForProject(projectDir, agentHome)
+    : null;
+  if (!endpoint) {
+    const fallback = await readEndpoint(agentHome);
+    if (
+      fallback &&
+      (!projectDir ||
+        !fallback.projectDir ||
+        normalizePath(fallback.projectDir) === normalizePath(projectDir))
+    ) {
+      endpoint = fallback;
+    }
+  }
+  if (!endpoint) return null;
+  // A crashed VS Code can leave the file behind; a dead PID means stale.
+  if (!isProcessAlive(endpoint.processId)) return null;
+
+  const url = liveUrlForTool(endpoint.url, projectDir, toolName, args);
+  if (!url) return null;
+  const controller = new AbortController();
+  const timeout = setTimeout(() => controller.abort(), 1500);
+  try {
+    const res = await fetch(url, {
+      headers: { authorization: `Bearer ${endpoint.token}` },
+      signal: controller.signal,
+    });
+    if (!res.ok) {
+      if (res.status === 401 || res.status === 404) return null;
+      return null;
+    }
+    const payload: unknown = await res.json();
+    return { payload, mismatched: responseProjectMismatch(payload, projectDir) };
+  } catch {
+    liveUnavailableUntil = Date.now() + LIVE_RETRY_COOLDOWN_MS;
+    return null;
+  } finally {
+    clearTimeout(timeout);
+  }
+}
+
+function sendMessage(message: unknown): void {
+  process.stdout.write(`${JSON.stringify(message)}\n`);
+}
+
+function resultFor(id: unknown, result: unknown): unknown {
+  return { jsonrpc: "2.0", id, result };
+}
+
+function errorFor(id: unknown, code: number, message: string): unknown {
+  return { jsonrpc: "2.0", id, error: { code, message } };
+}
+
+async function handleRequest(
+  message: Record<string, unknown>,
+  snapshot: AgentIndexSnapshot | null,
+  projectDir: string | null,
+  agentHome?: string,
+): Promise<unknown | null> {
+  const method = String(message.method ?? "");
+  const id = message.id;
+  const params = (message.params ?? {}) as Record<string, unknown>;
+
+  switch (method) {
+    case "initialize":
+      return resultFor(id, {
+        protocolVersion: params.protocolVersion ?? "2024-11-05",
+        capabilities: { tools: {} },
+        serverInfo: { name: "ra3-mod-xml", version: "0.1.0" },
+      });
+    case "ping":
+      return resultFor(id, {});
+    case "tools/list":
+      return resultFor(id, {
+        tools: TOOLS.map((tool) => ({
+          name: tool.name,
+          description: tool.description,
+          inputSchema: tool.inputSchema,
+        })),
+      });
+    case "tools/call": {
+      const toolName = String(params.name ?? "");
+      const tool = TOOLS.find((t) => t.name === toolName);
+      if (!tool) return errorFor(id, -32602, `Unknown tool: ${toolName}`);
+      const args = (params.arguments ?? {}) as Record<string, unknown>;
+      const live = await tryLiveQuery(toolName, args, projectDir, agentHome);
+
+      if (live?.mismatched) {
+        // The server answered for another project. Refuse it: a plausible
+        // wrong answer is worse than an explicit failure.
+        liveUnavailableUntil = Date.now() + LIVE_RETRY_COOLDOWN_MS;
+        return textResult(id, {
+          index: { state: "error", projectDir: projectDir ?? undefined },
+          error: `The live server answered for a different project than "${projectDir}"; refusing the result. Re-run "RA3 Mod XML: Enable AI Agent access…" for this project.`,
+        });
+      }
+
+      if (LIVE_ONLY_TOOLS.has(toolName)) {
+        // Report a clear reason instead of an empty result, so the agent does
+        // not conclude "this asset has no references".
+        return textResult(
+          id,
+          live?.payload ?? {
+            index: { state: "no_index", projectDir: projectDir ?? undefined },
+            error: `"${toolName}" requires a live index. Open the project in VS Code with AI Agent access enabled, then retry.`,
+          },
+        );
+      }
+
+      const output = live?.payload ?? tool.handler(args, snapshot);
+      return textResult(id, output);
+    }
+    default:
+      // Notifications have no id; ignore them.
+      if (id === undefined) return null;
+      return errorFor(id, -32601, `Method not found: ${method}`);
+  }
+}
+
+/** Wraps any tool payload into an MCP text content result. */
+function textResult(id: unknown, payload: unknown): unknown {
+  const text = typeof payload === "string" ? payload : JSON.stringify(payload, null, 2);
+  return resultFor(id, { content: [{ type: "text", text }] });
+}
+
+async function main(): Promise<void> {
+  const args = process.argv.slice(2);
+  let projectDir: string | null = null;
+  let snapshotPath: string | null = null;
+  let agentHome: string | undefined;
+  for (let i = 0; i < args.length; i++) {
+    if (args[i] === "--project" || args[i] === "-p") projectDir = args[++i] ?? null;
+    else if (args[i] === "--snapshot" || args[i] === "-s") snapshotPath = args[++i] ?? null;
+    else if (args[i] === "--agent-home") agentHome = args[++i] ?? undefined;
+  }
+  const resolvedSnapshotPath =
+    snapshotPath ??
+    (projectDir ? snapshotPathForProject(projectDir, agentHome) : null);
+  let snapshot: AgentIndexSnapshot | null = null;
+  if (resolvedSnapshotPath) snapshot = await readSnapshotFile(resolvedSnapshotPath);
+
+  const rl = createInterface({
+    input: process.stdin,
+    crlfDelay: Infinity,
+  });
+  rl.on("line", (line) => {
+    if (!line.trim()) return;
+    let message: Record<string, unknown>;
+    try {
+      message = JSON.parse(line) as Record<string, unknown>;
+    } catch {
+      return;
+    }
+    void handleRequest(message, snapshot, projectDir, agentHome).then((response) => {
+      if (response != null) sendMessage(response);
+    });
+  });
+}
+
+// Only run the stdio loop when executed directly, so the module stays
+// importable by tests.
+if (typeof require !== "undefined" && require.main === module) {
+  void main();
+}