Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,7 @@ and adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
- An MCP tool call whose `projectPath` names a file inside a project, such as the file an agent is looking at, now answers from that project. Before, it failed with "Path is not a directory", even though a path that doesn't exist at all worked. Sensitive folders such as `/etc` are still refused when the path names a file inside them. Thanks @inth3shadows for the report and @maxmilian for the fix. (#2982)
- An index left incomplete by an indexing run that was stopped partway, by a crash, an out-of-memory kill or the server's own hang check, no longer gets the background server killed with "Main thread unresponsive" every time it starts: on a large project, restoring the database indexes the run had set aside took longer than the hang check allows, so each restart died at the same point and the index was never repaired. They are now restored in the background, and the project is served once they are back. Thanks @cbeaulieu-gt for the report and the first fix. (#1887)
- A new agent session no longer waits up to about 12 minutes for its first answer when the background server is stuck: a server whose process still holds its connection open but never responds made the session wait out a long series of retries. The session now gives up within seconds and answers on its own, and switches back to the background server once that has been restarted.
- The background server now recycles its query workers after five minutes idle, releasing the memory a burst of large queries left behind while keeping one fresh worker ready for the next call. Thanks @danusha2345 and @doomsday1906.
- An MCP session whose agent stops reading its answers, while keeping the session open, now shuts down the normal way the next time it answers: it stops its work, frees the project for the next session and exits cleanly, and other sessions sharing the background server carry on. Before, that answer crashed it with "Uncaught exception: Error: EPIPE". Thanks @iiijiashu for the report and the first fix. (#2991)
- A tool call whose arguments aren't an object, like `"arguments": "alpha"` from a hand-written MCP client or script, now gets the standard invalid-params error saying the arguments must be an object, and no tool runs. Before, a string or a number made most tools fail and come back as an internal CodeGraph error that told the agent to retry and then to continue without CodeGraph, and an array was run as if it were the arguments. Thanks @inth3shadows for the report and @kjgggggg for the first fix. (#2981)
- The bundled Node runtime moves to 24.22.1, Node's October security release, and every release now checks the Node it bundles against the checksums nodejs.org publishes. The new runtime also fixes a crash on Windows as a command like `codegraph sync` exits, after its work is done, printing `Assertion failed: !(handle->flags & UV_HANDLE_CLOSING)`, which made a script or CI run that had succeeded fail. Thanks @Ancienttwo for the report and the first fix, and @IdealBoy709 for the report on the security release. (#2617, #2780)
Expand Down
161 changes: 159 additions & 2 deletions __tests__/query-pool.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,10 @@ class FakeWorker implements PoolWorker {
private errorCb?: (e: Error) => void;
private exitCb?: (code: number) => void;
alive = true;
terminated = false;
exitCode: number | undefined;
/** Ignore the pool's 'close', to exercise its terminate fallback. */
ignoresClose = false;
constructor(private behavior: (m: CallMsg) => Action, readyOk: boolean | null = true) {
setTimeout(() => { if (this.alive && readyOk !== null) this.emitMessage({ type: 'ready', ok: readyOk }); }, 0);
}
Expand All @@ -51,7 +55,14 @@ class FakeWorker implements PoolWorker {
if (this.alive) this.msgCb?.({ type: 'result', id, result });
}
postMessage(msg: unknown): void {
const m = msg as CallMsg;
const m = msg as CallMsg | { type: 'close' };
if (m?.type === 'close') {
// Like the real worker: close up and exit by itself, with code 0.
if (this.ignoresClose) return;
this.alive = false;
setTimeout(() => { this.exitCode = 0; this.exitCb?.(0); }, 0);
return;
}
if (!m || m.type !== 'call') return;
const action = this.behavior(m);
if ('crash' in action) {
Expand All @@ -63,7 +74,7 @@ class FakeWorker implements PoolWorker {
if ('wait' in action) { void action.wait.then((r) => this.reply(m.id, r)); return; }
setTimeout(() => this.reply(m.id, action.result), 0);
}
terminate(): Promise<number> { this.alive = false; return Promise.resolve(0); }
terminate(): Promise<number> { this.alive = false; this.terminated = true; return Promise.resolve(0); }
}

const ok = (text: string): ToolResult => ({ content: [{ type: 'text', text }] });
Expand Down Expand Up @@ -126,6 +137,127 @@ describe('QueryPool', () => {
await pool.destroy();
});

it('recycles an idle burst back to one fresh worker', async () => {
let release!: () => void;
const gate = new Promise<void>((r) => { release = r; });
const workers: FakeWorker[] = [];
const pool = new QueryPool({
root: '/x', size: 4, idleShrinkMs: 20,
createWorker: () => {
const worker = new FakeWorker((m) => ({
wait: gate.then(() => ok(`r${m.id}`)),
}));
workers.push(worker);
return worker;
},
});

const calls = Promise.all(Array.from({ length: 4 }, (_, i) => pool.run('codegraph_search', { i })));
await sleep(40);
expect(pool.liveWorkers).toBe(4);
release();
await calls;
await sleep(40);

expect(pool.liveWorkers).toBe(1);
expect(workers).toHaveLength(5); // four used isolates replaced by one clean isolate
expect(workers.slice(0, 4).every((w) => !w.alive)).toBe(true);
// Each used worker exited by itself on 'close'; none was terminated.
expect(workers.slice(0, 4).every((w) => w.exitCode === 0 && !w.terminated)).toBe(true);
expect(pool.ready).toBe(true);
const again = await pool.run('codegraph_node', { symbol: 's' });
expect(again.isError).toBeFalsy();
await pool.destroy();
});

it('terminates a retired idle worker that does not exit by itself', async () => {
const workers: FakeWorker[] = [];
const pool = new QueryPool({
root: '/x', size: 1, idleShrinkMs: 20, retireTimeoutMs: 200,
createWorker: () => {
const worker = new FakeWorker(() => ({ result: ok('r') }));
worker.ignoresClose = workers.length === 0;
workers.push(worker);
return worker;
},
});

await pool.run('codegraph_search', {});
await vi.waitFor(() => expect(workers).toHaveLength(2)); // retired, replaced
expect(workers[0].terminated).toBe(false); // still has time to exit
await vi.waitFor(() => expect(workers[0].terminated).toBe(true));
expect(workers[1].terminated).toBe(false);
await pool.destroy();
});

it('recycles even when a burst worker reports ready after the last result', async () => {
const workers: FakeWorker[] = [];
let release!: () => void;
const gate = new Promise<void>((r) => { release = r; });
const pool = new QueryPool({
root: '/x', size: 2, idleShrinkMs: 20,
createWorker: () => {
// The second worker stays pending until the test lets it report ready.
const worker = new FakeWorker((m) => ({ wait: gate.then(() => ok(`r${m.id}`)) }), workers.length === 0 ? true : null);
workers.push(worker);
return worker;
},
});

await vi.waitFor(() => expect(pool.ready).toBe(true));
const calls = Promise.all([pool.run('codegraph_search', { i: 0 }), pool.run('codegraph_search', { i: 1 })]);
await vi.waitFor(() => expect(workers).toHaveLength(2));
release();
await sleep(40);
expect(workers).toHaveLength(2); // one worker is still starting, so nothing recycled
workers[1]!.emitMessage({ type: 'ready', ok: true });
await calls;
await vi.waitFor(() => expect(pool.liveWorkers).toBe(1));
expect(workers).toHaveLength(3);
await pool.destroy();
});

it('never recycles a warm worker that has served nothing', async () => {
const workers: FakeWorker[] = [];
const pool = new QueryPool({
root: '/x', size: 2, idleShrinkMs: 5,
createWorker: () => {
const worker = new FakeWorker(() => ({ result: ok('r') }));
workers.push(worker);
return worker;
},
});

await vi.waitFor(() => expect(pool.ready).toBe(true));
await sleep(40);
expect(workers).toHaveLength(1);

// A recycle leaves one clean worker, which then stays until it is used.
await pool.run('codegraph_search', {});
await vi.waitFor(() => expect(workers).toHaveLength(2));
await sleep(40);
expect(workers).toHaveLength(2);
await pool.destroy();
});

it('destroy waits for a retired worker still on its way out', async () => {
const workers: FakeWorker[] = [];
const pool = new QueryPool({
root: '/x', size: 1, idleShrinkMs: 5, retireTimeoutMs: 60,
createWorker: () => {
const worker = new FakeWorker(() => ({ result: ok('r') }));
worker.ignoresClose = workers.length === 0;
workers.push(worker);
return worker;
},
});

await pool.run('codegraph_search', {});
await vi.waitFor(() => expect(workers).toHaveLength(2));
await pool.destroy();
expect(workers[0]!.terminated).toBe(true); // the timeout ended it, inside destroy()
});

it('recovers from a worker crash: retries the in-flight call and respawns', async () => {
let calls = 0;
const pool = new QueryPool({
Expand Down Expand Up @@ -465,6 +597,31 @@ describe('MCP query pool with real projects (#1465)', () => {
}
}, 30000);

it('retires an idle real worker that exits by itself with code 0', async () => {
const alpha = await indexProject('alpha', 'alphaSymbol');
const terminated: boolean[] = [];
const exits: Promise<number>[] = [];
pool = new QueryPool({
root: alpha, size: 1, idleShrinkMs: 50,
createWorker: () => {
const worker = new Worker(path.resolve(__dirname, '../dist/mcp/query-worker.js'), { workerData: { root: alpha } });
const i = terminated.push(false) - 1;
const terminate = worker.terminate.bind(worker);
worker.terminate = () => {
terminated[i] = true;
return terminate();
};
exits.push(new Promise((resolve) => worker.once('exit', resolve)));
return worker;
},
});
const result = await pool.run('codegraph_search', { query: 'alphaSymbol' });
expect(result.isError).toBeFalsy();
expect(await exits[0]).toBe(0);
expect(terminated[0]).toBe(false);
expect(exits).toHaveLength(2); // a fresh worker replaced it
}, 30000);

it('never terminates a real worker before it has started', async () => {
const alpha = await indexProject('alpha', 'alphaSymbol');
const startedWhenEnded: boolean[] = [];
Expand Down
108 changes: 103 additions & 5 deletions src/mcp/query-pool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,15 @@ export interface PoolWorker {
/** Default linger before a queued call is answered with busy-guidance. */
const DEFAULT_BUSY_TIMEOUT_MS = 45_000; // < the ~60s MCP client request timeout

/** Recycle a quiescent pool after a burst so worker-local graph/SQLite caches
* do not remain resident for the daemon's whole lifetime. Long enough that an
* agent pausing to edit or build keeps its warm workers: the first explore on
* a fresh one takes about 200 ms more. */
const DEFAULT_IDLE_SHRINK_MS = 300_000;

/** How long a retired idle worker gets to exit on 'close' before it is terminated. */
const RETIRE_TIMEOUT_MS = 5_000;

/** Hard ceiling on pool size regardless of core count / env. */
const MAX_POOL_SIZE = 16;

Expand Down Expand Up @@ -108,6 +117,10 @@ export interface QueryPoolOptions {
softTimeoutMs?: number;
/** Retries for an in-flight call whose worker crashed. Default 1. */
maxRetries?: number;
/** Idle delay before recycling back to one fresh worker. Default 5 minutes; 0 disables. */
idleShrinkMs?: number;
/** How long a retired idle worker gets to exit by itself (tests shorten it). Default 5s. */
retireTimeoutMs?: number;
/** Worker factory (tests inject a fake). Defaults to a real `worker_threads` Worker. */
createWorker?: () => PoolWorker;
/** How long destroy() waits on a worker still starting up (tests shorten it). Default 15s. */
Expand Down Expand Up @@ -174,14 +187,25 @@ export class QueryPool {
private readonly maxSize: number;
private readonly softTimeoutMs: number;
private readonly maxRetries: number;
private readonly idleShrinkMs: number;
private readonly retireTimeoutMs: number;
private readonly createWorker: () => PoolWorker;
private readonly startSettleMs: number;
private idleShrinkTimer?: NodeJS.Timeout;
// Retired idle workers asked to exit: the timer that terminates one that
// hasn't by retireTimeoutMs, and what destroy() waits on.
private retiring = new Map<PoolWorker, { cap: NodeJS.Timeout; gone: Promise<void>; settle: () => void }>();
// Set by a served call and cleared by a recycle, so a clean warm worker that
// has answered nothing is never recycled for its own sake.
private servedSinceRecycle = false;

constructor(opts: QueryPoolOptions) {
this.root = opts.root;
this.maxSize = Math.max(1, Math.min(opts.size ?? Math.max(1, os.cpus().length - 1), MAX_POOL_SIZE));
this.softTimeoutMs = opts.softTimeoutMs ?? resolveBusyTimeoutMs();
this.maxRetries = opts.maxRetries ?? 1;
this.idleShrinkMs = Math.max(0, opts.idleShrinkMs ?? DEFAULT_IDLE_SHRINK_MS);
this.retireTimeoutMs = opts.retireTimeoutMs ?? RETIRE_TIMEOUT_MS;
this.createWorker = opts.createWorker ?? (() => new Worker(WORKER_FILE, { workerData: { root: this.root } }));
this.startSettleMs = opts.startSettleMs ?? WORKER_START_SETTLE_MS;
this.spawnOne(); // one eager warm worker, ready for the first call
Expand Down Expand Up @@ -211,8 +235,9 @@ export class QueryPool {
* hitting that window was the recurring #662 test flake (and a real
* first-call stall for agents). The pool exists for CONCURRENT load, which
* by definition arrives after warm-up; the pre-pool in-process path is
* strictly better while nothing is warm. Stays true for the pool's
* lifetime — later crash-respawn gaps are covered by retry + backstop.
* strictly better while nothing is warm, which also holds while an idle
* recycle's fresh worker starts. Stays true across crash-respawn gaps, which
* are covered by retry + backstop.
*/
get ready(): boolean {
return this.everReady && !this.destroyed;
Expand All @@ -235,7 +260,7 @@ export class QueryPool {
this.starts.set(w, { settled, settle });
w.on('message', (m) => this.onMessage(w, (m ?? {}) as WorkerMessage));
w.on('error', () => { this.startSettled(w); this.onWorkerGone(w); });
w.on('exit', (code) => { this.startSettled(w); if (code !== 0) this.onWorkerGone(w); });
w.on('exit', (code) => { this.startSettled(w); this.retired(w); if (code !== 0) this.onWorkerGone(w); });
}

private onMessage(w: PoolWorker, m: WorkerMessage): void {
Expand All @@ -252,17 +277,84 @@ export class QueryPool {
this.everReady = true;
this.idle.push(w);
this.drain();
this.armIdleShrink(); // the last result may have come while this one started
return;
}
if (m.type === 'result') {
const job = this.inflight.get(w);
this.inflight.delete(w);
this.idle.push(w);
this.servedSinceRecycle = true;
if (job) this.settle(job, m.result ?? busyGuidance(0));
this.drain();
this.armIdleShrink();
}
}

private clearIdleShrink(): void {
if (this.idleShrinkTimer) clearTimeout(this.idleShrinkTimer);
this.idleShrinkTimer = undefined;
}

private armIdleShrink(): void {
this.clearIdleShrink();
if (
this.idleShrinkMs === 0 || this.destroyed || !this.servedSinceRecycle ||
this.queue.length > 0 || this.inflight.size > 0 || this.pendingWorkers.size > 0 ||
this.idle.length !== this.workers.size
) return;
this.idleShrinkTimer = setTimeout(() => this.shrinkIdle(), this.idleShrinkMs);
this.idleShrinkTimer.unref?.();
}

private shrinkIdle(): void {
this.idleShrinkTimer = undefined;
if (
this.destroyed || this.queue.length > 0 || this.inflight.size > 0 ||
this.pendingWorkers.size > 0 || this.idle.length !== this.workers.size
) return;

// Drop every used isolate, including the last one: a single worker can hold
// hundreds of MB in graph and SQLite caches after a large query. Replace it
// with one clean warm worker so the next burst still avoids a cold queue.
const stale = [...this.workers];
this.workers.clear();
this.idle = [];
this.everReady = false;
this.servedSinceRecycle = false;
for (const w of stale) this.retire(w);
this.spawnOne();
}

/**
* Ask an idle worker to exit by itself: it closes its connections and
* collects garbage first, so no marking is in flight when its thread ends
* (see worker-teardown.ts). One that hasn't exited by `retireTimeoutMs` is
* terminated.
*/
private retire(w: PoolWorker): void {
const terminate = (): void => { void this.terminateStarted(w).then(() => this.retired(w)); };
const cap = setTimeout(terminate, this.retireTimeoutMs);
cap.unref?.();
let settle!: () => void;
const gone = new Promise<void>((resolve) => { settle = resolve; });
this.retiring.set(w, { cap, gone, settle });
try {
w.postMessage({ type: 'close' });
} catch {
clearTimeout(cap);
terminate();
}
}

private retired(w: PoolWorker): void {
const retirement = this.retiring.get(w);
if (!retirement) return;
clearTimeout(retirement.cap);
this.retiring.delete(w);
retirement.settle();
}

// A worker died (crash hook, OOM, segfault, exit≠0). Respawn a replacement and
// retry its in-flight job once; a job that keeps crashing workers fails
// gracefully so it can't loop the pool forever.
Expand Down Expand Up @@ -319,6 +411,7 @@ export class QueryPool {

/** Run a read tool on the pool. Always resolves (never rejects). */
run(toolName: string, args: Record<string, unknown>): Promise<ToolResult> {
this.clearIdleShrink();
return new Promise<ToolResult>((resolve) => {
const job: Job = {
id: this.nextId++, toolName, args, resolve,
Expand Down Expand Up @@ -364,11 +457,13 @@ export class QueryPool {
/**
* Terminate all workers and answer any outstanding calls gracefully. A worker
* still starting up is terminated once it has started (see
* WORKER_START_SETTLE_MS); outstanding calls are answered first.
* WORKER_START_SETTLE_MS); outstanding calls are answered first. A retired
* worker still on its way out is waited for, not left to end with the process.
*/
async destroy(): Promise<void> {
if (this.destroyed) return;
this.destroyed = true;
this.clearIdleShrink();
const ws = [...this.workers];
this.workers.clear();
this.pendingWorkers.clear();
Expand All @@ -378,6 +473,9 @@ export class QueryPool {
}
this.inflight.clear();
this.queue = [];
await Promise.all(ws.map((w) => this.terminateStarted(w)));
await Promise.all([
...ws.map((w) => this.terminateStarted(w)),
...[...this.retiring.values()].map((retirement) => retirement.gone),
]);
}
}
Loading