import { afterEach, describe, expect, it } from "bun:test"; import { mkdirSync, mkdtempSync, realpathSync, rmSync, symlinkSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { handleAstGrepMcpRequest } from "./mcp"; import type { RunOptions } from "./runner"; import type { SgResult } from "./types"; const emptyResult: SgResult = { matches: [], totalMatches: 0, truncated: false, }; const temporaryDirectories: string[] = []; function createTemporaryDirectory(prefix: string): string { const directory = mkdtempSync(join(tmpdir(), prefix)); temporaryDirectories.push(directory); return directory; } afterEach(() => { for (const directory of temporaryDirectories.splice(0)) { rmSync(directory, { recursive: true, force: true }); } }); describe("ast-grep MCP", () => { it("#given initialize request #when handled #then advertises tools capability", async () => { const response = await handleAstGrepMcpRequest({ jsonrpc: "2.0", id: 1, method: "initialize", params: { protocolVersion: "2024-11-05" }, }); expect(response).toEqual({ jsonrpc: "2.0", id: 1, result: { capabilities: { tools: { listChanged: false } }, serverInfo: { name: "ast_grep", version: "0.1.0" }, protocolVersion: "2024-11-05", }, }); }); it("#given tools list request #when handled #then exposes search and replace tools", async () => { const response = await handleAstGrepMcpRequest({ jsonrpc: "2.0", id: "tools", method: "tools/list" }); expect(response?.result?.tools?.map((tool) => tool.name)).toEqual(["search", "replace"]); }); it("#given search call without paths #when handled #then defaults paths to workspace directory", async () => { const captured: { value?: RunOptions } = {}; const workspaceDirectory = createTemporaryDirectory("omo-ast-grep-workspace-"); const response = await handleAstGrepMcpRequest( { jsonrpc: "2.0", id: "search", method: "tools/call", params: { name: "search", arguments: { pattern: "console.log($$$)", lang: "typescript" } }, }, { workspaceDirectory, runSg: async (options) => { captured.value = options; return emptyResult; }, }, ); expect(captured.value).toEqual({ pattern: "console.log($$$)", lang: "typescript", cwd: realpathSync(workspaceDirectory), paths: ["."], globs: undefined, context: undefined }); expect(response?.result?.content?.[0]?.text).toContain("No matches found"); }); it("#given replace call without dryRun #when handled #then keeps dry-run default", async () => { const captured: { value?: RunOptions } = {}; const workspaceDirectory = createTemporaryDirectory("omo-ast-grep-replace-workspace-"); mkdirSync(join(workspaceDirectory, "src")); await handleAstGrepMcpRequest( { jsonrpc: "2.0", id: "replace", method: "tools/call", params: { name: "replace", arguments: { pattern: "console.log($MSG)", rewrite: "logger.info($MSG)", lang: "typescript", paths: ["src"] }, }, }, { workspaceDirectory, runSg: async (options) => { captured.value = options; return emptyResult; }, }, ); expect(captured.value).toEqual({ pattern: "console.log($MSG)", rewrite: "logger.info($MSG)", lang: "typescript", cwd: realpathSync(workspaceDirectory), paths: ["src"], globs: undefined, updateAll: false, }); }); it("#given disabled replace tool #when listed and called #then hides and rejects it", async () => { const listResponse = await handleAstGrepMcpRequest({ jsonrpc: "2.0", id: "tools", method: "tools/list" }, { disabledTools: ["replace"] }); expect(listResponse?.result?.tools?.map((tool) => tool.name)).toEqual(["search"]); const callResponse = await handleAstGrepMcpRequest( { jsonrpc: "2.0", id: "replace", method: "tools/call", params: { name: "replace", arguments: { pattern: "console.log($MSG)", rewrite: "logger.info($MSG)", lang: "typescript", paths: ["src"] }, }, }, { disabledTools: ["replace"] }, ); expect(callResponse?.result?.isError).toBe(true); expect(callResponse?.result?.content?.[0]?.text).toContain("ast-grep tool is disabled: replace"); }); it("#given unsafe paths #when search is called #then rejects before running ast-grep", async () => { const workspaceDirectory = createTemporaryDirectory("omo-ast-grep-sandbox-"); const outsideDirectory = createTemporaryDirectory("omo-ast-grep-outside-"); symlinkSync(outsideDirectory, join(workspaceDirectory, "outside-link")); let didRun = false; for (const path of ["../outside", "/tmp", "--update-all", "outside-link"]) { const response = await handleAstGrepMcpRequest( { jsonrpc: "2.0", id: path, method: "tools/call", params: { name: "search", arguments: { pattern: "console.log($$$)", lang: "typescript", paths: [path] } }, }, { workspaceDirectory, runSg: async () => { didRun = true; return emptyResult; }, }, ); expect(response?.result?.isError).toBe(true); } expect(didRun).toBe(false); }); it("#given paths contains an absolute path inside the workspace #when search and replace are called #then it is normalized to relative and processed normally", async () => { const captured: RunOptions[] = []; const workspaceDirectory = createTemporaryDirectory("omo-ast-grep-absolute-inside-"); const sourcePath = join(workspaceDirectory, "foo.ts"); writeFileSync(sourcePath, "console.log('hello')\n"); const runSg = async (options: RunOptions): Promise => { captured.push(options); return emptyResult; }; const searchResponse = await handleAstGrepMcpRequest( { jsonrpc: "2.0", id: "search-absolute-inside", method: "tools/call", params: { name: "search", arguments: { pattern: "console.log($$$)", lang: "typescript", paths: [sourcePath] } }, }, { workspaceDirectory, runSg }, ); const replaceResponse = await handleAstGrepMcpRequest( { jsonrpc: "2.0", id: "replace-absolute-inside", method: "tools/call", params: { name: "replace", arguments: { pattern: "console.log($MSG)", rewrite: "logger.info($MSG)", lang: "typescript", paths: [sourcePath] } }, }, { workspaceDirectory, runSg }, ); expect(searchResponse?.result?.isError).toBe(false); expect(replaceResponse?.result?.isError).toBe(false); expect(captured.map((options) => options.paths)).toEqual([["foo.ts"], ["foo.ts"]]); }); it("#given paths contains an absolute path outside the workspace #when search is called #then it is rejected", async () => { const workspaceDirectory = createTemporaryDirectory("omo-ast-grep-absolute-outside-workspace-"); const outsideDirectory = createTemporaryDirectory("omo-ast-grep-absolute-outside-"); const outsidePath = join(outsideDirectory, "foo.ts"); writeFileSync(outsidePath, "console.log('outside')\n"); let didRun = false; const response = await handleAstGrepMcpRequest( { jsonrpc: "2.0", id: "absolute-outside", method: "tools/call", params: { name: "search", arguments: { pattern: "console.log($$$)", lang: "typescript", paths: [outsidePath] } }, }, { workspaceDirectory, runSg: async () => { didRun = true; return emptyResult; }, }, ); expect(response?.result?.isError).toBe(true); expect(response?.result?.content?.[0]?.text).toContain("stay inside the workspace"); expect(didRun).toBe(false); }); it("#given paths contains an absolute path whose realpath resolves inside workspace through a symlink #when search is called #then it is allowed", async () => { const captured: { value?: RunOptions } = {}; const workspaceDirectory = createTemporaryDirectory("omo-ast-grep-symlink-target-"); const symlinkParent = createTemporaryDirectory("omo-ast-grep-symlink-parent-"); const sourcePath = join(workspaceDirectory, "foo.ts"); const symlinkPath = join(symlinkParent, "workspace-link"); writeFileSync(sourcePath, "console.log('via symlink')\n"); symlinkSync(workspaceDirectory, symlinkPath); const response = await handleAstGrepMcpRequest( { jsonrpc: "2.0", id: "absolute-symlink-inside", method: "tools/call", params: { name: "search", arguments: { pattern: "console.log($$$)", lang: "typescript", paths: [join(symlinkPath, "foo.ts")] } }, }, { workspaceDirectory, runSg: async (options) => { captured.value = options; return emptyResult; }, }, ); expect(response?.result?.isError).toBe(false); expect(captured.value?.paths).toEqual(["foo.ts"]); }); it("#given tools list request #when handled #then preserves detailed ast-grep guidance", async () => { const response = await handleAstGrepMcpRequest({ jsonrpc: "2.0", id: "tools", method: "tools/list" }); const searchTool = response?.result?.tools?.find((tool) => tool.name === "search"); expect(searchTool?.description).toContain("This is NOT regex"); expect(searchTool?.description).toContain("Meta-variables"); }); });