feat(packages): add mcp-stdio-transport.ts

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<ArrayBufferLike>;
+    };
+
+const HEADER_SEPARATOR = Buffer.from("\r\n\r\n");
+
+export async function* readStdioJsonRpcMessages(input: Readable): AsyncGenerator<StdioJsonRpcMessage> {
+  let buffer: Buffer<ArrayBufferLike> = 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<ArrayBufferLike>): ReadResult {
+  if (buffer.length === 0) return { kind: "incomplete" };
+  return startsWithContentLength(buffer) ? readFramedMessage(buffer) : readLineMessage(buffer);
+}
+
+function readLineMessage(buffer: Buffer<ArrayBufferLike>): 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<ArrayBufferLike>): 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<ArrayBufferLike>): 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<ArrayBufferLike> {
+  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);
+}