fix: bound MCP session refresh and request handling

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<WorkspaceSymbolsResponse> {
-  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<RenamePreviewRequest, "root" | "buildOptions">,
   runtimeOptions: ToolRenamePreviewRuntimeOptions = {},
 ): Promise<RenamePreviewResponse> {
-  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<RefactorPlanRequest, "root" | "buildOptions">,
   runtimeOptions: ToolRefactorPlanRuntimeOptions = {},
 ): Promise<RefactorPlanResponse> {
-  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,

src/agent/query-index/sessionStore.ts

@@ -10,17 +10,18 @@ type SessionQueryIndexState = {
 };
 
 const QUERY_INDEX_BY_SESSION = new WeakMap<AgentSession, SessionQueryIndexState>();
+const MAX_QUERY_INDEX_GENERATION_RETRIES = 3;
 
 function closeHandle(handle: QueryIndexHandle): void {
   handle.store?.close();
 }
 
 function closeState(state: SessionQueryIndexState): void {
+  state.closing = true;
   if (state.resolved) {
     closeHandle(state.resolved);
     return;
   }
-  state.closing = true;
   void state.handle.then(closeHandle, () => undefined);
 }
 
@@ -29,29 +30,36 @@ export async function ensureSessionQueryIndex(
   snapshot: AgentProjectSnapshot,
 ): Promise<QueryIndexHandle> {
   const identity = snapshot.index.projectSnapshotIdentity ?? "";
-  const existing = QUERY_INDEX_BY_SESSION.get(session);
-  if (existing?.identity === identity) return await existing.handle;
-  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.resolved = resolved;
-  } else if (!state.closing) {
+    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);
   }
-  if (snapshot.buildReport) snapshot.buildReport.queryIndex = resolved.diagnostics;
-  if (snapshot.index.buildReport) snapshot.index.buildReport.queryIndex = resolved.diagnostics;
-  return resolved;
+  throw new Error(
+    "Query index generation changed repeatedly while loading; retry the request after refresh completes.",
+  );
 }
 
 export function disposeSessionQueryIndex(session: AgentSession): void {

src/agent/search.ts

@@ -202,6 +202,8 @@ const NATURAL_LANGUAGE_SYNTAX_TERMS = new Set([
   "or",
 ]);
 const SEARCH_CACHES = new WeakMap<AgentProjectSnapshot, SearchCache>();
+// 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<AgentSession, Map<string, Promise<AgentSearchResponse>>>();
 const SEARCH_RANKING_VERSION = 2;
 
@@ -234,6 +236,7 @@ export async function searchCodegraphWithSession(
   const cacheKey = searchResultCacheKey(snapshot, request);
   const existing = resultCache.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,7 +248,7 @@ 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);
   });
@@ -352,6 +355,20 @@ function getSessionSearchResultCache(session: AgentSession): Map<string, Promise
   return created;
 }
 
+function promoteSessionSearchResult(
+  cache: Map<string, Promise<AgentSearchResponse>>,
+  key: string,
+  result: Promise<AgentSearchResponse>,
+): void {
+  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.delete(oldest);
+  }
+}
+
 function searchResultCacheKey(snapshot: AgentProjectSnapshot, request: AgentSearchRequest): string {
   const { rankTokens, normalizedRankPhrase, identifierLike } = buildQueryTerms(request.query, snapshot);
   return JSON.stringify({

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: [],

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);
   };
 }
 

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<string>;
@@ -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<ParsedJsonBody> {
+export async function readJsonRequestBody(
+  request: IncomingMessage,
+  maxBytes: number,
+  timeoutMs: number,
+): Promise<ParsedJsonBody> {
   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 {