From 223822463ce8a48d1c87958802000a5f83a1eed0 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sat, 15 Aug 2026 12:14:22 -0400 Subject: [PATCH 1/3] fix: bound and cancel impact analysis streaming Impact streaming's internal chunk queue buffered every produced item in an unbounded array whenever the consumer was not actively reading, and the background analyzeImpact() producer kept running to completion even after a consumer abandoned the stream (broke out of iteration, or the generator was otherwise returned early), retaining index snapshots, sets, and closures for work nobody would ever read. Add an internal AbortController to analyzeImpactStreaming(); the generator's finally block aborts it on every exit path, including the async-generator return protocol triggered by early consumer cancellation. The onImpactItem producer callback checks the signal and throws once it fires, unwinding analyzeImpact's in-progress work so no further symbol batches or transitive passes start. This needs no public API change: yield* delegation already forwards a caller's early return into the inner generator, so session.ts's analyzeImpactStream benefits without any code change there. Cap the internal queue at a bounded number of unread chunks (DEFAULT_MAX_IMPACT_STREAM_QUEUED_CHUNKS, overridable via the internal ImpactStreamingContext.maxQueuedChunks test seam). True backpressure would require making the onImpactItem emission callback awaitable at every synchronous call site in direct.ts/transitive.ts, which is outside this module. On overflow the stream surfaces an explicit ImpactStreamOverflowError as a terminal error chunk instead of silently dropping items or completing as if nothing were missed. --- src/impact/streaming.ts | 73 +++++++++++- tests/impact-streaming.test.ts | 202 ++++++++++++++++++++++++++++++++- tests/session.test.ts | 68 +++++++++++ 3 files changed, 339 insertions(+), 4 deletions(-) diff --git a/src/impact/streaming.ts b/src/impact/streaming.ts index 86cb59be..925d4756 100644 --- a/src/impact/streaming.ts +++ b/src/impact/streaming.ts @@ -71,13 +71,54 @@ function validateImpactStreamingOptions(options: ImpactStreamingOptions): "full" return streamSummary; } +/** + * Raised when a stream consumer falls behind the producer far enough that the buffered, + * unread chunk count would grow without bound. The producer stops (see + * `analyzeImpactStreaming`'s `onImpactItem` wiring) instead of silently dropping chunks, + * so a stalled consumer learns the stream could not keep up rather than quietly receiving + * a truncated-but-apparently-successful result. + */ +export class ImpactStreamOverflowError extends Error { + constructor(maxQueuedChunks: number) { + super( + `Impact stream consumer fell behind the producer: more than ${maxQueuedChunks} chunks were buffered ` + + "without being read. The stream was stopped instead of dropping chunks silently.", + ); + this.name = "ImpactStreamOverflowError"; + } +} + +/** + * Thrown from the `onImpactItem` producer callback once the consumer has abandoned the + * stream (see the `analyzeImpactStreaming` cancellation note below). Unwinds + * `analyzeImpact`'s in-progress work through its normal promise-rejection path; nothing + * outside this module ever observes it, since by construction nobody is listening to the + * stream anymore once it fires. + */ +class ImpactStreamAbandonedError extends Error { + constructor() { + super("Impact stream consumer stopped reading; cancelling in-progress analysis."); + this.name = "ImpactStreamAbandonedError"; + } +} + +/** + * Default cap on buffered-but-unread stream chunks before `ImpactStreamOverflowError` is + * raised. True backpressure (pausing the producer until the consumer catches up) would + * require the `onImpactItem` emission callback to be genuinely awaitable, which means + * awaiting it at every synchronous call site in `direct.ts`/`transitive.ts` — an invasive + * redesign of code outside this module. A hard cap is the non-invasive alternative: it + * turns unbounded memory growth into an explicit, surfaced failure instead. + */ +export const DEFAULT_MAX_IMPACT_STREAM_QUEUED_CHUNKS = 10_000; + type AsyncQueue = { push: (value: T) => void; close: () => void; next: () => Promise>; }; -function createAsyncQueue(): AsyncQueue { +function createAsyncQueue(maxQueuedChunks: number): AsyncQueue { const values: T[] = []; const waiters: Array<(result: IteratorResult) => void> = []; let closed = false; @@ -90,6 +131,9 @@ function createAsyncQueue(): AsyncQueue { waiter({ value, done: false }); return; } + if (values.length >= maxQueuedChunks) { + throw new ImpactStreamOverflowError(maxQueuedChunks); + } values.push(value); }, close() { @@ -175,6 +219,9 @@ export function impactItemEmissionKey(item: ImpactItem, partial: boolean): strin export type ImpactStreamingContext = { buildReport?: BuildReport | undefined; + /** @internal Overrides the buffered-chunk cap (`DEFAULT_MAX_IMPACT_STREAM_QUEUED_CHUNKS`). + * Test seam for deterministically exercising `ImpactStreamOverflowError`. */ + maxQueuedChunks?: number | undefined; }; /** @@ -187,6 +234,16 @@ export type ImpactStreamingContext = { * chains, graph edges, cycles, diagnostics, and schema metadata. Use * `streamSummary: "light"` when a caller only needs the progressive chunks and * a cheap terminal count/detail summary. + * + * Cancellation: if the consumer stops iterating early — a `for await` `break`, or an + * explicit `.return()` on the generator — the async-generator return protocol resumes + * this function's execution at its `finally` block, which aborts an internal + * `AbortController`. The background `analyzeImpact()` producer's `onImpactItem` callback + * checks that signal and throws once it fires, unwinding `analyzeImpact`'s in-progress + * work (no further batches or transitive passes run) instead of letting the whole + * analysis complete unread. This needs no signal parameter on the public API: `yield*` + * delegation (used by `session.ts`'s `analyzeImpactStream`) forwards a caller's + * `.return()` through automatically. */ export async function* analyzeImpactStreaming( projectRoot: string, @@ -194,6 +251,7 @@ export async function* analyzeImpactStreaming( options: ImpactStreamingOptions, context: ImpactStreamingContext = {}, ): AsyncGenerator { + const abortController = new AbortController(); try { const streamSummary = validateImpactStreamingOptions(options); const impactOptions = toImpactOptions(options); @@ -253,7 +311,9 @@ export async function* analyzeImpactStreaming( const normalizedChanges = normalizedDiff.files; const fileLevelFallback = impactOptions.fileLevelFallback ?? true; const fileLevelFallbackPaths = listFileLevelFallbackPaths(normalizedChanges, filesWithSymbols); - const impactQueue = createAsyncQueue(); + const impactQueue = createAsyncQueue( + context.maxQueuedChunks ?? DEFAULT_MAX_IMPACT_STREAM_QUEUED_CHUNKS, + ); const emittedSignatures = new Set(); let impactedItems: ImpactItem[] = []; let impactError: string | null = null; @@ -278,6 +338,9 @@ export async function* analyzeImpactStreaming( fileLevelFallbackPaths, diagnostics, onImpactItem: (item, phase) => { + if (abortController.signal.aborted) { + throw new ImpactStreamAbandonedError(); + } queueImpactItem(item, phase === "partial"); }, }) @@ -355,6 +418,12 @@ export async function* analyzeImpactStreaming( type: "error", error: errorMessage(error), }; + } finally { + // Runs on normal completion, on a caught error, and — via the async-generator return + // protocol — when the consumer stops iterating early. In the early-abandonment case + // this is what actually stops the background analysis: it flips the shared abort + // signal that the `onImpactItem` producer callback above checks and throws from. + abortController.abort(); } } diff --git a/tests/impact-streaming.test.ts b/tests/impact-streaming.test.ts index a40b2ff3..ab990552 100644 --- a/tests/impact-streaming.test.ts +++ b/tests/impact-streaming.test.ts @@ -1,4 +1,4 @@ -import { describe, it, expect } from "vitest"; +import { describe, it, expect, vi } from "vitest"; import path from "node:path"; import os from "node:os"; import fsp from "node:fs/promises"; @@ -8,8 +8,9 @@ import { type ImpactStreamChunk, type ImpactStreamSummaryReport, } from "../src/impact/index.js"; -import { impactItemEmissionKey } from "../src/impact/streaming.js"; +import { impactItemEmissionKey, ImpactStreamOverflowError } from "../src/impact/streaming.js"; import { buildProjectIndex } from "../src/index.js"; +import * as navigation from "../src/indexer/navigation.js"; import { runGit as git } from "./helpers/git.js"; async function mkTmpDir(prefix: string): Promise { @@ -587,3 +588,200 @@ index 1234567..abcdef0 100644 } }); }); + +/** Builds `symbolCount` distinct top-level exported functions in one file, plus a raw + * unified diff with one hunk per function so `mapChangedFileSymbols` reports + * `symbolCount` distinct changed symbols. */ +async function writeManySymbolFixture(root: string, symbolCount: number): Promise<{ diffText: string }> { + const lines = Array.from({ length: symbolCount }, (_, i) => `export function fn${i}() { return ${i}; }`); + await fsp.writeFile(path.join(root, "feature.ts"), `${lines.join("\n")}\n`, "utf8"); + const hunks = lines + .map((line, i) => { + const updated = line.replace(`return ${i};`, `return ${i + 1000};`); + return `@@ -${i + 1} +${i + 1} @@\n-${line}\n+${updated}\n`; + }) + .join(""); + const diffText = `diff --git a/feature.ts b/feature.ts +index 1234567..abcdef0 100644 +--- a/feature.ts ++++ b/feature.ts +${hunks}`; + return { diffText }; +} + +/** Polls a spy's call count until it stops changing, instead of awaiting a fixed real + * delay: the background analysis chain abandoned by the stream consumer settles + * asynchronously and this test holds no promise handle for it, so there is no signal to + * await other than the observable side effect (spy calls) itself. */ +async function waitForStableCallCount( + spy: { mock: { calls: unknown[] } }, + quietMs = 150, + timeoutMs = 5_000, +): Promise { + const deadline = Date.now() + timeoutMs; + let lastCount = spy.mock.calls.length; + let lastChangeAt = Date.now(); + while (Date.now() < deadline) { + const { promise, resolve } = Promise.withResolvers(); + setTimeout(resolve, 20); + await promise; + const count = spy.mock.calls.length; + if (count !== lastCount) { + lastCount = count; + lastChangeAt = Date.now(); + } else if (Date.now() - lastChangeAt >= quietMs) { + return lastCount; + } + } + return lastCount; +} + +describe("Impact streaming resource bounds", () => { + it("surfaces a bounded overflow error instead of silently truncating when a producer burst outruns the queue cap", async () => { + const root = await mkTmpDir("dg-stream-overflow-"); + await fsp.writeFile(path.join(root, "feature.ts"), "export function helper() { return 1; }\n", "utf8"); + const consumerCount = 6; + for (let i = 0; i < consumerCount; i += 1) { + await fsp.writeFile( + path.join(root, `consumer${i}.ts`), + `import { helper } from "./feature";\nexport function run${i}() { return helper(); }\n`, + "utf8", + ); + } + const index = await buildProjectIndex(root); + + try { + const diffText = `diff --git a/feature.ts b/feature.ts +index 1234567..abcdef0 100644 +--- a/feature.ts ++++ b/feature.ts +@@ -1 +1 @@ +-export function helper() { return 1; } ++export function helper() { return 2; } +`; + + // All 6 references to `helper` are emitted inside one synchronous loop in + // direct.ts (no `await` between them), so a cap of 2 overflows deterministically + // on every run regardless of machine speed: the consumer cannot possibly dequeue + // mid-burst. + const chunkTypes: string[] = []; + const errors: string[] = []; + const impactFiles: string[] = []; + for await (const chunk of analyzeImpactStreaming( + root, + index, + { provider: "raw", diffText }, + { maxQueuedChunks: 2 }, + )) { + chunkTypes.push(chunk.type); + if (chunk.type === "impactItem") impactFiles.push(chunk.item.file); + if (chunk.type === "error") errors.push(chunk.error); + } + + expect(chunkTypes).toContain("error"); + expect(chunkTypes).not.toContain("complete"); + expect(errors).toHaveLength(1); + expect(errors[0]).toMatch(/fell behind the producer/); + expect(errors[0]).toMatch(/more than 2 chunks/); + // The consumer learns exactly how far the stream got before it failed, not nothing. + expect(impactFiles.length).toBeGreaterThan(0); + expect(impactFiles.length).toBeLessThan(consumerCount); + } finally { + await fsp.rm(root, { recursive: true, force: true }); + } + }); + + it("does not overflow the same fixture under the default buffered-chunk cap", async () => { + const root = await mkTmpDir("dg-stream-no-overflow-"); + await fsp.writeFile(path.join(root, "feature.ts"), "export function helper() { return 1; }\n", "utf8"); + const consumerCount = 6; + for (let i = 0; i < consumerCount; i += 1) { + await fsp.writeFile( + path.join(root, `consumer${i}.ts`), + `import { helper } from "./feature";\nexport function run${i}() { return helper(); }\n`, + "utf8", + ); + } + const index = await buildProjectIndex(root); + + try { + const diffText = `diff --git a/feature.ts b/feature.ts +index 1234567..abcdef0 100644 +--- a/feature.ts ++++ b/feature.ts +@@ -1 +1 @@ +-export function helper() { return 1; } ++export function helper() { return 2; } +`; + + const chunkTypes: string[] = []; + const impactFiles: string[] = []; + for await (const chunk of analyzeImpactStreaming(root, index, { provider: "raw", diffText })) { + chunkTypes.push(chunk.type); + if (chunk.type === "impactItem") impactFiles.push(chunk.item.file); + } + + expect(chunkTypes).toContain("complete"); + expect(chunkTypes).not.toContain("error"); + const impactedFileSet = new Set(impactFiles); + for (let i = 0; i < consumerCount; i += 1) { + expect(impactedFileSet.has(`consumer${i}.ts`)).toBe(true); + } + } finally { + await fsp.rm(root, { recursive: true, force: true }); + } + }); + + it("stops the background analyzer once the consumer abandons the stream mid-analysis", async () => { + const root = await mkTmpDir("dg-stream-cancel-"); + const symbolCount = 40; + const { diffText } = await writeManySymbolFixture(root, symbolCount); + const index = await buildProjectIndex(root); + + try { + const findReferencesSpy = vi.spyOn(navigation, "findReferences"); + + let sawImpactItem = false; + for await (const chunk of analyzeImpactStreaming(root, index, { provider: "raw", diffText })) { + if (chunk.type === "impactItem") { + sawImpactItem = true; + break; + } + } + expect(sawImpactItem).toBe(true); + + const settledCalls = await waitForStableCallCount(findReferencesSpy); + // Changed symbols are analyzed in fixed batches of 8 (IMPACT_SYMBOL_BATCH_SIZE); + // each batch is awaited fully before the next starts. Cancelling mid-first-batch + // must prevent every later batch from ever starting: comfortably fewer than half + // of the 40 symbols should ever reach a reference lookup. + expect(settledCalls).toBeGreaterThan(0); + expect(settledCalls).toBeLessThan(symbolCount / 2); + } finally { + vi.restoreAllMocks(); + await fsp.rm(root, { recursive: true, force: true }); + } + }); + + it("analyzes every changed symbol when the same stream is consumed to completion", async () => { + const root = await mkTmpDir("dg-stream-nocancel-"); + const symbolCount = 40; + const { diffText } = await writeManySymbolFixture(root, symbolCount); + const index = await buildProjectIndex(root); + + try { + const findReferencesSpy = vi.spyOn(navigation, "findReferences"); + + const chunkTypes: string[] = []; + for await (const chunk of analyzeImpactStreaming(root, index, { provider: "raw", diffText })) { + chunkTypes.push(chunk.type); + } + + expect(chunkTypes).toContain("complete"); + expect(findReferencesSpy).toHaveBeenCalledTimes(symbolCount); + } finally { + vi.restoreAllMocks(); + await fsp.rm(root, { recursive: true, force: true }); + } + }); +}); diff --git a/tests/session.test.ts b/tests/session.test.ts index e6f6d3d2..f73779f7 100644 --- a/tests/session.test.ts +++ b/tests/session.test.ts @@ -3,6 +3,7 @@ import type { ICodeReviewSession } from "../src/index.js"; import type { BuildOptions, BuildReport, LanguageExtensionMap } from "../src/indexer/types.js"; import { CodeReviewSession, SessionManager, createCodeReviewSession } from "../src/session.js"; import * as indexerBuild from "../src/indexer/build-index.js"; +import * as navigation from "../src/indexer/navigation.js"; import path from "node:path"; import os from "node:os"; import fs from "node:fs"; @@ -1789,3 +1790,70 @@ describe("SessionManager", () => { expect(manager.getSession("shared")).toBe(existing); }); }); + +describe("CodeReviewSession impact stream cancellation", () => { + test("stops the background analyzer once a session.analyzeImpactStream consumer abandons the stream", async () => { + const root = await fsp.mkdtemp(path.join(os.tmpdir(), "dg-session-impact-stream-cancel-")); + try { + const symbolCount = 40; + const lines = Array.from({ length: symbolCount }, (_, i) => `export function fn${i}() { return ${i}; }`); + await fsp.writeFile(path.join(root, "feature.ts"), `${lines.join("\n")}\n`, "utf8"); + const hunks = lines + .map((line, i) => { + const updated = line.replace(`return ${i};`, `return ${i + 1000};`); + return `@@ -${i + 1} +${i + 1} @@\n-${line}\n+${updated}\n`; + }) + .join(""); + const diffText = `diff --git a/feature.ts b/feature.ts +index 1234567..abcdef0 100644 +--- a/feature.ts ++++ b/feature.ts +${hunks}`; + + const session = await createCodeReviewSession({ + root, + buildOptions: { cache: "memory", useBloomFilters: true }, + }); + const findReferencesSpy = vi.spyOn(navigation, "findReferences"); + try { + let sawImpactItem = false; + for await (const chunk of session.analyzeImpactStream({ provider: "raw", diffText })) { + if (chunk.type === "impactItem") { + sawImpactItem = true; + break; + } + } + expect(sawImpactItem).toBe(true); + + // session.analyzeImpactStream is `yield* analyzeImpactStreaming(...)`: this proves + // that delegation forwards the consumer's early `break` (an async-generator + // `.return()` call) through to the inner generator without any extra plumbing. + // Poll instead of a fixed sleep: the abandoned background chain settles + // asynchronously and this test holds no promise handle for it. + const deadline = Date.now() + 5_000; + let lastCount = findReferencesSpy.mock.calls.length; + let lastChangeAt = Date.now(); + while (Date.now() < deadline && Date.now() - lastChangeAt < 150) { + const { promise, resolve } = Promise.withResolvers(); + setTimeout(resolve, 20); + await promise; + const count = findReferencesSpy.mock.calls.length; + if (count !== lastCount) { + lastCount = count; + lastChangeAt = Date.now(); + } + } + + expect(lastCount).toBeGreaterThan(0); + // Changed symbols are analyzed in fixed batches of 8 (IMPACT_SYMBOL_BATCH_SIZE); + // cancelling mid-first-batch must prevent every later batch from ever starting. + expect(lastCount).toBeLessThan(symbolCount / 2); + } finally { + findReferencesSpy.mockRestore(); + session.dispose(); + } + } finally { + await fsp.rm(root, { recursive: true, force: true }); + } + }); +}); From ba5bd671182c61adfa25a830c8edaea7412d8244 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sat, 15 Aug 2026 12:14:32 -0400 Subject: [PATCH 2/3] fix: deadline raw SQLite query execution Raw query_sqlite reads ran a synchronous DatabaseSync iteration with row and byte caps but no time budget or cancellation: a non-recursive but expensive statement (a large join, ORDER BY random() with no matching index, ...) could hold the host event loop for as long as SQLite took to produce a row, and a client disconnect did not stop it. Run the query in a dedicated Piscina worker thread (same pattern as the existing query-index worker pool) with a hard deadline. On expiry the worker thread is force-terminated and the call rejects with SqliteQueryDeadlineExceededError immediately; the host event loop is never blocked regardless of how long the underlying query actually runs, and a subsequent query against the same file succeeds right away since concurrent read-only SQLite connections do not block each other. Pool teardown is fire-and-forget on deadline expiry rather than awaited, so the caller is never delayed by an orphaned worker thread still finishing a single already-in-flight synchronous native call (worker termination cannot preempt one in-progress call the same way it can prevent further JS from running) -- documented in rawQueryWorkerPool.ts, verified directly against a 200M-row recursive CTE. If the compiled worker asset cannot be located (a corrupted or partial install), the query falls back to running in-process under a per-row elapsed-time budget instead of refusing outright. That fallback is strictly weaker and is documented as such: the budget is only checked between rows the native iterator has already produced, so a statement that is slow to produce its very first row is not bounded by it. Preserves the existing row/byte cap contracts and normalizeSqliteRowLimit reuse in this file. --- src/sqlite/query.ts | 84 ++++++++++++++++++++++++++++--- src/sqlite/rawQueryWorker.ts | 45 +++++++++++++++++ src/sqlite/rawQueryWorkerPool.ts | 77 ++++++++++++++++++++++++++++ tests/sqlite-query-bounds.test.ts | 80 ++++++++++++++++++++++++++++- 4 files changed, 276 insertions(+), 10 deletions(-) create mode 100644 src/sqlite/rawQueryWorker.ts create mode 100644 src/sqlite/rawQueryWorkerPool.ts diff --git a/src/sqlite/query.ts b/src/sqlite/query.ts index d1e740ff..aecce786 100644 --- a/src/sqlite/query.ts +++ b/src/sqlite/query.ts @@ -8,35 +8,91 @@ import { MAX_SQLITE_ROW_LIMIT, normalizeSqliteRowLimit, } from "./rowBounds.js"; +import { + resolveRawSqlQueryWorkerPath, + runRawSqlQueryInWorker, + SqliteQueryDeadlineExceededError, +} from "./rawQueryWorkerPool.js"; export { queryGraphSqlite } from "./canned-query.js"; +export { SqliteQueryDeadlineExceededError }; + +/** Hard wall-clock budget for a single raw `query_sqlite` execution. */ +export const DEFAULT_SQLITE_QUERY_DEADLINE_MS = 10_000; export type QueryGraphSqliteRawOptions = { maxRows?: number | undefined; maxBytes?: number | undefined; maxCellBytes?: number | undefined; + deadlineMs?: number | undefined; }; +/** + * Runs a bounded read-only raw SQL query. + * + * Preferred path: the query executes in a dedicated worker thread with a hard + * `deadlineMs` budget (`rawQueryWorkerPool.ts`). On expiry the worker thread is + * terminated outright, which stops the query even while it is blocked inside a single + * synchronous `DatabaseSync` call — a client disconnect or a slow non-recursive + * statement (large join, `ORDER BY random()`, ...) can no longer hold the host event + * loop hostage indefinitely. + * + * Degraded fallback: if the compiled worker asset cannot be located (a corrupted or + * partial install — the normal build/publish/standalone pipelines all ship it), the + * query instead runs in-process under a *per-row* elapsed-time budget. That fallback is + * strictly weaker: the budget is only checked between rows the native iterator has + * already produced, so a statement that is slow to produce its very first row (for + * example a full-table scan with no matching rows) is not bounded by it. It exists to + * keep the common case usable in a degraded install, not as a substitute for the worker + * deadline. + */ export async function queryGraphSqliteRaw( outputPath: string, sql: string, params: Array = [], options?: QueryGraphSqliteRawOptions, +): Promise { + const maxRows = normalizeSqliteRowLimit(options?.maxRows ?? MAX_SQLITE_ROW_LIMIT); + const maxBytes = options?.maxBytes ?? DEFAULT_SQLITE_BYTE_LIMIT; + const maxCellBytes = options?.maxCellBytes ?? MAX_SQLITE_CELL_BYTES; + const deadlineMs = options?.deadlineMs ?? DEFAULT_SQLITE_QUERY_DEADLINE_MS; + + try { + resolveRawSqlQueryWorkerPath(); + } catch { + return await queryGraphSqliteRawInProcessBounded(outputPath, sql, params, { + maxRows, + maxBytes, + maxCellBytes, + deadlineMs, + }); + } + + return await runRawSqlQueryInWorker({ outputPath, sql, params, maxRows, maxBytes, maxCellBytes }, deadlineMs); +} + +async function queryGraphSqliteRawInProcessBounded( + outputPath: string, + sql: string, + params: Array, + bounds: { maxRows: number; maxBytes: number; maxCellBytes: number; deadlineMs: number }, ): Promise { return await withReadOnlySqliteDatabase(outputPath, (db) => { try { const stmt = db.prepare(sql); assertReadOnlyQueryStatement(stmt); const columns = stmt.columns().map((col) => col.name); - const maxRows = normalizeSqliteRowLimit(options?.maxRows ?? MAX_SQLITE_ROW_LIMIT); - const maxBytes = options?.maxBytes ?? DEFAULT_SQLITE_BYTE_LIMIT; - const maxCellBytes = options?.maxCellBytes ?? MAX_SQLITE_CELL_BYTES; - + const deadlineAt = Date.now() + bounds.deadlineMs; + const rows = withPerRowDeadline( + stmt.raw().iterate(params) as Iterable>, + deadlineAt, + bounds.deadlineMs, + ); // Always stream via iterate so per-cell and cumulative budgets apply before append. - return collectBoundedRawSqlRows(columns, stmt.raw().iterate(params) as Iterable>, { - maxRows, - maxBytes, - maxCellBytes, + return collectBoundedRawSqlRows(columns, rows, { + maxRows: bounds.maxRows, + maxBytes: bounds.maxBytes, + maxCellBytes: bounds.maxCellBytes, }); } catch (error) { if (isReadOnlySqliteError(error)) { @@ -46,3 +102,15 @@ export async function queryGraphSqliteRaw( } }); } + +/** Throws once the wall-clock deadline has passed between two already-produced rows. See + * the fallback caveat on `queryGraphSqliteRaw`: a slow-before-first-row query is not + * caught here, only slow-*between*-rows iteration is. */ +function* withPerRowDeadline(rows: Iterable, deadlineAt: number, deadlineMs: number): Generator { + for (const row of rows) { + if (Date.now() > deadlineAt) { + throw new SqliteQueryDeadlineExceededError(deadlineMs); + } + yield row; + } +} diff --git a/src/sqlite/rawQueryWorker.ts b/src/sqlite/rawQueryWorker.ts new file mode 100644 index 00000000..6b5f70c4 --- /dev/null +++ b/src/sqlite/rawQueryWorker.ts @@ -0,0 +1,45 @@ +import { isReadOnlySqliteError } from "../sqlite-driver.js"; +import type { RawSqlResult } from "./types.js"; +import { assertReadOnlyQueryStatement, withReadOnlySqliteDatabase } from "./database.js"; +import { collectBoundedRawSqlRows } from "./rowBounds.js"; + +/** + * Task payload for a single bounded raw SQL read, executed inside a dedicated worker + * thread (see `rawQueryWorkerPool.ts`). Every field must be structured-clone safe. + */ +export type RawQueryWorkerTask = { + outputPath: string; + sql: string; + params: Array; + maxRows: number; + maxBytes: number | undefined; + maxCellBytes: number | undefined; +}; + +/** + * Worker entry point. Mirrors the previous in-process body of `queryGraphSqliteRaw` + * exactly: open the database read-only, assert the statement is read-only, and stream + * rows through the shared row/byte-bounded collector. Running this inside a worker + * thread lets the pool enforce a hard execution deadline by terminating the thread — + * which works even mid-synchronous-iteration, since thread termination does not need + * the blocked thread's cooperation. + */ +export default async function runRawQueryWorkerTask(task: RawQueryWorkerTask): Promise { + return await withReadOnlySqliteDatabase(task.outputPath, (db) => { + try { + const stmt = db.prepare(task.sql); + assertReadOnlyQueryStatement(stmt); + const columns = stmt.columns().map((col) => col.name); + return collectBoundedRawSqlRows(columns, stmt.raw().iterate(task.params) as Iterable>, { + maxRows: task.maxRows, + maxBytes: task.maxBytes, + maxCellBytes: task.maxCellBytes, + }); + } catch (error) { + if (isReadOnlySqliteError(error)) { + throw new Error("Raw SQLite queries must be read-only result-producing statements such as SELECT or PRAGMA."); + } + throw error; + } + }); +} diff --git a/src/sqlite/rawQueryWorkerPool.ts b/src/sqlite/rawQueryWorkerPool.ts new file mode 100644 index 00000000..3af4594b --- /dev/null +++ b/src/sqlite/rawQueryWorkerPool.ts @@ -0,0 +1,77 @@ +import fs from "node:fs"; +import path from "node:path"; +import { fileURLToPath } from "node:url"; +import { Piscina } from "piscina"; +import { findPackageRoot } from "../util/packageInfo.js"; +import type { RawSqlResult } from "./types.js"; +import type { RawQueryWorkerTask } from "./rawQueryWorker.js"; + +export class SqliteQueryDeadlineExceededError extends Error { + constructor(deadlineMs: number) { + super(`SQLite query exceeded its ${deadlineMs}ms execution budget and was terminated.`); + this.name = "SqliteQueryDeadlineExceededError"; + } +} + +/** Resolves the compiled worker entry the same way `queryIndexWorker.js` is resolved: + * a compiled sibling next to this module (production/standalone layouts, where the + * whole `dist/` tree ships), falling back to the package-root-relative compiled path + * (vitest running this module from `src/`, where only `dist/` is built). */ +export function resolveRawSqlQueryWorkerPath(): string { + const selfDir = path.dirname(fileURLToPath(import.meta.url)); + const sibling = path.resolve(selfDir, "rawQueryWorker.js"); + if (fs.existsSync(sibling)) return sibling; + const packageRoot = findPackageRoot(selfDir); + const compiled = path.join(packageRoot, "dist", "sqlite", "rawQueryWorker.js"); + if (fs.existsSync(compiled)) return compiled; + throw new Error(`Raw SQLite query worker file not found: ${compiled}`); +} + +/** + * Runs a single bounded raw SQL read in a dedicated worker thread with a hard execution + * deadline. A fresh single-thread pool is created per call and destroyed afterward — + * matching the existing `prepareQueryIndexFilesInWorker` pattern — since `query_sqlite` + * calls are interactive, not a hot loop, and a persistent pool would need a shutdown hook + * this module has no access to register. + * + * On deadline expiry, Piscina's `signal` option force-terminates the worker thread + * (`worker.terminate()`) and rejects immediately — the caller never waits longer than + * `deadlineMs`, and the host event loop is never blocked by the query regardless of how + * long it runs. Cancellation is real (no further JS runs on that thread and the query + * can never touch this process's caller again), but it has one unavoidable limit shared + * by every in-process cancellation mechanism: `terminate()` cannot preempt a single + * already-in-flight synchronous native call. A query whose entire cost is inside one + * `sqlite3_step()` — a recursive CTE, or a plan that must fully sort/scan before it can + * produce a first row — keeps running on the orphaned worker thread in the background + * until that native call returns naturally; only then does the thread actually exit. + * We do not make the caller wait for that: `pool.destroy()` is fired and forgotten here, + * not awaited, so a subsequent query (in its own fresh pool) is never delayed by it. + * Concurrent read-only SQLite connections against the same file do not block each other, + * so the lingering background reader does not stop that subsequent query from succeeding. + * Verified directly: an aborted 200M-row recursive-CTE count rejects this call in + * ~`deadlineMs`, and an immediately following query against the same file succeeds in + * milliseconds. The one place the orphaned thread is still observable is process + * shutdown: Node cannot fully tear a process down while one of its Worker threads is + * blocked in native code, so a process exit racing a runaway query can itself be + * delayed until that native call returns — a platform limit of `worker_threads`, not of + * this module, and orthogonal to the per-call deadline this function guarantees. + */ +export async function runRawSqlQueryInWorker(task: RawQueryWorkerTask, deadlineMs: number): Promise { + const workerPath = resolveRawSqlQueryWorkerPath(); + const pool = new Piscina({ + filename: workerPath, + minThreads: 1, + maxThreads: 1, + idleTimeout: 5_000, + }); + try { + return (await pool.run(task, { signal: AbortSignal.timeout(deadlineMs) })) as RawSqlResult; + } catch (error) { + if (error instanceof Error && error.name === "AbortError") { + throw new SqliteQueryDeadlineExceededError(deadlineMs); + } + throw error; + } finally { + void pool.destroy().catch(() => {}); + } +} diff --git a/tests/sqlite-query-bounds.test.ts b/tests/sqlite-query-bounds.test.ts index 24654e99..29d38026 100644 --- a/tests/sqlite-query-bounds.test.ts +++ b/tests/sqlite-query-bounds.test.ts @@ -10,7 +10,31 @@ import { MAX_SQLITE_ROW_LIMIT, SQLITE_TRUNCATED_MARKER, } from "../src/mcp/sqliteGuard.js"; -import { queryGraphSqliteRaw } from "../src/sqlite/query.js"; +import { queryGraphSqliteRaw, SqliteQueryDeadlineExceededError } from "../src/sqlite/query.js"; + +// A deadline-terminated query's worker thread is force-terminated but, if it was blocked +// inside a single synchronous native SQLite call, keeps running that call in the +// background until it returns naturally (see rawQueryWorkerPool.ts). On Windows this can +// hold the temp db file open for a short window after the deadline test's assertions +// already ran. This is a real platform race (an actual lingering OS file lock, not +// simulated timing logic), so it is retried against the real clock instead of being +// modeled with fake timers. +async function removeWithRetry(root: string): Promise { + const deadline = Date.now() + 10_000; + for (;;) { + try { + await fsp.rm(root, { recursive: true, force: true }); + return; + } catch (error) { + const code = (error as NodeJS.ErrnoException).code; + if (code !== "EBUSY" && code !== "ENOTEMPTY") throw error; + if (Date.now() > deadline) throw error; + const { promise, resolve } = Promise.withResolvers(); + setTimeout(resolve, 100); + await promise; + } + } +} async function withTempDb(run: (dbPath: string) => Promise): Promise { const root = await fsp.mkdtemp(path.join(os.tmpdir(), "cg-sqlite-bounds-")); @@ -18,7 +42,7 @@ async function withTempDb(run: (dbPath: string) => Promise): Promise try { await run(dbPath); } finally { - await fsp.rm(root, { recursive: true, force: true }); + await removeWithRetry(root); } } @@ -102,3 +126,55 @@ describe("SQLite query byte/cell bounds during iterate", () => { } }); }); + +describe("SQLite raw query execution deadline", () => { + it("terminates an over-budget query with a bounded error while a subsequent query on the same file still succeeds", async () => { + await withTempDb(async (dbPath) => { + const db = new DatabaseSync(dbPath); + db.exec("CREATE TABLE t (n INTEGER);"); + db.prepare("INSERT INTO t (n) VALUES (?)").run(42); + db.close(); + + // The whole cost of this query is inside one synchronous native step (see + // rawQueryWorkerPool.ts): SQLite must finish counting before it can return the + // single aggregate row, so this reliably runs well past a short deadline without + // depending on machine speed for a *count* of loop iterations. + const slowSql = + "WITH RECURSIVE spin(x) AS (SELECT 1 UNION ALL SELECT x + 1 FROM spin WHERE x < 8000000) " + + "SELECT count(*) FROM spin;"; + + const start = Date.now(); + await expect(queryGraphSqliteRaw(dbPath, slowSql, [], { deadlineMs: 100 })).rejects.toMatchObject({ + name: "SqliteQueryDeadlineExceededError", + message: expect.stringMatching(/exceeded its 100ms execution budget/), + }); + const elapsed = Date.now() - start; + // The caller is bounded by the deadline, not by how long the runaway query + // actually takes to finish in the background (calibrated well above 100ms). + expect(elapsed).toBeLessThan(2_000); + + const result = await queryGraphSqliteRaw(dbPath, "SELECT n FROM t;", [], { deadlineMs: 5_000 }); + expect(result.rows).toEqual([[42]]); + expect(result.truncated).toBeFalsy(); + }); + }); + + it("does not reject an ordinary query that finishes comfortably inside its deadline", async () => { + await withTempDb(async (dbPath) => { + const db = new DatabaseSync(dbPath); + db.exec("CREATE TABLE t (n INTEGER);"); + db.prepare("INSERT INTO t (n) VALUES (?)").run(7); + db.close(); + + const result = await queryGraphSqliteRaw(dbPath, "SELECT n FROM t;", [], { deadlineMs: 5_000 }); + expect(result.rows).toEqual([[7]]); + }); + }); + + it("exposes SqliteQueryDeadlineExceededError as a named export for callers to distinguish deadline failures", () => { + const error = new SqliteQueryDeadlineExceededError(250); + expect(error).toBeInstanceOf(Error); + expect(error.name).toBe("SqliteQueryDeadlineExceededError"); + expect(error.message).toBe("SQLite query exceeded its 250ms execution budget and was terminated."); + }); +}); From bd463d88975e42e887f027f0429b1863b88a1373 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sat, 15 Aug 2026 23:00:56 -0400 Subject: [PATCH 3/3] fix: bound raw MCP query worker cleanup --- scripts/bundle-cli-lib.mjs | 18 ++- scripts/ensure-dist-for-tests-lib.mjs | 8 +- scripts/stage-core-package-lib.mjs | 5 +- src/mcp/server.ts | 38 ++++-- src/sqlite/query.ts | 23 +++- src/sqlite/rawQueryWorkerPool.ts | 147 ++++++++++++++++++++--- tests/cli-bundle-entry.test.ts | 2 + tests/core-package-surface.test.ts | 1 + tests/ensure-dist-for-tests.test.ts | 4 + tests/mcp-stream-cancellation.test.ts | 51 ++++++++ tests/query-index-worker-path.test.ts | 13 ++ tests/raw-query-worker-lifecycle.test.ts | 63 ++++++++++ 12 files changed, 336 insertions(+), 37 deletions(-) create mode 100644 tests/mcp-stream-cancellation.test.ts create mode 100644 tests/raw-query-worker-lifecycle.test.ts diff --git a/scripts/bundle-cli-lib.mjs b/scripts/bundle-cli-lib.mjs index cbed4db8..d683b203 100644 --- a/scripts/bundle-cli-lib.mjs +++ b/scripts/bundle-cli-lib.mjs @@ -12,10 +12,12 @@ export function getBundlePaths(rootDir = defaultRootDir) { rootDir, entryPoint: path.join(rootDir, "dist", "cliBootstrap.js"), workerEntryPoint: path.join(rootDir, "dist", "agent", "query-index", "queryIndexWorker.js"), + rawQueryWorkerEntryPoint: path.join(rootDir, "dist", "sqlite", "rawQueryWorker.js"), unbundledCli: path.join(rootDir, "dist", "cli.js"), outdir: path.join(rootDir, "dist", "bin"), bundledEntry: path.join(rootDir, "dist", "bin", "cli.js"), bundledWorker: path.join(rootDir, "dist", "bin", "queryIndexWorker.js"), + bundledRawQueryWorker: path.join(rootDir, "dist", "bin", "rawQueryWorker.js"), }; } @@ -37,12 +39,21 @@ export async function bundleCli({ rootDir = defaultRootDir, logLevel = "warning" if (!fs.existsSync(paths.workerEntryPoint)) { throw new Error(`Missing query worker build input: ${paths.workerEntryPoint}. Run tsc before bundling.`); } + if (!fs.existsSync(paths.rawQueryWorkerEntryPoint)) { + throw new Error( + `Missing raw SQLite worker build input: ${paths.rawQueryWorkerEntryPoint}. Run tsc before bundling.`, + ); + } fs.rmSync(paths.outdir, { recursive: true, force: true }); fs.mkdirSync(paths.outdir, { recursive: true }); const result = await esbuild.build({ - entryPoints: { cli: paths.entryPoint, queryIndexWorker: paths.workerEntryPoint }, + entryPoints: { + cli: paths.entryPoint, + queryIndexWorker: paths.workerEntryPoint, + rawQueryWorker: paths.rawQueryWorkerEntryPoint, + }, bundle: true, platform: "node", format: "esm", @@ -59,7 +70,7 @@ export async function bundleCli({ rootDir = defaultRootDir, logLevel = "warning" }); const outputFiles = Object.keys(result.metafile.outputs).sort(); - const selfContainedEntries = new Set([paths.bundledEntry, paths.bundledWorker]); + const selfContainedEntries = new Set([paths.bundledEntry, paths.bundledWorker, paths.bundledRawQueryWorker]); const unexpectedOutputs = outputFiles.filter((file) => !selfContainedEntries.has(path.resolve(file))); if (unexpectedOutputs.length) { throw new Error( @@ -73,6 +84,9 @@ export async function bundleCli({ rootDir = defaultRootDir, logLevel = "warning" if (!fs.existsSync(paths.bundledWorker)) { throw new Error(`Bundled query worker was not written to ${paths.bundledWorker}`); } + if (!fs.existsSync(paths.bundledRawQueryWorker)) { + throw new Error(`Bundled raw SQLite worker was not written to ${paths.bundledRawQueryWorker}`); + } return { ...paths, diff --git a/scripts/ensure-dist-for-tests-lib.mjs b/scripts/ensure-dist-for-tests-lib.mjs index 8d6edc49..f9a02f83 100644 --- a/scripts/ensure-dist-for-tests-lib.mjs +++ b/scripts/ensure-dist-for-tests-lib.mjs @@ -1,7 +1,13 @@ import fs from "node:fs"; import path from "node:path"; -const requiredDistEntries = ["dist/index.js", "dist/cli.js", "dist/bin/cli.js", "dist/bin/queryIndexWorker.js"]; +const requiredDistEntries = [ + "dist/index.js", + "dist/cli.js", + "dist/bin/cli.js", + "dist/bin/queryIndexWorker.js", + "dist/bin/rawQueryWorker.js", +]; const freshnessInputs = [ "package.json", "tsconfig.json", diff --git a/scripts/stage-core-package-lib.mjs b/scripts/stage-core-package-lib.mjs index 3ecc2f97..03a2234c 100644 --- a/scripts/stage-core-package-lib.mjs +++ b/scripts/stage-core-package-lib.mjs @@ -10,7 +10,10 @@ export const CORE_PACKAGE_ENTRIES = Object.freeze([ "languages.js", ]); -export const CORE_PACKAGE_EXTRA_FILES = Object.freeze(["agent/query-index/queryIndexWorker.js"]); +export const CORE_PACKAGE_EXTRA_FILES = Object.freeze([ + "agent/query-index/queryIndexWorker.js", + "sqlite/rawQueryWorker.js", +]); const IMPORT_PATTERN = /(?:import|export)\s+(?:type\s+)?(?:[^;]*?\s+from\s+)?["'](\.[^"']+)["']|import\(["'](\.[^"']+)["']\)/g; diff --git a/src/mcp/server.ts b/src/mcp/server.ts index 012828ce..4af564b4 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -172,6 +172,10 @@ type OriginValidator = (request: IncomingMessage, response: ServerResponse) => b export type CodegraphMcpFreshResult = T & { freshness: AgentFreshnessResult }; +type McpToolExecutionOptions = { + signal?: AbortSignal | undefined; +}; + /** * Truncation metadata for a capped collection response, per finding #44: * lets a machine caller tell a complete result apart from a capped prefix. @@ -321,11 +325,14 @@ export type CodegraphMcpHandlers = { refreshed: true; warmup: CodegraphMcpWarmupMode; }>; - query_sqlite: (request: { - query: string; - params?: Array | undefined; - limit?: number | undefined; - }) => Promise>; + query_sqlite: ( + request: { + query: string; + params?: Array | undefined; + limit?: number | undefined; + }, + options?: McpToolExecutionOptions, + ) => Promise>; 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 { +async function callMcpTool( + handlers: CodegraphMcpHandlers, + name: string, + input: unknown, + signal?: AbortSignal, +): Promise { 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": diff --git a/src/sqlite/query.ts b/src/sqlite/query.ts index aecce786..b477c026 100644 --- a/src/sqlite/query.ts +++ b/src/sqlite/query.ts @@ -11,11 +11,12 @@ import { import { resolveRawSqlQueryWorkerPath, runRawSqlQueryInWorker, + SqliteQueryCancelledError, SqliteQueryDeadlineExceededError, } from "./rawQueryWorkerPool.js"; export { queryGraphSqlite } from "./canned-query.js"; -export { SqliteQueryDeadlineExceededError }; +export { SqliteQueryCancelledError, SqliteQueryDeadlineExceededError }; /** Hard wall-clock budget for a single raw `query_sqlite` execution. */ export const DEFAULT_SQLITE_QUERY_DEADLINE_MS = 10_000; @@ -25,6 +26,7 @@ export type QueryGraphSqliteRawOptions = { maxBytes?: number | undefined; maxCellBytes?: number | undefined; deadlineMs?: number | undefined; + signal?: AbortSignal | undefined; }; /** @@ -65,18 +67,24 @@ export async function queryGraphSqliteRaw( maxBytes, maxCellBytes, deadlineMs, + ...(options?.signal ? { signal: options.signal } : {}), }); } - return await runRawSqlQueryInWorker({ outputPath, sql, params, maxRows, maxBytes, maxCellBytes }, deadlineMs); + return await runRawSqlQueryInWorker( + { outputPath, sql, params, maxRows, maxBytes, maxCellBytes }, + deadlineMs, + options?.signal, + ); } async function queryGraphSqliteRawInProcessBounded( outputPath: string, sql: string, params: Array, - bounds: { maxRows: number; maxBytes: number; maxCellBytes: number; deadlineMs: number }, + bounds: { maxRows: number; maxBytes: number; maxCellBytes: number; deadlineMs: number; signal?: AbortSignal }, ): Promise { + if (bounds.signal?.aborted) throw new SqliteQueryCancelledError(); return await withReadOnlySqliteDatabase(outputPath, (db) => { try { const stmt = db.prepare(sql); @@ -87,6 +95,7 @@ async function queryGraphSqliteRawInProcessBounded( stmt.raw().iterate(params) as Iterable>, deadlineAt, bounds.deadlineMs, + bounds.signal, ); // Always stream via iterate so per-cell and cumulative budgets apply before append. return collectBoundedRawSqlRows(columns, rows, { @@ -106,8 +115,14 @@ async function queryGraphSqliteRawInProcessBounded( /** Throws once the wall-clock deadline has passed between two already-produced rows. See * the fallback caveat on `queryGraphSqliteRaw`: a slow-before-first-row query is not * caught here, only slow-*between*-rows iteration is. */ -function* withPerRowDeadline(rows: Iterable, deadlineAt: number, deadlineMs: number): Generator { +function* withPerRowDeadline( + rows: Iterable, + deadlineAt: number, + deadlineMs: number, + signal: AbortSignal | undefined, +): Generator { for (const row of rows) { + if (signal?.aborted) throw new SqliteQueryCancelledError(); if (Date.now() > deadlineAt) { throw new SqliteQueryDeadlineExceededError(deadlineMs); } diff --git a/src/sqlite/rawQueryWorkerPool.ts b/src/sqlite/rawQueryWorkerPool.ts index 3af4594b..aae0711e 100644 --- a/src/sqlite/rawQueryWorkerPool.ts +++ b/src/sqlite/rawQueryWorkerPool.ts @@ -6,6 +6,20 @@ import { findPackageRoot } from "../util/packageInfo.js"; import type { RawSqlResult } from "./types.js"; import type { RawQueryWorkerTask } from "./rawQueryWorker.js"; +export const MAX_RAW_SQL_QUERY_WORKERS = 2; + +export type RawSqlQueryWorkerPool = { + run(task: RawQueryWorkerTask, options: { signal: AbortSignal }): Promise; + destroy(): Promise; +}; + +type RawSqlQueryWorkerPoolFactory = () => RawSqlQueryWorkerPool; + +export type RawSqlQueryWorkerLifecycleState = { + activeWorkers: number; + maxWorkers: number; +}; + export class SqliteQueryDeadlineExceededError extends Error { constructor(deadlineMs: number) { super(`SQLite query exceeded its ${deadlineMs}ms execution budget and was terminated.`); @@ -13,6 +27,96 @@ export class SqliteQueryDeadlineExceededError extends Error { } } +export class SqliteQueryCancelledError extends Error { + constructor() { + super("SQLite query was cancelled by the MCP client."); + this.name = "SqliteQueryCancelledError"; + } +} + +export class SqliteQueryWorkerCleanupCapacityExceededError extends Error { + constructor(maxWorkers: number) { + super( + `SQLite query cleanup capacity is exhausted: ${maxWorkers} terminated worker${maxWorkers === 1 ? " is" : "s are"} still exiting. Retry after cleanup completes.`, + ); + this.name = "SqliteQueryWorkerCleanupCapacityExceededError"; + } +} + +/** + * Keeps raw-query worker cleanup bounded. A timed-out native SQLite call can delay + * `Piscina.destroy()` until its current synchronous step returns, so each slot remains + * reserved until that destroy promise settles instead of being forgotten in the background. + */ +export class RawSqlQueryWorkerLifecycle { + private readonly activeWorkerSlots = new Set(); + + constructor(private readonly maxWorkers = MAX_RAW_SQL_QUERY_WORKERS) {} + + state(): RawSqlQueryWorkerLifecycleState { + return { activeWorkers: this.activeWorkerSlots.size, maxWorkers: this.maxWorkers }; + } + + async run( + task: RawQueryWorkerTask, + deadlineMs: number, + signal: AbortSignal | undefined, + createPool: RawSqlQueryWorkerPoolFactory, + ): Promise { + if (signal?.aborted) throw new SqliteQueryCancelledError(); + if (this.activeWorkerSlots.size >= this.maxWorkers) { + throw new SqliteQueryWorkerCleanupCapacityExceededError(this.maxWorkers); + } + + const slot = Symbol("raw-sql-query-worker"); + this.activeWorkerSlots.add(slot); + let pool: RawSqlQueryWorkerPool | undefined; + let cleanupInBackground = false; + const deadlineSignal = AbortSignal.timeout(deadlineMs); + const combinedSignal = signal ? AbortSignal.any([signal, deadlineSignal]) : deadlineSignal; + + try { + pool = createPool(); + return await pool.run(task, { signal: combinedSignal }); + } catch (error) { + if (combinedSignal.aborted || (error instanceof Error && error.name === "AbortError")) { + cleanupInBackground = true; + if (signal?.aborted) throw new SqliteQueryCancelledError(); + throw new SqliteQueryDeadlineExceededError(deadlineMs); + } + throw error; + } finally { + if (!pool) { + this.activeWorkerSlots.delete(slot); + } else { + const cleanup = pool.destroy(); + if (cleanupInBackground) { + void cleanup.then( + () => { + this.activeWorkerSlots.delete(slot); + }, + () => { + this.activeWorkerSlots.delete(slot); + }, + ); + } else { + try { + await cleanup; + } finally { + this.activeWorkerSlots.delete(slot); + } + } + } + } + } +} + +const rawSqlQueryWorkerLifecycle = new RawSqlQueryWorkerLifecycle(); + +export function getRawSqlQueryWorkerLifecycleState(): RawSqlQueryWorkerLifecycleState { + return rawSqlQueryWorkerLifecycle.state(); +} + /** Resolves the compiled worker entry the same way `queryIndexWorker.js` is resolved: * a compiled sibling next to this module (production/standalone layouts, where the * whole `dist/` tree ships), falling back to the package-root-relative compiled path @@ -24,7 +128,9 @@ export function resolveRawSqlQueryWorkerPath(): string { const packageRoot = findPackageRoot(selfDir); const compiled = path.join(packageRoot, "dist", "sqlite", "rawQueryWorker.js"); if (fs.existsSync(compiled)) return compiled; - throw new Error(`Raw SQLite query worker file not found: ${compiled}`); + const bundled = path.join(packageRoot, "dist", "bin", "rawQueryWorker.js"); + if (fs.existsSync(bundled)) return bundled; + throw new Error(`Raw SQLite query worker file not found: ${bundled}`); } /** @@ -44,8 +150,9 @@ export function resolveRawSqlQueryWorkerPath(): string { * `sqlite3_step()` — a recursive CTE, or a plan that must fully sort/scan before it can * produce a first row — keeps running on the orphaned worker thread in the background * until that native call returns naturally; only then does the thread actually exit. - * We do not make the caller wait for that: `pool.destroy()` is fired and forgotten here, - * not awaited, so a subsequent query (in its own fresh pool) is never delayed by it. + * The caller does not wait for that cleanup, but the lifecycle retains its worker slot + * until `pool.destroy()` settles. The bounded slot count makes delayed cleanup observable + * and prevents repeated cancellation from accumulating an unbounded number of workers. * Concurrent read-only SQLite connections against the same file do not block each other, * so the lingering background reader does not stop that subsequent query from succeeding. * Verified directly: an aborted 200M-row recursive-CTE count rejects this call in @@ -56,22 +163,22 @@ export function resolveRawSqlQueryWorkerPath(): string { * delayed until that native call returns — a platform limit of `worker_threads`, not of * this module, and orthogonal to the per-call deadline this function guarantees. */ -export async function runRawSqlQueryInWorker(task: RawQueryWorkerTask, deadlineMs: number): Promise { +export async function runRawSqlQueryInWorker( + task: RawQueryWorkerTask, + deadlineMs: number, + signal?: AbortSignal, +): Promise { const workerPath = resolveRawSqlQueryWorkerPath(); - const pool = new Piscina({ - filename: workerPath, - minThreads: 1, - maxThreads: 1, - idleTimeout: 5_000, - }); - try { - return (await pool.run(task, { signal: AbortSignal.timeout(deadlineMs) })) as RawSqlResult; - } catch (error) { - if (error instanceof Error && error.name === "AbortError") { - throw new SqliteQueryDeadlineExceededError(deadlineMs); - } - throw error; - } finally { - void pool.destroy().catch(() => {}); - } + return await rawSqlQueryWorkerLifecycle.run( + task, + deadlineMs, + signal, + () => + new Piscina({ + filename: workerPath, + minThreads: 1, + maxThreads: 1, + idleTimeout: 5_000, + }), + ); } diff --git a/tests/cli-bundle-entry.test.ts b/tests/cli-bundle-entry.test.ts index fa4a3a4e..4d61446c 100644 --- a/tests/cli-bundle-entry.test.ts +++ b/tests/cli-bundle-entry.test.ts @@ -7,6 +7,7 @@ import { fileURLToPath } from "node:url"; const rootDir = path.resolve(path.dirname(fileURLToPath(import.meta.url)), ".."); const bundledCli = path.join(rootDir, "dist", "bin", "cli.js"); +const bundledRawQueryWorker = path.join(rootDir, "dist", "bin", "rawQueryWorker.js"); const unbundledCli = path.join(rootDir, "dist", "cli.js"); function run(entry: string, args: string[], cwd: string = rootDir, env: NodeJS.ProcessEnv = process.env) { @@ -21,6 +22,7 @@ function run(entry: string, args: string[], cwd: string = rootDir, env: NodeJS.P describe("bundled CLI entry", () => { it("ships a split ESM bin entry that matches unbundled --version", () => { expect(fs.existsSync(bundledCli)).toBe(true); + expect(fs.existsSync(bundledRawQueryWorker)).toBe(true); expect(fs.existsSync(unbundledCli)).toBe(true); const bundled = run(bundledCli, ["--version"]); diff --git a/tests/core-package-surface.test.ts b/tests/core-package-surface.test.ts index f7e64748..def36129 100644 --- a/tests/core-package-surface.test.ts +++ b/tests/core-package-surface.test.ts @@ -20,6 +20,7 @@ describe("codegraph-core package surface", () => { expect(files.length).toBeGreaterThan(100); expect(files.some(isForbiddenCorePackagePath)).toBe(false); expect(files).toContain("agent/query-index/queryIndexWorker.js"); + expect(files).toContain("sqlite/rawQueryWorker.js"); expect(files).toContain("graphs/types.d.ts"); expect(files).toContain("agent/semantic.d.ts"); expect(files).toContain("chunking/types.d.ts"); diff --git a/tests/ensure-dist-for-tests.test.ts b/tests/ensure-dist-for-tests.test.ts index a12f3ac9..7b2cdc2c 100644 --- a/tests/ensure-dist-for-tests.test.ts +++ b/tests/ensure-dist-for-tests.test.ts @@ -53,6 +53,7 @@ describe("inspectDistForTests", () => { await fsp.mkdir(path.join(root, "dist", "bin"), { recursive: true }); await setFileMtime(path.join(root, "dist", "bin", "cli.js"), distTime, "export {};\n"); await setFileMtime(path.join(root, "dist", "bin", "queryIndexWorker.js"), distTime, "export {};\n"); + await setFileMtime(path.join(root, "dist", "bin", "rawQueryWorker.js"), distTime, "export {};\n"); await setFileMtime(path.join(root, "src", "index.ts"), srcTime); expect(inspectDistForTests(root)).toMatchObject({ @@ -77,6 +78,7 @@ describe("inspectDistForTests", () => { await fsp.mkdir(path.join(root, "dist", "bin"), { recursive: true }); await setFileMtime(path.join(root, "dist", "bin", "cli.js"), distTime, "export {};\n"); await setFileMtime(path.join(root, "dist", "bin", "queryIndexWorker.js"), distTime, "export {};\n"); + await setFileMtime(path.join(root, "dist", "bin", "rawQueryWorker.js"), distTime, "export {};\n"); expect(inspectDistForTests(root)).toMatchObject({ needsBuild: false, @@ -101,6 +103,7 @@ describe("inspectDistForTests", () => { await fsp.mkdir(path.join(root, "dist", "bin"), { recursive: true }); await setFileMtime(path.join(root, "dist", "bin", "cli.js"), distTime, "export {};\n"); await setFileMtime(path.join(root, "dist", "bin", "queryIndexWorker.js"), distTime, "export {};\n"); + await setFileMtime(path.join(root, "dist", "bin", "rawQueryWorker.js"), distTime, "export {};\n"); expect(inspectDistForTests(root)).toMatchObject({ needsBuild: false, reason: "fresh" }); @@ -130,6 +133,7 @@ describe("inspectDistForTests", () => { await fsp.mkdir(path.join(root, "dist", "bin"), { recursive: true }); await setFileMtime(path.join(root, "dist", "bin", "cli.js"), distTime, "export {};\n"); await setFileMtime(path.join(root, "dist", "bin", "queryIndexWorker.js"), distTime, "export {};\n"); + await setFileMtime(path.join(root, "dist", "bin", "rawQueryWorker.js"), distTime, "export {};\n"); expect(inspectDistForTests(root)).toMatchObject({ needsBuild: false, reason: "fresh" }); diff --git a/tests/mcp-stream-cancellation.test.ts b/tests/mcp-stream-cancellation.test.ts new file mode 100644 index 00000000..f90eb1c8 --- /dev/null +++ b/tests/mcp-stream-cancellation.test.ts @@ -0,0 +1,51 @@ +import fsp from "node:fs/promises"; +import os from "node:os"; +import path from "node:path"; +import { Client, InMemoryTransport } from "@modelcontextprotocol/client"; +import { describe, expect, it } from "vitest"; +import { createCodegraphMcpHandlers, createCodegraphMcpProtocolServer } from "../src/mcp/server.js"; + +describe("MCP query_sqlite cancellation", () => { + it("forwards a cancelled tool stream to the raw query handler", async () => { + const root = await fsp.mkdtemp(path.join(os.tmpdir(), "cg-mcp-query-cancel-")); + const handlers = createCodegraphMcpHandlers({ root }); + const started = Promise.withResolvers(); + const cancelled = Promise.withResolvers(); + handlers.query_sqlite = async (_request, executionOptions) => { + const signal = executionOptions?.signal; + if (!signal) throw new Error("MCP query_sqlite did not receive a cancellation signal."); + started.resolve(); + await new Promise((_resolve, reject) => { + signal.addEventListener( + "abort", + () => { + cancelled.resolve(); + reject(new Error("raw query cancelled")); + }, + { once: true }, + ); + }); + }; + + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + const server = createCodegraphMcpProtocolServer(handlers); + const client = new Client({ name: "mcp-stream-cancellation-test", version: "1.0.0" }); + const controller = new AbortController(); + try { + await server.connect(serverTransport); + await client.connect(clientTransport); + const call = client.callTool( + { name: "query_sqlite", arguments: { query: "SELECT 1;" } }, + { signal: controller.signal }, + ); + await started.promise; + controller.abort(); + + await expect(call).rejects.toThrow(); + await cancelled.promise; + } finally { + await Promise.allSettled([client.close(), server.close()]); + await fsp.rm(root, { recursive: true, force: true }); + } + }); +}); diff --git a/tests/query-index-worker-path.test.ts b/tests/query-index-worker-path.test.ts index 0a64da40..495133e7 100644 --- a/tests/query-index-worker-path.test.ts +++ b/tests/query-index-worker-path.test.ts @@ -2,11 +2,24 @@ import fs from "node:fs"; import path from "node:path"; import { afterEach, describe, expect, it, vi } from "vitest"; import { resolveQueryIndexWorkerPath } from "../src/agent/query-index/workerPool.js"; +import { resolveRawSqlQueryWorkerPath } from "../src/sqlite/rawQueryWorkerPool.js"; afterEach(() => { vi.restoreAllMocks(); }); +describe("resolveRawSqlQueryWorkerPath", () => { + it("falls back to the bundled worker when the compiled dist worker is missing", () => { + const bundledSuffix = path.normalize(path.join("dist", "bin", "rawQueryWorker.js")); + vi.spyOn(fs, "existsSync").mockImplementation((candidate) => { + const filePath = path.normalize(typeof candidate === "string" ? candidate : String(candidate)); + return filePath.endsWith(bundledSuffix); + }); + + expect(path.normalize(resolveRawSqlQueryWorkerPath())).toContain(bundledSuffix); + }); +}); + describe("resolveQueryIndexWorkerPath", () => { it("falls back to the bundled worker when the compiled dist worker is missing", () => { const bundledSuffix = path.normalize(path.join("dist", "bin", "queryIndexWorker.js")); diff --git a/tests/raw-query-worker-lifecycle.test.ts b/tests/raw-query-worker-lifecycle.test.ts new file mode 100644 index 00000000..cbbe8d26 --- /dev/null +++ b/tests/raw-query-worker-lifecycle.test.ts @@ -0,0 +1,63 @@ +import { describe, expect, it } from "vitest"; +import { + RawSqlQueryWorkerLifecycle, + SqliteQueryCancelledError, + type RawSqlQueryWorkerPool, +} from "../src/sqlite/rawQueryWorkerPool.js"; +import type { RawQueryWorkerTask } from "../src/sqlite/rawQueryWorker.js"; + +const task: RawQueryWorkerTask = { + outputPath: "fixture.sqlite", + sql: "SELECT 1;", + params: [], + maxRows: 1, + maxBytes: 1024, + maxCellBytes: 1024, +}; + +function createAbortablePool(cleanup: Promise): RawSqlQueryWorkerPool { + return { + run: async (_task, options) => + await new Promise((_, reject) => { + options.signal?.addEventListener( + "abort", + () => { + const error = new Error("aborted"); + error.name = "AbortError"; + reject(error); + }, + { once: true }, + ); + }), + destroy: async () => await cleanup, + }; +} + +describe("RawSqlQueryWorkerLifecycle", () => { + it("caps cancelled worker cleanup slots until each worker has actually exited", async () => { + const lifecycle = new RawSqlQueryWorkerLifecycle(2); + const firstCleanup = Promise.withResolvers(); + const secondCleanup = Promise.withResolvers(); + const firstAbort = new AbortController(); + const secondAbort = new AbortController(); + + const first = lifecycle.run(task, 10_000, firstAbort.signal, () => createAbortablePool(firstCleanup.promise)); + firstAbort.abort(); + await expect(first).rejects.toBeInstanceOf(SqliteQueryCancelledError); + + const second = lifecycle.run(task, 10_000, secondAbort.signal, () => createAbortablePool(secondCleanup.promise)); + secondAbort.abort(); + await expect(second).rejects.toBeInstanceOf(SqliteQueryCancelledError); + + expect(lifecycle.state()).toEqual({ activeWorkers: 2, maxWorkers: 2 }); + await expect(lifecycle.run(task, 10_000, undefined, () => createAbortablePool(Promise.resolve()))).rejects.toThrow( + /cleanup capacity/i, + ); + + firstCleanup.resolve(); + secondCleanup.resolve(); + await Promise.resolve(); + await Promise.resolve(); + expect(lifecycle.state()).toEqual({ activeWorkers: 0, maxWorkers: 2 }); + }); +});