diff --git a/packages/ast-grep-mcp/src/cli.ts b/packages/ast-grep-mcp/src/cli.ts index 397ffdb15..2661ca4cb 100644 --- a/packages/ast-grep-mcp/src/cli.ts +++ b/packages/ast-grep-mcp/src/cli.ts @@ -1,14 +1,20 @@ #!/usr/bin/env node import { argv, stderr } from "node:process"; +import { writeMcpLifecycleLog } from "./mcp-lifecycle-log"; import { runMcpStdioServer } from "./mcp"; async function main(): Promise { const [command = "mcp"] = argv.slice(2); if (command === "mcp") { - await runMcpStdioServer(); + await runMcpStdioServer(process.stdin, process.stdout, {}, { + log: writeMcpLifecycleLog, + onIdleTimeout: () => { + process.exit(0); + }, + }); return; } - stderr.write("Usage: ast-grep-mcp [mcp]\n"); + stderr.write("Usage: omo-ast-grep [mcp]\n"); process.exitCode = 2; } diff --git a/packages/ast-grep-mcp/src/mcp.test.ts b/packages/ast-grep-mcp/src/mcp.test.ts index 5bf869fca..892143a52 100644 --- a/packages/ast-grep-mcp/src/mcp.test.ts +++ b/packages/ast-grep-mcp/src/mcp.test.ts @@ -2,7 +2,8 @@ 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 { PassThrough } from "node:stream"; +import { handleAstGrepMcpRequest, runMcpStdioServer } from "./mcp"; import type { RunOptions } from "./runner"; import type { SgResult } from "./types"; @@ -259,4 +260,20 @@ describe("ast-grep MCP", () => { expect(searchTool?.description).toContain("This is NOT regex"); expect(searchTool?.description).toContain("Meta-variables"); }); + + it("#given idle stdio connection #when no request arrives before timeout #then server exits through idle callback", async () => { + const input = new PassThrough(); + const output = new PassThrough(); + let idleCallCount = 0; + + await runMcpStdioServer(input, output, {}, { + idleTimeoutMs: 1, + onIdleTimeout: () => { + idleCallCount++; + input.end(); + }, + }); + + expect(idleCallCount).toBe(1); + }); }); diff --git a/packages/ast-grep-mcp/src/mcp.ts b/packages/ast-grep-mcp/src/mcp.ts index 20bb86e87..1d3477375 100644 --- a/packages/ast-grep-mcp/src/mcp.ts +++ b/packages/ast-grep-mcp/src/mcp.ts @@ -1,5 +1,6 @@ -import { createInterface } from "node:readline"; +import type { Readable, Writable } from "node:stream"; import { CLI_LANGUAGES } from "./constants"; +import { runJsonRpcStdioServer, type McpStdioServerOptions } from "./mcp-stdio-server"; import { getPatternHint } from "./pattern-hints"; import { formatReplaceResult, formatSearchResult } from "./result-formatter"; import { runSg, type RunOptions } from "./runner"; @@ -116,23 +117,12 @@ export async function handleAstGrepMcpRequest(input: unknown, options: AstGrepMc } export async function runMcpStdioServer( - input: NodeJS.ReadableStream = process.stdin, - output: NodeJS.WritableStream = process.stdout, + input: Readable = process.stdin, + output: Writable = process.stdout, options: AstGrepMcpOptions = {}, + stdioOptions: McpStdioServerOptions = {}, ): Promise { - 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`); - } + await runJsonRpcStdioServer(handleAstGrepMcpRequest, input, output, options, stdioOptions); } async function handleToolCall(id: JsonRpcId, params: unknown, options: AstGrepMcpOptions): Promise {