feat(mcp): add package-backed ast-grep MCP

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
2026-05-18 21:19:15 +09:00
parent 4ea29e2c94
commit 499aff011a
19 changed files with 1372 additions and 0 deletions
@@ -0,0 +1,275 @@
+import { createInterface } from "node:readline";
+import { CLI_LANGUAGES } from "./constants";
+import { getPatternHint } from "./pattern-hints";
+import { formatReplaceResult, formatSearchResult } from "./result-formatter";
+import { runSg, type RunOptions } from "./runner";
+import { AST_GREP_REPLACE_DESCRIPTION, AST_GREP_SEARCH_DESCRIPTION, AST_GREP_SEARCH_PATTERN_PARAM } from "./tool-descriptions";
+import type { CliLanguage, SgResult } from "./types";
+import { normalizeWorkspaceDirectory, resolveWorkspacePaths } from "./workspace-paths";
+
+export type JsonRpcId = string | number | null;
+
+export interface TextContent {
+  readonly type: "text";
+  readonly text: string;
+}
+
+export interface McpToolDescriptor {
+  readonly name: string;
+  readonly title: string;
+  readonly description: string;
+  readonly inputSchema: unknown;
+}
+
+export interface JsonRpcError {
+  readonly code: number;
+  readonly message: string;
+  readonly data?: unknown;
+}
+
+export interface JsonRpcResult {
+  readonly capabilities?: Record<string, unknown>;
+  readonly serverInfo?: Record<string, unknown>;
+  readonly protocolVersion?: string;
+  readonly tools?: readonly McpToolDescriptor[];
+  readonly content?: readonly TextContent[];
+  readonly isError?: boolean;
+  readonly [key: string]: unknown;
+}
+
+export interface JsonRpcResponse {
+  readonly jsonrpc: "2.0";
+  readonly id: JsonRpcId;
+  readonly result?: JsonRpcResult;
+  readonly error?: JsonRpcError;
+}
+
+export interface AstGrepMcpOptions {
+  readonly workspaceDirectory?: string;
+  readonly disabledTools?: readonly string[];
+  readonly runSg?: (options: RunOptions) => Promise<SgResult>;
+}
+
+type ToolCallResult = {
+  readonly content: readonly TextContent[];
+  readonly isError?: boolean;
+};
+
+const SERVER_NAME = "ast_grep";
+const SERVER_VERSION = "0.1.0";
+const LANGUAGE_VALUES: readonly string[] = CLI_LANGUAGES;
+const DISABLED_TOOLS_ENV = "OMO_AST_GREP_DISABLED_TOOLS";
+
+const AST_GREP_MCP_TOOLS = [
+  {
+    name: "search",
+    title: "AST grep search",
+    description: AST_GREP_SEARCH_DESCRIPTION,
+    inputSchema: {
+      type: "object",
+      properties: {
+        pattern: { type: "string", description: AST_GREP_SEARCH_PATTERN_PARAM },
+        lang: { type: "string", enum: CLI_LANGUAGES, description: "Target language" },
+        paths: { type: "array", items: { type: "string" }, description: "Paths to search" },
+        globs: { type: "array", items: { type: "string" }, description: "Include/exclude globs" },
+        context: { type: "number", description: "Context lines around each match" },
+      },
+      required: ["pattern", "lang"],
+      additionalProperties: false,
+    },
+  },
+  {
+    name: "replace",
+    title: "AST grep replace",
+    description: AST_GREP_REPLACE_DESCRIPTION,
+    inputSchema: {
+      type: "object",
+      properties: {
+        pattern: { type: "string", description: "AST pattern to match" },
+        rewrite: { type: "string", description: "Replacement pattern" },
+        lang: { type: "string", enum: CLI_LANGUAGES, description: "Target language" },
+        paths: { type: "array", items: { type: "string" }, description: "Paths to search" },
+        globs: { type: "array", items: { type: "string" }, description: "Include/exclude globs" },
+        dryRun: { type: "boolean", description: "Preview changes without applying. Defaults to true." },
+      },
+      required: ["pattern", "rewrite", "lang"],
+      additionalProperties: false,
+    },
+  },
+] as const satisfies readonly McpToolDescriptor[];
+
+export async function handleAstGrepMcpRequest(input: unknown, options: AstGrepMcpOptions = {}): Promise<JsonRpcResponse | undefined> {
+  if (!isRecord(input)) return errorResponse(null, -32600, "Invalid Request");
+  const id = jsonRpcId(input.id);
+  if (input.method === "notifications/initialized") return undefined;
+  if (input.method === "ping") return successResponse(id, {});
+  if (input.method === "initialize") {
+    return successResponse(id, {
+      capabilities: { tools: { listChanged: false } },
+      serverInfo: { name: SERVER_NAME, version: SERVER_VERSION },
+      protocolVersion: requestedProtocolVersion(input.params),
+    });
+  }
+  if (input.method === "tools/list") return successResponse(id, { tools: enabledTools(options) });
+  if (input.method === "tools/call") return handleToolCall(id, input.params, options);
+  return errorResponse(id, -32601, `Method not found: ${String(input.method)}`);
+}
+
+export async function runMcpStdioServer(
+  input: NodeJS.ReadableStream = process.stdin,
+  output: NodeJS.WritableStream = process.stdout,
+  options: AstGrepMcpOptions = {},
+): Promise<void> {
+  const lines = createInterface({ input, crlfDelay: Number.POSITIVE_INFINITY });
+  for await (const line of lines) {
+    if (!line.trim()) continue;
+    let parsed: unknown;
+    try {
+      parsed = JSON.parse(line);
+    } catch (error) {
+      output.write(`${JSON.stringify(errorResponse(null, -32700, "Parse error", messageFromError(error)))}\n`);
+      continue;
+    }
+    const response = await handleAstGrepMcpRequest(parsed, options);
+    if (response) output.write(`${JSON.stringify(response)}\n`);
+  }
+}
+
+async function handleToolCall(id: JsonRpcId, params: unknown, options: AstGrepMcpOptions): Promise<JsonRpcResponse> {
+  if (!isRecord(params) || typeof params.name !== "string") return errorResponse(id, -32602, "tools/call requires params.name");
+  try {
+    const result = await executeAstGrepTool(params.name, params.arguments, options);
+    return successResponse(id, { content: result.content, isError: result.isError ?? false });
+  } catch (error) {
+    return successResponse(id, { content: [{ type: "text", text: messageFromError(error) }], isError: true });
+  }
+}
+
+async function executeAstGrepTool(name: string, args: unknown, options: AstGrepMcpOptions): Promise<ToolCallResult> {
+  if (disabledToolNames(options).has(name)) throw new Error(`ast-grep tool is disabled: ${name}`);
+  const runner = options.runSg ?? runSg;
+  const workspaceDirectory = normalizeWorkspaceDirectory(options.workspaceDirectory ?? process.env.OMO_AST_GREP_WORKSPACE ?? process.cwd());
+  if (name === "search") {
+    const input = parseSearchArgs(args, workspaceDirectory);
+    const result = await runner(input);
+    let output = formatSearchResult(result);
+    if (result.matches.length === 0 && !result.error) {
+      const hint = getPatternHint(input.pattern, input.lang);
+      if (hint) output += `\n\n${hint}`;
+    }
+    return { content: [{ type: "text", text: output }], isError: Boolean(result.error) };
+  }
+  if (name === "replace") {
+    const input = parseReplaceArgs(args, workspaceDirectory);
+    const result = await runner(input.options);
+    return { content: [{ type: "text", text: formatReplaceResult(result, input.dryRun) }], isError: Boolean(result.error) };
+  }
+  throw new Error(`Unknown ast-grep tool: ${name}`);
+}
+
+function parseSearchArgs(args: unknown, workspaceDirectory: string): RunOptions {
+  const input = requireRecord(args);
+  return {
+    pattern: requireString(input, "pattern"),
+    lang: requireLanguage(input, "lang"),
+    cwd: workspaceDirectory,
+    paths: resolveWorkspacePaths(optionalStringArray(input, "paths"), workspaceDirectory),
+    globs: optionalStringArray(input, "globs"),
+    context: optionalNumber(input, "context"),
+  };
+}
+
+function parseReplaceArgs(args: unknown, workspaceDirectory: string): { readonly options: RunOptions; readonly dryRun: boolean } {
+  const input = requireRecord(args);
+  const dryRun = optionalBoolean(input, "dryRun") ?? true;
+  return {
+    dryRun,
+    options: {
+      pattern: requireString(input, "pattern"),
+      rewrite: requireString(input, "rewrite"),
+      lang: requireLanguage(input, "lang"),
+      cwd: workspaceDirectory,
+      paths: resolveWorkspacePaths(optionalStringArray(input, "paths"), workspaceDirectory),
+      globs: optionalStringArray(input, "globs"),
+      updateAll: !dryRun,
+    },
+  };
+}
+
+function requireRecord(value: unknown): Record<string, unknown> {
+  if (!isRecord(value)) throw new Error("Tool arguments must be an object");
+  return value;
+}
+
+function requireString(input: Record<string, unknown>, key: string): string {
+  const value = input[key];
+  if (typeof value !== "string" || value.length === 0) throw new Error(`${key} must be a non-empty string`);
+  return value;
+}
+
+function requireLanguage(input: Record<string, unknown>, key: string): CliLanguage {
+  const value = requireString(input, key);
+  if (!isCliLanguage(value)) throw new Error(`${key} must be one of: ${LANGUAGE_VALUES.join(", ")}`);
+  return value;
+}
+
+function isCliLanguage(value: string): value is CliLanguage {
+  return LANGUAGE_VALUES.includes(value);
+}
+
+function optionalStringArray(input: Record<string, unknown>, key: string): string[] | undefined {
+  const value = input[key];
+  if (value === undefined) return undefined;
+  if (!Array.isArray(value) || !value.every((item) => typeof item === "string")) throw new Error(`${key} must be an array of strings`);
+  return value;
+}
+
+function enabledTools(options: AstGrepMcpOptions): readonly McpToolDescriptor[] {
+  const disabled = disabledToolNames(options);
+  return AST_GREP_MCP_TOOLS.filter((tool) => !disabled.has(tool.name));
+}
+
+function disabledToolNames(options: AstGrepMcpOptions): ReadonlySet<string> {
+  const fromOptions = options.disabledTools ?? [];
+  const fromEnv = process.env[DISABLED_TOOLS_ENV]?.split(",") ?? [];
+  return new Set([...fromOptions, ...fromEnv].map((tool) => tool.trim()).filter(Boolean));
+}
+
+function optionalNumber(input: Record<string, unknown>, key: string): number | undefined {
+  const value = input[key];
+  if (value === undefined) return undefined;
+  if (typeof value !== "number") throw new Error(`${key} must be a number`);
+  return value;
+}
+
+function optionalBoolean(input: Record<string, unknown>, key: string): boolean | undefined {
+  const value = input[key];
+  if (value === undefined) return undefined;
+  if (typeof value !== "boolean") throw new Error(`${key} must be a boolean`);
+  return value;
+}
+
+function successResponse(id: JsonRpcId, result: JsonRpcResult): JsonRpcResponse {
+  return { jsonrpc: "2.0", id, result };
+}
+
+function errorResponse(id: JsonRpcId, code: number, message: string, data?: unknown): JsonRpcResponse {
+  return { jsonrpc: "2.0", id, error: data === undefined ? { code, message } : { code, message, data } };
+}
+
+function requestedProtocolVersion(params: unknown): string {
+  if (!isRecord(params) || typeof params.protocolVersion !== "string") return "2024-11-05";
+  return params.protocolVersion;
+}
+
+function jsonRpcId(value: unknown): JsonRpcId {
+  return typeof value === "string" || typeof value === "number" || value === null ? value : null;
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
+function messageFromError(error: unknown): string {
+  return error instanceof Error ? error.message : String(error);
+}