From 20a9458d52c887c7770a8c8137ee948985033c23 Mon Sep 17 00:00:00 2001 From: mixxer Date: Fri, 25 Sep 2026 12:40:52 +0900 Subject: [PATCH 1/2] fix(mcp): trim idle project connections during active calls --- __tests__/mcp-project-cache.test.ts | 68 +++++++++++++++++++++++++++++ src/mcp/tools.ts | 27 ++++++++++-- 2 files changed, 92 insertions(+), 3 deletions(-) create mode 100644 __tests__/mcp-project-cache.test.ts diff --git a/__tests__/mcp-project-cache.test.ts b/__tests__/mcp-project-cache.test.ts new file mode 100644 index 0000000000..93513627ec --- /dev/null +++ b/__tests__/mcp-project-cache.test.ts @@ -0,0 +1,68 @@ +import fs from 'fs'; +import os from 'os'; +import path from 'path'; +import { afterEach, expect, it, vi } from 'vitest'; +import CodeGraph from '../src/index'; +import { ToolHandler, MAX_CACHED_PROJECTS, __setLoadCodeGraphForTests } from '../src/mcp/tools'; + +const roots: string[] = []; +let handler: ToolHandler | null = null; + +afterEach(async () => { + await handler?.closeAll(); + handler = null; + __setLoadCodeGraphForTests(null); + for (const root of roots.splice(0)) fs.rmSync(root, { recursive: true, force: true }); +}); + +it('trims idle projects while another project call remains active', async () => { + const root = fs.mkdtempSync(path.join(os.tmpdir(), 'cg-project-cache-')); + roots.push(root); + __setLoadCodeGraphForTests(CodeGraph); + handler = new ToolHandler(null); + const cache = handler as unknown as { + projectCache: Map; + getCodeGraph(projectPath: string): CodeGraph; + trimProjects(): void; + }; + const first = path.join(root, 'project-0'); + fs.mkdirSync(first); + CodeGraph.initSync(first).close(); + const oldest = cache.getCodeGraph(first); + let finish!: (result: { content: [{ type: 'text'; text: string }] }) => void; + const pending = new Promise<{ content: [{ type: 'text'; text: string }] }>(resolve => { finish = resolve; }); + vi.spyOn(cache as unknown as { dispatchTool(): typeof pending }, 'dispatchTool') + .mockReturnValue(pending); + const alias = path.join(root, 'active-alias'); + if (process.platform !== 'win32') fs.symlinkSync(first, alias, 'dir'); + const activeCall = handler.executeReadTool('codegraph_search', { + projectPath: process.platform === 'win32' ? first : alias, + }); + for (let i = 1; i <= MAX_CACHED_PROJECTS; i++) { + const project = path.join(root, `project-${i}`); + fs.mkdirSync(project); + CodeGraph.initSync(project).close(); + cache.getCodeGraph(project); + } + const close = vi.spyOn(oldest!, 'close'); + cache.trimProjects(); + expect(cache.projectCache.size).toBe(MAX_CACHED_PROJECTS); + expect(close).not.toHaveBeenCalled(); + + const extra = path.join(root, 'project-extra'); + fs.mkdirSync(extra); + CodeGraph.initSync(extra).close(); + cache.getCodeGraph(extra); + cache.trimProjects(); + expect(cache.projectCache.size).toBe(MAX_CACHED_PROJECTS); + expect(close).not.toHaveBeenCalled(); + + finish({ content: [{ type: 'text', text: 'done' }] }); + await activeCall; + const last = path.join(root, 'project-last'); + fs.mkdirSync(last); + CodeGraph.initSync(last).close(); + cache.getCodeGraph(last); + cache.trimProjects(); + expect(close).toHaveBeenCalledOnce(); +}); diff --git a/src/mcp/tools.ts b/src/mcp/tools.ts index c8fec8d3e5..1202d9a395 100644 --- a/src/mcp/tools.ts +++ b/src/mcp/tools.ts @@ -1937,6 +1937,7 @@ export class ToolHandler { // agent instead of "no project loaded" — today only the Windows/WSL // shared-index error (#995). Engine-maintained; cleared by a successful open. private defaultOpenFailure: WslSharedIndexError | null = null; + private activeProjectRoots = new Map(); // Per-start-path cache of the git worktree/index mismatch (issue #155). The // mismatch is a fixed property of (where the request came from → which // .codegraph/ it resolves to), so the up-to-two `git rev-parse` spawns run @@ -2381,14 +2382,27 @@ export class ToolHandler { * Evicts over the LRU bound, on close, and once idle past the timeout * (#2087). The cache is in last-use order, so idle entries lead it. */ + private pinProject(projectPath: unknown): string | null { + const root = typeof projectPath === 'string' ? findNearestCodeGraphRoot(projectPath) : null; + if (root) this.activeProjectRoots.set(root, (this.activeProjectRoots.get(root) ?? 0) + 1); + return root; + } + + private unpinProject(root: string | null): void { + if (!root) return; + const count = this.activeProjectRoots.get(root)!; + if (count === 1) this.activeProjectRoots.delete(root); + else this.activeProjectRoots.set(root, count - 1); + } + private trimProjects(): void { - if (this.activeCalls > 0) return; + if (this.activeCalls > 0 && (this.closing || this.activeProjectRoots.size === 0)) return; const idleMs = resolveProjectIdleTimeoutMs(); const now = Date.now(); for (const [root, cg] of this.projectCache) { const idle = idleMs > 0 && now - (this.projectUsedAt.get(root) ?? now) >= idleMs; if (!this.closing && this.projectCache.size <= MAX_CACHED_PROJECTS && !idle) break; - if (this.projectGates.has(cg)) continue; + if (this.projectGates.has(cg) || [...this.activeProjectRoots.keys()].some(active => isSameIndexRoot(root, active))) continue; this.projectCache.delete(root); this.projectUsedAt.delete(root); if (this.projectLifecycle) { @@ -2413,7 +2427,7 @@ export class ToolHandler { private scheduleIdleRelease(idleMs: number): void { if (this.idleReleaseTimer || this.closing || idleMs <= 0) return; for (const [root, cg] of this.projectCache) { - if (this.projectGates.has(cg)) continue; + if (this.projectGates.has(cg) || [...this.activeProjectRoots.keys()].some(active => isSameIndexRoot(root, active))) continue; const due = (this.projectUsedAt.get(root) ?? Date.now()) + idleMs - Date.now(); this.idleReleaseTimer = setTimeout(() => { this.idleReleaseTimer = null; @@ -2761,6 +2775,7 @@ export class ToolHandler { ): Promise { if (this.closing) return this.textResult('This MCP session is closing; retry with a connected session.'); this.activeCalls++; + let pinnedRoot: string | null = null; try { // Block the first tool call on the engine's post-open reconcile so we // never serve rows for files deleted/edited while no MCP server was @@ -2794,6 +2809,7 @@ export class ToolHandler { if (typeof pathCheck === 'string') { await this.awaitProjectGate(pathCheck); } + pinnedRoot = this.pinProject(pathCheck); // The `path` and `pattern` properties used by codegraph_files are // also path-shaped — apply the same cap. if (args.path !== undefined) { @@ -2905,6 +2921,7 @@ export class ToolHandler { 'continue without codegraph for this task.' ); } finally { + this.unpinProject(pinnedRoot); this.activeCalls--; this.trimProjects(); } @@ -2975,6 +2992,7 @@ export class ToolHandler { * path validation already ran in {@link execute} before routing here. */ async executeReadTool(toolName: string, args: Record): Promise { + const pinnedRoot = this.pinProject(args.projectPath); try { return await this.dispatchTool(toolName, args); } catch (err) { @@ -2992,6 +3010,9 @@ export class ToolHandler { 'This is an internal codegraph error — retry the call once; if it persists, ' + 'continue without codegraph for this task.' ); + } finally { + this.unpinProject(pinnedRoot); + this.trimProjects(); } } From 7398a498102d58cff96e296e8ae6ce45e3109c5b Mon Sep 17 00:00:00 2001 From: mixxer Date: Mon, 5 Oct 2026 00:21:47 +0900 Subject: [PATCH 2/2] fix(mcp): retain pinned project shutdown guarantees --- __tests__/mcp-project-cache.test.ts | 19 +++++++++++++++++++ src/mcp/tools.ts | 22 ++++++++++++---------- 2 files changed, 31 insertions(+), 10 deletions(-) diff --git a/__tests__/mcp-project-cache.test.ts b/__tests__/mcp-project-cache.test.ts index 93513627ec..4c07646e4e 100644 --- a/__tests__/mcp-project-cache.test.ts +++ b/__tests__/mcp-project-cache.test.ts @@ -66,3 +66,22 @@ it('trims idle projects while another project call remains active', async () => cache.trimProjects(); expect(close).toHaveBeenCalledOnce(); }); + +it('waits for a pinned read before finishing closeAll even with an empty cache', async () => { + const root = fs.mkdtempSync(path.join(os.tmpdir(), 'codegraph-cache-close-')); + roots.push(root); + CodeGraph.initSync(root).close(); + handler = new ToolHandler(null); + let finish!: (result: { content: [{ type: 'text'; text: string }] }) => void; + const pending = new Promise<{ content: [{ type: 'text'; text: string }] }>(resolve => { finish = resolve; }); + vi.spyOn(handler as unknown as { dispatchTool(): typeof pending }, 'dispatchTool').mockReturnValue(pending); + const activeCall = handler.executeReadTool('codegraph_search', { projectPath: root }); + let closed = false; + const closing = handler.closeAll().then(() => { closed = true; }); + await Promise.resolve(); + expect(closed).toBe(false); + finish({ content: [{ type: 'text', text: 'done' }] }); + await activeCall; + await closing; + expect(closed).toBe(true); +}); diff --git a/src/mcp/tools.ts b/src/mcp/tools.ts index 1202d9a395..d8ee38b56b 100644 --- a/src/mcp/tools.ts +++ b/src/mcp/tools.ts @@ -2377,14 +2377,11 @@ export class ToolHandler { await this.awaitCatchUpGate(gate); } - /** - * Never evict a graph while a tool call or its timed-out reconcile uses it. - * Evicts over the LRU bound, on close, and once idle past the timeout - * (#2087). The cache is in last-use order, so idle entries lead it. - */ private pinProject(projectPath: unknown): string | null { - const root = typeof projectPath === 'string' ? findNearestCodeGraphRoot(projectPath) : null; - if (root) this.activeProjectRoots.set(root, (this.activeProjectRoots.get(root) ?? 0) + 1); + const resolved = typeof projectPath === 'string' ? findNearestCodeGraphRoot(projectPath) : null; + if (!resolved) return null; + const root = canonicalPath(resolved); + this.activeProjectRoots.set(root, (this.activeProjectRoots.get(root) ?? 0) + 1); return root; } @@ -2395,6 +2392,11 @@ export class ToolHandler { else this.activeProjectRoots.set(root, count - 1); } + /** + * Never evict a graph while a tool call or its timed-out reconcile uses it. + * Evicts over the LRU bound, on close, and once idle past the timeout + * (#2087). The cache is in last-use order, so idle entries lead it. + */ private trimProjects(): void { if (this.activeCalls > 0 && (this.closing || this.activeProjectRoots.size === 0)) return; const idleMs = resolveProjectIdleTimeoutMs(); @@ -2413,7 +2415,7 @@ export class ToolHandler { }); } else cg.close(); } - if (this.closing && this.projectCache.size === 0 && this.pendingCloses === 0) { + if (this.closing && this.projectCache.size === 0 && this.activeCalls === 0 && this.activeProjectRoots.size === 0 && this.pendingCloses === 0) { for (const resolve of this.closeWaiters.splice(0)) resolve(); } this.scheduleIdleRelease(idleMs); @@ -2474,7 +2476,7 @@ export class ToolHandler { this.idleReleaseTimer = null; this.nestedRepoCache.clear(); this.trimProjects(); - if (this.projectCache.size === 0 && this.activeCalls === 0 && this.pendingCloses === 0) return Promise.resolve(); + if (this.projectCache.size === 0 && this.activeCalls === 0 && this.activeProjectRoots.size === 0 && this.pendingCloses === 0) return Promise.resolve(); return new Promise((resolve) => this.closeWaiters.push(resolve)); } @@ -2802,6 +2804,7 @@ export class ToolHandler { if (typeof pathCheck === 'object' && pathCheck !== undefined) { return pathCheck; } + pinnedRoot = this.pinProject(pathCheck); // An explicit project gets the same first-call guarantee as the default // (#1835): its post-open catch-up sync finishes (time-boxed) before we // serve it. Resolved on the main thread so the watcher lives here even @@ -2809,7 +2812,6 @@ export class ToolHandler { if (typeof pathCheck === 'string') { await this.awaitProjectGate(pathCheck); } - pinnedRoot = this.pinProject(pathCheck); // The `path` and `pattern` properties used by codegraph_files are // also path-shaped — apply the same cap. if (args.path !== undefined) {