diff --git a/codegraph-skill/codegraph/SKILL.md b/codegraph-skill/codegraph/SKILL.md index 3f2c1868..7400b312 100644 --- a/codegraph-skill/codegraph/SKILL.md +++ b/codegraph-skill/codegraph/SKILL.md @@ -164,7 +164,7 @@ Sensitive-file rules: If MCP tools are available, prefer them over repeated CLI invocations. Use `explore`, `orient`, `workspace_symbols`, `search`, `get_file`, `get_symbol`, `packet_get`, `goto`, `refs`, `rename_preview`, `refactor_plan`, `calls`, `type_hierarchy`, `implementations`, `file_deps`, `path`, `impact`, `review`, `query_sqlite`, `refresh_index`, and `artifact_build`; fall back to the CLI when MCP is unavailable. Legacy `callers`/`callees`, `supertypes`/`subtypes`, and `deps`/`rdeps` names remain valid `tools/call` aliases. -codegraph uses the official MCP SDK v2 to serve current 2026-07-28 clients while retaining compatibility with 2025-era clients. MCP protocol connections and HTTP protocol sessions keep separate transport state, but all share the server's one warm codegraph analysis session for the configured root. Tool schemas reject unknown fields, and idle HTTP protocol sessions are evicted with a bounded session count. +codegraph uses the official MCP SDK v2 to serve current 2026-07-28 clients while retaining compatibility with 2025-era clients. MCP protocol connections and HTTP protocol sessions keep separate transport state, but all share the server's one warm codegraph analysis session for the configured root. Tool schemas reject unknown fields, each protocol session caps tool concurrency at 4 with a retryable busy error, HTTP request bodies time out after 30 seconds, and idle HTTP protocol sessions are evicted with a bounded session count. Cancellation responds to the caller promptly but retains the occupied slot until shared work settles, so abandoned calls cannot bypass that resource bound. Concurrent `refresh_index` calls serialize and honor each request's requested warmup. On the first `tools/call`, codegraph can emit `notifications/message` and, when the request includes `_meta.progressToken`, `notifications/progress` before the final result. Stdio carries them inline, and modern Streamable HTTP clients that accept `text/event-stream` receive them as a stream until the terminal result frame. HTTP enforces Host and Origin policies. A missing `Origin` is accepted for non-browser clients; unapproved, malformed, and opaque origins are rejected. This is not authentication: binding `--host` to a non-loopback address exposes an unauthenticated endpoint intended only for trusted networks or containers. diff --git a/docs/agent-workflows.md b/docs/agent-workflows.md index d5bbf869..f026ca1b 100644 --- a/docs/agent-workflows.md +++ b/docs/agent-workflows.md @@ -290,9 +290,10 @@ Available presets: ### Managing multiple sessions ```ts -import { SessionManager } from "@lzehrung/codegraph"; +import { SessionManager, type SessionManagerOptions } from "@lzehrung/codegraph"; -const manager = new SessionManager(); +const options: SessionManagerOptions = { maxSessions: 32 }; +const manager = new SessionManager(options); const pr1Session = await manager.getOrCreateSession("pr-123", { root: "/path/to/repo", }); @@ -306,8 +307,16 @@ const sameSession = await manager.getOrCreateSession("pr-123", { manager.cleanupExpired(); const allStats = manager.getAllStats(); console.log(Boolean(pr1Session), Boolean(pr2Session), Boolean(sameSession), allStats); + +manager.dispose(); ``` +`SessionManager` defaults to 32 live or initializing sessions and scans for expired sessions every 60 seconds. Set `{ maxSessions, evictionIntervalMs }` to tune those bounds; `maxSessions` also applies to net-new `warmup()` sessions. + +- Capacity frees immediately after a ready session is disposed or expires, or after a canceled or failed initialization settles. +- Set `evictionIntervalMs: 0` to disable periodic cleanup. +- Call `manager.dispose()` when the manager is no longer needed; it disposes all sessions, stops the interval, and is terminal. `disposeAll()` remains reusable. + ## Streaming impact analysis Stream impact results as they are discovered so the agent can start reasoning before the full pass completes: @@ -337,6 +346,8 @@ for await (const chunk of analyzeImpactStreaming(root, index, { } ``` +Handle `error` as terminal: an overfull bounded queue does not emit `complete`. Breaking iteration or calling `.return()` cancels background work at its next analysis boundary; a synchronous lookup already in progress cannot be interrupted. + Use the same pattern through a warm session when repeated review passes matter: ```ts diff --git a/docs/coverage/js.md b/docs/coverage/js.md index fcdfef37..1d3d4598 100644 --- a/docs/coverage/js.md +++ b/docs/coverage/js.md @@ -6,20 +6,20 @@ Source: `coverage/js/lcov.info` | Metric | Hit | Found | Coverage | | --------- | ----: | ----: | -------: | -| Lines | 27491 | 30264 | 90.84% | -| Functions | 4565 | 4844 | 94.24% | -| Branches | 20716 | 26123 | 79.30% | +| Lines | 28189 | 31001 | 90.93% | +| Functions | 4695 | 4990 | 94.09% | +| Branches | 21301 | 26826 | 79.40% | ## Least-covered Files | File | Lines | Functions | Branches | | ---------------------------------------------- | -----: | --------: | -------: | +| `src/sqlite/rawQueryWorker.ts` | 0.00% | 0.00% | 0.00% | | `src/languages/definitions/htmlStub.ts` | 50.00% | 50.00% | n/a | | `src/languages/definitions/javascript.ts` | 62.50% | 83.33% | 57.89% | | `src/cli/explore.ts` | 62.50% | 100.00% | 50.00% | | `src/languages/definitions/typescript.ts` | 62.86% | 81.82% | 42.86% | | `src/impact/call-compatibility/textScanner.ts` | 66.04% | 100.00% | 70.75% | -| `src/cli/index.ts` | 66.67% | 100.00% | 58.97% | | `src/indexer.ts` | 66.67% | 66.67% | 100.00% | | `src/cli/bootstrap.ts` | 66.67% | 50.00% | n/a | | `src/languages/definitions/css.ts` | 66.67% | 0.00% | n/a | @@ -28,12 +28,12 @@ Source: `coverage/js/lcov.info` | `src/languages/definitions/sql.ts` | 66.67% | 0.00% | n/a | | `src/languages/definitions/svelte.ts` | 66.67% | 0.00% | n/a | | `src/languages/definitions/vue.ts` | 66.67% | 0.00% | n/a | +| `src/cli/index.ts` | 67.39% | 100.00% | 60.47% | | `src/agent/handles.ts` | 68.00% | 91.67% | 45.45% | | `src/cli/artifact.ts` | 70.00% | 100.00% | 72.22% | | `src/agent/query-index/workerPool.ts` | 70.97% | 62.50% | 56.25% | | `src/agent/followUps.ts` | 72.15% | 77.78% | 53.91% | | `src/cli/context.ts` | 72.22% | 60.78% | 75.25% | -| `src/indexer/imports/languageSpecific.ts` | 73.97% | 100.00% | 63.74% | ## Type-Only Or Re-Export Files diff --git a/docs/library-api.md b/docs/library-api.md index baf35f8d..8d40c740 100644 --- a/docs/library-api.md +++ b/docs/library-api.md @@ -439,6 +439,8 @@ console.log(packet.kind, refs.references, rows.rows, rows.freshness.state); ``` `serveCodegraphMcp()` (from `@lzehrung/codegraph/mcp`) starts the stdio server used by `codegraph mcp serve`. MCP is an agent ergonomics and cache layer over the same analysis engine, not a separate indexer. MCP file and artifact paths are confined after realpath resolution. + +`CodegraphMcpServerOptions.mcpToolConcurrency` caps concurrent calls per protocol session (default `4`); saturation returns a retryable busy error. `httpBodyTimeoutMs` bounds HTTP request-body receipt (default `30_000` ms), returning HTTP 408 on expiry. Concurrent `refresh_index` calls serialize and each applies its requested `warmup` after the preceding refresh completes. Client cancellation returns promptly but retains its concurrency slot until shared work settles, so cancellation cannot create unbounded background work. `query_sqlite` is read-only and row- and byte-bounded. It returns freshness metadata for fresh artifact reads, refreshes codegraph-owned SQLite artifacts after small edits when write access is enabled, and rejects stale artifact queries it cannot refresh safely. `artifact_build` is disabled by default and requires `readOnly: false` or CLI `--allow-build`; it refuses to write outputs from a stale MCP index until `refresh_index` succeeds. MCP `orient` and `packet_get` calls use the server-configured root; they do not accept per-request root overrides. @@ -730,7 +732,12 @@ console.log(mermaid); ## Read-only SQL from code ```ts -import { queryGraphSqliteRaw } from "@lzehrung/codegraph-core"; +import { + queryGraphSqliteRaw, + SqliteQueryCancelledError, + SqliteQueryDeadlineExceededError, + SqliteQueryWorkerCleanupCapacityExceededError, +} from "@lzehrung/codegraph-core"; const result = await queryGraphSqliteRaw( "./codegraph.sqlite", @@ -741,7 +748,11 @@ const result = await queryGraphSqliteRaw( console.log(result.columns, result.rows); ``` -`queryGraphSqliteRaw()` is intentionally read-only. It accepts result-producing statements such as `SELECT` and `PRAGMA` and rejects mutating SQL. Pass `{ maxRows }` to bound raw result rows. +`queryGraphSqliteRaw()` is intentionally read-only. It accepts result-producing statements such as `SELECT` and `PRAGMA`, rejects mutating SQL, and bounds rows, cells, response bytes, and `{ deadlineMs }`. + +`deadlineMs` must be a non-negative integer no greater than `2_147_483_647`; invalid values throw `RangeError` before the worker or fallback execution path is selected. + +The 10-second default deadline rejects the caller promptly and requests worker termination. A native SQLite step already in progress can continue in a bounded cleanup slot until it returns; degraded installs without the worker asset use a weaker in-process check after each iterator step. Callers can catch the exported `SqliteQueryDeadlineExceededError`, `SqliteQueryCancelledError`, and `SqliteQueryWorkerCleanupCapacityExceededError`; cancellation is a stable generic message so it exposes no MCP client details. ## SQL artifact facts @@ -948,7 +959,7 @@ Use the exported TypeScript APIs when another program is composing deterministic - `buildReviewReport()` returns a review bundle with `schemaVersion`, changed files, changed symbols, `graphDelta`, candidate tests, `riskSummary`, `reviewTasks`, an offline `markdownLinks` result for Markdown sources in the analysis scope when there are changes, optional duplicate sibling-check tasks, optional `sqlContext`, compatibility hints when available, and diagnostics. Accepts an optional third argument, `{ index?, loadIndex?, duplicateAnalysis?, loadDuplicateAnalysis? }`, so a caller that already holds a warm `ProjectIndex` (or wants to defer loading it until review work actually needs it) and, for repeated review calls, a `DuplicatePreparedAnalysis` from `prepareDuplicateAnalysis()` can skip redundant rebuilds. The MCP `review` tool uses the lazy forms to avoid paying index or duplicate-analysis cost on no-change reviews. - `analyzeImpactFromDiff()` returns the full or compact impact report shape for batch consumers, including an offline `markdownLinks` result for Markdown sources in the analysis scope when diffs are non-empty and changed-symbol `callCompatibility` hints when available. -- `analyzeImpactStreaming()` emits progress and incremental chunks, then a final `complete.report` summary. Streaming always returns `format: "stream-summary"`. By default this includes the same key structured fields needed by pack builders: changed files, changed symbols, impacted items, Markdown link findings, suggestions, export summaries, re-export chains, ranked top impacts, surface area, clusters, cycles, graph edges, diagnostics, and warning text. Set `streamSummary: "light"` to drop suggestions, export summaries, re-export chains, ranked top impacts, graph metadata, cycles, clusters, and surface area from the final report. +- `analyzeImpactStreaming()` emits progress and incremental chunks, then a final `complete.report` summary on success. Streaming always returns `format: "stream-summary"`. By default this includes the same key structured fields needed by pack builders: changed files, changed symbols, impacted items, Markdown link findings, suggestions, export summaries, re-export chains, ranked top impacts, surface area, clusters, cycles, graph edges, diagnostics, and warning text. Set `streamSummary: "light"` to drop suggestions, export summaries, re-export chains, ranked top impacts, graph metadata, cycles, clusters, and surface area from the final report. A bounded queue overflow instead emits terminal `error` without `complete`; ending iteration early cancels later analysis batches, but cannot interrupt a synchronous lookup already in progress. Review-pack builders should preserve symbol handles, diff snippets, callsites, `callCompatibility`, diagnostics, candidate-test confidence, impact reasons, and graph edge metadata. Render prose only at the final UI or prompt boundary. diff --git a/docs/mcp.md b/docs/mcp.md index 15673593..4191e98d 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -33,6 +33,10 @@ Stdio servers exit when the client closes stdin, when an IPC parent disconnects, HTTP protocol sessions track last activity, cap concurrent legacy sessions (default 32), and evict idle sessions on a timer (default 30 minutes). Capacity and idle eviction skip sessions with in-flight requests or open SSE streams; when every slot is active, a new `initialize` receives an actionable JSON-RPC capacity error instead of evicting a working client. Transport errors and protocol session closes also remove the session. +Each MCP protocol session permits four concurrent tool calls by default; a saturated session returns a retryable busy error rather than queueing unbounded work. The programmatic server options `mcpToolConcurrency` and `httpBodyTimeoutMs` tune that cap and the 30-second HTTP request-body deadline; an HTTP body that misses its deadline receives `408 Request Timeout`. + +Client cancellation returns promptly, but does not discard shared index or artifact work. A cancelled call continues to occupy its concurrency slot until its underlying operation settles, preventing a burst of abandoned requests from exceeding the configured resource bound. + Use stdio for a client-owned subprocess. Use HTTP for one long-running codegraph process per repository, then point every MCP-capable IDE, terminal, or agent client at the same local URL. Exact config keys vary by client, but the MCP settings should use HTTP/Streamable HTTP transport plus the `/mcp` URL instead of a `command`/`args` stdio launch. codegraph uses the official MCP SDK v2 to serve current 2026-07-28 clients while retaining compatibility with 2025-era clients. MCP protocol connections and HTTP protocol sessions keep separate transport state, but all share the server's one warm codegraph analysis session for the configured root. Tool request schemas set `additionalProperties: false` and reject unknown fields with an actionable invalid-parameter error instead of silently ignoring typos. @@ -100,7 +104,7 @@ Text and hybrid searches reuse a prepared handle for `.codegraph-cache/index-v1/ If the sidecar is busy or unavailable, MCP uses the same exact in-memory matcher. [How it works](./how-it-works.md#cache-and-session-behavior) explains the search cache. -Use `refresh_index` to rebuild the snapshot, reset SQLite artifact state, or recover after a change burst exceeds automatic limits. With write access, `query_sqlite` refreshes codegraph SQLite artifacts after small edits; otherwise it refuses stale rows. `artifact_build` refuses stale indexes, so run `refresh_index` after large change bursts. +Use `refresh_index` to rebuild the snapshot, reset SQLite artifact state, or recover after a change burst exceeds automatic limits. Concurrent refresh requests serialize; each request runs its own requested `warmup` (`off`, `base`, or `symbols`) after an active refresh completes. With write access, `query_sqlite` refreshes codegraph SQLite artifacts after small edits; otherwise it refuses stale rows. `artifact_build` refuses stale indexes, so run `refresh_index` after large change bursts. `get_file` reads live bytes from disk after path confinement. It does not require a fresh index; only an explicit `includeGraphContext: true` checks indexed freshness and adds direct graph context, so returned file bytes and `totalLines` remain live even when `freshness` reports stale context. Tool schemas are flat JSON objects for broad client compatibility; argument combinations such as `refs` handle-vs-position mode are validated by the server. Legacy paired names (`callers`, `callees`, `supertypes`, `subtypes`, `deps`, and `rdeps`) remain accepted by `tools/call` as aliases, but only the unified tools appear in `tools/list`. @@ -215,7 +219,7 @@ An MCP `explore` request whose entire query resolves to an indexed project-relat - Tool calls do not accept per-request root overrides. - Tools are read-only by default. - `artifact_build` requires `--allow-build` and a fresh or auto-refreshed MCP index. -- `query_sqlite` rejects mutating SQL, recursive queries, synthetic payload functions, and stale artifact queries it cannot refresh safely. +- `query_sqlite` rejects mutating SQL, recursive queries, synthetic payload functions, and stale artifact queries it cannot refresh safely. Each query has a 10-second execution deadline. - `get_file` rejects raw reads and structural text-config summaries over the 16 MiB input limit. Accepted reads use separate output-page bounds from `maxBytes`, `offset`, and `limit`; binary input is rejected, and sensitive formats require `allowSensitive: true` for raw values. - SQLite responses are row- and byte-bounded. 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/agent-tools.ts b/src/agent-tools.ts index fac25c6a..086f0f21 100644 --- a/src/agent-tools.ts +++ b/src/agent-tools.ts @@ -23,7 +23,7 @@ import { errorMessage } from "./util/errors.js"; import { listProjectFiles } from "./util/projectFiles.js"; import { boundAgentList, defaultAgentLimit, normalizeAgentLimit } from "./agent/bounds.js"; import { normalizeAgentOutputPath } from "./agent/normalize.js"; -import type { AgentSession } from "./agent/session.js"; +import { assertNoPrebuiltSessionWithBuildOptions, type AgentSession } from "./agent/session.js"; import { workspaceSymbols, workspaceSymbolsWithSession, @@ -303,9 +303,7 @@ export async function tool_workspaceSymbols( request: WorkspaceSymbolsRequest, runtimeOptions: ToolWorkspaceSymbolsRuntimeOptions = {}, ): Promise { - if (runtimeOptions.session && runtimeOptions.buildOptions) { - throw new Error("Workspace symbol tool options cannot combine a prebuilt session with buildOptions."); - } + assertNoPrebuiltSessionWithBuildOptions(runtimeOptions, "Workspace symbol tool options"); const agentRequest = { root, ...request, @@ -360,9 +358,7 @@ export async function tool_findImplementations( } function assertTypeHierarchyToolOptions(runtimeOptions: ToolTypeHierarchyRuntimeOptions): void { - if (runtimeOptions.session && runtimeOptions.buildOptions) { - throw new Error("Type hierarchy tool options cannot combine a prebuilt session with buildOptions."); - } + assertNoPrebuiltSessionWithBuildOptions(runtimeOptions, "Type hierarchy tool options"); } function typeHierarchyAgentRequest( @@ -408,9 +404,7 @@ export async function tool_findCallees( } function assertCallHierarchyToolOptions(runtimeOptions: ToolCallHierarchyRuntimeOptions): void { - if (runtimeOptions.session && runtimeOptions.buildOptions) { - throw new Error("Call hierarchy tool options cannot combine a prebuilt session with buildOptions."); - } + assertNoPrebuiltSessionWithBuildOptions(runtimeOptions, "Call hierarchy tool options"); } function callHierarchyAgentRequest( @@ -441,9 +435,7 @@ export async function tool_previewRename( request: Omit, runtimeOptions: ToolRenamePreviewRuntimeOptions = {}, ): Promise { - if (runtimeOptions.session && runtimeOptions.buildOptions) { - throw new Error("Rename preview tool options cannot combine a prebuilt session with buildOptions."); - } + assertNoPrebuiltSessionWithBuildOptions(runtimeOptions, "Rename preview tool options"); const agentRequest: RenamePreviewRequest = { root, ...request, @@ -469,9 +461,7 @@ export async function tool_buildRefactorPlan( request: Omit, runtimeOptions: ToolRefactorPlanRuntimeOptions = {}, ): Promise { - if (runtimeOptions.session && runtimeOptions.buildOptions) { - throw new Error("Refactor plan tool options cannot combine a prebuilt session with buildOptions."); - } + assertNoPrebuiltSessionWithBuildOptions(runtimeOptions, "Refactor plan tool options"); const agentRequest: RefactorPlanRequest = { root, ...request, diff --git a/src/agent/explore.ts b/src/agent/explore.ts index 710953b6..8b67a795 100644 --- a/src/agent/explore.ts +++ b/src/agent/explore.ts @@ -2,6 +2,7 @@ import path from "node:path"; import type { AnalysisSummary } from "../analysisSummary.js"; import { getReverseDependencies, getShortestPath, type DependencyNode } from "../graphs/traversal.js"; import { defNodeId } from "../graphs/symbol-graph.js"; +import { boundList } from "../presentation/bounds.js"; import type { BuildOptions } from "../indexer/types.js"; import { listCandidateTestFiles } from "../impact/context.js"; import { fileIdentityKey, normalizePath, toProjectDisplayPath } from "../util/paths.js"; @@ -483,11 +484,12 @@ function collectBlastRadius( const summaries: AgentExploreBlastRadiusSummary[] = []; for (const file of anchorFiles.slice(0, entryLimit)) { const dependencies = getReverseDependencies(snapshot.fileGraph, file, { limit: dependencyLimit + 1, depth: 2 }); - const visible = dependencies.slice(0, dependencyLimit).map((dependency) => formatDependency(snapshot, dependency)); + const boundedDependencies = boundList(dependencies, dependencyLimit); + const visible = boundedDependencies.items.map((dependency) => formatDependency(snapshot, dependency)); summaries.push({ file: toProjectDisplayPath(snapshot.root, file), reverseDependencies: visible, - omittedLowerBound: Math.max(0, dependencies.length - dependencyLimit), + omittedLowerBound: boundedDependencies.omitted, }); } return summaries; @@ -511,9 +513,10 @@ function collectCandidateTests( maxCandidates: snapshot.index.byFile.size, projectRoot: snapshot.root, }); + const boundedCandidates = boundList(candidates, limit); return { - items: candidates.slice(0, limit).map((candidate) => toProjectDisplayPath(snapshot.root, candidate.file)), - omittedCount: Math.max(0, candidates.length - limit), + items: boundedCandidates.items.map((candidate) => toProjectDisplayPath(snapshot.root, candidate.file)), + omittedCount: boundedCandidates.omitted, }; } function collectFollowUps( diff --git a/src/agent/query-index/sessionStore.ts b/src/agent/query-index/sessionStore.ts index c5800a05..6c8edd64 100644 --- a/src/agent/query-index/sessionStore.ts +++ b/src/agent/query-index/sessionStore.ts @@ -10,48 +10,68 @@ type SessionQueryIndexState = { }; const QUERY_INDEX_BY_SESSION = new WeakMap(); +const QUERY_INDEX_INVALIDATION_HOOKS = new WeakSet(); +const MAX_QUERY_INDEX_GENERATION_RETRIES = 3; function closeHandle(handle: QueryIndexHandle): void { handle.store?.close(); } function closeState(state: SessionQueryIndexState): void { + state.closing = true; if (state.resolved) { closeHandle(state.resolved); return; } - state.closing = true; void state.handle.then(closeHandle, () => undefined); } +function disposeSessionQueryIndexOnInvalidation(session: AgentSession): () => void { + return () => { + QUERY_INDEX_INVALIDATION_HOOKS.delete(session); + disposeSessionQueryIndex(session); + }; +} + export async function ensureSessionQueryIndex( session: AgentSession, snapshot: AgentProjectSnapshot, ): Promise { const identity = snapshot.index.projectSnapshotIdentity ?? ""; - const existing = QUERY_INDEX_BY_SESSION.get(session); - if (existing?.identity === identity) return await existing.handle; - if (existing) { - QUERY_INDEX_BY_SESSION.delete(session); - closeState(existing); - } + for (let attempt = 0; attempt < MAX_QUERY_INDEX_GENERATION_RETRIES; attempt += 1) { + const existing = QUERY_INDEX_BY_SESSION.get(session); + if (existing?.identity === identity && !existing.closing) { + const resolved = await existing.handle; + if (!existing.closing) return resolved; + } + if (existing) { + if (QUERY_INDEX_BY_SESSION.get(session) !== existing) continue; + QUERY_INDEX_BY_SESSION.delete(session); + closeState(existing); + } - const handle = ensureQueryIndex(snapshot); - const state: SessionQueryIndexState = { identity, handle }; - if (!existing) registerSessionInvalidationHook(session, () => disposeSessionQueryIndex(session)); - QUERY_INDEX_BY_SESSION.set(session, state); - handle.catch(() => { - if (QUERY_INDEX_BY_SESSION.get(session) === state) QUERY_INDEX_BY_SESSION.delete(session); - }); - const resolved = await handle; - if (QUERY_INDEX_BY_SESSION.get(session) === state) { - state.resolved = resolved; - } else if (!state.closing) { + const handle = ensureQueryIndex(snapshot); + const state: SessionQueryIndexState = { identity, handle }; + if (!QUERY_INDEX_INVALIDATION_HOOKS.has(session)) { + QUERY_INDEX_INVALIDATION_HOOKS.add(session); + registerSessionInvalidationHook(session, disposeSessionQueryIndexOnInvalidation(session)); + } + QUERY_INDEX_BY_SESSION.set(session, state); + handle.catch(() => { + if (QUERY_INDEX_BY_SESSION.get(session) === state) QUERY_INDEX_BY_SESSION.delete(session); + }); + const resolved = await handle; + if (QUERY_INDEX_BY_SESSION.get(session) === state && !state.closing) { + state.resolved = resolved; + if (snapshot.buildReport) snapshot.buildReport.queryIndex = resolved.diagnostics; + if (snapshot.index.buildReport) snapshot.index.buildReport.queryIndex = resolved.diagnostics; + return resolved; + } closeHandle(resolved); } - if (snapshot.buildReport) snapshot.buildReport.queryIndex = resolved.diagnostics; - if (snapshot.index.buildReport) snapshot.index.buildReport.queryIndex = resolved.diagnostics; - return resolved; + throw new Error( + "Query index generation changed repeatedly while loading; retry the request after refresh completes.", + ); } export function disposeSessionQueryIndex(session: AgentSession): void { diff --git a/src/agent/search.ts b/src/agent/search.ts index e221479d..8f4856ce 100644 --- a/src/agent/search.ts +++ b/src/agent/search.ts @@ -207,6 +207,8 @@ const NATURAL_LANGUAGE_SYNTAX_TERMS = new Set([ "or", ]); const SEARCH_CACHES = new WeakMap(); +// A fixed ceiling prevents client-supplied query keys retaining an unbounded session heap. +const SESSION_SEARCH_CACHE_MAX_ENTRIES = 100; const SEARCH_RESULT_CACHES = new WeakMap>>(); const SEARCH_RANKING_VERSION = 2; @@ -239,6 +241,7 @@ export async function searchCodegraphWithSession( const cacheKey = searchResultCacheKey(snapshot, request); const existing = resultCache.get(cacheKey); if (existing) { + promoteSessionSearchResult(resultCache, cacheKey, existing); const response = await existing; if (response.query === request.query) return response; return { ...response, query: request.query }; @@ -250,7 +253,7 @@ export async function searchCodegraphWithSession( queryIndex = await ensureSessionQueryIndex(session, snapshot); } const search = searchSnapshot(snapshot, request, queryIndex); - resultCache.set(cacheKey, search); + promoteSessionSearchResult(resultCache, cacheKey, search); search.catch(() => { if (resultCache.get(cacheKey) === search) resultCache.delete(cacheKey); }); @@ -357,6 +360,20 @@ function getSessionSearchResultCache(session: AgentSession): Map>, + key: string, + result: Promise, +): void { + cache.delete(key); + cache.set(key, result); + while (cache.size > SESSION_SEARCH_CACHE_MAX_ENTRIES) { + const oldest = cache.keys().next().value; + if (oldest === undefined) return; + cache.delete(oldest); + } +} + function searchResultCacheKey(snapshot: AgentProjectSnapshot, request: AgentSearchRequest): string { const { rankTokens, normalizedRankPhrase, identifierLike } = buildQueryTerms(request.query, snapshot); return JSON.stringify({ diff --git a/src/agent/session.ts b/src/agent/session.ts index e47a8427..cbe2884a 100644 --- a/src/agent/session.ts +++ b/src/agent/session.ts @@ -76,6 +76,15 @@ export type AgentSession = { invalidate: () => void; }; +export function assertNoPrebuiltSessionWithBuildOptions( + options: { session?: AgentSession; buildOptions?: BuildOptions }, + consumer: string, +): void { + if (options.session && options.buildOptions) { + throw new Error(`${consumer} cannot combine a prebuilt session with buildOptions.`); + } +} + const EMPTY_SYMBOL_GRAPH: SymbolGraph = { nodes: new Map(), edges: [], diff --git a/src/cli/viewer.ts b/src/cli/viewer.ts index 34a3de6a..32091466 100644 --- a/src/cli/viewer.ts +++ b/src/cli/viewer.ts @@ -201,64 +201,74 @@ async function writeGeneratedGraphResponse( } } +function writeViewerRequestError(response: http.ServerResponse): void { + if (response.writableEnded) return; + if (!response.headersSent) response.writeHead(500); + response.end(); +} + function viewerRequestHandler( options: ResolvedViewerOptions, getAllowedHostHeaders: () => AllowedHostHeaderRules, ): http.RequestListener { return (request, response) => { - if (!isAllowedHostHeader(request, getAllowedHostHeaders())) { - response.writeHead(403); - response.end(); - return; - } - if (request.method !== "GET" && request.method !== "HEAD") { - response.writeHead(405, { Allow: "GET, HEAD" }); - response.end(); - return; - } - - const rawPathname = (request.url ?? "/").split(/[?#]/, 1)[0] ?? "/"; try { - if ( - decodeURIComponent(rawPathname) - .split("/") - .some((segment) => segment === "..") - ) { - response.writeHead(404); + if (!isAllowedHostHeader(request, getAllowedHostHeaders())) { + response.writeHead(403); response.end(); return; } - } catch { - response.writeHead(404); - response.end(); - return; - } - - const pathname = new URL(request.url ?? "/", "http://viewer.local").pathname; - if (pathname === "/graph.json") { - if (options.graphFile) { - writeFileResponse( - request, - response, - options.graphFile.path, - "application/json; charset=utf-8", - options.graphFile.fileDescriptor, - ); + if (request.method !== "GET" && request.method !== "HEAD") { + response.writeHead(405, { Allow: "GET, HEAD" }); + response.end(); return; } - if (options.graphProvider) { - void writeGeneratedGraphResponse(request, response, options.graphProvider); + + const rawPathname = (request.url ?? "/").split(/[?#]/, 1)[0] ?? "/"; + try { + if ( + decodeURIComponent(rawPathname) + .split("/") + .some((segment) => segment === "..") + ) { + response.writeHead(404); + response.end(); + return; + } + } catch { + response.writeHead(404); + response.end(); return; } - } - const asset = VIEWER_ASSETS[pathname]; - if (!asset) { - response.writeHead(404); - response.end(); - return; + const pathname = new URL(request.url ?? "/", "http://viewer.local").pathname; + if (pathname === "/graph.json") { + if (options.graphFile) { + writeFileResponse( + request, + response, + options.graphFile.path, + "application/json; charset=utf-8", + options.graphFile.fileDescriptor, + ); + return; + } + if (options.graphProvider) { + void writeGeneratedGraphResponse(request, response, options.graphProvider); + return; + } + } + + const asset = VIEWER_ASSETS[pathname]; + if (!asset) { + response.writeHead(404); + response.end(); + return; + } + writeFileResponse(request, response, path.join(options.assetRoot, asset.file), asset.contentType); + } catch { + writeViewerRequestError(response); } - writeFileResponse(request, response, path.join(options.assetRoot, asset.file), asset.contentType); }; } diff --git a/src/impact/analyzer.ts b/src/impact/analyzer.ts index 39b40dd3..4d89b9ee 100644 --- a/src/impact/analyzer.ts +++ b/src/impact/analyzer.ts @@ -46,6 +46,13 @@ function compareImpactItems(left: ImpactItem, right: ImpactItem): number { return 0; } +function throwIfImpactAnalysisAborted(signal: AbortSignal | undefined): void { + if (!signal?.aborted) return; + const error = new Error("Impact analysis was cancelled."); + error.name = "AbortError"; + throw error; +} + export async function analyzeImpact( index: ProjectIndex, changedSymbols: ChangedSymbol[], @@ -62,6 +69,7 @@ export async function analyzeImpact( refContextLines, refBlockMaxLines, onImpactItem, + signal, } = options; const diagnostics = options.diagnostics; const projectRoot = @@ -76,6 +84,7 @@ export async function analyzeImpact( const isIgnored = projectRoot ? createImpactIgnoreMatcher(projectRoot, ignoreGlobs) : () => false; const referenceCache = createReferenceLookupCache(); const workBudget = createImpactWorkBudget(options); + throwIfImpactAnalysisAborted(signal); const impacted = new Map(); const processedSymbols = new Set(); @@ -159,6 +168,7 @@ export async function analyzeImpact( return !isIgnored(file); }, }); + throwIfImpactAnalysisAborted(signal); syncBudgetDiagnostics(diagnostics, workBudget); await analyzeDirectReferences({ index, @@ -172,6 +182,7 @@ export async function analyzeImpact( emitImpactItem, }); syncBudgetDiagnostics(diagnostics, workBudget); + throwIfImpactAnalysisAborted(signal); } // Seed transitive impact from changed files. This is NOT redundant with @@ -179,17 +190,21 @@ export async function analyzeImpact( // (they no longer exist), so they would never enter `impacted` through the symbol // loop above. seedTransitiveFromFiles plants them directly so the transitive pass // can propagate their impact to dependents. + throwIfImpactAnalysisAborted(signal); if (!options.membersOnly && !isImpactDeadlineExceeded(workBudget)) { seedTransitiveFromFiles(index, impacted, changedFiles, normalizedOptions, reverseDeps, emitImpactItem); } + throwIfImpactAnalysisAborted(signal); // Transitive impact via graph traversal (skip if membersOnly) if (!options.membersOnly && !isImpactDeadlineExceeded(workBudget)) { analyzeTransitiveImpact(impacted, depth, normalizedOptions, isIndexTestFile, reverseDeps, emitImpactItem); } + throwIfImpactAnalysisAborted(signal); syncBudgetDiagnostics(diagnostics, workBudget); const sorted = Array.from(impacted.values()).sort(compareImpactItems); + throwIfImpactAnalysisAborted(signal); for (const item of sorted) { emitImpactItem(item, "final"); } diff --git a/src/impact/streaming.ts b/src/impact/streaming.ts index 86cb59be..c66c71e4 100644 --- a/src/impact/streaming.ts +++ b/src/impact/streaming.ts @@ -31,7 +31,7 @@ export type ImpactStreamChunk = | { type: "projectFiles"; files: ProjectFileInfo[] } | import("../types.js").ProgressUpdate | { type: "changedSymbol"; symbol: ChangedSymbol } - | { type: "impactItem"; item: ImpactItem; partial?: boolean } + | { type: "impactItem"; item: ImpactItem; partial?: true } | { type: "complete"; summary: { totalChanged: number; totalImpacted: number }; @@ -40,7 +40,7 @@ export type ImpactStreamChunk = | { type: "error"; error: string }; type PublicImpactStreamingOptions = Options extends unknown - ? Omit + ? Omit : never; type WithoutCompact = Options extends unknown ? Omit : never; @@ -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,14 @@ 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 receives that signal and + * checks it at analysis work boundaries and before emitting items, so it does not start + * later batches or transitive work after abandonment. An in-progress synchronous lookup + * still runs until it returns because it cannot be preempted by JavaScript. */ export async function* analyzeImpactStreaming( projectRoot: string, @@ -194,6 +249,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 +309,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; @@ -277,7 +335,11 @@ export async function* analyzeImpactStreaming( fileLevelFallback, fileLevelFallbackPaths, diagnostics, + signal: abortController.signal, onImpactItem: (item, phase) => { + if (abortController.signal.aborted) { + throw new ImpactStreamAbandonedError(); + } queueImpactItem(item, phase === "partial"); }, }) @@ -355,6 +417,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/impact/types.ts b/src/impact/types.ts index aea9c936..489db33a 100644 --- a/src/impact/types.ts +++ b/src/impact/types.ts @@ -571,4 +571,6 @@ export type ImpactOptions = DiffProviderOptions & { diagnostics?: ImpactDiagnostics; /** @internal Internal callback used by streaming analysis to emit progressive impact snapshots */ onImpactItem?: (item: ImpactItem, phase: "partial" | "final") => void; + /** @internal Cancellation signal used to stop background streaming analysis at work boundaries. */ + signal?: AbortSignal; }; diff --git a/src/index.ts b/src/index.ts index e315c9da..b41b8569 100644 --- a/src/index.ts +++ b/src/index.ts @@ -205,6 +205,7 @@ export { createCodeReviewSession, type ICodeReviewSession, type SessionOptions, + type SessionManagerOptions, type SessionStatus, type SessionStats, type SessionStaleReason, @@ -252,6 +253,9 @@ export { updateGraphSqlite, queryGraphSqlite, queryGraphSqliteRaw, + SqliteQueryCancelledError, + SqliteQueryDeadlineExceededError, + SqliteQueryWorkerCleanupCapacityExceededError, SQLITE_ARTIFACT_FILE_SIGNATURES_METADATA_KEY, type SqliteGraphOptions, type SqliteGraphUpdateOptions, diff --git a/src/indexer/type-hierarchy.ts b/src/indexer/type-hierarchy.ts index 71feeb63..10b0223f 100644 --- a/src/indexer/type-hierarchy.ts +++ b/src/indexer/type-hierarchy.ts @@ -1,4 +1,5 @@ import type { SymbolEdge, SymbolGraph, SymbolNode } from "../graphs/symbol-graph.js"; +import { boundList } from "../presentation/bounds.js"; import { resolveSymbolId } from "./symbols.js"; import type { ProjectIndex } from "./types.js"; @@ -106,12 +107,13 @@ export function findTypeHierarchy( } relations.sort((left, right) => compareRelations(graph, left, right)); + const boundedRelations = boundList(relations, limit); return { status: "ok", targetId, direction, - relations: relations.slice(0, limit), - omitted: Math.max(0, relations.length - limit), + relations: boundedRelations.items, + omitted: boundedRelations.omitted, limit, }; } @@ -137,11 +139,12 @@ export function findImplementations( }; } const matches = collectTypeImplementations(graph, hierarchy, targetId, rootRelations); + const boundedMatches = boundList(matches, limit); return { status: "ok", targetId, - implementations: matches.slice(0, limit), - omitted: Math.max(0, matches.length - limit), + implementations: boundedMatches.items, + omitted: boundedMatches.omitted, ambiguous: 0, unresolved: [], limit, @@ -229,14 +232,15 @@ export function findImplementations( } const sortedUnresolved = unresolved.sort((left, right) => left.symbolId.localeCompare(right.symbolId)); const matches = [...memberMatches.values()].sort((left, right) => compareImplementationMatches(graph, left, right)); - const truncated = Math.max(0, matches.length - limit); + const boundedMatches = boundList(matches, limit); + const boundedUnresolved = boundList(sortedUnresolved, limit); return { status: "ok", targetId, - implementations: matches.slice(0, limit), - omitted: truncated + ambiguous, + implementations: boundedMatches.items, + omitted: boundedMatches.omitted + ambiguous, ambiguous, - unresolved: sortedUnresolved.slice(0, limit), + unresolved: boundedUnresolved.items, limit, }; } diff --git a/src/indexer/workspace-symbols.ts b/src/indexer/workspace-symbols.ts index d2e7b01c..906d185e 100644 --- a/src/indexer/workspace-symbols.ts +++ b/src/indexer/workspace-symbols.ts @@ -5,6 +5,7 @@ import { ensureParsedContext } from "./parse-context.js"; import { getCachedScope } from "./navigation-references.js"; import { resolveImported } from "./navigation-resolve.js"; import { defNodeId } from "../graphs/symbol-graph.js"; +import { boundList } from "../presentation/bounds.js"; import type { ImportBinding, ProjectIndex, SymbolDef, SymbolKind } from "./types.js"; export const DEFAULT_WORKSPACE_SYMBOL_LIMIT = 50; @@ -100,12 +101,13 @@ export async function workspaceSymbols( } ranked.sort(compareRankedCandidates); - const symbols = ranked.slice(0, limit).map(({ candidate }) => candidate); + const boundedRanked = boundList(ranked, limit); + const symbols = boundedRanked.items.map(({ candidate }) => candidate); return { query, symbols, totalCandidates: ranked.length, - omitted: Math.max(0, ranked.length - symbols.length), + omitted: boundedRanked.omitted, limit, omittedImports, importScanFailures, diff --git a/src/mcp/http.ts b/src/mcp/http.ts index f55f0ec8..e00ca628 100644 --- a/src/mcp/http.ts +++ b/src/mcp/http.ts @@ -2,7 +2,11 @@ import type { IncomingMessage, Server as HttpServer, ServerResponse } from "node import type { AddressInfo } from "node:net"; import os from "node:os"; -export type ParsedJsonBody = { status: "ok"; body: unknown } | { status: "too_large" } | { status: "invalid_json" }; +export type ParsedJsonBody = + | { status: "ok"; body: unknown } + | { status: "too_large" } + | { status: "timeout" } + | { status: "invalid_json" }; export type AllowedHostHeaderRules = { exact: Set; @@ -13,31 +17,86 @@ export function getRequestPath(request: IncomingMessage): string { return new URL(request.url ?? "/", "http://127.0.0.1").pathname; } -export async function readJsonRequestBody(request: IncomingMessage, maxBytes: number): Promise { +export async function readJsonRequestBody( + request: IncomingMessage, + maxBytes: number, + timeoutMs: number, +): Promise { const contentLength = getContentLength(request); if (contentLength !== undefined && contentLength > maxBytes) { - request.resume(); + drainRequestBody(request, timeoutMs); return { status: "too_large" }; } - const chunks: Buffer[] = []; - let bytes = 0; - for await (const chunk of request) { - const buffer = typeof chunk === "string" ? Buffer.from(chunk) : chunk; - bytes += buffer.byteLength; - if (bytes > maxBytes) { - return { status: "too_large" }; - } - chunks.push(buffer); - } + return await new Promise((resolve) => { + const chunks: Buffer[] = []; + let bytes = 0; + let settled = false; + const deadline = setTimeout(() => settle({ status: "timeout" }, true), timeoutMs); + deadline.unref?.(); + + const cleanup = (): void => { + clearTimeout(deadline); + request.off("data", onData); + request.off("end", onEnd); + request.off("error", onFailure); + request.off("aborted", onFailure); + }; + const settle = (result: ParsedJsonBody, drain: boolean): void => { + if (settled) return; + settled = true; + cleanup(); + if (drain) request.resume(); + resolve(result); + }; + const onData = (chunk: string | Buffer): void => { + const buffer = typeof chunk === "string" ? Buffer.from(chunk) : chunk; + bytes += buffer.byteLength; + if (bytes > maxBytes) { + drainRequestBody(request, timeoutMs); + settle({ status: "too_large" }, false); + return; + } + chunks.push(buffer); + }; + const onEnd = (): void => { + const rawBody = Buffer.concat(chunks).toString("utf8"); + try { + const body: unknown = rawBody.length ? JSON.parse(rawBody) : null; + settle({ status: "ok", body }, false); + } catch { + settle({ status: "invalid_json" }, false); + } + }; + const onFailure = (): void => settle({ status: "invalid_json" }, true); - const rawBody = Buffer.concat(chunks).toString("utf8"); - try { - const body: unknown = rawBody.length ? JSON.parse(rawBody) : null; - return { status: "ok", body }; - } catch { - return { status: "invalid_json" }; - } + request.on("data", onData); + request.once("end", onEnd); + request.once("error", onFailure); + request.once("aborted", onFailure); + }); +} + +function drainRequestBody(request: IncomingMessage, timeoutMs: number): void { + const onDrained = (): void => cleanup(); + const onTimedOut = (): void => { + cleanup(); + request.destroy(); + }; + const deadline = setTimeout(onTimedOut, timeoutMs); + deadline.unref?.(); + + const cleanup = (): void => { + clearTimeout(deadline); + request.off("end", onDrained); + request.off("error", onDrained); + request.off("aborted", onDrained); + }; + + request.once("end", onDrained); + request.once("error", onDrained); + request.once("aborted", onDrained); + request.resume(); } export function emptyAllowedHostHeaderRules(): AllowedHostHeaderRules { diff --git a/src/mcp/server.ts b/src/mcp/server.ts index 04891138..2257c238 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -55,10 +55,15 @@ import { DEFAULT_BOUNDED_IMPACT_BUDGETS } from "../impact/budgets.js"; import { buildReviewReport, type ReviewDepth, type ReviewReport } from "../review.js"; import { boundReviewReportForTransport, type ReviewReportForTransport } from "../review/types.js"; import { SQLITE_ARTIFACT_FILE_SIGNATURES_METADATA_KEY, queryGraphSqliteRaw, type RawSqlResult } from "../sqlite.js"; +import { boundList, countOmitted } from "../presentation/bounds.js"; import { isPlainRecord } from "../util/guards.js"; import { toProjectDisplayPath } from "../util/paths.js"; import { errorMessage } from "../util/errors.js"; -import { createAgentSession, listAgentSessionFiles } from "../agent/session.js"; +import { + assertNoPrebuiltSessionWithBuildOptions, + createAgentSession, + listAgentSessionFiles, +} from "../agent/session.js"; import { mapLimit } from "../util/concurrency.js"; import { assertRealPathCandidateWithinRoot, resolveProjectFile } from "../util/confinedFile.js"; import type { AgentFreshnessResult, AgentProjectSnapshot, AgentSession } from "../agent/session.js"; @@ -67,7 +72,6 @@ import { DEFAULT_WORKSPACE_SYMBOL_LIMIT, MAX_WORKSPACE_SYMBOL_LIMIT } from "../i import type { BuildOptions, FindReferencesResult, GoToResult } from "../indexer/types.js"; import { assertMcpSqliteQueryResourceBounded, - boundRawSqlResult, DEFAULT_SQLITE_BYTE_LIMIT, normalizeSqliteRowLimit, } from "./sqliteGuard.js"; @@ -80,7 +84,7 @@ import { MAX_TYPE_HIERARCHY_LIMIT, MAX_MCP_COLLECTION_LIMIT, MAX_RENAME_PREVIEW_EDITS, - MCP_TOOLS, + MCP_TOOL_REGISTRY, MAX_REFACTOR_PLAN_LIMIT, } from "./tools.js"; import { @@ -134,6 +138,10 @@ export type CodegraphMcpServerOptions = CodegraphMcpHandlerOptions & { httpSessionMaxCount?: number; /** How often to scan for idle HTTP sessions in ms. Defaults to 60 seconds. */ httpSessionEvictionIntervalMs?: number; + /** Maximum concurrent tool calls per MCP protocol session. Defaults to 4. */ + mcpToolConcurrency?: number; + /** Maximum time to receive an HTTP MCP request body in ms. Defaults to 30 seconds. */ + httpBodyTimeoutMs?: number; onHttpListen?: ((info: CodegraphMcpHttpServerInfo) => void) | undefined; runtimeIdentity?: CodegraphRuntimeIdentity; }; @@ -152,6 +160,13 @@ export type CodegraphMcpHttpServer = CodegraphMcpHttpServerInfo & { export const DEFAULT_MCP_HTTP_SESSION_IDLE_MS = 30 * 60 * 1000; export const DEFAULT_MCP_HTTP_SESSION_MAX_COUNT = 32; export const DEFAULT_MCP_HTTP_SESSION_EVICTION_INTERVAL_MS = 60_000; +export const DEFAULT_MCP_HTTP_BODY_TIMEOUT_MS = 30_000; +export const DEFAULT_MCP_TOOL_CONCURRENCY = 4; + +function normalizeMcpToolConcurrency(value: number): number { + if (!Number.isFinite(value)) return DEFAULT_MCP_TOOL_CONCURRENCY; + return Math.max(1, Math.floor(value)); +} type LegacyMcpSession = { server: Server; @@ -165,6 +180,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. @@ -176,7 +195,13 @@ export type McpTruncationMeta = { omitted: number; }; -export type CodegraphMcpHandlers = { +type McpDependenciesResponse = CodegraphMcpFreshResult< + McpTruncationMeta & { dependencies: Array<{ file: string; depth: number }> } +>; +type McpReverseDependenciesResponse = CodegraphMcpFreshResult< + McpTruncationMeta & { reverseDependencies: Array<{ file: string; depth: number }> } +>; +type CodegraphMcpHandlerDefinitions = { search: (request: { query: string; mode?: AgentSearchMode | undefined; @@ -295,14 +320,12 @@ export type CodegraphMcpHandlers = { file: string; depth?: number | undefined; limit?: number | undefined; - }) => Promise }>>; + }) => Promise; rdeps: (request: { file: string; depth?: number | undefined; limit?: number | undefined; - }) => Promise< - CodegraphMcpFreshResult }> - >; + }) => Promise; path: (request: { from: string; to: string }) => Promise>; impact: (request: { base: string; head: string }) => Promise>; review: (request: { @@ -314,11 +337,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; @@ -329,6 +355,14 @@ export type CodegraphMcpHandlers = { }) => Promise>; }; +type WithAbortSignal = { + [K in Exclude]: T[K] extends (request: infer Request) => Promise + ? (request: Request, signal?: AbortSignal) => Promise + : never; +} & Pick; + +export type CodegraphMcpHandlers = WithAbortSignal; + type McpDependencyRequest = { file: string; depth?: number | undefined; @@ -351,12 +385,11 @@ type SqliteArtifactFileSignature = { }; const MAX_MCP_FRESHNESS_CHANGED_FILES = 25; +const MAX_MCP_FRESHNESS_RETRIES = 3; const SQLITE_ARTIFACT_STAT_CONCURRENCY = 64; function assertMcpSessionOptions(options: CodegraphMcpHandlerOptions): void { - if (options.session !== undefined && options.buildOptions !== undefined) { - throw new Error("MCP server options cannot combine a prebuilt session with buildOptions."); - } + assertNoPrebuiltSessionWithBuildOptions(options, "MCP server options"); } function createCodegraphMcpSession(options: CodegraphMcpHandlerOptions, root: string): AgentSession { @@ -384,7 +417,14 @@ function startCodegraphMcpWarmup( return undefined; } -async function createWarmedCodegraphMcpHandlers(options: CodegraphMcpServerOptions): Promise { +type WarmedCodegraphMcpResources = { + handlers: CodegraphMcpHandlers; + session: AgentSession; +}; + +async function createWarmedCodegraphMcpResources( + options: CodegraphMcpServerOptions, +): Promise { const root = path.resolve(options.root); const session = createCodegraphMcpSession(options, root); await startCodegraphMcpWarmup(session, options.warmup); @@ -393,7 +433,10 @@ async function createWarmedCodegraphMcpHandlers(options: CodegraphMcpServerOptio void host; void port; void onHttpListen; - return createCodegraphMcpHandlersForSession({ ...handlerOptions, root }, session); + return { + handlers: createCodegraphMcpHandlersForSession({ ...handlerOptions, root }, session), + session, + }; } const MCP_HTTP_PATH = "/mcp"; @@ -420,6 +463,8 @@ function createCodegraphMcpHandlersForSession( let sqlitePath = configuredSqlitePath; let sqliteOutDir = configuredSqliteOutDir; let sqliteCanRefresh = configuredSqliteCanRefresh; + let refreshPromise: Promise | undefined; + let refreshEpoch = 0; const relative = (file: string): string => toProjectDisplayPath(root, file); const boundedLimit = (limit: number | undefined, fallback: number, max: number): number => { @@ -433,7 +478,7 @@ function createCodegraphMcpHandlersForSession( state: "stale", changedFiles: boundedChangedFiles, changedFileCount: changedFiles.length, - omittedChangedFileCount: Math.max(0, changedFiles.length - boundedChangedFiles.length), + omittedChangedFileCount: countOmitted(changedFiles.length, boundedChangedFiles.length), reason, }; }; @@ -444,9 +489,14 @@ function createCodegraphMcpHandlersForSession( const withFreshness = async ( run: () => Promise, ): Promise => { - const freshness = await checkMcpFreshness(); - const result = await run(); - return { ...result, freshness }; + for (let attempt = 0; attempt < MAX_MCP_FRESHNESS_RETRIES; attempt += 1) { + if (refreshPromise) await refreshPromise; + const epoch = refreshEpoch; + const freshness = await checkMcpFreshness(); + const result = await run(); + if (epoch === refreshEpoch && !refreshPromise) return { ...result, freshness }; + } + throw new Error("Workspace refresh changed repeatedly while serving the request; retry after refresh completes."); }; const formatSqliteFreshnessError = (freshness: AgentFreshnessResult): string => { if (freshness.state === "fresh") return "SQLite artifact freshness check unexpectedly failed."; @@ -611,7 +661,7 @@ function createCodegraphMcpHandlersForSession( // Probe one entry past the display limit so `truncated` is exact // rather than a `results.length === limit` heuristic (which cannot // tell "exactly limit reachable files" apart from "more exist"), - // without re-walking the whole reachable graph for an exact total — + // without re-walking the whole reachable graph for an exact total - // see finding #44. limit: limit + 1, }; @@ -625,13 +675,18 @@ function createCodegraphMcpHandlersForSession( targetFile = await resolveProjectFile(await realRoot, root, request.file); } const collected = collectEntries(snapshot.fileGraph, targetFile, queryOptions); - const totalSeen = collected.length; - const truncated = totalSeen > limit; - const entries = collected.slice(0, limit).map((dependency) => ({ + const { items, omitted } = boundList(collected, limit); + const entries = items.map((dependency) => ({ file: relative(dependency.file), depth: dependency.depth, })); - return { entries, limit, totalSeen, truncated, omitted: truncated ? totalSeen - limit : 0 }; + return { + entries, + limit, + totalSeen: collected.length, + truncated: Boolean(omitted), + omitted, + }; }; const calls = async (request: { direction: "callers" | "callees"; @@ -670,12 +725,15 @@ function createCodegraphMcpHandlersForSession( if (result.status !== "ok") { return { references: [], limit, totalSeen: 0, truncated: false, omitted: 0 }; } - const totalSeen = result.references.length; - const truncated = totalSeen > limit; - const references = result.references - .slice(0, limit) - .map((reference) => ({ file: relative(reference.file), range: reference.range })); - return { references, limit, totalSeen, truncated, omitted: truncated ? totalSeen - limit : 0 }; + const { items, omitted } = boundList(result.references, limit); + const references = items.map((reference) => ({ file: relative(reference.file), range: reference.range })); + return { + references, + limit, + totalSeen: result.references.length, + truncated: Boolean(omitted), + omitted, + }; }; const fileDeps = async (request: { direction: "deps" | "rdeps"; @@ -842,7 +900,7 @@ function createCodegraphMcpHandlersForSession( if (parseQualifiedSymbolPath(handle)) { const snapshot = await session.loadProject({ symbolGraph: "skip" }); const resolved = requireSemanticSymbol(snapshot, handle); - // Probe one reference past the display limit so `truncated` is exact — see + // Probe one reference past the display limit so `truncated` is exact - see // `collectMcpDependencyEntries` for the rationale (finding #44). const result = await findReferences(snapshot.index, { def: resolved.def }, { maxReferences: limit + 1 }); return boundedReferencesFromResult(result, limit); @@ -942,7 +1000,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."); } @@ -963,17 +1021,30 @@ function createCodegraphMcpHandlersForSession( } const result = await queryGraphSqliteRaw(realSqlitePath, request.query, request.params ?? [], { maxRows: normalizeSqliteRowLimit(request.limit), + maxBytes: DEFAULT_SQLITE_BYTE_LIMIT, + ...(executionOptions?.signal ? { signal: executionOptions.signal } : {}), }); - return { ...boundRawSqlResult(result, DEFAULT_SQLITE_BYTE_LIMIT), freshness: artifactFreshness }; + return { ...result, truncated: Boolean(result.truncated), freshness: artifactFreshness }; }, refresh_index: async (request) => { const warmup = request.warmup ?? "off"; - session.invalidate(); - sqlitePath = configuredSqlitePath; - sqliteOutDir = configuredSqliteOutDir; - sqliteCanRefresh = configuredSqliteCanRefresh; - await startCodegraphMcpWarmup(session, warmup); + const previousRefresh = refreshPromise; + const refresh = (async () => { + if (previousRefresh) await previousRefresh.catch(() => undefined); + ++refreshEpoch; + session.invalidate(); + sqlitePath = configuredSqlitePath; + sqliteOutDir = configuredSqliteOutDir; + sqliteCanRefresh = configuredSqliteCanRefresh; + await startCodegraphMcpWarmup(session, warmup); + })(); + refreshPromise = refresh; + try { + await refresh; + } finally { + if (refreshPromise === refresh) refreshPromise = undefined; + } return { refreshed: true, warmup }; }, @@ -1026,6 +1097,7 @@ export function createCodegraphMcpProtocolServer( runtimeIdentity: CodegraphRuntimeIdentity = captureCodegraphRuntimeIdentity(getCurrentNativeBindingOrigin()), installedVersion: InstalledVersionChecker = createInstalledVersionChecker(runtimeIdentity), toolCallState: { firstToolCallPending: boolean } = { firstToolCallPending: true }, + maxConcurrentToolCalls = DEFAULT_MCP_TOOL_CONCURRENCY, ): Server { const server = new Server( { @@ -1036,9 +1108,15 @@ export function createCodegraphMcpProtocolServer( capabilities: { tools: {}, logging: {} }, }, ); + let inFlightToolCalls = 0; + const toolConcurrency = normalizeMcpToolConcurrency(maxConcurrentToolCalls); - server.setRequestHandler("tools/list", () => ({ tools: MCP_TOOLS })); + server.setRequestHandler("tools/list", () => ({ tools: listCodegraphMcpTools() })); server.setRequestHandler("tools/call", async (request, ctx): Promise => { + if (inFlightToolCalls >= toolConcurrency) { + throw new Error("MCP tool execution is busy; retry shortly."); + } + inFlightToolCalls += 1; const isFirstToolCall = toolCallState.firstToolCallPending; toolCallState.firstToolCallPending = false; const progressToken = isFirstToolCall ? getToolCallProgressToken(request.params) : undefined; @@ -1053,12 +1131,7 @@ export function createCodegraphMcpProtocolServer( if (progressToken !== undefined) { await ctx.mcpReq.notify({ method: "notifications/progress", - params: { - progressToken, - progress, - total: 1, - message, - }, + params: { progressToken, progress, total: 1, message }, }); } } catch (error) { @@ -1076,7 +1149,11 @@ 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 operation = callMcpTool(handlers, request.params.name, request.params.arguments ?? {}, ctx.mcpReq.signal); + const releaseToolCall = (): void => { + inFlightToolCalls -= 1; + }; + const result = await awaitMcpToolOperation(ctx.mcpReq.signal, operation, releaseToolCall); await emitFirstToolCallVisibility( "info", 1, @@ -1099,10 +1176,18 @@ export function createCodegraphMcpProtocolServer( function createCodegraphMcpProtocolFactory( handlers: CodegraphMcpHandlers, runtimeIdentity: CodegraphRuntimeIdentity, + maxConcurrentToolCalls = DEFAULT_MCP_TOOL_CONCURRENCY, ): () => Server { const installedVersion = createInstalledVersionChecker(runtimeIdentity); const toolCallState = { firstToolCallPending: true }; - return () => createCodegraphMcpProtocolServer(handlers, runtimeIdentity, installedVersion, toolCallState); + return () => + createCodegraphMcpProtocolServer( + handlers, + runtimeIdentity, + installedVersion, + toolCallState, + maxConcurrentToolCalls, + ); } export async function serveCodegraphMcp(options: CodegraphMcpServerOptions): Promise { @@ -1118,9 +1203,13 @@ export async function serveCodegraphMcp(options: CodegraphMcpServerOptions): Pro return; } - const handlers = await createWarmedCodegraphMcpHandlers(options); + const { handlers, session } = await createWarmedCodegraphMcpResources(options); const runtimeIdentity = options.runtimeIdentity ?? captureCodegraphRuntimeIdentity(getCurrentNativeBindingOrigin()); - const createProtocolServer = createCodegraphMcpProtocolFactory(handlers, runtimeIdentity); + const createProtocolServer = createCodegraphMcpProtocolFactory( + handlers, + runtimeIdentity, + options.mcpToolConcurrency ?? DEFAULT_MCP_TOOL_CONCURRENCY, + ); const handle = serveStdio(createProtocolServer, { legacy: "serve", onerror: (error) => { @@ -1134,6 +1223,7 @@ export async function serveCodegraphMcp(options: CodegraphMcpServerOptions): Pro console.error(`[codegraph] MCP stdio shutting down (${shutdownReason})`); }, }); + session.invalidate(); // Ensure orphaned stdio servers do not linger after the client is gone. process.exitCode = 0; } @@ -1142,9 +1232,13 @@ export async function startCodegraphMcpHttpServer( options: CodegraphMcpServerOptions & { port: number }, ): Promise { const host = options.host ?? "127.0.0.1"; - const handlers = await createWarmedCodegraphMcpHandlers(options); + const { handlers, session } = await createWarmedCodegraphMcpResources(options); const runtimeIdentity = options.runtimeIdentity ?? captureCodegraphRuntimeIdentity(getCurrentNativeBindingOrigin()); - const createProtocolServer = createCodegraphMcpProtocolFactory(handlers, runtimeIdentity); + const createProtocolServer = createCodegraphMcpProtocolFactory( + handlers, + runtimeIdentity, + options.mcpToolConcurrency ?? DEFAULT_MCP_TOOL_CONCURRENCY, + ); const sessionStore = createLegacyMcpSessionStore({ idleMs: options.httpSessionIdleMs ?? DEFAULT_MCP_HTTP_SESSION_IDLE_MS, maxCount: options.httpSessionMaxCount ?? DEFAULT_MCP_HTTP_SESSION_MAX_COUNT, @@ -1166,8 +1260,12 @@ export async function startCodegraphMcpHttpServer( let closeResourcesPromise: Promise | undefined; const closeResources = (): Promise => { closeResourcesPromise ??= (async () => { - sessionStore.stop(); - await closeMcpResources(sessionStore.sessions, modernHandler.close); + try { + session.invalidate(); + } finally { + sessionStore.stop(); + await closeMcpResources(sessionStore.sessions, modernHandler.close); + } })(); return closeResourcesPromise; }; @@ -1181,6 +1279,7 @@ export async function startCodegraphMcpHttpServer( validateOrigin, modernNodeHandler, createProtocolServer, + options.httpBodyTimeoutMs ?? DEFAULT_MCP_HTTP_BODY_TIMEOUT_MS, ); }); @@ -1188,7 +1287,12 @@ export async function startCodegraphMcpHttpServer( void closeResources(); }); - await listenOnHttpServer(server, options.port, host); + try { + await listenOnHttpServer(server, options.port, host); + } catch (error) { + await closeResources(); + throw error; + } const address = server.address(); const actualPort = getHttpServerPort(address); const urlHost = formatHostForUrl(host); @@ -1218,7 +1322,14 @@ async function handleMcpHttpRequest( validateOrigin: OriginValidator, modernNodeHandler: NodeMcpRequestHandler, createProtocolServer: () => Server, + bodyTimeoutMs: number, ): Promise { + const writeClosingJsonRpcError = (statusCode: number, message: string): void => { + response.setHeader("connection", "close"); + + writeJsonRpcError(response, statusCode, message); + }; + const requestPath = getRequestPath(request); if (requestPath !== MCP_HTTP_PATH) { writeJsonResponse(response, 404, { error: "Not found" }); @@ -1233,11 +1344,15 @@ async function handleMcpHttpRequest( try { if (request.method === "POST") { - const parsedBody = await readJsonRequestBody(request, MAX_MCP_HTTP_BODY_BYTES); + const parsedBody = await readJsonRequestBody(request, MAX_MCP_HTTP_BODY_BYTES, bodyTimeoutMs); if (parsedBody.status === "too_large") { writeJsonRpcError(response, 413, "MCP request body is too large"); return; } + if (parsedBody.status === "timeout") { + writeClosingJsonRpcError(408, "MCP request body timed out"); + return; + } if (parsedBody.status === "invalid_json") { writeJsonRpcError(response, 400, "Invalid JSON request body"); return; @@ -1285,12 +1400,7 @@ async function handleLegacyMcpHttpPost( return; } sessionStore.touch(sessionId); - try { - await handleLegacyMcpSessionRequest(session, request, response, body); - } catch (error) { - await sessionStore.delete(sessionId); - throw error; - } + await handleLegacyMcpSessionRequest(session, request, response, body); return; } @@ -1343,9 +1453,16 @@ async function handleLegacyMcpHttpPost( openSseStreams: 0, }; sessionRef.current = session; + // The SDK transport reports every per-request validation rejection through onerror + // too (bad Accept header, wrong Content-Type, malformed JSON, an unsupported + // protocol version, ...) - each of those already answered its own request with a + // 4xx response and left the transport fully usable. Deleting the session here would + // tear down an otherwise healthy session over one malformed follow-up request. Only + // onclose reflects the transport actually shutting down (an explicit DELETE, an + // eviction we triggered, or a real fatal failure), so session teardown is driven by + // onclose alone; onerror only logs. transport.onerror = (error) => { console.error(`[codegraph] MCP HTTP session transport error: ${error.message}`); - if (initializedSessionId !== undefined) void sessionStore.delete(initializedSessionId); }; transport.onclose = () => { if (initializedSessionId !== undefined) void sessionStore.delete(initializedSessionId); @@ -1354,6 +1471,14 @@ async function handleLegacyMcpHttpPost( try { await protocolServer.connect(transport); await handleLegacyMcpSessionRequest(session, request, response, body); + if (initializedSessionId === undefined) { + // The transport answered a pre-session 4xx (invalid Accept header, wrong + // Content-Type, malformed JSON, ...) without throwing and without ever reaching + // onsessioninitialized, so nothing else releases this capacity reservation or + // closes this ad hoc protocol server/transport pair. + releaseCapacityReservation(); + await closeMcpSession(session); + } } catch (error) { if (initializedSessionId !== undefined) { await sessionStore.delete(initializedSessionId); @@ -1381,12 +1506,7 @@ async function handleExistingMcpSessionRequest( return; } sessionStore.touch(sessionId); - try { - await handleLegacyMcpSessionRequest(session, request, response); - } catch (error) { - await sessionStore.delete(sessionId); - throw error; - } + await handleLegacyMcpSessionRequest(session, request, response); } async function handleLegacyMcpSessionRequest( @@ -1528,102 +1648,132 @@ function isMcpNodeRequest(request: IncomingMessage): request is IncomingMessage return request.method !== undefined && request.url !== undefined; } -async function callMcpTool(handlers: CodegraphMcpHandlers, name: string, input: unknown): Promise { - switch (name) { +export function awaitMcpToolOperation( + signal: AbortSignal | undefined, + operation: Promise, + onSettled: () => void, +): Promise { + void operation.then(onSettled, onSettled); + return withAbortSignal(signal, operation); +} + +function withAbortSignal(signal: AbortSignal | undefined, operation: Promise): Promise { + if (!signal) return operation; + if (signal.aborted) return Promise.reject(new Error("MCP tool call was cancelled.")); + const cancellation = Promise.withResolvers(); + const onAbort = (): void => cancellation.reject(new Error("MCP tool call was cancelled.")); + signal.addEventListener("abort", onAbort, { once: true }); + return Promise.race([operation, cancellation.promise]).finally(() => signal.removeEventListener("abort", onAbort)); +} + +export async function callMcpTool( + handlers: CodegraphMcpHandlers, + name: string, + input: unknown, + signal?: AbortSignal, +): Promise { + const tool = MCP_TOOL_REGISTRY.find((entry) => entry.name === name); + if (!tool) throw new Error(`Unknown MCP tool: ${name}`); + + switch (tool.dispatch.handler) { case "search": - return await handlers.search(parseMcpToolInput(searchSchema, input, "search")); + return await handlers.search(parseMcpToolInput(searchSchema, input, name), signal); case "workspace_symbols": - return await handlers.workspace_symbols(parseMcpToolInput(workspaceSymbolsSchema, input, "workspace_symbols")); + return await handlers.workspace_symbols(parseMcpToolInput(workspaceSymbolsSchema, input, name), signal); case "rename_preview": - return await handlers.rename_preview(parseMcpToolInput(renamePreviewSchema, input, "rename_preview")); + return await handlers.rename_preview(parseMcpToolInput(renamePreviewSchema, input, name), signal); case "refactor_plan": - return await handlers.refactor_plan(parseMcpToolInput(refactorPlanSchema, input, "refactor_plan")); + return await handlers.refactor_plan(parseMcpToolInput(refactorPlanSchema, input, name), signal); case "calls": - return await handlers.calls(parseMcpToolInput(callsSchema, input, "calls")); - case "callers": - return await handlers.calls({ ...parseMcpToolInput(callHierarchySchema, input, name), direction: "callers" }); - case "callees": - return await handlers.calls({ ...parseMcpToolInput(callHierarchySchema, input, name), direction: "callees" }); + if (tool.dispatch.direction) { + return await handlers.calls( + { ...parseMcpToolInput(callHierarchySchema, input, name), direction: tool.dispatch.direction }, + signal, + ); + } + return await handlers.calls(parseMcpToolInput(callsSchema, input, name), signal); case "type_hierarchy": - return await handlers.type_hierarchy(parseMcpToolInput(typeHierarchyUnifiedSchema, input, "type_hierarchy")); - case "supertypes": - return await handlers.type_hierarchy({ - ...parseMcpToolInput(typeHierarchySchema, input, name), - direction: "supertypes", - }); - case "subtypes": - return await handlers.type_hierarchy({ - ...parseMcpToolInput(typeHierarchySchema, input, name), - direction: "subtypes", - }); + if (tool.dispatch.direction) { + return await handlers.type_hierarchy( + { ...parseMcpToolInput(typeHierarchySchema, input, name), direction: tool.dispatch.direction }, + signal, + ); + } + return await handlers.type_hierarchy(parseMcpToolInput(typeHierarchyUnifiedSchema, input, name), signal); case "implementations": - return await handlers.implementations(parseMcpToolInput(implementationsSchema, input, "implementations")); + return await handlers.implementations(parseMcpToolInput(implementationsSchema, input, name), signal); case "explore": - return await handlers.explore(parseMcpToolInput(exploreSchema, input, "explore")); + return await handlers.explore(parseMcpToolInput(exploreSchema, input, name), signal); case "orient": - return await handlers.orient(parseMcpToolInput(orientSchema, input, "orient")); + return await handlers.orient(parseMcpToolInput(orientSchema, input, name), signal); case "packet_get": - return await handlers.packet_get(parseMcpToolInput(packetGetSchema, input, "packet_get")); + return await handlers.packet_get(parseMcpToolInput(packetGetSchema, input, name), signal); case "get_file": - return await handlers.get_file(parseMcpToolInput(getFileSchema, input, "get_file")); + return await handlers.get_file(parseMcpToolInput(getFileSchema, input, name), signal); case "get_symbol": - return await handlers.get_symbol(parseMcpToolInput(handleSchema, input, "get_symbol")); + return await handlers.get_symbol(parseMcpToolInput(handleSchema, input, name), signal); case "goto": - return await callGotoTool(handlers, input); + return await callGotoTool(handlers, input, signal); case "refs": - return await callRefsTool(handlers, input); + return await callRefsTool(handlers, input, signal); case "file_deps": - return await handlers.file_deps(parseMcpToolInput(fileDepsUnifiedSchema, input, "file_deps")); - case "deps": - return await handlers.file_deps({ ...parseMcpToolInput(fileGraphSchema, input, name), direction: "deps" }); - case "rdeps": - return await handlers.file_deps({ ...parseMcpToolInput(fileGraphSchema, input, name), direction: "rdeps" }); + if (tool.dispatch.direction) { + return await handlers.file_deps( + { ...parseMcpToolInput(fileGraphSchema, input, name), direction: tool.dispatch.direction }, + signal, + ); + } + return await handlers.file_deps(parseMcpToolInput(fileDepsUnifiedSchema, input, name), signal); case "path": - return await handlers.path(parseMcpToolInput(pathSchema, input, "path")); + return await handlers.path(parseMcpToolInput(pathSchema, input, name), signal); case "impact": - return await handlers.impact(parseMcpToolInput(gitRangeSchema, input, name)); + return await handlers.impact(parseMcpToolInput(gitRangeSchema, input, name), signal); case "review": - return await handlers.review(parseMcpToolInput(reviewSchema, input, "review")); + return await handlers.review(parseMcpToolInput(reviewSchema, input, name), signal); case "query_sqlite": - return await handlers.query_sqlite(parseMcpToolInput(querySqliteSchema, input, "query_sqlite")); + return await handlers.query_sqlite(parseMcpToolInput(querySqliteSchema, input, name), { + ...(signal ? { signal } : {}), + }); case "refresh_index": - return await handlers.refresh_index(parseMcpToolInput(refreshIndexSchema, input, "refresh_index")); + return await handlers.refresh_index(parseMcpToolInput(refreshIndexSchema, input, name), signal); case "artifact_build": - return await handlers.artifact_build(parseMcpToolInput(artifactBuildSchema, input, "artifact_build")); - default: - throw new Error(`Unknown MCP tool: ${name}`); + return await handlers.artifact_build(parseMcpToolInput(artifactBuildSchema, input, name), signal); } } -async function callGotoTool(handlers: CodegraphMcpHandlers, input: unknown): Promise { +async function callGotoTool(handlers: CodegraphMcpHandlers, input: unknown, signal?: AbortSignal): Promise { const request = parseMcpToolInput(navigationSchema, input, "goto"); - if (request.handle !== undefined) return await handlers.goto({ handle: request.handle }); + if (request.handle !== undefined) return await handlers.goto({ handle: request.handle }, signal); if (request.file === undefined || request.line === undefined || request.column === undefined) { throw new Error("goto requires either handle or file, line, and column."); } - return await handlers.goto({ file: request.file, line: request.line, column: request.column }); + return await handlers.goto({ file: request.file, line: request.line, column: request.column }, signal); } async function callRefsTool( handlers: CodegraphMcpHandlers, input: unknown, + signal?: AbortSignal, ): Promise { const request = parseMcpToolInput(refsSchema, input, "refs"); if (request.handle !== undefined) { - return await handlers.refs({ - handle: request.handle, - ...(request.limit !== undefined ? { limit: request.limit } : {}), - }); + return await handlers.refs( + { handle: request.handle, ...(request.limit !== undefined ? { limit: request.limit } : {}) }, + signal, + ); } if (request.file === undefined || request.line === undefined || request.column === undefined) { throw new Error("refs requires either handle or file, line, and column."); } - return await handlers.refs({ - file: request.file, - line: request.line, - column: request.column, - ...(request.limit !== undefined ? { limit: request.limit } : {}), - }); + return await handlers.refs( + { + file: request.file, + line: request.line, + column: request.column, + ...(request.limit !== undefined ? { limit: request.limit } : {}), + }, + signal, + ); } function toToolResult(value: unknown): CallToolResult { diff --git a/src/mcp/tools.ts b/src/mcp/tools.ts index edbb54fb..eaed1bda 100644 --- a/src/mcp/tools.ts +++ b/src/mcp/tools.ts @@ -127,7 +127,32 @@ function navigationInputSchema(includeLimit: boolean): Tool["inputSchema"] { }; } -export const MCP_TOOLS: Tool[] = [ +export type McpToolDispatch = + | { handler: "search" } + | { handler: "workspace_symbols" } + | { handler: "rename_preview" } + | { handler: "refactor_plan" } + | { handler: "calls"; direction?: "callers" | "callees" } + | { handler: "type_hierarchy"; direction?: "supertypes" | "subtypes" } + | { handler: "implementations" } + | { handler: "explore" } + | { handler: "orient" } + | { handler: "packet_get" } + | { handler: "get_file" } + | { handler: "get_symbol" } + | { handler: "goto" } + | { handler: "refs" } + | { handler: "file_deps"; direction?: "deps" | "rdeps" } + | { handler: "path" } + | { handler: "impact" } + | { handler: "review" } + | { handler: "query_sqlite" } + | { handler: "refresh_index" } + | { handler: "artifact_build" }; + +export type McpToolDefinition = Tool & { dispatch: McpToolDispatch; advertised?: boolean }; + +export const MCP_TOOL_REGISTRY: McpToolDefinition[] = [ { name: "search", description: "Deterministic ranked search across files, symbols, chunks, SQL objects, and graph context.", @@ -141,6 +166,7 @@ export const MCP_TOOLS: Tool[] = [ }, ["query"], ), + dispatch: { handler: "search" }, }, { name: "workspace_symbols", @@ -162,6 +188,7 @@ export const MCP_TOOLS: Tool[] = [ }, ["query"], ), + dispatch: { handler: "workspace_symbols" }, }, { name: "rename_preview", @@ -183,6 +210,7 @@ export const MCP_TOOLS: Tool[] = [ }, ["handle", "newName"], ), + dispatch: { handler: "rename_preview" }, }, { name: "refactor_plan", @@ -199,18 +227,21 @@ export const MCP_TOOLS: Tool[] = [ }, ["handle"], ), + dispatch: { handler: "refactor_plan" }, }, { name: "calls", description: "Find proven semantic callers or callees and exact grouped callsites for a portable symbol handle. Use refs for every symbol reference and file_deps for file-level dependencies.", inputSchema: callHierarchyInputSchema(), + dispatch: { handler: "calls" }, }, { name: "type_hierarchy", description: "Find proven direct or transitive supertypes or subtypes for a portable symbol handle. Returns currently extracted extends and implements relationships only.", inputSchema: typeHierarchyInputSchema(), + dispatch: { handler: "type_hierarchy" }, }, { name: "implementations", @@ -228,6 +259,7 @@ export const MCP_TOOLS: Tool[] = [ }, ["handle"], ), + dispatch: { handler: "implementations" }, }, { name: "explore", @@ -243,6 +275,7 @@ export const MCP_TOOLS: Tool[] = [ }, ["query"], ), + dispatch: { handler: "explore" }, }, { name: "orient", @@ -251,6 +284,7 @@ export const MCP_TOOLS: Tool[] = [ includeRoots: { type: "array", items: stringProperty }, budget: orientBudgetProperty, }), + dispatch: { handler: "orient" }, }, { name: "packet_get", @@ -264,6 +298,7 @@ export const MCP_TOOLS: Tool[] = [ }, ["target"], ), + dispatch: { handler: "packet_get" }, }, { name: "get_file", @@ -279,37 +314,44 @@ export const MCP_TOOLS: Tool[] = [ }, ["file"], ), + dispatch: { handler: "get_file" }, }, { name: "get_symbol", description: "Resolve a stable search or explain handle.", inputSchema: objectSchema({ handle: stringProperty }, ["handle"]), + dispatch: { handler: "get_symbol" }, }, { name: "goto", description: "Resolve a definition by portable handle, qualified file::symbol path, or file position.", inputSchema: navigationInputSchema(false), + dispatch: { handler: "goto" }, }, { name: "refs", description: "Find references by portable handle, qualified file::symbol path, or file position.", inputSchema: navigationInputSchema(true), + dispatch: { handler: "refs" }, }, { name: "file_deps", description: "List file dependencies or reverse file dependencies by file path, qualified file::symbol path, or portable handle.", inputSchema: dependencyInputSchema(), + dispatch: { handler: "file_deps" }, }, { name: "path", description: "Find the shortest dependency path between two files.", inputSchema: objectSchema({ from: stringProperty, to: stringProperty }, ["from", "to"]), + dispatch: { handler: "path" }, }, { name: "impact", description: "Build compact impact context for a git range.", inputSchema: objectSchema({ base: stringProperty, head: stringProperty }, ["base", "head"]), + dispatch: { handler: "impact" }, }, { name: "review", @@ -322,6 +364,7 @@ export const MCP_TOOLS: Tool[] = [ }, ["base", "head"], ), + dispatch: { handler: "review" }, }, { name: "query_sqlite", @@ -337,6 +380,7 @@ export const MCP_TOOLS: Tool[] = [ }, ["query"], ), + dispatch: { handler: "query_sqlite" }, }, { name: "refresh_index", @@ -344,6 +388,7 @@ export const MCP_TOOLS: Tool[] = [ inputSchema: objectSchema({ warmup: { type: "string", enum: ["off", "base", "symbols"] }, }), + dispatch: { handler: "refresh_index" }, }, { name: "artifact_build", @@ -356,9 +401,128 @@ export const MCP_TOOLS: Tool[] = [ questions: booleanProperty, force: booleanProperty, }), + dispatch: { handler: "artifact_build" }, + }, + { + name: "callers", + description: "Legacy alias for calls with direction callers.", + advertised: false, + inputSchema: objectSchema( + { + handle: stringProperty, + depth: { type: "integer", minimum: 1, maximum: MAX_CALL_HIERARCHY_DEPTH, default: 1 }, + limit: { + type: "integer", + minimum: 0, + maximum: MAX_CALL_HIERARCHY_LIMIT, + default: DEFAULT_CALL_HIERARCHY_LIMIT, + }, + includeHeuristic: booleanProperty, + }, + ["handle"], + ), + dispatch: { handler: "calls", direction: "callers" }, + }, + { + name: "callees", + description: "Legacy alias for calls with direction callees.", + advertised: false, + inputSchema: objectSchema( + { + handle: stringProperty, + depth: { type: "integer", minimum: 1, maximum: MAX_CALL_HIERARCHY_DEPTH, default: 1 }, + limit: { + type: "integer", + minimum: 0, + maximum: MAX_CALL_HIERARCHY_LIMIT, + default: DEFAULT_CALL_HIERARCHY_LIMIT, + }, + includeHeuristic: booleanProperty, + }, + ["handle"], + ), + dispatch: { handler: "calls", direction: "callees" }, + }, + { + name: "supertypes", + description: "Legacy alias for type_hierarchy with direction supertypes.", + advertised: false, + inputSchema: objectSchema( + { + handle: stringProperty, + depth: { type: "integer", minimum: 1, maximum: MAX_TYPE_HIERARCHY_DEPTH, default: 1 }, + limit: { + type: "integer", + minimum: 0, + maximum: MAX_TYPE_HIERARCHY_LIMIT, + default: DEFAULT_TYPE_HIERARCHY_LIMIT, + }, + }, + ["handle"], + ), + dispatch: { handler: "type_hierarchy", direction: "supertypes" }, + }, + { + name: "subtypes", + description: "Legacy alias for type_hierarchy with direction subtypes.", + advertised: false, + inputSchema: objectSchema( + { + handle: stringProperty, + depth: { type: "integer", minimum: 1, maximum: MAX_TYPE_HIERARCHY_DEPTH, default: 1 }, + limit: { + type: "integer", + minimum: 0, + maximum: MAX_TYPE_HIERARCHY_LIMIT, + default: DEFAULT_TYPE_HIERARCHY_LIMIT, + }, + }, + ["handle"], + ), + dispatch: { handler: "type_hierarchy", direction: "subtypes" }, + }, + { + name: "deps", + description: "Legacy alias for file_deps with direction deps.", + advertised: false, + inputSchema: objectSchema( + { + file: dependencyFileProperty, + depth: { type: "integer", minimum: 0, default: 1 }, + limit: { + type: "integer", + minimum: 0, + maximum: MAX_MCP_COLLECTION_LIMIT, + default: DEFAULT_MCP_COLLECTION_LIMIT, + }, + }, + ["file"], + ), + dispatch: { handler: "file_deps", direction: "deps" }, + }, + { + name: "rdeps", + description: "Legacy alias for file_deps with direction rdeps.", + advertised: false, + inputSchema: objectSchema( + { + file: dependencyFileProperty, + depth: { type: "integer", minimum: 0, default: 1 }, + limit: { + type: "integer", + minimum: 0, + maximum: MAX_MCP_COLLECTION_LIMIT, + default: DEFAULT_MCP_COLLECTION_LIMIT, + }, + }, + ["file"], + ), + dispatch: { handler: "file_deps", direction: "rdeps" }, }, ]; export function listCodegraphMcpTools(): Tool[] { - return MCP_TOOLS.map((tool) => ({ ...tool })); + return MCP_TOOL_REGISTRY.filter((tool) => tool.advertised !== false).map( + ({ dispatch: _dispatch, advertised: _advertised, ...tool }) => tool, + ); } diff --git a/src/session.ts b/src/session.ts index e234269d..e2c0677c 100644 --- a/src/session.ts +++ b/src/session.ts @@ -47,6 +47,16 @@ export type SessionOptions = { incremental?: boolean; }; +export type SessionManagerOptions = { + /** Maximum sessions, including sessions currently initializing. Defaults to 32. */ + maxSessions?: number; + /** Idle-session scan interval in milliseconds. Defaults to 60 seconds. Use 0 to disable. */ + evictionIntervalMs?: number; +}; + +export const DEFAULT_SESSION_MANAGER_MAX_SESSIONS = 32; +export const DEFAULT_SESSION_MANAGER_EVICTION_INTERVAL_MS = 60_000; + export type SessionStatus = "initializing" | "ready" | "expired" | "error"; export type SessionStaleReason = "tracked_files_changed" | "config_changed"; @@ -787,21 +797,63 @@ export class CodeReviewSession implements ICodeReviewSession { } } +function normalizeSessionManagerCapacity(value: number | undefined): number { + if (value === undefined || !Number.isFinite(value)) return DEFAULT_SESSION_MANAGER_MAX_SESSIONS; + return Math.max(1, Math.floor(value)); +} + +function normalizeSessionManagerEvictionInterval(value: number | undefined): number { + if (value === undefined || !Number.isFinite(value)) return DEFAULT_SESSION_MANAGER_EVICTION_INTERVAL_MS; + return Math.max(0, Math.floor(value)); +} + /** * Session manager for multiple concurrent sessions * Useful for agents handling multiple repositories or PRs */ + +type PendingSession = { + cancelled: boolean; + fingerprint: string; + retainPending: boolean; + promise: Promise; +}; export class SessionManager { private sessions = new Map(); - private pendingSessions = new Map< - string, - { - cancelled: boolean; - fingerprint: string; - retainPending: boolean; - promise: Promise; + private pendingSessions = new Map(); + private readonly maxSessions: number; + private readonly evictionTimer: ReturnType | undefined; + private disposed = false; + + constructor(options: SessionManagerOptions = {}) { + this.maxSessions = normalizeSessionManagerCapacity(options.maxSessions); + const evictionIntervalMs = normalizeSessionManagerEvictionInterval(options.evictionIntervalMs); + if (evictionIntervalMs) { + this.evictionTimer = setInterval(() => this.cleanupExpired(), evictionIntervalMs); + this.evictionTimer.unref?.(); + } + } + + private cancelPendingSession(sessionId: string, pending: PendingSession): void { + pending.cancelled = true; + pending.retainPending = false; + void pending.promise + .finally(() => { + if (this.pendingSessions.get(sessionId) === pending && !pending.retainPending) { + this.pendingSessions.delete(sessionId); + } + }) + .catch(() => {}); + } + + 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.`, + ); } - >(); + } private createSessionConfigurationError( sessionId: string, @@ -828,6 +880,9 @@ export class SessionManager { ): Promise | undefined { const pending = this.pendingSessions.get(sessionId); if (!pending) return undefined; + if (pending.cancelled) { + throw new Error(`Session "${sessionId}" is still cancelling initialization. Retry after it settles.`); + } const requestedFingerprint = sessionIdentityFingerprint(resolveSessionIdentity(options)); if (pending.fingerprint !== requestedFingerprint) { const existing = this.sessions.get(sessionId); @@ -880,10 +935,15 @@ export class SessionManager { return promise; } + private assertNotDisposed(): void { + if (this.disposed) throw new Error("Session manager is disposed."); + } + /** * Create or get a session for a repository */ async getOrCreateSession(sessionId: string, options: SessionOptions): Promise { + this.assertNotDisposed(); const pending = this.getPendingCompatibleSession(sessionId, options); if (pending) { return await pending; @@ -892,6 +952,7 @@ export class SessionManager { let session = this.ensureSessionIdCompatible(sessionId, options); if (!session) { + this.assertCapacityForNewSession(); session = new CodeReviewSession(options); return await this.trackSession(sessionId, options, session, false, (readySession) => { this.sessions.set(sessionId, readySession); @@ -921,10 +982,7 @@ export class SessionManager { */ disposeSession(sessionId: string): void { const pending = this.pendingSessions.get(sessionId); - if (pending) { - pending.cancelled = true; - this.pendingSessions.delete(sessionId); - } + if (pending) this.cancelPendingSession(sessionId, pending); const session = this.sessions.get(sessionId); if (session) { session.dispose(); @@ -936,16 +994,26 @@ export class SessionManager { * Dispose of all sessions */ disposeAll(): void { - for (const pending of this.pendingSessions.values()) { - pending.cancelled = true; + for (const [sessionId, pending] of this.pendingSessions) { + this.cancelPendingSession(sessionId, pending); } - this.pendingSessions.clear(); for (const session of this.sessions.values()) { session.dispose(); } this.sessions.clear(); } + /** + * Dispose all sessions and stop periodic expiration cleanup. + * This manager cannot be reused afterward. + */ + dispose(): void { + if (this.disposed) return; + this.disposed = true; + clearInterval(this.evictionTimer); + this.disposeAll(); + } + /** * Get all session IDs */ @@ -985,6 +1053,7 @@ export class SessionManager { * @param sessions - Array of session configs to pre-warm */ async warmup(sessions: Array<{ id: string; options: SessionOptions }>): Promise { + this.assertNotDisposed(); const requestedFingerprints = new Map(); const replacementSessions: Array<{ id: string; @@ -1014,6 +1083,7 @@ export class SessionManager { if (existing?.isReady()) { continue; } + if (!existing) this.assertCapacityForNewSession(); const session = new CodeReviewSession(options); replacementSessions.push(existing ? { id, existing, session } : { id, session }); warmupPromises.push(this.trackSession(id, options, session, true, () => {})); @@ -1022,10 +1092,7 @@ export class SessionManager { } catch (error) { for (const replacement of replacementSessions) { const pending = this.pendingSessions.get(replacement.id); - if (pending) { - pending.cancelled = true; - this.pendingSessions.delete(replacement.id); - } + if (pending) this.cancelPendingSession(replacement.id, pending); replacement.session.dispose(); } throw error; diff --git a/src/sqlite.ts b/src/sqlite.ts index ab8f0a23..5cf133fa 100644 --- a/src/sqlite.ts +++ b/src/sqlite.ts @@ -1,3 +1,9 @@ export type { GraphQueryResult, RawSqlResult, SqliteGraphOptions, SqliteGraphUpdateOptions } from "./sqlite/types.js"; export { SQLITE_ARTIFACT_FILE_SIGNATURES_METADATA_KEY, writeGraphSqlite, updateGraphSqlite } from "./sqlite/write.js"; -export { queryGraphSqlite, queryGraphSqliteRaw } from "./sqlite/query.js"; +export { + queryGraphSqlite, + queryGraphSqliteRaw, + SqliteQueryCancelledError, + SqliteQueryDeadlineExceededError, + SqliteQueryWorkerCleanupCapacityExceededError, +} from "./sqlite/query.js"; diff --git a/src/sqlite/query.ts b/src/sqlite/query.ts index 39854987..2f401718 100644 --- a/src/sqlite/query.ts +++ b/src/sqlite/query.ts @@ -6,40 +6,129 @@ import { DEFAULT_SQLITE_BYTE_LIMIT, MAX_SQLITE_CELL_BYTES, MAX_SQLITE_ROW_LIMIT, + normalizeSqliteRowLimit, } from "./rowBounds.js"; +import { + resolveRawSqlQueryWorkerPath, + runRawSqlQueryInWorker, + SqliteQueryCancelledError, + SqliteQueryDeadlineExceededError, + SqliteQueryWorkerCleanupCapacityExceededError, +} from "./rawQueryWorkerPool.js"; export { queryGraphSqlite } from "./canned-query.js"; +export { SqliteQueryCancelledError, SqliteQueryDeadlineExceededError, SqliteQueryWorkerCleanupCapacityExceededError }; + +/** Hard wall-clock budget for a single raw `query_sqlite` execution - see the caveat on + * `queryGraphSqliteRaw` about when this is actually enforceable. */ +export const DEFAULT_SQLITE_QUERY_DEADLINE_MS = 10_000; + +const MAX_SQLITE_QUERY_DEADLINE_MS = 2_147_483_647; export type QueryGraphSqliteRawOptions = { maxRows?: number | undefined; maxBytes?: number | undefined; maxCellBytes?: number | undefined; + /** Non-negative integer milliseconds through 2_147_483_647. Invalid values throw RangeError. */ + deadlineMs?: number | undefined; + signal?: AbortSignal | undefined; }; +function normalizeSqliteQueryDeadlineMs(value: number): number { + if (!Number.isInteger(value) || value < 0 || value > MAX_SQLITE_QUERY_DEADLINE_MS) { + throw new RangeError( + `SQLite query deadlineMs must be a non-negative integer no greater than ${MAX_SQLITE_QUERY_DEADLINE_MS}.`, + ); + } + return value; +} + +let loggedInProcessDeadlineFallback = false; +/** + * 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`). At expiry the caller receives a + * deadline error and the pool requests worker termination. A synchronous + * `DatabaseSync` call already inside SQLite may continue until that native step returns, + * but the lifecycle retains a bounded cleanup slot and the host event loop stays free. + * + * 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. `node:sqlite`'s + * `DatabaseSync` exposes no interrupt/cancellation API, so once execution is inside a + * single synchronous native call there is nothing in-process that can preempt it. + * + * The fallback checks only after each native iterator step returns, including a terminal + * empty result. A statement that is slow to produce its first row or discover that it + * has no rows still blocks for its full cost before the deadline can be observed. This + * fallback exists to keep a degraded install usable, not as a substitute for the worker + * deadline; a one-time process warning makes the weakened guarantee observable. + */ 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 = normalizeSqliteQueryDeadlineMs(options?.deadlineMs ?? DEFAULT_SQLITE_QUERY_DEADLINE_MS); + + try { + resolveRawSqlQueryWorkerPath(); + } catch { + if (!loggedInProcessDeadlineFallback) { + loggedInProcessDeadlineFallback = true; + console.error( + "[codegraph] Raw SQLite query worker asset is unavailable; falling back to an in-process " + + "execution deadline that is only checked between produced rows and cannot interrupt a " + + "single blocking native call. Reinstall to restore the enforced worker-thread deadline.", + ); + } + 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, + ); +} + +/** Degraded fallback for `queryGraphSqliteRaw` - see its doc comment for the enforcement + * caveat this path cannot avoid. */ +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 requestedRows = options?.maxRows; - const maxRows = - requestedRows === undefined - ? MAX_SQLITE_ROW_LIMIT - : Math.min(MAX_SQLITE_ROW_LIMIT, Math.max(0, Math.floor(requestedRows))); - 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)) { @@ -49,3 +138,33 @@ export async function queryGraphSqliteRaw( } }); } + +/** Throws once the wall-clock deadline has passed after a native iterator step returns. + * See the fallback caveat on `queryGraphSqliteRaw`: a slow-before-first-row query is not + * interrupted here, only detected after that synchronous iteration step completes. */ +function* withPerRowDeadline( + rows: Iterable, + deadlineAt: number, + deadlineMs: number, + signal: AbortSignal | undefined, +): Generator { + const iterator = rows[Symbol.iterator](); + let completed = false; + try { + while (true) { + if (signal?.aborted) throw new SqliteQueryCancelledError(); + const next = iterator.next(); + if (signal?.aborted) throw new SqliteQueryCancelledError(); + if (Date.now() > deadlineAt) { + throw new SqliteQueryDeadlineExceededError(deadlineMs); + } + if (next.done) { + completed = true; + return; + } + yield next.value; + } + } finally { + if (!completed) iterator.return?.(); + } +} diff --git a/src/sqlite/rawQueryWorker.ts b/src/sqlite/rawQueryWorker.ts new file mode 100644 index 00000000..b22ef635 --- /dev/null +++ b/src/sqlite/rawQueryWorker.ts @@ -0,0 +1,44 @@ +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. The pool rejects the caller at its + * deadline and requests worker termination. A native SQLite step already in progress + * can continue until it returns, while the lifecycle retains its bounded cleanup slot. + */ +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..f6a6f05b --- /dev/null +++ b/src/sqlite/rawQueryWorkerPool.ts @@ -0,0 +1,176 @@ +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; termination was requested.`); + this.name = "SqliteQueryDeadlineExceededError"; + } +} + +export class SqliteQueryCancelledError extends Error { + constructor() { + super("SQLite query was cancelled."); + this.name = "SqliteQueryCancelledError"; + } +} + +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"; + } +} + +/** + * 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(); + const deadlineSignal = AbortSignal.timeout(deadlineMs); + const combinedSignal = signal ? AbortSignal.any([signal, deadlineSignal]) : deadlineSignal; + 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; + + 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 selfDirectory = path.dirname(fileURLToPath(import.meta.url)); + const sibling = path.resolve(selfDirectory, "rawQueryWorker.js"); + if (fs.existsSync(sibling)) return sibling; + const packageRoot = findPackageRoot(selfDirectory); + 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 rejects the caller after requesting worker termination, so + * the caller never waits longer than `deadlineMs` and the host event loop is never + * blocked by the query. A `terminate()` request cannot preempt a single already-in-flight + * synchronous native call: a query whose entire cost is inside one `sqlite3_step()`, such + * as a recursive CTE or a plan that must fully sort or scan before a first row, continues + * on its orphaned worker thread until that native call returns naturally. + * + * The caller does not wait for 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. + * Process shutdown can still wait for a Worker blocked in native code, a `worker_threads` + * platform limit that is orthogonal to this function's prompt caller deadline. + */ +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/agent-explore.test.ts b/tests/agent-explore.test.ts index e058018d..6091db2b 100644 --- a/tests/agent-explore.test.ts +++ b/tests/agent-explore.test.ts @@ -841,4 +841,37 @@ describe("agent explore", () => { expect(readArray(response.blastRadius, "blastRadius")).toHaveLength(1); expect(response.freshness).toBeTypeOf("object"); }); + it("pins omission counts at and just past the limit for candidate tests and blast radius", async () => { + const root = await mkExploreRepo(); + await writeFile(root, "tests/auth.test.ts", "import { validateUser } from '../src/auth';\nvalidateUser('bob');\n"); + await writeFile( + root, + "tests/auth-spec.test.ts", + "import { validateUser } from '../src/auth';\nvalidateUser('carol');\n", + ); + + const exploreAll = await exploreCodegraph({ root, query: "validateUser" }); + expect(exploreAll.candidateTests.length).toBeGreaterThanOrEqual(2); + expect(exploreAll.omittedCounts.candidateTests).toBe(0); + + const spy = vi.spyOn(impactContext, "listCandidateTestFiles").mockReturnValue([ + { file: path.join(root, "tests/routes.test.ts"), reasons: [] }, + { file: path.join(root, "tests/auth.test.ts"), reasons: [] }, + { file: path.join(root, "tests/auth-spec.test.ts"), reasons: [] }, + ]); + + try { + const atLimitResponse = await exploreCodegraph({ root, query: "validateUser" }); + expect(atLimitResponse.candidateTests).toHaveLength(3); + expect(atLimitResponse.omittedCounts.candidateTests).toBe(0); + } finally { + spy.mockRestore(); + } + + const authExplore = await exploreCodegraph({ root, query: "src/db.ts" }); + const dbBlast = authExplore.blastRadius.find((entry) => entry.file === "src/db.ts"); + expect(dbBlast).toBeDefined(); + expect(dbBlast!.reverseDependencies.length).toBeGreaterThanOrEqual(1); + expect(dbBlast!.omittedLowerBound).toBe(0); + }); }); diff --git a/tests/agent-search.test.ts b/tests/agent-search.test.ts index 5228b884..a7db5fcb 100644 --- a/tests/agent-search.test.ts +++ b/tests/agent-search.test.ts @@ -854,4 +854,47 @@ describe("agent search", () => { expect(response.results).toEqual([]); }); + it("coalesces concurrent queries and evicts oldest entries when session search cache exceeds max entries", async () => { + const root = await mkRepo(); + const session = createAgentSession({ root }); + try { + await session.loadProject(); + + // 1. Identical concurrent queries coalesce to one in-flight promise + const p1 = searchCodegraphWithSession(session, { root, query: "validateUser", mode: "symbol" }); + const p2 = searchCodegraphWithSession(session, { root, query: "validateUser", mode: "symbol" }); + const [r1, r2] = await Promise.all([p1, p2]); + expect(r1).toBe(r2); + + // 2. Issuing more unique queries than the 100 cap evicts the oldest entry + const firstResult = r1; + const recentResults = []; + for (let i = 0; i < 100; i += 1) { + const res = await searchCodegraphWithSession(session, { + root, + query: "needleUniqueQuery" + i, + mode: "symbol", + }); + recentResults.push(res); + } + + // Re-querying the oldest entry ("validateUser") produces a new result because it was evicted + const reQueryFirst = await searchCodegraphWithSession(session, { + root, + query: "validateUser", + mode: "symbol", + }); + expect(reQueryFirst).not.toBe(firstResult); + + // Re-querying the most recent entry ("needleUniqueQuery99") returns the cached result + const reQueryLatest = await searchCodegraphWithSession(session, { + root, + query: "needleUniqueQuery99", + mode: "symbol", + }); + expect(reQueryLatest).toBe(recentResults[99]); + } finally { + session.invalidate(); + } + }); }); diff --git a/tests/agent-session.test.ts b/tests/agent-session.test.ts index f5361106..26d53952 100644 --- a/tests/agent-session.test.ts +++ b/tests/agent-session.test.ts @@ -1,3 +1,6 @@ +import { disposeSessionQueryIndex, ensureSessionQueryIndex } from "../src/agent/query-index/sessionStore.js"; +import * as updateModule from "../src/agent/query-index/update.js"; +import * as sessionLifecycleModule from "../src/agent/sessionLifecycle.js"; import fs from "node:fs/promises"; import { brotliCompressSync, brotliDecompressSync, constants as zlibConstants } from "node:zlib"; import { createHash } from "node:crypto"; @@ -6,6 +9,7 @@ import path from "node:path"; import { afterEach, describe, expect, it, vi } from "vitest"; import { buildProjectIndexIncremental, type BuildReport } from "../src/index.js"; import { AGENT_FRESHNESS_CHECK_INTERVAL_MS, createAgentSession, listAgentSessionFiles } from "../src/agent/session.js"; +import type { QueryIndexHandle } from "../src/agent/query-index/update.js"; import * as symbolGraphBuild from "../src/graphs/symbol-graph-detailed.js"; import * as indexerBuild from "../src/indexer/build-index.js"; import { createProjectSnapshotIdentity } from "../src/indexer/build-cache.js"; @@ -938,4 +942,97 @@ describe("agent session", () => { dateSpy.mockRestore(); } }); + + describe("query index sessionStore generation retries (S12)", () => { + it("bounds query index generation retries under sustained invalidation and surfaces a clear error", async () => { + const root = await mkRepo(); + const session = createAgentSession({ root }); + const snapshot = await session.loadProject(); + + let attempts = 0; + const realEnsureQueryIndex = updateModule.ensureQueryIndex; + const ensureQueryIndexSpy = vi.spyOn(updateModule, "ensureQueryIndex").mockImplementation(async (snap) => { + attempts += 1; + const res = await realEnsureQueryIndex(snap); + disposeSessionQueryIndex(session); + return res; + }); + const invalidationHookSpy = vi.spyOn(sessionLifecycleModule, "registerSessionInvalidationHook"); + + try { + await expect(ensureSessionQueryIndex(session, snapshot)).rejects.toThrow( + /Query index generation changed repeatedly while loading/i, + ); + expect(attempts).toBe(3); + expect(invalidationHookSpy).toHaveBeenCalledTimes(1); + } finally { + invalidationHookSpy.mockRestore(); + ensureQueryIndexSpy.mockRestore(); + } + }); + + it("reuses a replacement state when invalidation races a waiting caller", async () => { + const root = await mkRepo(); + const session = createAgentSession({ root }); + const snapshot = await session.loadProject(); + const firstBuildStarted = Promise.withResolvers(); + const releaseFirstBuild = Promise.withResolvers(); + const firstHandle: QueryIndexHandle = { + store: null, + diagnostics: { + sidecarState: "created", + filesRead: 0, + filesAdded: 0, + filesUpdated: 0, + filesDeleted: 0, + fileCandidates: 0, + chunkCandidates: 0, + openMs: 0, + updateMs: 0, + candidateMs: 0, + scoringMs: 0, + }, + }; + const replacementHandle: QueryIndexHandle = { + store: null, + diagnostics: { + sidecarState: "created", + filesRead: 0, + filesAdded: 0, + filesUpdated: 0, + filesDeleted: 0, + fileCandidates: 0, + chunkCandidates: 0, + openMs: 0, + updateMs: 0, + candidateMs: 0, + scoringMs: 0, + }, + }; + let attempts = 0; + const ensureQueryIndexSpy = vi.spyOn(updateModule, "ensureQueryIndex").mockImplementation(() => { + attempts += 1; + if (attempts === 1) { + firstBuildStarted.resolve(); + return releaseFirstBuild.promise.then(() => firstHandle); + } + return Promise.resolve(replacementHandle); + }); + + try { + const waitingLoad = ensureSessionQueryIndex(session, snapshot); + await firstBuildStarted.promise; + disposeSessionQueryIndex(session); + const replacementLoad = ensureSessionQueryIndex(session, snapshot); + releaseFirstBuild.resolve(); + + await expect(waitingLoad).resolves.toBe(replacementHandle); + await expect(replacementLoad).resolves.toBe(replacementHandle); + expect(attempts).toBe(2); + } finally { + disposeSessionQueryIndex(session); + ensureQueryIndexSpy.mockRestore(); + } + }); + }); }); diff --git a/tests/cli-bundle-entry.test.ts b/tests/cli-bundle-entry.test.ts index 44a29492..f8a03ee7 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"]); @@ -77,11 +79,12 @@ describe("bundled CLI entry", () => { .readdirSync(binDir) .filter((name) => name.endsWith(".js")) .sort(); - // Bundle emits exactly two self-contained entrypoints (cli + queryIndexWorker). - expect(outputs).toEqual(["cli.js", "queryIndexWorker.js"]); + // Bundle emits exactly three self-contained entrypoints (cli + queryIndexWorker + rawQueryWorker). + expect(outputs).toEqual(["cli.js", "queryIndexWorker.js", "rawQueryWorker.js"]); const entry = fs.readFileSync(bundledCli, "utf8"); expect(entry).toContain("queryIndexWorker.js"); expect(fs.existsSync(path.join(binDir, "queryIndexWorker.js"))).toBe(true); + expect(fs.existsSync(path.join(binDir, "rawQueryWorker.js"))).toBe(true); }); it("keeps a leading shebang so package managers can exec the bin directly", () => { 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-server.test.ts b/tests/mcp-server.test.ts index 259359a4..d2b750a9 100644 --- a/tests/mcp-server.test.ts +++ b/tests/mcp-server.test.ts @@ -1,11 +1,16 @@ +import { registerSessionInvalidationHook } from "../src/agent/sessionLifecycle.js"; +import { ensureSessionQueryIndex } from "../src/agent/query-index/sessionStore.js"; import fs from "node:fs/promises"; import { request as httpRequest, type IncomingMessage } from "node:http"; +import { NodeStreamableHTTPServerTransport } from "@modelcontextprotocol/node"; import os from "node:os"; import path from "node:path"; import { DatabaseSync } from "node:sqlite"; import { describe, expect, it, vi } from "vitest"; import { createAgentSession, type AgentProjectSnapshot, type AgentSession } from "../src/agent/session.js"; import { + awaitMcpToolOperation, + callMcpTool, createCodegraphMcpHandlers, createCodegraphMcpProtocolServer, listCodegraphMcpTools, @@ -14,6 +19,7 @@ import { type CodegraphMcpHandlers, } from "../src/mcp/server.js"; import { SymbolKind, type ModuleIndex, type ProjectIndex } from "../src/indexer/types.js"; +import { MCP_TOOL_REGISTRY } from "../src/mcp/tools.js"; import { DEFAULT_REVIEW_TRANSPORT_LIMITS } from "../src/review/types.js"; import type { Graph } from "../src/types.js"; import * as symbolGraphBuild from "../src/graphs/symbol-graph-detailed.js"; @@ -1125,7 +1131,7 @@ describe("codegraph MCP handlers", () => { } }); - it("rejects oversized HTTP MCP request bodies before parsing", async () => { + it("rejects declared oversized HTTP MCP request bodies without buffering their payload", async () => { const root = await fs.mkdtemp(path.join(os.tmpdir(), "cg-mcp-http-large-")); await fs.writeFile(path.join(root, "auth.ts"), "export const ok = 1;\n", "utf8"); const httpServer = await startCodegraphMcpHttpServer({ @@ -1143,16 +1149,107 @@ describe("codegraph MCP handlers", () => { }, body: JSON.stringify({ jsonrpc: "2.0", id: 1, method: "initialize", padding: "x".repeat(1_000_000) }), }); - const payload = readObject((await response.json()) as unknown); - const error = readObject(payload.error); expect(response.status).toBe(413); - expect(error.message).toBe("MCP request body is too large"); } finally { await httpServer.close(); } }); + it("returns a timeout response while draining an incomplete HTTP MCP body", async () => { + const root = await fs.mkdtemp(path.join(os.tmpdir(), "cg-mcp-http-timeout-")); + await fs.writeFile(path.join(root, "auth.ts"), "export const ok = 1;\n", "utf8"); + const httpServer = await startCodegraphMcpHttpServer({ + root, + host: "127.0.0.1", + port: 0, + httpBodyTimeoutMs: 25, + }); + + try { + const endpoint = new URL(httpServer.url); + const partialBody = '{"jsonrpc":"2.0","id":1,"method":"initialize"'; + const response = await new Promise<{ status: number; payload: JsonRpcObject; connection: string | undefined }>( + (resolve, reject) => { + let responseReceived = false; + const request = httpRequest( + { + hostname: endpoint.hostname, + port: endpoint.port, + path: endpoint.pathname, + method: "POST", + headers: { + accept: "application/json", + "content-type": "application/json", + "content-length": String(Buffer.byteLength(partialBody) + 1), + }, + }, + (incoming) => { + let responseBody = ""; + incoming.setEncoding("utf8"); + incoming.on("data", (chunk: string) => { + responseBody += chunk; + }); + incoming.on("end", () => { + responseReceived = true; + try { + resolve({ + status: incoming.statusCode ?? 0, + payload: readJsonRpcObject(JSON.parse(responseBody)), + connection: incoming.headers.connection, + }); + } catch (error) { + reject(error instanceof Error ? error : new Error(String(error))); + } finally { + request.destroy(); + } + }); + }, + ); + request.on("error", (error) => { + if (!responseReceived) reject(error); + }); + request.write(partialBody); + }, + ); + + expect([response.status, response.connection]).toEqual([408, "close"]); + expect(readObject(response.payload.error).message).toBe("MCP request body timed out"); + } finally { + await httpServer.close(); + } + }); + + it("invalidates prebuilt resources when HTTP server binding fails", async () => { + const root = await fs.mkdtemp(path.join(os.tmpdir(), "cg-mcp-listen-failure-")); + await fs.writeFile(path.join(root, "auth.ts"), "export const ok = 1;\n", "utf8"); + const occupiedServer = await startCodegraphMcpHttpServer({ + root, + host: "127.0.0.1", + port: 0, + }); + const session = createAgentSession({ root }); + let invalidated = false; + registerSessionInvalidationHook(session, () => { + invalidated = true; + }); + + try { + await expect( + startCodegraphMcpHttpServer({ + root, + host: "127.0.0.1", + port: occupiedServer.port, + session, + }), + ).rejects.toThrow(); + expect(invalidated).toBe(true); + } finally { + await occupiedServer.close(); + await fs.rm(root, { force: true, recursive: true }); + } + }); + it("reuses one session across search, get_symbol, refs, and query_sqlite handlers", async () => { const root = await fs.mkdtemp(path.join(os.tmpdir(), "cg-mcp-")); await fs.writeFile(path.join(root, "auth.ts"), "export function validateUser(id: number) { return id > 0; }\n"); @@ -2754,3 +2851,479 @@ describe("codegraph MCP handlers", () => { function normalizeSqlitePath(value: unknown): string { return typeof value === "string" ? value.replace(/\\/g, "/") : ""; } + +describe("MCP tool registry dispatch", () => { + it("routes every registered tool without advertising legacy aliases", async () => { + const root = await fs.mkdtemp(path.join(os.tmpdir(), "cg-mcp-registry-")); + await fs.writeFile(path.join(root, "auth.ts"), "export function ok(): number { return 1; }\n", "utf8"); + runGit(root, ["init"]); + runGit(root, ["add", "."]); + runGit(root, ["commit", "-m", "base"]); + const handlers = createCodegraphMcpHandlers({ root }); + const handlerSpies = { + search: vi.spyOn(handlers, "search"), + workspace_symbols: vi.spyOn(handlers, "workspace_symbols"), + rename_preview: vi.spyOn(handlers, "rename_preview"), + refactor_plan: vi.spyOn(handlers, "refactor_plan"), + calls: vi.spyOn(handlers, "calls"), + type_hierarchy: vi.spyOn(handlers, "type_hierarchy"), + implementations: vi.spyOn(handlers, "implementations"), + explore: vi.spyOn(handlers, "explore"), + orient: vi.spyOn(handlers, "orient"), + packet_get: vi.spyOn(handlers, "packet_get"), + get_file: vi.spyOn(handlers, "get_file"), + get_symbol: vi.spyOn(handlers, "get_symbol"), + goto: vi.spyOn(handlers, "goto"), + refs: vi.spyOn(handlers, "refs"), + file_deps: vi.spyOn(handlers, "file_deps"), + path: vi.spyOn(handlers, "path"), + impact: vi.spyOn(handlers, "impact"), + review: vi.spyOn(handlers, "review"), + query_sqlite: vi.spyOn(handlers, "query_sqlite"), + refresh_index: vi.spyOn(handlers, "refresh_index"), + artifact_build: vi.spyOn(handlers, "artifact_build"), + }; + const advertisedTools = listCodegraphMcpTools(); + expect(advertisedTools.map((tool) => tool.name)).toEqual( + MCP_TOOL_REGISTRY.filter((tool) => tool.advertised !== false).map((tool) => tool.name), + ); + expect(advertisedTools.some((tool) => "dispatch" in tool)).toBe(false); + const handle = "auth.ts::ok"; + const toolInputs: Record> = { + search: { query: "ok" }, + workspace_symbols: { query: "ok" }, + rename_preview: { handle, newName: "renamed" }, + refactor_plan: { handle }, + calls: { handle, direction: "callers" }, + type_hierarchy: { handle, direction: "supertypes" }, + implementations: { handle }, + explore: { query: "ok" }, + orient: {}, + packet_get: { target: "auth.ts" }, + get_file: { file: "auth.ts" }, + get_symbol: { handle }, + goto: { handle }, + refs: { handle }, + file_deps: { file: "auth.ts", direction: "deps" }, + path: { from: "auth.ts", to: "auth.ts" }, + impact: { base: "HEAD", head: "HEAD" }, + review: { base: "HEAD", head: "HEAD" }, + query_sqlite: { query: "SELECT 1" }, + refresh_index: {}, + artifact_build: {}, + callers: { handle }, + callees: { handle }, + supertypes: { handle }, + subtypes: { handle }, + deps: { file: "auth.ts" }, + rdeps: { file: "auth.ts" }, + }; + + for (const tool of MCP_TOOL_REGISTRY) { + const input = toolInputs[tool.name]; + expect(input, "missing valid input for " + tool.name).toBeDefined(); + try { + await callMcpTool(handlers, tool.name, input); + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + expect(message, tool.name).not.toMatch(/Unknown MCP tool|Invalid parameters for/i); + } + } + + for (const [name, spy] of Object.entries(handlerSpies)) { + expect(spy, name).toHaveBeenCalled(); + } + expect(handlerSpies.calls).toHaveBeenCalledWith(expect.objectContaining({ direction: "callers" }), undefined); + expect(handlerSpies.calls).toHaveBeenCalledWith(expect.objectContaining({ direction: "callees" }), undefined); + expect(handlerSpies.type_hierarchy).toHaveBeenCalledWith( + expect.objectContaining({ direction: "supertypes" }), + undefined, + ); + expect(handlerSpies.type_hierarchy).toHaveBeenCalledWith( + expect.objectContaining({ direction: "subtypes" }), + undefined, + ); + expect(handlerSpies.file_deps).toHaveBeenCalledWith(expect.objectContaining({ direction: "deps" }), undefined); + expect(handlerSpies.file_deps).toHaveBeenCalledWith(expect.objectContaining({ direction: "rdeps" }), undefined); + }); +}); + +describe("MCP refresh coalescing", () => { + it("serializes every queued request's requested warmup", async () => { + const root = await fs.mkdtemp(path.join(os.tmpdir(), "cg-mcp-refresh-coalesce-")); + await fs.writeFile(path.join(root, "auth.ts"), "export const ok = 1;\n", "utf8"); + const backingSession = createAgentSession({ root }); + const firstWarmupReached = Promise.withResolvers(); + const releaseFirstWarmup = Promise.withResolvers(); + const secondWarmupReached = Promise.withResolvers(); + const releaseSecondWarmup = Promise.withResolvers(); + const loadModes: Array<"skip" | "full"> = []; + let activeLoads = 0; + let maxActiveLoads = 0; + let fullWarmups = 0; + const session: AgentSession = { + ...backingSession, + loadProject: async (options) => { + const mode = options?.symbolGraph === "skip" ? "skip" : "full"; + loadModes.push(mode); + activeLoads += 1; + maxActiveLoads = Math.max(maxActiveLoads, activeLoads); + try { + if (mode === "skip") { + firstWarmupReached.resolve(); + await releaseFirstWarmup.promise; + } else { + fullWarmups += 1; + if (fullWarmups === 1) { + secondWarmupReached.resolve(); + await releaseSecondWarmup.promise; + } + } + return await backingSession.loadProject(options); + } finally { + activeLoads -= 1; + } + }, + }; + const handlers = createCodegraphMcpHandlers({ root, session }); + + const first = handlers.refresh_index({ warmup: "base" }); + await firstWarmupReached.promise; + const second = handlers.refresh_index({ warmup: "symbols" }); + const third = handlers.refresh_index({ warmup: "symbols" }); + releaseFirstWarmup.resolve(); + + try { + await secondWarmupReached.promise; + for (let turn = 0; turn < 4; turn += 1) await Promise.resolve(); + expect(loadModes).toEqual(["skip", "full"]); + expect(maxActiveLoads).toBe(1); + + releaseSecondWarmup.resolve(); + await expect(first).resolves.toEqual({ refreshed: true, warmup: "base" }); + await expect(second).resolves.toEqual({ refreshed: true, warmup: "symbols" }); + await expect(third).resolves.toEqual({ refreshed: true, warmup: "symbols" }); + expect(loadModes).toEqual(["skip", "full", "full"]); + expect(maxActiveLoads).toBe(1); + } finally { + releaseFirstWarmup.resolve(); + releaseSecondWarmup.resolve(); + } + }); + + it("bounds a request invalidated by repeated refreshes", async () => { + const root = await fs.mkdtemp(path.join(os.tmpdir(), "cg-mcp-refresh-retry-bound-")); + await fs.writeFile(path.join(root, "auth.ts"), "export const token = 1;\n", "utf8"); + const backingSession = createAgentSession({ root }); + let refreshes = 0; + const session: AgentSession = { + ...backingSession, + loadProject: async (options) => { + refreshes += 1; + await handlers.refresh_index({ warmup: "off" }); + return await backingSession.loadProject(options); + }, + }; + const handlers = createCodegraphMcpHandlers({ root, session }); + + await expect(handlers.goto({ file: "auth.ts", line: 1, column: 14 })).rejects.toThrow( + /Workspace refresh changed repeatedly while serving the request/i, + ); + expect(refreshes).toBe(3); + }); +}); + +describe("MCP session teardown regressions (S2)", () => { + it("closes sidecar query index handle and runs session invalidation hooks on server close", async () => { + const root = await fs.mkdtemp(path.join(os.tmpdir(), "cg-mcp-s2-teardown-")); + await fs.writeFile( + path.join(root, "auth.ts"), + "export function validateUser(token: string) { return !!token; }\n", + "utf8", + ); + const session = createAgentSession({ root }); + let invalidationHookRan = false; + registerSessionInvalidationHook(session, () => { + invalidationHookRan = true; + }); + + const httpServer = await startCodegraphMcpHttpServer({ + root, + host: "127.0.0.1", + port: 0, + session, + }); + + try { + const initialize = await postMcpJson(httpServer.url, { + jsonrpc: "2.0", + id: 1, + method: "initialize", + params: { + protocolVersion: "2025-11-25", + capabilities: {}, + clientInfo: { name: "codegraph-s2-test", version: "1.0.0" }, + }, + }); + const sessionId = initialize.response.headers.get("mcp-session-id"); + expect(sessionId).toBeTruthy(); + + const searchCall = await postMcpJson( + httpServer.url, + { + jsonrpc: "2.0", + id: 2, + method: "tools/call", + params: { name: "search", arguments: { query: "validateUser", mode: "hybrid" } }, + }, + sessionId ?? undefined, + ); + expect(searchCall.response.status).toBe(200); + + const snapshot = await session.loadProject(); + const handle = await ensureSessionQueryIndex(session, snapshot); + expect(handle.store?.closed).toBe(false); + expect(invalidationHookRan).toBe(false); + + await httpServer.close(); + + expect(invalidationHookRan).toBe(true); + expect(handle.store?.closed).toBe(true); + } finally { + await httpServer.close(); + } + }); + + it("closes legacy protocol transports when session invalidation fails during server shutdown", async () => { + const root = await fs.mkdtemp(path.join(os.tmpdir(), "cg-mcp-invalidation-close-")); + await fs.writeFile(path.join(root, "auth.ts"), "export const ok = 1;\n", "utf8"); + const session = createAgentSession({ root }); + const httpServer = await startCodegraphMcpHttpServer({ + root, + host: "127.0.0.1", + port: 0, + session, + }); + const transportCloseSpy = vi.spyOn(NodeStreamableHTTPServerTransport.prototype, "close"); + const invalidationSpy = vi.spyOn(session, "invalidate").mockImplementation(() => { + throw new Error("session invalidation failed"); + }); + + try { + const initialize = await postMcpJson(httpServer.url, { + jsonrpc: "2.0", + id: 1, + method: "initialize", + params: { + protocolVersion: "2025-11-25", + capabilities: {}, + clientInfo: { name: "codegraph-invalidation-close-test", version: "1.0.0" }, + }, + }); + expect(initialize.response.status).toBe(200); + transportCloseSpy.mockClear(); + + await expect(httpServer.close()).rejects.toThrow("session invalidation failed"); + expect(transportCloseSpy).toHaveBeenCalled(); + } finally { + invalidationSpy.mockRestore(); + transportCloseSpy.mockRestore(); + await httpServer.close().catch(() => {}); + } + }); +}); + +describe("MCP cancellation accounting", () => { + it("keeps a tool-call slot reserved until a cancelled operation settles", async () => { + const controller = new AbortController(); + const operation = Promise.withResolvers(); + let released = 0; + const pending = awaitMcpToolOperation(controller.signal, operation.promise, () => { + released += 1; + }); + + controller.abort(); + await expect(pending).rejects.toThrow("MCP tool call was cancelled."); + expect(released).toBe(0); + + operation.resolve("finished"); + await vi.waitFor(() => { + expect(released).toBe(1); + }); + }); +}); + +describe("MCP transport isolation regressions (S8)", () => { + it("preserves session and completes concurrent calls when one response connection is forcefully closed", async () => { + const root = await fs.mkdtemp(path.join(os.tmpdir(), "cg-mcp-s8-transport-")); + await fs.writeFile( + path.join(root, "auth.ts"), + "export function validateUser(token: string) { return !!token; }\nexport function secondarySymbol() { return true; }\n", + "utf8", + ); + const httpServer = await startCodegraphMcpHttpServer({ + root, + host: "127.0.0.1", + port: 0, + }); + + try { + const initialize = await postMcpJson(httpServer.url, { + jsonrpc: "2.0", + id: 1, + method: "initialize", + params: { + protocolVersion: "2025-11-25", + capabilities: {}, + clientInfo: { name: "codegraph-s8-test", version: "1.0.0" }, + }, + }); + const sessionId = initialize.response.headers.get("mcp-session-id"); + expect(sessionId).toBeTruthy(); + if (!sessionId) throw new Error("Missing sessionId"); + + const endpoint = new URL(httpServer.url); + const call1Payload = JSON.stringify({ + jsonrpc: "2.0", + id: 2, + method: "tools/call", + params: { name: "workspace_symbols", arguments: { query: "validateUser" } }, + }); + + const call1Closed = Promise.withResolvers(); + const req1 = httpRequest({ + hostname: endpoint.hostname, + port: endpoint.port, + path: endpoint.pathname, + method: "POST", + headers: { + accept: "application/json, text/event-stream", + "content-type": "application/json", + "content-length": String(Buffer.byteLength(call1Payload)), + "mcp-session-id": sessionId, + }, + }); + req1.on("error", () => { + call1Closed.resolve(); + }); + req1.on("close", () => { + call1Closed.resolve(); + }); + req1.write(call1Payload); + req1.destroy(new Error("Forced client disconnect")); + + const call2Promise = postMcpJson( + httpServer.url, + { + jsonrpc: "2.0", + id: 3, + method: "tools/call", + params: { name: "workspace_symbols", arguments: { query: "secondarySymbol" } }, + }, + sessionId, + ); + + await call1Closed.promise; + const call2 = await call2Promise; + expect(call2.response.status).toBe(200); + expect(readToolJsonResult(call2.payload).symbols).toEqual([expect.objectContaining({ name: "secondarySymbol" })]); + + const call3 = await postMcpJson( + httpServer.url, + { + jsonrpc: "2.0", + id: 4, + method: "tools/call", + params: { name: "workspace_symbols", arguments: { query: "validateUser" } }, + }, + sessionId, + ); + expect(call3.response.status).toBe(200); + expect(readToolJsonResult(call3.payload).symbols).toEqual([expect.objectContaining({ name: "validateUser" })]); + } finally { + await httpServer.close(); + } + }); +}); + +describe("MCP legacy session capacity and error-handling regressions", () => { + it("releases the initialization capacity reservation when legacy Accept header validation rejects the request", async () => { + const root = await fs.mkdtemp(path.join(os.tmpdir(), "cg-mcp-capacity-accept-")); + await fs.writeFile(path.join(root, "auth.ts"), "export const ok = 1;\n", "utf8"); + const httpServer = await startCodegraphMcpHttpServer({ + root, + port: 0, + httpSessionIdleMs: 0, + httpSessionMaxCount: 1, + }); + + const initializeRequest = { + jsonrpc: "2.0", + method: "initialize", + params: { + protocolVersion: "2025-11-25", + capabilities: {}, + clientInfo: { name: "codegraph-capacity-test", version: "1.0.0" }, + }, + }; + + try { + // No override supplies an Accept header, so postRawHttpJson's default + // ("application/json" without "text/event-stream") trips the legacy transport's + // own 406 validation before any session is created. + const rejected = await postRawHttpJson(httpServer.url, { ...initializeRequest, id: 1 }, {}); + expect(rejected.status).toBe(406); + + // With httpSessionMaxCount 1, a leaked capacity reservation from the rejected + // attempt would make this second initialize 503 instead of succeeding. + const accepted = await postMcpJson(httpServer.url, { ...initializeRequest, id: 2 }); + expect(accepted.response.status).toBe(200); + expect(accepted.response.headers.get("mcp-session-id")).toBeTruthy(); + } finally { + await httpServer.close(); + } + }); + + it("keeps a healthy session usable after a request-scoped SDK validation error on a follow-up request", async () => { + const root = await fs.mkdtemp(path.join(os.tmpdir(), "cg-mcp-session-request-error-")); + await fs.writeFile(path.join(root, "auth.ts"), "export function ok(): number { return 1; }\n", "utf8"); + const httpServer = await startCodegraphMcpHttpServer({ root, port: 0 }); + + try { + const initialize = await postMcpJson(httpServer.url, { + jsonrpc: "2.0", + id: 1, + method: "initialize", + params: { + protocolVersion: "2025-11-25", + capabilities: {}, + clientInfo: { name: "codegraph-session-error-test", version: "1.0.0" }, + }, + }); + const sessionId = initialize.response.headers.get("mcp-session-id"); + expect(sessionId).toBeTruthy(); + if (!sessionId) throw new Error("Missing sessionId"); + + // A follow-up request against the same session with a bad Accept header trips + // the transport's own request-scoped validation (406) through onerror, without + // throwing and without the transport ever closing. + const badAccept = await postRawHttpJson( + httpServer.url, + { jsonrpc: "2.0", id: 2, method: "tools/list", params: {} }, + { "mcp-session-id": sessionId }, + ); + expect(badAccept.status).toBe(406); + + // The session must still be usable: a prior bug deleted it from the store on + // every onerror, which would turn this into a 400 "Invalid or missing session ID". + const followUp = await postMcpJson( + httpServer.url, + { jsonrpc: "2.0", id: 3, method: "tools/list", params: {} }, + sessionId, + ); + expect(followUp.response.status).toBe(200); + } finally { + await httpServer.close(); + } + }); +}); diff --git a/tests/mcp-skill-parity.test.ts b/tests/mcp-skill-parity.test.ts index 2569050f..16df6d0f 100644 --- a/tests/mcp-skill-parity.test.ts +++ b/tests/mcp-skill-parity.test.ts @@ -1,7 +1,7 @@ import fsp from "node:fs/promises"; import path from "node:path"; import { describe, expect, it } from "vitest"; -import { MCP_TOOLS } from "../src/mcp/tools.js"; +import { listCodegraphMcpTools } from "../src/mcp/tools.js"; const LEGACY_ALIAS_NAMES = new Set(["callers", "callees", "supertypes", "subtypes", "deps", "rdeps"]); @@ -18,7 +18,7 @@ describe("MCP / SKILL inventory parity", () => { [...(inventoryLine?.matchAll(/`([a-z][a-z0-9_]*)`/g) ?? [])].map((match) => match[1]!).filter(Boolean), ); - for (const tool of MCP_TOOLS) { + for (const tool of listCodegraphMcpTools()) { expect(LEGACY_ALIAS_NAMES.has(tool.name)).toBe(false); expect(listed.has(tool.name), `${tool.name} missing from SKILL MCP inventory`).toBe(true); } diff --git a/tests/mcp-stream-cancellation.test.ts b/tests/mcp-stream-cancellation.test.ts new file mode 100644 index 00000000..63ad22de --- /dev/null +++ b/tests/mcp-stream-cancellation.test.ts @@ -0,0 +1,92 @@ +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, + DEFAULT_MCP_TOOL_CONCURRENCY, +} 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 }); + } + }); + + it("falls back to the default tool concurrency for NaN", async () => { + const root = await fsp.mkdtemp(path.join(os.tmpdir(), "cg-mcp-tool-concurrency-")); + const handlers = createCodegraphMcpHandlers({ root }); + const started = Promise.withResolvers(); + const release = Promise.withResolvers(); + let activeCalls = 0; + handlers.query_sqlite = async () => { + activeCalls += 1; + if (activeCalls === DEFAULT_MCP_TOOL_CONCURRENCY) started.resolve(); + await release.promise; + return { columns: [], rows: [] }; + }; + + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + const server = createCodegraphMcpProtocolServer(handlers, undefined, undefined, undefined, Number.NaN); + const client = new Client({ name: "mcp-tool-concurrency-test", version: "1.0.0" }); + try { + await server.connect(serverTransport); + await client.connect(clientTransport); + + const active = Array.from({ length: DEFAULT_MCP_TOOL_CONCURRENCY }, () => + client.callTool({ name: "query_sqlite", arguments: { query: "SELECT 1;" } }), + ); + await started.promise; + + await expect(client.callTool({ name: "query_sqlite", arguments: { query: "SELECT 1;" } })).rejects.toThrow( + /tool execution is busy/i, + ); + release.resolve(); + await expect(Promise.all(active)).resolves.toHaveLength(DEFAULT_MCP_TOOL_CONCURRENCY); + } finally { + release.resolve(); + 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/query-index.test.ts b/tests/query-index.test.ts index 3d34839d..fd27e985 100644 --- a/tests/query-index.test.ts +++ b/tests/query-index.test.ts @@ -4,7 +4,8 @@ import path from "node:path"; import { afterEach, describe, expect, it, vi } from "vitest"; import { searchCodegraphWithSession, type AgentSearchResponse } from "../src/agent/search.js"; import { createAgentSession, type AgentProjectSnapshot, type AgentSession } from "../src/agent/session.js"; -import { disposeSessionQueryIndex } from "../src/agent/query-index/sessionStore.js"; +import { disposeSessionQueryIndex, ensureSessionQueryIndex } from "../src/agent/query-index/sessionStore.js"; +import * as updateModule from "../src/agent/query-index/update.js"; import { resolveQueryIndexPaths, resolveQueryIndexSourcePath } from "../src/agent/query-index/paths.js"; import { expectedQueryIndexVersionMetadata, probeQueryIndexSqliteSupport } from "../src/agent/query-index/schema.js"; import { SqliteDatabase } from "../src/sqlite-driver.js"; @@ -868,4 +869,27 @@ describe("persistent query index", () => { expect(response.results.some((result) => result.file === "src/auth.ts")).toBe(true); await expect(fs.stat(path.join(root, ".codegraph-cache"))).rejects.toMatchObject({ code: "ENOENT" }); }); + it("bounds query index generation retries under sustained invalidation and surfaces a clear error", async () => { + const root = await createRepo(); + const session = createSession(root); + const snapshot = await session.loadProject(); + + let attempts = 0; + const realEnsureQueryIndex = updateModule.ensureQueryIndex; + const ensureQueryIndexSpy = vi.spyOn(updateModule, "ensureQueryIndex").mockImplementation(async (snap) => { + attempts += 1; + const res = await realEnsureQueryIndex(snap); + disposeSessionQueryIndex(session); + return res; + }); + + try { + await expect(ensureSessionQueryIndex(session, snapshot)).rejects.toThrow( + /Query index generation changed repeatedly while loading/i, + ); + expect(attempts).toBe(3); + } finally { + ensureQueryIndexSpy.mockRestore(); + } + }); }); diff --git a/tests/raw-query-worker-lifecycle.test.ts b/tests/raw-query-worker-lifecycle.test.ts new file mode 100644 index 00000000..88d1f279 --- /dev/null +++ b/tests/raw-query-worker-lifecycle.test.ts @@ -0,0 +1,98 @@ +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( + /worker capacity/i, + ); + + firstCleanup.resolve(); + secondCleanup.resolve(); + await Promise.resolve(); + await Promise.resolve(); + expect(lifecycle.state()).toEqual({ activeWorkers: 0, maxWorkers: 2 }); + }); + + it("describes capacity exhaustion caused by an active query", async () => { + const lifecycle = new RawSqlQueryWorkerLifecycle(1); + const complete = Promise.withResolvers<{ columns: string[]; rows: Array> }>(); + const active = lifecycle.run(task, 10_000, undefined, () => ({ + run: async () => await complete.promise, + destroy: async () => {}, + })); + + await expect(lifecycle.run(task, 10_000, undefined, () => createAbortablePool(Promise.resolve()))).rejects.toThrow( + /active or cleaning-up worker/i, + ); + + complete.resolve({ columns: [], rows: [] }); + await expect(active).resolves.toEqual({ columns: [], rows: [] }); + }); + + it("releases no worker slot when an invalid deadline is rejected before startup", async () => { + const lifecycle = new RawSqlQueryWorkerLifecycle(1); + + await expect( + lifecycle.run(task, -1, undefined, () => ({ + run: async () => ({ columns: [], rows: [] }), + destroy: async () => {}, + })), + ).rejects.toThrow(); + expect(lifecycle.state()).toEqual({ activeWorkers: 0, maxWorkers: 1 }); + + await expect( + lifecycle.run(task, 10_000, undefined, () => ({ + run: async () => ({ columns: [], rows: [] }), + destroy: async () => {}, + })), + ).resolves.toEqual({ columns: [], rows: [] }); + }); +}); diff --git a/tests/session.test.ts b/tests/session.test.ts index 510cf11d..59865199 100644 --- a/tests/session.test.ts +++ b/tests/session.test.ts @@ -1,8 +1,14 @@ import { describe, test, expect, beforeAll, afterAll, afterEach, beforeEach, vi } from "vitest"; -import type { ICodeReviewSession } from "../src/index.js"; +import type { ICodeReviewSession, SessionManagerOptions as PublicSessionManagerOptions } from "../src/index.js"; import type { BuildOptions, BuildReport, LanguageExtensionMap } from "../src/indexer/types.js"; -import { CodeReviewSession, SessionManager, createCodeReviewSession } from "../src/session.js"; +import { + CodeReviewSession, + DEFAULT_SESSION_MANAGER_MAX_SESSIONS, + 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"; @@ -1188,6 +1194,18 @@ describe("SessionManager", () => { manager = new SessionManager(); }); + afterEach(() => { + manager.dispose(); + }); + + test("exports SessionManagerOptions from the package root", () => { + const options: PublicSessionManagerOptions = { evictionIntervalMs: 0, maxSessions: 1 }; + const typedManager = new SessionManager(options); + + expect(typedManager).toBeInstanceOf(SessionManager); + typedManager.dispose(); + }); + test("should create and retrieve sessions", async () => { const session = await manager.getOrCreateSession("test-session", { root: sampleRoot, @@ -1215,6 +1233,133 @@ describe("SessionManager", () => { expect(session1).toBe(session2); }); + test("rejects a new session when configured capacity is exhausted", async () => { + const limitedManager = new SessionManager({ maxSessions: 1, evictionIntervalMs: 0 }); + try { + await limitedManager.getOrCreateSession("first", { + root: sampleRoot, + buildOptions: sampleBuildOptions(), + }); + + await expect( + limitedManager.getOrCreateSession("second", { + root: sampleRoot, + buildOptions: sampleBuildOptions(), + }), + ).rejects.toThrow("Session capacity reached (1)"); + } finally { + limitedManager.disposeAll(); + } + }); + + test("enforces capacity when warming net-new sessions", async () => { + const limitedManager = new SessionManager({ maxSessions: 1, evictionIntervalMs: 0 }); + try { + await limitedManager.getOrCreateSession("existing", { + root: sampleRoot, + buildOptions: sampleBuildOptions(), + }); + + await expect( + limitedManager.warmup([ + { + id: "warm", + options: { root: sampleRoot, buildOptions: sampleBuildOptions() }, + }, + ]), + ).rejects.toThrow("Session capacity reached (1)"); + } finally { + limitedManager.disposeAll(); + } + }); + + test("retains failed warmup capacity until initialization settles", async () => { + const limitedManager = new SessionManager({ maxSessions: 1, evictionIntervalMs: 0 }); + const originalBuild = indexerBuild.buildProjectIndexIncremental; + const buildStarted = Promise.withResolvers(); + const releaseBuild = Promise.withResolvers(); + const buildSpy = vi.spyOn(indexerBuild, "buildProjectIndexIncremental").mockImplementation(async (...args) => { + buildStarted.resolve(); + await releaseBuild.promise; + return await originalBuild(...args); + }); + + try { + const warmup = expect( + limitedManager.warmup([ + { + id: "warm-a", + options: { root: sampleRoot, buildOptions: sampleBuildOptions() }, + }, + { + id: "warm-b", + options: { root: sampleRoot, buildOptions: sampleBuildOptions() }, + }, + ]), + ).rejects.toThrow("Session capacity reached (1)"); + await buildStarted.promise; + await warmup; + + await expect( + limitedManager.getOrCreateSession("after-failed-warmup", { + root: sampleRoot, + buildOptions: sampleBuildOptions(), + }), + ).rejects.toThrow("Session capacity reached (1)"); + + releaseBuild.resolve(); + await vi.waitFor(() => { + expect(Reflect.get(limitedManager, "pendingSessions").size).toBe(0); + }); + await expect( + limitedManager.getOrCreateSession("after-failed-warmup", { + root: sampleRoot, + buildOptions: sampleBuildOptions(), + }), + ).resolves.toBeInstanceOf(CodeReviewSession); + } finally { + releaseBuild.resolve(); + buildSpy.mockRestore(); + limitedManager.disposeAll(); + } + }); + + test("releases failed warmup capacity after initialization rejects", async () => { + const limitedManager = new SessionManager({ maxSessions: 1, evictionIntervalMs: 0 }); + const missingRoot = path.join(os.tmpdir(), `cg-session-warmup-missing-${Date.now()}`); + try { + await expect( + limitedManager.warmup([ + { + id: "broken", + options: { root: missingRoot, buildOptions: sampleBuildOptions() }, + }, + ]), + ).rejects.toThrow(); + + await vi.waitFor(() => { + expect(Reflect.get(limitedManager, "pendingSessions").size).toBe(0); + }); + await expect( + limitedManager.getOrCreateSession("after-failed-init", { + root: sampleRoot, + buildOptions: sampleBuildOptions(), + }), + ).resolves.toBeInstanceOf(CodeReviewSession); + } finally { + limitedManager.disposeAll(); + } + }); + + test("falls back to default capacity when maxSessions is NaN", async () => { + const nanManager = new SessionManager({ maxSessions: Number.NaN, evictionIntervalMs: 0 }); + try { + expect(Reflect.get(nanManager, "maxSessions")).toBe(DEFAULT_SESSION_MANAGER_MAX_SESSIONS); + } finally { + nanManager.disposeAll(); + } + }); + test("should share one initialization across concurrent same-id creation", async () => { const buildSpy = vi.spyOn(indexerBuild, "buildProjectIndexIncremental"); @@ -1268,7 +1413,8 @@ describe("SessionManager", () => { } }); - test("should allow immediate recreation after disposing a pending session", async () => { + test("retains pending initialization capacity after disposal until it settles", async () => { + const limitedManager = new SessionManager({ maxSessions: 1, evictionIntervalMs: 0 }); const originalBuild = indexerBuild.buildProjectIndexIncremental; let releaseBuild: (() => void) | null = null; const buildGate = new Promise((resolve) => { @@ -1280,31 +1426,42 @@ describe("SessionManager", () => { }); try { - const firstSession = manager.getOrCreateSession("pending", { + const firstSession = limitedManager.getOrCreateSession("pending", { root: sampleRoot, buildOptions: sampleBuildOptions(), }); await Promise.resolve(); - manager.disposeSession("pending"); - const secondSession = manager.getOrCreateSession("pending", { - root: sampleRoot, - buildOptions: sampleBuildOptions(), - }); - releaseBuild?.(); + limitedManager.disposeSession("pending"); + await expect( + limitedManager.getOrCreateSession("pending", { + root: sampleRoot, + buildOptions: sampleBuildOptions(), + }), + ).rejects.toThrow(/still cancelling initialization/); + await expect( + limitedManager.getOrCreateSession("replacement", { + root: sampleRoot, + buildOptions: sampleBuildOptions(), + }), + ).rejects.toThrow("Session capacity reached (1)"); + releaseBuild?.(); await expect(firstSession).rejects.toThrow(/disposed during initialization/); - await expect(secondSession).resolves.toMatchObject({ - getStatus: expect.any(Function), + + const replacement = await limitedManager.getOrCreateSession("replacement", { + root: sampleRoot, + buildOptions: sampleBuildOptions(), }); - expect((await secondSession).getStatus()).toBe("ready"); - expect(manager.getSession("pending")).toBe(await secondSession); + expect(replacement.getStatus()).toBe("ready"); } finally { buildSpy.mockRestore(); + limitedManager.disposeAll(); } }); - test("should allow immediate recreation after disposeAll cancels a pending session", async () => { + test("retains pending initialization capacity after disposeAll until it settles", async () => { + const limitedManager = new SessionManager({ maxSessions: 1, evictionIntervalMs: 0 }); const originalBuild = indexerBuild.buildProjectIndexIncremental; let releaseBuild: (() => void) | null = null; const buildGate = new Promise((resolve) => { @@ -1316,27 +1473,31 @@ describe("SessionManager", () => { }); try { - const firstSession = manager.getOrCreateSession("pending", { + const firstSession = limitedManager.getOrCreateSession("pending", { root: sampleRoot, buildOptions: sampleBuildOptions(), }); await Promise.resolve(); - manager.disposeAll(); - const secondSession = manager.getOrCreateSession("pending", { - root: sampleRoot, - buildOptions: sampleBuildOptions(), - }); - releaseBuild?.(); + limitedManager.disposeAll(); + await expect( + limitedManager.getOrCreateSession("replacement", { + root: sampleRoot, + buildOptions: sampleBuildOptions(), + }), + ).rejects.toThrow("Session capacity reached (1)"); + releaseBuild?.(); await expect(firstSession).rejects.toThrow(/disposed during initialization/); - await expect(secondSession).resolves.toMatchObject({ - getStatus: expect.any(Function), + + const replacement = await limitedManager.getOrCreateSession("replacement", { + root: sampleRoot, + buildOptions: sampleBuildOptions(), }); - expect((await secondSession).getStatus()).toBe("ready"); - expect(manager.getSession("pending")).toBe(await secondSession); + expect(replacement.getStatus()).toBe("ready"); } finally { buildSpy.mockRestore(); + limitedManager.disposeAll(); } }); @@ -1560,6 +1721,41 @@ describe("SessionManager", () => { expect(manager.getSessionIds()).toHaveLength(0); }); + test("keeps periodic expiration cleanup after disposeAll", () => { + vi.useFakeTimers(); + const reusableManager = new SessionManager({ evictionIntervalMs: 10 }); + const cleanupSpy = vi.spyOn(reusableManager, "cleanupExpired"); + + try { + reusableManager.disposeAll(); + vi.advanceTimersByTime(10); + expect(cleanupSpy).toHaveBeenCalledTimes(1); + } finally { + vi.clearAllTimers(); + } + }); + + test("stops periodic cleanup and prevents reuse after terminal disposal", async () => { + vi.useFakeTimers(); + const disposableManager = new SessionManager({ evictionIntervalMs: 10 }); + const cleanupSpy = vi.spyOn(disposableManager, "cleanupExpired"); + + try { + disposableManager.dispose(); + vi.advanceTimersByTime(10); + expect(cleanupSpy).not.toHaveBeenCalled(); + await expect( + disposableManager.getOrCreateSession("replacement", { + root: sampleRoot, + buildOptions: sampleBuildOptions(), + }), + ).rejects.toThrow("Session manager is disposed."); + await expect(disposableManager.warmup([])).rejects.toThrow("Session manager is disposed."); + } finally { + vi.useRealTimers(); + } + }); + test("should cleanup expired sessions", async () => { await manager.getOrCreateSession("session-1", { root: sampleRoot, @@ -1811,3 +2007,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..2f71c872 100644 --- a/tests/sqlite-query-bounds.test.ts +++ b/tests/sqlite-query-bounds.test.ts @@ -1,7 +1,7 @@ import fsp from "node:fs/promises"; import os from "node:os"; import path from "node:path"; -import { describe, expect, it } from "vitest"; +import { describe, expect, it, vi } from "vitest"; import { DatabaseSync } from "node:sqlite"; import { @@ -10,7 +10,15 @@ 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-exceeded query requests worker termination 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 withTempDb(run: (dbPath: string) => Promise): Promise { const root = await fsp.mkdtemp(path.join(os.tmpdir(), "cg-sqlite-bounds-")); @@ -18,11 +26,35 @@ async function withTempDb(run: (dbPath: string) => Promise): Promise try { await run(dbPath); } finally { - await fsp.rm(root, { recursive: true, force: true }); + await removeWithRetry(root); + } +} + +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) { + if (!(error instanceof Error) || !("code" in error)) throw error; + if (error.code !== "EBUSY" && error.code !== "ENOTEMPTY" && error.code !== "EPERM") throw error; + if (Date.now() > deadline) throw error; + await new Promise((resolve) => setTimeout(resolve, 100)); + } } } describe("SQLite query byte/cell bounds during iterate", () => { + it("retries Windows-style EPERM cleanup races", async () => { + const removeSpy = vi.spyOn(fsp, "rm").mockRejectedValueOnce(Object.assign(new Error("locked"), { code: "EPERM" })); + try { + await withTempDb(async () => {}); + expect(removeSpy).toHaveBeenCalledTimes(2); + } finally { + removeSpy.mockRestore(); + } + }); it("applies per-cell and cumulative caps before appending huge existing TEXT cells", async () => { await withTempDb(async (dbPath) => { const db = new DatabaseSync(dbPath); @@ -102,3 +134,64 @@ describe("SQLite query byte/cell bounds during iterate", () => { } }); }); + +describe("SQLite raw query execution deadline", () => { + it("rejects invalid deadlineMs values before selecting the worker execution path", async () => { + for (const deadlineMs of [NaN, -1, 1.5, 2_147_483_648, Infinity]) { + await expect(queryGraphSqliteRaw("missing.sqlite", "SELECT 1;", [], { deadlineMs })).rejects.toMatchObject({ + name: "RangeError", + message: "SQLite query deadlineMs must be a non-negative integer no greater than 2147483647.", + }); + } + }); + + 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; termination was requested."); + }); +}); diff --git a/tests/sqlite-query-deadline-fallback.test.ts b/tests/sqlite-query-deadline-fallback.test.ts new file mode 100644 index 00000000..e93189ae --- /dev/null +++ b/tests/sqlite-query-deadline-fallback.test.ts @@ -0,0 +1,153 @@ +import fsp from "node:fs/promises"; +import os from "node:os"; +import path from "node:path"; +import { DatabaseSync } from "node:sqlite"; +import { describe, expect, it, vi } from "vitest"; + +// Force every query in this file through the in-process fallback (as if the compiled +// worker asset were missing) so its deadline behavior -- and its documented +// limitation -- can be exercised directly, without disturbing the worker-backed +// deadline tests in sqlite-query-bounds.test.ts. +vi.mock("../src/sqlite/rawQueryWorkerPool.js", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + resolveRawSqlQueryWorkerPath: () => { + throw new Error("worker asset unavailable in this test"); + }, + }; +}); + +import { + SqliteQueryCancelledError, + SqliteQueryDeadlineExceededError as PublicSqliteQueryDeadlineExceededError, + SqliteQueryWorkerCleanupCapacityExceededError as PublicSqliteQueryWorkerCleanupCapacityExceededError, +} from "../src/sqlite.js"; +import { + SqliteQueryCancelledError as RootSqliteQueryCancelledError, + SqliteQueryDeadlineExceededError as RootSqliteQueryDeadlineExceededError, + SqliteQueryWorkerCleanupCapacityExceededError as RootSqliteQueryWorkerCleanupCapacityExceededError, +} from "../src/index.js"; +import { + queryGraphSqliteRaw, + SqliteQueryDeadlineExceededError, + SqliteQueryWorkerCleanupCapacityExceededError, +} from "../src/sqlite/query.js"; + +async function withTempDb(run: (dbPath: string) => Promise): Promise { + const root = await fsp.mkdtemp(path.join(os.tmpdir(), "cg-sqlite-deadline-fallback-")); + const dbPath = path.join(root, "graph.sqlite"); + try { + await run(dbPath); + } finally { + await fsp.rm(root, { recursive: true, force: true }); + } +} + +describe("SQLite raw query in-process deadline fallback", () => { + it("rejects invalid deadlineMs values before selecting the fallback execution path", async () => { + for (const deadlineMs of [NaN, -1, 1.5, 2_147_483_648, Infinity]) { + await expect(queryGraphSqliteRaw("missing.sqlite", "SELECT 1;", [], { deadlineMs })).rejects.toMatchObject({ + name: "RangeError", + message: "SQLite query deadlineMs must be a non-negative integer no greater than 2147483647.", + }); + } + }); + + it("still enforces the deadline between rows when the worker asset is unavailable", async () => { + await withTempDb(async (dbPath) => { + const db = new DatabaseSync(dbPath); + db.exec("CREATE TABLE t (x INTEGER);"); + db.close(); + + // Every outer row pays for a large nested recursive scan, so successive rows are + // spaced far enough apart in wall-clock time that a short deadline is guaranteed + // to trip between rows -- well before the query would otherwise finish. + const perRowSlowSql = + "WITH RECURSIVE outer_r(x) AS (SELECT 1 UNION ALL SELECT x + 1 FROM outer_r WHERE x < 500) " + + "SELECT x, (" + + " WITH RECURSIVE inner_r(y) AS (SELECT 1 UNION ALL SELECT y + 1 FROM inner_r WHERE y < 200000 + outer_r.x) " + + " SELECT count(*) FROM inner_r" + + ") FROM outer_r;"; + + await expect(queryGraphSqliteRaw(dbPath, perRowSlowSql, [], { deadlineMs: 20 })).rejects.toMatchObject({ + name: "SqliteQueryDeadlineExceededError", + message: expect.stringMatching(/exceeded its 20ms execution budget/), + }); + }); + }); + + it("does not interrupt a single blocking call whose entire cost is before the first row", async () => { + await withTempDb(async (dbPath) => { + const db = new DatabaseSync(dbPath); + db.exec("CREATE TABLE t (n INTEGER); INSERT INTO t (n) VALUES (1);"); + db.close(); + + // The whole cost of this query is inside one synchronous native step: SQLite + // must finish counting before it can return the single aggregate row. The + // fallback's per-row check cannot fire until that call returns, so -- unlike the + // worker-backed path -- this rejects only after running to completion, not + // within the deadline. That gap is the documented, unavoidable limitation of the + // fallback (see the doc comment on queryGraphSqliteRaw). + const slowBeforeFirstRowSql = + "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, slowBeforeFirstRowSql, [], { deadlineMs: 20 })).rejects.toMatchObject({ + name: "SqliteQueryDeadlineExceededError", + }); + const elapsed = Date.now() - start; + // A true execution deadline would reject close to 20ms; the fallback instead + // blocks for close to the query's full running time before it can even check. + expect(elapsed).toBeGreaterThan(200); + }); + }); + + it("rejects a zero-row query that completes after the fallback deadline", async () => { + await withTempDb(async (dbPath) => { + const db = new DatabaseSync(dbPath); + db.exec("CREATE TABLE t (n INTEGER);"); + db.close(); + + const slowEmptySql = + "WITH RECURSIVE spin(x) AS (SELECT 1 UNION ALL SELECT x + 1 FROM spin WHERE x < 8000000) " + + "SELECT x FROM spin WHERE x < 0;"; + + const start = Date.now(); + await expect(queryGraphSqliteRaw(dbPath, slowEmptySql, [], { deadlineMs: 20 })).rejects.toMatchObject({ + name: "SqliteQueryDeadlineExceededError", + }); + expect(Date.now() - start).toBeGreaterThan(200); + }); + }); + + it("still succeeds for 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); INSERT INTO t (n) VALUES (7);"); + db.close(); + + const result = await queryGraphSqliteRaw(dbPath, "SELECT n FROM t;", [], { deadlineMs: 5_000 }); + expect(result.rows).toEqual([[7]]); + expect(result.truncated).toBeFalsy(); + }); + }); + + it("exports named cancellation, deadline, and capacity errors from public library barrels", () => { + expect(PublicSqliteQueryDeadlineExceededError).toBe(SqliteQueryDeadlineExceededError); + expect(PublicSqliteQueryWorkerCleanupCapacityExceededError).toBe(SqliteQueryWorkerCleanupCapacityExceededError); + expect(RootSqliteQueryDeadlineExceededError).toBe(SqliteQueryDeadlineExceededError); + expect(RootSqliteQueryCancelledError).toBe(SqliteQueryCancelledError); + expect(RootSqliteQueryWorkerCleanupCapacityExceededError).toBe(SqliteQueryWorkerCleanupCapacityExceededError); + + const deadline = new PublicSqliteQueryDeadlineExceededError(250); + expect(deadline).toBeInstanceOf(Error); + expect(deadline.name).toBe("SqliteQueryDeadlineExceededError"); + expect(deadline.message).toBe("SQLite query exceeded its 250ms execution budget; termination was requested."); + expect(new SqliteQueryCancelledError().message).toBe("SQLite query was cancelled."); + expect(new PublicSqliteQueryWorkerCleanupCapacityExceededError(2).message).toContain( + "SQLite query worker capacity is exhausted", + ); + }); +}); diff --git a/tests/type-hierarchy.test.ts b/tests/type-hierarchy.test.ts index f38013c2..9ccf96fd 100644 --- a/tests/type-hierarchy.test.ts +++ b/tests/type-hierarchy.test.ts @@ -205,4 +205,36 @@ describe("type hierarchy", () => { reason: expect.stringContaining("abstract"), }); }); + it("pins omission counts at and just past the limit for type hierarchy and implementations", async () => { + const { index, graph, byName } = await hierarchyFixture(); + const specialized = byName.get("SpecializedWorker"); + expect(specialized).toBeDefined(); + + const atSuperLimit = findTypeHierarchy(graph, specialized!.id, "super", { depth: 3, limit: 3 }); + expect(atSuperLimit).toMatchObject({ status: "ok", omitted: 0 }); + if (atSuperLimit.status === "ok") { + expect(atSuperLimit.relations).toHaveLength(3); + } + + const pastSuperLimit = findTypeHierarchy(graph, specialized!.id, "super", { depth: 3, limit: 2 }); + expect(pastSuperLimit).toMatchObject({ status: "ok", omitted: 1 }); + if (pastSuperLimit.status === "ok") { + expect(pastSuperLimit.relations).toHaveLength(2); + } + + const service = byName.get("Service"); + expect(service).toBeDefined(); + + const atImplLimit = findImplementations(index, graph, service!.id, { limit: 2 }); + expect(atImplLimit).toMatchObject({ status: "ok", omitted: 0 }); + if (atImplLimit.status === "ok") { + expect(atImplLimit.implementations).toHaveLength(2); + } + + const pastImplLimit = findImplementations(index, graph, service!.id, { limit: 1 }); + expect(pastImplLimit).toMatchObject({ status: "ok", omitted: 1 }); + if (pastImplLimit.status === "ok") { + expect(pastImplLimit.implementations).toHaveLength(1); + } + }); }); diff --git a/tests/viewer.test.ts b/tests/viewer.test.ts index 5e6a2f7f..b88863b0 100644 --- a/tests/viewer.test.ts +++ b/tests/viewer.test.ts @@ -278,4 +278,52 @@ describe("viewer server", () => { expect(() => createViewerServer({ graph: escapedGraph, root })).toThrow(/outside project root/i); }); + test("returns 500 when statSync or fstatSync throws during GET and continues serving subsequent requests", async () => { + const { root, graphPath } = await createViewerFixture(); + const server = await startViewerServer({ graph: graphPath, port: 0, root }); + servers.push(server.server); + + // 1. Test fstatSync throwing during GET /graph.json + let throwFstat = true; + const originalFstatSync = fs.fstatSync; + const fstatSpy = vi.spyOn(fs, "fstatSync").mockImplementation((...args) => { + if (throwFstat) { + throw new Error("Simulated filesystem fstatSync error"); + } + return originalFstatSync(...args); + }); + + try { + const firstFstatResponse = await request(server.server, "/graph.json"); + expect(firstFstatResponse.statusCode).toBe(500); + + throwFstat = false; + const secondFstatResponse = await request(server.server, "/graph.json"); + expect(secondFstatResponse.statusCode).toBe(200); + expect(secondFstatResponse.body).toContain('"nodes":[]'); + } finally { + fstatSpy.mockRestore(); + } + + // 2. Test statSync throwing during GET / + let throwStat = true; + const originalStatSync = fs.statSync; + const statSpy = vi.spyOn(fs, "statSync").mockImplementation((...args) => { + if (throwStat) { + throw new Error("Simulated filesystem statSync error"); + } + return originalStatSync(...args); + }); + + try { + const firstStatResponse = await request(server.server, "/"); + expect(firstStatResponse.statusCode).toBe(500); + + throwStat = false; + const secondStatResponse = await request(server.server, "/"); + expect(secondStatResponse.statusCode).toBe(200); + } finally { + statSpy.mockRestore(); + } + }); }); diff --git a/tests/workspace-symbols.test.ts b/tests/workspace-symbols.test.ts index 202a10d7..109dea6f 100644 --- a/tests/workspace-symbols.test.ts +++ b/tests/workspace-symbols.test.ts @@ -412,4 +412,18 @@ describe("workspace symbol lookup", () => { tool_workspaceSymbols(root, { query: "Service" }, { session, buildOptions: { cache: "off" } }), ).rejects.toThrow("cannot combine a prebuilt session with buildOptions"); }); + it("pins omission counts at and just past the limit for workspace symbols", async () => { + const all = await workspaceSymbols(index, { query: "Service", limit: 50 }); + const total = all.symbols.length; + expect(total).toBeGreaterThanOrEqual(2); + expect(all.omitted).toBe(0); + + const atLimit = await workspaceSymbols(index, { query: "Service", limit: total }); + expect(atLimit.symbols).toHaveLength(total); + expect(atLimit.omitted).toBe(0); + + const pastLimit = await workspaceSymbols(index, { query: "Service", limit: total - 1 }); + expect(pastLimit.symbols).toHaveLength(total - 1); + expect(pastLimit.omitted).toBe(1); + }); });