Skip to content

feat: reuse CodeGraph daemon across tool calls #57

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
zimoo354 wants to merge 1 commit into from
Closed

feat: reuse CodeGraph daemon across tool calls #57

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions .changeset/daemon-reuse.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
114 changes: 53 additions & 61 deletions __tests__/codegraph.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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";
Expand Down Expand Up @@ -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");
Expand All @@ -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");

Expand All @@ -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");

Expand Down Expand Up @@ -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();
Expand All @@ -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<string, unknown> | undefined;

getDaemonManager().killAll();

vi.mocked(spawn).mockImplementationOnce(() => {
const child = new EventEmitter() as any;
child.stdin = new PassThrough();
Expand Down
Loading