From 6532ca6960f903d36b7d908e12ac294fe2f8bfd3 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sat, 15 Aug 2026 00:49:13 -0400 Subject: [PATCH 01/37] fix: bound and clean up MCP session resources --- src/agent-tools.ts | 22 ++---- src/agent/query-index/sessionStore.ts | 19 +++--- src/agent/search.ts | 32 +++++++-- src/agent/session.ts | 9 +++ src/cli/viewer.ts | 96 +++++++++++++++------------ src/mcp/server.ts | 44 ++++++------ src/session.ts | 42 ++++++++++++ src/sqlite/query.ts | 7 +- tests/session.test.ts | 19 ++++++ 9 files changed, 192 insertions(+), 98 deletions(-) diff --git a/src/agent-tools.ts b/src/agent-tools.ts index fac25c6a..086f0f21 100644 --- a/src/agent-tools.ts +++ b/src/agent-tools.ts @@ -23,7 +23,7 @@ import { errorMessage } from "./util/errors.js"; import { listProjectFiles } from "./util/projectFiles.js"; import { boundAgentList, defaultAgentLimit, normalizeAgentLimit } from "./agent/bounds.js"; import { normalizeAgentOutputPath } from "./agent/normalize.js"; -import type { AgentSession } from "./agent/session.js"; +import { assertNoPrebuiltSessionWithBuildOptions, type AgentSession } from "./agent/session.js"; import { workspaceSymbols, workspaceSymbolsWithSession, @@ -303,9 +303,7 @@ export async function tool_workspaceSymbols( request: WorkspaceSymbolsRequest, runtimeOptions: ToolWorkspaceSymbolsRuntimeOptions = {}, ): Promise { - if (runtimeOptions.session && runtimeOptions.buildOptions) { - throw new Error("Workspace symbol tool options cannot combine a prebuilt session with buildOptions."); - } + assertNoPrebuiltSessionWithBuildOptions(runtimeOptions, "Workspace symbol tool options"); const agentRequest = { root, ...request, @@ -360,9 +358,7 @@ export async function tool_findImplementations( } function assertTypeHierarchyToolOptions(runtimeOptions: ToolTypeHierarchyRuntimeOptions): void { - if (runtimeOptions.session && runtimeOptions.buildOptions) { - throw new Error("Type hierarchy tool options cannot combine a prebuilt session with buildOptions."); - } + assertNoPrebuiltSessionWithBuildOptions(runtimeOptions, "Type hierarchy tool options"); } function typeHierarchyAgentRequest( @@ -408,9 +404,7 @@ export async function tool_findCallees( } function assertCallHierarchyToolOptions(runtimeOptions: ToolCallHierarchyRuntimeOptions): void { - if (runtimeOptions.session && runtimeOptions.buildOptions) { - throw new Error("Call hierarchy tool options cannot combine a prebuilt session with buildOptions."); - } + assertNoPrebuiltSessionWithBuildOptions(runtimeOptions, "Call hierarchy tool options"); } function callHierarchyAgentRequest( @@ -441,9 +435,7 @@ export async function tool_previewRename( request: Omit, runtimeOptions: ToolRenamePreviewRuntimeOptions = {}, ): Promise { - if (runtimeOptions.session && runtimeOptions.buildOptions) { - throw new Error("Rename preview tool options cannot combine a prebuilt session with buildOptions."); - } + assertNoPrebuiltSessionWithBuildOptions(runtimeOptions, "Rename preview tool options"); const agentRequest: RenamePreviewRequest = { root, ...request, @@ -469,9 +461,7 @@ export async function tool_buildRefactorPlan( request: Omit, runtimeOptions: ToolRefactorPlanRuntimeOptions = {}, ): Promise { - if (runtimeOptions.session && runtimeOptions.buildOptions) { - throw new Error("Refactor plan tool options cannot combine a prebuilt session with buildOptions."); - } + assertNoPrebuiltSessionWithBuildOptions(runtimeOptions, "Refactor plan tool options"); const agentRequest: RefactorPlanRequest = { root, ...request, diff --git a/src/agent/query-index/sessionStore.ts b/src/agent/query-index/sessionStore.ts index c5800a05..6679f29f 100644 --- a/src/agent/query-index/sessionStore.ts +++ b/src/agent/query-index/sessionStore.ts @@ -16,11 +16,11 @@ function closeHandle(handle: QueryIndexHandle): void { } function closeState(state: SessionQueryIndexState): void { + state.closing = true; if (state.resolved) { closeHandle(state.resolved); return; } - state.closing = true; void state.handle.then(closeHandle, () => undefined); } @@ -30,7 +30,10 @@ export async function ensureSessionQueryIndex( ): Promise { const identity = snapshot.index.projectSnapshotIdentity ?? ""; const existing = QUERY_INDEX_BY_SESSION.get(session); - if (existing?.identity === identity) return await existing.handle; + if (existing?.identity === identity && !existing.closing) { + const resolved = await existing.handle; + if (!existing.closing) return resolved; + } if (existing) { QUERY_INDEX_BY_SESSION.delete(session); closeState(existing); @@ -44,14 +47,14 @@ export async function ensureSessionQueryIndex( if (QUERY_INDEX_BY_SESSION.get(session) === state) QUERY_INDEX_BY_SESSION.delete(session); }); const resolved = await handle; - if (QUERY_INDEX_BY_SESSION.get(session) === state) { + if (QUERY_INDEX_BY_SESSION.get(session) === state && !state.closing) { state.resolved = resolved; - } else if (!state.closing) { - closeHandle(resolved); + if (snapshot.buildReport) snapshot.buildReport.queryIndex = resolved.diagnostics; + if (snapshot.index.buildReport) snapshot.index.buildReport.queryIndex = resolved.diagnostics; + return resolved; } - if (snapshot.buildReport) snapshot.buildReport.queryIndex = resolved.diagnostics; - if (snapshot.index.buildReport) snapshot.index.buildReport.queryIndex = resolved.diagnostics; - return resolved; + closeHandle(resolved); + return await ensureSessionQueryIndex(session, snapshot); } export function disposeSessionQueryIndex(session: AgentSession): void { diff --git a/src/agent/search.ts b/src/agent/search.ts index e221479d..4f6e5335 100644 --- a/src/agent/search.ts +++ b/src/agent/search.ts @@ -207,7 +207,8 @@ const NATURAL_LANGUAGE_SYNTAX_TERMS = new Set([ "or", ]); const SEARCH_CACHES = new WeakMap(); -const SEARCH_RESULT_CACHES = new WeakMap>>(); +export const DEFAULT_SESSION_SEARCH_CACHE_MAX_ENTRIES = 100; +const SEARCH_RESULT_CACHES = new WeakMap(); const SEARCH_RANKING_VERSION = 2; export async function searchCodegraph(request: AgentSearchRequest): Promise { @@ -237,8 +238,9 @@ export async function searchCodegraphWithSession( }); const resultCache = getSessionSearchResultCache(session); const cacheKey = searchResultCacheKey(snapshot, request); - const existing = resultCache.get(cacheKey); + const existing = resultCache.entries.get(cacheKey); if (existing) { + promoteSessionSearchResult(resultCache, cacheKey, existing); const response = await existing; if (response.query === request.query) return response; return { ...response, query: request.query }; @@ -250,9 +252,9 @@ export async function searchCodegraphWithSession( queryIndex = await ensureSessionQueryIndex(session, snapshot); } const search = searchSnapshot(snapshot, request, queryIndex); - resultCache.set(cacheKey, search); + promoteSessionSearchResult(resultCache, cacheKey, search); search.catch(() => { - if (resultCache.get(cacheKey) === search) resultCache.delete(cacheKey); + if (resultCache.entries.get(cacheKey) === search) resultCache.entries.delete(cacheKey); }); return await search; } @@ -348,15 +350,33 @@ function canUsePathFastPath(request: AgentSearchRequest): boolean { return (request.mode ?? "hybrid") === "path" && request.from === undefined; } -function getSessionSearchResultCache(session: AgentSession): Map> { +type SessionSearchResultCache = { + entries: Map>; +}; + +function getSessionSearchResultCache(session: AgentSession): SessionSearchResultCache { const existing = SEARCH_RESULT_CACHES.get(session); if (existing) return existing; - const created = new Map>(); + const created: SessionSearchResultCache = { entries: new Map() }; SEARCH_RESULT_CACHES.set(session, created); registerSessionInvalidationHook(session, () => SEARCH_RESULT_CACHES.delete(session)); return created; } +function promoteSessionSearchResult( + cache: SessionSearchResultCache, + key: string, + result: Promise, +): void { + cache.entries.delete(key); + cache.entries.set(key, result); + while (cache.entries.size > DEFAULT_SESSION_SEARCH_CACHE_MAX_ENTRIES) { + const oldest = cache.entries.keys().next().value; + if (oldest === undefined) return; + cache.entries.delete(oldest); + } +} + function searchResultCacheKey(snapshot: AgentProjectSnapshot, request: AgentSearchRequest): string { const { rankTokens, normalizedRankPhrase, identifierLike } = buildQueryTerms(request.query, snapshot); return JSON.stringify({ diff --git a/src/agent/session.ts b/src/agent/session.ts index e47a8427..cbe2884a 100644 --- a/src/agent/session.ts +++ b/src/agent/session.ts @@ -76,6 +76,15 @@ export type AgentSession = { invalidate: () => void; }; +export function assertNoPrebuiltSessionWithBuildOptions( + options: { session?: AgentSession; buildOptions?: BuildOptions }, + consumer: string, +): void { + if (options.session && options.buildOptions) { + throw new Error(`${consumer} cannot combine a prebuilt session with buildOptions.`); + } +} + const EMPTY_SYMBOL_GRAPH: SymbolGraph = { nodes: new Map(), edges: [], diff --git a/src/cli/viewer.ts b/src/cli/viewer.ts index 34a3de6a..32091466 100644 --- a/src/cli/viewer.ts +++ b/src/cli/viewer.ts @@ -201,64 +201,74 @@ async function writeGeneratedGraphResponse( } } +function writeViewerRequestError(response: http.ServerResponse): void { + if (response.writableEnded) return; + if (!response.headersSent) response.writeHead(500); + response.end(); +} + function viewerRequestHandler( options: ResolvedViewerOptions, getAllowedHostHeaders: () => AllowedHostHeaderRules, ): http.RequestListener { return (request, response) => { - if (!isAllowedHostHeader(request, getAllowedHostHeaders())) { - response.writeHead(403); - response.end(); - return; - } - if (request.method !== "GET" && request.method !== "HEAD") { - response.writeHead(405, { Allow: "GET, HEAD" }); - response.end(); - return; - } - - const rawPathname = (request.url ?? "/").split(/[?#]/, 1)[0] ?? "/"; try { - if ( - decodeURIComponent(rawPathname) - .split("/") - .some((segment) => segment === "..") - ) { - response.writeHead(404); + if (!isAllowedHostHeader(request, getAllowedHostHeaders())) { + response.writeHead(403); response.end(); return; } - } catch { - response.writeHead(404); - response.end(); - return; - } - - const pathname = new URL(request.url ?? "/", "http://viewer.local").pathname; - if (pathname === "/graph.json") { - if (options.graphFile) { - writeFileResponse( - request, - response, - options.graphFile.path, - "application/json; charset=utf-8", - options.graphFile.fileDescriptor, - ); + if (request.method !== "GET" && request.method !== "HEAD") { + response.writeHead(405, { Allow: "GET, HEAD" }); + response.end(); return; } - if (options.graphProvider) { - void writeGeneratedGraphResponse(request, response, options.graphProvider); + + const rawPathname = (request.url ?? "/").split(/[?#]/, 1)[0] ?? "/"; + try { + if ( + decodeURIComponent(rawPathname) + .split("/") + .some((segment) => segment === "..") + ) { + response.writeHead(404); + response.end(); + return; + } + } catch { + response.writeHead(404); + response.end(); return; } - } - const asset = VIEWER_ASSETS[pathname]; - if (!asset) { - response.writeHead(404); - response.end(); - return; + const pathname = new URL(request.url ?? "/", "http://viewer.local").pathname; + if (pathname === "/graph.json") { + if (options.graphFile) { + writeFileResponse( + request, + response, + options.graphFile.path, + "application/json; charset=utf-8", + options.graphFile.fileDescriptor, + ); + return; + } + if (options.graphProvider) { + void writeGeneratedGraphResponse(request, response, options.graphProvider); + return; + } + } + + const asset = VIEWER_ASSETS[pathname]; + if (!asset) { + response.writeHead(404); + response.end(); + return; + } + writeFileResponse(request, response, path.join(options.assetRoot, asset.file), asset.contentType); + } catch { + writeViewerRequestError(response); } - writeFileResponse(request, response, path.join(options.assetRoot, asset.file), asset.contentType); }; } diff --git a/src/mcp/server.ts b/src/mcp/server.ts index 04891138..0aa4e7dd 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -58,7 +58,11 @@ import { SQLITE_ARTIFACT_FILE_SIGNATURES_METADATA_KEY, queryGraphSqliteRaw, type import { isPlainRecord } from "../util/guards.js"; import { toProjectDisplayPath } from "../util/paths.js"; import { errorMessage } from "../util/errors.js"; -import { createAgentSession, listAgentSessionFiles } from "../agent/session.js"; +import { + assertNoPrebuiltSessionWithBuildOptions, + createAgentSession, + listAgentSessionFiles, +} from "../agent/session.js"; import { mapLimit } from "../util/concurrency.js"; import { assertRealPathCandidateWithinRoot, resolveProjectFile } from "../util/confinedFile.js"; import type { AgentFreshnessResult, AgentProjectSnapshot, AgentSession } from "../agent/session.js"; @@ -354,9 +358,7 @@ const MAX_MCP_FRESHNESS_CHANGED_FILES = 25; const SQLITE_ARTIFACT_STAT_CONCURRENCY = 64; function assertMcpSessionOptions(options: CodegraphMcpHandlerOptions): void { - if (options.session !== undefined && options.buildOptions !== undefined) { - throw new Error("MCP server options cannot combine a prebuilt session with buildOptions."); - } + assertNoPrebuiltSessionWithBuildOptions(options, "MCP server options"); } function createCodegraphMcpSession(options: CodegraphMcpHandlerOptions, root: string): AgentSession { @@ -384,7 +386,14 @@ function startCodegraphMcpWarmup( return undefined; } -async function createWarmedCodegraphMcpHandlers(options: CodegraphMcpServerOptions): Promise { +type WarmedCodegraphMcpResources = { + handlers: CodegraphMcpHandlers; + session: AgentSession; +}; + +async function createWarmedCodegraphMcpResources( + options: CodegraphMcpServerOptions, +): Promise { const root = path.resolve(options.root); const session = createCodegraphMcpSession(options, root); await startCodegraphMcpWarmup(session, options.warmup); @@ -393,7 +402,10 @@ async function createWarmedCodegraphMcpHandlers(options: CodegraphMcpServerOptio void host; void port; void onHttpListen; - return createCodegraphMcpHandlersForSession({ ...handlerOptions, root }, session); + return { + handlers: createCodegraphMcpHandlersForSession({ ...handlerOptions, root }, session), + session, + }; } const MCP_HTTP_PATH = "/mcp"; @@ -1118,7 +1130,7 @@ export async function serveCodegraphMcp(options: CodegraphMcpServerOptions): Pro return; } - const handlers = await createWarmedCodegraphMcpHandlers(options); + const { handlers, session } = await createWarmedCodegraphMcpResources(options); const runtimeIdentity = options.runtimeIdentity ?? captureCodegraphRuntimeIdentity(getCurrentNativeBindingOrigin()); const createProtocolServer = createCodegraphMcpProtocolFactory(handlers, runtimeIdentity); const handle = serveStdio(createProtocolServer, { @@ -1134,6 +1146,7 @@ export async function serveCodegraphMcp(options: CodegraphMcpServerOptions): Pro console.error(`[codegraph] MCP stdio shutting down (${shutdownReason})`); }, }); + session.invalidate(); // Ensure orphaned stdio servers do not linger after the client is gone. process.exitCode = 0; } @@ -1142,7 +1155,7 @@ export async function startCodegraphMcpHttpServer( options: CodegraphMcpServerOptions & { port: number }, ): Promise { const host = options.host ?? "127.0.0.1"; - const handlers = await createWarmedCodegraphMcpHandlers(options); + const { handlers, session } = await createWarmedCodegraphMcpResources(options); const runtimeIdentity = options.runtimeIdentity ?? captureCodegraphRuntimeIdentity(getCurrentNativeBindingOrigin()); const createProtocolServer = createCodegraphMcpProtocolFactory(handlers, runtimeIdentity); const sessionStore = createLegacyMcpSessionStore({ @@ -1166,6 +1179,7 @@ export async function startCodegraphMcpHttpServer( let closeResourcesPromise: Promise | undefined; const closeResources = (): Promise => { closeResourcesPromise ??= (async () => { + session.invalidate(); sessionStore.stop(); await closeMcpResources(sessionStore.sessions, modernHandler.close); })(); @@ -1285,12 +1299,7 @@ async function handleLegacyMcpHttpPost( return; } sessionStore.touch(sessionId); - try { - await handleLegacyMcpSessionRequest(session, request, response, body); - } catch (error) { - await sessionStore.delete(sessionId); - throw error; - } + await handleLegacyMcpSessionRequest(session, request, response, body); return; } @@ -1381,12 +1390,7 @@ async function handleExistingMcpSessionRequest( return; } sessionStore.touch(sessionId); - try { - await handleLegacyMcpSessionRequest(session, request, response); - } catch (error) { - await sessionStore.delete(sessionId); - throw error; - } + await handleLegacyMcpSessionRequest(session, request, response); } async function handleLegacyMcpSessionRequest( diff --git a/src/session.ts b/src/session.ts index e234269d..cff61702 100644 --- a/src/session.ts +++ b/src/session.ts @@ -47,6 +47,16 @@ export type SessionOptions = { incremental?: boolean; }; +export type SessionManagerOptions = { + /** Maximum sessions, including sessions currently initializing. Defaults to 32. */ + maxSessions?: number; + /** Idle-session scan interval in milliseconds. Defaults to 60 seconds. Use 0 to disable. */ + evictionIntervalMs?: number; +}; + +export const DEFAULT_SESSION_MANAGER_MAX_SESSIONS = 32; +export const DEFAULT_SESSION_MANAGER_EVICTION_INTERVAL_MS = 60_000; + export type SessionStatus = "initializing" | "ready" | "expired" | "error"; export type SessionStaleReason = "tracked_files_changed" | "config_changed"; @@ -787,6 +797,16 @@ export class CodeReviewSession implements ICodeReviewSession { } } +function normalizeSessionManagerCapacity(value: number | undefined): number { + if (value === undefined) return DEFAULT_SESSION_MANAGER_MAX_SESSIONS; + return Math.max(1, Math.floor(value)); +} + +function normalizeSessionManagerEvictionInterval(value: number | undefined): number { + if (value === undefined) return DEFAULT_SESSION_MANAGER_EVICTION_INTERVAL_MS; + return Math.max(0, Math.floor(value)); +} + /** * Session manager for multiple concurrent sessions * Useful for agents handling multiple repositories or PRs @@ -802,6 +822,26 @@ export class SessionManager { promise: Promise; } >(); + private readonly maxSessions: number; + private readonly evictionTimer: ReturnType | undefined; + + constructor(options: SessionManagerOptions = {}) { + this.maxSessions = normalizeSessionManagerCapacity(options.maxSessions); + const evictionIntervalMs = normalizeSessionManagerEvictionInterval(options.evictionIntervalMs); + if (evictionIntervalMs) { + this.evictionTimer = setInterval(() => this.cleanupExpired(), evictionIntervalMs); + this.evictionTimer.unref?.(); + } + } + + private assertCapacityForNewSession(): void { + this.cleanupExpired(); + if (this.sessions.size + this.pendingSessions.size >= this.maxSessions) { + throw new Error( + `Session capacity reached (${this.maxSessions}). Dispose an existing session before creating another.`, + ); + } + } private createSessionConfigurationError( sessionId: string, @@ -892,6 +932,7 @@ export class SessionManager { let session = this.ensureSessionIdCompatible(sessionId, options); if (!session) { + this.assertCapacityForNewSession(); session = new CodeReviewSession(options); return await this.trackSession(sessionId, options, session, false, (readySession) => { this.sessions.set(sessionId, readySession); @@ -944,6 +985,7 @@ export class SessionManager { session.dispose(); } this.sessions.clear(); + clearInterval(this.evictionTimer); } /** diff --git a/src/sqlite/query.ts b/src/sqlite/query.ts index 39854987..d1e740ff 100644 --- a/src/sqlite/query.ts +++ b/src/sqlite/query.ts @@ -6,6 +6,7 @@ import { DEFAULT_SQLITE_BYTE_LIMIT, MAX_SQLITE_CELL_BYTES, MAX_SQLITE_ROW_LIMIT, + normalizeSqliteRowLimit, } from "./rowBounds.js"; export { queryGraphSqlite } from "./canned-query.js"; @@ -27,11 +28,7 @@ export async function queryGraphSqliteRaw( const stmt = db.prepare(sql); assertReadOnlyQueryStatement(stmt); const columns = stmt.columns().map((col) => col.name); - const requestedRows = options?.maxRows; - const maxRows = - requestedRows === undefined - ? MAX_SQLITE_ROW_LIMIT - : Math.min(MAX_SQLITE_ROW_LIMIT, Math.max(0, Math.floor(requestedRows))); + const maxRows = normalizeSqliteRowLimit(options?.maxRows ?? MAX_SQLITE_ROW_LIMIT); const maxBytes = options?.maxBytes ?? DEFAULT_SQLITE_BYTE_LIMIT; const maxCellBytes = options?.maxCellBytes ?? MAX_SQLITE_CELL_BYTES; diff --git a/tests/session.test.ts b/tests/session.test.ts index 510cf11d..9a31b4e1 100644 --- a/tests/session.test.ts +++ b/tests/session.test.ts @@ -1215,6 +1215,25 @@ describe("SessionManager", () => { expect(session1).toBe(session2); }); + test("rejects a new session when configured capacity is exhausted", async () => { + const limitedManager = new SessionManager({ maxSessions: 1, evictionIntervalMs: 0 }); + try { + await limitedManager.getOrCreateSession("first", { + root: sampleRoot, + buildOptions: sampleBuildOptions(), + }); + + await expect( + limitedManager.getOrCreateSession("second", { + root: sampleRoot, + buildOptions: sampleBuildOptions(), + }), + ).rejects.toThrow("Session capacity reached (1)"); + } finally { + limitedManager.disposeAll(); + } + }); + test("should share one initialization across concurrent same-id creation", async () => { const buildSpy = vi.spyOn(indexerBuild, "buildProjectIndexIncremental"); From bba947c2fad78cf88d21a6ca26e8356ee6d04f3d Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sat, 15 Aug 2026 00:50:19 -0400 Subject: [PATCH 02/37] fix: coalesce MCP index refreshes --- src/mcp/server.ts | 26 +++++++++++++++++++++----- 1 file changed, 21 insertions(+), 5 deletions(-) diff --git a/src/mcp/server.ts b/src/mcp/server.ts index 0aa4e7dd..1d316262 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -432,6 +432,8 @@ function createCodegraphMcpHandlersForSession( let sqlitePath = configuredSqlitePath; let sqliteOutDir = configuredSqliteOutDir; let sqliteCanRefresh = configuredSqliteCanRefresh; + let refreshPromise: Promise | undefined; + let refreshEpoch = 0; const relative = (file: string): string => toProjectDisplayPath(root, file); const boundedLimit = (limit: number | undefined, fallback: number, max: number): number => { @@ -981,11 +983,25 @@ function createCodegraphMcpHandlersForSession( refresh_index: async (request) => { const warmup = request.warmup ?? "off"; - session.invalidate(); - sqlitePath = configuredSqlitePath; - sqliteOutDir = configuredSqliteOutDir; - sqliteCanRefresh = configuredSqliteCanRefresh; - await startCodegraphMcpWarmup(session, warmup); + if (refreshPromise) { + await refreshPromise; + return { refreshed: true, warmup }; + } + const epoch = ++refreshEpoch; + const refresh = (async () => { + session.invalidate(); + sqlitePath = configuredSqlitePath; + sqliteOutDir = configuredSqliteOutDir; + sqliteCanRefresh = configuredSqliteCanRefresh; + await startCodegraphMcpWarmup(session, warmup); + if (epoch !== refreshEpoch) throw new Error("MCP index refresh was superseded."); + })(); + refreshPromise = refresh; + try { + await refresh; + } finally { + if (refreshPromise === refresh) refreshPromise = undefined; + } return { refreshed: true, warmup }; }, From 4293dfded36864e7783cbd1ca8f72d199926367b Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sat, 15 Aug 2026 00:55:56 -0400 Subject: [PATCH 03/37] fix: bound query index generation retries --- src/agent/query-index/sessionStore.ts | 53 +++++++++++++++------------ src/agent/search.ts | 29 +++++++-------- 2 files changed, 42 insertions(+), 40 deletions(-) diff --git a/src/agent/query-index/sessionStore.ts b/src/agent/query-index/sessionStore.ts index 6679f29f..9abe093f 100644 --- a/src/agent/query-index/sessionStore.ts +++ b/src/agent/query-index/sessionStore.ts @@ -10,6 +10,7 @@ type SessionQueryIndexState = { }; const QUERY_INDEX_BY_SESSION = new WeakMap(); +const MAX_QUERY_INDEX_GENERATION_RETRIES = 3; function closeHandle(handle: QueryIndexHandle): void { handle.store?.close(); @@ -29,32 +30,36 @@ export async function ensureSessionQueryIndex( snapshot: AgentProjectSnapshot, ): Promise { const identity = snapshot.index.projectSnapshotIdentity ?? ""; - const existing = QUERY_INDEX_BY_SESSION.get(session); - if (existing?.identity === identity && !existing.closing) { - const resolved = await existing.handle; - if (!existing.closing) return resolved; - } - if (existing) { - QUERY_INDEX_BY_SESSION.delete(session); - closeState(existing); - } + for (let attempt = 0; attempt < MAX_QUERY_INDEX_GENERATION_RETRIES; attempt += 1) { + const existing = QUERY_INDEX_BY_SESSION.get(session); + if (existing?.identity === identity && !existing.closing) { + const resolved = await existing.handle; + if (!existing.closing) return resolved; + } + if (existing) { + QUERY_INDEX_BY_SESSION.delete(session); + closeState(existing); + } - const handle = ensureQueryIndex(snapshot); - const state: SessionQueryIndexState = { identity, handle }; - if (!existing) registerSessionInvalidationHook(session, () => disposeSessionQueryIndex(session)); - QUERY_INDEX_BY_SESSION.set(session, state); - handle.catch(() => { - if (QUERY_INDEX_BY_SESSION.get(session) === state) QUERY_INDEX_BY_SESSION.delete(session); - }); - const resolved = await handle; - if (QUERY_INDEX_BY_SESSION.get(session) === state && !state.closing) { - state.resolved = resolved; - if (snapshot.buildReport) snapshot.buildReport.queryIndex = resolved.diagnostics; - if (snapshot.index.buildReport) snapshot.index.buildReport.queryIndex = resolved.diagnostics; - return resolved; + const handle = ensureQueryIndex(snapshot); + const state: SessionQueryIndexState = { identity, handle }; + if (!existing) registerSessionInvalidationHook(session, () => disposeSessionQueryIndex(session)); + QUERY_INDEX_BY_SESSION.set(session, state); + handle.catch(() => { + if (QUERY_INDEX_BY_SESSION.get(session) === state) QUERY_INDEX_BY_SESSION.delete(session); + }); + const resolved = await handle; + if (QUERY_INDEX_BY_SESSION.get(session) === state && !state.closing) { + state.resolved = resolved; + if (snapshot.buildReport) snapshot.buildReport.queryIndex = resolved.diagnostics; + if (snapshot.index.buildReport) snapshot.index.buildReport.queryIndex = resolved.diagnostics; + return resolved; + } + closeHandle(resolved); } - closeHandle(resolved); - return await ensureSessionQueryIndex(session, snapshot); + throw new Error( + "Query index generation changed repeatedly while loading; retry the request after refresh completes.", + ); } export function disposeSessionQueryIndex(session: AgentSession): void { diff --git a/src/agent/search.ts b/src/agent/search.ts index 4f6e5335..8f4856ce 100644 --- a/src/agent/search.ts +++ b/src/agent/search.ts @@ -207,8 +207,9 @@ const NATURAL_LANGUAGE_SYNTAX_TERMS = new Set([ "or", ]); const SEARCH_CACHES = new WeakMap(); -export const DEFAULT_SESSION_SEARCH_CACHE_MAX_ENTRIES = 100; -const SEARCH_RESULT_CACHES = new WeakMap(); +// A fixed ceiling prevents client-supplied query keys retaining an unbounded session heap. +const SESSION_SEARCH_CACHE_MAX_ENTRIES = 100; +const SEARCH_RESULT_CACHES = new WeakMap>>(); const SEARCH_RANKING_VERSION = 2; export async function searchCodegraph(request: AgentSearchRequest): Promise { @@ -238,7 +239,7 @@ export async function searchCodegraphWithSession( }); const resultCache = getSessionSearchResultCache(session); const cacheKey = searchResultCacheKey(snapshot, request); - const existing = resultCache.entries.get(cacheKey); + const existing = resultCache.get(cacheKey); if (existing) { promoteSessionSearchResult(resultCache, cacheKey, existing); const response = await existing; @@ -254,7 +255,7 @@ export async function searchCodegraphWithSession( const search = searchSnapshot(snapshot, request, queryIndex); promoteSessionSearchResult(resultCache, cacheKey, search); search.catch(() => { - if (resultCache.entries.get(cacheKey) === search) resultCache.entries.delete(cacheKey); + if (resultCache.get(cacheKey) === search) resultCache.delete(cacheKey); }); return await search; } @@ -350,30 +351,26 @@ function canUsePathFastPath(request: AgentSearchRequest): boolean { return (request.mode ?? "hybrid") === "path" && request.from === undefined; } -type SessionSearchResultCache = { - entries: Map>; -}; - -function getSessionSearchResultCache(session: AgentSession): SessionSearchResultCache { +function getSessionSearchResultCache(session: AgentSession): Map> { const existing = SEARCH_RESULT_CACHES.get(session); if (existing) return existing; - const created: SessionSearchResultCache = { entries: new Map() }; + const created = new Map>(); SEARCH_RESULT_CACHES.set(session, created); registerSessionInvalidationHook(session, () => SEARCH_RESULT_CACHES.delete(session)); return created; } function promoteSessionSearchResult( - cache: SessionSearchResultCache, + cache: Map>, key: string, result: Promise, ): void { - cache.entries.delete(key); - cache.entries.set(key, result); - while (cache.entries.size > DEFAULT_SESSION_SEARCH_CACHE_MAX_ENTRIES) { - const oldest = cache.entries.keys().next().value; + cache.delete(key); + cache.set(key, result); + while (cache.size > SESSION_SEARCH_CACHE_MAX_ENTRIES) { + const oldest = cache.keys().next().value; if (oldest === undefined) return; - cache.entries.delete(oldest); + cache.delete(oldest); } } From 8edc7004fe9ede25141fe783b0cb33837a767fc9 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sat, 15 Aug 2026 00:57:33 -0400 Subject: [PATCH 04/37] fix: bound MCP HTTP body reads --- src/mcp/http.ts | 39 +++++++++++++++++++++++++++++++-------- src/mcp/server.ts | 11 ++++++++++- 2 files changed, 41 insertions(+), 9 deletions(-) diff --git a/src/mcp/http.ts b/src/mcp/http.ts index f55f0ec8..633963ca 100644 --- a/src/mcp/http.ts +++ b/src/mcp/http.ts @@ -2,7 +2,11 @@ import type { IncomingMessage, Server as HttpServer, ServerResponse } from "node import type { AddressInfo } from "node:net"; import os from "node:os"; -export type ParsedJsonBody = { status: "ok"; body: unknown } | { status: "too_large" } | { status: "invalid_json" }; +export type ParsedJsonBody = + | { status: "ok"; body: unknown } + | { status: "too_large" } + | { status: "timeout" } + | { status: "invalid_json" }; export type AllowedHostHeaderRules = { exact: Set; @@ -13,7 +17,11 @@ export function getRequestPath(request: IncomingMessage): string { return new URL(request.url ?? "/", "http://127.0.0.1").pathname; } -export async function readJsonRequestBody(request: IncomingMessage, maxBytes: number): Promise { +export async function readJsonRequestBody( + request: IncomingMessage, + maxBytes: number, + timeoutMs: number, +): Promise { const contentLength = getContentLength(request); if (contentLength !== undefined && contentLength > maxBytes) { request.resume(); @@ -22,14 +30,29 @@ export async function readJsonRequestBody(request: IncomingMessage, maxBytes: nu const chunks: Buffer[] = []; let bytes = 0; - for await (const chunk of request) { - const buffer = typeof chunk === "string" ? Buffer.from(chunk) : chunk; - bytes += buffer.byteLength; - if (bytes > maxBytes) { - return { status: "too_large" }; + let timedOut = false; + const deadline = setTimeout(() => { + timedOut = true; + request.destroy(); + }, timeoutMs); + deadline.unref?.(); + try { + for await (const chunk of request) { + const buffer = typeof chunk === "string" ? Buffer.from(chunk) : chunk; + bytes += buffer.byteLength; + if (bytes > maxBytes) { + request.resume(); + return { status: "too_large" }; + } + chunks.push(buffer); } - chunks.push(buffer); + } catch { + if (timedOut) return { status: "timeout" }; + return { status: "invalid_json" }; + } finally { + clearTimeout(deadline); } + if (timedOut) return { status: "timeout" }; const rawBody = Buffer.concat(chunks).toString("utf8"); try { diff --git a/src/mcp/server.ts b/src/mcp/server.ts index 1d316262..012828ce 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -138,6 +138,8 @@ export type CodegraphMcpServerOptions = CodegraphMcpHandlerOptions & { httpSessionMaxCount?: number; /** How often to scan for idle HTTP sessions in ms. Defaults to 60 seconds. */ httpSessionEvictionIntervalMs?: number; + /** Maximum time to receive an HTTP MCP request body in ms. Defaults to 30 seconds. */ + httpBodyTimeoutMs?: number; onHttpListen?: ((info: CodegraphMcpHttpServerInfo) => void) | undefined; runtimeIdentity?: CodegraphRuntimeIdentity; }; @@ -156,6 +158,7 @@ export type CodegraphMcpHttpServer = CodegraphMcpHttpServerInfo & { export const DEFAULT_MCP_HTTP_SESSION_IDLE_MS = 30 * 60 * 1000; export const DEFAULT_MCP_HTTP_SESSION_MAX_COUNT = 32; export const DEFAULT_MCP_HTTP_SESSION_EVICTION_INTERVAL_MS = 60_000; +export const DEFAULT_MCP_HTTP_BODY_TIMEOUT_MS = 30_000; type LegacyMcpSession = { server: Server; @@ -1211,6 +1214,7 @@ export async function startCodegraphMcpHttpServer( validateOrigin, modernNodeHandler, createProtocolServer, + options.httpBodyTimeoutMs ?? DEFAULT_MCP_HTTP_BODY_TIMEOUT_MS, ); }); @@ -1248,6 +1252,7 @@ async function handleMcpHttpRequest( validateOrigin: OriginValidator, modernNodeHandler: NodeMcpRequestHandler, createProtocolServer: () => Server, + bodyTimeoutMs: number, ): Promise { const requestPath = getRequestPath(request); if (requestPath !== MCP_HTTP_PATH) { @@ -1263,11 +1268,15 @@ async function handleMcpHttpRequest( try { if (request.method === "POST") { - const parsedBody = await readJsonRequestBody(request, MAX_MCP_HTTP_BODY_BYTES); + const parsedBody = await readJsonRequestBody(request, MAX_MCP_HTTP_BODY_BYTES, bodyTimeoutMs); if (parsedBody.status === "too_large") { writeJsonRpcError(response, 413, "MCP request body is too large"); return; } + if (parsedBody.status === "timeout") { + writeJsonRpcError(response, 408, "MCP request body timed out"); + return; + } if (parsedBody.status === "invalid_json") { writeJsonRpcError(response, 400, "Invalid JSON request body"); return; From 52aac765b2a3f3423a47938a744d1c6309aeb664 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sat, 15 Aug 2026 11:51:45 -0400 Subject: [PATCH 05/37] fix: bound MCP tool execution --- src/mcp/server.ts | 302 ++++++++++++++++++++++++++++------------------ src/mcp/tools.ts | 104 ++++++++++++++++ 2 files changed, 287 insertions(+), 119 deletions(-) diff --git a/src/mcp/server.ts b/src/mcp/server.ts index 012828ce..721202b4 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -55,6 +55,7 @@ import { DEFAULT_BOUNDED_IMPACT_BUDGETS } from "../impact/budgets.js"; import { buildReviewReport, type ReviewDepth, type ReviewReport } from "../review.js"; import { boundReviewReportForTransport, type ReviewReportForTransport } from "../review/types.js"; import { SQLITE_ARTIFACT_FILE_SIGNATURES_METADATA_KEY, queryGraphSqliteRaw, type RawSqlResult } from "../sqlite.js"; +import { boundList, countOmitted } from "../presentation/bounds.js"; import { isPlainRecord } from "../util/guards.js"; import { toProjectDisplayPath } from "../util/paths.js"; import { errorMessage } from "../util/errors.js"; @@ -138,6 +139,8 @@ export type CodegraphMcpServerOptions = CodegraphMcpHandlerOptions & { httpSessionMaxCount?: number; /** How often to scan for idle HTTP sessions in ms. Defaults to 60 seconds. */ httpSessionEvictionIntervalMs?: number; + /** Maximum concurrent tool calls per MCP protocol session. Defaults to 4. */ + mcpToolConcurrency?: number; /** Maximum time to receive an HTTP MCP request body in ms. Defaults to 30 seconds. */ httpBodyTimeoutMs?: number; onHttpListen?: ((info: CodegraphMcpHttpServerInfo) => void) | undefined; @@ -159,6 +162,7 @@ export const DEFAULT_MCP_HTTP_SESSION_IDLE_MS = 30 * 60 * 1000; export const DEFAULT_MCP_HTTP_SESSION_MAX_COUNT = 32; export const DEFAULT_MCP_HTTP_SESSION_EVICTION_INTERVAL_MS = 60_000; export const DEFAULT_MCP_HTTP_BODY_TIMEOUT_MS = 30_000; +export const DEFAULT_MCP_TOOL_CONCURRENCY = 4; type LegacyMcpSession = { server: Server; @@ -183,7 +187,13 @@ export type McpTruncationMeta = { omitted: number; }; -export type CodegraphMcpHandlers = { +type McpDependenciesResponse = CodegraphMcpFreshResult< + McpTruncationMeta & { dependencies: Array<{ file: string; depth: number }> } +>; +type McpReverseDependenciesResponse = CodegraphMcpFreshResult< + McpTruncationMeta & { reverseDependencies: Array<{ file: string; depth: number }> } +>; +type CodegraphMcpHandlerDefinitions = { search: (request: { query: string; mode?: AgentSearchMode | undefined; @@ -302,14 +312,12 @@ export type CodegraphMcpHandlers = { file: string; depth?: number | undefined; limit?: number | undefined; - }) => Promise }>>; + }) => Promise; rdeps: (request: { file: string; depth?: number | undefined; limit?: number | undefined; - }) => Promise< - CodegraphMcpFreshResult }> - >; + }) => Promise; path: (request: { from: string; to: string }) => Promise>; impact: (request: { base: string; head: string }) => Promise>; review: (request: { @@ -317,10 +325,9 @@ export type CodegraphMcpHandlers = { head: string; reviewDepth?: ReviewDepth | undefined; }) => Promise>; - refresh_index: (request: { warmup?: CodegraphMcpWarmupMode | undefined }) => Promise<{ - refreshed: true; - warmup: CodegraphMcpWarmupMode; - }>; + refresh_index: (request: { + warmup?: CodegraphMcpWarmupMode | undefined; + }) => Promise<{ refreshed: true; warmup: CodegraphMcpWarmupMode }>; query_sqlite: (request: { query: string; params?: Array | undefined; @@ -336,6 +343,14 @@ export type CodegraphMcpHandlers = { }) => Promise>; }; +type WithAbortSignal = { + [K in keyof T]: T[K] extends (request: infer Request) => Promise + ? (request: Request, signal?: AbortSignal) => Promise + : never; +}; + +export type CodegraphMcpHandlers = WithAbortSignal; + type McpDependencyRequest = { file: string; depth?: number | undefined; @@ -450,7 +465,7 @@ function createCodegraphMcpHandlersForSession( state: "stale", changedFiles: boundedChangedFiles, changedFileCount: changedFiles.length, - omittedChangedFileCount: Math.max(0, changedFiles.length - boundedChangedFiles.length), + omittedChangedFileCount: countOmitted(changedFiles.length, boundedChangedFiles.length), reason, }; }; @@ -461,9 +476,13 @@ function createCodegraphMcpHandlersForSession( const withFreshness = async ( run: () => Promise, ): Promise => { - const freshness = await checkMcpFreshness(); - const result = await run(); - return { ...result, freshness }; + while (true) { + if (refreshPromise) await refreshPromise; + const epoch = refreshEpoch; + const freshness = await checkMcpFreshness(); + const result = await run(); + if (epoch === refreshEpoch && !refreshPromise) return { ...result, freshness }; + } }; const formatSqliteFreshnessError = (freshness: AgentFreshnessResult): string => { if (freshness.state === "fresh") return "SQLite artifact freshness check unexpectedly failed."; @@ -642,13 +661,18 @@ function createCodegraphMcpHandlersForSession( targetFile = await resolveProjectFile(await realRoot, root, request.file); } const collected = collectEntries(snapshot.fileGraph, targetFile, queryOptions); - const totalSeen = collected.length; - const truncated = totalSeen > limit; - const entries = collected.slice(0, limit).map((dependency) => ({ + const { items, omitted } = boundList(collected, limit); + const entries = items.map((dependency) => ({ file: relative(dependency.file), depth: dependency.depth, })); - return { entries, limit, totalSeen, truncated, omitted: truncated ? totalSeen - limit : 0 }; + return { + entries, + limit, + totalSeen: collected.length, + truncated: Boolean(omitted), + omitted, + }; }; const calls = async (request: { direction: "callers" | "callees"; @@ -687,12 +711,15 @@ function createCodegraphMcpHandlersForSession( if (result.status !== "ok") { return { references: [], limit, totalSeen: 0, truncated: false, omitted: 0 }; } - const totalSeen = result.references.length; - const truncated = totalSeen > limit; - const references = result.references - .slice(0, limit) - .map((reference) => ({ file: relative(reference.file), range: reference.range })); - return { references, limit, totalSeen, truncated, omitted: truncated ? totalSeen - limit : 0 }; + const { items, omitted } = boundList(result.references, limit); + const references = items.map((reference) => ({ file: relative(reference.file), range: reference.range })); + return { + references, + limit, + totalSeen: result.references.length, + truncated: Boolean(omitted), + omitted, + }; }; const fileDeps = async (request: { direction: "deps" | "rdeps"; @@ -1057,6 +1084,7 @@ export function createCodegraphMcpProtocolServer( runtimeIdentity: CodegraphRuntimeIdentity = captureCodegraphRuntimeIdentity(getCurrentNativeBindingOrigin()), installedVersion: InstalledVersionChecker = createInstalledVersionChecker(runtimeIdentity), toolCallState: { firstToolCallPending: boolean } = { firstToolCallPending: true }, + maxConcurrentToolCalls = DEFAULT_MCP_TOOL_CONCURRENCY, ): Server { const server = new Server( { @@ -1067,60 +1095,68 @@ export function createCodegraphMcpProtocolServer( capabilities: { tools: {}, logging: {} }, }, ); + let inFlightToolCalls = 0; server.setRequestHandler("tools/list", () => ({ tools: MCP_TOOLS })); server.setRequestHandler("tools/call", async (request, ctx): Promise => { - const isFirstToolCall = toolCallState.firstToolCallPending; - toolCallState.firstToolCallPending = false; - const progressToken = isFirstToolCall ? getToolCallProgressToken(request.params) : undefined; - const emitFirstToolCallVisibility = async ( - level: "info" | "error", - progress: number, - message: string, - ): Promise => { - if (!isFirstToolCall) return; - try { - await ctx.mcpReq.log(level, message, "codegraph"); - if (progressToken !== undefined) { - await ctx.mcpReq.notify({ - method: "notifications/progress", - params: { - progressToken, - progress, - total: 1, - message, - }, - }); - } - } catch (error) { - console.error(`[codegraph] MCP cold-start visibility failed: ${errorMessage(error)}`); - } - }; - await emitFirstToolCallVisibility( - "info", - 0, - `Codegraph is warming the first tool call for '${request.params.name}'.`, - ); - try { - installedVersion.check(); - } catch (error) { - console.error(`[codegraph] installed-version check failed: ${errorMessage(error)}`); + if (inFlightToolCalls >= maxConcurrentToolCalls) { + throw new Error("MCP tool execution is busy; retry shortly."); } + inFlightToolCalls += 1; try { - const result = await callMcpTool(handlers, request.params.name, request.params.arguments ?? {}); + const isFirstToolCall = toolCallState.firstToolCallPending; + toolCallState.firstToolCallPending = false; + const progressToken = isFirstToolCall ? getToolCallProgressToken(request.params) : undefined; + const emitFirstToolCallVisibility = async ( + level: "info" | "error", + progress: number, + message: string, + ): Promise => { + if (!isFirstToolCall) return; + try { + await ctx.mcpReq.log(level, message, "codegraph"); + if (progressToken !== undefined) { + await ctx.mcpReq.notify({ + method: "notifications/progress", + params: { progressToken, progress, total: 1, message }, + }); + } + } catch (error) { + console.error(`[codegraph] MCP cold-start visibility failed: ${errorMessage(error)}`); + } + }; await emitFirstToolCallVisibility( "info", - 1, - `Codegraph finished warming the first tool call for '${request.params.name}'.`, - ); - return toToolResult(result); - } catch (error) { - await emitFirstToolCallVisibility( - "error", - 1, - `Codegraph failed while warming the first tool call for '${request.params.name}'.`, + 0, + `Codegraph is warming the first tool call for '${request.params.name}'.`, ); - throw error; + try { + installedVersion.check(); + } catch (error) { + console.error(`[codegraph] installed-version check failed: ${errorMessage(error)}`); + } + try { + const result = await withAbortSignal( + ctx.mcpReq.signal, + async () => + await callMcpTool(handlers, request.params.name, request.params.arguments ?? {}, ctx.mcpReq.signal), + ); + await emitFirstToolCallVisibility( + "info", + 1, + `Codegraph finished warming the first tool call for '${request.params.name}'.`, + ); + return toToolResult(result); + } catch (error) { + await emitFirstToolCallVisibility( + "error", + 1, + `Codegraph failed while warming the first tool call for '${request.params.name}'.`, + ); + throw error; + } + } finally { + inFlightToolCalls -= 1; } }); @@ -1130,10 +1166,18 @@ export function createCodegraphMcpProtocolServer( function createCodegraphMcpProtocolFactory( handlers: CodegraphMcpHandlers, runtimeIdentity: CodegraphRuntimeIdentity, + maxConcurrentToolCalls = DEFAULT_MCP_TOOL_CONCURRENCY, ): () => Server { const installedVersion = createInstalledVersionChecker(runtimeIdentity); const toolCallState = { firstToolCallPending: true }; - return () => createCodegraphMcpProtocolServer(handlers, runtimeIdentity, installedVersion, toolCallState); + return () => + createCodegraphMcpProtocolServer( + handlers, + runtimeIdentity, + installedVersion, + toolCallState, + maxConcurrentToolCalls, + ); } export async function serveCodegraphMcp(options: CodegraphMcpServerOptions): Promise { @@ -1151,7 +1195,11 @@ export async function serveCodegraphMcp(options: CodegraphMcpServerOptions): Pro const { handlers, session } = await createWarmedCodegraphMcpResources(options); const runtimeIdentity = options.runtimeIdentity ?? captureCodegraphRuntimeIdentity(getCurrentNativeBindingOrigin()); - const createProtocolServer = createCodegraphMcpProtocolFactory(handlers, runtimeIdentity); + const createProtocolServer = createCodegraphMcpProtocolFactory( + handlers, + runtimeIdentity, + options.mcpToolConcurrency ?? DEFAULT_MCP_TOOL_CONCURRENCY, + ); const handle = serveStdio(createProtocolServer, { legacy: "serve", onerror: (error) => { @@ -1176,7 +1224,11 @@ export async function startCodegraphMcpHttpServer( const host = options.host ?? "127.0.0.1"; const { handlers, session } = await createWarmedCodegraphMcpResources(options); const runtimeIdentity = options.runtimeIdentity ?? captureCodegraphRuntimeIdentity(getCurrentNativeBindingOrigin()); - const createProtocolServer = createCodegraphMcpProtocolFactory(handlers, runtimeIdentity); + const createProtocolServer = createCodegraphMcpProtocolFactory( + handlers, + runtimeIdentity, + options.mcpToolConcurrency ?? DEFAULT_MCP_TOOL_CONCURRENCY, + ); const sessionStore = createLegacyMcpSessionStore({ idleMs: options.httpSessionIdleMs ?? DEFAULT_MCP_HTTP_SESSION_IDLE_MS, maxCount: options.httpSessionMaxCount ?? DEFAULT_MCP_HTTP_SESSION_MAX_COUNT, @@ -1557,102 +1609,114 @@ function isMcpNodeRequest(request: IncomingMessage): request is IncomingMessage return request.method !== undefined && request.url !== undefined; } -async function callMcpTool(handlers: CodegraphMcpHandlers, name: string, input: unknown): Promise { +function withAbortSignal(signal: AbortSignal | undefined, run: () => Promise): Promise { + if (!signal) return run(); + if (signal.aborted) return Promise.reject(new Error("MCP tool call was cancelled.")); + const cancellation = Promise.withResolvers(); + const onAbort = (): void => cancellation.reject(new Error("MCP tool call was cancelled.")); + signal.addEventListener("abort", onAbort, { once: true }); + return Promise.race([run(), cancellation.promise]).finally(() => signal.removeEventListener("abort", onAbort)); +} + +async function callMcpTool( + handlers: CodegraphMcpHandlers, + name: string, + input: unknown, + signal?: AbortSignal, +): Promise { switch (name) { case "search": - return await handlers.search(parseMcpToolInput(searchSchema, input, "search")); + return await handlers.search(parseMcpToolInput(searchSchema, input, name), signal); case "workspace_symbols": - return await handlers.workspace_symbols(parseMcpToolInput(workspaceSymbolsSchema, input, "workspace_symbols")); + return await handlers.workspace_symbols(parseMcpToolInput(workspaceSymbolsSchema, input, name), signal); case "rename_preview": - return await handlers.rename_preview(parseMcpToolInput(renamePreviewSchema, input, "rename_preview")); + return await handlers.rename_preview(parseMcpToolInput(renamePreviewSchema, input, name), signal); case "refactor_plan": - return await handlers.refactor_plan(parseMcpToolInput(refactorPlanSchema, input, "refactor_plan")); + return await handlers.refactor_plan(parseMcpToolInput(refactorPlanSchema, input, name), signal); case "calls": - return await handlers.calls(parseMcpToolInput(callsSchema, input, "calls")); + return await handlers.calls(parseMcpToolInput(callsSchema, input, name), signal); case "callers": - return await handlers.calls({ ...parseMcpToolInput(callHierarchySchema, input, name), direction: "callers" }); case "callees": - return await handlers.calls({ ...parseMcpToolInput(callHierarchySchema, input, name), direction: "callees" }); + return await handlers.calls({ ...parseMcpToolInput(callHierarchySchema, input, name), direction: name }, signal); case "type_hierarchy": - return await handlers.type_hierarchy(parseMcpToolInput(typeHierarchyUnifiedSchema, input, "type_hierarchy")); + return await handlers.type_hierarchy(parseMcpToolInput(typeHierarchyUnifiedSchema, input, name), signal); case "supertypes": - return await handlers.type_hierarchy({ - ...parseMcpToolInput(typeHierarchySchema, input, name), - direction: "supertypes", - }); case "subtypes": - return await handlers.type_hierarchy({ - ...parseMcpToolInput(typeHierarchySchema, input, name), - direction: "subtypes", - }); + return await handlers.type_hierarchy( + { ...parseMcpToolInput(typeHierarchySchema, input, name), direction: name }, + signal, + ); case "implementations": - return await handlers.implementations(parseMcpToolInput(implementationsSchema, input, "implementations")); + return await handlers.implementations(parseMcpToolInput(implementationsSchema, input, name), signal); case "explore": - return await handlers.explore(parseMcpToolInput(exploreSchema, input, "explore")); + return await handlers.explore(parseMcpToolInput(exploreSchema, input, name), signal); case "orient": - return await handlers.orient(parseMcpToolInput(orientSchema, input, "orient")); + return await handlers.orient(parseMcpToolInput(orientSchema, input, name), signal); case "packet_get": - return await handlers.packet_get(parseMcpToolInput(packetGetSchema, input, "packet_get")); + return await handlers.packet_get(parseMcpToolInput(packetGetSchema, input, name), signal); case "get_file": - return await handlers.get_file(parseMcpToolInput(getFileSchema, input, "get_file")); + return await handlers.get_file(parseMcpToolInput(getFileSchema, input, name), signal); case "get_symbol": - return await handlers.get_symbol(parseMcpToolInput(handleSchema, input, "get_symbol")); + return await handlers.get_symbol(parseMcpToolInput(handleSchema, input, name), signal); case "goto": - return await callGotoTool(handlers, input); + return await callGotoTool(handlers, input, signal); case "refs": - return await callRefsTool(handlers, input); + return await callRefsTool(handlers, input, signal); case "file_deps": - return await handlers.file_deps(parseMcpToolInput(fileDepsUnifiedSchema, input, "file_deps")); + return await handlers.file_deps(parseMcpToolInput(fileDepsUnifiedSchema, input, name), signal); case "deps": - return await handlers.file_deps({ ...parseMcpToolInput(fileGraphSchema, input, name), direction: "deps" }); case "rdeps": - return await handlers.file_deps({ ...parseMcpToolInput(fileGraphSchema, input, name), direction: "rdeps" }); + return await handlers.file_deps({ ...parseMcpToolInput(fileGraphSchema, input, name), direction: name }, signal); case "path": - return await handlers.path(parseMcpToolInput(pathSchema, input, "path")); + return await handlers.path(parseMcpToolInput(pathSchema, input, name), signal); case "impact": - return await handlers.impact(parseMcpToolInput(gitRangeSchema, input, name)); + return await handlers.impact(parseMcpToolInput(gitRangeSchema, input, name), signal); case "review": - return await handlers.review(parseMcpToolInput(reviewSchema, input, "review")); + return await handlers.review(parseMcpToolInput(reviewSchema, input, name), signal); case "query_sqlite": - return await handlers.query_sqlite(parseMcpToolInput(querySqliteSchema, input, "query_sqlite")); + return await handlers.query_sqlite(parseMcpToolInput(querySqliteSchema, input, name), signal); case "refresh_index": - return await handlers.refresh_index(parseMcpToolInput(refreshIndexSchema, input, "refresh_index")); + return await handlers.refresh_index(parseMcpToolInput(refreshIndexSchema, input, name), signal); case "artifact_build": - return await handlers.artifact_build(parseMcpToolInput(artifactBuildSchema, input, "artifact_build")); + return await handlers.artifact_build(parseMcpToolInput(artifactBuildSchema, input, name), signal); default: throw new Error(`Unknown MCP tool: ${name}`); } } -async function callGotoTool(handlers: CodegraphMcpHandlers, input: unknown): Promise { +async function callGotoTool(handlers: CodegraphMcpHandlers, input: unknown, signal?: AbortSignal): Promise { const request = parseMcpToolInput(navigationSchema, input, "goto"); - if (request.handle !== undefined) return await handlers.goto({ handle: request.handle }); + if (request.handle !== undefined) return await handlers.goto({ handle: request.handle }, signal); if (request.file === undefined || request.line === undefined || request.column === undefined) { throw new Error("goto requires either handle or file, line, and column."); } - return await handlers.goto({ file: request.file, line: request.line, column: request.column }); + return await handlers.goto({ file: request.file, line: request.line, column: request.column }, signal); } async function callRefsTool( handlers: CodegraphMcpHandlers, input: unknown, + signal?: AbortSignal, ): Promise { const request = parseMcpToolInput(refsSchema, input, "refs"); if (request.handle !== undefined) { - return await handlers.refs({ - handle: request.handle, - ...(request.limit !== undefined ? { limit: request.limit } : {}), - }); + return await handlers.refs( + { handle: request.handle, ...(request.limit !== undefined ? { limit: request.limit } : {}) }, + signal, + ); } if (request.file === undefined || request.line === undefined || request.column === undefined) { throw new Error("refs requires either handle or file, line, and column."); } - return await handlers.refs({ - file: request.file, - line: request.line, - column: request.column, - ...(request.limit !== undefined ? { limit: request.limit } : {}), - }); + return await handlers.refs( + { + file: request.file, + line: request.line, + column: request.column, + ...(request.limit !== undefined ? { limit: request.limit } : {}), + }, + signal, + ); } function toToolResult(value: unknown): CallToolResult { diff --git a/src/mcp/tools.ts b/src/mcp/tools.ts index edbb54fb..e40bbe36 100644 --- a/src/mcp/tools.ts +++ b/src/mcp/tools.ts @@ -357,6 +357,110 @@ export const MCP_TOOLS: Tool[] = [ force: booleanProperty, }), }, + { + name: "callers", + description: "Legacy alias for calls with direction callers.", + inputSchema: objectSchema( + { + handle: stringProperty, + depth: { type: "integer", minimum: 1, maximum: MAX_CALL_HIERARCHY_DEPTH, default: 1 }, + limit: { + type: "integer", + minimum: 0, + maximum: MAX_CALL_HIERARCHY_LIMIT, + default: DEFAULT_CALL_HIERARCHY_LIMIT, + }, + includeHeuristic: booleanProperty, + }, + ["handle"], + ), + }, + { + name: "callees", + description: "Legacy alias for calls with direction callees.", + inputSchema: objectSchema( + { + handle: stringProperty, + depth: { type: "integer", minimum: 1, maximum: MAX_CALL_HIERARCHY_DEPTH, default: 1 }, + limit: { + type: "integer", + minimum: 0, + maximum: MAX_CALL_HIERARCHY_LIMIT, + default: DEFAULT_CALL_HIERARCHY_LIMIT, + }, + includeHeuristic: booleanProperty, + }, + ["handle"], + ), + }, + { + name: "supertypes", + description: "Legacy alias for type_hierarchy with direction supertypes.", + inputSchema: objectSchema( + { + handle: stringProperty, + depth: { type: "integer", minimum: 1, maximum: MAX_TYPE_HIERARCHY_DEPTH, default: 1 }, + limit: { + type: "integer", + minimum: 0, + maximum: MAX_TYPE_HIERARCHY_LIMIT, + default: DEFAULT_TYPE_HIERARCHY_LIMIT, + }, + }, + ["handle"], + ), + }, + { + name: "subtypes", + description: "Legacy alias for type_hierarchy with direction subtypes.", + inputSchema: objectSchema( + { + handle: stringProperty, + depth: { type: "integer", minimum: 1, maximum: MAX_TYPE_HIERARCHY_DEPTH, default: 1 }, + limit: { + type: "integer", + minimum: 0, + maximum: MAX_TYPE_HIERARCHY_LIMIT, + default: DEFAULT_TYPE_HIERARCHY_LIMIT, + }, + }, + ["handle"], + ), + }, + { + name: "deps", + description: "Legacy alias for file_deps with direction deps.", + inputSchema: objectSchema( + { + file: dependencyFileProperty, + depth: { type: "integer", minimum: 0, default: 1 }, + limit: { + type: "integer", + minimum: 0, + maximum: MAX_MCP_COLLECTION_LIMIT, + default: DEFAULT_MCP_COLLECTION_LIMIT, + }, + }, + ["file"], + ), + }, + { + name: "rdeps", + description: "Legacy alias for file_deps with direction rdeps.", + inputSchema: objectSchema( + { + file: dependencyFileProperty, + depth: { type: "integer", minimum: 0, default: 1 }, + limit: { + type: "integer", + minimum: 0, + maximum: MAX_MCP_COLLECTION_LIMIT, + default: DEFAULT_MCP_COLLECTION_LIMIT, + }, + }, + ["file"], + ), + }, ]; export function listCodegraphMcpTools(): Tool[] { From 6d4d39c0bcedf9b6a32fde4896f7f0a202ecb459 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sat, 15 Aug 2026 11:53:01 -0400 Subject: [PATCH 06/37] refactor: expose MCP tool registry --- src/mcp/server.ts | 4 ++++ src/mcp/tools.ts | 4 +++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/src/mcp/server.ts b/src/mcp/server.ts index 721202b4..e4cd6c95 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -85,6 +85,7 @@ import { MAX_TYPE_HIERARCHY_LIMIT, MAX_MCP_COLLECTION_LIMIT, MAX_RENAME_PREVIEW_EDITS, + MCP_TOOL_REGISTRY, MCP_TOOLS, MAX_REFACTOR_PLAN_LIMIT, } from "./tools.js"; @@ -1624,6 +1625,9 @@ async function callMcpTool( input: unknown, signal?: AbortSignal, ): Promise { + if (!MCP_TOOL_REGISTRY.some((tool) => tool.name === name)) { + throw new Error(`Unknown MCP tool: ${name}`); + } switch (name) { case "search": return await handlers.search(parseMcpToolInput(searchSchema, input, name), signal); diff --git a/src/mcp/tools.ts b/src/mcp/tools.ts index e40bbe36..bfbc8d0b 100644 --- a/src/mcp/tools.ts +++ b/src/mcp/tools.ts @@ -127,7 +127,7 @@ function navigationInputSchema(includeLimit: boolean): Tool["inputSchema"] { }; } -export const MCP_TOOLS: Tool[] = [ +export const MCP_TOOL_REGISTRY: Tool[] = [ { name: "search", description: "Deterministic ranked search across files, symbols, chunks, SQL objects, and graph context.", @@ -463,6 +463,8 @@ export const MCP_TOOLS: Tool[] = [ }, ]; +export const MCP_TOOLS = MCP_TOOL_REGISTRY; + export function listCodegraphMcpTools(): Tool[] { return MCP_TOOLS.map((tool) => ({ ...tool })); } From e0c73fed66c51da2617b95583be5aefdcb81e016 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sat, 15 Aug 2026 19:27:59 -0400 Subject: [PATCH 07/37] fix: derive MCP dispatch from tool registry --- src/mcp/server.ts | 45 ++++++++++++++++--------------- src/mcp/tools.ts | 58 +++++++++++++++++++++++++++++++++++++--- tests/mcp-server.test.ts | 57 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 135 insertions(+), 25 deletions(-) diff --git a/src/mcp/server.ts b/src/mcp/server.ts index e4cd6c95..0bac860b 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -86,7 +86,6 @@ import { MAX_MCP_COLLECTION_LIMIT, MAX_RENAME_PREVIEW_EDITS, MCP_TOOL_REGISTRY, - MCP_TOOLS, MAX_REFACTOR_PLAN_LIMIT, } from "./tools.js"; import { @@ -1098,7 +1097,7 @@ export function createCodegraphMcpProtocolServer( ); let inFlightToolCalls = 0; - server.setRequestHandler("tools/list", () => ({ tools: MCP_TOOLS })); + server.setRequestHandler("tools/list", () => ({ tools: listCodegraphMcpTools() })); server.setRequestHandler("tools/call", async (request, ctx): Promise => { if (inFlightToolCalls >= maxConcurrentToolCalls) { throw new Error("MCP tool execution is busy; retry shortly."); @@ -1619,16 +1618,16 @@ function withAbortSignal(signal: AbortSignal | undefined, run: () => Promise< return Promise.race([run(), cancellation.promise]).finally(() => signal.removeEventListener("abort", onAbort)); } -async function callMcpTool( +export async function callMcpTool( handlers: CodegraphMcpHandlers, name: string, input: unknown, signal?: AbortSignal, ): Promise { - if (!MCP_TOOL_REGISTRY.some((tool) => tool.name === name)) { - throw new Error(`Unknown MCP tool: ${name}`); - } - switch (name) { + const tool = MCP_TOOL_REGISTRY.find((entry) => entry.name === name); + if (!tool) throw new Error(`Unknown MCP tool: ${name}`); + + switch (tool.dispatch.handler) { case "search": return await handlers.search(parseMcpToolInput(searchSchema, input, name), signal); case "workspace_symbols": @@ -1638,18 +1637,21 @@ async function callMcpTool( case "refactor_plan": return await handlers.refactor_plan(parseMcpToolInput(refactorPlanSchema, input, name), signal); case "calls": + if (tool.dispatch.direction) { + return await handlers.calls( + { ...parseMcpToolInput(callHierarchySchema, input, name), direction: tool.dispatch.direction }, + signal, + ); + } return await handlers.calls(parseMcpToolInput(callsSchema, input, name), signal); - case "callers": - case "callees": - return await handlers.calls({ ...parseMcpToolInput(callHierarchySchema, input, name), direction: name }, signal); case "type_hierarchy": + if (tool.dispatch.direction) { + return await handlers.type_hierarchy( + { ...parseMcpToolInput(typeHierarchySchema, input, name), direction: tool.dispatch.direction }, + signal, + ); + } return await handlers.type_hierarchy(parseMcpToolInput(typeHierarchyUnifiedSchema, input, name), signal); - case "supertypes": - case "subtypes": - return await handlers.type_hierarchy( - { ...parseMcpToolInput(typeHierarchySchema, input, name), direction: name }, - signal, - ); case "implementations": return await handlers.implementations(parseMcpToolInput(implementationsSchema, input, name), signal); case "explore": @@ -1667,10 +1669,13 @@ async function callMcpTool( case "refs": return await callRefsTool(handlers, input, signal); case "file_deps": + if (tool.dispatch.direction) { + return await handlers.file_deps( + { ...parseMcpToolInput(fileGraphSchema, input, name), direction: tool.dispatch.direction }, + signal, + ); + } return await handlers.file_deps(parseMcpToolInput(fileDepsUnifiedSchema, input, name), signal); - case "deps": - case "rdeps": - return await handlers.file_deps({ ...parseMcpToolInput(fileGraphSchema, input, name), direction: name }, signal); case "path": return await handlers.path(parseMcpToolInput(pathSchema, input, name), signal); case "impact": @@ -1683,8 +1688,6 @@ async function callMcpTool( return await handlers.refresh_index(parseMcpToolInput(refreshIndexSchema, input, name), signal); case "artifact_build": return await handlers.artifact_build(parseMcpToolInput(artifactBuildSchema, input, name), signal); - default: - throw new Error(`Unknown MCP tool: ${name}`); } } diff --git a/src/mcp/tools.ts b/src/mcp/tools.ts index bfbc8d0b..119ea208 100644 --- a/src/mcp/tools.ts +++ b/src/mcp/tools.ts @@ -127,7 +127,32 @@ function navigationInputSchema(includeLimit: boolean): Tool["inputSchema"] { }; } -export const MCP_TOOL_REGISTRY: Tool[] = [ +export type McpToolDispatch = + | { handler: "search" } + | { handler: "workspace_symbols" } + | { handler: "rename_preview" } + | { handler: "refactor_plan" } + | { handler: "calls"; direction?: "callers" | "callees" } + | { handler: "type_hierarchy"; direction?: "supertypes" | "subtypes" } + | { handler: "implementations" } + | { handler: "explore" } + | { handler: "orient" } + | { handler: "packet_get" } + | { handler: "get_file" } + | { handler: "get_symbol" } + | { handler: "goto" } + | { handler: "refs" } + | { handler: "file_deps"; direction?: "deps" | "rdeps" } + | { handler: "path" } + | { handler: "impact" } + | { handler: "review" } + | { handler: "query_sqlite" } + | { handler: "refresh_index" } + | { handler: "artifact_build" }; + +export type McpToolDefinition = Tool & { dispatch: McpToolDispatch }; + +export const MCP_TOOL_REGISTRY: McpToolDefinition[] = [ { name: "search", description: "Deterministic ranked search across files, symbols, chunks, SQL objects, and graph context.", @@ -141,6 +166,7 @@ export const MCP_TOOL_REGISTRY: Tool[] = [ }, ["query"], ), + dispatch: { handler: "search" }, }, { name: "workspace_symbols", @@ -162,6 +188,7 @@ export const MCP_TOOL_REGISTRY: Tool[] = [ }, ["query"], ), + dispatch: { handler: "workspace_symbols" }, }, { name: "rename_preview", @@ -183,6 +210,7 @@ export const MCP_TOOL_REGISTRY: Tool[] = [ }, ["handle", "newName"], ), + dispatch: { handler: "rename_preview" }, }, { name: "refactor_plan", @@ -199,18 +227,21 @@ export const MCP_TOOL_REGISTRY: Tool[] = [ }, ["handle"], ), + dispatch: { handler: "refactor_plan" }, }, { name: "calls", description: "Find proven semantic callers or callees and exact grouped callsites for a portable symbol handle. Use refs for every symbol reference and file_deps for file-level dependencies.", inputSchema: callHierarchyInputSchema(), + dispatch: { handler: "calls" }, }, { name: "type_hierarchy", description: "Find proven direct or transitive supertypes or subtypes for a portable symbol handle. Returns currently extracted extends and implements relationships only.", inputSchema: typeHierarchyInputSchema(), + dispatch: { handler: "type_hierarchy" }, }, { name: "implementations", @@ -228,6 +259,7 @@ export const MCP_TOOL_REGISTRY: Tool[] = [ }, ["handle"], ), + dispatch: { handler: "implementations" }, }, { name: "explore", @@ -243,6 +275,7 @@ export const MCP_TOOL_REGISTRY: Tool[] = [ }, ["query"], ), + dispatch: { handler: "explore" }, }, { name: "orient", @@ -251,6 +284,7 @@ export const MCP_TOOL_REGISTRY: Tool[] = [ includeRoots: { type: "array", items: stringProperty }, budget: orientBudgetProperty, }), + dispatch: { handler: "orient" }, }, { name: "packet_get", @@ -264,6 +298,7 @@ export const MCP_TOOL_REGISTRY: Tool[] = [ }, ["target"], ), + dispatch: { handler: "packet_get" }, }, { name: "get_file", @@ -279,37 +314,44 @@ export const MCP_TOOL_REGISTRY: Tool[] = [ }, ["file"], ), + dispatch: { handler: "get_file" }, }, { name: "get_symbol", description: "Resolve a stable search or explain handle.", inputSchema: objectSchema({ handle: stringProperty }, ["handle"]), + dispatch: { handler: "get_symbol" }, }, { name: "goto", description: "Resolve a definition by portable handle, qualified file::symbol path, or file position.", inputSchema: navigationInputSchema(false), + dispatch: { handler: "goto" }, }, { name: "refs", description: "Find references by portable handle, qualified file::symbol path, or file position.", inputSchema: navigationInputSchema(true), + dispatch: { handler: "refs" }, }, { name: "file_deps", description: "List file dependencies or reverse file dependencies by file path, qualified file::symbol path, or portable handle.", inputSchema: dependencyInputSchema(), + dispatch: { handler: "file_deps" }, }, { name: "path", description: "Find the shortest dependency path between two files.", inputSchema: objectSchema({ from: stringProperty, to: stringProperty }, ["from", "to"]), + dispatch: { handler: "path" }, }, { name: "impact", description: "Build compact impact context for a git range.", inputSchema: objectSchema({ base: stringProperty, head: stringProperty }, ["base", "head"]), + dispatch: { handler: "impact" }, }, { name: "review", @@ -322,6 +364,7 @@ export const MCP_TOOL_REGISTRY: Tool[] = [ }, ["base", "head"], ), + dispatch: { handler: "review" }, }, { name: "query_sqlite", @@ -337,6 +380,7 @@ export const MCP_TOOL_REGISTRY: Tool[] = [ }, ["query"], ), + dispatch: { handler: "query_sqlite" }, }, { name: "refresh_index", @@ -344,6 +388,7 @@ export const MCP_TOOL_REGISTRY: Tool[] = [ inputSchema: objectSchema({ warmup: { type: "string", enum: ["off", "base", "symbols"] }, }), + dispatch: { handler: "refresh_index" }, }, { name: "artifact_build", @@ -356,6 +401,7 @@ export const MCP_TOOL_REGISTRY: Tool[] = [ questions: booleanProperty, force: booleanProperty, }), + dispatch: { handler: "artifact_build" }, }, { name: "callers", @@ -374,6 +420,7 @@ export const MCP_TOOL_REGISTRY: Tool[] = [ }, ["handle"], ), + dispatch: { handler: "calls", direction: "callers" }, }, { name: "callees", @@ -392,6 +439,7 @@ export const MCP_TOOL_REGISTRY: Tool[] = [ }, ["handle"], ), + dispatch: { handler: "calls", direction: "callees" }, }, { name: "supertypes", @@ -409,6 +457,7 @@ export const MCP_TOOL_REGISTRY: Tool[] = [ }, ["handle"], ), + dispatch: { handler: "type_hierarchy", direction: "supertypes" }, }, { name: "subtypes", @@ -426,6 +475,7 @@ export const MCP_TOOL_REGISTRY: Tool[] = [ }, ["handle"], ), + dispatch: { handler: "type_hierarchy", direction: "subtypes" }, }, { name: "deps", @@ -443,6 +493,7 @@ export const MCP_TOOL_REGISTRY: Tool[] = [ }, ["file"], ), + dispatch: { handler: "file_deps", direction: "deps" }, }, { name: "rdeps", @@ -460,11 +511,10 @@ export const MCP_TOOL_REGISTRY: Tool[] = [ }, ["file"], ), + dispatch: { handler: "file_deps", direction: "rdeps" }, }, ]; -export const MCP_TOOLS = MCP_TOOL_REGISTRY; - export function listCodegraphMcpTools(): Tool[] { - return MCP_TOOLS.map((tool) => ({ ...tool })); + return MCP_TOOL_REGISTRY.map(({ dispatch: _dispatch, ...tool }) => tool); } diff --git a/tests/mcp-server.test.ts b/tests/mcp-server.test.ts index 259359a4..138bc28c 100644 --- a/tests/mcp-server.test.ts +++ b/tests/mcp-server.test.ts @@ -6,6 +6,7 @@ import { DatabaseSync } from "node:sqlite"; import { describe, expect, it, vi } from "vitest"; import { createAgentSession, type AgentProjectSnapshot, type AgentSession } from "../src/agent/session.js"; import { + callMcpTool, createCodegraphMcpHandlers, createCodegraphMcpProtocolServer, listCodegraphMcpTools, @@ -14,6 +15,7 @@ import { type CodegraphMcpHandlers, } from "../src/mcp/server.js"; import { SymbolKind, type ModuleIndex, type ProjectIndex } from "../src/indexer/types.js"; +import { MCP_TOOL_REGISTRY } from "../src/mcp/tools.js"; import { DEFAULT_REVIEW_TRANSPORT_LIMITS } from "../src/review/types.js"; import type { Graph } from "../src/types.js"; import * as symbolGraphBuild from "../src/graphs/symbol-graph-detailed.js"; @@ -2754,3 +2756,58 @@ describe("codegraph MCP handlers", () => { function normalizeSqlitePath(value: unknown): string { return typeof value === "string" ? value.replace(/\\/g, "/") : ""; } + +describe("MCP tool registry dispatch", () => { + it("routes every advertised tool to a schema-valid handler", async () => { + const root = await fs.mkdtemp(path.join(os.tmpdir(), "cg-mcp-registry-")); + await fs.writeFile(path.join(root, "auth.ts"), "export function ok(): number { return 1; }\n", "utf8"); + runGit(root, ["init"]); + runGit(root, ["add", "."]); + runGit(root, ["commit", "-m", "base"]); + const handlers = createCodegraphMcpHandlers({ root }); + const advertisedTools = listCodegraphMcpTools(); + expect(advertisedTools.map((tool) => tool.name)).toEqual(MCP_TOOL_REGISTRY.map((tool) => tool.name)); + expect(advertisedTools.some((tool) => "dispatch" in tool)).toBe(false); + const handle = "auth.ts::ok"; + const toolInputs: Record> = { + search: { query: "ok" }, + workspace_symbols: { query: "ok" }, + rename_preview: { handle, newName: "renamed" }, + refactor_plan: { handle }, + calls: { handle, direction: "callers" }, + type_hierarchy: { handle, direction: "supertypes" }, + implementations: { handle }, + explore: { query: "ok" }, + orient: {}, + packet_get: { target: "auth.ts" }, + get_file: { file: "auth.ts" }, + get_symbol: { handle }, + goto: { handle }, + refs: { handle }, + file_deps: { file: "auth.ts", direction: "deps" }, + path: { from: "auth.ts", to: "auth.ts" }, + impact: { base: "HEAD", head: "HEAD" }, + review: { base: "HEAD", head: "HEAD" }, + query_sqlite: { query: "SELECT 1" }, + refresh_index: {}, + artifact_build: {}, + callers: { handle }, + callees: { handle }, + supertypes: { handle }, + subtypes: { handle }, + deps: { file: "auth.ts" }, + rdeps: { file: "auth.ts" }, + }; + + for (const tool of advertisedTools) { + const input = toolInputs[tool.name]; + expect(input, "missing valid input for " + tool.name).toBeDefined(); + try { + await callMcpTool(handlers, tool.name, input); + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + expect(message, tool.name).not.toMatch(/Unknown MCP tool|Invalid parameters for/i); + } + } + }); +}); From c3ad27b0d7ebd3093900c641133b20baf0be0a94 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sat, 15 Aug 2026 20:06:56 -0400 Subject: [PATCH 08/37] fix: serialize MCP refresh requests --- codegraph-skill/codegraph/SKILL.md | 2 +- docs/library-api.md | 2 + docs/mcp.md | 6 +- src/mcp/server.ts | 118 +++++++++++++-------------- tests/mcp-server.test.ts | 124 +++++++++++++++++++++++++++++ 5 files changed, 192 insertions(+), 60 deletions(-) diff --git a/codegraph-skill/codegraph/SKILL.md b/codegraph-skill/codegraph/SKILL.md index 3f2c1868..7400b312 100644 --- a/codegraph-skill/codegraph/SKILL.md +++ b/codegraph-skill/codegraph/SKILL.md @@ -164,7 +164,7 @@ Sensitive-file rules: If MCP tools are available, prefer them over repeated CLI invocations. Use `explore`, `orient`, `workspace_symbols`, `search`, `get_file`, `get_symbol`, `packet_get`, `goto`, `refs`, `rename_preview`, `refactor_plan`, `calls`, `type_hierarchy`, `implementations`, `file_deps`, `path`, `impact`, `review`, `query_sqlite`, `refresh_index`, and `artifact_build`; fall back to the CLI when MCP is unavailable. Legacy `callers`/`callees`, `supertypes`/`subtypes`, and `deps`/`rdeps` names remain valid `tools/call` aliases. -codegraph uses the official MCP SDK v2 to serve current 2026-07-28 clients while retaining compatibility with 2025-era clients. MCP protocol connections and HTTP protocol sessions keep separate transport state, but all share the server's one warm codegraph analysis session for the configured root. Tool schemas reject unknown fields, and idle HTTP protocol sessions are evicted with a bounded session count. +codegraph uses the official MCP SDK v2 to serve current 2026-07-28 clients while retaining compatibility with 2025-era clients. MCP protocol connections and HTTP protocol sessions keep separate transport state, but all share the server's one warm codegraph analysis session for the configured root. Tool schemas reject unknown fields, each protocol session caps tool concurrency at 4 with a retryable busy error, HTTP request bodies time out after 30 seconds, and idle HTTP protocol sessions are evicted with a bounded session count. Cancellation responds to the caller promptly but retains the occupied slot until shared work settles, so abandoned calls cannot bypass that resource bound. Concurrent `refresh_index` calls serialize and honor each request's requested warmup. On the first `tools/call`, codegraph can emit `notifications/message` and, when the request includes `_meta.progressToken`, `notifications/progress` before the final result. Stdio carries them inline, and modern Streamable HTTP clients that accept `text/event-stream` receive them as a stream until the terminal result frame. HTTP enforces Host and Origin policies. A missing `Origin` is accepted for non-browser clients; unapproved, malformed, and opaque origins are rejected. This is not authentication: binding `--host` to a non-loopback address exposes an unauthenticated endpoint intended only for trusted networks or containers. diff --git a/docs/library-api.md b/docs/library-api.md index baf35f8d..102575b4 100644 --- a/docs/library-api.md +++ b/docs/library-api.md @@ -439,6 +439,8 @@ console.log(packet.kind, refs.references, rows.rows, rows.freshness.state); ``` `serveCodegraphMcp()` (from `@lzehrung/codegraph/mcp`) starts the stdio server used by `codegraph mcp serve`. MCP is an agent ergonomics and cache layer over the same analysis engine, not a separate indexer. MCP file and artifact paths are confined after realpath resolution. + +`CodegraphMcpServerOptions.mcpToolConcurrency` caps concurrent calls per protocol session (default `4`); saturation returns a retryable busy error. `httpBodyTimeoutMs` bounds HTTP request-body receipt (default `30_000` ms), returning HTTP 408 on expiry. Concurrent `refresh_index` calls serialize and each applies its requested `warmup` after the preceding refresh completes. Client cancellation returns promptly but retains its concurrency slot until shared work settles, so cancellation cannot create unbounded background work. `query_sqlite` is read-only and row- and byte-bounded. It returns freshness metadata for fresh artifact reads, refreshes codegraph-owned SQLite artifacts after small edits when write access is enabled, and rejects stale artifact queries it cannot refresh safely. `artifact_build` is disabled by default and requires `readOnly: false` or CLI `--allow-build`; it refuses to write outputs from a stale MCP index until `refresh_index` succeeds. MCP `orient` and `packet_get` calls use the server-configured root; they do not accept per-request root overrides. diff --git a/docs/mcp.md b/docs/mcp.md index 15673593..ec37bf62 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -33,6 +33,10 @@ Stdio servers exit when the client closes stdin, when an IPC parent disconnects, HTTP protocol sessions track last activity, cap concurrent legacy sessions (default 32), and evict idle sessions on a timer (default 30 minutes). Capacity and idle eviction skip sessions with in-flight requests or open SSE streams; when every slot is active, a new `initialize` receives an actionable JSON-RPC capacity error instead of evicting a working client. Transport errors and protocol session closes also remove the session. +Each MCP protocol session permits four concurrent tool calls by default; a saturated session returns a retryable busy error rather than queueing unbounded work. The programmatic server options `mcpToolConcurrency` and `httpBodyTimeoutMs` tune that cap and the 30-second HTTP request-body deadline; an HTTP body that misses its deadline receives `408 Request Timeout`. + +Client cancellation returns promptly, but does not discard shared index or artifact work. A cancelled call continues to occupy its concurrency slot until its underlying operation settles, preventing a burst of abandoned requests from exceeding the configured resource bound. + Use stdio for a client-owned subprocess. Use HTTP for one long-running codegraph process per repository, then point every MCP-capable IDE, terminal, or agent client at the same local URL. Exact config keys vary by client, but the MCP settings should use HTTP/Streamable HTTP transport plus the `/mcp` URL instead of a `command`/`args` stdio launch. codegraph uses the official MCP SDK v2 to serve current 2026-07-28 clients while retaining compatibility with 2025-era clients. MCP protocol connections and HTTP protocol sessions keep separate transport state, but all share the server's one warm codegraph analysis session for the configured root. Tool request schemas set `additionalProperties: false` and reject unknown fields with an actionable invalid-parameter error instead of silently ignoring typos. @@ -100,7 +104,7 @@ Text and hybrid searches reuse a prepared handle for `.codegraph-cache/index-v1/ If the sidecar is busy or unavailable, MCP uses the same exact in-memory matcher. [How it works](./how-it-works.md#cache-and-session-behavior) explains the search cache. -Use `refresh_index` to rebuild the snapshot, reset SQLite artifact state, or recover after a change burst exceeds automatic limits. With write access, `query_sqlite` refreshes codegraph SQLite artifacts after small edits; otherwise it refuses stale rows. `artifact_build` refuses stale indexes, so run `refresh_index` after large change bursts. +Use `refresh_index` to rebuild the snapshot, reset SQLite artifact state, or recover after a change burst exceeds automatic limits. Concurrent refresh requests serialize; each request runs its own requested `warmup` (`off`, `base`, or `symbols`) after an active refresh completes. With write access, `query_sqlite` refreshes codegraph SQLite artifacts after small edits; otherwise it refuses stale rows. `artifact_build` refuses stale indexes, so run `refresh_index` after large change bursts. `get_file` reads live bytes from disk after path confinement. It does not require a fresh index; only an explicit `includeGraphContext: true` checks indexed freshness and adds direct graph context, so returned file bytes and `totalLines` remain live even when `freshness` reports stale context. Tool schemas are flat JSON objects for broad client compatibility; argument combinations such as `refs` handle-vs-position mode are validated by the server. Legacy paired names (`callers`, `callees`, `supertypes`, `subtypes`, `deps`, and `rdeps`) remain accepted by `tools/call` as aliases, but only the unified tools appear in `tools/list`. diff --git a/src/mcp/server.ts b/src/mcp/server.ts index 0bac860b..9e33d492 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -1013,18 +1013,15 @@ function createCodegraphMcpHandlersForSession( refresh_index: async (request) => { const warmup = request.warmup ?? "off"; - if (refreshPromise) { - await refreshPromise; - return { refreshed: true, warmup }; - } - const epoch = ++refreshEpoch; + const previousRefresh = refreshPromise; const refresh = (async () => { + if (previousRefresh) await previousRefresh.catch(() => undefined); + ++refreshEpoch; session.invalidate(); sqlitePath = configuredSqlitePath; sqliteOutDir = configuredSqliteOutDir; sqliteCanRefresh = configuredSqliteCanRefresh; await startCodegraphMcpWarmup(session, warmup); - if (epoch !== refreshEpoch) throw new Error("MCP index refresh was superseded."); })(); refreshPromise = refresh; try { @@ -1103,60 +1100,56 @@ export function createCodegraphMcpProtocolServer( throw new Error("MCP tool execution is busy; retry shortly."); } inFlightToolCalls += 1; - try { - const isFirstToolCall = toolCallState.firstToolCallPending; - toolCallState.firstToolCallPending = false; - const progressToken = isFirstToolCall ? getToolCallProgressToken(request.params) : undefined; - const emitFirstToolCallVisibility = async ( - level: "info" | "error", - progress: number, - message: string, - ): Promise => { - if (!isFirstToolCall) return; - try { - await ctx.mcpReq.log(level, message, "codegraph"); - if (progressToken !== undefined) { - await ctx.mcpReq.notify({ - method: "notifications/progress", - params: { progressToken, progress, total: 1, message }, - }); - } - } catch (error) { - console.error(`[codegraph] MCP cold-start visibility failed: ${errorMessage(error)}`); + const isFirstToolCall = toolCallState.firstToolCallPending; + toolCallState.firstToolCallPending = false; + const progressToken = isFirstToolCall ? getToolCallProgressToken(request.params) : undefined; + const emitFirstToolCallVisibility = async ( + level: "info" | "error", + progress: number, + message: string, + ): Promise => { + if (!isFirstToolCall) return; + try { + await ctx.mcpReq.log(level, message, "codegraph"); + if (progressToken !== undefined) { + await ctx.mcpReq.notify({ + method: "notifications/progress", + params: { progressToken, progress, total: 1, message }, + }); } + } catch (error) { + console.error(`[codegraph] MCP cold-start visibility failed: ${errorMessage(error)}`); + } + }; + await emitFirstToolCallVisibility( + "info", + 0, + `Codegraph is warming the first tool call for '${request.params.name}'.`, + ); + try { + installedVersion.check(); + } catch (error) { + console.error(`[codegraph] installed-version check failed: ${errorMessage(error)}`); + } + try { + const operation = callMcpTool(handlers, request.params.name, request.params.arguments ?? {}, ctx.mcpReq.signal); + const releaseToolCall = (): void => { + inFlightToolCalls -= 1; }; + const result = await awaitMcpToolOperation(ctx.mcpReq.signal, operation, releaseToolCall); await emitFirstToolCallVisibility( "info", - 0, - `Codegraph is warming the first tool call for '${request.params.name}'.`, + 1, + `Codegraph finished warming the first tool call for '${request.params.name}'.`, ); - try { - installedVersion.check(); - } catch (error) { - console.error(`[codegraph] installed-version check failed: ${errorMessage(error)}`); - } - try { - const result = await withAbortSignal( - ctx.mcpReq.signal, - async () => - await callMcpTool(handlers, request.params.name, request.params.arguments ?? {}, ctx.mcpReq.signal), - ); - await emitFirstToolCallVisibility( - "info", - 1, - `Codegraph finished warming the first tool call for '${request.params.name}'.`, - ); - return toToolResult(result); - } catch (error) { - await emitFirstToolCallVisibility( - "error", - 1, - `Codegraph failed while warming the first tool call for '${request.params.name}'.`, - ); - throw error; - } - } finally { - inFlightToolCalls -= 1; + return toToolResult(result); + } catch (error) { + await emitFirstToolCallVisibility( + "error", + 1, + `Codegraph failed while warming the first tool call for '${request.params.name}'.`, + ); + throw error; } }); @@ -1609,13 +1602,22 @@ function isMcpNodeRequest(request: IncomingMessage): request is IncomingMessage return request.method !== undefined && request.url !== undefined; } -function withAbortSignal(signal: AbortSignal | undefined, run: () => Promise): Promise { - if (!signal) return run(); +export function awaitMcpToolOperation( + signal: AbortSignal | undefined, + operation: Promise, + onSettled: () => void, +): Promise { + void operation.then(onSettled, onSettled); + return withAbortSignal(signal, operation); +} + +function withAbortSignal(signal: AbortSignal | undefined, operation: Promise): Promise { + if (!signal) return operation; if (signal.aborted) return Promise.reject(new Error("MCP tool call was cancelled.")); const cancellation = Promise.withResolvers(); const onAbort = (): void => cancellation.reject(new Error("MCP tool call was cancelled.")); signal.addEventListener("abort", onAbort, { once: true }); - return Promise.race([run(), cancellation.promise]).finally(() => signal.removeEventListener("abort", onAbort)); + return Promise.race([operation, cancellation.promise]).finally(() => signal.removeEventListener("abort", onAbort)); } export async function callMcpTool( diff --git a/tests/mcp-server.test.ts b/tests/mcp-server.test.ts index 138bc28c..b8f5fe79 100644 --- a/tests/mcp-server.test.ts +++ b/tests/mcp-server.test.ts @@ -6,6 +6,7 @@ import { DatabaseSync } from "node:sqlite"; import { describe, expect, it, vi } from "vitest"; import { createAgentSession, type AgentProjectSnapshot, type AgentSession } from "../src/agent/session.js"; import { + awaitMcpToolOperation, callMcpTool, createCodegraphMcpHandlers, createCodegraphMcpProtocolServer, @@ -2765,6 +2766,29 @@ describe("MCP tool registry dispatch", () => { runGit(root, ["add", "."]); runGit(root, ["commit", "-m", "base"]); const handlers = createCodegraphMcpHandlers({ root }); + const handlerSpies = { + search: vi.spyOn(handlers, "search"), + workspace_symbols: vi.spyOn(handlers, "workspace_symbols"), + rename_preview: vi.spyOn(handlers, "rename_preview"), + refactor_plan: vi.spyOn(handlers, "refactor_plan"), + calls: vi.spyOn(handlers, "calls"), + type_hierarchy: vi.spyOn(handlers, "type_hierarchy"), + implementations: vi.spyOn(handlers, "implementations"), + explore: vi.spyOn(handlers, "explore"), + orient: vi.spyOn(handlers, "orient"), + packet_get: vi.spyOn(handlers, "packet_get"), + get_file: vi.spyOn(handlers, "get_file"), + get_symbol: vi.spyOn(handlers, "get_symbol"), + goto: vi.spyOn(handlers, "goto"), + refs: vi.spyOn(handlers, "refs"), + file_deps: vi.spyOn(handlers, "file_deps"), + path: vi.spyOn(handlers, "path"), + impact: vi.spyOn(handlers, "impact"), + review: vi.spyOn(handlers, "review"), + query_sqlite: vi.spyOn(handlers, "query_sqlite"), + refresh_index: vi.spyOn(handlers, "refresh_index"), + artifact_build: vi.spyOn(handlers, "artifact_build"), + }; const advertisedTools = listCodegraphMcpTools(); expect(advertisedTools.map((tool) => tool.name)).toEqual(MCP_TOOL_REGISTRY.map((tool) => tool.name)); expect(advertisedTools.some((tool) => "dispatch" in tool)).toBe(false); @@ -2809,5 +2833,105 @@ describe("MCP tool registry dispatch", () => { expect(message, tool.name).not.toMatch(/Unknown MCP tool|Invalid parameters for/i); } } + + for (const [name, spy] of Object.entries(handlerSpies)) { + expect(spy, name).toHaveBeenCalled(); + } + expect(handlerSpies.calls).toHaveBeenCalledWith(expect.objectContaining({ direction: "callers" }), undefined); + expect(handlerSpies.calls).toHaveBeenCalledWith(expect.objectContaining({ direction: "callees" }), undefined); + expect(handlerSpies.type_hierarchy).toHaveBeenCalledWith( + expect.objectContaining({ direction: "supertypes" }), + undefined, + ); + expect(handlerSpies.type_hierarchy).toHaveBeenCalledWith( + expect.objectContaining({ direction: "subtypes" }), + undefined, + ); + expect(handlerSpies.file_deps).toHaveBeenCalledWith(expect.objectContaining({ direction: "deps" }), undefined); + expect(handlerSpies.file_deps).toHaveBeenCalledWith(expect.objectContaining({ direction: "rdeps" }), undefined); + }); +}); + +describe("MCP refresh coalescing", () => { + it("serializes every queued request's requested warmup", async () => { + const root = await fs.mkdtemp(path.join(os.tmpdir(), "cg-mcp-refresh-coalesce-")); + await fs.writeFile(path.join(root, "auth.ts"), "export const ok = 1;\n", "utf8"); + const backingSession = createAgentSession({ root }); + const firstWarmupReached = Promise.withResolvers(); + const releaseFirstWarmup = Promise.withResolvers(); + const secondWarmupReached = Promise.withResolvers(); + const releaseSecondWarmup = Promise.withResolvers(); + const loadModes: Array<"skip" | "full"> = []; + let activeLoads = 0; + let maxActiveLoads = 0; + let fullWarmups = 0; + const session: AgentSession = { + ...backingSession, + loadProject: async (options) => { + const mode = options?.symbolGraph === "skip" ? "skip" : "full"; + loadModes.push(mode); + activeLoads += 1; + maxActiveLoads = Math.max(maxActiveLoads, activeLoads); + try { + if (mode === "skip") { + firstWarmupReached.resolve(); + await releaseFirstWarmup.promise; + } else { + fullWarmups += 1; + if (fullWarmups === 1) { + secondWarmupReached.resolve(); + await releaseSecondWarmup.promise; + } + } + return await backingSession.loadProject(options); + } finally { + activeLoads -= 1; + } + }, + }; + const handlers = createCodegraphMcpHandlers({ root, session }); + + const first = handlers.refresh_index({ warmup: "base" }); + await firstWarmupReached.promise; + const second = handlers.refresh_index({ warmup: "symbols" }); + const third = handlers.refresh_index({ warmup: "symbols" }); + releaseFirstWarmup.resolve(); + + try { + await secondWarmupReached.promise; + for (let turn = 0; turn < 4; turn += 1) await Promise.resolve(); + expect(loadModes).toEqual(["skip", "full"]); + expect(maxActiveLoads).toBe(1); + + releaseSecondWarmup.resolve(); + await expect(first).resolves.toEqual({ refreshed: true, warmup: "base" }); + await expect(second).resolves.toEqual({ refreshed: true, warmup: "symbols" }); + await expect(third).resolves.toEqual({ refreshed: true, warmup: "symbols" }); + expect(loadModes).toEqual(["skip", "full", "full"]); + expect(maxActiveLoads).toBe(1); + } finally { + releaseFirstWarmup.resolve(); + releaseSecondWarmup.resolve(); + } + }); +}); + +describe("MCP cancellation accounting", () => { + it("keeps a tool-call slot reserved until a cancelled operation settles", async () => { + const controller = new AbortController(); + const operation = Promise.withResolvers(); + let released = 0; + const pending = awaitMcpToolOperation(controller.signal, operation.promise, () => { + released += 1; + }); + + controller.abort(); + await expect(pending).rejects.toThrow("MCP tool call was cancelled."); + expect(released).toBe(0); + + operation.resolve("finished"); + await vi.waitFor(() => { + expect(released).toBe(1); + }); }); }); From 96879eda0d3d54aee24b04adbd1767e0f52c81c2 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sat, 15 Aug 2026 20:15:10 -0400 Subject: [PATCH 09/37] fix: hide legacy MCP aliases from tool listing --- src/mcp/tools.ts | 12 ++++++++++-- tests/mcp-server.test.ts | 8 +++++--- 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/src/mcp/tools.ts b/src/mcp/tools.ts index 119ea208..eaed1bda 100644 --- a/src/mcp/tools.ts +++ b/src/mcp/tools.ts @@ -150,7 +150,7 @@ export type McpToolDispatch = | { handler: "refresh_index" } | { handler: "artifact_build" }; -export type McpToolDefinition = Tool & { dispatch: McpToolDispatch }; +export type McpToolDefinition = Tool & { dispatch: McpToolDispatch; advertised?: boolean }; export const MCP_TOOL_REGISTRY: McpToolDefinition[] = [ { @@ -406,6 +406,7 @@ export const MCP_TOOL_REGISTRY: McpToolDefinition[] = [ { name: "callers", description: "Legacy alias for calls with direction callers.", + advertised: false, inputSchema: objectSchema( { handle: stringProperty, @@ -425,6 +426,7 @@ export const MCP_TOOL_REGISTRY: McpToolDefinition[] = [ { name: "callees", description: "Legacy alias for calls with direction callees.", + advertised: false, inputSchema: objectSchema( { handle: stringProperty, @@ -444,6 +446,7 @@ export const MCP_TOOL_REGISTRY: McpToolDefinition[] = [ { name: "supertypes", description: "Legacy alias for type_hierarchy with direction supertypes.", + advertised: false, inputSchema: objectSchema( { handle: stringProperty, @@ -462,6 +465,7 @@ export const MCP_TOOL_REGISTRY: McpToolDefinition[] = [ { name: "subtypes", description: "Legacy alias for type_hierarchy with direction subtypes.", + advertised: false, inputSchema: objectSchema( { handle: stringProperty, @@ -480,6 +484,7 @@ export const MCP_TOOL_REGISTRY: McpToolDefinition[] = [ { name: "deps", description: "Legacy alias for file_deps with direction deps.", + advertised: false, inputSchema: objectSchema( { file: dependencyFileProperty, @@ -498,6 +503,7 @@ export const MCP_TOOL_REGISTRY: McpToolDefinition[] = [ { name: "rdeps", description: "Legacy alias for file_deps with direction rdeps.", + advertised: false, inputSchema: objectSchema( { file: dependencyFileProperty, @@ -516,5 +522,7 @@ export const MCP_TOOL_REGISTRY: McpToolDefinition[] = [ ]; export function listCodegraphMcpTools(): Tool[] { - return MCP_TOOL_REGISTRY.map(({ dispatch: _dispatch, ...tool }) => tool); + return MCP_TOOL_REGISTRY.filter((tool) => tool.advertised !== false).map( + ({ dispatch: _dispatch, advertised: _advertised, ...tool }) => tool, + ); } diff --git a/tests/mcp-server.test.ts b/tests/mcp-server.test.ts index b8f5fe79..c1accbea 100644 --- a/tests/mcp-server.test.ts +++ b/tests/mcp-server.test.ts @@ -2759,7 +2759,7 @@ function normalizeSqlitePath(value: unknown): string { } describe("MCP tool registry dispatch", () => { - it("routes every advertised tool to a schema-valid handler", async () => { + it("routes every registered tool without advertising legacy aliases", async () => { const root = await fs.mkdtemp(path.join(os.tmpdir(), "cg-mcp-registry-")); await fs.writeFile(path.join(root, "auth.ts"), "export function ok(): number { return 1; }\n", "utf8"); runGit(root, ["init"]); @@ -2790,7 +2790,9 @@ describe("MCP tool registry dispatch", () => { artifact_build: vi.spyOn(handlers, "artifact_build"), }; const advertisedTools = listCodegraphMcpTools(); - expect(advertisedTools.map((tool) => tool.name)).toEqual(MCP_TOOL_REGISTRY.map((tool) => tool.name)); + expect(advertisedTools.map((tool) => tool.name)).toEqual( + MCP_TOOL_REGISTRY.filter((tool) => tool.advertised !== false).map((tool) => tool.name), + ); expect(advertisedTools.some((tool) => "dispatch" in tool)).toBe(false); const handle = "auth.ts::ok"; const toolInputs: Record> = { @@ -2823,7 +2825,7 @@ describe("MCP tool registry dispatch", () => { rdeps: { file: "auth.ts" }, }; - for (const tool of advertisedTools) { + for (const tool of MCP_TOOL_REGISTRY) { const input = toolInputs[tool.name]; expect(input, "missing valid input for " + tool.name).toBeDefined(); try { From a68bdae5d0a9cdc46a5303597bbccc2a54fbc4f6 Mon Sep 17 00:00:00 2001 From: codegraph Date: Sat, 15 Aug 2026 12:17:09 -0400 Subject: [PATCH 10/37] refactor: route omission counting through shared bound helper --- src/agent/explore.ts | 11 +++++++---- src/indexer/type-hierarchy.ts | 20 ++++++++++++-------- src/indexer/workspace-symbols.ts | 6 ++++-- tests/agent-explore.test.ts | 29 +++++++++++++++++++++++++++++ tests/type-hierarchy.test.ts | 32 ++++++++++++++++++++++++++++++++ tests/workspace-symbols.test.ts | 14 ++++++++++++++ 6 files changed, 98 insertions(+), 14 deletions(-) diff --git a/src/agent/explore.ts b/src/agent/explore.ts index 710953b6..8b67a795 100644 --- a/src/agent/explore.ts +++ b/src/agent/explore.ts @@ -2,6 +2,7 @@ import path from "node:path"; import type { AnalysisSummary } from "../analysisSummary.js"; import { getReverseDependencies, getShortestPath, type DependencyNode } from "../graphs/traversal.js"; import { defNodeId } from "../graphs/symbol-graph.js"; +import { boundList } from "../presentation/bounds.js"; import type { BuildOptions } from "../indexer/types.js"; import { listCandidateTestFiles } from "../impact/context.js"; import { fileIdentityKey, normalizePath, toProjectDisplayPath } from "../util/paths.js"; @@ -483,11 +484,12 @@ function collectBlastRadius( const summaries: AgentExploreBlastRadiusSummary[] = []; for (const file of anchorFiles.slice(0, entryLimit)) { const dependencies = getReverseDependencies(snapshot.fileGraph, file, { limit: dependencyLimit + 1, depth: 2 }); - const visible = dependencies.slice(0, dependencyLimit).map((dependency) => formatDependency(snapshot, dependency)); + const boundedDependencies = boundList(dependencies, dependencyLimit); + const visible = boundedDependencies.items.map((dependency) => formatDependency(snapshot, dependency)); summaries.push({ file: toProjectDisplayPath(snapshot.root, file), reverseDependencies: visible, - omittedLowerBound: Math.max(0, dependencies.length - dependencyLimit), + omittedLowerBound: boundedDependencies.omitted, }); } return summaries; @@ -511,9 +513,10 @@ function collectCandidateTests( maxCandidates: snapshot.index.byFile.size, projectRoot: snapshot.root, }); + const boundedCandidates = boundList(candidates, limit); return { - items: candidates.slice(0, limit).map((candidate) => toProjectDisplayPath(snapshot.root, candidate.file)), - omittedCount: Math.max(0, candidates.length - limit), + items: boundedCandidates.items.map((candidate) => toProjectDisplayPath(snapshot.root, candidate.file)), + omittedCount: boundedCandidates.omitted, }; } function collectFollowUps( diff --git a/src/indexer/type-hierarchy.ts b/src/indexer/type-hierarchy.ts index 71feeb63..10b0223f 100644 --- a/src/indexer/type-hierarchy.ts +++ b/src/indexer/type-hierarchy.ts @@ -1,4 +1,5 @@ import type { SymbolEdge, SymbolGraph, SymbolNode } from "../graphs/symbol-graph.js"; +import { boundList } from "../presentation/bounds.js"; import { resolveSymbolId } from "./symbols.js"; import type { ProjectIndex } from "./types.js"; @@ -106,12 +107,13 @@ export function findTypeHierarchy( } relations.sort((left, right) => compareRelations(graph, left, right)); + const boundedRelations = boundList(relations, limit); return { status: "ok", targetId, direction, - relations: relations.slice(0, limit), - omitted: Math.max(0, relations.length - limit), + relations: boundedRelations.items, + omitted: boundedRelations.omitted, limit, }; } @@ -137,11 +139,12 @@ export function findImplementations( }; } const matches = collectTypeImplementations(graph, hierarchy, targetId, rootRelations); + const boundedMatches = boundList(matches, limit); return { status: "ok", targetId, - implementations: matches.slice(0, limit), - omitted: Math.max(0, matches.length - limit), + implementations: boundedMatches.items, + omitted: boundedMatches.omitted, ambiguous: 0, unresolved: [], limit, @@ -229,14 +232,15 @@ export function findImplementations( } const sortedUnresolved = unresolved.sort((left, right) => left.symbolId.localeCompare(right.symbolId)); const matches = [...memberMatches.values()].sort((left, right) => compareImplementationMatches(graph, left, right)); - const truncated = Math.max(0, matches.length - limit); + const boundedMatches = boundList(matches, limit); + const boundedUnresolved = boundList(sortedUnresolved, limit); return { status: "ok", targetId, - implementations: matches.slice(0, limit), - omitted: truncated + ambiguous, + implementations: boundedMatches.items, + omitted: boundedMatches.omitted + ambiguous, ambiguous, - unresolved: sortedUnresolved.slice(0, limit), + unresolved: boundedUnresolved.items, limit, }; } diff --git a/src/indexer/workspace-symbols.ts b/src/indexer/workspace-symbols.ts index d2e7b01c..906d185e 100644 --- a/src/indexer/workspace-symbols.ts +++ b/src/indexer/workspace-symbols.ts @@ -5,6 +5,7 @@ import { ensureParsedContext } from "./parse-context.js"; import { getCachedScope } from "./navigation-references.js"; import { resolveImported } from "./navigation-resolve.js"; import { defNodeId } from "../graphs/symbol-graph.js"; +import { boundList } from "../presentation/bounds.js"; import type { ImportBinding, ProjectIndex, SymbolDef, SymbolKind } from "./types.js"; export const DEFAULT_WORKSPACE_SYMBOL_LIMIT = 50; @@ -100,12 +101,13 @@ export async function workspaceSymbols( } ranked.sort(compareRankedCandidates); - const symbols = ranked.slice(0, limit).map(({ candidate }) => candidate); + const boundedRanked = boundList(ranked, limit); + const symbols = boundedRanked.items.map(({ candidate }) => candidate); return { query, symbols, totalCandidates: ranked.length, - omitted: Math.max(0, ranked.length - symbols.length), + omitted: boundedRanked.omitted, limit, omittedImports, importScanFailures, diff --git a/tests/agent-explore.test.ts b/tests/agent-explore.test.ts index e058018d..0f002e4b 100644 --- a/tests/agent-explore.test.ts +++ b/tests/agent-explore.test.ts @@ -841,4 +841,33 @@ describe("agent explore", () => { expect(readArray(response.blastRadius, "blastRadius")).toHaveLength(1); expect(response.freshness).toBeTypeOf("object"); }); + it("pins omission counts at and just past the limit for candidate tests and blast radius", async () => { + const root = await mkExploreRepo(); + await writeFile(root, "tests/auth.test.ts", "import { validateUser } from '../src/auth';\nvalidateUser('bob');\n"); + await writeFile(root, "tests/auth-spec.test.ts", "import { validateUser } from '../src/auth';\nvalidateUser('carol');\n"); + + const exploreAll = await exploreCodegraph({ root, query: "validateUser" }); + expect(exploreAll.candidateTests.length).toBeGreaterThanOrEqual(2); + expect(exploreAll.omittedCounts.candidateTests).toBe(0); + + const spy = vi.spyOn(impactContext, "listCandidateTestFiles").mockReturnValue([ + { file: path.join(root, "tests/routes.test.ts"), reasons: [] }, + { file: path.join(root, "tests/auth.test.ts"), reasons: [] }, + { file: path.join(root, "tests/auth-spec.test.ts"), reasons: [] }, + ]); + + try { + const atLimitResponse = await exploreCodegraph({ root, query: "validateUser" }); + expect(atLimitResponse.candidateTests).toHaveLength(3); + expect(atLimitResponse.omittedCounts.candidateTests).toBe(0); + } finally { + spy.mockRestore(); + } + + const authExplore = await exploreCodegraph({ root, query: "src/db.ts" }); + const dbBlast = authExplore.blastRadius.find((entry) => entry.file === "src/db.ts"); + expect(dbBlast).toBeDefined(); + expect(dbBlast!.reverseDependencies.length).toBeGreaterThanOrEqual(1); + expect(dbBlast!.omittedLowerBound).toBe(0); + }); }); diff --git a/tests/type-hierarchy.test.ts b/tests/type-hierarchy.test.ts index f38013c2..9ccf96fd 100644 --- a/tests/type-hierarchy.test.ts +++ b/tests/type-hierarchy.test.ts @@ -205,4 +205,36 @@ describe("type hierarchy", () => { reason: expect.stringContaining("abstract"), }); }); + it("pins omission counts at and just past the limit for type hierarchy and implementations", async () => { + const { index, graph, byName } = await hierarchyFixture(); + const specialized = byName.get("SpecializedWorker"); + expect(specialized).toBeDefined(); + + const atSuperLimit = findTypeHierarchy(graph, specialized!.id, "super", { depth: 3, limit: 3 }); + expect(atSuperLimit).toMatchObject({ status: "ok", omitted: 0 }); + if (atSuperLimit.status === "ok") { + expect(atSuperLimit.relations).toHaveLength(3); + } + + const pastSuperLimit = findTypeHierarchy(graph, specialized!.id, "super", { depth: 3, limit: 2 }); + expect(pastSuperLimit).toMatchObject({ status: "ok", omitted: 1 }); + if (pastSuperLimit.status === "ok") { + expect(pastSuperLimit.relations).toHaveLength(2); + } + + const service = byName.get("Service"); + expect(service).toBeDefined(); + + const atImplLimit = findImplementations(index, graph, service!.id, { limit: 2 }); + expect(atImplLimit).toMatchObject({ status: "ok", omitted: 0 }); + if (atImplLimit.status === "ok") { + expect(atImplLimit.implementations).toHaveLength(2); + } + + const pastImplLimit = findImplementations(index, graph, service!.id, { limit: 1 }); + expect(pastImplLimit).toMatchObject({ status: "ok", omitted: 1 }); + if (pastImplLimit.status === "ok") { + expect(pastImplLimit.implementations).toHaveLength(1); + } + }); }); diff --git a/tests/workspace-symbols.test.ts b/tests/workspace-symbols.test.ts index 202a10d7..109dea6f 100644 --- a/tests/workspace-symbols.test.ts +++ b/tests/workspace-symbols.test.ts @@ -412,4 +412,18 @@ describe("workspace symbol lookup", () => { tool_workspaceSymbols(root, { query: "Service" }, { session, buildOptions: { cache: "off" } }), ).rejects.toThrow("cannot combine a prebuilt session with buildOptions"); }); + it("pins omission counts at and just past the limit for workspace symbols", async () => { + const all = await workspaceSymbols(index, { query: "Service", limit: 50 }); + const total = all.symbols.length; + expect(total).toBeGreaterThanOrEqual(2); + expect(all.omitted).toBe(0); + + const atLimit = await workspaceSymbols(index, { query: "Service", limit: total }); + expect(atLimit.symbols).toHaveLength(total); + expect(atLimit.omitted).toBe(0); + + const pastLimit = await workspaceSymbols(index, { query: "Service", limit: total - 1 }); + expect(pastLimit.symbols).toHaveLength(total - 1); + expect(pastLimit.omitted).toBe(1); + }); }); From f3970ecbffbfcc27c38665a9a7c24672e9673c04 Mon Sep 17 00:00:00 2001 From: codegraph Date: Sat, 15 Aug 2026 12:17:10 -0400 Subject: [PATCH 11/37] test: cover MCP session resource bounds and teardown --- tests/agent-search.test.ts | 43 ++++++++++ tests/agent-session.test.ts | 29 +++++++ tests/mcp-server.test.ts | 162 ++++++++++++++++++++++++++++++++++++ tests/query-index.test.ts | 26 +++++- tests/viewer.test.ts | 48 +++++++++++ 5 files changed, 307 insertions(+), 1 deletion(-) diff --git a/tests/agent-search.test.ts b/tests/agent-search.test.ts index 5228b884..a7db5fcb 100644 --- a/tests/agent-search.test.ts +++ b/tests/agent-search.test.ts @@ -854,4 +854,47 @@ describe("agent search", () => { expect(response.results).toEqual([]); }); + it("coalesces concurrent queries and evicts oldest entries when session search cache exceeds max entries", async () => { + const root = await mkRepo(); + const session = createAgentSession({ root }); + try { + await session.loadProject(); + + // 1. Identical concurrent queries coalesce to one in-flight promise + const p1 = searchCodegraphWithSession(session, { root, query: "validateUser", mode: "symbol" }); + const p2 = searchCodegraphWithSession(session, { root, query: "validateUser", mode: "symbol" }); + const [r1, r2] = await Promise.all([p1, p2]); + expect(r1).toBe(r2); + + // 2. Issuing more unique queries than the 100 cap evicts the oldest entry + const firstResult = r1; + const recentResults = []; + for (let i = 0; i < 100; i += 1) { + const res = await searchCodegraphWithSession(session, { + root, + query: "needleUniqueQuery" + i, + mode: "symbol", + }); + recentResults.push(res); + } + + // Re-querying the oldest entry ("validateUser") produces a new result because it was evicted + const reQueryFirst = await searchCodegraphWithSession(session, { + root, + query: "validateUser", + mode: "symbol", + }); + expect(reQueryFirst).not.toBe(firstResult); + + // Re-querying the most recent entry ("needleUniqueQuery99") returns the cached result + const reQueryLatest = await searchCodegraphWithSession(session, { + root, + query: "needleUniqueQuery99", + mode: "symbol", + }); + expect(reQueryLatest).toBe(recentResults[99]); + } finally { + session.invalidate(); + } + }); }); diff --git a/tests/agent-session.test.ts b/tests/agent-session.test.ts index f5361106..7405d5e7 100644 --- a/tests/agent-session.test.ts +++ b/tests/agent-session.test.ts @@ -1,3 +1,5 @@ +import { disposeSessionQueryIndex, ensureSessionQueryIndex } from "../src/agent/query-index/sessionStore.js"; +import * as updateModule from "../src/agent/query-index/update.js"; import fs from "node:fs/promises"; import { brotliCompressSync, brotliDecompressSync, constants as zlibConstants } from "node:zlib"; import { createHash } from "node:crypto"; @@ -938,4 +940,31 @@ describe("agent session", () => { dateSpy.mockRestore(); } }); + +describe("query index sessionStore generation retries (S12)", () => { + it("bounds query index generation retries under sustained invalidation and surfaces a clear error", async () => { + const root = await mkRepo(); + const session = createAgentSession({ root }); + const snapshot = await session.loadProject(); + + let attempts = 0; + const realEnsureQueryIndex = updateModule.ensureQueryIndex; + const ensureQueryIndexSpy = vi.spyOn(updateModule, "ensureQueryIndex").mockImplementation(async (snap) => { + attempts += 1; + const res = await realEnsureQueryIndex(snap); + disposeSessionQueryIndex(session); + return res; + }); + + try { + await expect(ensureSessionQueryIndex(session, snapshot)).rejects.toThrow( + /Query index generation changed repeatedly while loading/i, + ); + expect(attempts).toBe(3); + } finally { + ensureQueryIndexSpy.mockRestore(); + } + }); }); + +}); \ No newline at end of file diff --git a/tests/mcp-server.test.ts b/tests/mcp-server.test.ts index c1accbea..32f2aaf1 100644 --- a/tests/mcp-server.test.ts +++ b/tests/mcp-server.test.ts @@ -1,3 +1,5 @@ +import { registerSessionInvalidationHook } from "../src/agent/sessionLifecycle.js"; +import { ensureSessionQueryIndex } from "../src/agent/query-index/sessionStore.js"; import fs from "node:fs/promises"; import { request as httpRequest, type IncomingMessage } from "node:http"; import os from "node:os"; @@ -2918,6 +2920,69 @@ describe("MCP refresh coalescing", () => { }); }); + +describe("MCP session teardown regressions (S2)", () => { + it("closes sidecar query index handle and runs session invalidation hooks on server close", async () => { + const root = await fs.mkdtemp(path.join(os.tmpdir(), "cg-mcp-s2-teardown-")); + await fs.writeFile( + path.join(root, "auth.ts"), + "export function validateUser(token: string) { return !!token; }\n", + "utf8", + ); + const session = createAgentSession({ root }); + let invalidationHookRan = false; + registerSessionInvalidationHook(session, () => { + invalidationHookRan = true; + }); + + const httpServer = await startCodegraphMcpHttpServer({ + root, + host: "127.0.0.1", + port: 0, + session, + }); + + try { + const initialize = await postMcpJson(httpServer.url, { + jsonrpc: "2.0", + id: 1, + method: "initialize", + params: { + protocolVersion: "2025-11-25", + capabilities: {}, + clientInfo: { name: "codegraph-s2-test", version: "1.0.0" }, + }, + }); + const sessionId = initialize.response.headers.get("mcp-session-id"); + expect(sessionId).toBeTruthy(); + + const searchCall = await postMcpJson( + httpServer.url, + { + jsonrpc: "2.0", + id: 2, + method: "tools/call", + params: { name: "search", arguments: { query: "validateUser", mode: "hybrid" } }, + }, + sessionId ?? undefined, + ); + expect(searchCall.response.status).toBe(200); + + const snapshot = await session.loadProject(); + const handle = await ensureSessionQueryIndex(session, snapshot); + expect(handle.store?.closed).toBe(false); + expect(invalidationHookRan).toBe(false); + + await httpServer.close(); + + expect(invalidationHookRan).toBe(true); + expect(handle.store?.closed).toBe(true); + } finally { + await httpServer.close(); + } + }); +}); + describe("MCP cancellation accounting", () => { it("keeps a tool-call slot reserved until a cancelled operation settles", async () => { const controller = new AbortController(); @@ -2937,3 +3002,100 @@ describe("MCP cancellation accounting", () => { }); }); }); + +describe("MCP transport isolation regressions (S8)", () => { + it("preserves session and completes concurrent calls when one response connection is forcefully closed", async () => { + const root = await fs.mkdtemp(path.join(os.tmpdir(), "cg-mcp-s8-transport-")); + await fs.writeFile( + path.join(root, "auth.ts"), + "export function validateUser(token: string) { return !!token; }\nexport function secondarySymbol() { return true; }\n", + "utf8", + ); + const httpServer = await startCodegraphMcpHttpServer({ + root, + host: "127.0.0.1", + port: 0, + }); + + try { + const initialize = await postMcpJson(httpServer.url, { + jsonrpc: "2.0", + id: 1, + method: "initialize", + params: { + protocolVersion: "2025-11-25", + capabilities: {}, + clientInfo: { name: "codegraph-s8-test", version: "1.0.0" }, + }, + }); + const sessionId = initialize.response.headers.get("mcp-session-id"); + expect(sessionId).toBeTruthy(); + if (!sessionId) throw new Error("Missing sessionId"); + + const endpoint = new URL(httpServer.url); + const call1Payload = JSON.stringify({ + jsonrpc: "2.0", + id: 2, + method: "tools/call", + params: { name: "workspace_symbols", arguments: { query: "validateUser" } }, + }); + + const call1Closed = Promise.withResolvers(); + const req1 = httpRequest({ + hostname: endpoint.hostname, + port: endpoint.port, + path: endpoint.pathname, + method: "POST", + headers: { + accept: "application/json, text/event-stream", + "content-type": "application/json", + "content-length": String(Buffer.byteLength(call1Payload)), + "mcp-session-id": sessionId, + }, + }); + req1.on("error", () => { + call1Closed.resolve(); + }); + req1.on("close", () => { + call1Closed.resolve(); + }); + req1.write(call1Payload); + req1.destroy(new Error("Forced client disconnect")); + + const call2Promise = postMcpJson( + httpServer.url, + { + jsonrpc: "2.0", + id: 3, + method: "tools/call", + params: { name: "workspace_symbols", arguments: { query: "secondarySymbol" } }, + }, + sessionId, + ); + + await call1Closed.promise; + const call2 = await call2Promise; + expect(call2.response.status).toBe(200); + expect(readToolJsonResult(call2.payload).symbols).toEqual([ + expect.objectContaining({ name: "secondarySymbol" }), + ]); + + const call3 = await postMcpJson( + httpServer.url, + { + jsonrpc: "2.0", + id: 4, + method: "tools/call", + params: { name: "workspace_symbols", arguments: { query: "validateUser" } }, + }, + sessionId, + ); + expect(call3.response.status).toBe(200); + expect(readToolJsonResult(call3.payload).symbols).toEqual([ + expect.objectContaining({ name: "validateUser" }), + ]); + } finally { + await httpServer.close(); + } + }); +}); diff --git a/tests/query-index.test.ts b/tests/query-index.test.ts index 3d34839d..fd27e985 100644 --- a/tests/query-index.test.ts +++ b/tests/query-index.test.ts @@ -4,7 +4,8 @@ import path from "node:path"; import { afterEach, describe, expect, it, vi } from "vitest"; import { searchCodegraphWithSession, type AgentSearchResponse } from "../src/agent/search.js"; import { createAgentSession, type AgentProjectSnapshot, type AgentSession } from "../src/agent/session.js"; -import { disposeSessionQueryIndex } from "../src/agent/query-index/sessionStore.js"; +import { disposeSessionQueryIndex, ensureSessionQueryIndex } from "../src/agent/query-index/sessionStore.js"; +import * as updateModule from "../src/agent/query-index/update.js"; import { resolveQueryIndexPaths, resolveQueryIndexSourcePath } from "../src/agent/query-index/paths.js"; import { expectedQueryIndexVersionMetadata, probeQueryIndexSqliteSupport } from "../src/agent/query-index/schema.js"; import { SqliteDatabase } from "../src/sqlite-driver.js"; @@ -868,4 +869,27 @@ describe("persistent query index", () => { expect(response.results.some((result) => result.file === "src/auth.ts")).toBe(true); await expect(fs.stat(path.join(root, ".codegraph-cache"))).rejects.toMatchObject({ code: "ENOENT" }); }); + it("bounds query index generation retries under sustained invalidation and surfaces a clear error", async () => { + const root = await createRepo(); + const session = createSession(root); + const snapshot = await session.loadProject(); + + let attempts = 0; + const realEnsureQueryIndex = updateModule.ensureQueryIndex; + const ensureQueryIndexSpy = vi.spyOn(updateModule, "ensureQueryIndex").mockImplementation(async (snap) => { + attempts += 1; + const res = await realEnsureQueryIndex(snap); + disposeSessionQueryIndex(session); + return res; + }); + + try { + await expect(ensureSessionQueryIndex(session, snapshot)).rejects.toThrow( + /Query index generation changed repeatedly while loading/i, + ); + expect(attempts).toBe(3); + } finally { + ensureQueryIndexSpy.mockRestore(); + } + }); }); diff --git a/tests/viewer.test.ts b/tests/viewer.test.ts index 5e6a2f7f..b88863b0 100644 --- a/tests/viewer.test.ts +++ b/tests/viewer.test.ts @@ -278,4 +278,52 @@ describe("viewer server", () => { expect(() => createViewerServer({ graph: escapedGraph, root })).toThrow(/outside project root/i); }); + test("returns 500 when statSync or fstatSync throws during GET and continues serving subsequent requests", async () => { + const { root, graphPath } = await createViewerFixture(); + const server = await startViewerServer({ graph: graphPath, port: 0, root }); + servers.push(server.server); + + // 1. Test fstatSync throwing during GET /graph.json + let throwFstat = true; + const originalFstatSync = fs.fstatSync; + const fstatSpy = vi.spyOn(fs, "fstatSync").mockImplementation((...args) => { + if (throwFstat) { + throw new Error("Simulated filesystem fstatSync error"); + } + return originalFstatSync(...args); + }); + + try { + const firstFstatResponse = await request(server.server, "/graph.json"); + expect(firstFstatResponse.statusCode).toBe(500); + + throwFstat = false; + const secondFstatResponse = await request(server.server, "/graph.json"); + expect(secondFstatResponse.statusCode).toBe(200); + expect(secondFstatResponse.body).toContain('"nodes":[]'); + } finally { + fstatSpy.mockRestore(); + } + + // 2. Test statSync throwing during GET / + let throwStat = true; + const originalStatSync = fs.statSync; + const statSpy = vi.spyOn(fs, "statSync").mockImplementation((...args) => { + if (throwStat) { + throw new Error("Simulated filesystem statSync error"); + } + return originalStatSync(...args); + }); + + try { + const firstStatResponse = await request(server.server, "/"); + expect(firstStatResponse.statusCode).toBe(500); + + throwStat = false; + const secondStatResponse = await request(server.server, "/"); + expect(secondStatResponse.statusCode).toBe(200); + } finally { + statSpy.mockRestore(); + } + }); }); From 6328bbf7d6c84a4b8136abc7e4ec1493853db927 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sat, 15 Aug 2026 23:01:11 -0400 Subject: [PATCH 12/37] fix: bound MCP HTTP and SQLite execution --- docs/library-api.md | 2 +- docs/mcp.md | 2 +- src/mcp/http.ts | 69 ++++++++++++++++++------------- src/mcp/server.ts | 4 +- src/sqlite/query.ts | 60 +++++++++++++++++++++++---- src/sqlite/rawQueryWorker.ts | 33 +++++++++++++++ src/sqlite/rawQueryWorkerPool.ts | 43 +++++++++++++++++++ tests/mcp-server.test.ts | 67 ++++++++++++++++++++++++++---- tests/sqlite-query-bounds.test.ts | 52 ++++++++++++++++++++++- 9 files changed, 283 insertions(+), 49 deletions(-) create mode 100644 src/sqlite/rawQueryWorker.ts create mode 100644 src/sqlite/rawQueryWorkerPool.ts diff --git a/docs/library-api.md b/docs/library-api.md index 102575b4..5e795034 100644 --- a/docs/library-api.md +++ b/docs/library-api.md @@ -743,7 +743,7 @@ const result = await queryGraphSqliteRaw( console.log(result.columns, result.rows); ``` -`queryGraphSqliteRaw()` is intentionally read-only. It accepts result-producing statements such as `SELECT` and `PRAGMA` and rejects mutating SQL. Pass `{ maxRows }` to bound raw result rows. +`queryGraphSqliteRaw()` is intentionally read-only. It accepts result-producing statements such as `SELECT` and `PRAGMA` and rejects mutating SQL. Its defaults bound rows, cells, response bytes, and execution to 10 seconds; callers can further tighten `{ maxRows, maxBytes, maxCellBytes, deadlineMs }`. ## SQL artifact facts diff --git a/docs/mcp.md b/docs/mcp.md index ec37bf62..4191e98d 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -219,7 +219,7 @@ An MCP `explore` request whose entire query resolves to an indexed project-relat - Tool calls do not accept per-request root overrides. - Tools are read-only by default. - `artifact_build` requires `--allow-build` and a fresh or auto-refreshed MCP index. -- `query_sqlite` rejects mutating SQL, recursive queries, synthetic payload functions, and stale artifact queries it cannot refresh safely. +- `query_sqlite` rejects mutating SQL, recursive queries, synthetic payload functions, and stale artifact queries it cannot refresh safely. Each query has a 10-second execution deadline. - `get_file` rejects raw reads and structural text-config summaries over the 16 MiB input limit. Accepted reads use separate output-page bounds from `maxBytes`, `offset`, and `limit`; binary input is rejected, and sensitive formats require `allowSensitive: true` for raw values. - SQLite responses are row- and byte-bounded. diff --git a/src/mcp/http.ts b/src/mcp/http.ts index 633963ca..c1870fcb 100644 --- a/src/mcp/http.ts +++ b/src/mcp/http.ts @@ -28,39 +28,52 @@ export async function readJsonRequestBody( return { status: "too_large" }; } - const chunks: Buffer[] = []; - let bytes = 0; - let timedOut = false; - const deadline = setTimeout(() => { - timedOut = true; - request.destroy(); - }, timeoutMs); - deadline.unref?.(); - try { - for await (const chunk of request) { + return await new Promise((resolve) => { + const chunks: Buffer[] = []; + let bytes = 0; + let settled = false; + const deadline = setTimeout(() => settle({ status: "timeout" }, true), timeoutMs); + deadline.unref?.(); + + const cleanup = (): void => { + clearTimeout(deadline); + request.off("data", onData); + request.off("end", onEnd); + request.off("error", onFailure); + request.off("aborted", onFailure); + }; + const settle = (result: ParsedJsonBody, drain: boolean): void => { + if (settled) return; + settled = true; + cleanup(); + if (drain) request.resume(); + resolve(result); + }; + const onData = (chunk: string | Buffer): void => { const buffer = typeof chunk === "string" ? Buffer.from(chunk) : chunk; bytes += buffer.byteLength; if (bytes > maxBytes) { - request.resume(); - return { status: "too_large" }; + settle({ status: "too_large" }, true); + return; } chunks.push(buffer); - } - } catch { - if (timedOut) return { status: "timeout" }; - return { status: "invalid_json" }; - } finally { - clearTimeout(deadline); - } - if (timedOut) return { status: "timeout" }; - - const rawBody = Buffer.concat(chunks).toString("utf8"); - try { - const body: unknown = rawBody.length ? JSON.parse(rawBody) : null; - return { status: "ok", body }; - } catch { - return { status: "invalid_json" }; - } + }; + const onEnd = (): void => { + const rawBody = Buffer.concat(chunks).toString("utf8"); + try { + const body: unknown = rawBody.length ? JSON.parse(rawBody) : null; + settle({ status: "ok", body }, false); + } catch { + settle({ status: "invalid_json" }, false); + } + }; + const onFailure = (): void => settle({ status: "invalid_json" }, true); + + request.on("data", onData); + request.once("end", onEnd); + request.once("error", onFailure); + request.once("aborted", onFailure); + }); } export function emptyAllowedHostHeaderRules(): AllowedHostHeaderRules { diff --git a/src/mcp/server.ts b/src/mcp/server.ts index 9e33d492..45505d28 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -72,7 +72,6 @@ import { DEFAULT_WORKSPACE_SYMBOL_LIMIT, MAX_WORKSPACE_SYMBOL_LIMIT } from "../i import type { BuildOptions, FindReferencesResult, GoToResult } from "../indexer/types.js"; import { assertMcpSqliteQueryResourceBounded, - boundRawSqlResult, DEFAULT_SQLITE_BYTE_LIMIT, normalizeSqliteRowLimit, } from "./sqliteGuard.js"; @@ -1007,8 +1006,9 @@ function createCodegraphMcpHandlersForSession( } const result = await queryGraphSqliteRaw(realSqlitePath, request.query, request.params ?? [], { maxRows: normalizeSqliteRowLimit(request.limit), + maxBytes: DEFAULT_SQLITE_BYTE_LIMIT, }); - return { ...boundRawSqlResult(result, DEFAULT_SQLITE_BYTE_LIMIT), freshness: artifactFreshness }; + return { ...result, truncated: Boolean(result.truncated), freshness: artifactFreshness }; }, refresh_index: async (request) => { diff --git a/src/sqlite/query.ts b/src/sqlite/query.ts index d1e740ff..6cf062a7 100644 --- a/src/sqlite/query.ts +++ b/src/sqlite/query.ts @@ -8,13 +8,22 @@ import { MAX_SQLITE_ROW_LIMIT, normalizeSqliteRowLimit, } from "./rowBounds.js"; +import { + resolveRawSqlQueryWorkerPath, + runRawSqlQueryInWorker, + SqliteQueryDeadlineExceededError, +} from "./rawQueryWorkerPool.js"; export { queryGraphSqlite } from "./canned-query.js"; +export { SqliteQueryDeadlineExceededError }; + +export const DEFAULT_SQLITE_QUERY_DEADLINE_MS = 10_000; export type QueryGraphSqliteRawOptions = { maxRows?: number | undefined; maxBytes?: number | undefined; maxCellBytes?: number | undefined; + deadlineMs?: number | undefined; }; export async function queryGraphSqliteRaw( @@ -22,21 +31,47 @@ export async function queryGraphSqliteRaw( sql: string, params: Array = [], options?: QueryGraphSqliteRawOptions, +): Promise { + const maxRows = normalizeSqliteRowLimit(options?.maxRows ?? MAX_SQLITE_ROW_LIMIT); + const maxBytes = options?.maxBytes ?? DEFAULT_SQLITE_BYTE_LIMIT; + const maxCellBytes = options?.maxCellBytes ?? MAX_SQLITE_CELL_BYTES; + const deadlineMs = options?.deadlineMs ?? DEFAULT_SQLITE_QUERY_DEADLINE_MS; + + try { + resolveRawSqlQueryWorkerPath(); + } catch { + return await queryGraphSqliteRawInProcessBounded(outputPath, sql, params, { + maxRows, + maxBytes, + maxCellBytes, + deadlineMs, + }); + } + + return await runRawSqlQueryInWorker({ outputPath, sql, params, maxRows, maxBytes, maxCellBytes }, deadlineMs); +} + +async function queryGraphSqliteRawInProcessBounded( + outputPath: string, + sql: string, + params: Array, + bounds: { maxRows: number; maxBytes: number; maxCellBytes: number; deadlineMs: number }, ): Promise { return await withReadOnlySqliteDatabase(outputPath, (db) => { try { const stmt = db.prepare(sql); assertReadOnlyQueryStatement(stmt); const columns = stmt.columns().map((col) => col.name); - const maxRows = normalizeSqliteRowLimit(options?.maxRows ?? MAX_SQLITE_ROW_LIMIT); - const maxBytes = options?.maxBytes ?? DEFAULT_SQLITE_BYTE_LIMIT; - const maxCellBytes = options?.maxCellBytes ?? MAX_SQLITE_CELL_BYTES; - - // Always stream via iterate so per-cell and cumulative budgets apply before append. - return collectBoundedRawSqlRows(columns, stmt.raw().iterate(params) as Iterable>, { - maxRows, - maxBytes, - maxCellBytes, + const deadlineAt = Date.now() + bounds.deadlineMs; + const rows = withPerRowDeadline( + stmt.raw().iterate(params) as Iterable>, + deadlineAt, + bounds.deadlineMs, + ); + return collectBoundedRawSqlRows(columns, rows, { + maxRows: bounds.maxRows, + maxBytes: bounds.maxBytes, + maxCellBytes: bounds.maxCellBytes, }); } catch (error) { if (isReadOnlySqliteError(error)) { @@ -46,3 +81,10 @@ export async function queryGraphSqliteRaw( } }); } + +function* withPerRowDeadline(rows: Iterable, deadlineAt: number, deadlineMs: number): Generator { + for (const row of rows) { + if (Date.now() > deadlineAt) throw new SqliteQueryDeadlineExceededError(deadlineMs); + yield row; + } +} diff --git a/src/sqlite/rawQueryWorker.ts b/src/sqlite/rawQueryWorker.ts new file mode 100644 index 00000000..dc74a7a7 --- /dev/null +++ b/src/sqlite/rawQueryWorker.ts @@ -0,0 +1,33 @@ +import { isReadOnlySqliteError } from "../sqlite-driver.js"; +import type { RawSqlResult } from "./types.js"; +import { assertReadOnlyQueryStatement, withReadOnlySqliteDatabase } from "./database.js"; +import { collectBoundedRawSqlRows } from "./rowBounds.js"; + +export type RawQueryWorkerTask = { + outputPath: string; + sql: string; + params: Array; + maxRows: number; + maxBytes: number | undefined; + maxCellBytes: number | undefined; +}; + +export default async function runRawQueryWorkerTask(task: RawQueryWorkerTask): Promise { + return await withReadOnlySqliteDatabase(task.outputPath, (db) => { + try { + const statement = db.prepare(task.sql); + assertReadOnlyQueryStatement(statement); + const columns = statement.columns().map((column) => column.name); + return collectBoundedRawSqlRows(columns, statement.raw().iterate(task.params) as Iterable>, { + maxRows: task.maxRows, + maxBytes: task.maxBytes, + maxCellBytes: task.maxCellBytes, + }); + } catch (error) { + if (isReadOnlySqliteError(error)) { + throw new Error("Raw SQLite queries must be read-only result-producing statements such as SELECT or PRAGMA."); + } + throw error; + } + }); +} diff --git a/src/sqlite/rawQueryWorkerPool.ts b/src/sqlite/rawQueryWorkerPool.ts new file mode 100644 index 00000000..e40432f4 --- /dev/null +++ b/src/sqlite/rawQueryWorkerPool.ts @@ -0,0 +1,43 @@ +import fs from "node:fs"; +import path from "node:path"; +import { fileURLToPath } from "node:url"; +import { Piscina } from "piscina"; +import { findPackageRoot } from "../util/packageInfo.js"; +import type { RawSqlResult } from "./types.js"; +import type { RawQueryWorkerTask } from "./rawQueryWorker.js"; + +export class SqliteQueryDeadlineExceededError extends Error { + constructor(deadlineMs: number) { + super(`SQLite query exceeded its ${deadlineMs}ms execution budget and was terminated.`); + this.name = "SqliteQueryDeadlineExceededError"; + } +} + +export function resolveRawSqlQueryWorkerPath(): string { + const selfDirectory = path.dirname(fileURLToPath(import.meta.url)); + const sibling = path.resolve(selfDirectory, "rawQueryWorker.js"); + if (fs.existsSync(sibling)) return sibling; + const packageRoot = findPackageRoot(selfDirectory); + const compiled = path.join(packageRoot, "dist", "sqlite", "rawQueryWorker.js"); + if (fs.existsSync(compiled)) return compiled; + throw new Error(`Raw SQLite query worker file not found: ${compiled}`); +} + +export async function runRawSqlQueryInWorker(task: RawQueryWorkerTask, deadlineMs: number): Promise { + const pool = new Piscina({ + filename: resolveRawSqlQueryWorkerPath(), + minThreads: 1, + maxThreads: 1, + idleTimeout: 5_000, + }); + try { + return (await pool.run(task, { signal: AbortSignal.timeout(deadlineMs) })) as RawSqlResult; + } catch (error) { + if (error instanceof Error && error.name === "AbortError") { + throw new SqliteQueryDeadlineExceededError(deadlineMs); + } + throw error; + } finally { + void pool.destroy().catch(() => {}); + } +} diff --git a/tests/mcp-server.test.ts b/tests/mcp-server.test.ts index 32f2aaf1..f032258c 100644 --- a/tests/mcp-server.test.ts +++ b/tests/mcp-server.test.ts @@ -1158,6 +1158,64 @@ describe("codegraph MCP handlers", () => { } }); + it("returns a timeout response while draining an incomplete HTTP MCP body", async () => { + const root = await fs.mkdtemp(path.join(os.tmpdir(), "cg-mcp-http-timeout-")); + await fs.writeFile(path.join(root, "auth.ts"), "export const ok = 1;\n", "utf8"); + const httpServer = await startCodegraphMcpHttpServer({ + root, + host: "127.0.0.1", + port: 0, + httpBodyTimeoutMs: 25, + }); + + try { + const endpoint = new URL(httpServer.url); + const partialBody = '{"jsonrpc":"2.0","id":1,"method":"initialize"'; + const response = await new Promise<{ status: number; payload: JsonRpcObject }>((resolve, reject) => { + let responseReceived = false; + const request = httpRequest( + { + hostname: endpoint.hostname, + port: endpoint.port, + path: endpoint.pathname, + method: "POST", + headers: { + accept: "application/json", + "content-type": "application/json", + "content-length": String(Buffer.byteLength(partialBody) + 1), + }, + }, + (incoming) => { + let responseBody = ""; + incoming.setEncoding("utf8"); + incoming.on("data", (chunk: string) => { + responseBody += chunk; + }); + incoming.on("end", () => { + responseReceived = true; + try { + resolve({ status: incoming.statusCode ?? 0, payload: readJsonRpcObject(JSON.parse(responseBody)) }); + } catch (error) { + reject(error instanceof Error ? error : new Error(String(error))); + } finally { + request.destroy(); + } + }); + }, + ); + request.on("error", (error) => { + if (!responseReceived) reject(error); + }); + request.write(partialBody); + }); + + expect(response.status).toBe(408); + expect(readObject(response.payload.error).message).toBe("MCP request body timed out"); + } finally { + await httpServer.close(); + } + }); + it("reuses one session across search, get_symbol, refs, and query_sqlite handlers", async () => { const root = await fs.mkdtemp(path.join(os.tmpdir(), "cg-mcp-")); await fs.writeFile(path.join(root, "auth.ts"), "export function validateUser(id: number) { return id > 0; }\n"); @@ -2920,7 +2978,6 @@ describe("MCP refresh coalescing", () => { }); }); - describe("MCP session teardown regressions (S2)", () => { it("closes sidecar query index handle and runs session invalidation hooks on server close", async () => { const root = await fs.mkdtemp(path.join(os.tmpdir(), "cg-mcp-s2-teardown-")); @@ -3076,9 +3133,7 @@ describe("MCP transport isolation regressions (S8)", () => { await call1Closed.promise; const call2 = await call2Promise; expect(call2.response.status).toBe(200); - expect(readToolJsonResult(call2.payload).symbols).toEqual([ - expect.objectContaining({ name: "secondarySymbol" }), - ]); + expect(readToolJsonResult(call2.payload).symbols).toEqual([expect.objectContaining({ name: "secondarySymbol" })]); const call3 = await postMcpJson( httpServer.url, @@ -3091,9 +3146,7 @@ describe("MCP transport isolation regressions (S8)", () => { sessionId, ); expect(call3.response.status).toBe(200); - expect(readToolJsonResult(call3.payload).symbols).toEqual([ - expect.objectContaining({ name: "validateUser" }), - ]); + expect(readToolJsonResult(call3.payload).symbols).toEqual([expect.objectContaining({ name: "validateUser" })]); } finally { await httpServer.close(); } diff --git a/tests/sqlite-query-bounds.test.ts b/tests/sqlite-query-bounds.test.ts index 24654e99..539205da 100644 --- a/tests/sqlite-query-bounds.test.ts +++ b/tests/sqlite-query-bounds.test.ts @@ -10,7 +10,7 @@ import { MAX_SQLITE_ROW_LIMIT, SQLITE_TRUNCATED_MARKER, } from "../src/mcp/sqliteGuard.js"; -import { queryGraphSqliteRaw } from "../src/sqlite/query.js"; +import { queryGraphSqliteRaw, SqliteQueryDeadlineExceededError } from "../src/sqlite/query.js"; async function withTempDb(run: (dbPath: string) => Promise): Promise { const root = await fsp.mkdtemp(path.join(os.tmpdir(), "cg-sqlite-bounds-")); @@ -22,6 +22,21 @@ async function withTempDb(run: (dbPath: string) => Promise): Promise } } +async function removeWithRetry(root: string): Promise { + const deadline = Date.now() + 10_000; + for (;;) { + try { + await fsp.rm(root, { recursive: true, force: true }); + return; + } catch (error) { + if (!(error instanceof Error) || !("code" in error)) throw error; + if (error.code !== "EBUSY" && error.code !== "ENOTEMPTY") throw error; + if (Date.now() > deadline) throw error; + await new Promise((resolve) => setTimeout(resolve, 100)); + } + } +} + describe("SQLite query byte/cell bounds during iterate", () => { it("applies per-cell and cumulative caps before appending huge existing TEXT cells", async () => { await withTempDb(async (dbPath) => { @@ -102,3 +117,38 @@ describe("SQLite query byte/cell bounds during iterate", () => { } }); }); + +describe("SQLite raw query execution deadline", () => { + it("terminates an over-budget query without delaying a following query", async () => { + const root = await fsp.mkdtemp(path.join(os.tmpdir(), "cg-sqlite-deadline-")); + const dbPath = path.join(root, "graph.sqlite"); + try { + const db = new DatabaseSync(dbPath); + db.exec("CREATE TABLE values_table (n INTEGER);"); + db.prepare("INSERT INTO values_table (n) VALUES (?)").run(42); + db.close(); + + const slowSql = + "WITH RECURSIVE spin(x) AS (SELECT 1 UNION ALL SELECT x + 1 FROM spin WHERE x < 8000000) " + + "SELECT count(*) FROM spin;"; + const startedAt = Date.now(); + await expect(queryGraphSqliteRaw(dbPath, slowSql, [], { deadlineMs: 100 })).rejects.toMatchObject({ + name: "SqliteQueryDeadlineExceededError", + message: expect.stringMatching(/exceeded its 100ms execution budget/), + }); + expect(Date.now() - startedAt).toBeLessThan(2_000); + + const result = await queryGraphSqliteRaw(dbPath, "SELECT n FROM values_table;", [], { deadlineMs: 5_000 }); + expect(result.rows).toEqual([[42]]); + expect(result.truncated).toBeFalsy(); + } finally { + await removeWithRetry(root); + } + }); + + it("exports a named deadline error for callers", () => { + const error = new SqliteQueryDeadlineExceededError(250); + expect(error.name).toBe("SqliteQueryDeadlineExceededError"); + expect(error.message).toBe("SQLite query exceeded its 250ms execution budget and was terminated."); + }); +}); From 28ae0518f408ecb6b3c067274ade3f70f451f1ef Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sun, 16 Aug 2026 00:03:59 -0400 Subject: [PATCH 13/37] fix: repair MCP session bounds and SQLite deadline contract regressions Three post-review defects in the MCP HTTP transport and raw SQLite query path: 1. Legacy initialize capacity reservations leaked whenever the SDK transport answered a pre-session 4xx without throwing (most notably Accept header validation): onsessioninitialized never fired to release it, and the catch block only covered thrown errors. Release the reservation whenever handleLegacyMcpSessionRequest resolves without a session having been initialized. 2. transport.onerror deleted the whole legacy session for any per-request SDK validation error (bad Accept, wrong Content-Type, malformed JSON, unsupported protocol version, ...), even though those already answered their own request and left the transport healthy. Session teardown is now driven by onclose alone; onerror only logs. 3. The in-process SQLite query fallback (used when the compiled worker asset can't be located) only checks its deadline between already-produced rows, so a statement slow to produce its first row isn't bounded by it. node:sqlite has no interrupt API, so true enforcement requires the worker thread this fallback exists because it couldn't find; corrected the public contract instead (JSDoc, docs/library-api.md, and a one-time degraded-mode log) so the gap is documented and observable rather than silently implied away. --- docs/library-api.md | 2 +- src/mcp/server.ts | 17 +++- src/sqlite/query.ts | 41 ++++++++ src/sqlite/rawQueryWorkerPool.ts | 15 +++ tests/mcp-server.test.ts | 82 +++++++++++++++ tests/sqlite-query-deadline-fallback.test.ts | 102 +++++++++++++++++++ 6 files changed, 257 insertions(+), 2 deletions(-) create mode 100644 tests/sqlite-query-deadline-fallback.test.ts diff --git a/docs/library-api.md b/docs/library-api.md index 5e795034..f7dcac79 100644 --- a/docs/library-api.md +++ b/docs/library-api.md @@ -743,7 +743,7 @@ const result = await queryGraphSqliteRaw( console.log(result.columns, result.rows); ``` -`queryGraphSqliteRaw()` is intentionally read-only. It accepts result-producing statements such as `SELECT` and `PRAGMA` and rejects mutating SQL. Its defaults bound rows, cells, response bytes, and execution to 10 seconds; callers can further tighten `{ maxRows, maxBytes, maxCellBytes, deadlineMs }`. +`queryGraphSqliteRaw()` is intentionally read-only. It accepts result-producing statements such as `SELECT` and `PRAGMA` and rejects mutating SQL. Its defaults bound rows, cells, and response bytes, and callers can further tighten `{ maxRows, maxBytes, maxCellBytes, deadlineMs }`. The 10-second default execution budget (`deadlineMs`) is enforced by running the query in a dedicated worker thread that is force-terminated on expiry, so it interrupts a query even mid-execution; in a degraded install where that worker asset cannot be located, the query instead runs in-process under a weaker per-row check that cannot interrupt a single blocking native call (a logged, one-time-per-process condition). ## SQL artifact facts diff --git a/src/mcp/server.ts b/src/mcp/server.ts index 45505d28..178a0d75 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -1422,9 +1422,16 @@ async function handleLegacyMcpHttpPost( openSseStreams: 0, }; sessionRef.current = session; + // The SDK transport reports every per-request validation rejection through onerror + // too (bad Accept header, wrong Content-Type, malformed JSON, an unsupported + // protocol version, ...) — each of those already answered its own request with a + // 4xx response and left the transport fully usable. Deleting the session here would + // tear down an otherwise healthy session over one malformed follow-up request. Only + // onclose reflects the transport actually shutting down (an explicit DELETE, an + // eviction we triggered, or a real fatal failure), so session teardown is driven by + // onclose alone; onerror only logs. transport.onerror = (error) => { console.error(`[codegraph] MCP HTTP session transport error: ${error.message}`); - if (initializedSessionId !== undefined) void sessionStore.delete(initializedSessionId); }; transport.onclose = () => { if (initializedSessionId !== undefined) void sessionStore.delete(initializedSessionId); @@ -1433,6 +1440,14 @@ async function handleLegacyMcpHttpPost( try { await protocolServer.connect(transport); await handleLegacyMcpSessionRequest(session, request, response, body); + if (initializedSessionId === undefined) { + // The transport answered a pre-session 4xx (invalid Accept header, wrong + // Content-Type, malformed JSON, ...) without throwing and without ever reaching + // onsessioninitialized, so nothing else releases this capacity reservation or + // closes this ad hoc protocol server/transport pair. + releaseCapacityReservation(); + await closeMcpSession(session); + } } catch (error) { if (initializedSessionId !== undefined) { await sessionStore.delete(initializedSessionId); diff --git a/src/sqlite/query.ts b/src/sqlite/query.ts index 6cf062a7..79ea826b 100644 --- a/src/sqlite/query.ts +++ b/src/sqlite/query.ts @@ -17,6 +17,8 @@ import { export { queryGraphSqlite } from "./canned-query.js"; export { SqliteQueryDeadlineExceededError }; +/** Hard wall-clock budget for a single raw `query_sqlite` execution — see the caveat on + * `queryGraphSqliteRaw` about when this is actually enforceable. */ export const DEFAULT_SQLITE_QUERY_DEADLINE_MS = 10_000; export type QueryGraphSqliteRawOptions = { @@ -26,6 +28,32 @@ export type QueryGraphSqliteRawOptions = { deadlineMs?: number | undefined; }; +let loggedInProcessDeadlineFallback = false; + +/** + * Runs a bounded read-only raw SQL query. + * + * Preferred path: the query executes in a dedicated worker thread with a hard + * `deadlineMs` budget (`rawQueryWorkerPool.ts`). On expiry the worker thread is + * terminated outright, which stops the query even while it is blocked inside a single + * synchronous `DatabaseSync` call — a slow non-recursive statement (large join, + * `ORDER BY random()`, a recursive CTE, ...) cannot hold the deadline hostage. + * + * Degraded fallback: if the compiled worker asset cannot be located (a corrupted or + * partial install — the normal build/publish/standalone pipelines all ship it), the + * query instead runs in-process under a *per-row* elapsed-time budget. `node:sqlite`'s + * `DatabaseSync` exposes no interrupt/cancellation API, so once execution is inside a + * single synchronous native call there is nothing in-process that can preempt it — + * true enforcement genuinely requires the separate worker thread this fallback exists + * because it could not find. The per-row check is therefore strictly weaker, not just a + * smaller budget: it is only evaluated between rows the native iterator has already + * produced, so a statement that is slow to produce its very first row (a full scan + * before any match, an aggregate over a large recursive CTE, ...) blocks for its full + * cost before the deadline is ever checked. This fallback exists to keep the common + * case usable in a degraded install, not as a substitute for the worker deadline; a + * warning is logged once per process when it activates so a degraded install is + * observable rather than silently under-enforcing its documented time budget. + */ export async function queryGraphSqliteRaw( outputPath: string, sql: string, @@ -40,6 +68,14 @@ export async function queryGraphSqliteRaw( try { resolveRawSqlQueryWorkerPath(); } catch { + if (!loggedInProcessDeadlineFallback) { + loggedInProcessDeadlineFallback = true; + console.error( + "[codegraph] Raw SQLite query worker asset is unavailable; falling back to an in-process " + + "execution deadline that is only checked between produced rows and cannot interrupt a " + + "single blocking native call. Reinstall to restore the enforced worker-thread deadline.", + ); + } return await queryGraphSqliteRawInProcessBounded(outputPath, sql, params, { maxRows, maxBytes, @@ -51,6 +87,8 @@ export async function queryGraphSqliteRaw( return await runRawSqlQueryInWorker({ outputPath, sql, params, maxRows, maxBytes, maxCellBytes }, deadlineMs); } +/** Degraded fallback for `queryGraphSqliteRaw` — see its doc comment for the enforcement + * caveat this path cannot avoid. */ async function queryGraphSqliteRawInProcessBounded( outputPath: string, sql: string, @@ -82,6 +120,9 @@ async function queryGraphSqliteRawInProcessBounded( }); } +/** Throws once the wall-clock deadline has passed between two already-produced rows. + * See the fallback caveat on `queryGraphSqliteRaw`: a statement slow to produce its + * first row is not bounded here — only slow-*between*-rows iteration is caught. */ function* withPerRowDeadline(rows: Iterable, deadlineAt: number, deadlineMs: number): Generator { for (const row of rows) { if (Date.now() > deadlineAt) throw new SqliteQueryDeadlineExceededError(deadlineMs); diff --git a/src/sqlite/rawQueryWorkerPool.ts b/src/sqlite/rawQueryWorkerPool.ts index e40432f4..6cf01c12 100644 --- a/src/sqlite/rawQueryWorkerPool.ts +++ b/src/sqlite/rawQueryWorkerPool.ts @@ -13,6 +13,11 @@ export class SqliteQueryDeadlineExceededError extends Error { } } +/** Resolves the compiled worker entry: a compiled sibling next to this module + * (production/standalone layouts, where the whole `dist/` tree ships), falling back to + * the package-root-relative compiled path (running this module from `src/`, where only + * `dist/` is built). Throwing here is the trigger `queryGraphSqliteRaw` uses to fall + * back to the strictly weaker in-process per-row deadline check. */ export function resolveRawSqlQueryWorkerPath(): string { const selfDirectory = path.dirname(fileURLToPath(import.meta.url)); const sibling = path.resolve(selfDirectory, "rawQueryWorker.js"); @@ -23,6 +28,16 @@ export function resolveRawSqlQueryWorkerPath(): string { throw new Error(`Raw SQLite query worker file not found: ${compiled}`); } +/** + * Runs a single bounded raw SQL read in a dedicated worker thread with a hard execution + * deadline. On expiry, Piscina's `signal` option force-terminates the worker thread and + * rejects immediately — the caller never waits longer than `deadlineMs`, regardless of + * how long the underlying query actually takes, because termination does not need the + * blocked thread's cooperation. `pool.destroy()` is fire-and-forget rather than awaited, + * so an orphaned worker still finishing one already-in-flight synchronous native call + * never delays this call's rejection or a subsequent query against the same file + * (concurrent read-only SQLite connections do not block each other). + */ export async function runRawSqlQueryInWorker(task: RawQueryWorkerTask, deadlineMs: number): Promise { const pool = new Piscina({ filename: resolveRawSqlQueryWorkerPath(), diff --git a/tests/mcp-server.test.ts b/tests/mcp-server.test.ts index f032258c..ea71e528 100644 --- a/tests/mcp-server.test.ts +++ b/tests/mcp-server.test.ts @@ -3152,3 +3152,85 @@ describe("MCP transport isolation regressions (S8)", () => { } }); }); + +describe("MCP legacy session capacity and error-handling regressions", () => { + it("releases the initialization capacity reservation when legacy Accept header validation rejects the request", async () => { + const root = await fs.mkdtemp(path.join(os.tmpdir(), "cg-mcp-capacity-accept-")); + await fs.writeFile(path.join(root, "auth.ts"), "export const ok = 1;\n", "utf8"); + const httpServer = await startCodegraphMcpHttpServer({ + root, + port: 0, + httpSessionIdleMs: 0, + httpSessionMaxCount: 1, + }); + + const initializeRequest = { + jsonrpc: "2.0", + method: "initialize", + params: { + protocolVersion: "2025-11-25", + capabilities: {}, + clientInfo: { name: "codegraph-capacity-test", version: "1.0.0" }, + }, + }; + + try { + // No override supplies an Accept header, so postRawHttpJson's default + // ("application/json" without "text/event-stream") trips the legacy transport's + // own 406 validation before any session is created. + const rejected = await postRawHttpJson(httpServer.url, { ...initializeRequest, id: 1 }, {}); + expect(rejected.status).toBe(406); + + // With httpSessionMaxCount 1, a leaked capacity reservation from the rejected + // attempt would make this second initialize 503 instead of succeeding. + const accepted = await postMcpJson(httpServer.url, { ...initializeRequest, id: 2 }); + expect(accepted.response.status).toBe(200); + expect(accepted.response.headers.get("mcp-session-id")).toBeTruthy(); + } finally { + await httpServer.close(); + } + }); + + it("keeps a healthy session usable after a request-scoped SDK validation error on a follow-up request", async () => { + const root = await fs.mkdtemp(path.join(os.tmpdir(), "cg-mcp-session-request-error-")); + await fs.writeFile(path.join(root, "auth.ts"), "export function ok(): number { return 1; }\n", "utf8"); + const httpServer = await startCodegraphMcpHttpServer({ root, port: 0 }); + + try { + const initialize = await postMcpJson(httpServer.url, { + jsonrpc: "2.0", + id: 1, + method: "initialize", + params: { + protocolVersion: "2025-11-25", + capabilities: {}, + clientInfo: { name: "codegraph-session-error-test", version: "1.0.0" }, + }, + }); + const sessionId = initialize.response.headers.get("mcp-session-id"); + expect(sessionId).toBeTruthy(); + if (!sessionId) throw new Error("Missing sessionId"); + + // A follow-up request against the same session with a bad Accept header trips + // the transport's own request-scoped validation (406) through onerror, without + // throwing and without the transport ever closing. + const badAccept = await postRawHttpJson( + httpServer.url, + { jsonrpc: "2.0", id: 2, method: "tools/list", params: {} }, + { "mcp-session-id": sessionId }, + ); + expect(badAccept.status).toBe(406); + + // The session must still be usable: a prior bug deleted it from the store on + // every onerror, which would turn this into a 400 "Invalid or missing session ID". + const followUp = await postMcpJson( + httpServer.url, + { jsonrpc: "2.0", id: 3, method: "tools/list", params: {} }, + sessionId, + ); + expect(followUp.response.status).toBe(200); + } finally { + await httpServer.close(); + } + }); +}); diff --git a/tests/sqlite-query-deadline-fallback.test.ts b/tests/sqlite-query-deadline-fallback.test.ts new file mode 100644 index 00000000..917bb8a2 --- /dev/null +++ b/tests/sqlite-query-deadline-fallback.test.ts @@ -0,0 +1,102 @@ +import fsp from "node:fs/promises"; +import os from "node:os"; +import path from "node:path"; +import { DatabaseSync } from "node:sqlite"; +import { describe, expect, it, vi } from "vitest"; + +// Force every query in this file through the in-process fallback (as if the compiled +// worker asset were missing) so its deadline behavior -- and its documented +// limitation -- can be exercised directly, without disturbing the worker-backed +// deadline tests in sqlite-query-bounds.test.ts. +vi.mock("../src/sqlite/rawQueryWorkerPool.js", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + resolveRawSqlQueryWorkerPath: () => { + throw new Error("worker asset unavailable in this test"); + }, + }; +}); + +import { queryGraphSqliteRaw, SqliteQueryDeadlineExceededError } from "../src/sqlite/query.js"; + +async function withTempDb(run: (dbPath: string) => Promise): Promise { + const root = await fsp.mkdtemp(path.join(os.tmpdir(), "cg-sqlite-deadline-fallback-")); + const dbPath = path.join(root, "graph.sqlite"); + try { + await run(dbPath); + } finally { + await fsp.rm(root, { recursive: true, force: true }); + } +} + +describe("SQLite raw query in-process deadline fallback", () => { + it("still enforces the deadline between rows when the worker asset is unavailable", async () => { + await withTempDb(async (dbPath) => { + const db = new DatabaseSync(dbPath); + db.exec("CREATE TABLE t (x INTEGER);"); + db.close(); + + // Every outer row pays for a large nested recursive scan, so successive rows are + // spaced far enough apart in wall-clock time that a short deadline is guaranteed + // to trip between rows -- well before the query would otherwise finish. + const perRowSlowSql = + "WITH RECURSIVE outer_r(x) AS (SELECT 1 UNION ALL SELECT x + 1 FROM outer_r WHERE x < 500) " + + "SELECT x, (" + + " WITH RECURSIVE inner_r(y) AS (SELECT 1 UNION ALL SELECT y + 1 FROM inner_r WHERE y < 200000) " + + " SELECT count(*) FROM inner_r" + + ") FROM outer_r;"; + + await expect(queryGraphSqliteRaw(dbPath, perRowSlowSql, [], { deadlineMs: 20 })).rejects.toMatchObject({ + name: "SqliteQueryDeadlineExceededError", + message: expect.stringMatching(/exceeded its 20ms execution budget/), + }); + }); + }); + + it("does not interrupt a single blocking call whose entire cost is before the first row", async () => { + await withTempDb(async (dbPath) => { + const db = new DatabaseSync(dbPath); + db.exec("CREATE TABLE t (n INTEGER); INSERT INTO t (n) VALUES (1);"); + db.close(); + + // The whole cost of this query is inside one synchronous native step: SQLite + // must finish counting before it can return the single aggregate row. The + // fallback's per-row check cannot fire until that call returns, so -- unlike the + // worker-backed path -- this rejects only after running to completion, not + // within the deadline. That gap is the documented, unavoidable limitation of the + // fallback (see the doc comment on queryGraphSqliteRaw). + const slowBeforeFirstRowSql = + "WITH RECURSIVE spin(x) AS (SELECT 1 UNION ALL SELECT x + 1 FROM spin WHERE x < 8000000) " + + "SELECT count(*) FROM spin;"; + + const start = Date.now(); + await expect(queryGraphSqliteRaw(dbPath, slowBeforeFirstRowSql, [], { deadlineMs: 20 })).rejects.toMatchObject({ + name: "SqliteQueryDeadlineExceededError", + }); + const elapsed = Date.now() - start; + // A true execution deadline would reject close to 20ms; the fallback instead + // blocks for close to the query's full running time before it can even check. + expect(elapsed).toBeGreaterThan(200); + }); + }); + + it("still succeeds for an ordinary query that finishes comfortably inside its deadline", async () => { + await withTempDb(async (dbPath) => { + const db = new DatabaseSync(dbPath); + db.exec("CREATE TABLE t (n INTEGER); INSERT INTO t (n) VALUES (7);"); + db.close(); + + const result = await queryGraphSqliteRaw(dbPath, "SELECT n FROM t;", [], { deadlineMs: 5_000 }); + expect(result.rows).toEqual([[7]]); + expect(result.truncated).toBeFalsy(); + }); + }); + + it("exports a named deadline error for callers regardless of which path produced it", () => { + const error = new SqliteQueryDeadlineExceededError(250); + expect(error).toBeInstanceOf(Error); + expect(error.name).toBe("SqliteQueryDeadlineExceededError"); + expect(error.message).toBe("SQLite query exceeded its 250ms execution budget and was terminated."); + }); +}); From 86b949032b43422d421c7ca1b460947e91918cc0 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sat, 15 Aug 2026 12:14:22 -0400 Subject: [PATCH 14/37] fix: bound and cancel impact analysis streaming Impact streaming's internal chunk queue buffered every produced item in an unbounded array whenever the consumer was not actively reading, and the background analyzeImpact() producer kept running to completion even after a consumer abandoned the stream (broke out of iteration, or the generator was otherwise returned early), retaining index snapshots, sets, and closures for work nobody would ever read. Add an internal AbortController to analyzeImpactStreaming(); the generator's finally block aborts it on every exit path, including the async-generator return protocol triggered by early consumer cancellation. The onImpactItem producer callback checks the signal and throws once it fires, unwinding analyzeImpact's in-progress work so no further symbol batches or transitive passes start. This needs no public API change: yield* delegation already forwards a caller's early return into the inner generator, so session.ts's analyzeImpactStream benefits without any code change there. Cap the internal queue at a bounded number of unread chunks (DEFAULT_MAX_IMPACT_STREAM_QUEUED_CHUNKS, overridable via the internal ImpactStreamingContext.maxQueuedChunks test seam). True backpressure would require making the onImpactItem emission callback awaitable at every synchronous call site in direct.ts/transitive.ts, which is outside this module. On overflow the stream surfaces an explicit ImpactStreamOverflowError as a terminal error chunk instead of silently dropping items or completing as if nothing were missed. --- src/impact/streaming.ts | 73 +++++++++++- tests/impact-streaming.test.ts | 202 ++++++++++++++++++++++++++++++++- tests/session.test.ts | 68 +++++++++++ 3 files changed, 339 insertions(+), 4 deletions(-) diff --git a/src/impact/streaming.ts b/src/impact/streaming.ts index 86cb59be..925d4756 100644 --- a/src/impact/streaming.ts +++ b/src/impact/streaming.ts @@ -71,13 +71,54 @@ function validateImpactStreamingOptions(options: ImpactStreamingOptions): "full" return streamSummary; } +/** + * Raised when a stream consumer falls behind the producer far enough that the buffered, + * unread chunk count would grow without bound. The producer stops (see + * `analyzeImpactStreaming`'s `onImpactItem` wiring) instead of silently dropping chunks, + * so a stalled consumer learns the stream could not keep up rather than quietly receiving + * a truncated-but-apparently-successful result. + */ +export class ImpactStreamOverflowError extends Error { + constructor(maxQueuedChunks: number) { + super( + `Impact stream consumer fell behind the producer: more than ${maxQueuedChunks} chunks were buffered ` + + "without being read. The stream was stopped instead of dropping chunks silently.", + ); + this.name = "ImpactStreamOverflowError"; + } +} + +/** + * Thrown from the `onImpactItem` producer callback once the consumer has abandoned the + * stream (see the `analyzeImpactStreaming` cancellation note below). Unwinds + * `analyzeImpact`'s in-progress work through its normal promise-rejection path; nothing + * outside this module ever observes it, since by construction nobody is listening to the + * stream anymore once it fires. + */ +class ImpactStreamAbandonedError extends Error { + constructor() { + super("Impact stream consumer stopped reading; cancelling in-progress analysis."); + this.name = "ImpactStreamAbandonedError"; + } +} + +/** + * Default cap on buffered-but-unread stream chunks before `ImpactStreamOverflowError` is + * raised. True backpressure (pausing the producer until the consumer catches up) would + * require the `onImpactItem` emission callback to be genuinely awaitable, which means + * awaiting it at every synchronous call site in `direct.ts`/`transitive.ts` — an invasive + * redesign of code outside this module. A hard cap is the non-invasive alternative: it + * turns unbounded memory growth into an explicit, surfaced failure instead. + */ +export const DEFAULT_MAX_IMPACT_STREAM_QUEUED_CHUNKS = 10_000; + type AsyncQueue = { push: (value: T) => void; close: () => void; next: () => Promise>; }; -function createAsyncQueue(): AsyncQueue { +function createAsyncQueue(maxQueuedChunks: number): AsyncQueue { const values: T[] = []; const waiters: Array<(result: IteratorResult) => void> = []; let closed = false; @@ -90,6 +131,9 @@ function createAsyncQueue(): AsyncQueue { waiter({ value, done: false }); return; } + if (values.length >= maxQueuedChunks) { + throw new ImpactStreamOverflowError(maxQueuedChunks); + } values.push(value); }, close() { @@ -175,6 +219,9 @@ export function impactItemEmissionKey(item: ImpactItem, partial: boolean): strin export type ImpactStreamingContext = { buildReport?: BuildReport | undefined; + /** @internal Overrides the buffered-chunk cap (`DEFAULT_MAX_IMPACT_STREAM_QUEUED_CHUNKS`). + * Test seam for deterministically exercising `ImpactStreamOverflowError`. */ + maxQueuedChunks?: number | undefined; }; /** @@ -187,6 +234,16 @@ export type ImpactStreamingContext = { * chains, graph edges, cycles, diagnostics, and schema metadata. Use * `streamSummary: "light"` when a caller only needs the progressive chunks and * a cheap terminal count/detail summary. + * + * Cancellation: if the consumer stops iterating early — a `for await` `break`, or an + * explicit `.return()` on the generator — the async-generator return protocol resumes + * this function's execution at its `finally` block, which aborts an internal + * `AbortController`. The background `analyzeImpact()` producer's `onImpactItem` callback + * checks that signal and throws once it fires, unwinding `analyzeImpact`'s in-progress + * work (no further batches or transitive passes run) instead of letting the whole + * analysis complete unread. This needs no signal parameter on the public API: `yield*` + * delegation (used by `session.ts`'s `analyzeImpactStream`) forwards a caller's + * `.return()` through automatically. */ export async function* analyzeImpactStreaming( projectRoot: string, @@ -194,6 +251,7 @@ export async function* analyzeImpactStreaming( options: ImpactStreamingOptions, context: ImpactStreamingContext = {}, ): AsyncGenerator { + const abortController = new AbortController(); try { const streamSummary = validateImpactStreamingOptions(options); const impactOptions = toImpactOptions(options); @@ -253,7 +311,9 @@ export async function* analyzeImpactStreaming( const normalizedChanges = normalizedDiff.files; const fileLevelFallback = impactOptions.fileLevelFallback ?? true; const fileLevelFallbackPaths = listFileLevelFallbackPaths(normalizedChanges, filesWithSymbols); - const impactQueue = createAsyncQueue(); + const impactQueue = createAsyncQueue( + context.maxQueuedChunks ?? DEFAULT_MAX_IMPACT_STREAM_QUEUED_CHUNKS, + ); const emittedSignatures = new Set(); let impactedItems: ImpactItem[] = []; let impactError: string | null = null; @@ -278,6 +338,9 @@ export async function* analyzeImpactStreaming( fileLevelFallbackPaths, diagnostics, onImpactItem: (item, phase) => { + if (abortController.signal.aborted) { + throw new ImpactStreamAbandonedError(); + } queueImpactItem(item, phase === "partial"); }, }) @@ -355,6 +418,12 @@ export async function* analyzeImpactStreaming( type: "error", error: errorMessage(error), }; + } finally { + // Runs on normal completion, on a caught error, and — via the async-generator return + // protocol — when the consumer stops iterating early. In the early-abandonment case + // this is what actually stops the background analysis: it flips the shared abort + // signal that the `onImpactItem` producer callback above checks and throws from. + abortController.abort(); } } diff --git a/tests/impact-streaming.test.ts b/tests/impact-streaming.test.ts index a40b2ff3..ab990552 100644 --- a/tests/impact-streaming.test.ts +++ b/tests/impact-streaming.test.ts @@ -1,4 +1,4 @@ -import { describe, it, expect } from "vitest"; +import { describe, it, expect, vi } from "vitest"; import path from "node:path"; import os from "node:os"; import fsp from "node:fs/promises"; @@ -8,8 +8,9 @@ import { type ImpactStreamChunk, type ImpactStreamSummaryReport, } from "../src/impact/index.js"; -import { impactItemEmissionKey } from "../src/impact/streaming.js"; +import { impactItemEmissionKey, ImpactStreamOverflowError } from "../src/impact/streaming.js"; import { buildProjectIndex } from "../src/index.js"; +import * as navigation from "../src/indexer/navigation.js"; import { runGit as git } from "./helpers/git.js"; async function mkTmpDir(prefix: string): Promise { @@ -587,3 +588,200 @@ index 1234567..abcdef0 100644 } }); }); + +/** Builds `symbolCount` distinct top-level exported functions in one file, plus a raw + * unified diff with one hunk per function so `mapChangedFileSymbols` reports + * `symbolCount` distinct changed symbols. */ +async function writeManySymbolFixture(root: string, symbolCount: number): Promise<{ diffText: string }> { + const lines = Array.from({ length: symbolCount }, (_, i) => `export function fn${i}() { return ${i}; }`); + await fsp.writeFile(path.join(root, "feature.ts"), `${lines.join("\n")}\n`, "utf8"); + const hunks = lines + .map((line, i) => { + const updated = line.replace(`return ${i};`, `return ${i + 1000};`); + return `@@ -${i + 1} +${i + 1} @@\n-${line}\n+${updated}\n`; + }) + .join(""); + const diffText = `diff --git a/feature.ts b/feature.ts +index 1234567..abcdef0 100644 +--- a/feature.ts ++++ b/feature.ts +${hunks}`; + return { diffText }; +} + +/** Polls a spy's call count until it stops changing, instead of awaiting a fixed real + * delay: the background analysis chain abandoned by the stream consumer settles + * asynchronously and this test holds no promise handle for it, so there is no signal to + * await other than the observable side effect (spy calls) itself. */ +async function waitForStableCallCount( + spy: { mock: { calls: unknown[] } }, + quietMs = 150, + timeoutMs = 5_000, +): Promise { + const deadline = Date.now() + timeoutMs; + let lastCount = spy.mock.calls.length; + let lastChangeAt = Date.now(); + while (Date.now() < deadline) { + const { promise, resolve } = Promise.withResolvers(); + setTimeout(resolve, 20); + await promise; + const count = spy.mock.calls.length; + if (count !== lastCount) { + lastCount = count; + lastChangeAt = Date.now(); + } else if (Date.now() - lastChangeAt >= quietMs) { + return lastCount; + } + } + return lastCount; +} + +describe("Impact streaming resource bounds", () => { + it("surfaces a bounded overflow error instead of silently truncating when a producer burst outruns the queue cap", async () => { + const root = await mkTmpDir("dg-stream-overflow-"); + await fsp.writeFile(path.join(root, "feature.ts"), "export function helper() { return 1; }\n", "utf8"); + const consumerCount = 6; + for (let i = 0; i < consumerCount; i += 1) { + await fsp.writeFile( + path.join(root, `consumer${i}.ts`), + `import { helper } from "./feature";\nexport function run${i}() { return helper(); }\n`, + "utf8", + ); + } + const index = await buildProjectIndex(root); + + try { + const diffText = `diff --git a/feature.ts b/feature.ts +index 1234567..abcdef0 100644 +--- a/feature.ts ++++ b/feature.ts +@@ -1 +1 @@ +-export function helper() { return 1; } ++export function helper() { return 2; } +`; + + // All 6 references to `helper` are emitted inside one synchronous loop in + // direct.ts (no `await` between them), so a cap of 2 overflows deterministically + // on every run regardless of machine speed: the consumer cannot possibly dequeue + // mid-burst. + const chunkTypes: string[] = []; + const errors: string[] = []; + const impactFiles: string[] = []; + for await (const chunk of analyzeImpactStreaming( + root, + index, + { provider: "raw", diffText }, + { maxQueuedChunks: 2 }, + )) { + chunkTypes.push(chunk.type); + if (chunk.type === "impactItem") impactFiles.push(chunk.item.file); + if (chunk.type === "error") errors.push(chunk.error); + } + + expect(chunkTypes).toContain("error"); + expect(chunkTypes).not.toContain("complete"); + expect(errors).toHaveLength(1); + expect(errors[0]).toMatch(/fell behind the producer/); + expect(errors[0]).toMatch(/more than 2 chunks/); + // The consumer learns exactly how far the stream got before it failed, not nothing. + expect(impactFiles.length).toBeGreaterThan(0); + expect(impactFiles.length).toBeLessThan(consumerCount); + } finally { + await fsp.rm(root, { recursive: true, force: true }); + } + }); + + it("does not overflow the same fixture under the default buffered-chunk cap", async () => { + const root = await mkTmpDir("dg-stream-no-overflow-"); + await fsp.writeFile(path.join(root, "feature.ts"), "export function helper() { return 1; }\n", "utf8"); + const consumerCount = 6; + for (let i = 0; i < consumerCount; i += 1) { + await fsp.writeFile( + path.join(root, `consumer${i}.ts`), + `import { helper } from "./feature";\nexport function run${i}() { return helper(); }\n`, + "utf8", + ); + } + const index = await buildProjectIndex(root); + + try { + const diffText = `diff --git a/feature.ts b/feature.ts +index 1234567..abcdef0 100644 +--- a/feature.ts ++++ b/feature.ts +@@ -1 +1 @@ +-export function helper() { return 1; } ++export function helper() { return 2; } +`; + + const chunkTypes: string[] = []; + const impactFiles: string[] = []; + for await (const chunk of analyzeImpactStreaming(root, index, { provider: "raw", diffText })) { + chunkTypes.push(chunk.type); + if (chunk.type === "impactItem") impactFiles.push(chunk.item.file); + } + + expect(chunkTypes).toContain("complete"); + expect(chunkTypes).not.toContain("error"); + const impactedFileSet = new Set(impactFiles); + for (let i = 0; i < consumerCount; i += 1) { + expect(impactedFileSet.has(`consumer${i}.ts`)).toBe(true); + } + } finally { + await fsp.rm(root, { recursive: true, force: true }); + } + }); + + it("stops the background analyzer once the consumer abandons the stream mid-analysis", async () => { + const root = await mkTmpDir("dg-stream-cancel-"); + const symbolCount = 40; + const { diffText } = await writeManySymbolFixture(root, symbolCount); + const index = await buildProjectIndex(root); + + try { + const findReferencesSpy = vi.spyOn(navigation, "findReferences"); + + let sawImpactItem = false; + for await (const chunk of analyzeImpactStreaming(root, index, { provider: "raw", diffText })) { + if (chunk.type === "impactItem") { + sawImpactItem = true; + break; + } + } + expect(sawImpactItem).toBe(true); + + const settledCalls = await waitForStableCallCount(findReferencesSpy); + // Changed symbols are analyzed in fixed batches of 8 (IMPACT_SYMBOL_BATCH_SIZE); + // each batch is awaited fully before the next starts. Cancelling mid-first-batch + // must prevent every later batch from ever starting: comfortably fewer than half + // of the 40 symbols should ever reach a reference lookup. + expect(settledCalls).toBeGreaterThan(0); + expect(settledCalls).toBeLessThan(symbolCount / 2); + } finally { + vi.restoreAllMocks(); + await fsp.rm(root, { recursive: true, force: true }); + } + }); + + it("analyzes every changed symbol when the same stream is consumed to completion", async () => { + const root = await mkTmpDir("dg-stream-nocancel-"); + const symbolCount = 40; + const { diffText } = await writeManySymbolFixture(root, symbolCount); + const index = await buildProjectIndex(root); + + try { + const findReferencesSpy = vi.spyOn(navigation, "findReferences"); + + const chunkTypes: string[] = []; + for await (const chunk of analyzeImpactStreaming(root, index, { provider: "raw", diffText })) { + chunkTypes.push(chunk.type); + } + + expect(chunkTypes).toContain("complete"); + expect(findReferencesSpy).toHaveBeenCalledTimes(symbolCount); + } finally { + vi.restoreAllMocks(); + await fsp.rm(root, { recursive: true, force: true }); + } + }); +}); diff --git a/tests/session.test.ts b/tests/session.test.ts index 9a31b4e1..9f75e8e4 100644 --- a/tests/session.test.ts +++ b/tests/session.test.ts @@ -3,6 +3,7 @@ import type { ICodeReviewSession } from "../src/index.js"; import type { BuildOptions, BuildReport, LanguageExtensionMap } from "../src/indexer/types.js"; import { CodeReviewSession, SessionManager, createCodeReviewSession } from "../src/session.js"; import * as indexerBuild from "../src/indexer/build-index.js"; +import * as navigation from "../src/indexer/navigation.js"; import path from "node:path"; import os from "node:os"; import fs from "node:fs"; @@ -1830,3 +1831,70 @@ describe("SessionManager", () => { expect(manager.getSession("shared")).toBe(existing); }); }); + +describe("CodeReviewSession impact stream cancellation", () => { + test("stops the background analyzer once a session.analyzeImpactStream consumer abandons the stream", async () => { + const root = await fsp.mkdtemp(path.join(os.tmpdir(), "dg-session-impact-stream-cancel-")); + try { + const symbolCount = 40; + const lines = Array.from({ length: symbolCount }, (_, i) => `export function fn${i}() { return ${i}; }`); + await fsp.writeFile(path.join(root, "feature.ts"), `${lines.join("\n")}\n`, "utf8"); + const hunks = lines + .map((line, i) => { + const updated = line.replace(`return ${i};`, `return ${i + 1000};`); + return `@@ -${i + 1} +${i + 1} @@\n-${line}\n+${updated}\n`; + }) + .join(""); + const diffText = `diff --git a/feature.ts b/feature.ts +index 1234567..abcdef0 100644 +--- a/feature.ts ++++ b/feature.ts +${hunks}`; + + const session = await createCodeReviewSession({ + root, + buildOptions: { cache: "memory", useBloomFilters: true }, + }); + const findReferencesSpy = vi.spyOn(navigation, "findReferences"); + try { + let sawImpactItem = false; + for await (const chunk of session.analyzeImpactStream({ provider: "raw", diffText })) { + if (chunk.type === "impactItem") { + sawImpactItem = true; + break; + } + } + expect(sawImpactItem).toBe(true); + + // session.analyzeImpactStream is `yield* analyzeImpactStreaming(...)`: this proves + // that delegation forwards the consumer's early `break` (an async-generator + // `.return()` call) through to the inner generator without any extra plumbing. + // Poll instead of a fixed sleep: the abandoned background chain settles + // asynchronously and this test holds no promise handle for it. + const deadline = Date.now() + 5_000; + let lastCount = findReferencesSpy.mock.calls.length; + let lastChangeAt = Date.now(); + while (Date.now() < deadline && Date.now() - lastChangeAt < 150) { + const { promise, resolve } = Promise.withResolvers(); + setTimeout(resolve, 20); + await promise; + const count = findReferencesSpy.mock.calls.length; + if (count !== lastCount) { + lastCount = count; + lastChangeAt = Date.now(); + } + } + + expect(lastCount).toBeGreaterThan(0); + // Changed symbols are analyzed in fixed batches of 8 (IMPACT_SYMBOL_BATCH_SIZE); + // cancelling mid-first-batch must prevent every later batch from ever starting. + expect(lastCount).toBeLessThan(symbolCount / 2); + } finally { + findReferencesSpy.mockRestore(); + session.dispose(); + } + } finally { + await fsp.rm(root, { recursive: true, force: true }); + } + }); +}); From f816539fb7e17f6a5d4220bd082071bd0dc3bf2d Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sat, 15 Aug 2026 12:14:32 -0400 Subject: [PATCH 15/37] fix: deadline raw SQLite query execution Raw query_sqlite reads ran a synchronous DatabaseSync iteration with row and byte caps but no time budget or cancellation: a non-recursive but expensive statement (a large join, ORDER BY random() with no matching index, ...) could hold the host event loop for as long as SQLite took to produce a row, and a client disconnect did not stop it. Run the query in a dedicated Piscina worker thread (same pattern as the existing query-index worker pool) with a hard deadline. On expiry the worker thread is force-terminated and the call rejects with SqliteQueryDeadlineExceededError immediately; the host event loop is never blocked regardless of how long the underlying query actually runs, and a subsequent query against the same file succeeds right away since concurrent read-only SQLite connections do not block each other. Pool teardown is fire-and-forget on deadline expiry rather than awaited, so the caller is never delayed by an orphaned worker thread still finishing a single already-in-flight synchronous native call (worker termination cannot preempt one in-progress call the same way it can prevent further JS from running) -- documented in rawQueryWorkerPool.ts, verified directly against a 200M-row recursive CTE. If the compiled worker asset cannot be located (a corrupted or partial install), the query falls back to running in-process under a per-row elapsed-time budget instead of refusing outright. That fallback is strictly weaker and is documented as such: the budget is only checked between rows the native iterator has already produced, so a statement that is slow to produce its very first row is not bounded by it. Preserves the existing row/byte cap contracts and normalizeSqliteRowLimit reuse in this file. --- src/sqlite/query.ts | 35 +++++++--------- src/sqlite/rawQueryWorker.ts | 20 +++++++-- src/sqlite/rawQueryWorkerPool.ts | 36 ++++++++++++---- tests/sqlite-query-bounds.test.ts | 69 ++++++++++++++++++++++++------- 4 files changed, 115 insertions(+), 45 deletions(-) diff --git a/src/sqlite/query.ts b/src/sqlite/query.ts index 79ea826b..a085a5ba 100644 --- a/src/sqlite/query.ts +++ b/src/sqlite/query.ts @@ -29,30 +29,24 @@ export type QueryGraphSqliteRawOptions = { }; let loggedInProcessDeadlineFallback = false; - /** * Runs a bounded read-only raw SQL query. * * Preferred path: the query executes in a dedicated worker thread with a hard * `deadlineMs` budget (`rawQueryWorkerPool.ts`). On expiry the worker thread is * terminated outright, which stops the query even while it is blocked inside a single - * synchronous `DatabaseSync` call — a slow non-recursive statement (large join, - * `ORDER BY random()`, a recursive CTE, ...) cannot hold the deadline hostage. + * synchronous `DatabaseSync` call — a client disconnect or a slow non-recursive + * statement (large join, `ORDER BY random()`, ...) can no longer hold the host event + * loop hostage indefinitely. * * Degraded fallback: if the compiled worker asset cannot be located (a corrupted or * partial install — the normal build/publish/standalone pipelines all ship it), the - * query instead runs in-process under a *per-row* elapsed-time budget. `node:sqlite`'s - * `DatabaseSync` exposes no interrupt/cancellation API, so once execution is inside a - * single synchronous native call there is nothing in-process that can preempt it — - * true enforcement genuinely requires the separate worker thread this fallback exists - * because it could not find. The per-row check is therefore strictly weaker, not just a - * smaller budget: it is only evaluated between rows the native iterator has already - * produced, so a statement that is slow to produce its very first row (a full scan - * before any match, an aggregate over a large recursive CTE, ...) blocks for its full - * cost before the deadline is ever checked. This fallback exists to keep the common - * case usable in a degraded install, not as a substitute for the worker deadline; a - * warning is logged once per process when it activates so a degraded install is - * observable rather than silently under-enforcing its documented time budget. + * query instead runs in-process under a *per-row* elapsed-time budget. That fallback is + * strictly weaker: the budget is only checked between rows the native iterator has + * already produced, so a statement that is slow to produce its very first row (for + * example a full-table scan with no matching rows) is not bounded by it. It exists to + * keep the common case usable in a degraded install, not as a substitute for the worker + * deadline. */ export async function queryGraphSqliteRaw( outputPath: string, @@ -106,6 +100,7 @@ async function queryGraphSqliteRawInProcessBounded( deadlineAt, bounds.deadlineMs, ); + // Always stream via iterate so per-cell and cumulative budgets apply before append. return collectBoundedRawSqlRows(columns, rows, { maxRows: bounds.maxRows, maxBytes: bounds.maxBytes, @@ -120,12 +115,14 @@ async function queryGraphSqliteRawInProcessBounded( }); } -/** Throws once the wall-clock deadline has passed between two already-produced rows. - * See the fallback caveat on `queryGraphSqliteRaw`: a statement slow to produce its - * first row is not bounded here — only slow-*between*-rows iteration is caught. */ +/** Throws once the wall-clock deadline has passed between two already-produced rows. See + * the fallback caveat on `queryGraphSqliteRaw`: a slow-before-first-row query is not + * caught here, only slow-*between*-rows iteration is. */ function* withPerRowDeadline(rows: Iterable, deadlineAt: number, deadlineMs: number): Generator { for (const row of rows) { - if (Date.now() > deadlineAt) throw new SqliteQueryDeadlineExceededError(deadlineMs); + if (Date.now() > deadlineAt) { + throw new SqliteQueryDeadlineExceededError(deadlineMs); + } yield row; } } diff --git a/src/sqlite/rawQueryWorker.ts b/src/sqlite/rawQueryWorker.ts index dc74a7a7..6b5f70c4 100644 --- a/src/sqlite/rawQueryWorker.ts +++ b/src/sqlite/rawQueryWorker.ts @@ -3,6 +3,10 @@ import type { RawSqlResult } from "./types.js"; import { assertReadOnlyQueryStatement, withReadOnlySqliteDatabase } from "./database.js"; import { collectBoundedRawSqlRows } from "./rowBounds.js"; +/** + * Task payload for a single bounded raw SQL read, executed inside a dedicated worker + * thread (see `rawQueryWorkerPool.ts`). Every field must be structured-clone safe. + */ export type RawQueryWorkerTask = { outputPath: string; sql: string; @@ -12,13 +16,21 @@ export type RawQueryWorkerTask = { maxCellBytes: number | undefined; }; +/** + * Worker entry point. Mirrors the previous in-process body of `queryGraphSqliteRaw` + * exactly: open the database read-only, assert the statement is read-only, and stream + * rows through the shared row/byte-bounded collector. Running this inside a worker + * thread lets the pool enforce a hard execution deadline by terminating the thread — + * which works even mid-synchronous-iteration, since thread termination does not need + * the blocked thread's cooperation. + */ export default async function runRawQueryWorkerTask(task: RawQueryWorkerTask): Promise { return await withReadOnlySqliteDatabase(task.outputPath, (db) => { try { - const statement = db.prepare(task.sql); - assertReadOnlyQueryStatement(statement); - const columns = statement.columns().map((column) => column.name); - return collectBoundedRawSqlRows(columns, statement.raw().iterate(task.params) as Iterable>, { + const stmt = db.prepare(task.sql); + assertReadOnlyQueryStatement(stmt); + const columns = stmt.columns().map((col) => col.name); + return collectBoundedRawSqlRows(columns, stmt.raw().iterate(task.params) as Iterable>, { maxRows: task.maxRows, maxBytes: task.maxBytes, maxCellBytes: task.maxCellBytes, diff --git a/src/sqlite/rawQueryWorkerPool.ts b/src/sqlite/rawQueryWorkerPool.ts index 6cf01c12..49b9ff09 100644 --- a/src/sqlite/rawQueryWorkerPool.ts +++ b/src/sqlite/rawQueryWorkerPool.ts @@ -30,17 +30,37 @@ export function resolveRawSqlQueryWorkerPath(): string { /** * Runs a single bounded raw SQL read in a dedicated worker thread with a hard execution - * deadline. On expiry, Piscina's `signal` option force-terminates the worker thread and - * rejects immediately — the caller never waits longer than `deadlineMs`, regardless of - * how long the underlying query actually takes, because termination does not need the - * blocked thread's cooperation. `pool.destroy()` is fire-and-forget rather than awaited, - * so an orphaned worker still finishing one already-in-flight synchronous native call - * never delays this call's rejection or a subsequent query against the same file - * (concurrent read-only SQLite connections do not block each other). + * deadline. A fresh single-thread pool is created per call and destroyed afterward — + * matching the existing `prepareQueryIndexFilesInWorker` pattern — since `query_sqlite` + * calls are interactive, not a hot loop, and a persistent pool would need a shutdown hook + * this module has no access to register. + * + * On deadline expiry, Piscina's `signal` option force-terminates the worker thread + * (`worker.terminate()`) and rejects immediately — the caller never waits longer than + * `deadlineMs`, and the host event loop is never blocked by the query regardless of how + * long it runs. Cancellation is real (no further JS runs on that thread and the query + * can never touch this process's caller again), but it has one unavoidable limit shared + * by every in-process cancellation mechanism: `terminate()` cannot preempt a single + * already-in-flight synchronous native call. A query whose entire cost is inside one + * `sqlite3_step()` — a recursive CTE, or a plan that must fully sort/scan before it can + * produce a first row — keeps running on the orphaned worker thread in the background + * until that native call returns naturally; only then does the thread actually exit. + * We do not make the caller wait for that: `pool.destroy()` is fired and forgotten here, + * not awaited, so a subsequent query (in its own fresh pool) is never delayed by it. + * Concurrent read-only SQLite connections against the same file do not block each other, + * so the lingering background reader does not stop that subsequent query from succeeding. + * Verified directly: an aborted 200M-row recursive-CTE count rejects this call in + * ~`deadlineMs`, and an immediately following query against the same file succeeds in + * milliseconds. The one place the orphaned thread is still observable is process + * shutdown: Node cannot fully tear a process down while one of its Worker threads is + * blocked in native code, so a process exit racing a runaway query can itself be + * delayed until that native call returns — a platform limit of `worker_threads`, not of + * this module, and orthogonal to the per-call deadline this function guarantees. */ export async function runRawSqlQueryInWorker(task: RawQueryWorkerTask, deadlineMs: number): Promise { + const workerPath = resolveRawSqlQueryWorkerPath(); const pool = new Piscina({ - filename: resolveRawSqlQueryWorkerPath(), + filename: workerPath, minThreads: 1, maxThreads: 1, idleTimeout: 5_000, diff --git a/tests/sqlite-query-bounds.test.ts b/tests/sqlite-query-bounds.test.ts index 539205da..49343e74 100644 --- a/tests/sqlite-query-bounds.test.ts +++ b/tests/sqlite-query-bounds.test.ts @@ -12,13 +12,37 @@ import { } from "../src/mcp/sqliteGuard.js"; import { queryGraphSqliteRaw, SqliteQueryDeadlineExceededError } from "../src/sqlite/query.js"; +// A deadline-terminated query's worker thread is force-terminated but, if it was blocked +// inside a single synchronous native SQLite call, keeps running that call in the +// background until it returns naturally (see rawQueryWorkerPool.ts). On Windows this can +// hold the temp db file open for a short window after the deadline test's assertions +// already ran. This is a real platform race (an actual lingering OS file lock, not +// simulated timing logic), so it is retried against the real clock instead of being +// modeled with fake timers. +async function removeWithRetry(root: string): Promise { + const deadline = Date.now() + 10_000; + for (;;) { + try { + await fsp.rm(root, { recursive: true, force: true }); + return; + } catch (error) { + const code = (error as NodeJS.ErrnoException).code; + if (code !== "EBUSY" && code !== "ENOTEMPTY") throw error; + if (Date.now() > deadline) throw error; + const { promise, resolve } = Promise.withResolvers(); + setTimeout(resolve, 100); + await promise; + } + } +} + async function withTempDb(run: (dbPath: string) => Promise): Promise { const root = await fsp.mkdtemp(path.join(os.tmpdir(), "cg-sqlite-bounds-")); const dbPath = path.join(root, "graph.sqlite"); try { await run(dbPath); } finally { - await fsp.rm(root, { recursive: true, force: true }); + await removeWithRetry(root); } } @@ -119,35 +143,52 @@ describe("SQLite query byte/cell bounds during iterate", () => { }); describe("SQLite raw query execution deadline", () => { - it("terminates an over-budget query without delaying a following query", async () => { - const root = await fsp.mkdtemp(path.join(os.tmpdir(), "cg-sqlite-deadline-")); - const dbPath = path.join(root, "graph.sqlite"); - try { + it("terminates an over-budget query with a bounded error while a subsequent query on the same file still succeeds", async () => { + await withTempDb(async (dbPath) => { const db = new DatabaseSync(dbPath); - db.exec("CREATE TABLE values_table (n INTEGER);"); - db.prepare("INSERT INTO values_table (n) VALUES (?)").run(42); + db.exec("CREATE TABLE t (n INTEGER);"); + db.prepare("INSERT INTO t (n) VALUES (?)").run(42); db.close(); + // The whole cost of this query is inside one synchronous native step (see + // rawQueryWorkerPool.ts): SQLite must finish counting before it can return the + // single aggregate row, so this reliably runs well past a short deadline without + // depending on machine speed for a *count* of loop iterations. const slowSql = "WITH RECURSIVE spin(x) AS (SELECT 1 UNION ALL SELECT x + 1 FROM spin WHERE x < 8000000) " + "SELECT count(*) FROM spin;"; - const startedAt = Date.now(); + + const start = Date.now(); await expect(queryGraphSqliteRaw(dbPath, slowSql, [], { deadlineMs: 100 })).rejects.toMatchObject({ name: "SqliteQueryDeadlineExceededError", message: expect.stringMatching(/exceeded its 100ms execution budget/), }); - expect(Date.now() - startedAt).toBeLessThan(2_000); + const elapsed = Date.now() - start; + // The caller is bounded by the deadline, not by how long the runaway query + // actually takes to finish in the background (calibrated well above 100ms). + expect(elapsed).toBeLessThan(2_000); - const result = await queryGraphSqliteRaw(dbPath, "SELECT n FROM values_table;", [], { deadlineMs: 5_000 }); + const result = await queryGraphSqliteRaw(dbPath, "SELECT n FROM t;", [], { deadlineMs: 5_000 }); expect(result.rows).toEqual([[42]]); expect(result.truncated).toBeFalsy(); - } finally { - await removeWithRetry(root); - } + }); + }); + + it("does not reject an ordinary query that finishes comfortably inside its deadline", async () => { + await withTempDb(async (dbPath) => { + const db = new DatabaseSync(dbPath); + db.exec("CREATE TABLE t (n INTEGER);"); + db.prepare("INSERT INTO t (n) VALUES (?)").run(7); + db.close(); + + const result = await queryGraphSqliteRaw(dbPath, "SELECT n FROM t;", [], { deadlineMs: 5_000 }); + expect(result.rows).toEqual([[7]]); + }); }); - it("exports a named deadline error for callers", () => { + it("exposes SqliteQueryDeadlineExceededError as a named export for callers to distinguish deadline failures", () => { const error = new SqliteQueryDeadlineExceededError(250); + expect(error).toBeInstanceOf(Error); expect(error.name).toBe("SqliteQueryDeadlineExceededError"); expect(error.message).toBe("SQLite query exceeded its 250ms execution budget and was terminated."); }); From 359a19858900429b998d16317802655c76db3ba3 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sat, 15 Aug 2026 23:00:56 -0400 Subject: [PATCH 16/37] fix: bound raw MCP query worker cleanup --- scripts/bundle-cli-lib.mjs | 18 ++- scripts/ensure-dist-for-tests-lib.mjs | 8 +- scripts/stage-core-package-lib.mjs | 5 +- src/mcp/server.ts | 31 +++-- src/sqlite/query.ts | 23 +++- src/sqlite/rawQueryWorkerPool.ts | 156 +++++++++++++++++++---- tests/cli-bundle-entry.test.ts | 2 + tests/core-package-surface.test.ts | 1 + tests/ensure-dist-for-tests.test.ts | 4 + tests/mcp-stream-cancellation.test.ts | 51 ++++++++ tests/query-index-worker-path.test.ts | 13 ++ tests/raw-query-worker-lifecycle.test.ts | 63 +++++++++ 12 files changed, 332 insertions(+), 43 deletions(-) create mode 100644 tests/mcp-stream-cancellation.test.ts create mode 100644 tests/raw-query-worker-lifecycle.test.ts diff --git a/scripts/bundle-cli-lib.mjs b/scripts/bundle-cli-lib.mjs index cbed4db8..d683b203 100644 --- a/scripts/bundle-cli-lib.mjs +++ b/scripts/bundle-cli-lib.mjs @@ -12,10 +12,12 @@ export function getBundlePaths(rootDir = defaultRootDir) { rootDir, entryPoint: path.join(rootDir, "dist", "cliBootstrap.js"), workerEntryPoint: path.join(rootDir, "dist", "agent", "query-index", "queryIndexWorker.js"), + rawQueryWorkerEntryPoint: path.join(rootDir, "dist", "sqlite", "rawQueryWorker.js"), unbundledCli: path.join(rootDir, "dist", "cli.js"), outdir: path.join(rootDir, "dist", "bin"), bundledEntry: path.join(rootDir, "dist", "bin", "cli.js"), bundledWorker: path.join(rootDir, "dist", "bin", "queryIndexWorker.js"), + bundledRawQueryWorker: path.join(rootDir, "dist", "bin", "rawQueryWorker.js"), }; } @@ -37,12 +39,21 @@ export async function bundleCli({ rootDir = defaultRootDir, logLevel = "warning" if (!fs.existsSync(paths.workerEntryPoint)) { throw new Error(`Missing query worker build input: ${paths.workerEntryPoint}. Run tsc before bundling.`); } + if (!fs.existsSync(paths.rawQueryWorkerEntryPoint)) { + throw new Error( + `Missing raw SQLite worker build input: ${paths.rawQueryWorkerEntryPoint}. Run tsc before bundling.`, + ); + } fs.rmSync(paths.outdir, { recursive: true, force: true }); fs.mkdirSync(paths.outdir, { recursive: true }); const result = await esbuild.build({ - entryPoints: { cli: paths.entryPoint, queryIndexWorker: paths.workerEntryPoint }, + entryPoints: { + cli: paths.entryPoint, + queryIndexWorker: paths.workerEntryPoint, + rawQueryWorker: paths.rawQueryWorkerEntryPoint, + }, bundle: true, platform: "node", format: "esm", @@ -59,7 +70,7 @@ export async function bundleCli({ rootDir = defaultRootDir, logLevel = "warning" }); const outputFiles = Object.keys(result.metafile.outputs).sort(); - const selfContainedEntries = new Set([paths.bundledEntry, paths.bundledWorker]); + const selfContainedEntries = new Set([paths.bundledEntry, paths.bundledWorker, paths.bundledRawQueryWorker]); const unexpectedOutputs = outputFiles.filter((file) => !selfContainedEntries.has(path.resolve(file))); if (unexpectedOutputs.length) { throw new Error( @@ -73,6 +84,9 @@ export async function bundleCli({ rootDir = defaultRootDir, logLevel = "warning" if (!fs.existsSync(paths.bundledWorker)) { throw new Error(`Bundled query worker was not written to ${paths.bundledWorker}`); } + if (!fs.existsSync(paths.bundledRawQueryWorker)) { + throw new Error(`Bundled raw SQLite worker was not written to ${paths.bundledRawQueryWorker}`); + } return { ...paths, diff --git a/scripts/ensure-dist-for-tests-lib.mjs b/scripts/ensure-dist-for-tests-lib.mjs index 8d6edc49..f9a02f83 100644 --- a/scripts/ensure-dist-for-tests-lib.mjs +++ b/scripts/ensure-dist-for-tests-lib.mjs @@ -1,7 +1,13 @@ import fs from "node:fs"; import path from "node:path"; -const requiredDistEntries = ["dist/index.js", "dist/cli.js", "dist/bin/cli.js", "dist/bin/queryIndexWorker.js"]; +const requiredDistEntries = [ + "dist/index.js", + "dist/cli.js", + "dist/bin/cli.js", + "dist/bin/queryIndexWorker.js", + "dist/bin/rawQueryWorker.js", +]; const freshnessInputs = [ "package.json", "tsconfig.json", diff --git a/scripts/stage-core-package-lib.mjs b/scripts/stage-core-package-lib.mjs index 3ecc2f97..03a2234c 100644 --- a/scripts/stage-core-package-lib.mjs +++ b/scripts/stage-core-package-lib.mjs @@ -10,7 +10,10 @@ export const CORE_PACKAGE_ENTRIES = Object.freeze([ "languages.js", ]); -export const CORE_PACKAGE_EXTRA_FILES = Object.freeze(["agent/query-index/queryIndexWorker.js"]); +export const CORE_PACKAGE_EXTRA_FILES = Object.freeze([ + "agent/query-index/queryIndexWorker.js", + "sqlite/rawQueryWorker.js", +]); const IMPORT_PATTERN = /(?:import|export)\s+(?:type\s+)?(?:[^;]*?\s+from\s+)?["'](\.[^"']+)["']|import\(["'](\.[^"']+)["']\)/g; diff --git a/src/mcp/server.ts b/src/mcp/server.ts index 178a0d75..58588b16 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -175,6 +175,10 @@ type OriginValidator = (request: IncomingMessage, response: ServerResponse) => b export type CodegraphMcpFreshResult = T & { freshness: AgentFreshnessResult }; +type McpToolExecutionOptions = { + signal?: AbortSignal | undefined; +}; + /** * Truncation metadata for a capped collection response, per finding #44: * lets a machine caller tell a complete result apart from a capped prefix. @@ -324,14 +328,18 @@ type CodegraphMcpHandlerDefinitions = { head: string; reviewDepth?: ReviewDepth | undefined; }) => Promise>; - refresh_index: (request: { - warmup?: CodegraphMcpWarmupMode | undefined; - }) => Promise<{ refreshed: true; warmup: CodegraphMcpWarmupMode }>; - query_sqlite: (request: { - query: string; - params?: Array | undefined; - limit?: number | undefined; - }) => Promise>; + refresh_index: (request: { warmup?: CodegraphMcpWarmupMode | undefined }) => Promise<{ + refreshed: true; + warmup: CodegraphMcpWarmupMode; + }>; + query_sqlite: ( + request: { + query: string; + params?: Array | undefined; + limit?: number | undefined; + }, + options?: McpToolExecutionOptions, + ) => Promise>; artifact_build: (request: { outDir?: string | undefined; sqlite?: boolean | undefined; @@ -985,7 +993,7 @@ function createCodegraphMcpHandlersForSession( return boundReviewReportForTransport(report); }), - query_sqlite: async (request) => { + query_sqlite: async (request, executionOptions) => { if (!sqlitePath) { throw new Error("No SQLite artifact is available. Run artifact_build first or pass artifactPath."); } @@ -1007,6 +1015,7 @@ function createCodegraphMcpHandlersForSession( const result = await queryGraphSqliteRaw(realSqlitePath, request.query, request.params ?? [], { maxRows: normalizeSqliteRowLimit(request.limit), maxBytes: DEFAULT_SQLITE_BYTE_LIMIT, + ...(executionOptions?.signal ? { signal: executionOptions.signal } : {}), }); return { ...result, truncated: Boolean(result.truncated), freshness: artifactFreshness }; }, @@ -1700,7 +1709,9 @@ export async function callMcpTool( case "review": return await handlers.review(parseMcpToolInput(reviewSchema, input, name), signal); case "query_sqlite": - return await handlers.query_sqlite(parseMcpToolInput(querySqliteSchema, input, name), signal); + return await handlers.query_sqlite(parseMcpToolInput(querySqliteSchema, input, name), { + ...(signal ? { signal } : {}), + }); case "refresh_index": return await handlers.refresh_index(parseMcpToolInput(refreshIndexSchema, input, name), signal); case "artifact_build": diff --git a/src/sqlite/query.ts b/src/sqlite/query.ts index a085a5ba..5a08b0a6 100644 --- a/src/sqlite/query.ts +++ b/src/sqlite/query.ts @@ -11,11 +11,12 @@ import { import { resolveRawSqlQueryWorkerPath, runRawSqlQueryInWorker, + SqliteQueryCancelledError, SqliteQueryDeadlineExceededError, } from "./rawQueryWorkerPool.js"; export { queryGraphSqlite } from "./canned-query.js"; -export { SqliteQueryDeadlineExceededError }; +export { SqliteQueryCancelledError, SqliteQueryDeadlineExceededError }; /** Hard wall-clock budget for a single raw `query_sqlite` execution — see the caveat on * `queryGraphSqliteRaw` about when this is actually enforceable. */ @@ -26,6 +27,7 @@ export type QueryGraphSqliteRawOptions = { maxBytes?: number | undefined; maxCellBytes?: number | undefined; deadlineMs?: number | undefined; + signal?: AbortSignal | undefined; }; let loggedInProcessDeadlineFallback = false; @@ -75,10 +77,15 @@ export async function queryGraphSqliteRaw( maxBytes, maxCellBytes, deadlineMs, + ...(options?.signal ? { signal: options.signal } : {}), }); } - return await runRawSqlQueryInWorker({ outputPath, sql, params, maxRows, maxBytes, maxCellBytes }, deadlineMs); + return await runRawSqlQueryInWorker( + { outputPath, sql, params, maxRows, maxBytes, maxCellBytes }, + deadlineMs, + options?.signal, + ); } /** Degraded fallback for `queryGraphSqliteRaw` — see its doc comment for the enforcement @@ -87,8 +94,9 @@ async function queryGraphSqliteRawInProcessBounded( outputPath: string, sql: string, params: Array, - bounds: { maxRows: number; maxBytes: number; maxCellBytes: number; deadlineMs: number }, + bounds: { maxRows: number; maxBytes: number; maxCellBytes: number; deadlineMs: number; signal?: AbortSignal }, ): Promise { + if (bounds.signal?.aborted) throw new SqliteQueryCancelledError(); return await withReadOnlySqliteDatabase(outputPath, (db) => { try { const stmt = db.prepare(sql); @@ -99,6 +107,7 @@ async function queryGraphSqliteRawInProcessBounded( stmt.raw().iterate(params) as Iterable>, deadlineAt, bounds.deadlineMs, + bounds.signal, ); // Always stream via iterate so per-cell and cumulative budgets apply before append. return collectBoundedRawSqlRows(columns, rows, { @@ -118,8 +127,14 @@ async function queryGraphSqliteRawInProcessBounded( /** Throws once the wall-clock deadline has passed between two already-produced rows. See * the fallback caveat on `queryGraphSqliteRaw`: a slow-before-first-row query is not * caught here, only slow-*between*-rows iteration is. */ -function* withPerRowDeadline(rows: Iterable, deadlineAt: number, deadlineMs: number): Generator { +function* withPerRowDeadline( + rows: Iterable, + deadlineAt: number, + deadlineMs: number, + signal: AbortSignal | undefined, +): Generator { for (const row of rows) { + if (signal?.aborted) throw new SqliteQueryCancelledError(); if (Date.now() > deadlineAt) { throw new SqliteQueryDeadlineExceededError(deadlineMs); } diff --git a/src/sqlite/rawQueryWorkerPool.ts b/src/sqlite/rawQueryWorkerPool.ts index 49b9ff09..c0019c41 100644 --- a/src/sqlite/rawQueryWorkerPool.ts +++ b/src/sqlite/rawQueryWorkerPool.ts @@ -6,6 +6,20 @@ import { findPackageRoot } from "../util/packageInfo.js"; import type { RawSqlResult } from "./types.js"; import type { RawQueryWorkerTask } from "./rawQueryWorker.js"; +export const MAX_RAW_SQL_QUERY_WORKERS = 2; + +export type RawSqlQueryWorkerPool = { + run(task: RawQueryWorkerTask, options: { signal: AbortSignal }): Promise; + destroy(): Promise; +}; + +type RawSqlQueryWorkerPoolFactory = () => RawSqlQueryWorkerPool; + +export type RawSqlQueryWorkerLifecycleState = { + activeWorkers: number; + maxWorkers: number; +}; + export class SqliteQueryDeadlineExceededError extends Error { constructor(deadlineMs: number) { super(`SQLite query exceeded its ${deadlineMs}ms execution budget and was terminated.`); @@ -13,11 +27,100 @@ export class SqliteQueryDeadlineExceededError extends Error { } } -/** Resolves the compiled worker entry: a compiled sibling next to this module - * (production/standalone layouts, where the whole `dist/` tree ships), falling back to - * the package-root-relative compiled path (running this module from `src/`, where only - * `dist/` is built). Throwing here is the trigger `queryGraphSqliteRaw` uses to fall - * back to the strictly weaker in-process per-row deadline check. */ +export class SqliteQueryCancelledError extends Error { + constructor() { + super("SQLite query was cancelled by the MCP client."); + this.name = "SqliteQueryCancelledError"; + } +} + +export class SqliteQueryWorkerCleanupCapacityExceededError extends Error { + constructor(maxWorkers: number) { + super( + `SQLite query cleanup capacity is exhausted: ${maxWorkers} terminated worker${maxWorkers === 1 ? " is" : "s are"} still exiting. Retry after cleanup completes.`, + ); + this.name = "SqliteQueryWorkerCleanupCapacityExceededError"; + } +} + +/** + * Keeps raw-query worker cleanup bounded. A timed-out native SQLite call can delay + * `Piscina.destroy()` until its current synchronous step returns, so each slot remains + * reserved until that destroy promise settles instead of being forgotten in the background. + */ +export class RawSqlQueryWorkerLifecycle { + private readonly activeWorkerSlots = new Set(); + + constructor(private readonly maxWorkers = MAX_RAW_SQL_QUERY_WORKERS) {} + + state(): RawSqlQueryWorkerLifecycleState { + return { activeWorkers: this.activeWorkerSlots.size, maxWorkers: this.maxWorkers }; + } + + async run( + task: RawQueryWorkerTask, + deadlineMs: number, + signal: AbortSignal | undefined, + createPool: RawSqlQueryWorkerPoolFactory, + ): Promise { + if (signal?.aborted) throw new SqliteQueryCancelledError(); + if (this.activeWorkerSlots.size >= this.maxWorkers) { + throw new SqliteQueryWorkerCleanupCapacityExceededError(this.maxWorkers); + } + + const slot = Symbol("raw-sql-query-worker"); + this.activeWorkerSlots.add(slot); + let pool: RawSqlQueryWorkerPool | undefined; + let cleanupInBackground = false; + const deadlineSignal = AbortSignal.timeout(deadlineMs); + const combinedSignal = signal ? AbortSignal.any([signal, deadlineSignal]) : deadlineSignal; + + try { + pool = createPool(); + return await pool.run(task, { signal: combinedSignal }); + } catch (error) { + if (combinedSignal.aborted || (error instanceof Error && error.name === "AbortError")) { + cleanupInBackground = true; + if (signal?.aborted) throw new SqliteQueryCancelledError(); + throw new SqliteQueryDeadlineExceededError(deadlineMs); + } + throw error; + } finally { + if (!pool) { + this.activeWorkerSlots.delete(slot); + } else { + const cleanup = pool.destroy(); + if (cleanupInBackground) { + void cleanup.then( + () => { + this.activeWorkerSlots.delete(slot); + }, + () => { + this.activeWorkerSlots.delete(slot); + }, + ); + } else { + try { + await cleanup; + } finally { + this.activeWorkerSlots.delete(slot); + } + } + } + } + } +} + +const rawSqlQueryWorkerLifecycle = new RawSqlQueryWorkerLifecycle(); + +export function getRawSqlQueryWorkerLifecycleState(): RawSqlQueryWorkerLifecycleState { + return rawSqlQueryWorkerLifecycle.state(); +} + +/** Resolves the compiled worker entry the same way `queryIndexWorker.js` is resolved: + * a compiled sibling next to this module (production/standalone layouts, where the + * whole `dist/` tree ships), falling back to the package-root-relative compiled path + * (vitest running this module from `src/`, where only `dist/` is built). */ export function resolveRawSqlQueryWorkerPath(): string { const selfDirectory = path.dirname(fileURLToPath(import.meta.url)); const sibling = path.resolve(selfDirectory, "rawQueryWorker.js"); @@ -25,7 +128,9 @@ export function resolveRawSqlQueryWorkerPath(): string { const packageRoot = findPackageRoot(selfDirectory); const compiled = path.join(packageRoot, "dist", "sqlite", "rawQueryWorker.js"); if (fs.existsSync(compiled)) return compiled; - throw new Error(`Raw SQLite query worker file not found: ${compiled}`); + const bundled = path.join(packageRoot, "dist", "bin", "rawQueryWorker.js"); + if (fs.existsSync(bundled)) return bundled; + throw new Error(`Raw SQLite query worker file not found: ${bundled}`); } /** @@ -45,8 +150,9 @@ export function resolveRawSqlQueryWorkerPath(): string { * `sqlite3_step()` — a recursive CTE, or a plan that must fully sort/scan before it can * produce a first row — keeps running on the orphaned worker thread in the background * until that native call returns naturally; only then does the thread actually exit. - * We do not make the caller wait for that: `pool.destroy()` is fired and forgotten here, - * not awaited, so a subsequent query (in its own fresh pool) is never delayed by it. + * The caller does not wait for that cleanup, but the lifecycle retains its worker slot + * until `pool.destroy()` settles. The bounded slot count makes delayed cleanup observable + * and prevents repeated cancellation from accumulating an unbounded number of workers. * Concurrent read-only SQLite connections against the same file do not block each other, * so the lingering background reader does not stop that subsequent query from succeeding. * Verified directly: an aborted 200M-row recursive-CTE count rejects this call in @@ -57,22 +163,22 @@ export function resolveRawSqlQueryWorkerPath(): string { * delayed until that native call returns — a platform limit of `worker_threads`, not of * this module, and orthogonal to the per-call deadline this function guarantees. */ -export async function runRawSqlQueryInWorker(task: RawQueryWorkerTask, deadlineMs: number): Promise { +export async function runRawSqlQueryInWorker( + task: RawQueryWorkerTask, + deadlineMs: number, + signal?: AbortSignal, +): Promise { const workerPath = resolveRawSqlQueryWorkerPath(); - const pool = new Piscina({ - filename: workerPath, - minThreads: 1, - maxThreads: 1, - idleTimeout: 5_000, - }); - try { - return (await pool.run(task, { signal: AbortSignal.timeout(deadlineMs) })) as RawSqlResult; - } catch (error) { - if (error instanceof Error && error.name === "AbortError") { - throw new SqliteQueryDeadlineExceededError(deadlineMs); - } - throw error; - } finally { - void pool.destroy().catch(() => {}); - } + return await rawSqlQueryWorkerLifecycle.run( + task, + deadlineMs, + signal, + () => + new Piscina({ + filename: workerPath, + minThreads: 1, + maxThreads: 1, + idleTimeout: 5_000, + }), + ); } diff --git a/tests/cli-bundle-entry.test.ts b/tests/cli-bundle-entry.test.ts index 44a29492..2cc932d5 100644 --- a/tests/cli-bundle-entry.test.ts +++ b/tests/cli-bundle-entry.test.ts @@ -7,6 +7,7 @@ import { fileURLToPath } from "node:url"; const rootDir = path.resolve(path.dirname(fileURLToPath(import.meta.url)), ".."); const bundledCli = path.join(rootDir, "dist", "bin", "cli.js"); +const bundledRawQueryWorker = path.join(rootDir, "dist", "bin", "rawQueryWorker.js"); const unbundledCli = path.join(rootDir, "dist", "cli.js"); function run(entry: string, args: string[], cwd: string = rootDir, env: NodeJS.ProcessEnv = process.env) { @@ -21,6 +22,7 @@ function run(entry: string, args: string[], cwd: string = rootDir, env: NodeJS.P describe("bundled CLI entry", () => { it("ships a split ESM bin entry that matches unbundled --version", () => { expect(fs.existsSync(bundledCli)).toBe(true); + expect(fs.existsSync(bundledRawQueryWorker)).toBe(true); expect(fs.existsSync(unbundledCli)).toBe(true); const bundled = run(bundledCli, ["--version"]); diff --git a/tests/core-package-surface.test.ts b/tests/core-package-surface.test.ts index f7e64748..def36129 100644 --- a/tests/core-package-surface.test.ts +++ b/tests/core-package-surface.test.ts @@ -20,6 +20,7 @@ describe("codegraph-core package surface", () => { expect(files.length).toBeGreaterThan(100); expect(files.some(isForbiddenCorePackagePath)).toBe(false); expect(files).toContain("agent/query-index/queryIndexWorker.js"); + expect(files).toContain("sqlite/rawQueryWorker.js"); expect(files).toContain("graphs/types.d.ts"); expect(files).toContain("agent/semantic.d.ts"); expect(files).toContain("chunking/types.d.ts"); diff --git a/tests/ensure-dist-for-tests.test.ts b/tests/ensure-dist-for-tests.test.ts index a12f3ac9..7b2cdc2c 100644 --- a/tests/ensure-dist-for-tests.test.ts +++ b/tests/ensure-dist-for-tests.test.ts @@ -53,6 +53,7 @@ describe("inspectDistForTests", () => { await fsp.mkdir(path.join(root, "dist", "bin"), { recursive: true }); await setFileMtime(path.join(root, "dist", "bin", "cli.js"), distTime, "export {};\n"); await setFileMtime(path.join(root, "dist", "bin", "queryIndexWorker.js"), distTime, "export {};\n"); + await setFileMtime(path.join(root, "dist", "bin", "rawQueryWorker.js"), distTime, "export {};\n"); await setFileMtime(path.join(root, "src", "index.ts"), srcTime); expect(inspectDistForTests(root)).toMatchObject({ @@ -77,6 +78,7 @@ describe("inspectDistForTests", () => { await fsp.mkdir(path.join(root, "dist", "bin"), { recursive: true }); await setFileMtime(path.join(root, "dist", "bin", "cli.js"), distTime, "export {};\n"); await setFileMtime(path.join(root, "dist", "bin", "queryIndexWorker.js"), distTime, "export {};\n"); + await setFileMtime(path.join(root, "dist", "bin", "rawQueryWorker.js"), distTime, "export {};\n"); expect(inspectDistForTests(root)).toMatchObject({ needsBuild: false, @@ -101,6 +103,7 @@ describe("inspectDistForTests", () => { await fsp.mkdir(path.join(root, "dist", "bin"), { recursive: true }); await setFileMtime(path.join(root, "dist", "bin", "cli.js"), distTime, "export {};\n"); await setFileMtime(path.join(root, "dist", "bin", "queryIndexWorker.js"), distTime, "export {};\n"); + await setFileMtime(path.join(root, "dist", "bin", "rawQueryWorker.js"), distTime, "export {};\n"); expect(inspectDistForTests(root)).toMatchObject({ needsBuild: false, reason: "fresh" }); @@ -130,6 +133,7 @@ describe("inspectDistForTests", () => { await fsp.mkdir(path.join(root, "dist", "bin"), { recursive: true }); await setFileMtime(path.join(root, "dist", "bin", "cli.js"), distTime, "export {};\n"); await setFileMtime(path.join(root, "dist", "bin", "queryIndexWorker.js"), distTime, "export {};\n"); + await setFileMtime(path.join(root, "dist", "bin", "rawQueryWorker.js"), distTime, "export {};\n"); expect(inspectDistForTests(root)).toMatchObject({ needsBuild: false, reason: "fresh" }); diff --git a/tests/mcp-stream-cancellation.test.ts b/tests/mcp-stream-cancellation.test.ts new file mode 100644 index 00000000..f90eb1c8 --- /dev/null +++ b/tests/mcp-stream-cancellation.test.ts @@ -0,0 +1,51 @@ +import fsp from "node:fs/promises"; +import os from "node:os"; +import path from "node:path"; +import { Client, InMemoryTransport } from "@modelcontextprotocol/client"; +import { describe, expect, it } from "vitest"; +import { createCodegraphMcpHandlers, createCodegraphMcpProtocolServer } from "../src/mcp/server.js"; + +describe("MCP query_sqlite cancellation", () => { + it("forwards a cancelled tool stream to the raw query handler", async () => { + const root = await fsp.mkdtemp(path.join(os.tmpdir(), "cg-mcp-query-cancel-")); + const handlers = createCodegraphMcpHandlers({ root }); + const started = Promise.withResolvers(); + const cancelled = Promise.withResolvers(); + handlers.query_sqlite = async (_request, executionOptions) => { + const signal = executionOptions?.signal; + if (!signal) throw new Error("MCP query_sqlite did not receive a cancellation signal."); + started.resolve(); + await new Promise((_resolve, reject) => { + signal.addEventListener( + "abort", + () => { + cancelled.resolve(); + reject(new Error("raw query cancelled")); + }, + { once: true }, + ); + }); + }; + + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + const server = createCodegraphMcpProtocolServer(handlers); + const client = new Client({ name: "mcp-stream-cancellation-test", version: "1.0.0" }); + const controller = new AbortController(); + try { + await server.connect(serverTransport); + await client.connect(clientTransport); + const call = client.callTool( + { name: "query_sqlite", arguments: { query: "SELECT 1;" } }, + { signal: controller.signal }, + ); + await started.promise; + controller.abort(); + + await expect(call).rejects.toThrow(); + await cancelled.promise; + } finally { + await Promise.allSettled([client.close(), server.close()]); + await fsp.rm(root, { recursive: true, force: true }); + } + }); +}); diff --git a/tests/query-index-worker-path.test.ts b/tests/query-index-worker-path.test.ts index 0a64da40..495133e7 100644 --- a/tests/query-index-worker-path.test.ts +++ b/tests/query-index-worker-path.test.ts @@ -2,11 +2,24 @@ import fs from "node:fs"; import path from "node:path"; import { afterEach, describe, expect, it, vi } from "vitest"; import { resolveQueryIndexWorkerPath } from "../src/agent/query-index/workerPool.js"; +import { resolveRawSqlQueryWorkerPath } from "../src/sqlite/rawQueryWorkerPool.js"; afterEach(() => { vi.restoreAllMocks(); }); +describe("resolveRawSqlQueryWorkerPath", () => { + it("falls back to the bundled worker when the compiled dist worker is missing", () => { + const bundledSuffix = path.normalize(path.join("dist", "bin", "rawQueryWorker.js")); + vi.spyOn(fs, "existsSync").mockImplementation((candidate) => { + const filePath = path.normalize(typeof candidate === "string" ? candidate : String(candidate)); + return filePath.endsWith(bundledSuffix); + }); + + expect(path.normalize(resolveRawSqlQueryWorkerPath())).toContain(bundledSuffix); + }); +}); + describe("resolveQueryIndexWorkerPath", () => { it("falls back to the bundled worker when the compiled dist worker is missing", () => { const bundledSuffix = path.normalize(path.join("dist", "bin", "queryIndexWorker.js")); diff --git a/tests/raw-query-worker-lifecycle.test.ts b/tests/raw-query-worker-lifecycle.test.ts new file mode 100644 index 00000000..cbbe8d26 --- /dev/null +++ b/tests/raw-query-worker-lifecycle.test.ts @@ -0,0 +1,63 @@ +import { describe, expect, it } from "vitest"; +import { + RawSqlQueryWorkerLifecycle, + SqliteQueryCancelledError, + type RawSqlQueryWorkerPool, +} from "../src/sqlite/rawQueryWorkerPool.js"; +import type { RawQueryWorkerTask } from "../src/sqlite/rawQueryWorker.js"; + +const task: RawQueryWorkerTask = { + outputPath: "fixture.sqlite", + sql: "SELECT 1;", + params: [], + maxRows: 1, + maxBytes: 1024, + maxCellBytes: 1024, +}; + +function createAbortablePool(cleanup: Promise): RawSqlQueryWorkerPool { + return { + run: async (_task, options) => + await new Promise((_, reject) => { + options.signal?.addEventListener( + "abort", + () => { + const error = new Error("aborted"); + error.name = "AbortError"; + reject(error); + }, + { once: true }, + ); + }), + destroy: async () => await cleanup, + }; +} + +describe("RawSqlQueryWorkerLifecycle", () => { + it("caps cancelled worker cleanup slots until each worker has actually exited", async () => { + const lifecycle = new RawSqlQueryWorkerLifecycle(2); + const firstCleanup = Promise.withResolvers(); + const secondCleanup = Promise.withResolvers(); + const firstAbort = new AbortController(); + const secondAbort = new AbortController(); + + const first = lifecycle.run(task, 10_000, firstAbort.signal, () => createAbortablePool(firstCleanup.promise)); + firstAbort.abort(); + await expect(first).rejects.toBeInstanceOf(SqliteQueryCancelledError); + + const second = lifecycle.run(task, 10_000, secondAbort.signal, () => createAbortablePool(secondCleanup.promise)); + secondAbort.abort(); + await expect(second).rejects.toBeInstanceOf(SqliteQueryCancelledError); + + expect(lifecycle.state()).toEqual({ activeWorkers: 2, maxWorkers: 2 }); + await expect(lifecycle.run(task, 10_000, undefined, () => createAbortablePool(Promise.resolve()))).rejects.toThrow( + /cleanup capacity/i, + ); + + firstCleanup.resolve(); + secondCleanup.resolve(); + await Promise.resolve(); + await Promise.resolve(); + expect(lifecycle.state()).toEqual({ activeWorkers: 0, maxWorkers: 2 }); + }); +}); From 869a94a25855550038d000fe2b643f99c43d1173 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sun, 16 Aug 2026 08:58:06 -0400 Subject: [PATCH 17/37] fix: preserve MCP SQLite cancellation wiring --- src/mcp/server.ts | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/src/mcp/server.ts b/src/mcp/server.ts index 58588b16..c568800d 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -350,11 +350,7 @@ type CodegraphMcpHandlerDefinitions = { }) => Promise>; }; -type WithAbortSignal = { - [K in keyof T]: T[K] extends (request: infer Request) => Promise - ? (request: Request, signal?: AbortSignal) => Promise - : never; -}; +type WithAbortSignal = { [K in Exclude]: T[K] extends (request: infer Request) => Promise ? (request: Request, signal?: AbortSignal) => Promise : never; } & Pick; export type CodegraphMcpHandlers = WithAbortSignal; From 7fb5c06ea0e165bbe6b6e79b17a5371fd3801879 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sun, 16 Aug 2026 09:03:19 -0400 Subject: [PATCH 18/37] test: correlate fallback deadline workload --- tests/sqlite-query-deadline-fallback.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/sqlite-query-deadline-fallback.test.ts b/tests/sqlite-query-deadline-fallback.test.ts index 917bb8a2..f7677a2d 100644 --- a/tests/sqlite-query-deadline-fallback.test.ts +++ b/tests/sqlite-query-deadline-fallback.test.ts @@ -43,7 +43,7 @@ describe("SQLite raw query in-process deadline fallback", () => { const perRowSlowSql = "WITH RECURSIVE outer_r(x) AS (SELECT 1 UNION ALL SELECT x + 1 FROM outer_r WHERE x < 500) " + "SELECT x, (" + - " WITH RECURSIVE inner_r(y) AS (SELECT 1 UNION ALL SELECT y + 1 FROM inner_r WHERE y < 200000) " + + " WITH RECURSIVE inner_r(y) AS (SELECT 1 UNION ALL SELECT y + 1 FROM inner_r WHERE y < 200000 + outer_r.x) " + " SELECT count(*) FROM inner_r" + ") FROM outer_r;"; From afd0880e1c09d7688087116040b1ddb38a4f5ba9 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sun, 16 Aug 2026 11:58:59 -0400 Subject: [PATCH 19/37] fix: bound MCP refresh retries --- src/agent/query-index/sessionStore.ts | 13 ++++++- src/mcp/server.ts | 10 ++++-- tests/agent-session.test.ts | 49 ++++++++++++++------------- tests/mcp-server.test.ts | 21 ++++++++++++ 4 files changed, 67 insertions(+), 26 deletions(-) diff --git a/src/agent/query-index/sessionStore.ts b/src/agent/query-index/sessionStore.ts index 9abe093f..1283c1e6 100644 --- a/src/agent/query-index/sessionStore.ts +++ b/src/agent/query-index/sessionStore.ts @@ -10,6 +10,7 @@ type SessionQueryIndexState = { }; const QUERY_INDEX_BY_SESSION = new WeakMap(); +const QUERY_INDEX_INVALIDATION_HOOKS = new WeakSet(); const MAX_QUERY_INDEX_GENERATION_RETRIES = 3; function closeHandle(handle: QueryIndexHandle): void { @@ -25,6 +26,13 @@ function closeState(state: SessionQueryIndexState): void { void state.handle.then(closeHandle, () => undefined); } +function disposeSessionQueryIndexOnInvalidation(session: AgentSession): () => void { + return () => { + QUERY_INDEX_INVALIDATION_HOOKS.delete(session); + disposeSessionQueryIndex(session); + }; +} + export async function ensureSessionQueryIndex( session: AgentSession, snapshot: AgentProjectSnapshot, @@ -43,7 +51,10 @@ export async function ensureSessionQueryIndex( const handle = ensureQueryIndex(snapshot); const state: SessionQueryIndexState = { identity, handle }; - if (!existing) registerSessionInvalidationHook(session, () => disposeSessionQueryIndex(session)); + if (!QUERY_INDEX_INVALIDATION_HOOKS.has(session)) { + QUERY_INDEX_INVALIDATION_HOOKS.add(session); + registerSessionInvalidationHook(session, disposeSessionQueryIndexOnInvalidation(session)); + } QUERY_INDEX_BY_SESSION.set(session, state); handle.catch(() => { if (QUERY_INDEX_BY_SESSION.get(session) === state) QUERY_INDEX_BY_SESSION.delete(session); diff --git a/src/mcp/server.ts b/src/mcp/server.ts index c568800d..8632be78 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -350,7 +350,11 @@ type CodegraphMcpHandlerDefinitions = { }) => Promise>; }; -type WithAbortSignal = { [K in Exclude]: T[K] extends (request: infer Request) => Promise ? (request: Request, signal?: AbortSignal) => Promise : never; } & Pick; +type WithAbortSignal = { + [K in Exclude]: T[K] extends (request: infer Request) => Promise + ? (request: Request, signal?: AbortSignal) => Promise + : never; +} & Pick; export type CodegraphMcpHandlers = WithAbortSignal; @@ -376,6 +380,7 @@ type SqliteArtifactFileSignature = { }; const MAX_MCP_FRESHNESS_CHANGED_FILES = 25; +const MAX_MCP_FRESHNESS_RETRIES = 3; const SQLITE_ARTIFACT_STAT_CONCURRENCY = 64; function assertMcpSessionOptions(options: CodegraphMcpHandlerOptions): void { @@ -479,13 +484,14 @@ function createCodegraphMcpHandlersForSession( const withFreshness = async ( run: () => Promise, ): Promise => { - while (true) { + for (let attempt = 0; attempt < MAX_MCP_FRESHNESS_RETRIES; attempt += 1) { if (refreshPromise) await refreshPromise; const epoch = refreshEpoch; const freshness = await checkMcpFreshness(); const result = await run(); if (epoch === refreshEpoch && !refreshPromise) return { ...result, freshness }; } + throw new Error("Workspace refresh changed repeatedly while serving the request; retry after refresh completes."); }; const formatSqliteFreshnessError = (freshness: AgentFreshnessResult): string => { if (freshness.state === "fresh") return "SQLite artifact freshness check unexpectedly failed."; diff --git a/tests/agent-session.test.ts b/tests/agent-session.test.ts index 7405d5e7..6b6ce147 100644 --- a/tests/agent-session.test.ts +++ b/tests/agent-session.test.ts @@ -1,5 +1,6 @@ import { disposeSessionQueryIndex, ensureSessionQueryIndex } from "../src/agent/query-index/sessionStore.js"; import * as updateModule from "../src/agent/query-index/update.js"; +import * as sessionLifecycleModule from "../src/agent/sessionLifecycle.js"; import fs from "node:fs/promises"; import { brotliCompressSync, brotliDecompressSync, constants as zlibConstants } from "node:zlib"; import { createHash } from "node:crypto"; @@ -941,30 +942,32 @@ describe("agent session", () => { } }); -describe("query index sessionStore generation retries (S12)", () => { - it("bounds query index generation retries under sustained invalidation and surfaces a clear error", async () => { - const root = await mkRepo(); - const session = createAgentSession({ root }); - const snapshot = await session.loadProject(); + describe("query index sessionStore generation retries (S12)", () => { + it("bounds query index generation retries under sustained invalidation and surfaces a clear error", async () => { + const root = await mkRepo(); + const session = createAgentSession({ root }); + const snapshot = await session.loadProject(); - let attempts = 0; - const realEnsureQueryIndex = updateModule.ensureQueryIndex; - const ensureQueryIndexSpy = vi.spyOn(updateModule, "ensureQueryIndex").mockImplementation(async (snap) => { - attempts += 1; - const res = await realEnsureQueryIndex(snap); - disposeSessionQueryIndex(session); - return res; + let attempts = 0; + const realEnsureQueryIndex = updateModule.ensureQueryIndex; + const ensureQueryIndexSpy = vi.spyOn(updateModule, "ensureQueryIndex").mockImplementation(async (snap) => { + attempts += 1; + const res = await realEnsureQueryIndex(snap); + disposeSessionQueryIndex(session); + return res; + }); + const invalidationHookSpy = vi.spyOn(sessionLifecycleModule, "registerSessionInvalidationHook"); + + try { + await expect(ensureSessionQueryIndex(session, snapshot)).rejects.toThrow( + /Query index generation changed repeatedly while loading/i, + ); + expect(attempts).toBe(3); + expect(invalidationHookSpy).toHaveBeenCalledTimes(1); + } finally { + invalidationHookSpy.mockRestore(); + ensureQueryIndexSpy.mockRestore(); + } }); - - try { - await expect(ensureSessionQueryIndex(session, snapshot)).rejects.toThrow( - /Query index generation changed repeatedly while loading/i, - ); - expect(attempts).toBe(3); - } finally { - ensureQueryIndexSpy.mockRestore(); - } }); }); - -}); \ No newline at end of file diff --git a/tests/mcp-server.test.ts b/tests/mcp-server.test.ts index ea71e528..d686df10 100644 --- a/tests/mcp-server.test.ts +++ b/tests/mcp-server.test.ts @@ -2976,6 +2976,27 @@ describe("MCP refresh coalescing", () => { releaseSecondWarmup.resolve(); } }); + + it("bounds a request invalidated by repeated refreshes", async () => { + const root = await fs.mkdtemp(path.join(os.tmpdir(), "cg-mcp-refresh-retry-bound-")); + await fs.writeFile(path.join(root, "auth.ts"), "export const token = 1;\n", "utf8"); + const backingSession = createAgentSession({ root }); + let refreshes = 0; + const session: AgentSession = { + ...backingSession, + loadProject: async (options) => { + refreshes += 1; + await handlers.refresh_index({ warmup: "off" }); + return await backingSession.loadProject(options); + }, + }; + const handlers = createCodegraphMcpHandlers({ root, session }); + + await expect(handlers.goto({ file: "auth.ts", line: 1, column: 14 })).rejects.toThrow( + /Workspace refresh changed repeatedly while serving the request/i, + ); + expect(refreshes).toBe(3); + }); }); describe("MCP session teardown regressions (S2)", () => { From 02aad5dcc5ab4d148dc916a403d3227b17f814fa Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sun, 16 Aug 2026 16:42:44 -0400 Subject: [PATCH 20/37] fix: close rejected MCP body connections --- src/mcp/server.ts | 10 ++++- tests/mcp-server.test.ts | 82 +++++++++++++++++++++------------------- 2 files changed, 52 insertions(+), 40 deletions(-) diff --git a/src/mcp/server.ts b/src/mcp/server.ts index 8632be78..34fed573 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -1310,6 +1310,12 @@ async function handleMcpHttpRequest( createProtocolServer: () => Server, bodyTimeoutMs: number, ): Promise { + const writeClosingJsonRpcError = (statusCode: number, message: string): void => { + response.setHeader("connection", "close"); + + writeJsonRpcError(response, statusCode, message); + }; + const requestPath = getRequestPath(request); if (requestPath !== MCP_HTTP_PATH) { writeJsonResponse(response, 404, { error: "Not found" }); @@ -1326,11 +1332,11 @@ async function handleMcpHttpRequest( if (request.method === "POST") { const parsedBody = await readJsonRequestBody(request, MAX_MCP_HTTP_BODY_BYTES, bodyTimeoutMs); if (parsedBody.status === "too_large") { - writeJsonRpcError(response, 413, "MCP request body is too large"); + writeClosingJsonRpcError(413, "MCP request body is too large"); return; } if (parsedBody.status === "timeout") { - writeJsonRpcError(response, 408, "MCP request body timed out"); + writeClosingJsonRpcError(408, "MCP request body timed out"); return; } if (parsedBody.status === "invalid_json") { diff --git a/tests/mcp-server.test.ts b/tests/mcp-server.test.ts index d686df10..b751084f 100644 --- a/tests/mcp-server.test.ts +++ b/tests/mcp-server.test.ts @@ -1151,7 +1151,7 @@ describe("codegraph MCP handlers", () => { const payload = readObject((await response.json()) as unknown); const error = readObject(payload.error); - expect(response.status).toBe(413); + expect([response.status, response.headers.get("connection")]).toEqual([413, "close"]); expect(error.message).toBe("MCP request body is too large"); } finally { await httpServer.close(); @@ -1171,45 +1171,51 @@ describe("codegraph MCP handlers", () => { try { const endpoint = new URL(httpServer.url); const partialBody = '{"jsonrpc":"2.0","id":1,"method":"initialize"'; - const response = await new Promise<{ status: number; payload: JsonRpcObject }>((resolve, reject) => { - let responseReceived = false; - const request = httpRequest( - { - hostname: endpoint.hostname, - port: endpoint.port, - path: endpoint.pathname, - method: "POST", - headers: { - accept: "application/json", - "content-type": "application/json", - "content-length": String(Buffer.byteLength(partialBody) + 1), + const response = await new Promise<{ status: number; payload: JsonRpcObject; connection: string | undefined }>( + (resolve, reject) => { + let responseReceived = false; + const request = httpRequest( + { + hostname: endpoint.hostname, + port: endpoint.port, + path: endpoint.pathname, + method: "POST", + headers: { + accept: "application/json", + "content-type": "application/json", + "content-length": String(Buffer.byteLength(partialBody) + 1), + }, }, - }, - (incoming) => { - let responseBody = ""; - incoming.setEncoding("utf8"); - incoming.on("data", (chunk: string) => { - responseBody += chunk; - }); - incoming.on("end", () => { - responseReceived = true; - try { - resolve({ status: incoming.statusCode ?? 0, payload: readJsonRpcObject(JSON.parse(responseBody)) }); - } catch (error) { - reject(error instanceof Error ? error : new Error(String(error))); - } finally { - request.destroy(); - } - }); - }, - ); - request.on("error", (error) => { - if (!responseReceived) reject(error); - }); - request.write(partialBody); - }); + (incoming) => { + let responseBody = ""; + incoming.setEncoding("utf8"); + incoming.on("data", (chunk: string) => { + responseBody += chunk; + }); + incoming.on("end", () => { + responseReceived = true; + try { + resolve({ + status: incoming.statusCode ?? 0, + payload: readJsonRpcObject(JSON.parse(responseBody)), + connection: incoming.headers.connection, + }); + } catch (error) { + reject(error instanceof Error ? error : new Error(String(error))); + } finally { + request.destroy(); + } + }); + }, + ); + request.on("error", (error) => { + if (!responseReceived) reject(error); + }); + request.write(partialBody); + }, + ); - expect(response.status).toBe(408); + expect([response.status, response.connection]).toEqual([408, "close"]); expect(readObject(response.payload.error).message).toBe("MCP request body timed out"); } finally { await httpServer.close(); From cb2e2e7a18d955479e4f956bc9b2ac94de8b5ba9 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sun, 16 Aug 2026 16:58:06 -0400 Subject: [PATCH 21/37] style: fix prettier formatting in agent-explore and sqlite-query-bounds tests --- tests/agent-explore.test.ts | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/agent-explore.test.ts b/tests/agent-explore.test.ts index 0f002e4b..6091db2b 100644 --- a/tests/agent-explore.test.ts +++ b/tests/agent-explore.test.ts @@ -844,7 +844,11 @@ describe("agent explore", () => { it("pins omission counts at and just past the limit for candidate tests and blast radius", async () => { const root = await mkExploreRepo(); await writeFile(root, "tests/auth.test.ts", "import { validateUser } from '../src/auth';\nvalidateUser('bob');\n"); - await writeFile(root, "tests/auth-spec.test.ts", "import { validateUser } from '../src/auth';\nvalidateUser('carol');\n"); + await writeFile( + root, + "tests/auth-spec.test.ts", + "import { validateUser } from '../src/auth';\nvalidateUser('carol');\n", + ); const exploreAll = await exploreCodegraph({ root, query: "validateUser" }); expect(exploreAll.candidateTests.length).toBeGreaterThanOrEqual(2); From 508fb09bebea4b49deda860e9053d6fafb6307f3 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sun, 16 Aug 2026 17:05:21 -0400 Subject: [PATCH 22/37] fix: repair main-merge test regressions for MCP tool registry and bundle entries --- tests/cli-bundle-entry.test.ts | 5 +++-- tests/mcp-skill-parity.test.ts | 4 ++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/cli-bundle-entry.test.ts b/tests/cli-bundle-entry.test.ts index 2cc932d5..f8a03ee7 100644 --- a/tests/cli-bundle-entry.test.ts +++ b/tests/cli-bundle-entry.test.ts @@ -79,11 +79,12 @@ describe("bundled CLI entry", () => { .readdirSync(binDir) .filter((name) => name.endsWith(".js")) .sort(); - // Bundle emits exactly two self-contained entrypoints (cli + queryIndexWorker). - expect(outputs).toEqual(["cli.js", "queryIndexWorker.js"]); + // Bundle emits exactly three self-contained entrypoints (cli + queryIndexWorker + rawQueryWorker). + expect(outputs).toEqual(["cli.js", "queryIndexWorker.js", "rawQueryWorker.js"]); const entry = fs.readFileSync(bundledCli, "utf8"); expect(entry).toContain("queryIndexWorker.js"); expect(fs.existsSync(path.join(binDir, "queryIndexWorker.js"))).toBe(true); + expect(fs.existsSync(path.join(binDir, "rawQueryWorker.js"))).toBe(true); }); it("keeps a leading shebang so package managers can exec the bin directly", () => { diff --git a/tests/mcp-skill-parity.test.ts b/tests/mcp-skill-parity.test.ts index 2569050f..16df6d0f 100644 --- a/tests/mcp-skill-parity.test.ts +++ b/tests/mcp-skill-parity.test.ts @@ -1,7 +1,7 @@ import fsp from "node:fs/promises"; import path from "node:path"; import { describe, expect, it } from "vitest"; -import { MCP_TOOLS } from "../src/mcp/tools.js"; +import { listCodegraphMcpTools } from "../src/mcp/tools.js"; const LEGACY_ALIAS_NAMES = new Set(["callers", "callees", "supertypes", "subtypes", "deps", "rdeps"]); @@ -18,7 +18,7 @@ describe("MCP / SKILL inventory parity", () => { [...(inventoryLine?.matchAll(/`([a-z][a-z0-9_]*)`/g) ?? [])].map((match) => match[1]!).filter(Boolean), ); - for (const tool of MCP_TOOLS) { + for (const tool of listCodegraphMcpTools()) { expect(LEGACY_ALIAS_NAMES.has(tool.name)).toBe(false); expect(listed.has(tool.name), `${tool.name} missing from SKILL MCP inventory`).toBe(true); } From 437212cf36d2b0c9c6783f40f952fc2a06b891c0 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sun, 16 Aug 2026 17:28:51 -0400 Subject: [PATCH 23/37] fix: guard session manager capacity/eviction normalizers against non-finite input --- src/session.ts | 4 ++-- tests/session.test.ts | 16 +++++++++++++++- 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/src/session.ts b/src/session.ts index cff61702..c40bc8d4 100644 --- a/src/session.ts +++ b/src/session.ts @@ -798,12 +798,12 @@ export class CodeReviewSession implements ICodeReviewSession { } function normalizeSessionManagerCapacity(value: number | undefined): number { - if (value === undefined) return DEFAULT_SESSION_MANAGER_MAX_SESSIONS; + if (value === undefined || !Number.isFinite(value)) return DEFAULT_SESSION_MANAGER_MAX_SESSIONS; return Math.max(1, Math.floor(value)); } function normalizeSessionManagerEvictionInterval(value: number | undefined): number { - if (value === undefined) return DEFAULT_SESSION_MANAGER_EVICTION_INTERVAL_MS; + if (value === undefined || !Number.isFinite(value)) return DEFAULT_SESSION_MANAGER_EVICTION_INTERVAL_MS; return Math.max(0, Math.floor(value)); } diff --git a/tests/session.test.ts b/tests/session.test.ts index 9f75e8e4..82305aca 100644 --- a/tests/session.test.ts +++ b/tests/session.test.ts @@ -1,7 +1,12 @@ import { describe, test, expect, beforeAll, afterAll, afterEach, beforeEach, vi } from "vitest"; import type { ICodeReviewSession } from "../src/index.js"; import type { BuildOptions, BuildReport, LanguageExtensionMap } from "../src/indexer/types.js"; -import { CodeReviewSession, SessionManager, createCodeReviewSession } from "../src/session.js"; +import { + CodeReviewSession, + DEFAULT_SESSION_MANAGER_MAX_SESSIONS, + SessionManager, + createCodeReviewSession, +} from "../src/session.js"; import * as indexerBuild from "../src/indexer/build-index.js"; import * as navigation from "../src/indexer/navigation.js"; import path from "node:path"; @@ -1235,6 +1240,15 @@ describe("SessionManager", () => { } }); + test("falls back to default capacity when maxSessions is NaN", async () => { + const nanManager = new SessionManager({ maxSessions: Number.NaN, evictionIntervalMs: 0 }); + try { + expect(Reflect.get(nanManager, "maxSessions")).toBe(DEFAULT_SESSION_MANAGER_MAX_SESSIONS); + } finally { + nanManager.disposeAll(); + } + }); + test("should share one initialization across concurrent same-id creation", async () => { const buildSpy = vi.spyOn(indexerBuild, "buildProjectIndexIncremental"); From dfc1656402333cfb828af38c6963ae36fea6421f Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sun, 16 Aug 2026 18:10:04 -0400 Subject: [PATCH 24/37] fix: address remaining MCP resource-bound review findings --- docs/library-api.md | 10 ++- src/agent/query-index/sessionStore.ts | 1 + src/impact/streaming.ts | 10 +-- src/index.ts | 2 + src/mcp/server.ts | 14 +++-- src/session.ts | 1 + src/sqlite.ts | 7 ++- src/sqlite/query.ts | 43 +++++++------ src/sqlite/rawQueryWorker.ts | 7 +-- src/sqlite/rawQueryWorkerPool.ts | 38 +++++------- tests/agent-session.test.ts | 65 ++++++++++++++++++++ tests/mcp-stream-cancellation.test.ts | 43 ++++++++++++- tests/raw-query-worker-lifecycle.test.ts | 18 +++++- tests/session.test.ts | 21 +++++++ tests/sqlite-query-deadline-fallback.test.ts | 34 +++++++++- 15 files changed, 252 insertions(+), 62 deletions(-) diff --git a/docs/library-api.md b/docs/library-api.md index f7dcac79..3b94d71d 100644 --- a/docs/library-api.md +++ b/docs/library-api.md @@ -732,7 +732,11 @@ console.log(mermaid); ## Read-only SQL from code ```ts -import { queryGraphSqliteRaw } from "@lzehrung/codegraph-core"; +import { + queryGraphSqliteRaw, + SqliteQueryCancelledError, + SqliteQueryDeadlineExceededError, +} from "@lzehrung/codegraph-core"; const result = await queryGraphSqliteRaw( "./codegraph.sqlite", @@ -743,7 +747,9 @@ const result = await queryGraphSqliteRaw( console.log(result.columns, result.rows); ``` -`queryGraphSqliteRaw()` is intentionally read-only. It accepts result-producing statements such as `SELECT` and `PRAGMA` and rejects mutating SQL. Its defaults bound rows, cells, and response bytes, and callers can further tighten `{ maxRows, maxBytes, maxCellBytes, deadlineMs }`. The 10-second default execution budget (`deadlineMs`) is enforced by running the query in a dedicated worker thread that is force-terminated on expiry, so it interrupts a query even mid-execution; in a degraded install where that worker asset cannot be located, the query instead runs in-process under a weaker per-row check that cannot interrupt a single blocking native call (a logged, one-time-per-process condition). +`queryGraphSqliteRaw()` is intentionally read-only. It accepts result-producing statements such as `SELECT` and `PRAGMA`, rejects mutating SQL, and bounds rows, cells, response bytes, and `{ deadlineMs }`. + +The 10-second default deadline rejects the caller promptly and requests worker termination. A native SQLite step already in progress can continue in a bounded cleanup slot until it returns; degraded installs without the worker asset use a weaker in-process check after each iterator step. Callers can catch the exported `SqliteQueryDeadlineExceededError` and `SqliteQueryCancelledError`. ## SQL artifact facts diff --git a/src/agent/query-index/sessionStore.ts b/src/agent/query-index/sessionStore.ts index 1283c1e6..6c8edd64 100644 --- a/src/agent/query-index/sessionStore.ts +++ b/src/agent/query-index/sessionStore.ts @@ -45,6 +45,7 @@ export async function ensureSessionQueryIndex( if (!existing.closing) return resolved; } if (existing) { + if (QUERY_INDEX_BY_SESSION.get(session) !== existing) continue; QUERY_INDEX_BY_SESSION.delete(session); closeState(existing); } diff --git a/src/impact/streaming.ts b/src/impact/streaming.ts index 925d4756..61790f25 100644 --- a/src/impact/streaming.ts +++ b/src/impact/streaming.ts @@ -106,7 +106,7 @@ class ImpactStreamAbandonedError extends Error { * Default cap on buffered-but-unread stream chunks before `ImpactStreamOverflowError` is * raised. True backpressure (pausing the producer until the consumer catches up) would * require the `onImpactItem` emission callback to be genuinely awaitable, which means - * awaiting it at every synchronous call site in `direct.ts`/`transitive.ts` — an invasive + * awaiting it at every synchronous call site in `direct.ts`/`transitive.ts` - an invasive * redesign of code outside this module. A hard cap is the non-invasive alternative: it * turns unbounded memory growth into an explicit, surfaced failure instead. */ @@ -235,8 +235,8 @@ export type ImpactStreamingContext = { * `streamSummary: "light"` when a caller only needs the progressive chunks and * a cheap terminal count/detail summary. * - * Cancellation: if the consumer stops iterating early — a `for await` `break`, or an - * explicit `.return()` on the generator — the async-generator return protocol resumes + * Cancellation: if the consumer stops iterating early - a `for await` `break`, or an + * explicit `.return()` on the generator - the async-generator return protocol resumes * this function's execution at its `finally` block, which aborts an internal * `AbortController`. The background `analyzeImpact()` producer's `onImpactItem` callback * checks that signal and throws once it fires, unwinding `analyzeImpact`'s in-progress @@ -419,8 +419,8 @@ export async function* analyzeImpactStreaming( error: errorMessage(error), }; } finally { - // Runs on normal completion, on a caught error, and — via the async-generator return - // protocol — when the consumer stops iterating early. In the early-abandonment case + // Runs on normal completion, on a caught error, and - via the async-generator return + // protocol - when the consumer stops iterating early. In the early-abandonment case // this is what actually stops the background analysis: it flips the shared abort // signal that the `onImpactItem` producer callback above checks and throws from. abortController.abort(); diff --git a/src/index.ts b/src/index.ts index e315c9da..17b71773 100644 --- a/src/index.ts +++ b/src/index.ts @@ -252,6 +252,8 @@ export { updateGraphSqlite, queryGraphSqlite, queryGraphSqliteRaw, + SqliteQueryCancelledError, + SqliteQueryDeadlineExceededError, SQLITE_ARTIFACT_FILE_SIGNATURES_METADATA_KEY, type SqliteGraphOptions, type SqliteGraphUpdateOptions, diff --git a/src/mcp/server.ts b/src/mcp/server.ts index 34fed573..b4194c63 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -163,6 +163,11 @@ export const DEFAULT_MCP_HTTP_SESSION_EVICTION_INTERVAL_MS = 60_000; export const DEFAULT_MCP_HTTP_BODY_TIMEOUT_MS = 30_000; export const DEFAULT_MCP_TOOL_CONCURRENCY = 4; +function normalizeMcpToolConcurrency(value: number): number { + if (!Number.isFinite(value)) return DEFAULT_MCP_TOOL_CONCURRENCY; + return Math.max(1, Math.floor(value)); +} + type LegacyMcpSession = { server: Server; transport: NodeStreamableHTTPServerTransport; @@ -656,7 +661,7 @@ function createCodegraphMcpHandlersForSession( // Probe one entry past the display limit so `truncated` is exact // rather than a `results.length === limit` heuristic (which cannot // tell "exactly limit reachable files" apart from "more exist"), - // without re-walking the whole reachable graph for an exact total — + // without re-walking the whole reachable graph for an exact total - // see finding #44. limit: limit + 1, }; @@ -895,7 +900,7 @@ function createCodegraphMcpHandlersForSession( if (parseQualifiedSymbolPath(handle)) { const snapshot = await session.loadProject({ symbolGraph: "skip" }); const resolved = requireSemanticSymbol(snapshot, handle); - // Probe one reference past the display limit so `truncated` is exact — see + // Probe one reference past the display limit so `truncated` is exact - see // `collectMcpDependencyEntries` for the rationale (finding #44). const result = await findReferences(snapshot.index, { def: resolved.def }, { maxReferences: limit + 1 }); return boundedReferencesFromResult(result, limit); @@ -1104,10 +1109,11 @@ export function createCodegraphMcpProtocolServer( }, ); let inFlightToolCalls = 0; + const toolConcurrency = normalizeMcpToolConcurrency(maxConcurrentToolCalls); server.setRequestHandler("tools/list", () => ({ tools: listCodegraphMcpTools() })); server.setRequestHandler("tools/call", async (request, ctx): Promise => { - if (inFlightToolCalls >= maxConcurrentToolCalls) { + if (inFlightToolCalls >= toolConcurrency) { throw new Error("MCP tool execution is busy; retry shortly."); } inFlightToolCalls += 1; @@ -1441,7 +1447,7 @@ async function handleLegacyMcpHttpPost( sessionRef.current = session; // The SDK transport reports every per-request validation rejection through onerror // too (bad Accept header, wrong Content-Type, malformed JSON, an unsupported - // protocol version, ...) — each of those already answered its own request with a + // protocol version, ...) - each of those already answered its own request with a // 4xx response and left the transport fully usable. Deleting the session here would // tear down an otherwise healthy session over one malformed follow-up request. Only // onclose reflects the transport actually shutting down (an explicit DELETE, an diff --git a/src/session.ts b/src/session.ts index c40bc8d4..eea9c0b7 100644 --- a/src/session.ts +++ b/src/session.ts @@ -1056,6 +1056,7 @@ export class SessionManager { if (existing?.isReady()) { continue; } + if (!existing) this.assertCapacityForNewSession(); const session = new CodeReviewSession(options); replacementSessions.push(existing ? { id, existing, session } : { id, session }); warmupPromises.push(this.trackSession(id, options, session, true, () => {})); diff --git a/src/sqlite.ts b/src/sqlite.ts index ab8f0a23..5bd1023a 100644 --- a/src/sqlite.ts +++ b/src/sqlite.ts @@ -1,3 +1,8 @@ export type { GraphQueryResult, RawSqlResult, SqliteGraphOptions, SqliteGraphUpdateOptions } from "./sqlite/types.js"; export { SQLITE_ARTIFACT_FILE_SIGNATURES_METADATA_KEY, writeGraphSqlite, updateGraphSqlite } from "./sqlite/write.js"; -export { queryGraphSqlite, queryGraphSqliteRaw } from "./sqlite/query.js"; +export { + queryGraphSqlite, + queryGraphSqliteRaw, + SqliteQueryCancelledError, + SqliteQueryDeadlineExceededError, +} from "./sqlite/query.js"; diff --git a/src/sqlite/query.ts b/src/sqlite/query.ts index 5a08b0a6..ace71a14 100644 --- a/src/sqlite/query.ts +++ b/src/sqlite/query.ts @@ -18,7 +18,7 @@ import { export { queryGraphSqlite } from "./canned-query.js"; export { SqliteQueryCancelledError, SqliteQueryDeadlineExceededError }; -/** Hard wall-clock budget for a single raw `query_sqlite` execution — see the caveat on +/** Hard wall-clock budget for a single raw `query_sqlite` execution - see the caveat on * `queryGraphSqliteRaw` about when this is actually enforceable. */ export const DEFAULT_SQLITE_QUERY_DEADLINE_MS = 10_000; @@ -35,20 +35,22 @@ let loggedInProcessDeadlineFallback = false; * Runs a bounded read-only raw SQL query. * * Preferred path: the query executes in a dedicated worker thread with a hard - * `deadlineMs` budget (`rawQueryWorkerPool.ts`). On expiry the worker thread is - * terminated outright, which stops the query even while it is blocked inside a single - * synchronous `DatabaseSync` call — a client disconnect or a slow non-recursive - * statement (large join, `ORDER BY random()`, ...) can no longer hold the host event - * loop hostage indefinitely. + * `deadlineMs` budget (`rawQueryWorkerPool.ts`). At expiry the caller receives a + * deadline error and the pool requests worker termination. A synchronous + * `DatabaseSync` call already inside SQLite may continue until that native step returns, + * but the lifecycle retains a bounded cleanup slot and the host event loop stays free. * * Degraded fallback: if the compiled worker asset cannot be located (a corrupted or - * partial install — the normal build/publish/standalone pipelines all ship it), the - * query instead runs in-process under a *per-row* elapsed-time budget. That fallback is - * strictly weaker: the budget is only checked between rows the native iterator has - * already produced, so a statement that is slow to produce its very first row (for - * example a full-table scan with no matching rows) is not bounded by it. It exists to - * keep the common case usable in a degraded install, not as a substitute for the worker - * deadline. + * partial install - the normal build/publish/standalone pipelines all ship it), the + * query instead runs in-process under a *per-row* elapsed-time budget. `node:sqlite`'s + * `DatabaseSync` exposes no interrupt/cancellation API, so once execution is inside a + * single synchronous native call there is nothing in-process that can preempt it. + * + * The fallback checks only after each native iterator step returns, including a terminal + * empty result. A statement that is slow to produce its first row or discover that it + * has no rows still blocks for its full cost before the deadline can be observed. This + * fallback exists to keep a degraded install usable, not as a substitute for the worker + * deadline; a one-time process warning makes the weakened guarantee observable. */ export async function queryGraphSqliteRaw( outputPath: string, @@ -88,7 +90,7 @@ export async function queryGraphSqliteRaw( ); } -/** Degraded fallback for `queryGraphSqliteRaw` — see its doc comment for the enforcement +/** Degraded fallback for `queryGraphSqliteRaw` - see its doc comment for the enforcement * caveat this path cannot avoid. */ async function queryGraphSqliteRawInProcessBounded( outputPath: string, @@ -124,20 +126,23 @@ async function queryGraphSqliteRawInProcessBounded( }); } -/** Throws once the wall-clock deadline has passed between two already-produced rows. See - * the fallback caveat on `queryGraphSqliteRaw`: a slow-before-first-row query is not - * caught here, only slow-*between*-rows iteration is. */ +/** Throws once the wall-clock deadline has passed after a native iterator step returns. + * See the fallback caveat on `queryGraphSqliteRaw`: a slow-before-first-row query is not + * interrupted here, only detected after that synchronous iteration step completes. */ function* withPerRowDeadline( rows: Iterable, deadlineAt: number, deadlineMs: number, signal: AbortSignal | undefined, ): Generator { - for (const row of rows) { + const iterator = rows[Symbol.iterator](); + while (true) { if (signal?.aborted) throw new SqliteQueryCancelledError(); + const next = iterator.next(); if (Date.now() > deadlineAt) { throw new SqliteQueryDeadlineExceededError(deadlineMs); } - yield row; + if (next.done) return; + yield next.value; } } diff --git a/src/sqlite/rawQueryWorker.ts b/src/sqlite/rawQueryWorker.ts index 6b5f70c4..b22ef635 100644 --- a/src/sqlite/rawQueryWorker.ts +++ b/src/sqlite/rawQueryWorker.ts @@ -19,10 +19,9 @@ export type RawQueryWorkerTask = { /** * Worker entry point. Mirrors the previous in-process body of `queryGraphSqliteRaw` * exactly: open the database read-only, assert the statement is read-only, and stream - * rows through the shared row/byte-bounded collector. Running this inside a worker - * thread lets the pool enforce a hard execution deadline by terminating the thread — - * which works even mid-synchronous-iteration, since thread termination does not need - * the blocked thread's cooperation. + * rows through the shared row/byte-bounded collector. The pool rejects the caller at its + * deadline and requests worker termination. A native SQLite step already in progress + * can continue until it returns, while the lifecycle retains its bounded cleanup slot. */ export default async function runRawQueryWorkerTask(task: RawQueryWorkerTask): Promise { return await withReadOnlySqliteDatabase(task.outputPath, (db) => { diff --git a/src/sqlite/rawQueryWorkerPool.ts b/src/sqlite/rawQueryWorkerPool.ts index c0019c41..9be5cf08 100644 --- a/src/sqlite/rawQueryWorkerPool.ts +++ b/src/sqlite/rawQueryWorkerPool.ts @@ -37,7 +37,7 @@ export class SqliteQueryCancelledError extends Error { export class SqliteQueryWorkerCleanupCapacityExceededError extends Error { constructor(maxWorkers: number) { super( - `SQLite query cleanup capacity is exhausted: ${maxWorkers} terminated worker${maxWorkers === 1 ? " is" : "s are"} still exiting. Retry after cleanup completes.`, + `SQLite query worker capacity is exhausted: ${maxWorkers} active or cleaning-up worker${maxWorkers === 1 ? " is" : "s are"} using the available slots. Retry after a query completes or cleanup finishes.`, ); this.name = "SqliteQueryWorkerCleanupCapacityExceededError"; } @@ -135,33 +135,25 @@ export function resolveRawSqlQueryWorkerPath(): string { /** * Runs a single bounded raw SQL read in a dedicated worker thread with a hard execution - * deadline. A fresh single-thread pool is created per call and destroyed afterward — - * matching the existing `prepareQueryIndexFilesInWorker` pattern — since `query_sqlite` + * deadline. A fresh single-thread pool is created per call and destroyed afterward - + * matching the existing `prepareQueryIndexFilesInWorker` pattern - since `query_sqlite` * calls are interactive, not a hot loop, and a persistent pool would need a shutdown hook * this module has no access to register. * - * On deadline expiry, Piscina's `signal` option force-terminates the worker thread - * (`worker.terminate()`) and rejects immediately — the caller never waits longer than - * `deadlineMs`, and the host event loop is never blocked by the query regardless of how - * long it runs. Cancellation is real (no further JS runs on that thread and the query - * can never touch this process's caller again), but it has one unavoidable limit shared - * by every in-process cancellation mechanism: `terminate()` cannot preempt a single - * already-in-flight synchronous native call. A query whose entire cost is inside one - * `sqlite3_step()` — a recursive CTE, or a plan that must fully sort/scan before it can - * produce a first row — keeps running on the orphaned worker thread in the background - * until that native call returns naturally; only then does the thread actually exit. - * The caller does not wait for that cleanup, but the lifecycle retains its worker slot - * until `pool.destroy()` settles. The bounded slot count makes delayed cleanup observable - * and prevents repeated cancellation from accumulating an unbounded number of workers. + * On deadline expiry, Piscina rejects the caller after requesting worker termination, so + * the caller never waits longer than `deadlineMs` and the host event loop is never + * blocked by the query. A `terminate()` request cannot preempt a single already-in-flight + * synchronous native call: a query whose entire cost is inside one `sqlite3_step()`, such + * as a recursive CTE or a plan that must fully sort or scan before a first row, continues + * on its orphaned worker thread until that native call returns naturally. + * + * The caller does not wait for cleanup, but the lifecycle retains its worker slot until + * `pool.destroy()` settles. The bounded slot count makes delayed cleanup observable and + * prevents repeated cancellation from accumulating an unbounded number of workers. * Concurrent read-only SQLite connections against the same file do not block each other, * so the lingering background reader does not stop that subsequent query from succeeding. - * Verified directly: an aborted 200M-row recursive-CTE count rejects this call in - * ~`deadlineMs`, and an immediately following query against the same file succeeds in - * milliseconds. The one place the orphaned thread is still observable is process - * shutdown: Node cannot fully tear a process down while one of its Worker threads is - * blocked in native code, so a process exit racing a runaway query can itself be - * delayed until that native call returns — a platform limit of `worker_threads`, not of - * this module, and orthogonal to the per-call deadline this function guarantees. + * Process shutdown can still wait for a Worker blocked in native code, a `worker_threads` + * platform limit that is orthogonal to this function's prompt caller deadline. */ export async function runRawSqlQueryInWorker( task: RawQueryWorkerTask, diff --git a/tests/agent-session.test.ts b/tests/agent-session.test.ts index 6b6ce147..26d53952 100644 --- a/tests/agent-session.test.ts +++ b/tests/agent-session.test.ts @@ -9,6 +9,7 @@ import path from "node:path"; import { afterEach, describe, expect, it, vi } from "vitest"; import { buildProjectIndexIncremental, type BuildReport } from "../src/index.js"; import { AGENT_FRESHNESS_CHECK_INTERVAL_MS, createAgentSession, listAgentSessionFiles } from "../src/agent/session.js"; +import type { QueryIndexHandle } from "../src/agent/query-index/update.js"; import * as symbolGraphBuild from "../src/graphs/symbol-graph-detailed.js"; import * as indexerBuild from "../src/indexer/build-index.js"; import { createProjectSnapshotIdentity } from "../src/indexer/build-cache.js"; @@ -969,5 +970,69 @@ describe("agent session", () => { ensureQueryIndexSpy.mockRestore(); } }); + + it("reuses a replacement state when invalidation races a waiting caller", async () => { + const root = await mkRepo(); + const session = createAgentSession({ root }); + const snapshot = await session.loadProject(); + const firstBuildStarted = Promise.withResolvers(); + const releaseFirstBuild = Promise.withResolvers(); + const firstHandle: QueryIndexHandle = { + store: null, + diagnostics: { + sidecarState: "created", + filesRead: 0, + filesAdded: 0, + filesUpdated: 0, + filesDeleted: 0, + fileCandidates: 0, + chunkCandidates: 0, + openMs: 0, + updateMs: 0, + candidateMs: 0, + scoringMs: 0, + }, + }; + const replacementHandle: QueryIndexHandle = { + store: null, + diagnostics: { + sidecarState: "created", + filesRead: 0, + filesAdded: 0, + filesUpdated: 0, + filesDeleted: 0, + fileCandidates: 0, + chunkCandidates: 0, + openMs: 0, + updateMs: 0, + candidateMs: 0, + scoringMs: 0, + }, + }; + let attempts = 0; + const ensureQueryIndexSpy = vi.spyOn(updateModule, "ensureQueryIndex").mockImplementation(() => { + attempts += 1; + if (attempts === 1) { + firstBuildStarted.resolve(); + return releaseFirstBuild.promise.then(() => firstHandle); + } + return Promise.resolve(replacementHandle); + }); + + try { + const waitingLoad = ensureSessionQueryIndex(session, snapshot); + await firstBuildStarted.promise; + disposeSessionQueryIndex(session); + const replacementLoad = ensureSessionQueryIndex(session, snapshot); + releaseFirstBuild.resolve(); + + await expect(waitingLoad).resolves.toBe(replacementHandle); + await expect(replacementLoad).resolves.toBe(replacementHandle); + expect(attempts).toBe(2); + } finally { + disposeSessionQueryIndex(session); + ensureQueryIndexSpy.mockRestore(); + } + }); }); }); diff --git a/tests/mcp-stream-cancellation.test.ts b/tests/mcp-stream-cancellation.test.ts index f90eb1c8..63ad22de 100644 --- a/tests/mcp-stream-cancellation.test.ts +++ b/tests/mcp-stream-cancellation.test.ts @@ -3,7 +3,11 @@ import os from "node:os"; import path from "node:path"; import { Client, InMemoryTransport } from "@modelcontextprotocol/client"; import { describe, expect, it } from "vitest"; -import { createCodegraphMcpHandlers, createCodegraphMcpProtocolServer } from "../src/mcp/server.js"; +import { + createCodegraphMcpHandlers, + createCodegraphMcpProtocolServer, + DEFAULT_MCP_TOOL_CONCURRENCY, +} from "../src/mcp/server.js"; describe("MCP query_sqlite cancellation", () => { it("forwards a cancelled tool stream to the raw query handler", async () => { @@ -48,4 +52,41 @@ describe("MCP query_sqlite cancellation", () => { await fsp.rm(root, { recursive: true, force: true }); } }); + + it("falls back to the default tool concurrency for NaN", async () => { + const root = await fsp.mkdtemp(path.join(os.tmpdir(), "cg-mcp-tool-concurrency-")); + const handlers = createCodegraphMcpHandlers({ root }); + const started = Promise.withResolvers(); + const release = Promise.withResolvers(); + let activeCalls = 0; + handlers.query_sqlite = async () => { + activeCalls += 1; + if (activeCalls === DEFAULT_MCP_TOOL_CONCURRENCY) started.resolve(); + await release.promise; + return { columns: [], rows: [] }; + }; + + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + const server = createCodegraphMcpProtocolServer(handlers, undefined, undefined, undefined, Number.NaN); + const client = new Client({ name: "mcp-tool-concurrency-test", version: "1.0.0" }); + try { + await server.connect(serverTransport); + await client.connect(clientTransport); + + const active = Array.from({ length: DEFAULT_MCP_TOOL_CONCURRENCY }, () => + client.callTool({ name: "query_sqlite", arguments: { query: "SELECT 1;" } }), + ); + await started.promise; + + await expect(client.callTool({ name: "query_sqlite", arguments: { query: "SELECT 1;" } })).rejects.toThrow( + /tool execution is busy/i, + ); + release.resolve(); + await expect(Promise.all(active)).resolves.toHaveLength(DEFAULT_MCP_TOOL_CONCURRENCY); + } finally { + release.resolve(); + await Promise.allSettled([client.close(), server.close()]); + await fsp.rm(root, { recursive: true, force: true }); + } + }); }); diff --git a/tests/raw-query-worker-lifecycle.test.ts b/tests/raw-query-worker-lifecycle.test.ts index cbbe8d26..fed1e2c0 100644 --- a/tests/raw-query-worker-lifecycle.test.ts +++ b/tests/raw-query-worker-lifecycle.test.ts @@ -51,7 +51,7 @@ describe("RawSqlQueryWorkerLifecycle", () => { expect(lifecycle.state()).toEqual({ activeWorkers: 2, maxWorkers: 2 }); await expect(lifecycle.run(task, 10_000, undefined, () => createAbortablePool(Promise.resolve()))).rejects.toThrow( - /cleanup capacity/i, + /worker capacity/i, ); firstCleanup.resolve(); @@ -60,4 +60,20 @@ describe("RawSqlQueryWorkerLifecycle", () => { await Promise.resolve(); expect(lifecycle.state()).toEqual({ activeWorkers: 0, maxWorkers: 2 }); }); + + it("describes capacity exhaustion caused by an active query", async () => { + const lifecycle = new RawSqlQueryWorkerLifecycle(1); + const complete = Promise.withResolvers<{ columns: string[]; rows: Array> }>(); + const active = lifecycle.run(task, 10_000, undefined, () => ({ + run: async () => await complete.promise, + destroy: async () => {}, + })); + + await expect(lifecycle.run(task, 10_000, undefined, () => createAbortablePool(Promise.resolve()))).rejects.toThrow( + /active or cleaning-up worker/i, + ); + + complete.resolve({ columns: [], rows: [] }); + await expect(active).resolves.toEqual({ columns: [], rows: [] }); + }); }); diff --git a/tests/session.test.ts b/tests/session.test.ts index 82305aca..79f093e6 100644 --- a/tests/session.test.ts +++ b/tests/session.test.ts @@ -1240,6 +1240,27 @@ describe("SessionManager", () => { } }); + test("enforces capacity when warming net-new sessions", async () => { + const limitedManager = new SessionManager({ maxSessions: 1, evictionIntervalMs: 0 }); + try { + await limitedManager.getOrCreateSession("existing", { + root: sampleRoot, + buildOptions: sampleBuildOptions(), + }); + + await expect( + limitedManager.warmup([ + { + id: "warm", + options: { root: sampleRoot, buildOptions: sampleBuildOptions() }, + }, + ]), + ).rejects.toThrow("Session capacity reached (1)"); + } finally { + limitedManager.disposeAll(); + } + }); + test("falls back to default capacity when maxSessions is NaN", async () => { const nanManager = new SessionManager({ maxSessions: Number.NaN, evictionIntervalMs: 0 }); try { diff --git a/tests/sqlite-query-deadline-fallback.test.ts b/tests/sqlite-query-deadline-fallback.test.ts index f7677a2d..1b99dcb5 100644 --- a/tests/sqlite-query-deadline-fallback.test.ts +++ b/tests/sqlite-query-deadline-fallback.test.ts @@ -18,6 +18,14 @@ vi.mock("../src/sqlite/rawQueryWorkerPool.js", async (importOriginal) => { }; }); +import { + SqliteQueryCancelledError, + SqliteQueryDeadlineExceededError as PublicSqliteQueryDeadlineExceededError, +} from "../src/sqlite.js"; +import { + SqliteQueryCancelledError as RootSqliteQueryCancelledError, + SqliteQueryDeadlineExceededError as RootSqliteQueryDeadlineExceededError, +} from "../src/index.js"; import { queryGraphSqliteRaw, SqliteQueryDeadlineExceededError } from "../src/sqlite/query.js"; async function withTempDb(run: (dbPath: string) => Promise): Promise { @@ -81,6 +89,24 @@ describe("SQLite raw query in-process deadline fallback", () => { }); }); + it("rejects a zero-row query that completes after the fallback deadline", async () => { + await withTempDb(async (dbPath) => { + const db = new DatabaseSync(dbPath); + db.exec("CREATE TABLE t (n INTEGER);"); + db.close(); + + const slowEmptySql = + "WITH RECURSIVE spin(x) AS (SELECT 1 UNION ALL SELECT x + 1 FROM spin WHERE x < 8000000) " + + "SELECT x FROM spin WHERE x < 0;"; + + const start = Date.now(); + await expect(queryGraphSqliteRaw(dbPath, slowEmptySql, [], { deadlineMs: 20 })).rejects.toMatchObject({ + name: "SqliteQueryDeadlineExceededError", + }); + expect(Date.now() - start).toBeGreaterThan(200); + }); + }); + it("still succeeds for an ordinary query that finishes comfortably inside its deadline", async () => { await withTempDb(async (dbPath) => { const db = new DatabaseSync(dbPath); @@ -93,8 +119,12 @@ describe("SQLite raw query in-process deadline fallback", () => { }); }); - it("exports a named deadline error for callers regardless of which path produced it", () => { - const error = new SqliteQueryDeadlineExceededError(250); + it("exports named cancellation and deadline errors from public library barrels", () => { + expect(PublicSqliteQueryDeadlineExceededError).toBe(SqliteQueryDeadlineExceededError); + expect(RootSqliteQueryDeadlineExceededError).toBe(SqliteQueryDeadlineExceededError); + expect(RootSqliteQueryCancelledError).toBe(SqliteQueryCancelledError); + + const error = new PublicSqliteQueryDeadlineExceededError(250); expect(error).toBeInstanceOf(Error); expect(error.name).toBe("SqliteQueryDeadlineExceededError"); expect(error.message).toBe("SQLite query exceeded its 250ms execution budget and was terminated."); From 4cd4f5931078cf890cf9f13548774285dac351ba Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sun, 16 Aug 2026 18:18:55 -0400 Subject: [PATCH 25/37] fix: retain warmup capacity during cancelled initialization --- docs/agent-workflows.md | 2 ++ src/session.ts | 2 +- tests/session.test.ts | 51 +++++++++++++++++++++++++++++++++++++++++ 3 files changed, 54 insertions(+), 1 deletion(-) diff --git a/docs/agent-workflows.md b/docs/agent-workflows.md index d5bbf869..dcb0461b 100644 --- a/docs/agent-workflows.md +++ b/docs/agent-workflows.md @@ -308,6 +308,8 @@ const allStats = manager.getAllStats(); console.log(Boolean(pr1Session), Boolean(pr2Session), Boolean(sameSession), allStats); ``` +`SessionManager` defaults to 32 live or initializing sessions and scans for expired sessions every 60 seconds. Set `{ maxSessions, evictionIntervalMs }` to tune those bounds; `maxSessions` also applies to net-new `warmup()` sessions, and a request over capacity throws until another session is disposed or expires. Set `evictionIntervalMs: 0` to disable periodic cleanup. + ## Streaming impact analysis Stream impact results as they are discovered so the agent can start reasoning before the full pass completes: diff --git a/src/session.ts b/src/session.ts index eea9c0b7..0b2c085a 100644 --- a/src/session.ts +++ b/src/session.ts @@ -1067,7 +1067,7 @@ export class SessionManager { const pending = this.pendingSessions.get(replacement.id); if (pending) { pending.cancelled = true; - this.pendingSessions.delete(replacement.id); + pending.retainPending = false; } replacement.session.dispose(); } diff --git a/tests/session.test.ts b/tests/session.test.ts index 79f093e6..7038a5ca 100644 --- a/tests/session.test.ts +++ b/tests/session.test.ts @@ -1261,6 +1261,57 @@ describe("SessionManager", () => { } }); + test("retains failed warmup capacity until initialization settles", async () => { + const limitedManager = new SessionManager({ maxSessions: 1, evictionIntervalMs: 0 }); + const originalBuild = indexerBuild.buildProjectIndexIncremental; + const buildStarted = Promise.withResolvers(); + const releaseBuild = Promise.withResolvers(); + const buildSpy = vi.spyOn(indexerBuild, "buildProjectIndexIncremental").mockImplementation(async (...args) => { + buildStarted.resolve(); + await releaseBuild.promise; + return await originalBuild(...args); + }); + + try { + const warmup = expect( + limitedManager.warmup([ + { + id: "warm-a", + options: { root: sampleRoot, buildOptions: sampleBuildOptions() }, + }, + { + id: "warm-b", + options: { root: sampleRoot, buildOptions: sampleBuildOptions() }, + }, + ]), + ).rejects.toThrow("Session capacity reached (1)"); + await buildStarted.promise; + await warmup; + + await expect( + limitedManager.getOrCreateSession("after-failed-warmup", { + root: sampleRoot, + buildOptions: sampleBuildOptions(), + }), + ).rejects.toThrow("Session capacity reached (1)"); + + releaseBuild.resolve(); + await vi.waitFor(() => { + expect(Reflect.get(limitedManager, "pendingSessions").size).toBe(0); + }); + await expect( + limitedManager.getOrCreateSession("after-failed-warmup", { + root: sampleRoot, + buildOptions: sampleBuildOptions(), + }), + ).resolves.toBeInstanceOf(CodeReviewSession); + } finally { + releaseBuild.resolve(); + buildSpy.mockRestore(); + limitedManager.disposeAll(); + } + }); + test("falls back to default capacity when maxSessions is NaN", async () => { const nanManager = new SessionManager({ maxSessions: Number.NaN, evictionIntervalMs: 0 }); try { From 8c5412a7305ea9eff8c1e86ee9e6578e6f173b94 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sun, 16 Aug 2026 18:20:16 -0400 Subject: [PATCH 26/37] docs: refresh JavaScript coverage report --- docs/coverage/js.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/coverage/js.md b/docs/coverage/js.md index fcdfef37..2e4954e8 100644 --- a/docs/coverage/js.md +++ b/docs/coverage/js.md @@ -6,14 +6,15 @@ Source: `coverage/js/lcov.info` | Metric | Hit | Found | Coverage | | --------- | ----: | ----: | -------: | -| Lines | 27491 | 30264 | 90.84% | -| Functions | 4565 | 4844 | 94.24% | -| Branches | 20716 | 26123 | 79.30% | +| Lines | 27559 | 30335 | 90.85% | +| Functions | 4598 | 4885 | 94.12% | +| Branches | 20761 | 26158 | 79.37% | ## Least-covered Files | File | Lines | Functions | Branches | | ---------------------------------------------- | -----: | --------: | -------: | +| `src/sqlite/rawQueryWorker.ts` | 0.00% | 0.00% | 0.00% | | `src/languages/definitions/htmlStub.ts` | 50.00% | 50.00% | n/a | | `src/languages/definitions/javascript.ts` | 62.50% | 83.33% | 57.89% | | `src/cli/explore.ts` | 62.50% | 100.00% | 50.00% | @@ -33,7 +34,6 @@ Source: `coverage/js/lcov.info` | `src/agent/query-index/workerPool.ts` | 70.97% | 62.50% | 56.25% | | `src/agent/followUps.ts` | 72.15% | 77.78% | 53.91% | | `src/cli/context.ts` | 72.22% | 60.78% | 75.25% | -| `src/indexer/imports/languageSpecific.ts` | 73.97% | 100.00% | 63.74% | ## Type-Only Or Re-Export Files From 73037eb286dca25d5e4eb268927817c21736a5a1 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sun, 16 Aug 2026 18:23:56 -0400 Subject: [PATCH 27/37] docs: refresh JavaScript coverage report --- docs/coverage/js.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/coverage/js.md b/docs/coverage/js.md index 2e4954e8..50fe23ad 100644 --- a/docs/coverage/js.md +++ b/docs/coverage/js.md @@ -6,7 +6,7 @@ Source: `coverage/js/lcov.info` | Metric | Hit | Found | Coverage | | --------- | ----: | ----: | -------: | -| Lines | 27559 | 30335 | 90.85% | +| Lines | 27558 | 30335 | 90.85% | | Functions | 4598 | 4885 | 94.12% | | Branches | 20761 | 26158 | 79.37% | From efe41a4a1a6106407168d8efd40a73baae0de357 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sun, 16 Aug 2026 18:29:20 -0400 Subject: [PATCH 28/37] fix: release capacity after failed warmup initialization --- docs/agent-workflows.md | 2 +- src/session.ts | 7 +++++++ tests/session.test.ts | 27 +++++++++++++++++++++++++++ 3 files changed, 35 insertions(+), 1 deletion(-) diff --git a/docs/agent-workflows.md b/docs/agent-workflows.md index dcb0461b..5b131bc2 100644 --- a/docs/agent-workflows.md +++ b/docs/agent-workflows.md @@ -308,7 +308,7 @@ const allStats = manager.getAllStats(); console.log(Boolean(pr1Session), Boolean(pr2Session), Boolean(sameSession), allStats); ``` -`SessionManager` defaults to 32 live or initializing sessions and scans for expired sessions every 60 seconds. Set `{ maxSessions, evictionIntervalMs }` to tune those bounds; `maxSessions` also applies to net-new `warmup()` sessions, and a request over capacity throws until another session is disposed or expires. Set `evictionIntervalMs: 0` to disable periodic cleanup. +`SessionManager` defaults to 32 live or initializing sessions and scans for expired sessions every 60 seconds. Set `{ maxSessions, evictionIntervalMs }` to tune those bounds; `maxSessions` also applies to net-new `warmup()` sessions. An over-capacity request throws until capacity frees after a session is disposed, expires, or a canceled or failed warmup initialization settles. Set `evictionIntervalMs: 0` to disable periodic cleanup. ## Streaming impact analysis diff --git a/src/session.ts b/src/session.ts index 0b2c085a..d82d3489 100644 --- a/src/session.ts +++ b/src/session.ts @@ -1068,6 +1068,13 @@ export class SessionManager { if (pending) { pending.cancelled = true; pending.retainPending = false; + void pending.promise + .finally(() => { + if (this.pendingSessions.get(replacement.id) === pending && !pending.retainPending) { + this.pendingSessions.delete(replacement.id); + } + }) + .catch(() => {}); } replacement.session.dispose(); } diff --git a/tests/session.test.ts b/tests/session.test.ts index 7038a5ca..b542ea19 100644 --- a/tests/session.test.ts +++ b/tests/session.test.ts @@ -1312,6 +1312,33 @@ describe("SessionManager", () => { } }); + test("releases failed warmup capacity after initialization rejects", async () => { + const limitedManager = new SessionManager({ maxSessions: 1, evictionIntervalMs: 0 }); + const missingRoot = path.join(os.tmpdir(), `cg-session-warmup-missing-${Date.now()}`); + try { + await expect( + limitedManager.warmup([ + { + id: "broken", + options: { root: missingRoot, buildOptions: sampleBuildOptions() }, + }, + ]), + ).rejects.toThrow(); + + await vi.waitFor(() => { + expect(Reflect.get(limitedManager, "pendingSessions").size).toBe(0); + }); + await expect( + limitedManager.getOrCreateSession("after-failed-init", { + root: sampleRoot, + buildOptions: sampleBuildOptions(), + }), + ).resolves.toBeInstanceOf(CodeReviewSession); + } finally { + limitedManager.disposeAll(); + } + }); + test("falls back to default capacity when maxSessions is NaN", async () => { const nanManager = new SessionManager({ maxSessions: Number.NaN, evictionIntervalMs: 0 }); try { From 8158c5d052b374a3e78d012fd2416ce1a53cec23 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sun, 16 Aug 2026 18:32:34 -0400 Subject: [PATCH 29/37] docs: refresh JavaScript coverage report --- docs/coverage/js.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/coverage/js.md b/docs/coverage/js.md index 50fe23ad..bdf308c2 100644 --- a/docs/coverage/js.md +++ b/docs/coverage/js.md @@ -6,9 +6,9 @@ Source: `coverage/js/lcov.info` | Metric | Hit | Found | Coverage | | --------- | ----: | ----: | -------: | -| Lines | 27558 | 30335 | 90.85% | -| Functions | 4598 | 4885 | 94.12% | -| Branches | 20761 | 26158 | 79.37% | +| Lines | 27561 | 30338 | 90.85% | +| Functions | 4600 | 4887 | 94.13% | +| Branches | 20765 | 26162 | 79.37% | ## Least-covered Files From e9f75b291a9c5b501ae2c60e0caa05759b0c5a46 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sun, 16 Aug 2026 23:27:46 -0400 Subject: [PATCH 30/37] fix: address review resource-bound findings --- docs/agent-workflows.md | 4 +- docs/coverage/js.md | 78 ++++++++++---------- docs/library-api.md | 5 +- src/impact/analyzer.ts | 15 ++++ src/impact/streaming.ts | 15 ++-- src/impact/types.ts | 2 + src/index.ts | 1 + src/mcp/http.ts | 14 ++-- src/session.ts | 54 +++++++------- src/sqlite.ts | 1 + src/sqlite/query.ts | 26 +++++-- src/sqlite/rawQueryWorkerPool.ts | 2 +- tests/session.test.ts | 64 ++++++++++------ tests/sqlite-query-deadline-fallback.test.ts | 26 +++++-- 14 files changed, 185 insertions(+), 122 deletions(-) diff --git a/docs/agent-workflows.md b/docs/agent-workflows.md index 5b131bc2..cc269667 100644 --- a/docs/agent-workflows.md +++ b/docs/agent-workflows.md @@ -308,7 +308,7 @@ const allStats = manager.getAllStats(); console.log(Boolean(pr1Session), Boolean(pr2Session), Boolean(sameSession), allStats); ``` -`SessionManager` defaults to 32 live or initializing sessions and scans for expired sessions every 60 seconds. Set `{ maxSessions, evictionIntervalMs }` to tune those bounds; `maxSessions` also applies to net-new `warmup()` sessions. An over-capacity request throws until capacity frees after a session is disposed, expires, or a canceled or failed warmup initialization settles. Set `evictionIntervalMs: 0` to disable periodic cleanup. +`SessionManager` defaults to 32 live or initializing sessions and scans for expired sessions every 60 seconds. Set `{ maxSessions, evictionIntervalMs }` to tune those bounds; `maxSessions` also applies to net-new `warmup()` sessions. Capacity frees immediately after a ready session is disposed or expires, or after a canceled or failed initialization settles. Set `evictionIntervalMs: 0` to disable periodic cleanup. ## Streaming impact analysis @@ -339,6 +339,8 @@ for await (const chunk of analyzeImpactStreaming(root, index, { } ``` +Handle `error` as terminal: an overfull bounded queue does not emit `complete`. Breaking iteration or calling `.return()` cancels background work at its next analysis boundary; a synchronous lookup already in progress cannot be interrupted. + Use the same pattern through a warm session when repeated review passes matter: ```ts diff --git a/docs/coverage/js.md b/docs/coverage/js.md index bdf308c2..4e9b8c2a 100644 --- a/docs/coverage/js.md +++ b/docs/coverage/js.md @@ -4,52 +4,52 @@ Source: `coverage/js/lcov.info` ## Summary -| Metric | Hit | Found | Coverage | -| --------- | ----: | ----: | -------: | -| Lines | 27561 | 30338 | 90.85% | -| Functions | 4600 | 4887 | 94.13% | -| Branches | 20765 | 26162 | 79.37% | +| Metric | Hit | Found | Coverage | +| --- | ---: | ---: | ---: | +| Lines | 27578 | 30358 | 90.84% | +| Functions | 4602 | 4889 | 94.13% | +| Branches | 20773 | 26172 | 79.37% | ## Least-covered Files -| File | Lines | Functions | Branches | -| ---------------------------------------------- | -----: | --------: | -------: | -| `src/sqlite/rawQueryWorker.ts` | 0.00% | 0.00% | 0.00% | -| `src/languages/definitions/htmlStub.ts` | 50.00% | 50.00% | n/a | -| `src/languages/definitions/javascript.ts` | 62.50% | 83.33% | 57.89% | -| `src/cli/explore.ts` | 62.50% | 100.00% | 50.00% | -| `src/languages/definitions/typescript.ts` | 62.86% | 81.82% | 42.86% | -| `src/impact/call-compatibility/textScanner.ts` | 66.04% | 100.00% | 70.75% | -| `src/cli/index.ts` | 66.67% | 100.00% | 58.97% | -| `src/indexer.ts` | 66.67% | 66.67% | 100.00% | -| `src/cli/bootstrap.ts` | 66.67% | 50.00% | n/a | -| `src/languages/definitions/css.ts` | 66.67% | 0.00% | n/a | -| `src/languages/definitions/html.ts` | 66.67% | 0.00% | n/a | -| `src/languages/definitions/less.ts` | 66.67% | 0.00% | n/a | -| `src/languages/definitions/sql.ts` | 66.67% | 0.00% | n/a | -| `src/languages/definitions/svelte.ts` | 66.67% | 0.00% | n/a | -| `src/languages/definitions/vue.ts` | 66.67% | 0.00% | n/a | -| `src/agent/handles.ts` | 68.00% | 91.67% | 45.45% | -| `src/cli/artifact.ts` | 70.00% | 100.00% | 72.22% | -| `src/agent/query-index/workerPool.ts` | 70.97% | 62.50% | 56.25% | -| `src/agent/followUps.ts` | 72.15% | 77.78% | 53.91% | -| `src/cli/context.ts` | 72.22% | 60.78% | 75.25% | +| File | Lines | Functions | Branches | +| --- | ---: | ---: | ---: | +| `src/sqlite/rawQueryWorker.ts` | 0.00% | 0.00% | 0.00% | +| `src/languages/definitions/htmlStub.ts` | 50.00% | 50.00% | n/a | +| `src/languages/definitions/javascript.ts` | 62.50% | 83.33% | 57.89% | +| `src/cli/explore.ts` | 62.50% | 100.00% | 50.00% | +| `src/languages/definitions/typescript.ts` | 62.86% | 81.82% | 42.86% | +| `src/impact/call-compatibility/textScanner.ts` | 66.04% | 100.00% | 70.75% | +| `src/cli/index.ts` | 66.67% | 100.00% | 58.97% | +| `src/indexer.ts` | 66.67% | 66.67% | 100.00% | +| `src/cli/bootstrap.ts` | 66.67% | 50.00% | n/a | +| `src/languages/definitions/css.ts` | 66.67% | 0.00% | n/a | +| `src/languages/definitions/html.ts` | 66.67% | 0.00% | n/a | +| `src/languages/definitions/less.ts` | 66.67% | 0.00% | n/a | +| `src/languages/definitions/sql.ts` | 66.67% | 0.00% | n/a | +| `src/languages/definitions/svelte.ts` | 66.67% | 0.00% | n/a | +| `src/languages/definitions/vue.ts` | 66.67% | 0.00% | n/a | +| `src/agent/handles.ts` | 68.00% | 91.67% | 45.45% | +| `src/cli/artifact.ts` | 70.00% | 100.00% | 72.22% | +| `src/agent/query-index/workerPool.ts` | 70.97% | 62.50% | 56.25% | +| `src/agent/followUps.ts` | 72.15% | 77.78% | 53.91% | +| `src/cli/context.ts` | 72.22% | 60.78% | 75.25% | ## Type-Only Or Re-Export Files These files have line records but no function or branch records, so they are tracked outside the runtime ranking. -| File | Lines | Functions | Branches | -| --------------------------------------- | ------: | --------: | -------: | -| `src/cliBootstrap.ts` | 0.00% | n/a | n/a | -| `src/util/identifiers.ts` | 100.00% | n/a | n/a | -| `src/languages/definitions/jsFamily.ts` | 100.00% | n/a | n/a | -| `src/duplicate-keywords.ts` | 100.00% | n/a | n/a | -| `src/impact/types.ts` | 100.00% | n/a | n/a | -| `src/languages/definitions/adoc.ts` | 100.00% | n/a | n/a | -| `src/languages/definitions/astro.ts` | 100.00% | n/a | n/a | -| `src/languages/definitions/hbs.ts` | 100.00% | n/a | n/a | -| `src/languages/definitions/markdown.ts` | 100.00% | n/a | n/a | -| `src/languages/definitions/mdx.ts` | 100.00% | n/a | n/a | +| File | Lines | Functions | Branches | +| --- | ---: | ---: | ---: | +| `src/cliBootstrap.ts` | 0.00% | n/a | n/a | +| `src/languages/definitions/jsFamily.ts` | 100.00% | n/a | n/a | +| `src/duplicate-keywords.ts` | 100.00% | n/a | n/a | +| `src/impact/types.ts` | 100.00% | n/a | n/a | +| `src/languages/definitions/adoc.ts` | 100.00% | n/a | n/a | +| `src/languages/definitions/astro.ts` | 100.00% | n/a | n/a | +| `src/languages/definitions/hbs.ts` | 100.00% | n/a | n/a | +| `src/languages/definitions/markdown.ts` | 100.00% | n/a | n/a | +| `src/languages/definitions/mdx.ts` | 100.00% | n/a | n/a | +| `src/languages/definitions/rst.ts` | 100.00% | n/a | n/a | Generated from LCOV by `node ./scripts/coverage-markdown.mjs`. diff --git a/docs/library-api.md b/docs/library-api.md index 3b94d71d..3f807cb0 100644 --- a/docs/library-api.md +++ b/docs/library-api.md @@ -736,6 +736,7 @@ import { queryGraphSqliteRaw, SqliteQueryCancelledError, SqliteQueryDeadlineExceededError, + SqliteQueryWorkerCleanupCapacityExceededError, } from "@lzehrung/codegraph-core"; const result = await queryGraphSqliteRaw( @@ -749,7 +750,7 @@ console.log(result.columns, result.rows); `queryGraphSqliteRaw()` is intentionally read-only. It accepts result-producing statements such as `SELECT` and `PRAGMA`, rejects mutating SQL, and bounds rows, cells, response bytes, and `{ deadlineMs }`. -The 10-second default deadline rejects the caller promptly and requests worker termination. A native SQLite step already in progress can continue in a bounded cleanup slot until it returns; degraded installs without the worker asset use a weaker in-process check after each iterator step. Callers can catch the exported `SqliteQueryDeadlineExceededError` and `SqliteQueryCancelledError`. +The 10-second default deadline rejects the caller promptly and requests worker termination. A native SQLite step already in progress can continue in a bounded cleanup slot until it returns; degraded installs without the worker asset use a weaker in-process check after each iterator step. Callers can catch the exported `SqliteQueryDeadlineExceededError`, `SqliteQueryCancelledError`, and `SqliteQueryWorkerCleanupCapacityExceededError`; cancellation is a stable generic message so it exposes no MCP client details. ## SQL artifact facts @@ -956,7 +957,7 @@ Use the exported TypeScript APIs when another program is composing deterministic - `buildReviewReport()` returns a review bundle with `schemaVersion`, changed files, changed symbols, `graphDelta`, candidate tests, `riskSummary`, `reviewTasks`, an offline `markdownLinks` result for Markdown sources in the analysis scope when there are changes, optional duplicate sibling-check tasks, optional `sqlContext`, compatibility hints when available, and diagnostics. Accepts an optional third argument, `{ index?, loadIndex?, duplicateAnalysis?, loadDuplicateAnalysis? }`, so a caller that already holds a warm `ProjectIndex` (or wants to defer loading it until review work actually needs it) and, for repeated review calls, a `DuplicatePreparedAnalysis` from `prepareDuplicateAnalysis()` can skip redundant rebuilds. The MCP `review` tool uses the lazy forms to avoid paying index or duplicate-analysis cost on no-change reviews. - `analyzeImpactFromDiff()` returns the full or compact impact report shape for batch consumers, including an offline `markdownLinks` result for Markdown sources in the analysis scope when diffs are non-empty and changed-symbol `callCompatibility` hints when available. -- `analyzeImpactStreaming()` emits progress and incremental chunks, then a final `complete.report` summary. Streaming always returns `format: "stream-summary"`. By default this includes the same key structured fields needed by pack builders: changed files, changed symbols, impacted items, Markdown link findings, suggestions, export summaries, re-export chains, ranked top impacts, surface area, clusters, cycles, graph edges, diagnostics, and warning text. Set `streamSummary: "light"` to drop suggestions, export summaries, re-export chains, ranked top impacts, graph metadata, cycles, clusters, and surface area from the final report. +- `analyzeImpactStreaming()` emits progress and incremental chunks, then a final `complete.report` summary on success. Streaming always returns `format: "stream-summary"`. By default this includes the same key structured fields needed by pack builders: changed files, changed symbols, impacted items, Markdown link findings, suggestions, export summaries, re-export chains, ranked top impacts, surface area, clusters, cycles, graph edges, diagnostics, and warning text. Set `streamSummary: "light"` to drop suggestions, export summaries, re-export chains, ranked top impacts, graph metadata, cycles, clusters, and surface area from the final report. A bounded queue overflow instead emits terminal `error` without `complete`; ending iteration early cancels later analysis batches, but cannot interrupt a synchronous lookup already in progress. Review-pack builders should preserve symbol handles, diff snippets, callsites, `callCompatibility`, diagnostics, candidate-test confidence, impact reasons, and graph edge metadata. Render prose only at the final UI or prompt boundary. diff --git a/src/impact/analyzer.ts b/src/impact/analyzer.ts index 39b40dd3..4d89b9ee 100644 --- a/src/impact/analyzer.ts +++ b/src/impact/analyzer.ts @@ -46,6 +46,13 @@ function compareImpactItems(left: ImpactItem, right: ImpactItem): number { return 0; } +function throwIfImpactAnalysisAborted(signal: AbortSignal | undefined): void { + if (!signal?.aborted) return; + const error = new Error("Impact analysis was cancelled."); + error.name = "AbortError"; + throw error; +} + export async function analyzeImpact( index: ProjectIndex, changedSymbols: ChangedSymbol[], @@ -62,6 +69,7 @@ export async function analyzeImpact( refContextLines, refBlockMaxLines, onImpactItem, + signal, } = options; const diagnostics = options.diagnostics; const projectRoot = @@ -76,6 +84,7 @@ export async function analyzeImpact( const isIgnored = projectRoot ? createImpactIgnoreMatcher(projectRoot, ignoreGlobs) : () => false; const referenceCache = createReferenceLookupCache(); const workBudget = createImpactWorkBudget(options); + throwIfImpactAnalysisAborted(signal); const impacted = new Map(); const processedSymbols = new Set(); @@ -159,6 +168,7 @@ export async function analyzeImpact( return !isIgnored(file); }, }); + throwIfImpactAnalysisAborted(signal); syncBudgetDiagnostics(diagnostics, workBudget); await analyzeDirectReferences({ index, @@ -172,6 +182,7 @@ export async function analyzeImpact( emitImpactItem, }); syncBudgetDiagnostics(diagnostics, workBudget); + throwIfImpactAnalysisAborted(signal); } // Seed transitive impact from changed files. This is NOT redundant with @@ -179,17 +190,21 @@ export async function analyzeImpact( // (they no longer exist), so they would never enter `impacted` through the symbol // loop above. seedTransitiveFromFiles plants them directly so the transitive pass // can propagate their impact to dependents. + throwIfImpactAnalysisAborted(signal); if (!options.membersOnly && !isImpactDeadlineExceeded(workBudget)) { seedTransitiveFromFiles(index, impacted, changedFiles, normalizedOptions, reverseDeps, emitImpactItem); } + throwIfImpactAnalysisAborted(signal); // Transitive impact via graph traversal (skip if membersOnly) if (!options.membersOnly && !isImpactDeadlineExceeded(workBudget)) { analyzeTransitiveImpact(impacted, depth, normalizedOptions, isIndexTestFile, reverseDeps, emitImpactItem); } + throwIfImpactAnalysisAborted(signal); syncBudgetDiagnostics(diagnostics, workBudget); const sorted = Array.from(impacted.values()).sort(compareImpactItems); + throwIfImpactAnalysisAborted(signal); for (const item of sorted) { emitImpactItem(item, "final"); } diff --git a/src/impact/streaming.ts b/src/impact/streaming.ts index 61790f25..c66c71e4 100644 --- a/src/impact/streaming.ts +++ b/src/impact/streaming.ts @@ -31,7 +31,7 @@ export type ImpactStreamChunk = | { type: "projectFiles"; files: ProjectFileInfo[] } | import("../types.js").ProgressUpdate | { type: "changedSymbol"; symbol: ChangedSymbol } - | { type: "impactItem"; item: ImpactItem; partial?: boolean } + | { type: "impactItem"; item: ImpactItem; partial?: true } | { type: "complete"; summary: { totalChanged: number; totalImpacted: number }; @@ -40,7 +40,7 @@ export type ImpactStreamChunk = | { type: "error"; error: string }; type PublicImpactStreamingOptions = Options extends unknown - ? Omit + ? Omit : never; type WithoutCompact = Options extends unknown ? Omit : never; @@ -238,12 +238,10 @@ export type ImpactStreamingContext = { * Cancellation: if the consumer stops iterating early - a `for await` `break`, or an * explicit `.return()` on the generator - the async-generator return protocol resumes * this function's execution at its `finally` block, which aborts an internal - * `AbortController`. The background `analyzeImpact()` producer's `onImpactItem` callback - * checks that signal and throws once it fires, unwinding `analyzeImpact`'s in-progress - * work (no further batches or transitive passes run) instead of letting the whole - * analysis complete unread. This needs no signal parameter on the public API: `yield*` - * delegation (used by `session.ts`'s `analyzeImpactStream`) forwards a caller's - * `.return()` through automatically. + * `AbortController`. The background `analyzeImpact()` producer receives that signal and + * checks it at analysis work boundaries and before emitting items, so it does not start + * later batches or transitive work after abandonment. An in-progress synchronous lookup + * still runs until it returns because it cannot be preempted by JavaScript. */ export async function* analyzeImpactStreaming( projectRoot: string, @@ -337,6 +335,7 @@ export async function* analyzeImpactStreaming( fileLevelFallback, fileLevelFallbackPaths, diagnostics, + signal: abortController.signal, onImpactItem: (item, phase) => { if (abortController.signal.aborted) { throw new ImpactStreamAbandonedError(); diff --git a/src/impact/types.ts b/src/impact/types.ts index aea9c936..489db33a 100644 --- a/src/impact/types.ts +++ b/src/impact/types.ts @@ -571,4 +571,6 @@ export type ImpactOptions = DiffProviderOptions & { diagnostics?: ImpactDiagnostics; /** @internal Internal callback used by streaming analysis to emit progressive impact snapshots */ onImpactItem?: (item: ImpactItem, phase: "partial" | "final") => void; + /** @internal Cancellation signal used to stop background streaming analysis at work boundaries. */ + signal?: AbortSignal; }; diff --git a/src/index.ts b/src/index.ts index 17b71773..45db55bb 100644 --- a/src/index.ts +++ b/src/index.ts @@ -254,6 +254,7 @@ export { queryGraphSqliteRaw, SqliteQueryCancelledError, SqliteQueryDeadlineExceededError, + SqliteQueryWorkerCleanupCapacityExceededError, SQLITE_ARTIFACT_FILE_SIGNATURES_METADATA_KEY, type SqliteGraphOptions, type SqliteGraphUpdateOptions, diff --git a/src/mcp/http.ts b/src/mcp/http.ts index c1870fcb..f41156b0 100644 --- a/src/mcp/http.ts +++ b/src/mcp/http.ts @@ -23,14 +23,12 @@ export async function readJsonRequestBody( timeoutMs: number, ): Promise { const contentLength = getContentLength(request); - if (contentLength !== undefined && contentLength > maxBytes) { - request.resume(); - return { status: "too_large" }; - } + const knownTooLarge = contentLength !== undefined && contentLength > maxBytes; return await new Promise((resolve) => { const chunks: Buffer[] = []; let bytes = 0; + let tooLarge = knownTooLarge; let settled = false; const deadline = setTimeout(() => settle({ status: "timeout" }, true), timeoutMs); deadline.unref?.(); @@ -50,15 +48,21 @@ export async function readJsonRequestBody( resolve(result); }; const onData = (chunk: string | Buffer): void => { + if (tooLarge) return; const buffer = typeof chunk === "string" ? Buffer.from(chunk) : chunk; bytes += buffer.byteLength; if (bytes > maxBytes) { - settle({ status: "too_large" }, true); + tooLarge = true; + chunks.length = 0; return; } chunks.push(buffer); }; const onEnd = (): void => { + if (tooLarge) { + settle({ status: "too_large" }, false); + return; + } const rawBody = Buffer.concat(chunks).toString("utf8"); try { const body: unknown = rawBody.length ? JSON.parse(rawBody) : null; diff --git a/src/session.ts b/src/session.ts index d82d3489..4d85f9fa 100644 --- a/src/session.ts +++ b/src/session.ts @@ -811,17 +811,16 @@ function normalizeSessionManagerEvictionInterval(value: number | undefined): num * Session manager for multiple concurrent sessions * Useful for agents handling multiple repositories or PRs */ + +type PendingSession = { + cancelled: boolean; + fingerprint: string; + retainPending: boolean; + promise: Promise; +}; export class SessionManager { private sessions = new Map(); - private pendingSessions = new Map< - string, - { - cancelled: boolean; - fingerprint: string; - retainPending: boolean; - promise: Promise; - } - >(); + private pendingSessions = new Map(); private readonly maxSessions: number; private readonly evictionTimer: ReturnType | undefined; @@ -834,6 +833,18 @@ export class SessionManager { } } + private cancelPendingSession(sessionId: string, pending: PendingSession): void { + pending.cancelled = true; + pending.retainPending = false; + void pending.promise + .finally(() => { + if (this.pendingSessions.get(sessionId) === pending && !pending.retainPending) { + this.pendingSessions.delete(sessionId); + } + }) + .catch(() => {}); + } + private assertCapacityForNewSession(): void { this.cleanupExpired(); if (this.sessions.size + this.pendingSessions.size >= this.maxSessions) { @@ -868,6 +879,9 @@ export class SessionManager { ): Promise | undefined { const pending = this.pendingSessions.get(sessionId); if (!pending) return undefined; + if (pending.cancelled) { + throw new Error(`Session "${sessionId}" is still cancelling initialization. Retry after it settles.`); + } const requestedFingerprint = sessionIdentityFingerprint(resolveSessionIdentity(options)); if (pending.fingerprint !== requestedFingerprint) { const existing = this.sessions.get(sessionId); @@ -962,10 +976,7 @@ export class SessionManager { */ disposeSession(sessionId: string): void { const pending = this.pendingSessions.get(sessionId); - if (pending) { - pending.cancelled = true; - this.pendingSessions.delete(sessionId); - } + if (pending) this.cancelPendingSession(sessionId, pending); const session = this.sessions.get(sessionId); if (session) { session.dispose(); @@ -977,10 +988,9 @@ export class SessionManager { * Dispose of all sessions */ disposeAll(): void { - for (const pending of this.pendingSessions.values()) { - pending.cancelled = true; + for (const [sessionId, pending] of this.pendingSessions) { + this.cancelPendingSession(sessionId, pending); } - this.pendingSessions.clear(); for (const session of this.sessions.values()) { session.dispose(); } @@ -1065,17 +1075,7 @@ export class SessionManager { } catch (error) { for (const replacement of replacementSessions) { const pending = this.pendingSessions.get(replacement.id); - if (pending) { - pending.cancelled = true; - pending.retainPending = false; - void pending.promise - .finally(() => { - if (this.pendingSessions.get(replacement.id) === pending && !pending.retainPending) { - this.pendingSessions.delete(replacement.id); - } - }) - .catch(() => {}); - } + if (pending) this.cancelPendingSession(replacement.id, pending); replacement.session.dispose(); } throw error; diff --git a/src/sqlite.ts b/src/sqlite.ts index 5bd1023a..5cf133fa 100644 --- a/src/sqlite.ts +++ b/src/sqlite.ts @@ -5,4 +5,5 @@ export { queryGraphSqliteRaw, SqliteQueryCancelledError, SqliteQueryDeadlineExceededError, + SqliteQueryWorkerCleanupCapacityExceededError, } from "./sqlite/query.js"; diff --git a/src/sqlite/query.ts b/src/sqlite/query.ts index ace71a14..ff140a2d 100644 --- a/src/sqlite/query.ts +++ b/src/sqlite/query.ts @@ -13,10 +13,11 @@ import { runRawSqlQueryInWorker, SqliteQueryCancelledError, SqliteQueryDeadlineExceededError, + SqliteQueryWorkerCleanupCapacityExceededError, } from "./rawQueryWorkerPool.js"; export { queryGraphSqlite } from "./canned-query.js"; -export { SqliteQueryCancelledError, SqliteQueryDeadlineExceededError }; +export { SqliteQueryCancelledError, SqliteQueryDeadlineExceededError, SqliteQueryWorkerCleanupCapacityExceededError }; /** Hard wall-clock budget for a single raw `query_sqlite` execution - see the caveat on * `queryGraphSqliteRaw` about when this is actually enforceable. */ @@ -136,13 +137,22 @@ function* withPerRowDeadline( signal: AbortSignal | undefined, ): Generator { const iterator = rows[Symbol.iterator](); - while (true) { - if (signal?.aborted) throw new SqliteQueryCancelledError(); - const next = iterator.next(); - if (Date.now() > deadlineAt) { - throw new SqliteQueryDeadlineExceededError(deadlineMs); + let completed = false; + try { + while (true) { + if (signal?.aborted) throw new SqliteQueryCancelledError(); + const next = iterator.next(); + if (signal?.aborted) throw new SqliteQueryCancelledError(); + if (Date.now() > deadlineAt) { + throw new SqliteQueryDeadlineExceededError(deadlineMs); + } + if (next.done) { + completed = true; + return; + } + yield next.value; } - if (next.done) return; - yield next.value; + } finally { + if (!completed) iterator.return?.(); } } diff --git a/src/sqlite/rawQueryWorkerPool.ts b/src/sqlite/rawQueryWorkerPool.ts index 9be5cf08..3294bc73 100644 --- a/src/sqlite/rawQueryWorkerPool.ts +++ b/src/sqlite/rawQueryWorkerPool.ts @@ -29,7 +29,7 @@ export class SqliteQueryDeadlineExceededError extends Error { export class SqliteQueryCancelledError extends Error { constructor() { - super("SQLite query was cancelled by the MCP client."); + super("SQLite query was cancelled."); this.name = "SqliteQueryCancelledError"; } } diff --git a/tests/session.test.ts b/tests/session.test.ts index b542ea19..7116b20f 100644 --- a/tests/session.test.ts +++ b/tests/session.test.ts @@ -1401,7 +1401,8 @@ describe("SessionManager", () => { } }); - test("should allow immediate recreation after disposing a pending session", async () => { + test("retains pending initialization capacity after disposal until it settles", async () => { + const limitedManager = new SessionManager({ maxSessions: 1, evictionIntervalMs: 0 }); const originalBuild = indexerBuild.buildProjectIndexIncremental; let releaseBuild: (() => void) | null = null; const buildGate = new Promise((resolve) => { @@ -1413,31 +1414,42 @@ describe("SessionManager", () => { }); try { - const firstSession = manager.getOrCreateSession("pending", { + const firstSession = limitedManager.getOrCreateSession("pending", { root: sampleRoot, buildOptions: sampleBuildOptions(), }); await Promise.resolve(); - manager.disposeSession("pending"); - const secondSession = manager.getOrCreateSession("pending", { - root: sampleRoot, - buildOptions: sampleBuildOptions(), - }); - releaseBuild?.(); + limitedManager.disposeSession("pending"); + await expect( + limitedManager.getOrCreateSession("pending", { + root: sampleRoot, + buildOptions: sampleBuildOptions(), + }), + ).rejects.toThrow(/still cancelling initialization/); + await expect( + limitedManager.getOrCreateSession("replacement", { + root: sampleRoot, + buildOptions: sampleBuildOptions(), + }), + ).rejects.toThrow("Session capacity reached (1)"); + releaseBuild?.(); await expect(firstSession).rejects.toThrow(/disposed during initialization/); - await expect(secondSession).resolves.toMatchObject({ - getStatus: expect.any(Function), + + const replacement = await limitedManager.getOrCreateSession("replacement", { + root: sampleRoot, + buildOptions: sampleBuildOptions(), }); - expect((await secondSession).getStatus()).toBe("ready"); - expect(manager.getSession("pending")).toBe(await secondSession); + expect(replacement.getStatus()).toBe("ready"); } finally { buildSpy.mockRestore(); + limitedManager.disposeAll(); } }); - test("should allow immediate recreation after disposeAll cancels a pending session", async () => { + test("retains pending initialization capacity after disposeAll until it settles", async () => { + const limitedManager = new SessionManager({ maxSessions: 1, evictionIntervalMs: 0 }); const originalBuild = indexerBuild.buildProjectIndexIncremental; let releaseBuild: (() => void) | null = null; const buildGate = new Promise((resolve) => { @@ -1449,27 +1461,31 @@ describe("SessionManager", () => { }); try { - const firstSession = manager.getOrCreateSession("pending", { + const firstSession = limitedManager.getOrCreateSession("pending", { root: sampleRoot, buildOptions: sampleBuildOptions(), }); await Promise.resolve(); - manager.disposeAll(); - const secondSession = manager.getOrCreateSession("pending", { - root: sampleRoot, - buildOptions: sampleBuildOptions(), - }); - releaseBuild?.(); + limitedManager.disposeAll(); + await expect( + limitedManager.getOrCreateSession("replacement", { + root: sampleRoot, + buildOptions: sampleBuildOptions(), + }), + ).rejects.toThrow("Session capacity reached (1)"); + releaseBuild?.(); await expect(firstSession).rejects.toThrow(/disposed during initialization/); - await expect(secondSession).resolves.toMatchObject({ - getStatus: expect.any(Function), + + const replacement = await limitedManager.getOrCreateSession("replacement", { + root: sampleRoot, + buildOptions: sampleBuildOptions(), }); - expect((await secondSession).getStatus()).toBe("ready"); - expect(manager.getSession("pending")).toBe(await secondSession); + expect(replacement.getStatus()).toBe("ready"); } finally { buildSpy.mockRestore(); + limitedManager.disposeAll(); } }); diff --git a/tests/sqlite-query-deadline-fallback.test.ts b/tests/sqlite-query-deadline-fallback.test.ts index 1b99dcb5..cf20c255 100644 --- a/tests/sqlite-query-deadline-fallback.test.ts +++ b/tests/sqlite-query-deadline-fallback.test.ts @@ -21,12 +21,18 @@ vi.mock("../src/sqlite/rawQueryWorkerPool.js", async (importOriginal) => { import { SqliteQueryCancelledError, SqliteQueryDeadlineExceededError as PublicSqliteQueryDeadlineExceededError, + SqliteQueryWorkerCleanupCapacityExceededError as PublicSqliteQueryWorkerCleanupCapacityExceededError, } from "../src/sqlite.js"; import { SqliteQueryCancelledError as RootSqliteQueryCancelledError, SqliteQueryDeadlineExceededError as RootSqliteQueryDeadlineExceededError, + SqliteQueryWorkerCleanupCapacityExceededError as RootSqliteQueryWorkerCleanupCapacityExceededError, } from "../src/index.js"; -import { queryGraphSqliteRaw, SqliteQueryDeadlineExceededError } from "../src/sqlite/query.js"; +import { + queryGraphSqliteRaw, + SqliteQueryDeadlineExceededError, + SqliteQueryWorkerCleanupCapacityExceededError, +} from "../src/sqlite/query.js"; async function withTempDb(run: (dbPath: string) => Promise): Promise { const root = await fsp.mkdtemp(path.join(os.tmpdir(), "cg-sqlite-deadline-fallback-")); @@ -119,14 +125,20 @@ describe("SQLite raw query in-process deadline fallback", () => { }); }); - it("exports named cancellation and deadline errors from public library barrels", () => { + it("exports named cancellation, deadline, and capacity errors from public library barrels", () => { expect(PublicSqliteQueryDeadlineExceededError).toBe(SqliteQueryDeadlineExceededError); + expect(PublicSqliteQueryWorkerCleanupCapacityExceededError).toBe(SqliteQueryWorkerCleanupCapacityExceededError); expect(RootSqliteQueryDeadlineExceededError).toBe(SqliteQueryDeadlineExceededError); expect(RootSqliteQueryCancelledError).toBe(SqliteQueryCancelledError); - - const error = new PublicSqliteQueryDeadlineExceededError(250); - expect(error).toBeInstanceOf(Error); - expect(error.name).toBe("SqliteQueryDeadlineExceededError"); - expect(error.message).toBe("SQLite query exceeded its 250ms execution budget and was terminated."); + expect(RootSqliteQueryWorkerCleanupCapacityExceededError).toBe(SqliteQueryWorkerCleanupCapacityExceededError); + + const deadline = new PublicSqliteQueryDeadlineExceededError(250); + expect(deadline).toBeInstanceOf(Error); + expect(deadline.name).toBe("SqliteQueryDeadlineExceededError"); + expect(deadline.message).toBe("SQLite query exceeded its 250ms execution budget and was terminated."); + expect(new SqliteQueryCancelledError().message).toBe("SQLite query was cancelled."); + expect(new PublicSqliteQueryWorkerCleanupCapacityExceededError(2).message).toContain( + "SQLite query worker capacity is exhausted", + ); }); }); From 36e489bb3843b0e1e8837441f2cfe643e21cf138 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Mon, 17 Aug 2026 00:18:49 -0400 Subject: [PATCH 31/37] docs: normalize coverage report formatting --- docs/coverage/js.md | 78 ++++++++++++++++++++++----------------------- 1 file changed, 39 insertions(+), 39 deletions(-) diff --git a/docs/coverage/js.md b/docs/coverage/js.md index 4e9b8c2a..dfb72946 100644 --- a/docs/coverage/js.md +++ b/docs/coverage/js.md @@ -4,52 +4,52 @@ Source: `coverage/js/lcov.info` ## Summary -| Metric | Hit | Found | Coverage | -| --- | ---: | ---: | ---: | -| Lines | 27578 | 30358 | 90.84% | -| Functions | 4602 | 4889 | 94.13% | -| Branches | 20773 | 26172 | 79.37% | +| Metric | Hit | Found | Coverage | +| --------- | ----: | ----: | -------: | +| Lines | 27578 | 30358 | 90.84% | +| Functions | 4602 | 4889 | 94.13% | +| Branches | 20773 | 26172 | 79.37% | ## Least-covered Files -| File | Lines | Functions | Branches | -| --- | ---: | ---: | ---: | -| `src/sqlite/rawQueryWorker.ts` | 0.00% | 0.00% | 0.00% | -| `src/languages/definitions/htmlStub.ts` | 50.00% | 50.00% | n/a | -| `src/languages/definitions/javascript.ts` | 62.50% | 83.33% | 57.89% | -| `src/cli/explore.ts` | 62.50% | 100.00% | 50.00% | -| `src/languages/definitions/typescript.ts` | 62.86% | 81.82% | 42.86% | -| `src/impact/call-compatibility/textScanner.ts` | 66.04% | 100.00% | 70.75% | -| `src/cli/index.ts` | 66.67% | 100.00% | 58.97% | -| `src/indexer.ts` | 66.67% | 66.67% | 100.00% | -| `src/cli/bootstrap.ts` | 66.67% | 50.00% | n/a | -| `src/languages/definitions/css.ts` | 66.67% | 0.00% | n/a | -| `src/languages/definitions/html.ts` | 66.67% | 0.00% | n/a | -| `src/languages/definitions/less.ts` | 66.67% | 0.00% | n/a | -| `src/languages/definitions/sql.ts` | 66.67% | 0.00% | n/a | -| `src/languages/definitions/svelte.ts` | 66.67% | 0.00% | n/a | -| `src/languages/definitions/vue.ts` | 66.67% | 0.00% | n/a | -| `src/agent/handles.ts` | 68.00% | 91.67% | 45.45% | -| `src/cli/artifact.ts` | 70.00% | 100.00% | 72.22% | -| `src/agent/query-index/workerPool.ts` | 70.97% | 62.50% | 56.25% | -| `src/agent/followUps.ts` | 72.15% | 77.78% | 53.91% | -| `src/cli/context.ts` | 72.22% | 60.78% | 75.25% | +| File | Lines | Functions | Branches | +| ---------------------------------------------- | -----: | --------: | -------: | +| `src/sqlite/rawQueryWorker.ts` | 0.00% | 0.00% | 0.00% | +| `src/languages/definitions/htmlStub.ts` | 50.00% | 50.00% | n/a | +| `src/languages/definitions/javascript.ts` | 62.50% | 83.33% | 57.89% | +| `src/cli/explore.ts` | 62.50% | 100.00% | 50.00% | +| `src/languages/definitions/typescript.ts` | 62.86% | 81.82% | 42.86% | +| `src/impact/call-compatibility/textScanner.ts` | 66.04% | 100.00% | 70.75% | +| `src/cli/index.ts` | 66.67% | 100.00% | 58.97% | +| `src/indexer.ts` | 66.67% | 66.67% | 100.00% | +| `src/cli/bootstrap.ts` | 66.67% | 50.00% | n/a | +| `src/languages/definitions/css.ts` | 66.67% | 0.00% | n/a | +| `src/languages/definitions/html.ts` | 66.67% | 0.00% | n/a | +| `src/languages/definitions/less.ts` | 66.67% | 0.00% | n/a | +| `src/languages/definitions/sql.ts` | 66.67% | 0.00% | n/a | +| `src/languages/definitions/svelte.ts` | 66.67% | 0.00% | n/a | +| `src/languages/definitions/vue.ts` | 66.67% | 0.00% | n/a | +| `src/agent/handles.ts` | 68.00% | 91.67% | 45.45% | +| `src/cli/artifact.ts` | 70.00% | 100.00% | 72.22% | +| `src/agent/query-index/workerPool.ts` | 70.97% | 62.50% | 56.25% | +| `src/agent/followUps.ts` | 72.15% | 77.78% | 53.91% | +| `src/cli/context.ts` | 72.22% | 60.78% | 75.25% | ## Type-Only Or Re-Export Files These files have line records but no function or branch records, so they are tracked outside the runtime ranking. -| File | Lines | Functions | Branches | -| --- | ---: | ---: | ---: | -| `src/cliBootstrap.ts` | 0.00% | n/a | n/a | -| `src/languages/definitions/jsFamily.ts` | 100.00% | n/a | n/a | -| `src/duplicate-keywords.ts` | 100.00% | n/a | n/a | -| `src/impact/types.ts` | 100.00% | n/a | n/a | -| `src/languages/definitions/adoc.ts` | 100.00% | n/a | n/a | -| `src/languages/definitions/astro.ts` | 100.00% | n/a | n/a | -| `src/languages/definitions/hbs.ts` | 100.00% | n/a | n/a | -| `src/languages/definitions/markdown.ts` | 100.00% | n/a | n/a | -| `src/languages/definitions/mdx.ts` | 100.00% | n/a | n/a | -| `src/languages/definitions/rst.ts` | 100.00% | n/a | n/a | +| File | Lines | Functions | Branches | +| --------------------------------------- | ------: | --------: | -------: | +| `src/cliBootstrap.ts` | 0.00% | n/a | n/a | +| `src/languages/definitions/jsFamily.ts` | 100.00% | n/a | n/a | +| `src/duplicate-keywords.ts` | 100.00% | n/a | n/a | +| `src/impact/types.ts` | 100.00% | n/a | n/a | +| `src/languages/definitions/adoc.ts` | 100.00% | n/a | n/a | +| `src/languages/definitions/astro.ts` | 100.00% | n/a | n/a | +| `src/languages/definitions/hbs.ts` | 100.00% | n/a | n/a | +| `src/languages/definitions/markdown.ts` | 100.00% | n/a | n/a | +| `src/languages/definitions/mdx.ts` | 100.00% | n/a | n/a | +| `src/languages/definitions/rst.ts` | 100.00% | n/a | n/a | Generated from LCOV by `node ./scripts/coverage-markdown.mjs`. From eb7a1adfc2c38129d29e575a42147d49771a2f81 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Mon, 17 Aug 2026 01:55:49 -0400 Subject: [PATCH 32/37] fix: resolve latest resource review findings --- docs/coverage/js.md | 6 +++--- src/mcp/http.ts | 13 +++++-------- src/session.ts | 1 - tests/mcp-server.test.ts | 5 +---- tests/session.test.ts | 14 ++++++++++++++ tests/sqlite-query-bounds.test.ts | 13 +++++++++++-- 6 files changed, 34 insertions(+), 18 deletions(-) diff --git a/docs/coverage/js.md b/docs/coverage/js.md index dfb72946..6f364705 100644 --- a/docs/coverage/js.md +++ b/docs/coverage/js.md @@ -6,9 +6,9 @@ Source: `coverage/js/lcov.info` | Metric | Hit | Found | Coverage | | --------- | ----: | ----: | -------: | -| Lines | 27578 | 30358 | 90.84% | -| Functions | 4602 | 4889 | 94.13% | -| Branches | 20773 | 26172 | 79.37% | +| Lines | 27574 | 30354 | 90.84% | +| Functions | 4603 | 4889 | 94.15% | +| Branches | 20771 | 26170 | 79.37% | ## Least-covered Files diff --git a/src/mcp/http.ts b/src/mcp/http.ts index f41156b0..31f645ab 100644 --- a/src/mcp/http.ts +++ b/src/mcp/http.ts @@ -24,11 +24,14 @@ export async function readJsonRequestBody( ): Promise { const contentLength = getContentLength(request); const knownTooLarge = contentLength !== undefined && contentLength > maxBytes; + if (knownTooLarge) { + request.resume(); + return { status: "too_large" }; + } return await new Promise((resolve) => { const chunks: Buffer[] = []; let bytes = 0; - let tooLarge = knownTooLarge; let settled = false; const deadline = setTimeout(() => settle({ status: "timeout" }, true), timeoutMs); deadline.unref?.(); @@ -48,21 +51,15 @@ export async function readJsonRequestBody( resolve(result); }; const onData = (chunk: string | Buffer): void => { - if (tooLarge) return; const buffer = typeof chunk === "string" ? Buffer.from(chunk) : chunk; bytes += buffer.byteLength; if (bytes > maxBytes) { - tooLarge = true; - chunks.length = 0; + settle({ status: "too_large" }, true); return; } chunks.push(buffer); }; const onEnd = (): void => { - if (tooLarge) { - settle({ status: "too_large" }, false); - return; - } const rawBody = Buffer.concat(chunks).toString("utf8"); try { const body: unknown = rawBody.length ? JSON.parse(rawBody) : null; diff --git a/src/session.ts b/src/session.ts index 4d85f9fa..48519a71 100644 --- a/src/session.ts +++ b/src/session.ts @@ -995,7 +995,6 @@ export class SessionManager { session.dispose(); } this.sessions.clear(); - clearInterval(this.evictionTimer); } /** diff --git a/tests/mcp-server.test.ts b/tests/mcp-server.test.ts index b751084f..d2e061dd 100644 --- a/tests/mcp-server.test.ts +++ b/tests/mcp-server.test.ts @@ -1130,7 +1130,7 @@ describe("codegraph MCP handlers", () => { } }); - it("rejects oversized HTTP MCP request bodies before parsing", async () => { + it("closes declared oversized HTTP MCP request bodies before reading their payload", async () => { const root = await fs.mkdtemp(path.join(os.tmpdir(), "cg-mcp-http-large-")); await fs.writeFile(path.join(root, "auth.ts"), "export const ok = 1;\n", "utf8"); const httpServer = await startCodegraphMcpHttpServer({ @@ -1148,11 +1148,8 @@ describe("codegraph MCP handlers", () => { }, body: JSON.stringify({ jsonrpc: "2.0", id: 1, method: "initialize", padding: "x".repeat(1_000_000) }), }); - const payload = readObject((await response.json()) as unknown); - const error = readObject(payload.error); expect([response.status, response.headers.get("connection")]).toEqual([413, "close"]); - expect(error.message).toBe("MCP request body is too large"); } finally { await httpServer.close(); } diff --git a/tests/session.test.ts b/tests/session.test.ts index 7116b20f..63a8b657 100644 --- a/tests/session.test.ts +++ b/tests/session.test.ts @@ -1709,6 +1709,20 @@ describe("SessionManager", () => { expect(manager.getSessionIds()).toHaveLength(0); }); + test("keeps periodic expiration cleanup after disposeAll", () => { + vi.useFakeTimers(); + const reusableManager = new SessionManager({ evictionIntervalMs: 10 }); + const cleanupSpy = vi.spyOn(reusableManager, "cleanupExpired"); + + try { + reusableManager.disposeAll(); + vi.advanceTimersByTime(10); + expect(cleanupSpy).toHaveBeenCalledTimes(1); + } finally { + vi.clearAllTimers(); + } + }); + test("should cleanup expired sessions", async () => { await manager.getOrCreateSession("session-1", { root: sampleRoot, diff --git a/tests/sqlite-query-bounds.test.ts b/tests/sqlite-query-bounds.test.ts index 49343e74..b729c4f7 100644 --- a/tests/sqlite-query-bounds.test.ts +++ b/tests/sqlite-query-bounds.test.ts @@ -1,7 +1,7 @@ import fsp from "node:fs/promises"; import os from "node:os"; import path from "node:path"; -import { describe, expect, it } from "vitest"; +import { describe, expect, it, vi } from "vitest"; import { DatabaseSync } from "node:sqlite"; import { @@ -54,7 +54,7 @@ async function removeWithRetry(root: string): Promise { return; } catch (error) { if (!(error instanceof Error) || !("code" in error)) throw error; - if (error.code !== "EBUSY" && error.code !== "ENOTEMPTY") throw error; + if (error.code !== "EBUSY" && error.code !== "ENOTEMPTY" && error.code !== "EPERM") throw error; if (Date.now() > deadline) throw error; await new Promise((resolve) => setTimeout(resolve, 100)); } @@ -62,6 +62,15 @@ async function removeWithRetry(root: string): Promise { } describe("SQLite query byte/cell bounds during iterate", () => { + it("retries Windows-style EPERM cleanup races", async () => { + const removeSpy = vi.spyOn(fsp, "rm").mockRejectedValueOnce(Object.assign(new Error("locked"), { code: "EPERM" })); + try { + await withTempDb(async () => {}); + expect(removeSpy).toHaveBeenCalledTimes(2); + } finally { + removeSpy.mockRestore(); + } + }); it("applies per-cell and cumulative caps before appending huge existing TEXT cells", async () => { await withTempDb(async (dbPath) => { const db = new DatabaseSync(dbPath); From 75061e931082f5d6f2c2e13365dadafbe928c7a6 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Mon, 17 Aug 2026 14:43:05 -0400 Subject: [PATCH 33/37] fix: close MCP resources on shutdown --- docs/agent-workflows.md | 8 +++- docs/coverage/js.md | 6 +-- src/mcp/http.ts | 37 +++++++++++++++--- src/mcp/server.ts | 10 +++-- src/session.ts | 18 +++++++++ src/sqlite/rawQueryWorkerPool.ts | 2 +- tests/mcp-server.test.ts | 41 +++++++++++++++++++- tests/session.test.ts | 25 ++++++++++++ tests/sqlite-query-bounds.test.ts | 2 +- tests/sqlite-query-deadline-fallback.test.ts | 2 +- 10 files changed, 134 insertions(+), 17 deletions(-) diff --git a/docs/agent-workflows.md b/docs/agent-workflows.md index cc269667..a345969d 100644 --- a/docs/agent-workflows.md +++ b/docs/agent-workflows.md @@ -306,9 +306,15 @@ const sameSession = await manager.getOrCreateSession("pr-123", { manager.cleanupExpired(); const allStats = manager.getAllStats(); console.log(Boolean(pr1Session), Boolean(pr2Session), Boolean(sameSession), allStats); + +manager.dispose(); ``` -`SessionManager` defaults to 32 live or initializing sessions and scans for expired sessions every 60 seconds. Set `{ maxSessions, evictionIntervalMs }` to tune those bounds; `maxSessions` also applies to net-new `warmup()` sessions. Capacity frees immediately after a ready session is disposed or expires, or after a canceled or failed initialization settles. Set `evictionIntervalMs: 0` to disable periodic cleanup. +`SessionManager` defaults to 32 live or initializing sessions and scans for expired sessions every 60 seconds. Set `{ maxSessions, evictionIntervalMs }` to tune those bounds; `maxSessions` also applies to net-new `warmup()` sessions. + +- Capacity frees immediately after a ready session is disposed or expires, or after a canceled or failed initialization settles. +- Set `evictionIntervalMs: 0` to disable periodic cleanup. +- Call `manager.dispose()` when the manager is no longer needed; it disposes all sessions, stops the interval, and is terminal. `disposeAll()` remains reusable. ## Streaming impact analysis diff --git a/docs/coverage/js.md b/docs/coverage/js.md index 6f364705..92058a69 100644 --- a/docs/coverage/js.md +++ b/docs/coverage/js.md @@ -6,9 +6,9 @@ Source: `coverage/js/lcov.info` | Metric | Hit | Found | Coverage | | --------- | ----: | ----: | -------: | -| Lines | 27574 | 30354 | 90.84% | -| Functions | 4603 | 4889 | 94.15% | -| Branches | 20771 | 26170 | 79.37% | +| Lines | 27603 | 30382 | 90.85% | +| Functions | 4609 | 4895 | 94.16% | +| Branches | 20775 | 26176 | 79.37% | ## Least-covered Files diff --git a/src/mcp/http.ts b/src/mcp/http.ts index 31f645ab..e0f09d40 100644 --- a/src/mcp/http.ts +++ b/src/mcp/http.ts @@ -4,7 +4,7 @@ import os from "node:os"; export type ParsedJsonBody = | { status: "ok"; body: unknown } - | { status: "too_large" } + | { status: "too_large"; drained: Promise } | { status: "timeout" } | { status: "invalid_json" }; @@ -23,10 +23,8 @@ export async function readJsonRequestBody( timeoutMs: number, ): Promise { const contentLength = getContentLength(request); - const knownTooLarge = contentLength !== undefined && contentLength > maxBytes; - if (knownTooLarge) { - request.resume(); - return { status: "too_large" }; + if (contentLength !== undefined && contentLength > maxBytes) { + return { status: "too_large", drained: drainRequestBody(request, timeoutMs) }; } return await new Promise((resolve) => { @@ -54,7 +52,7 @@ export async function readJsonRequestBody( const buffer = typeof chunk === "string" ? Buffer.from(chunk) : chunk; bytes += buffer.byteLength; if (bytes > maxBytes) { - settle({ status: "too_large" }, true); + settle({ status: "too_large", drained: drainRequestBody(request, timeoutMs) }, false); return; } chunks.push(buffer); @@ -77,6 +75,33 @@ export async function readJsonRequestBody( }); } +function drainRequestBody(request: IncomingMessage, timeoutMs: number): Promise { + const { promise, resolve } = Promise.withResolvers(); + let settled = false; + const onDrained = (): void => settle(); + const deadline = setTimeout(onDrained, timeoutMs); + deadline.unref?.(); + + const cleanup = (): void => { + clearTimeout(deadline); + request.off("end", onDrained); + request.off("error", onDrained); + request.off("aborted", onDrained); + }; + const settle = (): void => { + if (settled) return; + settled = true; + cleanup(); + resolve(); + }; + + request.once("end", onDrained); + request.once("error", onDrained); + request.once("aborted", onDrained); + request.resume(); + return promise; +} + export function emptyAllowedHostHeaderRules(): AllowedHostHeaderRules { return { exact: new Set(), loopbackOnly: new Set() }; } diff --git a/src/mcp/server.ts b/src/mcp/server.ts index b4194c63..3daa9b19 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -1260,9 +1260,12 @@ export async function startCodegraphMcpHttpServer( let closeResourcesPromise: Promise | undefined; const closeResources = (): Promise => { closeResourcesPromise ??= (async () => { - session.invalidate(); - sessionStore.stop(); - await closeMcpResources(sessionStore.sessions, modernHandler.close); + try { + session.invalidate(); + } finally { + sessionStore.stop(); + await closeMcpResources(sessionStore.sessions, modernHandler.close); + } })(); return closeResourcesPromise; }; @@ -1338,6 +1341,7 @@ async function handleMcpHttpRequest( if (request.method === "POST") { const parsedBody = await readJsonRequestBody(request, MAX_MCP_HTTP_BODY_BYTES, bodyTimeoutMs); if (parsedBody.status === "too_large") { + await parsedBody.drained; writeClosingJsonRpcError(413, "MCP request body is too large"); return; } diff --git a/src/session.ts b/src/session.ts index 48519a71..e2c0677c 100644 --- a/src/session.ts +++ b/src/session.ts @@ -823,6 +823,7 @@ export class SessionManager { private pendingSessions = new Map(); private readonly maxSessions: number; private readonly evictionTimer: ReturnType | undefined; + private disposed = false; constructor(options: SessionManagerOptions = {}) { this.maxSessions = normalizeSessionManagerCapacity(options.maxSessions); @@ -934,10 +935,15 @@ export class SessionManager { return promise; } + private assertNotDisposed(): void { + if (this.disposed) throw new Error("Session manager is disposed."); + } + /** * Create or get a session for a repository */ async getOrCreateSession(sessionId: string, options: SessionOptions): Promise { + this.assertNotDisposed(); const pending = this.getPendingCompatibleSession(sessionId, options); if (pending) { return await pending; @@ -997,6 +1003,17 @@ export class SessionManager { this.sessions.clear(); } + /** + * Dispose all sessions and stop periodic expiration cleanup. + * This manager cannot be reused afterward. + */ + dispose(): void { + if (this.disposed) return; + this.disposed = true; + clearInterval(this.evictionTimer); + this.disposeAll(); + } + /** * Get all session IDs */ @@ -1036,6 +1053,7 @@ export class SessionManager { * @param sessions - Array of session configs to pre-warm */ async warmup(sessions: Array<{ id: string; options: SessionOptions }>): Promise { + this.assertNotDisposed(); const requestedFingerprints = new Map(); const replacementSessions: Array<{ id: string; diff --git a/src/sqlite/rawQueryWorkerPool.ts b/src/sqlite/rawQueryWorkerPool.ts index 3294bc73..ada2692b 100644 --- a/src/sqlite/rawQueryWorkerPool.ts +++ b/src/sqlite/rawQueryWorkerPool.ts @@ -22,7 +22,7 @@ export type RawSqlQueryWorkerLifecycleState = { export class SqliteQueryDeadlineExceededError extends Error { constructor(deadlineMs: number) { - super(`SQLite query exceeded its ${deadlineMs}ms execution budget and was terminated.`); + super(`SQLite query exceeded its ${deadlineMs}ms execution budget; termination was requested.`); this.name = "SqliteQueryDeadlineExceededError"; } } diff --git a/tests/mcp-server.test.ts b/tests/mcp-server.test.ts index d2e061dd..384ba20f 100644 --- a/tests/mcp-server.test.ts +++ b/tests/mcp-server.test.ts @@ -2,6 +2,7 @@ import { registerSessionInvalidationHook } from "../src/agent/sessionLifecycle.j import { ensureSessionQueryIndex } from "../src/agent/query-index/sessionStore.js"; import fs from "node:fs/promises"; import { request as httpRequest, type IncomingMessage } from "node:http"; +import { NodeStreamableHTTPServerTransport } from "@modelcontextprotocol/node"; import os from "node:os"; import path from "node:path"; import { DatabaseSync } from "node:sqlite"; @@ -1130,7 +1131,7 @@ describe("codegraph MCP handlers", () => { } }); - it("closes declared oversized HTTP MCP request bodies before reading their payload", async () => { + it("rejects declared oversized HTTP MCP request bodies without buffering their payload", async () => { const root = await fs.mkdtemp(path.join(os.tmpdir(), "cg-mcp-http-large-")); await fs.writeFile(path.join(root, "auth.ts"), "export const ok = 1;\n", "utf8"); const httpServer = await startCodegraphMcpHttpServer({ @@ -3062,6 +3063,44 @@ describe("MCP session teardown regressions (S2)", () => { await httpServer.close(); } }); + + it("closes legacy protocol transports when session invalidation fails during server shutdown", async () => { + const root = await fs.mkdtemp(path.join(os.tmpdir(), "cg-mcp-invalidation-close-")); + await fs.writeFile(path.join(root, "auth.ts"), "export const ok = 1;\n", "utf8"); + const session = createAgentSession({ root }); + const httpServer = await startCodegraphMcpHttpServer({ + root, + host: "127.0.0.1", + port: 0, + session, + }); + const transportCloseSpy = vi.spyOn(NodeStreamableHTTPServerTransport.prototype, "close"); + const invalidationSpy = vi.spyOn(session, "invalidate").mockImplementation(() => { + throw new Error("session invalidation failed"); + }); + + try { + const initialize = await postMcpJson(httpServer.url, { + jsonrpc: "2.0", + id: 1, + method: "initialize", + params: { + protocolVersion: "2025-11-25", + capabilities: {}, + clientInfo: { name: "codegraph-invalidation-close-test", version: "1.0.0" }, + }, + }); + expect(initialize.response.status).toBe(200); + transportCloseSpy.mockClear(); + + await expect(httpServer.close()).rejects.toThrow("session invalidation failed"); + expect(transportCloseSpy).toHaveBeenCalled(); + } finally { + invalidationSpy.mockRestore(); + transportCloseSpy.mockRestore(); + await httpServer.close().catch(() => {}); + } + }); }); describe("MCP cancellation accounting", () => { diff --git a/tests/session.test.ts b/tests/session.test.ts index 63a8b657..f566543d 100644 --- a/tests/session.test.ts +++ b/tests/session.test.ts @@ -1194,6 +1194,10 @@ describe("SessionManager", () => { manager = new SessionManager(); }); + afterEach(() => { + manager.dispose(); + }); + test("should create and retrieve sessions", async () => { const session = await manager.getOrCreateSession("test-session", { root: sampleRoot, @@ -1723,6 +1727,27 @@ describe("SessionManager", () => { } }); + test("stops periodic cleanup and prevents reuse after terminal disposal", async () => { + vi.useFakeTimers(); + const disposableManager = new SessionManager({ evictionIntervalMs: 10 }); + const cleanupSpy = vi.spyOn(disposableManager, "cleanupExpired"); + + try { + disposableManager.dispose(); + vi.advanceTimersByTime(10); + expect(cleanupSpy).not.toHaveBeenCalled(); + await expect( + disposableManager.getOrCreateSession("replacement", { + root: sampleRoot, + buildOptions: sampleBuildOptions(), + }), + ).rejects.toThrow("Session manager is disposed."); + await expect(disposableManager.warmup([])).rejects.toThrow("Session manager is disposed."); + } finally { + vi.useRealTimers(); + } + }); + test("should cleanup expired sessions", async () => { await manager.getOrCreateSession("session-1", { root: sampleRoot, diff --git a/tests/sqlite-query-bounds.test.ts b/tests/sqlite-query-bounds.test.ts index b729c4f7..b7b06349 100644 --- a/tests/sqlite-query-bounds.test.ts +++ b/tests/sqlite-query-bounds.test.ts @@ -199,6 +199,6 @@ describe("SQLite raw query execution deadline", () => { const error = new SqliteQueryDeadlineExceededError(250); expect(error).toBeInstanceOf(Error); expect(error.name).toBe("SqliteQueryDeadlineExceededError"); - expect(error.message).toBe("SQLite query exceeded its 250ms execution budget and was terminated."); + expect(error.message).toBe("SQLite query exceeded its 250ms execution budget; termination was requested."); }); }); diff --git a/tests/sqlite-query-deadline-fallback.test.ts b/tests/sqlite-query-deadline-fallback.test.ts index cf20c255..98c65c8a 100644 --- a/tests/sqlite-query-deadline-fallback.test.ts +++ b/tests/sqlite-query-deadline-fallback.test.ts @@ -135,7 +135,7 @@ describe("SQLite raw query in-process deadline fallback", () => { const deadline = new PublicSqliteQueryDeadlineExceededError(250); expect(deadline).toBeInstanceOf(Error); expect(deadline.name).toBe("SqliteQueryDeadlineExceededError"); - expect(deadline.message).toBe("SQLite query exceeded its 250ms execution budget and was terminated."); + expect(deadline.message).toBe("SQLite query exceeded its 250ms execution budget; termination was requested."); expect(new SqliteQueryCancelledError().message).toBe("SQLite query was cancelled."); expect(new PublicSqliteQueryWorkerCleanupCapacityExceededError(2).message).toContain( "SQLite query worker capacity is exhausted", From 8a9563a5bb40236132aff806dba5b5e42b0f7f3e Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Mon, 17 Aug 2026 14:56:17 -0400 Subject: [PATCH 34/37] fix: drain rejected MCP bodies before close --- docs/coverage/js.md | 6 +++--- src/mcp/http.ts | 27 ++++++++++++--------------- src/mcp/server.ts | 3 +-- tests/mcp-server.test.ts | 2 +- 4 files changed, 17 insertions(+), 21 deletions(-) diff --git a/docs/coverage/js.md b/docs/coverage/js.md index 92058a69..e9f54e4a 100644 --- a/docs/coverage/js.md +++ b/docs/coverage/js.md @@ -6,9 +6,9 @@ Source: `coverage/js/lcov.info` | Metric | Hit | Found | Coverage | | --------- | ----: | ----: | -------: | -| Lines | 27603 | 30382 | 90.85% | -| Functions | 4609 | 4895 | 94.16% | -| Branches | 20775 | 26176 | 79.37% | +| Lines | 27596 | 30378 | 90.84% | +| Functions | 4608 | 4895 | 94.14% | +| Branches | 20774 | 26174 | 79.37% | ## Least-covered Files diff --git a/src/mcp/http.ts b/src/mcp/http.ts index e0f09d40..e00ca628 100644 --- a/src/mcp/http.ts +++ b/src/mcp/http.ts @@ -4,7 +4,7 @@ import os from "node:os"; export type ParsedJsonBody = | { status: "ok"; body: unknown } - | { status: "too_large"; drained: Promise } + | { status: "too_large" } | { status: "timeout" } | { status: "invalid_json" }; @@ -24,7 +24,8 @@ export async function readJsonRequestBody( ): Promise { const contentLength = getContentLength(request); if (contentLength !== undefined && contentLength > maxBytes) { - return { status: "too_large", drained: drainRequestBody(request, timeoutMs) }; + drainRequestBody(request, timeoutMs); + return { status: "too_large" }; } return await new Promise((resolve) => { @@ -52,7 +53,8 @@ export async function readJsonRequestBody( const buffer = typeof chunk === "string" ? Buffer.from(chunk) : chunk; bytes += buffer.byteLength; if (bytes > maxBytes) { - settle({ status: "too_large", drained: drainRequestBody(request, timeoutMs) }, false); + drainRequestBody(request, timeoutMs); + settle({ status: "too_large" }, false); return; } chunks.push(buffer); @@ -75,11 +77,13 @@ export async function readJsonRequestBody( }); } -function drainRequestBody(request: IncomingMessage, timeoutMs: number): Promise { - const { promise, resolve } = Promise.withResolvers(); - let settled = false; - const onDrained = (): void => settle(); - const deadline = setTimeout(onDrained, timeoutMs); +function drainRequestBody(request: IncomingMessage, timeoutMs: number): void { + const onDrained = (): void => cleanup(); + const onTimedOut = (): void => { + cleanup(); + request.destroy(); + }; + const deadline = setTimeout(onTimedOut, timeoutMs); deadline.unref?.(); const cleanup = (): void => { @@ -88,18 +92,11 @@ function drainRequestBody(request: IncomingMessage, timeoutMs: number): Promise< request.off("error", onDrained); request.off("aborted", onDrained); }; - const settle = (): void => { - if (settled) return; - settled = true; - cleanup(); - resolve(); - }; request.once("end", onDrained); request.once("error", onDrained); request.once("aborted", onDrained); request.resume(); - return promise; } export function emptyAllowedHostHeaderRules(): AllowedHostHeaderRules { diff --git a/src/mcp/server.ts b/src/mcp/server.ts index 3daa9b19..bba7c261 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -1341,8 +1341,7 @@ async function handleMcpHttpRequest( if (request.method === "POST") { const parsedBody = await readJsonRequestBody(request, MAX_MCP_HTTP_BODY_BYTES, bodyTimeoutMs); if (parsedBody.status === "too_large") { - await parsedBody.drained; - writeClosingJsonRpcError(413, "MCP request body is too large"); + writeJsonRpcError(response, 413, "MCP request body is too large"); return; } if (parsedBody.status === "timeout") { diff --git a/tests/mcp-server.test.ts b/tests/mcp-server.test.ts index 384ba20f..524966de 100644 --- a/tests/mcp-server.test.ts +++ b/tests/mcp-server.test.ts @@ -1150,7 +1150,7 @@ describe("codegraph MCP handlers", () => { body: JSON.stringify({ jsonrpc: "2.0", id: 1, method: "initialize", padding: "x".repeat(1_000_000) }), }); - expect([response.status, response.headers.get("connection")]).toEqual([413, "close"]); + expect(response.status).toBe(413); } finally { await httpServer.close(); } From dade911af6af896b5a5df200e0716d1ad1effb8e Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Mon, 17 Aug 2026 16:36:19 -0400 Subject: [PATCH 35/37] fix: address latest resource review findings --- docs/agent-workflows.md | 5 +- docs/coverage/js.md | 78 ++++++++++++------------ src/index.ts | 1 + src/mcp/server.ts | 7 ++- src/sqlite/rawQueryWorkerPool.ts | 4 +- tests/mcp-server.test.ts | 30 +++++++++ tests/raw-query-worker-lifecycle.test.ts | 19 ++++++ tests/session.test.ts | 10 ++- 8 files changed, 109 insertions(+), 45 deletions(-) diff --git a/docs/agent-workflows.md b/docs/agent-workflows.md index a345969d..f026ca1b 100644 --- a/docs/agent-workflows.md +++ b/docs/agent-workflows.md @@ -290,9 +290,10 @@ Available presets: ### Managing multiple sessions ```ts -import { SessionManager } from "@lzehrung/codegraph"; +import { SessionManager, type SessionManagerOptions } from "@lzehrung/codegraph"; -const manager = new SessionManager(); +const options: SessionManagerOptions = { maxSessions: 32 }; +const manager = new SessionManager(options); const pr1Session = await manager.getOrCreateSession("pr-123", { root: "/path/to/repo", }); diff --git a/docs/coverage/js.md b/docs/coverage/js.md index e9f54e4a..df554521 100644 --- a/docs/coverage/js.md +++ b/docs/coverage/js.md @@ -4,52 +4,52 @@ Source: `coverage/js/lcov.info` ## Summary -| Metric | Hit | Found | Coverage | -| --------- | ----: | ----: | -------: | -| Lines | 27596 | 30378 | 90.84% | -| Functions | 4608 | 4895 | 94.14% | -| Branches | 20774 | 26174 | 79.37% | +| Metric | Hit | Found | Coverage | +| --- | ---: | ---: | ---: | +| Lines | 27601 | 30381 | 90.85% | +| Functions | 4609 | 4895 | 94.16% | +| Branches | 20774 | 26174 | 79.37% | ## Least-covered Files -| File | Lines | Functions | Branches | -| ---------------------------------------------- | -----: | --------: | -------: | -| `src/sqlite/rawQueryWorker.ts` | 0.00% | 0.00% | 0.00% | -| `src/languages/definitions/htmlStub.ts` | 50.00% | 50.00% | n/a | -| `src/languages/definitions/javascript.ts` | 62.50% | 83.33% | 57.89% | -| `src/cli/explore.ts` | 62.50% | 100.00% | 50.00% | -| `src/languages/definitions/typescript.ts` | 62.86% | 81.82% | 42.86% | -| `src/impact/call-compatibility/textScanner.ts` | 66.04% | 100.00% | 70.75% | -| `src/cli/index.ts` | 66.67% | 100.00% | 58.97% | -| `src/indexer.ts` | 66.67% | 66.67% | 100.00% | -| `src/cli/bootstrap.ts` | 66.67% | 50.00% | n/a | -| `src/languages/definitions/css.ts` | 66.67% | 0.00% | n/a | -| `src/languages/definitions/html.ts` | 66.67% | 0.00% | n/a | -| `src/languages/definitions/less.ts` | 66.67% | 0.00% | n/a | -| `src/languages/definitions/sql.ts` | 66.67% | 0.00% | n/a | -| `src/languages/definitions/svelte.ts` | 66.67% | 0.00% | n/a | -| `src/languages/definitions/vue.ts` | 66.67% | 0.00% | n/a | -| `src/agent/handles.ts` | 68.00% | 91.67% | 45.45% | -| `src/cli/artifact.ts` | 70.00% | 100.00% | 72.22% | -| `src/agent/query-index/workerPool.ts` | 70.97% | 62.50% | 56.25% | -| `src/agent/followUps.ts` | 72.15% | 77.78% | 53.91% | -| `src/cli/context.ts` | 72.22% | 60.78% | 75.25% | +| File | Lines | Functions | Branches | +| --- | ---: | ---: | ---: | +| `src/sqlite/rawQueryWorker.ts` | 0.00% | 0.00% | 0.00% | +| `src/languages/definitions/htmlStub.ts` | 50.00% | 50.00% | n/a | +| `src/languages/definitions/javascript.ts` | 62.50% | 83.33% | 57.89% | +| `src/cli/explore.ts` | 62.50% | 100.00% | 50.00% | +| `src/languages/definitions/typescript.ts` | 62.86% | 81.82% | 42.86% | +| `src/impact/call-compatibility/textScanner.ts` | 66.04% | 100.00% | 70.75% | +| `src/cli/index.ts` | 66.67% | 100.00% | 58.97% | +| `src/indexer.ts` | 66.67% | 66.67% | 100.00% | +| `src/cli/bootstrap.ts` | 66.67% | 50.00% | n/a | +| `src/languages/definitions/css.ts` | 66.67% | 0.00% | n/a | +| `src/languages/definitions/html.ts` | 66.67% | 0.00% | n/a | +| `src/languages/definitions/less.ts` | 66.67% | 0.00% | n/a | +| `src/languages/definitions/sql.ts` | 66.67% | 0.00% | n/a | +| `src/languages/definitions/svelte.ts` | 66.67% | 0.00% | n/a | +| `src/languages/definitions/vue.ts` | 66.67% | 0.00% | n/a | +| `src/agent/handles.ts` | 68.00% | 91.67% | 45.45% | +| `src/cli/artifact.ts` | 70.00% | 100.00% | 72.22% | +| `src/agent/query-index/workerPool.ts` | 70.97% | 62.50% | 56.25% | +| `src/agent/followUps.ts` | 72.15% | 77.78% | 53.91% | +| `src/cli/context.ts` | 72.22% | 60.78% | 75.25% | ## Type-Only Or Re-Export Files These files have line records but no function or branch records, so they are tracked outside the runtime ranking. -| File | Lines | Functions | Branches | -| --------------------------------------- | ------: | --------: | -------: | -| `src/cliBootstrap.ts` | 0.00% | n/a | n/a | -| `src/languages/definitions/jsFamily.ts` | 100.00% | n/a | n/a | -| `src/duplicate-keywords.ts` | 100.00% | n/a | n/a | -| `src/impact/types.ts` | 100.00% | n/a | n/a | -| `src/languages/definitions/adoc.ts` | 100.00% | n/a | n/a | -| `src/languages/definitions/astro.ts` | 100.00% | n/a | n/a | -| `src/languages/definitions/hbs.ts` | 100.00% | n/a | n/a | -| `src/languages/definitions/markdown.ts` | 100.00% | n/a | n/a | -| `src/languages/definitions/mdx.ts` | 100.00% | n/a | n/a | -| `src/languages/definitions/rst.ts` | 100.00% | n/a | n/a | +| File | Lines | Functions | Branches | +| --- | ---: | ---: | ---: | +| `src/cliBootstrap.ts` | 0.00% | n/a | n/a | +| `src/languages/definitions/jsFamily.ts` | 100.00% | n/a | n/a | +| `src/duplicate-keywords.ts` | 100.00% | n/a | n/a | +| `src/impact/types.ts` | 100.00% | n/a | n/a | +| `src/languages/definitions/adoc.ts` | 100.00% | n/a | n/a | +| `src/languages/definitions/astro.ts` | 100.00% | n/a | n/a | +| `src/languages/definitions/hbs.ts` | 100.00% | n/a | n/a | +| `src/languages/definitions/markdown.ts` | 100.00% | n/a | n/a | +| `src/languages/definitions/mdx.ts` | 100.00% | n/a | n/a | +| `src/languages/definitions/rst.ts` | 100.00% | n/a | n/a | Generated from LCOV by `node ./scripts/coverage-markdown.mjs`. diff --git a/src/index.ts b/src/index.ts index 45db55bb..b41b8569 100644 --- a/src/index.ts +++ b/src/index.ts @@ -205,6 +205,7 @@ export { createCodeReviewSession, type ICodeReviewSession, type SessionOptions, + type SessionManagerOptions, type SessionStatus, type SessionStats, type SessionStaleReason, diff --git a/src/mcp/server.ts b/src/mcp/server.ts index bba7c261..2257c238 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -1287,7 +1287,12 @@ export async function startCodegraphMcpHttpServer( void closeResources(); }); - await listenOnHttpServer(server, options.port, host); + try { + await listenOnHttpServer(server, options.port, host); + } catch (error) { + await closeResources(); + throw error; + } const address = server.address(); const actualPort = getHttpServerPort(address); const urlHost = formatHostForUrl(host); diff --git a/src/sqlite/rawQueryWorkerPool.ts b/src/sqlite/rawQueryWorkerPool.ts index ada2692b..f6a6f05b 100644 --- a/src/sqlite/rawQueryWorkerPool.ts +++ b/src/sqlite/rawQueryWorkerPool.ts @@ -64,6 +64,8 @@ export class RawSqlQueryWorkerLifecycle { createPool: RawSqlQueryWorkerPoolFactory, ): Promise { if (signal?.aborted) throw new SqliteQueryCancelledError(); + const deadlineSignal = AbortSignal.timeout(deadlineMs); + const combinedSignal = signal ? AbortSignal.any([signal, deadlineSignal]) : deadlineSignal; if (this.activeWorkerSlots.size >= this.maxWorkers) { throw new SqliteQueryWorkerCleanupCapacityExceededError(this.maxWorkers); } @@ -72,8 +74,6 @@ export class RawSqlQueryWorkerLifecycle { this.activeWorkerSlots.add(slot); let pool: RawSqlQueryWorkerPool | undefined; let cleanupInBackground = false; - const deadlineSignal = AbortSignal.timeout(deadlineMs); - const combinedSignal = signal ? AbortSignal.any([signal, deadlineSignal]) : deadlineSignal; try { pool = createPool(); diff --git a/tests/mcp-server.test.ts b/tests/mcp-server.test.ts index 524966de..d2b750a9 100644 --- a/tests/mcp-server.test.ts +++ b/tests/mcp-server.test.ts @@ -1220,6 +1220,36 @@ describe("codegraph MCP handlers", () => { } }); + it("invalidates prebuilt resources when HTTP server binding fails", async () => { + const root = await fs.mkdtemp(path.join(os.tmpdir(), "cg-mcp-listen-failure-")); + await fs.writeFile(path.join(root, "auth.ts"), "export const ok = 1;\n", "utf8"); + const occupiedServer = await startCodegraphMcpHttpServer({ + root, + host: "127.0.0.1", + port: 0, + }); + const session = createAgentSession({ root }); + let invalidated = false; + registerSessionInvalidationHook(session, () => { + invalidated = true; + }); + + try { + await expect( + startCodegraphMcpHttpServer({ + root, + host: "127.0.0.1", + port: occupiedServer.port, + session, + }), + ).rejects.toThrow(); + expect(invalidated).toBe(true); + } finally { + await occupiedServer.close(); + await fs.rm(root, { force: true, recursive: true }); + } + }); + it("reuses one session across search, get_symbol, refs, and query_sqlite handlers", async () => { const root = await fs.mkdtemp(path.join(os.tmpdir(), "cg-mcp-")); await fs.writeFile(path.join(root, "auth.ts"), "export function validateUser(id: number) { return id > 0; }\n"); diff --git a/tests/raw-query-worker-lifecycle.test.ts b/tests/raw-query-worker-lifecycle.test.ts index fed1e2c0..88d1f279 100644 --- a/tests/raw-query-worker-lifecycle.test.ts +++ b/tests/raw-query-worker-lifecycle.test.ts @@ -76,4 +76,23 @@ describe("RawSqlQueryWorkerLifecycle", () => { complete.resolve({ columns: [], rows: [] }); await expect(active).resolves.toEqual({ columns: [], rows: [] }); }); + + it("releases no worker slot when an invalid deadline is rejected before startup", async () => { + const lifecycle = new RawSqlQueryWorkerLifecycle(1); + + await expect( + lifecycle.run(task, -1, undefined, () => ({ + run: async () => ({ columns: [], rows: [] }), + destroy: async () => {}, + })), + ).rejects.toThrow(); + expect(lifecycle.state()).toEqual({ activeWorkers: 0, maxWorkers: 1 }); + + await expect( + lifecycle.run(task, 10_000, undefined, () => ({ + run: async () => ({ columns: [], rows: [] }), + destroy: async () => {}, + })), + ).resolves.toEqual({ columns: [], rows: [] }); + }); }); diff --git a/tests/session.test.ts b/tests/session.test.ts index f566543d..59865199 100644 --- a/tests/session.test.ts +++ b/tests/session.test.ts @@ -1,5 +1,5 @@ import { describe, test, expect, beforeAll, afterAll, afterEach, beforeEach, vi } from "vitest"; -import type { ICodeReviewSession } from "../src/index.js"; +import type { ICodeReviewSession, SessionManagerOptions as PublicSessionManagerOptions } from "../src/index.js"; import type { BuildOptions, BuildReport, LanguageExtensionMap } from "../src/indexer/types.js"; import { CodeReviewSession, @@ -1198,6 +1198,14 @@ describe("SessionManager", () => { manager.dispose(); }); + test("exports SessionManagerOptions from the package root", () => { + const options: PublicSessionManagerOptions = { evictionIntervalMs: 0, maxSessions: 1 }; + const typedManager = new SessionManager(options); + + expect(typedManager).toBeInstanceOf(SessionManager); + typedManager.dispose(); + }); + test("should create and retrieve sessions", async () => { const session = await manager.getOrCreateSession("test-session", { root: sampleRoot, From ad5ef5a5616cee01e16990ba3ba06f318cda9519 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Mon, 17 Aug 2026 16:39:08 -0400 Subject: [PATCH 36/37] docs: normalize coverage report formatting --- docs/coverage/js.md | 78 ++++++++++++++++++++++----------------------- 1 file changed, 39 insertions(+), 39 deletions(-) diff --git a/docs/coverage/js.md b/docs/coverage/js.md index df554521..e6102cf0 100644 --- a/docs/coverage/js.md +++ b/docs/coverage/js.md @@ -4,52 +4,52 @@ Source: `coverage/js/lcov.info` ## Summary -| Metric | Hit | Found | Coverage | -| --- | ---: | ---: | ---: | -| Lines | 27601 | 30381 | 90.85% | -| Functions | 4609 | 4895 | 94.16% | -| Branches | 20774 | 26174 | 79.37% | +| Metric | Hit | Found | Coverage | +| --------- | ----: | ----: | -------: | +| Lines | 27601 | 30381 | 90.85% | +| Functions | 4609 | 4895 | 94.16% | +| Branches | 20774 | 26174 | 79.37% | ## Least-covered Files -| File | Lines | Functions | Branches | -| --- | ---: | ---: | ---: | -| `src/sqlite/rawQueryWorker.ts` | 0.00% | 0.00% | 0.00% | -| `src/languages/definitions/htmlStub.ts` | 50.00% | 50.00% | n/a | -| `src/languages/definitions/javascript.ts` | 62.50% | 83.33% | 57.89% | -| `src/cli/explore.ts` | 62.50% | 100.00% | 50.00% | -| `src/languages/definitions/typescript.ts` | 62.86% | 81.82% | 42.86% | -| `src/impact/call-compatibility/textScanner.ts` | 66.04% | 100.00% | 70.75% | -| `src/cli/index.ts` | 66.67% | 100.00% | 58.97% | -| `src/indexer.ts` | 66.67% | 66.67% | 100.00% | -| `src/cli/bootstrap.ts` | 66.67% | 50.00% | n/a | -| `src/languages/definitions/css.ts` | 66.67% | 0.00% | n/a | -| `src/languages/definitions/html.ts` | 66.67% | 0.00% | n/a | -| `src/languages/definitions/less.ts` | 66.67% | 0.00% | n/a | -| `src/languages/definitions/sql.ts` | 66.67% | 0.00% | n/a | -| `src/languages/definitions/svelte.ts` | 66.67% | 0.00% | n/a | -| `src/languages/definitions/vue.ts` | 66.67% | 0.00% | n/a | -| `src/agent/handles.ts` | 68.00% | 91.67% | 45.45% | -| `src/cli/artifact.ts` | 70.00% | 100.00% | 72.22% | -| `src/agent/query-index/workerPool.ts` | 70.97% | 62.50% | 56.25% | -| `src/agent/followUps.ts` | 72.15% | 77.78% | 53.91% | -| `src/cli/context.ts` | 72.22% | 60.78% | 75.25% | +| File | Lines | Functions | Branches | +| ---------------------------------------------- | -----: | --------: | -------: | +| `src/sqlite/rawQueryWorker.ts` | 0.00% | 0.00% | 0.00% | +| `src/languages/definitions/htmlStub.ts` | 50.00% | 50.00% | n/a | +| `src/languages/definitions/javascript.ts` | 62.50% | 83.33% | 57.89% | +| `src/cli/explore.ts` | 62.50% | 100.00% | 50.00% | +| `src/languages/definitions/typescript.ts` | 62.86% | 81.82% | 42.86% | +| `src/impact/call-compatibility/textScanner.ts` | 66.04% | 100.00% | 70.75% | +| `src/cli/index.ts` | 66.67% | 100.00% | 58.97% | +| `src/indexer.ts` | 66.67% | 66.67% | 100.00% | +| `src/cli/bootstrap.ts` | 66.67% | 50.00% | n/a | +| `src/languages/definitions/css.ts` | 66.67% | 0.00% | n/a | +| `src/languages/definitions/html.ts` | 66.67% | 0.00% | n/a | +| `src/languages/definitions/less.ts` | 66.67% | 0.00% | n/a | +| `src/languages/definitions/sql.ts` | 66.67% | 0.00% | n/a | +| `src/languages/definitions/svelte.ts` | 66.67% | 0.00% | n/a | +| `src/languages/definitions/vue.ts` | 66.67% | 0.00% | n/a | +| `src/agent/handles.ts` | 68.00% | 91.67% | 45.45% | +| `src/cli/artifact.ts` | 70.00% | 100.00% | 72.22% | +| `src/agent/query-index/workerPool.ts` | 70.97% | 62.50% | 56.25% | +| `src/agent/followUps.ts` | 72.15% | 77.78% | 53.91% | +| `src/cli/context.ts` | 72.22% | 60.78% | 75.25% | ## Type-Only Or Re-Export Files These files have line records but no function or branch records, so they are tracked outside the runtime ranking. -| File | Lines | Functions | Branches | -| --- | ---: | ---: | ---: | -| `src/cliBootstrap.ts` | 0.00% | n/a | n/a | -| `src/languages/definitions/jsFamily.ts` | 100.00% | n/a | n/a | -| `src/duplicate-keywords.ts` | 100.00% | n/a | n/a | -| `src/impact/types.ts` | 100.00% | n/a | n/a | -| `src/languages/definitions/adoc.ts` | 100.00% | n/a | n/a | -| `src/languages/definitions/astro.ts` | 100.00% | n/a | n/a | -| `src/languages/definitions/hbs.ts` | 100.00% | n/a | n/a | -| `src/languages/definitions/markdown.ts` | 100.00% | n/a | n/a | -| `src/languages/definitions/mdx.ts` | 100.00% | n/a | n/a | -| `src/languages/definitions/rst.ts` | 100.00% | n/a | n/a | +| File | Lines | Functions | Branches | +| --------------------------------------- | ------: | --------: | -------: | +| `src/cliBootstrap.ts` | 0.00% | n/a | n/a | +| `src/languages/definitions/jsFamily.ts` | 100.00% | n/a | n/a | +| `src/duplicate-keywords.ts` | 100.00% | n/a | n/a | +| `src/impact/types.ts` | 100.00% | n/a | n/a | +| `src/languages/definitions/adoc.ts` | 100.00% | n/a | n/a | +| `src/languages/definitions/astro.ts` | 100.00% | n/a | n/a | +| `src/languages/definitions/hbs.ts` | 100.00% | n/a | n/a | +| `src/languages/definitions/markdown.ts` | 100.00% | n/a | n/a | +| `src/languages/definitions/mdx.ts` | 100.00% | n/a | n/a | +| `src/languages/definitions/rst.ts` | 100.00% | n/a | n/a | Generated from LCOV by `node ./scripts/coverage-markdown.mjs`. From 28d15d1300be464816da6da25546b5cfb4295f9e Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Mon, 17 Aug 2026 16:59:30 -0400 Subject: [PATCH 37/37] fix: preserve SQLite cleanup retry after rebase --- docs/coverage/js.md | 10 +++++----- tests/sqlite-query-bounds.test.ts | 18 +----------------- 2 files changed, 6 insertions(+), 22 deletions(-) diff --git a/docs/coverage/js.md b/docs/coverage/js.md index e6102cf0..41f2b61a 100644 --- a/docs/coverage/js.md +++ b/docs/coverage/js.md @@ -6,9 +6,9 @@ Source: `coverage/js/lcov.info` | Metric | Hit | Found | Coverage | | --------- | ----: | ----: | -------: | -| Lines | 27601 | 30381 | 90.85% | -| Functions | 4609 | 4895 | 94.16% | -| Branches | 20774 | 26174 | 79.37% | +| Lines | 28185 | 30997 | 90.93% | +| Functions | 4694 | 4989 | 94.09% | +| Branches | 21296 | 26821 | 79.40% | ## Least-covered Files @@ -20,7 +20,6 @@ Source: `coverage/js/lcov.info` | `src/cli/explore.ts` | 62.50% | 100.00% | 50.00% | | `src/languages/definitions/typescript.ts` | 62.86% | 81.82% | 42.86% | | `src/impact/call-compatibility/textScanner.ts` | 66.04% | 100.00% | 70.75% | -| `src/cli/index.ts` | 66.67% | 100.00% | 58.97% | | `src/indexer.ts` | 66.67% | 66.67% | 100.00% | | `src/cli/bootstrap.ts` | 66.67% | 50.00% | n/a | | `src/languages/definitions/css.ts` | 66.67% | 0.00% | n/a | @@ -29,6 +28,7 @@ Source: `coverage/js/lcov.info` | `src/languages/definitions/sql.ts` | 66.67% | 0.00% | n/a | | `src/languages/definitions/svelte.ts` | 66.67% | 0.00% | n/a | | `src/languages/definitions/vue.ts` | 66.67% | 0.00% | n/a | +| `src/cli/index.ts` | 67.39% | 100.00% | 60.47% | | `src/agent/handles.ts` | 68.00% | 91.67% | 45.45% | | `src/cli/artifact.ts` | 70.00% | 100.00% | 72.22% | | `src/agent/query-index/workerPool.ts` | 70.97% | 62.50% | 56.25% | @@ -42,6 +42,7 @@ These files have line records but no function or branch records, so they are tra | File | Lines | Functions | Branches | | --------------------------------------- | ------: | --------: | -------: | | `src/cliBootstrap.ts` | 0.00% | n/a | n/a | +| `src/util/identifiers.ts` | 100.00% | n/a | n/a | | `src/languages/definitions/jsFamily.ts` | 100.00% | n/a | n/a | | `src/duplicate-keywords.ts` | 100.00% | n/a | n/a | | `src/impact/types.ts` | 100.00% | n/a | n/a | @@ -50,6 +51,5 @@ These files have line records but no function or branch records, so they are tra | `src/languages/definitions/hbs.ts` | 100.00% | n/a | n/a | | `src/languages/definitions/markdown.ts` | 100.00% | n/a | n/a | | `src/languages/definitions/mdx.ts` | 100.00% | n/a | n/a | -| `src/languages/definitions/rst.ts` | 100.00% | n/a | n/a | Generated from LCOV by `node ./scripts/coverage-markdown.mjs`. diff --git a/tests/sqlite-query-bounds.test.ts b/tests/sqlite-query-bounds.test.ts index b7b06349..890fa9ba 100644 --- a/tests/sqlite-query-bounds.test.ts +++ b/tests/sqlite-query-bounds.test.ts @@ -12,29 +12,13 @@ import { } from "../src/mcp/sqliteGuard.js"; import { queryGraphSqliteRaw, SqliteQueryDeadlineExceededError } from "../src/sqlite/query.js"; -// A deadline-terminated query's worker thread is force-terminated but, if it was blocked +// A deadline-exceeded query requests worker termination but, if it was blocked // inside a single synchronous native SQLite call, keeps running that call in the // background until it returns naturally (see rawQueryWorkerPool.ts). On Windows this can // hold the temp db file open for a short window after the deadline test's assertions // already ran. This is a real platform race (an actual lingering OS file lock, not // simulated timing logic), so it is retried against the real clock instead of being // modeled with fake timers. -async function removeWithRetry(root: string): Promise { - const deadline = Date.now() + 10_000; - for (;;) { - try { - await fsp.rm(root, { recursive: true, force: true }); - return; - } catch (error) { - const code = (error as NodeJS.ErrnoException).code; - if (code !== "EBUSY" && code !== "ENOTEMPTY") throw error; - if (Date.now() > deadline) throw error; - const { promise, resolve } = Promise.withResolvers(); - setTimeout(resolve, 100); - await promise; - } - } -} async function withTempDb(run: (dbPath: string) => Promise): Promise { const root = await fsp.mkdtemp(path.join(os.tmpdir(), "cg-sqlite-bounds-"));