diff --git a/__tests__/mcp-daemon.test.ts b/__tests__/mcp-daemon.test.ts index a39feb1b2..add971cb3 100644 --- a/__tests__/mcp-daemon.test.ts +++ b/__tests__/mcp-daemon.test.ts @@ -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((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((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((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( diff --git a/__tests__/query-pool.test.ts b/__tests__/query-pool.test.ts index d7bf3be91..2a1744901 100644 --- a/__tests__/query-pool.test.ts +++ b/__tests__/query-pool.test.ts @@ -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 { + 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 { + 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 }).workers; + const exits = [...workers].map((worker) => new Promise((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); +}); diff --git a/src/mcp/engine.ts b/src/mcp/engine.ts index b88f1b121..ba95eb71e 100644 --- a/src/mcp/engine.ts +++ b/src/mcp/engine.ts @@ -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>; + private opts: Required> & Pick; 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); diff --git a/src/mcp/index.ts b/src/mcp/index.ts index c12ab8105..242cc5ba8 100644 --- a/src/mcp/index.ts +++ b/src/mcp/index.ts @@ -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, diff --git a/src/mcp/query-pool.ts b/src/mcp/query-pool.ts index 0575def38..7ee0f1335 100644 --- a/src/mcp/query-pool.ts +++ b/src/mcp/query-pool.ts @@ -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; diff --git a/src/mcp/query-worker.ts b/src/mcp/query-worker.ts index b732d2560..525e9334a 100644 --- a/src/mcp/query-worker.ts +++ b/src/mcp/query-worker.ts @@ -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 | 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); diff --git a/src/mcp/tools.ts b/src/mcp/tools.ts index ee7442021..867185595 100644 --- a/src/mcp/tools.ts +++ b/src/mcp/tools.ts @@ -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 | 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