From da5ccfb7568053300d81c5722e7ece72ceaf0b88 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sat, 15 Aug 2026 00:49:13 -0400 Subject: [PATCH 1/4] 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 b645a0fb..cb89d11e 100644 --- a/src/agent/search.ts +++ b/src/agent/search.ts @@ -202,7 +202,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 { @@ -232,8 +233,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 }; @@ -245,9 +247,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; } @@ -343,15 +345,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 3ac12602..57a3b36e 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 0450f8a8..c628dfb1 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 2f224fb1..a8bcdd68 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"; @@ -778,6 +788,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 @@ -793,6 +813,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, @@ -883,6 +923,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); @@ -935,6 +976,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 df171253..e6f6d3d2 100644 --- a/tests/session.test.ts +++ b/tests/session.test.ts @@ -1188,6 +1188,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 502d33d3e3567af8e0e08a48d192d1f7f5e68737 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sat, 15 Aug 2026 00:50:19 -0400 Subject: [PATCH 2/4] 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 2369d84fa8cb1e49ce8552a03c59e899956df5a6 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sat, 15 Aug 2026 00:55:56 -0400 Subject: [PATCH 3/4] 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 cb89d11e..b33cd067 100644 --- a/src/agent/search.ts +++ b/src/agent/search.ts @@ -202,8 +202,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 { @@ -233,7 +234,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; @@ -249,7 +250,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; } @@ -345,30 +346,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 355004152f98c2def500b0ce1ab5de4e7d954ab0 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sat, 15 Aug 2026 00:57:33 -0400 Subject: [PATCH 4/4] 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;