fix(mcp): enable query workers in direct mode (#1465) (#2022)

* fix(mcp): enable query workers in direct mode (#1465)

Direct and proxy-fallback engines never enabled the query pool, serializing concurrent reads despite CODEGRAPH_QUERY_POOL_SIZE.
Enable bounded workers in both paths and support sessions without a default index, including indexes discovered later.
Preserve main-thread catch-up, staleness notices, and session accounting, with regression coverage for concurrency, parity, configuration, and teardown.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(mcp): enable query workers in direct mode (#1465)

Direct and proxy-fallback engines never enabled the query pool, serializing concurrent reads despite CODEGRAPH_QUERY_POOL_SIZE.
Enable bounded workers in both paths and support sessions without a default index, including indexes discovered later.
Preserve main-thread catch-up, staleness notices, and session accounting, with regression coverage for concurrency, parity, configuration, and teardown.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(mcp): enable query workers in direct mode (#1465)

Direct and proxy-fallback engines never enabled the query pool, serializing concurrent reads despite CODEGRAPH_QUERY_POOL_SIZE.
Enable bounded workers in both paths and support sessions without a default index, including indexes discovered later.
Preserve main-thread catch-up, staleness notices, and session accounting, with regression coverage for concurrency, parity, configuration, and teardown.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Colby Mchenry
2026-09-27 13:23:38 +00:00
committed by GitHub
parent 0f44fdf8af
commit 73f7f07d9e
7 changed files with 270 additions and 40 deletions
+98 -2
View File
@@ -49,8 +49,8 @@ interface SpawnedServer {
stderr: string[];
}
function spawnServer(cwd: string, env: NodeJS.ProcessEnv = {}): SpawnedServer {
const child = spawn(process.execPath, [BIN, 'serve', '--mcp'], {
function spawnServer(cwd: string, env: NodeJS.ProcessEnv = {}, args: string[] = []): SpawnedServer {
const child = spawn(process.execPath, [BIN, 'serve', '--mcp', ...args], {
cwd,
stdio: ['pipe', 'pipe', 'pipe'],
// #618: the daemon-attach log line is now off by default; opt the test
@@ -310,6 +310,102 @@ describe('Shared MCP daemon (issue #411)', () => {
expect(fs.existsSync(path.join(realRoot, '.codegraph', 'daemon.log'))).toBe(false);
}, 20000);
it.each(['2', '0'])('direct stdio honors query pool size=%s and exits after concurrent reads', async (size) => {
fs.writeFileSync(path.join(tempDir, 'app.ts'), 'export function directSymbol() { return 42; }\n');
const cg = await CodeGraph.open(tempDir);
try { await cg.indexAll(); } finally { cg.close(); }
const server = spawnServer(tempDir, {
CODEGRAPH_NO_DAEMON: '1', CODEGRAPH_QUERY_POOL_SIZE: size,
}, ['--no-watch', '--path', tempDir]);
servers.push(server);
sendInitialize(server.child, `file://${tempDir}`, 1);
await waitFor(() => findResponse(server.stdout, 1), 10000);
await waitFor(() => server.stderr.some((l) => l.includes(size === '0'
? 'Query pool disabled' : 'Query pool: up to 2 worker thread(s)')), 10000, 25, 'direct query pool configuration');
const ids = [2, 3, 4, 5, 6, 7];
for (const id of ids) {
sendMessage(server.child, {
jsonrpc: '2.0', id, method: 'tools/call',
params: { name: 'codegraph_explore', arguments: { query: 'directSymbol' } },
});
}
const replies = await Promise.all(ids.map((id) => waitFor(() => findResponse(server.stdout, id), 15000)));
for (const reply of replies) {
expect(reply.error).toBeUndefined();
expect(reply.result.isError).toBeFalsy();
expect(reply.result.content[0].text).toContain('directSymbol');
expect(reply.result).not.toHaveProperty('_cgExploreEmission');
}
const exited = new Promise<number | null>((resolve) => server.child.once('exit', resolve));
server.child.stdin.end();
expect(await exited).toBe(0);
}, 30000);
it('direct stdio queries multiple projects without a default index', async () => {
fs.rmSync(path.join(tempDir, '.codegraph'), { recursive: true, force: true });
for (const name of ['alpha', 'beta']) {
const root = path.join(tempDir, name);
fs.mkdirSync(root);
fs.writeFileSync(path.join(root, 'app.ts'), `export function ${name}Symbol() { return 42; }\n`);
const cg = await CodeGraph.init(root);
try { await cg.indexAll(); } finally { cg.close(); }
}
const server = spawnServer(tempDir, {
CODEGRAPH_NO_DAEMON: '1', CODEGRAPH_QUERY_POOL_SIZE: '2',
}, ['--no-watch', '--path', tempDir]);
servers.push(server);
sendInitialize(server.child, `file://${tempDir}`, 1);
await waitFor(() => findResponse(server.stdout, 1), 10000);
await waitFor(() => server.stderr.some((l) => l.includes('Query pool: up to 2')), 10000);
for (const [i, name] of ['alpha', 'beta'].entries()) {
sendMessage(server.child, {
jsonrpc: '2.0', id: i + 2, method: 'tools/call',
params: { name: 'codegraph_explore', arguments: { query: `${name}Symbol`, projectPath: path.join(tempDir, name) } },
});
}
for (const [i, name] of ['alpha', 'beta'].entries()) {
const reply = await waitFor(() => findResponse(server.stdout, i + 2), 15000);
expect(reply.result.isError).toBeFalsy();
expect(reply.result.content[0].text).toContain(`${name}Symbol`);
}
sendMessage(server.child, {
jsonrpc: '2.0', id: 4, method: 'tools/call',
params: { name: 'codegraph_explore', arguments: { query: 'alphaSymbol' } },
});
const missing = await waitFor(() => findResponse(server.stdout, 4), 10000);
expect(missing.result.isError).toBeFalsy();
expect(missing.result.content[0].text).toContain('projectPath');
expect(fs.existsSync(path.join(tempDir, '.codegraph'))).toBe(false);
}, 30000);
it('proxy fallback enables the query pool when no live writer holds the project', async () => {
const net = await import('net');
const sockPath = getDaemonSocketPath(realRoot);
const miniServer = net.createServer((sock) => {
sock.end(JSON.stringify({
codegraph: '0.0.0-mismatch', pid: process.pid, socketPath: sockPath, protocol: 1,
}) + '\n');
});
await new Promise<void>((resolve) => miniServer.listen(sockPath, resolve));
try {
const server = spawnServer(tempDir, { CODEGRAPH_QUERY_POOL_SIZE: '2' }, ['--no-watch']);
servers.push(server);
sendInitialize(server.child, `file://${tempDir}`, 1);
await waitFor(() => findResponse(server.stdout, 1), 10000);
sendMessage(server.child, {
jsonrpc: '2.0', id: 2, method: 'tools/call',
params: { name: 'codegraph_files', arguments: {} },
});
const reply = await waitFor(() => findResponse(server.stdout, 2), 10000);
expect(reply.error).toBeUndefined();
expect(reply.result.isError).toBeFalsy();
expect(server.stderr.some((l) => l.includes('serving this session in-process'))).toBe(true);
expect(server.stderr.some((l) => l.includes('Query pool: up to 2'))).toBe(true);
} finally {
await new Promise<void>((resolve) => miniServer.close(() => resolve()));
}
}, 30000);
it('clears a stale (dead-pid) lockfile and a fresh daemon takes over', async () => {
// Plant a lockfile pointing at a definitely-dead pid + the real socket path.
fs.writeFileSync(
+127 -6
View File
@@ -1,15 +1,20 @@
/**
* QueryPool — the off-loop worker pool that keeps the shared daemon's main
* QueryPool — the off-loop worker pool that keeps the MCP server's main
* event loop free for the MCP transport under concurrent read load (the
* "10 subagents time out" report). These tests drive the pool's queue / growth /
* "10 subagents time out" report). Unit tests drive the pool's queue / growth /
* crash-recovery / backstop logic with INJECTED fake workers, so they exercise
* the real scheduling code without spawning threads or needing a built dist.
*
* End-to-end behavior with real worker threads (a worker opens its own WAL read
* connection and runs codegraph_explore) is validated separately against a real
* index; here we pin the orchestration that makes that safe and fair.
* Integration tests below use the built engine and real worker threads against
* real indexes to cover direct-mode routing, session accounting, and teardown.
*/
import { describe, it, expect } from 'vitest';
import { afterEach, beforeEach, describe, it, expect, vi } from 'vitest';
import * as fs from 'fs';
import * as os from 'os';
import * as path from 'path';
import { CodeGraph } from '../src';
import { ExploreSessionState, EXPLORE_EMISSION_KEY } from '../src/mcp/explore-session-state';
import type { MCPEngine } from '../src/mcp/engine';
import { QueryPool, resolvePoolSize, type PoolWorker } from '../src/mcp/query-pool';
import type { ToolResult } from '../src/mcp/tools';
@@ -196,3 +201,119 @@ describe('QueryPool', () => {
await pool.destroy();
});
});
// Use the built engine so its pool loads the real compiled worker sibling.
const BuiltEngine: typeof MCPEngine = require('../dist/mcp/engine').MCPEngine;
describe('MCP query pool with real projects (#1465)', () => {
let tempDir: string;
let engine: MCPEngine | undefined;
let pool: QueryPool | null;
beforeEach(() => {
tempDir = fs.realpathSync(fs.mkdtempSync(path.join(os.tmpdir(), 'codegraph-query-pool-')));
pool = null;
vi.stubEnv('CODEGRAPH_QUERY_POOL_SIZE', '2');
});
afterEach(async () => {
await pool?.destroy();
engine?.stop();
engine = undefined;
vi.unstubAllEnvs();
fs.rmSync(tempDir, { recursive: true, force: true });
});
async function indexProject(name: string, symbol: string): Promise<string> {
const root = path.join(tempDir, name);
fs.mkdirSync(root, { recursive: true });
fs.writeFileSync(path.join(root, 'app.ts'), `export function ${symbol}() { return 42; }\n`);
const cg = await CodeGraph.init(root);
try { await cg.indexAll(); } finally { cg.close(); }
return root;
}
async function start(root: string): Promise<MCPEngine> {
engine = new BuiltEngine({ watch: false, queryPool: true });
await engine.ensureInitialized(root);
pool = (engine as unknown as { queryPool: QueryPool | null }).queryPool;
return engine;
}
it.each([true, false])('routes concurrent reads across projects with default=%s, preserving output and session state', async (hasDefault) => {
const alpha = await indexProject('alpha', 'alphaSymbol');
const beta = await indexProject('beta', 'betaSymbol');
const workspace = path.join(tempDir, 'workspace');
fs.mkdirSync(workspace);
const activeEngine = await start(hasDefault ? alpha : workspace);
expect(pool).not.toBeNull();
await vi.waitFor(() => expect(pool!.ready).toBe(true), { timeout: 15000 });
const handler = activeEngine.getToolHandler();
// Drain catch-up before comparing the worker and in-process paths.
await handler.execute('codegraph_status', { projectPath: alpha });
const session = new ExploreSessionState();
const calls = Array.from({ length: 6 }, (_, i) => {
const projectPath = i % 2 ? beta : alpha;
const query = i % 2 ? 'betaSymbol' : 'alphaSymbol';
return handler.execute('codegraph_explore', { projectPath, query }, session);
});
expect(pool!.liveWorkers).toBe(2);
const results = await Promise.all(calls);
for (const [i, result] of results.entries()) {
expect(result.isError).toBeFalsy();
expect(result.content[0].text).toContain(i % 2 ? 'betaSymbol' : 'alphaSymbol');
expect(result.content[0].text).not.toContain(i % 2 ? 'alphaSymbol' : 'betaSymbol');
expect(result).not.toHaveProperty(EXPLORE_EMISSION_KEY);
}
expect(session.callCount(alpha)).toBe(3);
expect(session.callCount(beta)).toBe(3);
const args = { projectPath: beta, query: 'betaSymbol' };
const pooled = await handler.execute('codegraph_explore', args);
handler.setQueryPool(null);
expect(await handler.execute('codegraph_explore', args)).toEqual(pooled);
handler.setQueryPool(pool);
if (!hasDefault) {
const missing = await handler.execute('codegraph_explore', { query: 'alphaSymbol' });
expect(missing.isError).toBeFalsy();
expect(missing.content[0].text).toContain(workspace);
expect(missing.content[0].text).toContain('No CodeGraph project');
// A default index can appear after rootless workers are already warm.
await indexProject('workspace', 'lateSymbol');
activeEngine.retryInitializeSync(workspace);
const late = await handler.execute('codegraph_explore', { query: 'lateSymbol' });
expect(late.isError).toBeFalsy();
expect(late.content[0].text).toContain('lateSymbol');
}
}, 30000);
it('honors size=0 for projectPath-only sessions', async () => {
const alpha = await indexProject('alpha', 'alphaSymbol');
const workspace = path.join(tempDir, 'workspace');
fs.mkdirSync(workspace);
vi.stubEnv('CODEGRAPH_QUERY_POOL_SIZE', '0');
const activeEngine = await start(workspace);
expect(pool).toBeNull();
const result = await activeEngine.getToolHandler().execute('codegraph_explore', {
projectPath: alpha, query: 'alphaSymbol',
});
expect(result.isError).toBeFalsy();
expect(result.content[0].text).toContain('alphaSymbol');
}, 30000);
it('stops real workers and settles queued calls on engine teardown', async () => {
const alpha = await indexProject('alpha', 'alphaSymbol');
const activeEngine = await start(alpha);
expect(pool).not.toBeNull();
const calls = Array.from({ length: 6 }, () => pool!.run('codegraph_explore', { query: 'alphaSymbol' }));
const workers = (pool as unknown as { workers: Set<import('worker_threads').Worker> }).workers;
const exits = [...workers].map((worker) => new Promise<void>((resolve) => worker.once('exit', () => resolve())));
activeEngine.stop();
const results = await Promise.all(calls);
await Promise.all(exits);
expect(results.every((r) => r.content[0].text.includes('shutting down'))).toBe(true);
expect(pool!.liveWorkers).toBe(0);
expect(pool!.healthy).toBe(false);
await activeEngine.ensureInitialized(alpha);
expect((activeEngine as unknown as { queryPool: QueryPool | null }).queryPool).toBeNull();
}, 30000);
});
+22 -14
View File
@@ -39,14 +39,17 @@ export interface MCPEngineOptions {
*/
watch?: boolean;
/**
* Whether to off-load read-tool dispatch to a worker-thread pool. Only the
* SHARED daemon wants this — it serves many concurrent clients on one event
* loop, so without a pool concurrent explores serialize and starve the MCP
* transport. Direct mode (one stdio client, no concurrency) leaves it off so a
* single call never pays a worker round-trip. `CODEGRAPH_QUERY_POOL_SIZE=0`
* disables it even in daemon mode.
* Whether to off-load read-tool dispatch to a worker-thread pool. Both daemon
* and direct sessions can issue concurrent calls on one event loop.
* `CODEGRAPH_QUERY_POOL_SIZE=0` disables it in either mode.
*/
queryPool?: boolean;
/**
* Worker cap when `CODEGRAPH_QUERY_POOL_SIZE` is unset. A direct (single-client)
* session sets a small cap so every session doesn't hold one worker per core;
* the shared daemon leaves it unset and scales with the machine.
*/
queryPoolDefaultMax?: number;
/**
* Project root whose writer slot must be claimed synchronously before this
* engine can open the graph. Used by proxy fallback to fence catch-up sync,
@@ -76,14 +79,14 @@ export class MCPEngine {
private watcherStarted = false;
/** Set when this engine holds writer.pid (#1740). */
private writerLockRoot: string | null = null;
private opts: Required<Omit<MCPEngineOptions, 'writerLockRoot'>>;
private opts: Required<Omit<MCPEngineOptions, 'writerLockRoot' | 'queryPoolDefaultMax'>> & Pick<MCPEngineOptions, 'queryPoolDefaultMax'>;
private closed = false;
// Off-loop read-tool pool (daemon mode only). Created lazily once the default
// project is open — workers each hold their own WAL read connection.
// Off-loop read-tool pool. Workers each hold their own WAL read connections;
// sessions without a default index open projects lazily via projectPath.
private queryPool: QueryPool | null = null;
constructor(opts: MCPEngineOptions = {}) {
this.opts = { watch: opts.watch ?? true, queryPool: opts.queryPool ?? false };
this.opts = { watch: opts.watch ?? true, queryPool: opts.queryPool ?? false, queryPoolDefaultMax: opts.queryPoolDefaultMax };
this.toolHandler = new ToolHandler(null);
if (opts.writerLockRoot) {
const writer = tryAcquireWriterLock(opts.writerLockRoot, 'fallback');
@@ -95,14 +98,18 @@ export class MCPEngine {
}
/**
* Start the worker-thread query pool once a default project is open (daemon
* mode only; honors `CODEGRAPH_QUERY_POOL_SIZE`). Idempotent and best-effort:
* Start the worker-thread query pool after resolving the default project
* (which may be absent). Honors `CODEGRAPH_QUERY_POOL_SIZE`; best-effort:
* if workers can't spawn on this platform the ToolHandler keeps serving reads
* in-process, so the pool can only help, never break, tool calls.
*/
private maybeStartPool(root: string): void {
private maybeStartPool(root: string | null): void {
if (!this.opts.queryPool || this.queryPool || this.closed) return;
const size = resolvePoolSize(process.env.CODEGRAPH_QUERY_POOL_SIZE, os.cpus().length);
const envSize = process.env.CODEGRAPH_QUERY_POOL_SIZE;
let size = resolvePoolSize(envSize, os.cpus().length);
if ((envSize === undefined || envSize === '') && this.opts.queryPoolDefaultMax !== undefined) {
size = Math.min(size, this.opts.queryPoolDefaultMax);
}
if (size <= 0) {
process.stderr.write('[CodeGraph MCP] Query pool disabled (CODEGRAPH_QUERY_POOL_SIZE=0); serving reads in-process.\n');
return;
@@ -263,6 +270,7 @@ export class MCPEngine {
`[CodeGraph MCP] Indexed sub-projects found: ${rels.join(', ')}. Pass \`projectPath\` per call, or launch with --path.\n`
);
}
this.maybeStartPool(null);
return;
}
if (res.viaSubScan) this.logSubprojectAdoption(searchFrom, resolvedRoot);
+5 -2
View File
@@ -72,6 +72,9 @@ import { armStartupHandshakeTimeout } from './startup-handshake';
import { treatStdinFailureAsShutdown } from './stdin-teardown';
import { HOST_PPID_ENV } from '../extraction/wasm-runtime-flags';
/** Default worker cap for a direct (single-client) session; see MCPEngineOptions.queryPoolDefaultMax (#1465). */
const DIRECT_QUERY_POOL_MAX = 2;
/**
* Env var that marks a process as the *detached daemon* itself (set by
* {@link spawnDetachedDaemon} when it re-invokes the CLI). Without it a
@@ -121,7 +124,7 @@ function makeFallbackEngine(root: string): MCPEngine {
`Cannot start an in-process fallback while live daemon pid ${existing.pid} holds the project lock.`
);
}
return new MCPEngine({ writerLockRoot: root });
return new MCPEngine({ writerLockRoot: root, queryPool: true, queryPoolDefaultMax: DIRECT_QUERY_POOL_MAX });
}
/**
@@ -428,7 +431,7 @@ export class MCPServer {
this.writerLockRoot = writerRoot;
}
this.engine = new MCPEngine();
this.engine = new MCPEngine({ queryPool: true, queryPoolDefaultMax: DIRECT_QUERY_POOL_MAX });
const transport = new StdioTransport();
this.session = new MCPSession(transport, this.engine, {
explicitProjectPath: this.projectPath,
+3 -3
View File
@@ -89,8 +89,8 @@ interface Job {
}
export interface QueryPoolOptions {
/** Default project root each worker opens at spawn. */
root: string;
/** Default project root each worker opens at spawn, or null for per-call projects. */
root: string | null;
/** Max worker threads. Defaults to `clamp(cores-1, 1, 16)`. */
size?: number;
/** Linger before a queued call gets busy-guidance. Default 45s. */
@@ -153,7 +153,7 @@ export class QueryPool {
private nextId = 1;
private totalCrashes = 0;
private destroyed = false;
private readonly root: string;
private readonly root: string | null;
private readonly maxSize: number;
private readonly softTimeoutMs: number;
private readonly maxRetries: number;
+3 -3
View File
@@ -25,7 +25,7 @@ import { parentPort, workerData } from 'worker_threads';
import type { ToolResult } from './tools';
interface WorkerInit {
root: string;
root: string | null;
}
interface CallMessage {
@@ -46,14 +46,14 @@ if (parentPort) {
const port = parentPort;
const { root } = workerData as WorkerInit;
// Open the default project's READ connection once, at spawn. Other repos are
// Open the default project's READ connection once, if present. Other repos are
// opened lazily on first cross-project (projectPath) call by the ToolHandler's
// own per-handler cache. openSync does not start a watcher — workers are pure
// readers; the single watcher/writer stays on the daemon's main thread.
let handler: InstanceType<typeof import('./tools').ToolHandler> | null = null;
let initError: string | null = null;
try {
const cg = loadCodeGraph().openSync(root);
const cg = root ? loadCodeGraph().openSync(root) : null;
handler = new (loadToolHandler())(cg);
} catch (err) {
initError = err instanceof Error ? err.message : String(err);
+12 -10
View File
@@ -1531,17 +1531,15 @@ export class ToolHandler {
// huge repo can't hang the first call (#905); cleared on first await so
// subsequent calls don't pay any cost.
private catchUpGate: Promise<void> | null = null;
// Optional worker-thread pool for off-loop read-tool dispatch (daemon mode).
// When set + healthy, the heavy read tools run on a worker so the daemon's
// main loop stays free for the MCP transport under concurrent load. Null in
// direct/in-process mode (one client, no concurrency to parallelize).
// Optional worker-thread pool for off-loop read-tool dispatch. When ready +
// healthy, heavy reads leave the main loop free for the MCP transport.
private queryPool: QueryPool | null = null;
constructor(private cg: CodeGraph | null) {}
/**
* Engine-only: attach (or detach with null) the worker-thread query pool. The
* shared daemon sets this once its default project is open; the workers each
* engine sets this after resolving its default project; the workers each
* hold their own WAL read connection and run {@link executeReadTool}. A
* worker's own ToolHandler never has a pool, so there is no nested off-loading.
*/
@@ -2207,15 +2205,15 @@ export class ToolHandler {
}
// Read tools: off-load the CPU-heavy dispatch to the worker pool when one
// is attached, healthy, AND has finished its first cold start (daemon
// mode), so the daemon's single event loop stays free for the MCP
// is attached, healthy, AND has finished its first cold start,
// so the server's single event loop stays free for the MCP
// transport under concurrent load — otherwise N concurrent explores
// serialize AND starve the transport until the whole batch drains
// (clients then time out). Before the first worker is warm, calls run
// in-process: a call queued behind a cold start sat invisible until the
// 45s busy backstop — the daemon's first tool call stalling for however
// long a worker spawn takes on a loaded machine (the #662 flake). With
// no pool (direct mode) or a degraded one, dispatch runs in-process
// no pool or a degraded one, dispatch runs in-process
// exactly as before. Either way the result flows through the
// cross-cutting notices — worktree-index mismatch (#155) and per-file
// staleness (#403) — which need the watched MAIN instance and so are
@@ -2227,8 +2225,12 @@ export class ToolHandler {
// structured-clone boundary into a worker, where a closure or a handler
// field could not follow.
const dispatchArgs = this.withSessionView(toolName, args, sessionState);
const raw = (this.queryPool && this.queryPool.healthy && this.queryPool.ready)
? await this.queryPool.run(toolName, dispatchArgs)
// The default may have appeared after the workers started. Pass the
// main thread's current root explicitly; with no root or projectPath,
// keep the main handler's workspace-specific not-indexed guidance.
const projectPath = args.projectPath ?? this.cg?.getProjectRoot();
const raw = (projectPath && this.queryPool && this.queryPool.healthy && this.queryPool.ready)
? await this.queryPool.run(toolName, { ...dispatchArgs, projectPath })
: await this.executeReadTool(toolName, dispatchArgs);
// Record + STRIP before anything else touches the result: the emission is
// internal bookkeeping and must never reach the client, whether or not a