Ra3ModXmlExt/src/agent/mcpServer.ts

/**
 * 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 { LiveClient } from "./liveClient";
import { pruneInstances } from "./instances";
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
- list_projects() -> project roots the live extension has indexed
- 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: "list_projects",
    description:
      "Lists the project roots the live extension currently has indexed. Use it to discover which projects this server can answer for.",
    inputSchema: { type: "object", properties: {} },
    handler: () => ({
      index: { state: "no_index" },
      error:
        "list_projects requires a live index. Open the project in VS Code with AI Agent access enabled, then 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", "list_projects"]);

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,
  live: LiveClient,
): 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 result = await live.query(toolName, args);

      if (result?.mismatched) {
        // The server answered for another project. Refuse it: a plausible
        // wrong answer is worse than an explicit failure.
        live.markUnavailable();
        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,
          result?.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 = result?.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);

  // One-shot crash cleanup: a window that died without disposing leaves its
  // instance file behind, and whichever instance starts next prunes it.
  void pruneInstances(agentHome).catch(() => undefined);

  const live = new LiveClient({ projectDir, agentHome });

  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, live).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();
}