From e48b68e4404331ac8a33dc963a1fb341fd0452b6 Mon Sep 17 00:00:00 2001 From: YeonGyu-Kim Date: Sat, 30 May 2026 19:09:10 +0900 Subject: [PATCH] feat(packages): add mcp-stdio-transport.ts --- .../ast-grep-mcp/src/mcp-stdio-transport.ts | 138 ++++++++++++++++++ 1 file changed, 138 insertions(+) create mode 100644 packages/ast-grep-mcp/src/mcp-stdio-transport.ts diff --git a/packages/ast-grep-mcp/src/mcp-stdio-transport.ts b/packages/ast-grep-mcp/src/mcp-stdio-transport.ts new file mode 100644 index 000000000..bd2f5e905 --- /dev/null +++ b/packages/ast-grep-mcp/src/mcp-stdio-transport.ts @@ -0,0 +1,138 @@ +import type { Readable, Writable } from "node:stream"; + +export type StdioJsonRpcResponseMode = "line" | "framed"; + +export type StdioJsonRpcMessage = + | { + readonly kind: "request"; + readonly payload: unknown; + readonly responseMode: StdioJsonRpcResponseMode; + } + | { + readonly kind: "parse_error"; + readonly message: string; + readonly responseMode: StdioJsonRpcResponseMode; + }; + +type ReadResult = + | { readonly kind: "incomplete" } + | { + readonly kind: "complete"; + readonly message?: StdioJsonRpcMessage; + readonly remaining: Buffer; + }; + +const HEADER_SEPARATOR = Buffer.from("\r\n\r\n"); + +export async function* readStdioJsonRpcMessages(input: Readable): AsyncGenerator { + let buffer: Buffer = Buffer.alloc(0); + + for await (const chunk of input) { + buffer = Buffer.concat([buffer, bufferFromChunk(chunk)]); + while (true) { + const result = readNextMessage(buffer); + if (result.kind === "incomplete") break; + buffer = result.remaining; + if (result.message) yield result.message; + } + } + + const trailing = buffer.toString("utf8").trim(); + if (trailing.length > 0) { + yield parseJsonPayload(trailing, "line"); + } +} + +export function writeStdioJsonRpcResponse(output: Writable, response: unknown, responseMode: StdioJsonRpcResponseMode): void { + const body = JSON.stringify(response); + if (responseMode === "framed") { + output.write(`Content-Length: ${Buffer.byteLength(body, "utf8")}\r\n\r\n${body}`); + return; + } + output.write(`${body}\n`); +} + +function readNextMessage(buffer: Buffer): ReadResult { + if (buffer.length === 0) return { kind: "incomplete" }; + return startsWithContentLength(buffer) ? readFramedMessage(buffer) : readLineMessage(buffer); +} + +function readLineMessage(buffer: Buffer): ReadResult { + const newlineIndex = buffer.indexOf(0x0a); + if (newlineIndex === -1) return { kind: "incomplete" }; + const line = buffer.subarray(0, newlineIndex).toString("utf8").replace(/\r$/, ""); + if (line.trim().length === 0) { + return { + kind: "complete", + remaining: buffer.subarray(newlineIndex + 1), + }; + } + return { + kind: "complete", + message: parseJsonPayload(line, "line"), + remaining: buffer.subarray(newlineIndex + 1), + }; +} + +function readFramedMessage(buffer: Buffer): ReadResult { + const separatorIndex = buffer.indexOf(HEADER_SEPARATOR); + if (separatorIndex === -1) return { kind: "incomplete" }; + + const headers = buffer.subarray(0, separatorIndex).toString("ascii"); + const contentLength = parseContentLength(headers); + const bodyStart = separatorIndex + HEADER_SEPARATOR.length; + if (contentLength === undefined) { + return { + kind: "complete", + message: { + kind: "parse_error", + message: "Missing or invalid Content-Length header", + responseMode: "framed", + }, + remaining: buffer.subarray(bodyStart), + }; + } + + const bodyEnd = bodyStart + contentLength; + if (buffer.length < bodyEnd) return { kind: "incomplete" }; + const body = buffer.subarray(bodyStart, bodyEnd).toString("utf8"); + return { + kind: "complete", + message: parseJsonPayload(body, "framed"), + remaining: buffer.subarray(bodyEnd), + }; +} + +function startsWithContentLength(buffer: Buffer): boolean { + const prefix = buffer.subarray(0, "content-length:".length).toString("ascii").toLowerCase(); + return prefix === "content-length:"; +} + +function parseContentLength(headers: string): number | undefined { + for (const line of headers.split("\r\n")) { + const match = /^content-length:\s*(\d+)$/i.exec(line); + if (match === null) continue; + const value = match[1]; + if (value === undefined) return undefined; + return Number(value); + } + return undefined; +} + +function parseJsonPayload(payload: string, responseMode: StdioJsonRpcResponseMode): StdioJsonRpcMessage { + try { + return { kind: "request", payload: JSON.parse(payload), responseMode }; + } catch (error) { + return { kind: "parse_error", message: messageFromError(error), responseMode }; + } +} + +function bufferFromChunk(chunk: unknown): Buffer { + if (Buffer.isBuffer(chunk)) return chunk; + if (typeof chunk === "string") return Buffer.from(chunk); + throw new TypeError(`Unsupported stdio chunk type: ${typeof chunk}`); +} + +function messageFromError(error: unknown): string { + return error instanceof Error ? error.message : String(error); +}