diff --git a/packages/ast-grep-mcp/src/cli.ts b/packages/ast-grep-mcp/src/cli.ts index 397ffdb15..ca315c6d3 100644 --- a/packages/ast-grep-mcp/src/cli.ts +++ b/packages/ast-grep-mcp/src/cli.ts @@ -1,11 +1,17 @@ #!/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"); diff --git a/packages/ast-grep-mcp/src/mcp-lifecycle-log.ts b/packages/ast-grep-mcp/src/mcp-lifecycle-log.ts new file mode 100644 index 000000000..19fca8d88 --- /dev/null +++ b/packages/ast-grep-mcp/src/mcp-lifecycle-log.ts @@ -0,0 +1,33 @@ +import { appendFileSync, renameSync, statSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; + +type LogFieldValue = boolean | number | string | null; + +const LOG_FILE_NAME = "omo-ast-grep-mcp.log"; +const MAX_LOG_BYTES = 5 * 1024 * 1024; + +export function mcpLifecycleLogPath(): string { + return join(tmpdir(), LOG_FILE_NAME); +} + +export function writeMcpLifecycleLog(event: string, fields: Record = {}): void { + const path = mcpLifecycleLogPath(); + try { + rotateLogIfNeeded(path); + appendFileSync(path, `${JSON.stringify({ ts: new Date().toISOString(), event, pid: process.pid, ppid: process.ppid, ...fields })}\n`); + } catch (error) { + if (error instanceof Error) return; + return; + } +} + +function rotateLogIfNeeded(path: string): void { + try { + if (statSync(path).size < MAX_LOG_BYTES) return; + renameSync(path, `${path}.1`); + } catch (error) { + if (error instanceof Error) return; + return; + } +} diff --git a/packages/ast-grep-mcp/src/mcp-stdio-server.ts b/packages/ast-grep-mcp/src/mcp-stdio-server.ts new file mode 100644 index 000000000..ab49ba60b --- /dev/null +++ b/packages/ast-grep-mcp/src/mcp-stdio-server.ts @@ -0,0 +1,106 @@ +import { createInterface } from "node:readline"; + +import type { AstGrepMcpOptions, JsonRpcResponse } from "./mcp"; + +export type McpLifecycleLog = (event: string, fields?: Record) => void; + +export interface McpStdioServerOptions { + readonly idleTimeoutMs?: number; + readonly onIdleTimeout?: () => void | Promise; + readonly log?: McpLifecycleLog; +} + +export type McpRequestHandler = ( + input: unknown, + options: AstGrepMcpOptions, +) => Promise; + +const DEFAULT_IDLE_TIMEOUT_MS = 10 * 60_000; +const noopLog: McpLifecycleLog = () => {}; + +export async function runJsonRpcStdioServer( + handler: McpRequestHandler, + input: NodeJS.ReadableStream, + output: NodeJS.WritableStream, + options: AstGrepMcpOptions, + stdioOptions: McpStdioServerOptions = {}, +): Promise { + const log = stdioOptions.log ?? noopLog; + const idleTimeoutMs = stdioOptions.idleTimeoutMs ?? DEFAULT_IDLE_TIMEOUT_MS; + const idleTimer = createIdleTimer(idleTimeoutMs, log, stdioOptions.onIdleTimeout); + + log("stdio_started", { cwd: process.cwd(), idle_timeout_ms: idleTimeoutMs }); + idleTimer.arm(); + const lines = createInterface({ input, crlfDelay: Number.POSITIVE_INFINITY }); + try { + for await (const line of lines) { + if (idleTimer.closed()) break; + idleTimer.arm(); + if (!line.trim()) continue; + const parsed = parseJsonRpcLine(line, output, log); + if (parsed === undefined) continue; + const id = isRecord(parsed) ? jsonRpcId(parsed.id) : null; + const method = isRecord(parsed) && typeof parsed.method === "string" ? parsed.method : null; + log("request", { id: id === null ? null : String(id), method }); + const response = await handler(parsed, options); + if (response) { + output.write(`${JSON.stringify(response)}\n`); + log("response", { id: String(response.id), method, is_error: response.error !== undefined }); + } + } + } finally { + idleTimer.clear(); + log("stdio_stopped"); + } +} + +function createIdleTimer(idleTimeoutMs: number, log: McpLifecycleLog, onIdleTimeout?: () => void | Promise) { + let timer: NodeJS.Timeout | null = null; + let isClosed = false; + + return { + arm: () => { + if (timer !== null) clearTimeout(timer); + if (idleTimeoutMs <= 0) return; + timer = setTimeout(() => { + isClosed = true; + log("idle_timeout", { idle_timeout_ms: idleTimeoutMs }); + void onIdleTimeout?.(); + }, idleTimeoutMs); + timer.unref(); + }, + clear: () => { + if (timer === null) return; + clearTimeout(timer); + timer = null; + }, + closed: () => isClosed, + }; +} + +function parseJsonRpcLine(line: string, output: NodeJS.WritableStream, log: McpLifecycleLog): unknown | undefined { + try { + return JSON.parse(line); + } catch (error) { + const message = messageFromError(error); + log("parse_error", { message }); + output.write(`${JSON.stringify(errorResponse(null, -32700, "Parse error", message))}\n`); + return undefined; + } +} + +function errorResponse(id: string | number | null, code: number, message: string, data?: unknown): JsonRpcResponse { + return { jsonrpc: "2.0", id, error: data === undefined ? { code, message } : { code, message, data } }; +} + +function jsonRpcId(value: unknown): string | number | null { + return typeof value === "string" || typeof value === "number" || value === null ? value : null; +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function messageFromError(error: unknown): string { + return error instanceof Error ? error.message : String(error); +} 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..25260f3c3 100644 --- a/packages/ast-grep-mcp/src/mcp.ts +++ b/packages/ast-grep-mcp/src/mcp.ts @@ -1,5 +1,5 @@ -import { createInterface } from "node:readline"; 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"; @@ -119,20 +119,9 @@ export async function runMcpStdioServer( input: NodeJS.ReadableStream = process.stdin, output: NodeJS.WritableStream = 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 { diff --git a/packages/lsp-tools-mcp b/packages/lsp-tools-mcp index e7c65b04d..e3cae69bd 160000 --- a/packages/lsp-tools-mcp +++ b/packages/lsp-tools-mcp @@ -1 +1 @@ -Subproject commit e7c65b04d0cc549f0478d3b78b51714fc0f572b3 +Subproject commit e3cae69bd68357c85f65339194f8028934ee879c