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/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/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 d1e740ff..b477c026 100644 --- a/src/sqlite/query.ts +++ b/src/sqlite/query.ts @@ -8,35 +8,100 @@ import { MAX_SQLITE_ROW_LIMIT, normalizeSqliteRowLimit, } from "./rowBounds.js"; +import { + resolveRawSqlQueryWorkerPath, + runRawSqlQueryInWorker, + SqliteQueryCancelledError, + SqliteQueryDeadlineExceededError, +} from "./rawQueryWorkerPool.js"; export { queryGraphSqlite } from "./canned-query.js"; +export { SqliteQueryCancelledError, 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; + signal?: AbortSignal | 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, + ...(options?.signal ? { signal: options.signal } : {}), + }); + } + + 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; signal?: AbortSignal }, +): Promise { + if (bounds.signal?.aborted) throw new SqliteQueryCancelledError(); 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, + bounds.signal, + ); // 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 +111,21 @@ 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, + signal: AbortSignal | undefined, +): Generator { + for (const row of rows) { + if (signal?.aborted) throw new SqliteQueryCancelledError(); + 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..aae0711e --- /dev/null +++ b/src/sqlite/rawQueryWorkerPool.ts @@ -0,0 +1,184 @@ +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 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.`); + this.name = "SqliteQueryDeadlineExceededError"; + } +} + +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 + * (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; + 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}`); +} + +/** + * 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. + * 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 + * ~`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, + signal?: AbortSignal, +): Promise { + const workerPath = resolveRawSqlQueryWorkerPath(); + 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/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/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 }); + }); +}); 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 }); + } + }); +}); 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."); + }); +});