YeonGyu-Kim
95 lines
3.3 KiB
TypeScript
95 lines
3.3 KiB
TypeScript
import type { Readable, Writable } from "node:stream";
|
|
|
|
import type { AstGrepMcpOptions, JsonRpcResponse } from "./mcp";
|
|
import { readStdioJsonRpcMessages, writeStdioJsonRpcResponse } from "./mcp-stdio-transport";
|
|
|
|
export type McpLifecycleLog = (event: string, fields?: Record<string, boolean | number | string | null>) => void;
|
|
|
|
export interface McpStdioServerOptions {
|
|
readonly idleTimeoutMs?: number;
|
|
readonly onIdleTimeout?: () => void | Promise<void>;
|
|
readonly log?: McpLifecycleLog;
|
|
}
|
|
|
|
export type McpRequestHandler = (
|
|
input: unknown,
|
|
options: AstGrepMcpOptions,
|
|
) => Promise<JsonRpcResponse | undefined>;
|
|
|
|
const DEFAULT_IDLE_TIMEOUT_MS = 10 * 60_000;
|
|
const noopLog: McpLifecycleLog = () => {};
|
|
|
|
export async function runJsonRpcStdioServer(
|
|
handler: McpRequestHandler,
|
|
input: Readable,
|
|
output: Writable,
|
|
options: AstGrepMcpOptions,
|
|
stdioOptions: McpStdioServerOptions = {},
|
|
): Promise<void> {
|
|
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();
|
|
try {
|
|
for await (const message of readStdioJsonRpcMessages(input)) {
|
|
if (idleTimer.closed()) break;
|
|
idleTimer.arm();
|
|
if (message.kind === "parse_error") {
|
|
log("parse_error", { message: message.message });
|
|
writeStdioJsonRpcResponse(output, errorResponse(null, -32700, "Parse error", message.message), message.responseMode);
|
|
continue;
|
|
}
|
|
const parsed = message.payload;
|
|
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) {
|
|
writeStdioJsonRpcResponse(output, response, message.responseMode);
|
|
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<void>) {
|
|
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 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<string, unknown> {
|
|
return typeof value === "object" && value !== null && !Array.isArray(value);
|
|
}
|