diff --git a/.changeset/daemon-reuse.md b/.changeset/daemon-reuse.md new file mode 100644 index 0000000..812b295 --- /dev/null +++ b/.changeset/daemon-reuse.md @@ -0,0 +1,16 @@ +--- +"@vndv/pi-codegraph": minor +--- + +Reuse CodeGraph daemon across tool calls instead of spawning a new process per request + +**Before:** Each tool call spawned a new `codegraph serve --mcp` process, causing cold-start delays and timeouts on large projects. + +**After:** A single daemon is cached per project path and reused across calls. The daemon shuts down after 5 minutes of inactivity and is automatically restarted on the next request. + +Benefits: +- **First call:** spawns daemon (cold start) +- **Subsequent calls:** reuses daemon (~instant) +- **No calls for 5min:** daemon shuts down, next call respawns +- Concurrent requests share the same daemon safely (JSON-RPC multiplexing) +- Process exit cleans up all daemons diff --git a/__tests__/codegraph.test.ts b/__tests__/codegraph.test.ts index fd3bd60..33bd214 100644 --- a/__tests__/codegraph.test.ts +++ b/__tests__/codegraph.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, it, vi, afterEach } from "vitest"; +import { describe, expect, it, vi, afterEach, beforeEach } from "vitest"; import { EventEmitter } from "node:events"; import { PassThrough } from "node:stream"; import os from "node:os"; @@ -46,6 +46,12 @@ afterEach(() => { vi.restoreAllMocks(); }); +// Reset daemon manager between tests to avoid state leaking +beforeEach(async () => { + const mod = await import("../extensions/codegraph.js"); + mod.getDaemonManager().killAll(); +}); + describe("pi-codegraph extension", () => { it("exports all CodeGraph tool names", async () => { const mod = await import("../extensions/codegraph.js"); @@ -64,7 +70,10 @@ describe("pi-codegraph extension", () => { it("uses the direct codegraph executable outside Windows", async () => { vi.spyOn(process, "platform", "get").mockReturnValue("linux"); const { spawn } = await import("node:child_process"); - const { withCodeGraphMcp } = await import("../extensions/codegraph.js"); + const { withCodeGraphMcp, getDaemonManager } = await import("../extensions/codegraph.js"); + + // Kill any cached daemon so spawn is called fresh + getDaemonManager().killAll(); await withCodeGraphMcp(process.cwd(), undefined, async () => "success"); @@ -78,7 +87,10 @@ describe("pi-codegraph extension", () => { it("uses PowerShell command discovery for the CodeGraph executable on Windows", async () => { vi.spyOn(process, "platform", "get").mockReturnValue("win32"); const { spawn } = await import("node:child_process"); - const { withCodeGraphMcp } = await import("../extensions/codegraph.js"); + const { withCodeGraphMcp, getDaemonManager } = await import("../extensions/codegraph.js"); + + // Kill any cached daemon so spawn is called fresh + getDaemonManager().killAll(); await withCodeGraphMcp(process.cwd(), undefined, async () => "success"); @@ -212,46 +224,45 @@ describe("pi-codegraph extension", () => { }); }); - it("clears the timeout timer on successful session completion", async () => { + it("reuses daemon across multiple calls", async () => { const { spawn } = await import("node:child_process"); - const { withCodeGraphMcp } = await import("../extensions/codegraph.js"); + const { callCodeGraphTool, getDaemonManager } = await import("../extensions/codegraph.js"); - const setTimeoutSpy = vi.spyOn(globalThis, "setTimeout"); - const clearTimeoutSpy = vi.spyOn(globalThis, "clearTimeout"); + getDaemonManager().killAll(); + const spawnCallsBefore = vi.mocked(spawn).mock.calls.length; - vi.mocked(spawn).mockImplementationOnce(() => { - const child = new EventEmitter() as any; - child.stdin = new PassThrough(); - child.stdout = new PassThrough(); - child.stderr = new PassThrough(); - child.killed = false; - child.kill = vi.fn(() => { child.killed = true; }); + await callCodeGraphTool("codegraph_status", {}); + await callCodeGraphTool("codegraph_status", {}); - child.stdin.on("data", (chunk: Buffer) => { - const lines = chunk.toString("utf-8").trim().split("\n").filter(Boolean); - for (const line of lines) { - const msg = JSON.parse(line); - if (msg.method === "initialize") { - child.stdout.write(JSON.stringify({ jsonrpc: "2.0", id: msg.id, result: {} }) + "\n"); - } - } - }); + // Daemon should only spawn once, not twice + const spawnCallsAfter = vi.mocked(spawn).mock.calls.length; + expect(spawnCallsAfter - spawnCallsBefore).toBe(1); + }); - return child; - }); + it("kills daemon after idle timeout", async () => { + vi.useFakeTimers(); + const { spawn } = await import("node:child_process"); + const { callCodeGraphTool, getDaemonManager, DaemonIdleTimeoutMs } = await import("../extensions/codegraph.js"); - await withCodeGraphMcp(process.cwd(), undefined, async () => "success"); + getDaemonManager().killAll(); + await callCodeGraphTool("codegraph_status", {}); + + const child = vi.mocked(spawn).mock.results[0]?.value; + expect(child).toBeDefined(); + + // Advance past idle timeout + vi.advanceTimersByTime(DaemonIdleTimeoutMs + 1); + + // Daemon should be killed + expect(child.killed).toBe(true); - expect(setTimeoutSpy).toHaveBeenCalled(); - expect(clearTimeoutSpy).toHaveBeenCalled(); + vi.useRealTimers(); }); - it("clears the timeout timer on abort signal", async () => { + it("times out when the MCP request never completes", async () => { const { spawn } = await import("node:child_process"); - const { withCodeGraphMcp } = await import("../extensions/codegraph.js"); - - const clearTimeoutSpy = vi.spyOn(globalThis, "clearTimeout"); + // Create a process that responds to initialize but not tools/call vi.mocked(spawn).mockImplementationOnce(() => { const child = new EventEmitter() as any; child.stdin = new PassThrough(); @@ -267,49 +278,30 @@ describe("pi-codegraph extension", () => { if (msg.method === "initialize") { child.stdout.write(JSON.stringify({ jsonrpc: "2.0", id: msg.id, result: {} }) + "\n"); } + // Don't respond to tools/call - simulate timeout } }); return child; }); - const controller = new AbortController(); - const promise = withCodeGraphMcp(process.cwd(), controller.signal, async (request) => { - const toolPromise = request("tools/call", {}); - controller.abort(); - return toolPromise; - }); - - await expect(promise).rejects.toThrow("CodeGraph MCP process closed before responding."); - expect(clearTimeoutSpy).toHaveBeenCalled(); - }); - - it("times out when the MCP session never completes", async () => { - const { spawn } = await import("node:child_process"); - const { withCodeGraphMcp, SessionTimeoutMs } = await import("../extensions/codegraph.js"); + const { withCodeGraphMcp, getDaemonManager, SessionTimeoutMs } = await import("../extensions/codegraph.js"); + getDaemonManager().killAll(); - const killMock = vi.fn(() => {}); - vi.mocked(spawn).mockImplementationOnce(() => { - const child = new EventEmitter() as any; - child.stdin = new PassThrough(); - child.stdout = new PassThrough(); - child.stderr = new PassThrough(); - child.killed = false; - child.kill = killMock; - return child; + const promise = withCodeGraphMcp(process.cwd(), undefined, async (request) => { + return request("tools/call", { name: "test" }); }); - const promise = withCodeGraphMcp(process.cwd(), undefined, async () => "done"); - - await expect(promise).rejects.toThrow("CodeGraph MCP session timed out after " + SessionTimeoutMs); - expect(killMock).toHaveBeenCalled(); - }, 22000); + await expect(promise).rejects.toThrow("CodeGraph MCP request timed out after " + SessionTimeoutMs); + }, 25000); it("normalizes codegraph_files path before forwarding to the MCP server", async () => { const { spawn } = await import("node:child_process"); - const { callCodeGraphTool } = await import("../extensions/codegraph.js"); + const { callCodeGraphTool, getDaemonManager } = await import("../extensions/codegraph.js"); let capturedArgs: Record | undefined; + getDaemonManager().killAll(); + vi.mocked(spawn).mockImplementationOnce(() => { const child = new EventEmitter() as any; child.stdin = new PassThrough(); diff --git a/extensions/codegraph.ts b/extensions/codegraph.ts index 27b6c60..2ba3b00 100644 --- a/extensions/codegraph.ts +++ b/extensions/codegraph.ts @@ -122,6 +122,7 @@ type PendingJsonRpcRequests = Map tool.name); @@ -164,37 +165,216 @@ function spawnCodeGraphServer(cwd: string): ChildProcessWithoutNullStreams { }); } -export async function withCodeGraphMcp( +// ─── Daemon Manager ────────────────────────────────────────────────────────── + +interface DaemonEntry { + child: ChildProcessWithoutNullStreams; + pending: PendingJsonRpcRequests; + nextId: number; + initialized: boolean; + stdoutBuffer: string; + stderrBuffer: string; + refCount: number; + idleTimer: ReturnType | null; + dead: boolean; +} + +class DaemonManager { + private daemons = new Map(); + private cleanupRegistered = false; + + private ensureCleanup(): void { + if (this.cleanupRegistered) return; + this.cleanupRegistered = true; + + const cleanup = () => this.killAll(); + process.on("exit", cleanup); + process.on("SIGINT", () => { cleanup(); process.exit(130); }); + process.on("SIGTERM", () => { cleanup(); process.exit(143); }); + } + + async acquire(cwd: string): Promise { + this.ensureCleanup(); + + let entry = this.daemons.get(cwd); + + // Reuse existing daemon if alive + if (entry && !entry.dead) { + entry.refCount++; + if (entry.idleTimer) { + clearTimeout(entry.idleTimer); + entry.idleTimer = null; + } + return entry; + } + + // Spawn new daemon + const child = spawnCodeGraphServer(cwd); + entry = { + child, + pending: new Map(), + nextId: 1, + initialized: false, + stdoutBuffer: "", + stderrBuffer: "", + refCount: 1, + idleTimer: null, + dead: false, + }; + + this.attachHandlers(entry, cwd); + this.daemons.set(cwd, entry); + + // Initialize MCP session + await this.initialize(entry, cwd); + + return entry; + } + + release(cwd: string): void { + const entry = this.daemons.get(cwd); + if (!entry || entry.dead) return; + + entry.refCount--; + if (entry.refCount <= 0) { + entry.refCount = 0; + entry.idleTimer = setTimeout(() => this.kill(cwd), DaemonIdleTimeoutMs); + } + } + + kill(cwd: string): void { + const entry = this.daemons.get(cwd); + if (!entry) return; + + entry.dead = true; + if (entry.idleTimer) { + clearTimeout(entry.idleTimer); + entry.idleTimer = null; + } + rejectPendingRequests(entry.pending, new Error("CodeGraph daemon shut down.")); + if (!entry.child.killed) entry.child.kill(); + this.daemons.delete(cwd); + } + + killAll(): void { + for (const cwd of this.daemons.keys()) { + this.kill(cwd); + } + } + + private attachHandlers(entry: DaemonEntry, _cwd: string): void { + entry.child.stdout.on("data", (chunk: Buffer) => { + entry.stdoutBuffer += chunk.toString("utf-8"); + let newline; + while ((newline = entry.stdoutBuffer.indexOf("\n")) !== -1) { + const line = entry.stdoutBuffer.slice(0, newline).trim(); + entry.stdoutBuffer = entry.stdoutBuffer.slice(newline + 1); + if (line) resolveJsonRpcLine(line, entry.pending); + } + }); + + entry.child.stderr.on("data", (chunk: Buffer) => { + entry.stderrBuffer += chunk.toString("utf-8"); + }); + + entry.child.on("error", () => { + entry.dead = true; + rejectPendingRequests(entry.pending, new Error("CodeGraph daemon error.")); + }); + + entry.child.on("exit", (code) => { + entry.dead = true; + if (entry.pending.size > 0) { + const diagnostic = sanitizeDiagnostic(entry.stderrBuffer.trim()); + const msg = diagnostic || `CodeGraph daemon exited with code ${code}`; + rejectPendingRequests(entry.pending, new Error(msg)); + } + this.daemons.delete(_cwd); + }); + } + + private async initialize(entry: DaemonEntry, cwd: string): Promise { + const rootUri = pathToFileURL(cwd).href; + const sendRequest = this.createSender(entry); + + await sendRequest("initialize", { + protocolVersion: "2024-11-05", + rootUri, + workspaceFolders: [{ uri: rootUri, name: cwd.split(/[\\/]/).pop() || cwd }], + capabilities: {}, + clientInfo: { name: "pi-codegraph", version: "0.1.0" }, + }); + + // Send initialized notification + entry.child.stdin.write(JSON.stringify({ + jsonrpc: "2.0", + method: "initialized", + params: {}, + }) + "\n"); + + entry.initialized = true; + } + + createSender(entry: DaemonEntry): JsonRpcRequest { + return (method, params) => { + const id = entry.nextId++; + const payload = { jsonrpc: "2.0", id, method, params }; + const promise = new Promise((resolve, reject) => { + entry.pending.set(id, { resolve, reject }); + }); + entry.child.stdin.write(`${JSON.stringify(payload)}\n`); + return promise; + }; + } +} + +// Singleton manager +const daemonManager = new DaemonManager(); + +// Exported for testing +export function getDaemonManager(): DaemonManager { + return daemonManager; +} + +// ─── Session (reuses daemon) ───────────────────────────────────────────────── + +async function withDaemonSession( projectPath: string | undefined, signal: AbortSignal | undefined, fn: (request: JsonRpcRequest) => Promise, ): Promise { const cwd = await resolveProjectCwd(projectPath); - const child = spawnCodeGraphServer(cwd); - - const session = runJsonRpcSession(child, cwd, signal, fn); + const entry = await daemonManager.acquire(cwd); let timer: ReturnType | undefined; const onAbortClearTimer = () => clearTimeout(timer); const timeout = new Promise((_, reject) => { timer = setTimeout(() => { - if (!child.killed) child.kill(); reject(new Error( - "CodeGraph MCP session timed out after " + SessionTimeoutMs + "ms. " + + "CodeGraph MCP request timed out after " + SessionTimeoutMs + "ms. " + 'Try running "codegraph unlock" in the project directory, then restart pi.' )); }, SessionTimeoutMs); signal?.addEventListener("abort", onAbortClearTimer, { once: true }); }); - session.catch(() => {}); - timeout.catch(() => {}); - return Promise.race([session, timeout]).finally(() => { + const sendRequest = daemonManager.createSender(entry); + const task = fn(sendRequest); + + try { + return await Promise.race([task, timeout]); + } finally { clearTimeout(timer); signal?.removeEventListener("abort", onAbortClearTimer); - }); + daemonManager.release(cwd); + } } +// Keep backward-compatible export +export const withCodeGraphMcp = withDaemonSession; + +// ─── Utilities ─────────────────────────────────────────────────────────────── + export function normalizeWindowsPath(inputPath: string): string { let normalized = inputPath.trim(); @@ -274,80 +454,6 @@ export function sanitizeDiagnostic(value: string): string { : redacted; } -async function runJsonRpcSession( - child: ChildProcessWithoutNullStreams, - cwd: string, - signal: AbortSignal | undefined, - fn: (request: JsonRpcRequest) => Promise, -): Promise { - const pending: PendingJsonRpcRequests = new Map(); - const stderr = { value: "" }; - const cleanup = () => cleanupJsonRpcChild(child, pending); - const onAbort = () => cleanup(); - - signal?.addEventListener("abort", onAbort, { once: true }); - attachJsonRpcHandlers(child, pending, stderr); - - try { - const sendRequest = createJsonRpcRequestSender(child, pending); - await initializeJsonRpcSession(cwd, sendRequest, sendJsonRpcNotification.bind(undefined, child)); - return await fn(sendRequest); - } finally { - signal?.removeEventListener("abort", onAbort); - cleanup(); - } -} - -function cleanupJsonRpcChild( - child: ChildProcessWithoutNullStreams, - pending: PendingJsonRpcRequests, -): void { - rejectPendingJsonRpcRequests( - pending, - new Error("CodeGraph MCP process closed before responding."), - ); - if (!child.killed) child.kill(); -} - -function rejectPendingJsonRpcRequests( - pending: PendingJsonRpcRequests, - error: Error, -): void { - for (const entry of pending.values()) entry.reject(error); - pending.clear(); -} - -function attachJsonRpcHandlers( - child: ChildProcessWithoutNullStreams, - pending: PendingJsonRpcRequests, - stderr: { value: string }, -): void { - const stdout = { value: "" }; - - child.stdout.on("data", (chunk) => { - handleJsonRpcStdout(chunk, stdout, pending); - }); - child.stderr.on("data", (chunk) => { - stderr.value += chunk.toString("utf-8"); - }); - child.on("error", (err) => rejectPendingJsonRpcRequests(pending, err)); - child.on("exit", (code) => rejectPendingJsonRpcOnExit(pending, stderr.value, code)); -} - -function handleJsonRpcStdout( - chunk: Buffer, - stdout: { value: string }, - pending: PendingJsonRpcRequests, -): void { - stdout.value += chunk.toString("utf-8"); - let newline; - while ((newline = stdout.value.indexOf("\n")) !== -1) { - const line = stdout.value.slice(0, newline).trim(); - stdout.value = stdout.value.slice(newline + 1); - if (line) resolveJsonRpcLine(line, pending); - } -} - function resolveJsonRpcLine(line: string, pending: PendingJsonRpcRequests): void { let msg: any; try { @@ -363,55 +469,9 @@ function resolveJsonRpcLine(line: string, pending: PendingJsonRpcRequests): void else resolve(msg.result); } -function rejectPendingJsonRpcOnExit( - pending: PendingJsonRpcRequests, - stderr: string, - code: number | null, -): void { - if (pending.size === 0) return; - const diagnostic = sanitizeDiagnostic(stderr.trim()); - const msg = diagnostic || `CodeGraph MCP process exited with code ${code}`; - rejectPendingJsonRpcRequests(pending, new Error(msg)); -} - -function createJsonRpcRequestSender( - child: ChildProcessWithoutNullStreams, - pending: PendingJsonRpcRequests, -): JsonRpcRequest { - let nextId = 1; - return (method, params) => { - const id = nextId++; - const payload = { jsonrpc: "2.0", id, method, params }; - const promise = new Promise((resolve, reject) => { - pending.set(id, { resolve, reject }); - }); - child.stdin.write(`${JSON.stringify(payload)}\n`); - return promise; - }; -} - -function sendJsonRpcNotification( - child: ChildProcessWithoutNullStreams, - method: string, - params: Record, -): void { - child.stdin.write(`${JSON.stringify({ jsonrpc: "2.0", method, params })}\n`); -} - -async function initializeJsonRpcSession( - cwd: string, - sendRequest: JsonRpcRequest, - sendNotification: (method: string, params: Record) => void, -): Promise { - const rootUri = pathToFileURL(cwd).href; - await sendRequest("initialize", { - protocolVersion: "2024-11-05", - rootUri, - workspaceFolders: [{ uri: rootUri, name: cwd.split(/[\\/]/).pop() || cwd }], - capabilities: {}, - clientInfo: { name: "pi-codegraph", version: "0.1.0" }, - }); - sendNotification("initialized", {}); +function rejectPendingRequests(pending: PendingJsonRpcRequests, error: Error): void { + for (const entry of pending.values()) entry.reject(error); + pending.clear(); } async function prepareToolArguments(