@lzehrung/codegraph 2.1.0 → 2.1.1

package/dist/mcp/server.js

@@ -1,4 +1,5 @@
 import { randomUUID } from "node:crypto";
+import { AsyncLocalStorage } from "node:async_hooks";
 import fs from "node:fs/promises";
 import { createServer } from "node:http";
 import path from "node:path";
@@ -28,16 +29,17 @@ import { DEFAULT_BOUNDED_IMPACT_BUDGETS } from "../impact/budgets.js";
 import { buildReviewReport } from "../review.js";
 import { boundReviewReportForTransport } from "../review/types.js";
 import { SQLITE_ARTIFACT_FILE_SIGNATURES_METADATA_KEY, queryGraphSqliteRaw } 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";
-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 { SymbolKind } from "../indexer/types.js";
 import { DEFAULT_WORKSPACE_SYMBOL_LIMIT, MAX_WORKSPACE_SYMBOL_LIMIT } from "../indexer/workspace-symbols.js";
-import { assertMcpSqliteQueryResourceBounded, boundRawSqlResult, DEFAULT_SQLITE_BYTE_LIMIT, normalizeSqliteRowLimit, } from "./sqliteGuard.js";
-import { DEFAULT_CALL_HIERARCHY_LIMIT, MAX_CALL_HIERARCHY_LIMIT, DEFAULT_MCP_COLLECTION_LIMIT, DEFAULT_TYPE_HIERARCHY_LIMIT, listCodegraphMcpTools, MAX_TYPE_HIERARCHY_LIMIT, MAX_MCP_COLLECTION_LIMIT, MAX_RENAME_PREVIEW_EDITS, MCP_TOOLS, MAX_REFACTOR_PLAN_LIMIT, } from "./tools.js";
+import { assertMcpSqliteQueryResourceBounded, DEFAULT_SQLITE_BYTE_LIMIT, normalizeSqliteRowLimit, } from "./sqliteGuard.js";
+import { DEFAULT_CALL_HIERARCHY_LIMIT, MAX_CALL_HIERARCHY_LIMIT, DEFAULT_MCP_COLLECTION_LIMIT, DEFAULT_TYPE_HIERARCHY_LIMIT, listCodegraphMcpTools, MAX_TYPE_HIERARCHY_LIMIT, MAX_MCP_COLLECTION_LIMIT, MAX_RENAME_PREVIEW_EDITS, MCP_TOOL_REGISTRY, MAX_REFACTOR_PLAN_LIMIT, } from "./tools.js";
 import { buildAllowedHostHeaders, buildAllowedOriginHostnames, closeHttpServer, emptyAllowedHostHeaderRules, formatHostForUrl, getHeaderValue, getHttpServerPort, getRequestPath, isAllowedHostHeader, listenOnHttpServer, readJsonRequestBody, waitForHttpServerClose, writeJsonResponse, writeJsonRpcError, } from "./http.js";
 import { getCurrentNativeBindingOrigin } from "../native/runtime.js";
 import { captureCodegraphRuntimeIdentity, createInstalledVersionChecker, } from "../runtimeIdentity.js";
@@ -47,12 +49,34 @@ import { awaitStdioMcpLifecycle, DEFAULT_MCP_STDIO_IDLE_TIMEOUT_MS } from "./std
 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_TIMEOUT_MS = 30 * 60 * 1000;
+const MAX_MCP_HTTP_BODY_TIMEOUT_MS = 2_147_483_647;
+const MAX_MCP_TOOL_TIMEOUT_MS = 2_147_483_647;
+export const DEFAULT_MCP_TOOL_CONCURRENCY = 4;
+function normalizeMcpToolConcurrency(value) {
+    if (!Number.isFinite(value))
+        return DEFAULT_MCP_TOOL_CONCURRENCY;
+    return Math.max(1, Math.floor(value));
+}
+function assertMcpHttpBodyTimeout(value) {
+    if (!Number.isSafeInteger(value) || value < 1 || value > MAX_MCP_HTTP_BODY_TIMEOUT_MS) {
+        throw new RangeError(`httpBodyTimeoutMs must be a positive integer no greater than ${MAX_MCP_HTTP_BODY_TIMEOUT_MS}.`);
+    }
+    return value;
+}
+function assertMcpToolTimeout(value) {
+    if (!Number.isSafeInteger(value) || value < 0 || value > MAX_MCP_TOOL_TIMEOUT_MS) {
+        throw new RangeError(`mcpToolTimeoutMs must be a whole number from 0 through ${MAX_MCP_TOOL_TIMEOUT_MS}; 0 disables the deadline.`);
+    }
+    return value;
+}
+const legacyRequestAbortStorage = new AsyncLocalStorage();
 const MAX_MCP_FRESHNESS_CHANGED_FILES = 25;
+const MAX_MCP_FRESHNESS_RETRIES = 3;
 const SQLITE_ARTIFACT_STAT_CONCURRENCY = 64;
 function assertMcpSessionOptions(options) {
-    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, root) {
     assertMcpSessionOptions(options);
@@ -72,7 +96,7 @@ function startCodegraphMcpWarmup(session, warmup) {
     }
     return undefined;
 }
-async function createWarmedCodegraphMcpHandlers(options) {
+async function createWarmedCodegraphMcpResources(options) {
     const root = path.resolve(options.root);
     const session = createCodegraphMcpSession(options, root);
     await startCodegraphMcpWarmup(session, options.warmup);
@@ -81,7 +105,10 @@ async function createWarmedCodegraphMcpHandlers(options) {
     void host;
     void port;
     void onHttpListen;
-    return createCodegraphMcpHandlersForSession({ ...handlerOptions, root }, session);
+    return {
+        handlers: createCodegraphMcpHandlersForSession({ ...handlerOptions, root }, session),
+        session,
+    };
 }
 const MCP_HTTP_PATH = "/mcp";
 const MAX_MCP_HTTP_BODY_BYTES = 1_000_000;
@@ -102,6 +129,8 @@ function createCodegraphMcpHandlersForSession(options, session) {
     let sqlitePath = configuredSqlitePath;
     let sqliteOutDir = configuredSqliteOutDir;
     let sqliteCanRefresh = configuredSqliteCanRefresh;
+    let refreshPromise;
+    let refreshEpoch = 0;
     const relative = (file) => toProjectDisplayPath(root, file);
     const boundedLimit = (limit, fallback, max) => {
         if (typeof limit !== "number" || !Number.isFinite(limit))
@@ -115,7 +144,7 @@ function createCodegraphMcpHandlersForSession(options, session) {
             state: "stale",
             changedFiles: boundedChangedFiles,
             changedFileCount: changedFiles.length,
-            omittedChangedFileCount: Math.max(0, changedFiles.length - boundedChangedFiles.length),
+            omittedChangedFileCount: countOmitted(changedFiles.length, boundedChangedFiles.length),
             reason,
         };
     };
@@ -125,9 +154,16 @@ function createCodegraphMcpHandlersForSession(options, session) {
         return { state: "fresh" };
     };
     const withFreshness = async (run) => {
-        const freshness = await checkMcpFreshness();
-        const result = await run();
-        return { ...result, freshness };
+        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) => {
         if (freshness.state === "fresh")
@@ -286,7 +322,7 @@ function createCodegraphMcpHandlersForSession(options, session) {
             // 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,
         };
@@ -302,13 +338,18 @@ function createCodegraphMcpHandlersForSession(options, session) {
             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) => {
         const findCalls = request.direction === "callers" ? findCallersWithSession : findCalleesWithSession;
@@ -333,12 +374,15 @@ function createCodegraphMcpHandlersForSession(options, session) {
         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) => await withFreshness(async () => {
         const { entries, ...meta } = await collectMcpDependencyEntries(request, request.direction === "deps" ? getDependencies : getReverseDependencies);
@@ -458,7 +502,7 @@ function createCodegraphMcpHandlersForSession(options, session) {
                 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);
@@ -534,7 +578,7 @@ function createCodegraphMcpHandlersForSession(options, session) {
             // complete, unbounded report should call `buildReviewReport` directly.
             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.");
             }
@@ -555,16 +599,32 @@ function createCodegraphMcpHandlersForSession(options, session) {
             }
             const result = await queryGraphSqliteRaw(realSqlitePath, request.query, request.params ?? [], {
                 maxRows: normalizeSqliteRowLimit(request.limit),
+                maxBytes: DEFAULT_SQLITE_BYTE_LIMIT,
+                ...(executionOptions?.signal ? { signal: executionOptions.signal } : {}),
             });
-            return { ...boundRawSqlResult(result, DEFAULT_SQLITE_BYTE_LIMIT), freshness: artifactFreshness };
+            return { ...result, truncated: Boolean(result.truncated), freshness: artifactFreshness };
         },
         refresh_index: async (request) => {
             const warmup = request.warmup ?? "off";
-            session.invalidate();
-            sqlitePath = configuredSqlitePath;
-            sqliteOutDir = configuredSqliteOutDir;
-            sqliteCanRefresh = configuredSqliteCanRefresh;
-            await startCodegraphMcpWarmup(session, warmup);
+            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);
+            })();
+            refreshPromise = refresh;
+            try {
+                await refresh;
+            }
+            finally {
+                if (refreshPromise === refresh)
+                    refreshPromise = undefined;
+            }
             return { refreshed: true, warmup };
         },
         artifact_build: async (request) => {
@@ -602,15 +662,49 @@ function createCodegraphMcpHandlersForSession(options, session) {
         },
     };
 }
-export function createCodegraphMcpProtocolServer(handlers, runtimeIdentity = captureCodegraphRuntimeIdentity(getCurrentNativeBindingOrigin()), installedVersion = createInstalledVersionChecker(runtimeIdentity), toolCallState = { firstToolCallPending: true }) {
+function createMcpToolOperationTracker() {
+    let accepting = true;
+    const operations = new Set();
+    return {
+        isAccepting: () => accepting,
+        track: (operation) => {
+            if (!accepting)
+                return undefined;
+            const tracked = operation();
+            operations.add(tracked);
+            void tracked.then(() => operations.delete(tracked), () => operations.delete(tracked));
+            return tracked;
+        },
+        stop: () => {
+            accepting = false;
+        },
+        drain: async () => {
+            accepting = false;
+            await Promise.allSettled([...operations]);
+        },
+    };
+}
+export function createCodegraphMcpProtocolServer(handlers, runtimeIdentity = captureCodegraphRuntimeIdentity(getCurrentNativeBindingOrigin()), installedVersion = createInstalledVersionChecker(runtimeIdentity), toolCallState = { firstToolCallPending: true }, maxConcurrentToolCalls = DEFAULT_MCP_TOOL_CONCURRENCY, mcpToolTimeoutMs = DEFAULT_MCP_TOOL_TIMEOUT_MS) {
+    return createCodegraphMcpProtocolServerWithTracker(handlers, runtimeIdentity, installedVersion, toolCallState, maxConcurrentToolCalls, assertMcpToolTimeout(mcpToolTimeoutMs), createMcpToolOperationTracker());
+}
+function createCodegraphMcpProtocolServerWithTracker(handlers, runtimeIdentity, installedVersion, toolCallState, maxConcurrentToolCalls, mcpToolTimeoutMs, toolOperations) {
     const server = new Server({
         name: "codegraph",
         version: runtimeIdentity.runningVersion,
     }, {
         capabilities: { tools: {}, logging: {} },
     });
-    server.setRequestHandler("tools/list", () => ({ tools: MCP_TOOLS }));
+    let inFlightToolCalls = 0;
+    const toolConcurrency = normalizeMcpToolConcurrency(maxConcurrentToolCalls);
+    server.setRequestHandler("tools/list", () => ({ tools: listCodegraphMcpTools() }));
     server.setRequestHandler("tools/call", async (request, ctx) => {
+        if (!toolOperations.isAccepting()) {
+            throw new Error("MCP server is shutting down.");
+        }
+        if (inFlightToolCalls >= toolConcurrency) {
+            throw new Error("MCP tool execution is busy; retry shortly.");
+        }
+        inFlightToolCalls += 1;
         const isFirstToolCall = toolCallState.firstToolCallPending;
         toolCallState.firstToolCallPending = false;
         const progressToken = isFirstToolCall ? getToolCallProgressToken(request.params) : undefined;
@@ -622,12 +716,7 @@ export function createCodegraphMcpProtocolServer(handlers, runtimeIdentity = cap
                 if (progressToken !== undefined) {
                     await ctx.mcpReq.notify({
                         method: "notifications/progress",
-                        params: {
-                            progressToken,
-                            progress,
-                            total: 1,
-                            message,
-                        },
+                        params: { progressToken, progress, total: 1, message },
                     });
                 }
             }
@@ -643,9 +732,24 @@ export function createCodegraphMcpProtocolServer(handlers, runtimeIdentity = cap
             console.error(`[codegraph] installed-version check failed: ${errorMessage(error)}`);
         }
         try {
-            const result = await callMcpTool(handlers, request.params.name, request.params.arguments ?? {});
-            await emitFirstToolCallVisibility("info", 1, `Codegraph finished warming the first tool call for '${request.params.name}'.`);
-            return toToolResult(result);
+            const toolCallAbort = createMcpToolAbortSignal([ctx.mcpReq.signal, legacyRequestAbortStorage.getStore()], request.params.name, mcpToolTimeoutMs);
+            const operation = toolOperations.track(() => callMcpTool(handlers, request.params.name, request.params.arguments ?? {}, toolCallAbort.signal));
+            if (operation === undefined) {
+                inFlightToolCalls -= 1;
+                toolCallAbort.dispose();
+                throw new Error("MCP server is shutting down.");
+            }
+            const releaseToolCall = () => {
+                inFlightToolCalls -= 1;
+            };
+            try {
+                const result = await awaitMcpToolOperation(toolCallAbort.signal, operation, releaseToolCall);
+                await emitFirstToolCallVisibility("info", 1, `Codegraph finished warming the first tool call for '${request.params.name}'.`);
+                return toToolResult(result);
+            }
+            finally {
+                toolCallAbort.dispose();
+            }
         }
         catch (error) {
             await emitFirstToolCallVisibility("error", 1, `Codegraph failed while warming the first tool call for '${request.params.name}'.`);
@@ -654,12 +758,20 @@ export function createCodegraphMcpProtocolServer(handlers, runtimeIdentity = cap
     });
     return server;
 }
-function createCodegraphMcpProtocolFactory(handlers, runtimeIdentity) {
+function createCodegraphMcpProtocolFactory(handlers, runtimeIdentity, maxConcurrentToolCalls = DEFAULT_MCP_TOOL_CONCURRENCY, mcpToolTimeoutMs = DEFAULT_MCP_TOOL_TIMEOUT_MS) {
     const installedVersion = createInstalledVersionChecker(runtimeIdentity);
     const toolCallState = { firstToolCallPending: true };
-    return () => createCodegraphMcpProtocolServer(handlers, runtimeIdentity, installedVersion, toolCallState);
+    const toolOperations = createMcpToolOperationTracker();
+    const timeout = assertMcpToolTimeout(mcpToolTimeoutMs);
+    return {
+        create: () => createCodegraphMcpProtocolServerWithTracker(handlers, runtimeIdentity, installedVersion, toolCallState, maxConcurrentToolCalls, timeout, toolOperations),
+        stop: () => toolOperations.stop(),
+        drain: () => toolOperations.drain(),
+    };
 }
 export async function serveCodegraphMcp(options) {
+    const configuredMcpToolTimeout = options.mcpToolTimeoutMs === undefined ? DEFAULT_MCP_TOOL_TIMEOUT_MS : options.mcpToolTimeoutMs;
+    const mcpToolTimeoutMs = assertMcpToolTimeout(configuredMcpToolTimeout);
     const port = options.port;
     if (port !== undefined) {
         const started = await startCodegraphMcpHttpServer({ ...options, port });
@@ -671,36 +783,46 @@ export async function serveCodegraphMcp(options) {
         await waitForHttpServerClose(started.server);
         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, {
+    const protocolFactory = createCodegraphMcpProtocolFactory(handlers, runtimeIdentity, options.mcpToolConcurrency ?? DEFAULT_MCP_TOOL_CONCURRENCY, mcpToolTimeoutMs);
+    const handle = serveStdio(protocolFactory.create, {
         legacy: "serve",
         onerror: (error) => {
             console.error(`[codegraph] MCP stdio error: ${error.message}`);
         },
     });
     const idleTimeoutMs = options.idleTimeoutMs ?? DEFAULT_MCP_STDIO_IDLE_TIMEOUT_MS;
-    await awaitStdioMcpLifecycle(handle, {
-        idleTimeoutMs,
-        onShutdown: (shutdownReason) => {
-            console.error(`[codegraph] MCP stdio shutting down (${shutdownReason})`);
-        },
-    });
+    try {
+        await awaitStdioMcpLifecycle(handle, {
+            idleTimeoutMs,
+            onShutdown: (shutdownReason) => {
+                console.error(`[codegraph] MCP stdio shutting down (${shutdownReason})`);
+            },
+        });
+    }
+    finally {
+        protocolFactory.stop();
+        await protocolFactory.drain();
+        session.invalidate();
+    }
     // Ensure orphaned stdio servers do not linger after the client is gone.
     process.exitCode = 0;
 }
 export async function startCodegraphMcpHttpServer(options) {
+    const httpBodyTimeoutMs = assertMcpHttpBodyTimeout(options.httpBodyTimeoutMs === undefined ? DEFAULT_MCP_HTTP_BODY_TIMEOUT_MS : options.httpBodyTimeoutMs);
+    const configuredMcpToolTimeout = options.mcpToolTimeoutMs === undefined ? DEFAULT_MCP_TOOL_TIMEOUT_MS : options.mcpToolTimeoutMs;
+    const mcpToolTimeoutMs = assertMcpToolTimeout(configuredMcpToolTimeout);
     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 protocolFactory = createCodegraphMcpProtocolFactory(handlers, runtimeIdentity, options.mcpToolConcurrency ?? DEFAULT_MCP_TOOL_CONCURRENCY, mcpToolTimeoutMs);
     const sessionStore = createLegacyMcpSessionStore({
         idleMs: options.httpSessionIdleMs ?? DEFAULT_MCP_HTTP_SESSION_IDLE_MS,
         maxCount: options.httpSessionMaxCount ?? DEFAULT_MCP_HTTP_SESSION_MAX_COUNT,
         evictionIntervalMs: options.httpSessionEvictionIntervalMs ?? DEFAULT_MCP_HTTP_SESSION_EVICTION_INTERVAL_MS,
     });
-    const modernHandler = createMcpHandler(createProtocolServer, {
+    const modernHandler = createMcpHandler(protocolFactory.create, {
         legacy: "reject",
         onerror: (error) => {
             console.error(`[codegraph] MCP HTTP error: ${error.message}`);
@@ -716,18 +838,31 @@ export async function startCodegraphMcpHttpServer(options) {
     let closeResourcesPromise;
     const closeResources = () => {
         closeResourcesPromise ??= (async () => {
+            protocolFactory.stop();
             sessionStore.stop();
-            await closeMcpResources(sessionStore.sessions, modernHandler.close);
+            try {
+                await closeMcpResources(sessionStore.sessions, modernHandler.close);
+            }
+            finally {
+                await protocolFactory.drain();
+                session.invalidate();
+            }
         })();
         return closeResourcesPromise;
     };
     const server = createServer((request, response) => {
-        void handleMcpHttpRequest(request, response, sessionStore, () => allowedHostHeaders, validateOrigin, modernNodeHandler, createProtocolServer);
+        void handleMcpHttpRequest(request, response, sessionStore, () => allowedHostHeaders, validateOrigin, modernNodeHandler, protocolFactory.create, httpBodyTimeoutMs);
     });
     server.on("close", () => {
         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);
@@ -747,7 +882,11 @@ export async function startCodegraphMcpHttpServer(options) {
         },
     };
 }
-async function handleMcpHttpRequest(request, response, sessionStore, getAllowedHostHeaders, validateOrigin, modernNodeHandler, createProtocolServer) {
+async function handleMcpHttpRequest(request, response, sessionStore, getAllowedHostHeaders, validateOrigin, modernNodeHandler, createProtocolServer, bodyTimeoutMs) {
+    const writeClosingJsonRpcError = (statusCode, message) => {
+        response.setHeader("connection", "close");
+        writeJsonRpcError(response, statusCode, message);
+    };
     const requestPath = getRequestPath(request);
     if (requestPath !== MCP_HTTP_PATH) {
         writeJsonResponse(response, 404, { error: "Not found" });
@@ -761,11 +900,15 @@ async function handleMcpHttpRequest(request, response, sessionStore, getAllowedH
         return;
     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") {
+                writeClosingJsonRpcError(408, "MCP request body timed out");
+                return;
+            }
             if (parsedBody.status === "invalid_json") {
                 writeJsonRpcError(response, 400, "Invalid JSON request body");
                 return;
@@ -805,13 +948,7 @@ async function handleLegacyMcpHttpPost(request, response, body, sessionStore, cr
             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;
     }
     if (!isInitializeRequest(body)) {
@@ -859,10 +996,16 @@ async function handleLegacyMcpHttpPost(request, response, body, sessionStore, cr
         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)
@@ -871,6 +1014,14 @@ async function handleLegacyMcpHttpPost(request, response, body, sessionStore, cr
     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) {
@@ -895,12 +1046,33 @@ async function handleExistingMcpSessionRequest(request, response, sessionStore)
         return;
     }
     sessionStore.touch(sessionId);
+    await handleLegacyMcpSessionRequest(session, request, response);
+}
+/**
+ * Runs one legacy HTTP request with a signal that aborts when its connection closes.
+ * The signal is scoped to this request and is available to nested MCP tool dispatch.
+ */
+export async function runWithLegacyRequestAbortSignal(request, response, operation) {
+    const abortController = new AbortController();
+    let completed = false;
+    const abortDisconnectedRequest = () => {
+        if (completed)
+            return;
+        if (request.aborted || (!response.writableFinished && response.destroyed)) {
+            abortController.abort(new Error("MCP HTTP request connection closed."));
+        }
+    };
+    request.on("aborted", abortDisconnectedRequest);
+    request.on("close", abortDisconnectedRequest);
+    response.on("close", abortDisconnectedRequest);
     try {
-        await handleLegacyMcpSessionRequest(session, request, response);
+        return await legacyRequestAbortStorage.run(abortController.signal, operation);
     }
-    catch (error) {
-        await sessionStore.delete(sessionId);
-        throw error;
+    finally {
+        completed = true;
+        request.off("aborted", abortDisconnectedRequest);
+        request.off("close", abortDisconnectedRequest);
+        response.off("close", abortDisconnectedRequest);
     }
 }
 async function handleLegacyMcpSessionRequest(session, request, response, body) {
@@ -909,7 +1081,9 @@ async function handleLegacyMcpSessionRequest(session, request, response, body) {
     if (tracksSseStream)
         session.openSseStreams += 1;
     try {
-        await session.transport.handleRequest(request, response, body);
+        await runWithLegacyRequestAbortSignal(request, response, async () => {
+            await session.transport.handleRequest(request, response, body);
+        });
     }
     finally {
         session.inFlightRequests -= 1;
@@ -1019,88 +1193,129 @@ async function closeMcpResources(sessions, closeModernHandler) {
 function isMcpNodeRequest(request) {
     return request.method !== undefined && request.url !== undefined;
 }
-async function callMcpTool(handlers, name, input) {
-    switch (name) {
+function createMcpToolAbortSignal(requestSignals, toolName, timeoutMs) {
+    const controller = new AbortController();
+    const listeners = [];
+    for (const requestSignal of requestSignals) {
+        if (requestSignal === undefined)
+            continue;
+        const onAbort = () => {
+            controller.abort(requestSignal.reason);
+        };
+        listeners.push({ signal: requestSignal, onAbort });
+        if (requestSignal.aborted)
+            onAbort();
+        else
+            requestSignal.addEventListener("abort", onAbort, { once: true });
+    }
+    const timeout = timeoutMs > 0
+        ? setTimeout(() => {
+            controller.abort(new Error(`MCP tool '${toolName}' exceeded the configured deadline of ${timeoutMs} ms.`));
+        }, timeoutMs)
+        : undefined;
+    timeout?.unref?.();
+    return {
+        signal: controller.signal,
+        dispose: () => {
+            clearTimeout(timeout);
+            for (const listener of listeners)
+                listener.signal.removeEventListener("abort", listener.onAbort);
+        },
+    };
+}
+export function awaitMcpToolOperation(signal, operation, onSettled) {
+    void operation.then(onSettled, onSettled);
+    return withAbortSignal(signal, operation);
+}
+function withAbortSignal(signal, operation) {
+    if (!signal)
+        return operation;
+    const cancellationError = () => {
+        const reason = signal.reason;
+        if (reason instanceof Error && reason.name !== "AbortError")
+            return reason;
+        return new Error("MCP tool call was cancelled.");
+    };
+    if (signal.aborted)
+        return Promise.reject(cancellationError());
+    const cancellation = Promise.withResolvers();
+    const onAbort = () => cancellation.reject(cancellationError());
+    signal.addEventListener("abort", onAbort, { once: true });
+    return Promise.race([operation, cancellation.promise]).finally(() => signal.removeEventListener("abort", onAbort));
+}
+export async function callMcpTool(handlers, name, input, signal) {
+    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, "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"));
-        case "callers":
-            return await handlers.calls({ ...parseMcpToolInput(callHierarchySchema, input, name), direction: "callers" });
-        case "callees":
-            return await handlers.calls({ ...parseMcpToolInput(callHierarchySchema, input, name), direction: "callees" });
+            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 "type_hierarchy":
-            return await handlers.type_hierarchy(parseMcpToolInput(typeHierarchyUnifiedSchema, input, "type_hierarchy"));
-        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",
-            });
+            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 "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"));
-        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" });
+            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 "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 ? { 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"));
-        default:
-            throw new Error(`Unknown MCP tool: ${name}`);
+            return await handlers.artifact_build(parseMcpToolInput(artifactBuildSchema, input, name), signal);
     }
 }
-async function callGotoTool(handlers, input) {
+async function callGotoTool(handlers, input, signal) {
     const request = parseMcpToolInput(navigationSchema, input, "goto");
     if (request.handle !== undefined)
-        return await handlers.goto({ handle: request.handle });
+        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, input) {
+async function callRefsTool(handlers, input, signal) {
     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.");
@@ -1110,7 +1325,7 @@ async function callRefsTool(handlers, input) {
         line: request.line,
         column: request.column,
         ...(request.limit !== undefined ? { limit: request.limit } : {}),
-    });
+    }, signal);
 }
 function toToolResult(value) {
     return {