fix: bound MCP session and query resources - #264
Conversation
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.
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.
Three post-review defects in the MCP HTTP transport and raw SQLite query path: 1. Legacy initialize capacity reservations leaked whenever the SDK transport answered a pre-session 4xx without throwing (most notably Accept header validation): onsessioninitialized never fired to release it, and the catch block only covered thrown errors. Release the reservation whenever handleLegacyMcpSessionRequest resolves without a session having been initialized. 2. transport.onerror deleted the whole legacy session for any per-request SDK validation error (bad Accept, wrong Content-Type, malformed JSON, unsupported protocol version, ...), even though those already answered their own request and left the transport healthy. Session teardown is now driven by onclose alone; onerror only logs. 3. The in-process SQLite query fallback (used when the compiled worker asset can't be located) only checks its deadline between already-produced rows, so a statement slow to produce its first row isn't bounded by it. node:sqlite has no interrupt API, so true enforcement requires the worker thread this fallback exists because it couldn't find; corrected the public contract instead (JSDoc, docs/library-api.md, and a one-time degraded-mode log) so the gap is documented and observable rather than silently implied away.
…source-safety-consolidated # Conflicts: # tests/mcp-server.test.ts
…n-resource-safety-consolidated # Conflicts: # src/mcp/server.ts # src/sqlite/query.ts # src/sqlite/rawQueryWorker.ts # src/sqlite/rawQueryWorkerPool.ts # tests/sqlite-query-bounds.test.ts
There was a problem hiding this comment.
Pull request overview
This PR hardens the MCP server and related query/impact machinery by adding explicit resource bounds (HTTP body reads, tool concurrency, sessions, raw SQLite query deadlines, streaming buffers) and by tightening lifecycle teardown and dispatch behavior so long-running servers don’t leak capacity or background work.
Changes:
- Add worker-thread execution + deadline enforcement for raw SQLite queries (with a documented in-process fallback when the worker asset is missing).
- Introduce MCP tool-registry–driven dispatch, per-session tool concurrency limiting, HTTP body timeouts, and improved refresh/session teardown behavior.
- Bound streamed impact output (queue cap + overflow error) and extend test coverage for cancellation, omission-count bounds, retries, and packaging/bundling.
Reviewed changes
Copilot reviewed 40 out of 40 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/workspace-symbols.test.ts | Adds assertions for omitted-count behavior at/over symbol limits. |
| tests/viewer.test.ts | Verifies viewer continues serving after stat failures. |
| tests/type-hierarchy.test.ts | Adds omitted-count pinning for hierarchy + implementations. |
| tests/sqlite-query-deadline-fallback.test.ts | New tests for in-process raw-query deadline fallback semantics. |
| tests/sqlite-query-bounds.test.ts | Adds deadline tests + Windows cleanup retry for temp DB teardown. |
| tests/session.test.ts | Adds SessionManager capacity test + impact-stream cancellation test. |
| tests/raw-query-worker-lifecycle.test.ts | New test for bounded worker cleanup capacity behavior. |
| tests/query-index.test.ts | Adds bounded retry test for query-index generation under invalidation. |
| tests/query-index-worker-path.test.ts | Ensures raw SQL worker path resolution fallback behaves as expected. |
| tests/mcp-stream-cancellation.test.ts | New test proving MCP tool-call cancellation signal propagation. |
| tests/mcp-server.test.ts | Expands MCP regression coverage (timeouts, dispatch, refresh serialization, teardown, cancellation accounting, transport isolation, capacity). |
| tests/impact-streaming.test.ts | Adds overflow + cancellation + completion coverage for impact streaming bounds. |
| tests/ensure-dist-for-tests.test.ts | Ensures rawQueryWorker is included in dist test fixtures. |
| tests/core-package-surface.test.ts | Verifies core package surface includes rawQueryWorker asset. |
| tests/cli-bundle-entry.test.ts | Verifies bundled rawQueryWorker exists alongside bundled CLI. |
| tests/agent-session.test.ts | Adds query-index generation retry bounds regression test. |
| tests/agent-search.test.ts | Tests coalescing + LRU-like eviction for session search cache. |
| tests/agent-explore.test.ts | Pins omission-count behavior for explore candidate tests + blast radius. |
| src/sqlite/rawQueryWorkerPool.ts | New worker pool + lifecycle for bounded raw SQL queries with deadlines/cancellation. |
| src/sqlite/rawQueryWorker.ts | New Piscina worker entry for read-only bounded raw SQLite execution. |
| src/sqlite/query.ts | Adds deadlineMs/signal support; uses worker path with degraded in-process fallback. |
| src/session.ts | Adds SessionManager capacity + configurable eviction interval plumbing. |
| src/mcp/tools.ts | Introduces MCP_TOOL_REGISTRY with dispatch metadata and hides legacy aliases from tools/list. |
| src/mcp/server.ts | Adds bounded HTTP body reads, tool concurrency control, registry dispatch, refresh serialization, and improved teardown. |
| src/mcp/http.ts | Implements streaming body read with max-bytes + timeout behavior. |
| src/indexer/workspace-symbols.ts | Uses shared boundList helper for consistent omitted counting. |
| src/indexer/type-hierarchy.ts | Uses boundList for hierarchy/implementations omission calculations. |
| src/impact/streaming.ts | Adds queued-chunk cap + explicit overflow error + abandonment cancellation wiring. |
| src/cli/viewer.ts | Improves viewer request error containment to keep server alive after sync fs errors. |
| src/agent/session.ts | Centralizes “no prebuilt session with buildOptions” assertion. |
| src/agent/search.ts | Adds bounded per-session result cache with eviction + promotion on hit. |
| src/agent/query-index/sessionStore.ts | Adds bounded retries under invalidation while loading query index. |
| src/agent/explore.ts | Uses boundList for consistent omitted counts in explore outputs. |
| src/agent-tools.ts | Reuses centralized session/buildOptions assertion helper. |
| scripts/stage-core-package-lib.mjs | Stages rawQueryWorker into core package extra files. |
| scripts/ensure-dist-for-tests-lib.mjs | Requires rawQueryWorker in dist for tests. |
| scripts/bundle-cli-lib.mjs | Bundles rawQueryWorker into dist/bin alongside existing entries. |
| docs/mcp.md | Documents tool concurrency, body timeout, refresh serialization, cancellation semantics, and query deadline. |
| docs/library-api.md | Documents new MCP server options + raw SQLite query deadline behavior and fallback caveat. |
| codegraph-skill/codegraph/SKILL.md | Updates skill docs to match new MCP limits/concurrency/timeout behavior. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| while (true) { | ||
| if (refreshPromise) await refreshPromise; | ||
| const epoch = refreshEpoch; | ||
| const freshness = await checkMcpFreshness(); | ||
| const result = await run(); |
| const handle = ensureQueryIndex(snapshot); | ||
| const state: SessionQueryIndexState = { identity, handle }; | ||
| if (!existing) registerSessionInvalidationHook(session, () => disposeSessionQueryIndex(session)); | ||
| QUERY_INDEX_BY_SESSION.set(session, state); |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 40 out of 40 changed files in this pull request and generated 1 comment.
Suppressed comments (1)
src/mcp/server.ts:1335
- On HTTP body timeout, the handler returns 408 but does not close the underlying connection. If the client never finishes sending the declared body, the socket can remain open indefinitely, which is a common slowloris-style resource exhaustion vector. Consider closing the connection for the timeout case (set
Connection: closeand destroy the request after the response flushes).
if (parsedBody.status === "timeout") {
writeJsonRpcError(response, 408, "MCP request body timed out");
return;
}
| if (parsedBody.status === "too_large") { | ||
| writeJsonRpcError(response, 413, "MCP request body is too large"); | ||
| return; | ||
| } |
| function normalizeSessionManagerCapacity(value: number | undefined): number { | ||
| if (value === undefined) return DEFAULT_SESSION_MANAGER_MAX_SESSIONS; | ||
| return Math.max(1, Math.floor(value)); | ||
| } | ||
|
|
||
| function normalizeSessionManagerEvictionInterval(value: number | undefined): number { | ||
| if (value === undefined) return DEFAULT_SESSION_MANAGER_EVICTION_INTERVAL_MS; | ||
| return Math.max(0, Math.floor(value)); | ||
| } |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 41 out of 41 changed files in this pull request and generated 2 comments.
Suppressed comments (19)
src/sqlite/query.ts:44
- Replace the em dash with standard ASCII punctuation to comply with the repository text-character rule.
* partial install — the normal build/publish/standalone pipelines all ship it), the
src/sqlite/query.ts:47
- Replace the em dash with standard ASCII punctuation to comply with the repository text-character rule.
* single synchronous native call there is nothing in-process that can preempt it —
src/sqlite/query.ts:96
- Replace the em dash with standard ASCII punctuation to comply with the repository text-character rule.
/** Degraded fallback for `queryGraphSqliteRaw` — see its doc comment for the enforcement
src/impact/streaming.ts:239
- Replace the em dashes with standard ASCII punctuation to comply with the repository text-character rule.
* Cancellation: if the consumer stops iterating early — a `for await` `break`, or an
* explicit `.return()` on the generator — the async-generator return protocol resumes
src/impact/streaming.ts:423
- Replace the em dashes with standard ASCII punctuation to comply with the repository text-character rule.
// 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
src/sqlite/rawQueryWorkerPool.ts:143
- Replace the em dash with standard ASCII punctuation to comply with the repository text-character rule.
* (`worker.terminate()`) and rejects immediately — the caller never waits longer than
src/sqlite/rawQueryWorkerPool.ts:150
- Replace the em dashes with standard ASCII punctuation to comply with the repository text-character rule.
* `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
src/sqlite/rawQueryWorkerPool.ts:162
- Replace the em dash with standard ASCII punctuation to comply with the repository text-character rule.
* delayed until that native call returns — a platform limit of `worker_threads`, not of
src/mcp/server.ts:1096
- Normalize the public concurrency option before using it as a bound. Passing
NaNorInfinitymakesinFlightToolCalls >= maxConcurrentToolCallspermanently false, allowing unbounded tool work; fractional or non-positive values also produce surprising capacity behavior.
maxConcurrentToolCalls = DEFAULT_MCP_TOOL_CONCURRENCY,
src/sqlite/query.ts:148
- The fallback never checks the deadline when an expensive iterator step completes with
done: true. A slow query that returns zero rows can therefore exceeddeadlineMsand still resolve successfully; check after everyiterator.next(), including the terminal step.
for (const row of rows) {
if (signal?.aborted) throw new SqliteQueryCancelledError();
if (Date.now() > deadlineAt) {
throw new SqliteQueryDeadlineExceededError(deadlineMs);
}
yield row;
src/sqlite/query.ts:19
- These named errors are exported only from the internal
sqlite/query.jsmodule. The publicsrc/sqlite.tsand rootsrc/index.tsbarrels re-exportqueryGraphSqliteRawbut not these classes, so package consumers cannot use the named exports that the new tests describe as caller-facing.
export { SqliteQueryCancelledError, SqliteQueryDeadlineExceededError };
docs/library-api.md:726
- This public API documentation overstates worker termination. The worker lifecycle explicitly retains a slot because an in-flight synchronous
sqlite3_step()can continue after the caller's deadline; describe the prompt caller rejection and bounded background cleanup instead of claiming the query is interrupted mid-execution.
`queryGraphSqliteRaw()` is intentionally read-only. It accepts result-producing statements such as `SELECT` and `PRAGMA` and rejects mutating SQL. Its defaults bound rows, cells, and response bytes, and callers can further tighten `{ maxRows, maxBytes, maxCellBytes, deadlineMs }`. The 10-second default execution budget (`deadlineMs`) is enforced by running the query in a dedicated worker thread that is force-terminated on expiry, so it interrupts a query even mid-execution; in a degraded install where that worker asset cannot be located, the query instead runs in-process under a weaker per-row check that cannot interrupt a single blocking native call (a logged, one-time-per-process condition).
src/sqlite/query.ts:41
- This doc comment contradicts the lifecycle implementation below:
worker.terminate()cannot preempt an already-running synchronous native SQLite step. Document that the caller is released at the deadline while the native step may continue in the bounded cleanup slot.
* 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 slow non-recursive statement (large join,
* `ORDER BY random()`, a recursive CTE, ...) cannot hold the deadline hostage.
src/sqlite/rawQueryWorker.ts:25
- Thread termination does require cooperation from an in-flight native SQLite call, as documented in
rawQueryWorkerPool.tsand exercised by the cleanup tests. This comment should distinguish prompt caller cancellation from the native step that can continue until it returns.
* 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.
src/sqlite/rawQueryWorkerPool.ts:39
- This capacity error is also thrown when two ordinary queries are currently active, not only when terminated workers are exiting. The current message falsely diagnoses cleanup as the cause; make it describe both active and cleaning-up worker slots.
`SQLite query cleanup capacity is exhausted: ${maxWorkers} terminated worker${maxWorkers === 1 ? " is" : "s are"} still exiting. Retry after cleanup completes.`,
src/sqlite/query.ts:21
- Replace the em dash with standard ASCII punctuation to comply with the repository text-character rule.
This issue also appears in the following locations of the same file:
- line 44
- line 47
- line 96
/** Hard wall-clock budget for a single raw `query_sqlite` execution — see the caveat on
src/impact/streaming.ts:109
- Replace the em dash with standard ASCII punctuation to comply with the repository text-character rule.
This issue also appears in the following locations of the same file:
- line 238
- line 422
* awaiting it at every synchronous call site in `direct.ts`/`transitive.ts` — an invasive
src/sqlite/rawQueryWorkerPool.ts:138
- Replace the em dashes with standard ASCII punctuation to comply with the repository text-character rule.
This issue also appears in the following locations of the same file:
- line 143
- line 149
- line 162
* deadline. A fresh single-thread pool is created per call and destroyed afterward —
* matching the existing `prepareQueryIndexFilesInWorker` pattern — since `query_sqlite`
src/mcp/server.ts:1444
- Replace the em dash with standard ASCII punctuation to comply with the repository text-character rule.
// protocol version, ...) — each of those already answered its own request with a
| private assertCapacityForNewSession(): void { | ||
| this.cleanupExpired(); | ||
| if (this.sessions.size + this.pendingSessions.size >= this.maxSessions) { |
| if (existing) { | ||
| QUERY_INDEX_BY_SESSION.delete(session); | ||
| closeState(existing); | ||
| } |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 44 out of 45 changed files in this pull request and generated 2 comments.
Suppressed comments (4)
src/impact/streaming.ts:344
- Cancellation is observed only when
analyzeImpactemits another item. If the remaining batches find no impacts, or analysis is in work that emits nothing, this callback is never invoked and the abandoned analysis continues to completion while retaining the queue and other state. Pass the abort signal into the analyzer and check it at batch and transitive-work boundaries rather than relying on output events as cancellation checkpoints.
onImpactItem: (item, phase) => {
if (abortController.signal.aborted) {
throw new ImpactStreamAbandonedError();
}
queueImpactItem(item, phase === "partial");
src/sqlite/rawQueryWorkerPool.ts:34
queryGraphSqliteRaw()now exposes cancellation to general library callers, so attributing every abort to an MCP client is inaccurate and misleading outside MCP. Use caller-neutral wording.
export class SqliteQueryCancelledError extends Error {
constructor() {
super("SQLite query was cancelled by the MCP client.");
this.name = "SqliteQueryCancelledError";
}
src/sqlite/rawQueryWorkerPool.ts:42
- This error is part of the observable failure surface of the public
queryGraphSqliteRaw()API, but unlike the new cancellation and deadline errors it is not re-exported fromsqlite/query.ts,sqlite.ts, orindex.ts. Library callers therefore cannot reliably distinguish retryable worker saturation without unsupported deep imports or message parsing; re-export this named error through the public barrels and document it.
export class SqliteQueryWorkerCleanupCapacityExceededError extends Error {
constructor(maxWorkers: number) {
super(
`SQLite query worker capacity is exhausted: ${maxWorkers} active or cleaning-up worker${maxWorkers === 1 ? " is" : "s are"} using the available slots. Retry after a query completes or cleanup finishes.`,
);
this.name = "SqliteQueryWorkerCleanupCapacityExceededError";
}
src/impact/streaming.ts:113
- This cap changes the public
analyzeImpactStreaming()contract: a slow consumer can now receive anerrorchunk without a finalcompletechunk. The canonical streaming docs still state that the API emits progress/incremental chunks and thencomplete(docs/library-api.md:939) and do not describe overflow or early-consumer cancellation (docs/agent-workflows.md:313). Update those docs so callers know to handle this terminal failure mode.
/**
* 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;
| private assertCapacityForNewSession(): void { | ||
| this.cleanupExpired(); | ||
| if (this.sessions.size + this.pendingSessions.size >= this.maxSessions) { | ||
| throw new Error( | ||
| `Session capacity reached (${this.maxSessions}). Dispose an existing session before creating another.`, | ||
| ); |
| const iterator = rows[Symbol.iterator](); | ||
| while (true) { | ||
| if (signal?.aborted) throw new SqliteQueryCancelledError(); | ||
| const next = iterator.next(); | ||
| if (Date.now() > deadlineAt) { | ||
| throw new SqliteQueryDeadlineExceededError(deadlineMs); | ||
| } | ||
| if (next.done) return; | ||
| yield next.value; | ||
| } |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 46 out of 47 changed files in this pull request and generated no new comments.
Suppressed comments (3)
src/mcp/http.ts:31
- A request whose
Content-Lengthalready exceeds the cap is no longer rejected immediately. If the client sends only part of that body, this waits until the body deadline and returns 408 instead of the determinable 413, tying up the request for the full timeout. Preserve the previous early rejection and drain the request in the background.
const knownTooLarge = contentLength !== undefined && contentLength > maxBytes;
return await new Promise<ParsedJsonBody>((resolve) => {
const chunks: Buffer[] = [];
let bytes = 0;
let tooLarge = knownTooLarge;
src/session.ts:989
disposeAll()remains reusable (including the replacement session created in the updated tests), but it permanently clears the only eviction timer. Any sessions created afterward will therefore never be evicted periodically. Restart the timer when the manager is reused, or separate final manager disposal from reusabledisposeAll()semantics.
clearInterval(this.evictionTimer);
tests/sqlite-query-bounds.test.ts:42
- The retry is intended to handle Windows locks left by the worker, but Windows can report a locked file as
EPERMas well asEBUSY. In that case cleanup still fails immediately, making the new deadline test flaky on the platform this helper targets.
Summary
Verification
npm run test:coverage- 3,267 passed, 30 skipped.npm run lintnpm run fixtures:check-cleannpm run security:productionnode ./dist/cli.js review --base origin/main --head HEAD --duplicates off --jsonThe fallback deadline regression now uses a correlated scalar subquery so current SQLite cannot hoist the work out of the per-row path.