fix: bound raw MCP query worker lifecycle

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,

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",

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;

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<T> = {
   push: (value: T) => void;
   close: () => void;
   next: () => Promise<IteratorResult<T>>;
 };
 
-function createAsyncQueue<T>(): AsyncQueue<T> {
+function createAsyncQueue<T>(maxQueuedChunks: number): AsyncQueue<T> {
   const values: T[] = [];
   const waiters: Array<(result: IteratorResult<T>) => void> = [];
   let closed = false;
@@ -90,6 +131,9 @@ function createAsyncQueue<T>(): AsyncQueue<T> {
         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,13 +234,24 @@ 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,
   index: ProjectIndex,
   options: ImpactStreamingOptions,
   context: ImpactStreamingContext = {},
 ): AsyncGenerator<ImpactStreamChunk> {
+  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<ImpactStreamChunk>();
+    const impactQueue = createAsyncQueue<ImpactStreamChunk>(
+      context.maxQueuedChunks ?? DEFAULT_MAX_IMPACT_STREAM_QUEUED_CHUNKS,
+    );
     const emittedSignatures = new Set<string>();
     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();
   }
 }
 

src/mcp/server.ts

@@ -172,6 +172,10 @@ type OriginValidator = (request: IncomingMessage, response: ServerResponse) => b
 
 export type CodegraphMcpFreshResult<T extends object> = 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.
@@ -321,11 +325,14 @@ export type CodegraphMcpHandlers = {
     refreshed: true;
     warmup: CodegraphMcpWarmupMode;
   }>;
-  query_sqlite: (request: {
-    query: string;
-    params?: Array<string | number | null> | undefined;
-    limit?: number | undefined;
-  }) => Promise<CodegraphMcpFreshResult<RawSqlResult>>;
+  query_sqlite: (
+    request: {
+      query: string;
+      params?: Array<string | number | null> | undefined;
+      limit?: number | undefined;
+    },
+    options?: McpToolExecutionOptions,
+  ) => Promise<CodegraphMcpFreshResult<RawSqlResult>>;
   artifact_build: (request: {
     outDir?: string | undefined;
     sqlite?: boolean | undefined;
@@ -959,7 +966,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.");
       }
@@ -980,6 +987,7 @@ function createCodegraphMcpHandlersForSession(
       }
       const result = await queryGraphSqliteRaw(realSqlitePath, request.query, request.params ?? [], {
         maxRows: normalizeSqliteRowLimit(request.limit),
+        ...(executionOptions?.signal ? { signal: executionOptions.signal } : {}),
       });
       return { ...boundRawSqlResult(result, DEFAULT_SQLITE_BYTE_LIMIT), freshness: artifactFreshness };
     },
@@ -1107,7 +1115,12 @@ export function createCodegraphMcpProtocolServer(
       console.error(`[codegraph] installed-version check failed: ${errorMessage(error)}`);
     }
     try {
-      const result = await callMcpTool(handlers, request.params.name, request.params.arguments ?? {});
+      const result = await callMcpTool(
+        handlers,
+        request.params.name,
+        request.params.arguments ?? {},
+        ctx.mcpReq.signal,
+      );
       await emitFirstToolCallVisibility(
         "info",
         1,
@@ -1557,7 +1570,12 @@ function isMcpNodeRequest(request: IncomingMessage): request is IncomingMessage
   return request.method !== undefined && request.url !== undefined;
 }
 
-async function callMcpTool(handlers: CodegraphMcpHandlers, name: string, input: unknown): Promise<unknown> {
+async function callMcpTool(
+  handlers: CodegraphMcpHandlers,
+  name: string,
+  input: unknown,
+  signal?: AbortSignal,
+): Promise<unknown> {
   switch (name) {
     case "search":
       return await handlers.search(parseMcpToolInput(searchSchema, input, "search"));
@@ -1614,7 +1632,9 @@ async function callMcpTool(handlers: CodegraphMcpHandlers, name: string, input:
     case "review":
       return await handlers.review(parseMcpToolInput(reviewSchema, input, "review"));
     case "query_sqlite":
-      return await handlers.query_sqlite(parseMcpToolInput(querySqliteSchema, input, "query_sqlite"));
+      return await handlers.query_sqlite(parseMcpToolInput(querySqliteSchema, input, "query_sqlite"), {
+        ...(signal ? { signal } : {}),
+      });
     case "refresh_index":
       return await handlers.refresh_index(parseMcpToolInput(refreshIndexSchema, input, "refresh_index"));
     case "artifact_build":