From b6f1e9afb9e2ab935e4d4c00f114580320af80e6 Mon Sep 17 00:00:00 2001 From: codegraph Date: Sat, 15 Aug 2026 12:17:09 -0400 Subject: [PATCH 1/4] refactor: route omission counting through shared bound helper --- src/agent/explore.ts | 11 +++++++---- src/indexer/type-hierarchy.ts | 20 ++++++++++++-------- src/indexer/workspace-symbols.ts | 6 ++++-- tests/agent-explore.test.ts | 29 +++++++++++++++++++++++++++++ tests/type-hierarchy.test.ts | 32 ++++++++++++++++++++++++++++++++ tests/workspace-symbols.test.ts | 14 ++++++++++++++ 6 files changed, 98 insertions(+), 14 deletions(-) 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/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 15b95b26..046b7e7c 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/tests/agent-explore.test.ts b/tests/agent-explore.test.ts index e058018d..0f002e4b 100644 --- a/tests/agent-explore.test.ts +++ b/tests/agent-explore.test.ts @@ -841,4 +841,33 @@ 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/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/workspace-symbols.test.ts b/tests/workspace-symbols.test.ts index 636b5b9d..82c797c0 100644 --- a/tests/workspace-symbols.test.ts +++ b/tests/workspace-symbols.test.ts @@ -384,4 +384,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); + }); }); From e9b4352d71dc26a1636c6a10eb4b225de6c761a2 Mon Sep 17 00:00:00 2001 From: codegraph Date: Sat, 15 Aug 2026 12:17:10 -0400 Subject: [PATCH 2/4] test: cover MCP session resource bounds and teardown --- tests/agent-search.test.ts | 43 ++++++++++ tests/agent-session.test.ts | 29 +++++++ tests/mcp-server.test.ts | 162 ++++++++++++++++++++++++++++++++++++ tests/query-index.test.ts | 26 +++++- tests/viewer.test.ts | 48 +++++++++++ 5 files changed, 307 insertions(+), 1 deletion(-) diff --git a/tests/agent-search.test.ts b/tests/agent-search.test.ts index f9ff1fc2..89f8e95a 100644 --- a/tests/agent-search.test.ts +++ b/tests/agent-search.test.ts @@ -838,4 +838,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 86dd8f54..7e9dc130 100644 --- a/tests/agent-session.test.ts +++ b/tests/agent-session.test.ts @@ -1,3 +1,5 @@ +import { disposeSessionQueryIndex, ensureSessionQueryIndex } from "../src/agent/query-index/sessionStore.js"; +import * as updateModule from "../src/agent/query-index/update.js"; import fs from "node:fs/promises"; import { brotliCompressSync, brotliDecompressSync, constants as zlibConstants } from "node:zlib"; import { createHash } from "node:crypto"; @@ -874,4 +876,31 @@ 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; + }); + + try { + await expect(ensureSessionQueryIndex(session, snapshot)).rejects.toThrow( + /Query index generation changed repeatedly while loading/i, + ); + expect(attempts).toBe(3); + } finally { + ensureQueryIndexSpy.mockRestore(); + } + }); }); + +}); \ No newline at end of file diff --git a/tests/mcp-server.test.ts b/tests/mcp-server.test.ts index 259359a4..29a2274a 100644 --- a/tests/mcp-server.test.ts +++ b/tests/mcp-server.test.ts @@ -1,3 +1,5 @@ +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 os from "node:os"; @@ -2754,3 +2756,163 @@ describe("codegraph MCP handlers", () => { function normalizeSqlitePath(value: unknown): string { return typeof value === "string" ? value.replace(/\\/g, "/") : ""; } + + +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(); + } + }); +}); + +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(); + } + }); +}); diff --git a/tests/query-index.test.ts b/tests/query-index.test.ts index 6a14c732..303c1016 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"; @@ -808,4 +809,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/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(); + } + }); }); From 17699a854402ae88850da0be4048dc4274e020c1 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sat, 15 Aug 2026 23:01:11 -0400 Subject: [PATCH 3/4] fix: bound MCP HTTP and SQLite execution --- docs/library-api.md | 2 +- docs/mcp.md | 2 +- src/mcp/http.ts | 69 ++++++++++++++++++------------- src/mcp/server.ts | 4 +- src/sqlite/query.ts | 60 +++++++++++++++++++++++---- src/sqlite/rawQueryWorker.ts | 33 +++++++++++++++ src/sqlite/rawQueryWorkerPool.ts | 43 +++++++++++++++++++ tests/mcp-server.test.ts | 68 ++++++++++++++++++++++++++---- tests/sqlite-query-bounds.test.ts | 52 ++++++++++++++++++++++- 9 files changed, 283 insertions(+), 50 deletions(-) create mode 100644 src/sqlite/rawQueryWorker.ts create mode 100644 src/sqlite/rawQueryWorkerPool.ts diff --git a/docs/library-api.md b/docs/library-api.md index 098841c7..1a9e44a5 100644 --- a/docs/library-api.md +++ b/docs/library-api.md @@ -714,7 +714,7 @@ 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` and rejects mutating SQL. Its defaults bound rows, cells, response bytes, and execution to 10 seconds; callers can further tighten `{ maxRows, maxBytes, maxCellBytes, deadlineMs }`. ## SQL artifact facts diff --git a/docs/mcp.md b/docs/mcp.md index 8d5d716c..77f9e092 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -200,7 +200,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/src/mcp/http.ts b/src/mcp/http.ts index 633963ca..c1870fcb 100644 --- a/src/mcp/http.ts +++ b/src/mcp/http.ts @@ -28,39 +28,52 @@ export async function readJsonRequestBody( return { status: "too_large" }; } - const chunks: Buffer[] = []; - let bytes = 0; - let timedOut = false; - const deadline = setTimeout(() => { - timedOut = true; - request.destroy(); - }, timeoutMs); - deadline.unref?.(); - try { - for await (const chunk of request) { + 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) { - request.resume(); - return { status: "too_large" }; + settle({ status: "too_large" }, true); + return; } chunks.push(buffer); - } - } catch { - if (timedOut) return { status: "timeout" }; - return { status: "invalid_json" }; - } finally { - clearTimeout(deadline); - } - if (timedOut) return { status: "timeout" }; - - 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" }; - } + }; + 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); + + request.on("data", onData); + request.once("end", onEnd); + request.once("error", onFailure); + request.once("aborted", onFailure); + }); } export function emptyAllowedHostHeaderRules(): AllowedHostHeaderRules { diff --git a/src/mcp/server.ts b/src/mcp/server.ts index 012828ce..7700b06e 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -71,7 +71,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"; @@ -980,8 +979,9 @@ function createCodegraphMcpHandlersForSession( } const result = await queryGraphSqliteRaw(realSqlitePath, request.query, request.params ?? [], { maxRows: normalizeSqliteRowLimit(request.limit), + maxBytes: DEFAULT_SQLITE_BYTE_LIMIT, }); - return { ...boundRawSqlResult(result, DEFAULT_SQLITE_BYTE_LIMIT), freshness: artifactFreshness }; + return { ...result, truncated: Boolean(result.truncated), freshness: artifactFreshness }; }, refresh_index: async (request) => { diff --git a/src/sqlite/query.ts b/src/sqlite/query.ts index d1e740ff..6cf062a7 100644 --- a/src/sqlite/query.ts +++ b/src/sqlite/query.ts @@ -8,13 +8,22 @@ import { MAX_SQLITE_ROW_LIMIT, normalizeSqliteRowLimit, } from "./rowBounds.js"; +import { + resolveRawSqlQueryWorkerPath, + runRawSqlQueryInWorker, + SqliteQueryDeadlineExceededError, +} from "./rawQueryWorkerPool.js"; export { queryGraphSqlite } from "./canned-query.js"; +export { SqliteQueryDeadlineExceededError }; + +export const DEFAULT_SQLITE_QUERY_DEADLINE_MS = 10_000; export type QueryGraphSqliteRawOptions = { maxRows?: number | undefined; maxBytes?: number | undefined; maxCellBytes?: number | undefined; + deadlineMs?: number | undefined; }; export async function queryGraphSqliteRaw( @@ -22,21 +31,47 @@ export async function queryGraphSqliteRaw( sql: string, params: Array = [], options?: QueryGraphSqliteRawOptions, +): Promise { + const maxRows = normalizeSqliteRowLimit(options?.maxRows ?? MAX_SQLITE_ROW_LIMIT); + const maxBytes = options?.maxBytes ?? DEFAULT_SQLITE_BYTE_LIMIT; + const maxCellBytes = options?.maxCellBytes ?? MAX_SQLITE_CELL_BYTES; + const deadlineMs = options?.deadlineMs ?? DEFAULT_SQLITE_QUERY_DEADLINE_MS; + + try { + resolveRawSqlQueryWorkerPath(); + } catch { + return await queryGraphSqliteRawInProcessBounded(outputPath, sql, params, { + maxRows, + maxBytes, + maxCellBytes, + deadlineMs, + }); + } + + return await runRawSqlQueryInWorker({ outputPath, sql, params, maxRows, maxBytes, maxCellBytes }, deadlineMs); +} + +async function queryGraphSqliteRawInProcessBounded( + outputPath: string, + sql: string, + params: Array, + bounds: { maxRows: number; maxBytes: number; maxCellBytes: number; deadlineMs: number }, ): Promise { return await withReadOnlySqliteDatabase(outputPath, (db) => { try { const stmt = db.prepare(sql); assertReadOnlyQueryStatement(stmt); const columns = stmt.columns().map((col) => col.name); - const maxRows = normalizeSqliteRowLimit(options?.maxRows ?? MAX_SQLITE_ROW_LIMIT); - const maxBytes = options?.maxBytes ?? DEFAULT_SQLITE_BYTE_LIMIT; - const maxCellBytes = options?.maxCellBytes ?? MAX_SQLITE_CELL_BYTES; - - // 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, + const deadlineAt = Date.now() + bounds.deadlineMs; + const rows = withPerRowDeadline( + stmt.raw().iterate(params) as Iterable>, + deadlineAt, + bounds.deadlineMs, + ); + return collectBoundedRawSqlRows(columns, rows, { + maxRows: bounds.maxRows, + maxBytes: bounds.maxBytes, + maxCellBytes: bounds.maxCellBytes, }); } catch (error) { if (isReadOnlySqliteError(error)) { @@ -46,3 +81,10 @@ export async function queryGraphSqliteRaw( } }); } + +function* withPerRowDeadline(rows: Iterable, deadlineAt: number, deadlineMs: number): Generator { + for (const row of rows) { + if (Date.now() > deadlineAt) throw new SqliteQueryDeadlineExceededError(deadlineMs); + yield row; + } +} diff --git a/src/sqlite/rawQueryWorker.ts b/src/sqlite/rawQueryWorker.ts new file mode 100644 index 00000000..dc74a7a7 --- /dev/null +++ b/src/sqlite/rawQueryWorker.ts @@ -0,0 +1,33 @@ +import { isReadOnlySqliteError } from "../sqlite-driver.js"; +import type { RawSqlResult } from "./types.js"; +import { assertReadOnlyQueryStatement, withReadOnlySqliteDatabase } from "./database.js"; +import { collectBoundedRawSqlRows } from "./rowBounds.js"; + +export type RawQueryWorkerTask = { + outputPath: string; + sql: string; + params: Array; + maxRows: number; + maxBytes: number | undefined; + maxCellBytes: number | undefined; +}; + +export default async function runRawQueryWorkerTask(task: RawQueryWorkerTask): Promise { + return await withReadOnlySqliteDatabase(task.outputPath, (db) => { + try { + const statement = db.prepare(task.sql); + assertReadOnlyQueryStatement(statement); + const columns = statement.columns().map((column) => column.name); + return collectBoundedRawSqlRows(columns, statement.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..e40432f4 --- /dev/null +++ b/src/sqlite/rawQueryWorkerPool.ts @@ -0,0 +1,43 @@ +import fs from "node:fs"; +import path from "node:path"; +import { fileURLToPath } from "node:url"; +import { Piscina } from "piscina"; +import { findPackageRoot } from "../util/packageInfo.js"; +import type { RawSqlResult } from "./types.js"; +import type { RawQueryWorkerTask } from "./rawQueryWorker.js"; + +export class SqliteQueryDeadlineExceededError extends Error { + constructor(deadlineMs: number) { + super(`SQLite query exceeded its ${deadlineMs}ms execution budget and was terminated.`); + this.name = "SqliteQueryDeadlineExceededError"; + } +} + +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; + throw new Error(`Raw SQLite query worker file not found: ${compiled}`); +} + +export async function runRawSqlQueryInWorker(task: RawQueryWorkerTask, deadlineMs: number): Promise { + const pool = new Piscina({ + filename: resolveRawSqlQueryWorkerPath(), + minThreads: 1, + maxThreads: 1, + idleTimeout: 5_000, + }); + try { + return (await pool.run(task, { signal: AbortSignal.timeout(deadlineMs) })) as RawSqlResult; + } catch (error) { + if (error instanceof Error && error.name === "AbortError") { + throw new SqliteQueryDeadlineExceededError(deadlineMs); + } + throw error; + } finally { + void pool.destroy().catch(() => {}); + } +} diff --git a/tests/mcp-server.test.ts b/tests/mcp-server.test.ts index 29a2274a..da070078 100644 --- a/tests/mcp-server.test.ts +++ b/tests/mcp-server.test.ts @@ -1155,6 +1155,64 @@ describe("codegraph MCP handlers", () => { } }); + 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 }>((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)) }); + } 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).toBe(408); + expect(readObject(response.payload.error).message).toBe("MCP request body timed out"); + } finally { + await httpServer.close(); + } + }); + 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"); @@ -2756,8 +2814,6 @@ describe("codegraph MCP handlers", () => { function normalizeSqlitePath(value: unknown): string { return typeof value === "string" ? value.replace(/\\/g, "/") : ""; } - - 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-")); @@ -2893,9 +2949,7 @@ describe("MCP transport isolation regressions (S8)", () => { await call1Closed.promise; const call2 = await call2Promise; expect(call2.response.status).toBe(200); - expect(readToolJsonResult(call2.payload).symbols).toEqual([ - expect.objectContaining({ name: "secondarySymbol" }), - ]); + expect(readToolJsonResult(call2.payload).symbols).toEqual([expect.objectContaining({ name: "secondarySymbol" })]); const call3 = await postMcpJson( httpServer.url, @@ -2908,9 +2962,7 @@ describe("MCP transport isolation regressions (S8)", () => { sessionId, ); expect(call3.response.status).toBe(200); - expect(readToolJsonResult(call3.payload).symbols).toEqual([ - expect.objectContaining({ name: "validateUser" }), - ]); + expect(readToolJsonResult(call3.payload).symbols).toEqual([expect.objectContaining({ name: "validateUser" })]); } finally { await httpServer.close(); } diff --git a/tests/sqlite-query-bounds.test.ts b/tests/sqlite-query-bounds.test.ts index 24654e99..539205da 100644 --- a/tests/sqlite-query-bounds.test.ts +++ b/tests/sqlite-query-bounds.test.ts @@ -10,7 +10,7 @@ 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"; async function withTempDb(run: (dbPath: string) => Promise): Promise { const root = await fsp.mkdtemp(path.join(os.tmpdir(), "cg-sqlite-bounds-")); @@ -22,6 +22,21 @@ async function withTempDb(run: (dbPath: string) => Promise): Promise } } +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") throw error; + if (Date.now() > deadline) throw error; + await new Promise((resolve) => setTimeout(resolve, 100)); + } + } +} + describe("SQLite query byte/cell bounds during iterate", () => { it("applies per-cell and cumulative caps before appending huge existing TEXT cells", async () => { await withTempDb(async (dbPath) => { @@ -102,3 +117,38 @@ describe("SQLite query byte/cell bounds during iterate", () => { } }); }); + +describe("SQLite raw query execution deadline", () => { + it("terminates an over-budget query without delaying a following query", async () => { + const root = await fsp.mkdtemp(path.join(os.tmpdir(), "cg-sqlite-deadline-")); + const dbPath = path.join(root, "graph.sqlite"); + try { + const db = new DatabaseSync(dbPath); + db.exec("CREATE TABLE values_table (n INTEGER);"); + db.prepare("INSERT INTO values_table (n) VALUES (?)").run(42); + db.close(); + + const slowSql = + "WITH RECURSIVE spin(x) AS (SELECT 1 UNION ALL SELECT x + 1 FROM spin WHERE x < 8000000) " + + "SELECT count(*) FROM spin;"; + const startedAt = Date.now(); + await expect(queryGraphSqliteRaw(dbPath, slowSql, [], { deadlineMs: 100 })).rejects.toMatchObject({ + name: "SqliteQueryDeadlineExceededError", + message: expect.stringMatching(/exceeded its 100ms execution budget/), + }); + expect(Date.now() - startedAt).toBeLessThan(2_000); + + const result = await queryGraphSqliteRaw(dbPath, "SELECT n FROM values_table;", [], { deadlineMs: 5_000 }); + expect(result.rows).toEqual([[42]]); + expect(result.truncated).toBeFalsy(); + } finally { + await removeWithRetry(root); + } + }); + + it("exports a named deadline error for callers", () => { + const error = new SqliteQueryDeadlineExceededError(250); + expect(error.name).toBe("SqliteQueryDeadlineExceededError"); + expect(error.message).toBe("SQLite query exceeded its 250ms execution budget and was terminated."); + }); +}); From 07aa2c5670a9effdc3a4f91ce3d1ad3d38919699 Mon Sep 17 00:00:00 2001 From: Luke Zehrung Date: Sun, 16 Aug 2026 00:03:59 -0400 Subject: [PATCH 4/4] fix: repair MCP session bounds and SQLite deadline contract regressions Three post-review defects in the MCP HTTP transport and raw SQLite query path: 1. Legacy initialize capacity reservations leaked whenever the SDK transport answered a pre-session 4xx without throwing (most notably Accept header validation): onsessioninitialized never fired to release it, and the catch block only covered thrown errors. Release the reservation whenever handleLegacyMcpSessionRequest resolves without a session having been initialized. 2. transport.onerror deleted the whole legacy session for any per-request SDK validation error (bad Accept, wrong Content-Type, malformed JSON, unsupported protocol version, ...), even though those already answered their own request and left the transport healthy. Session teardown is now driven by onclose alone; onerror only logs. 3. The in-process SQLite query fallback (used when the compiled worker asset can't be located) only checks its deadline between already-produced rows, so a statement slow to produce its first row isn't bounded by it. node:sqlite has no interrupt API, so true enforcement requires the worker thread this fallback exists because it couldn't find; corrected the public contract instead (JSDoc, docs/library-api.md, and a one-time degraded-mode log) so the gap is documented and observable rather than silently implied away. --- docs/library-api.md | 2 +- src/mcp/server.ts | 17 +++- src/sqlite/query.ts | 41 ++++++++ src/sqlite/rawQueryWorkerPool.ts | 15 +++ tests/mcp-server.test.ts | 82 +++++++++++++++ tests/sqlite-query-deadline-fallback.test.ts | 102 +++++++++++++++++++ 6 files changed, 257 insertions(+), 2 deletions(-) create mode 100644 tests/sqlite-query-deadline-fallback.test.ts diff --git a/docs/library-api.md b/docs/library-api.md index 1a9e44a5..1a9b89c8 100644 --- a/docs/library-api.md +++ b/docs/library-api.md @@ -714,7 +714,7 @@ 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. Its defaults bound rows, cells, response bytes, and execution to 10 seconds; callers can further tighten `{ maxRows, maxBytes, maxCellBytes, deadlineMs }`. +`queryGraphSqliteRaw()` is intentionally read-only. It accepts result-producing statements such as `SELECT` and `PRAGMA` and rejects mutating SQL. Its defaults bound rows, cells, and response bytes, and callers can further tighten `{ maxRows, maxBytes, maxCellBytes, deadlineMs }`. The 10-second default execution budget (`deadlineMs`) is enforced by running the query in a dedicated worker thread that is force-terminated on expiry, so it interrupts a query even mid-execution; in a degraded install where that worker asset cannot be located, the query instead runs in-process under a weaker per-row check that cannot interrupt a single blocking native call (a logged, one-time-per-process condition). ## SQL artifact facts diff --git a/src/mcp/server.ts b/src/mcp/server.ts index 7700b06e..b3e33444 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -1377,9 +1377,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); @@ -1388,6 +1395,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); diff --git a/src/sqlite/query.ts b/src/sqlite/query.ts index 6cf062a7..79ea826b 100644 --- a/src/sqlite/query.ts +++ b/src/sqlite/query.ts @@ -17,6 +17,8 @@ import { export { queryGraphSqlite } from "./canned-query.js"; export { SqliteQueryDeadlineExceededError }; +/** 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; export type QueryGraphSqliteRawOptions = { @@ -26,6 +28,32 @@ export type QueryGraphSqliteRawOptions = { deadlineMs?: number | undefined; }; +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`). On expiry the worker thread is + * terminated outright, which stops the query even while it is blocked inside a single + * synchronous `DatabaseSync` call — a slow non-recursive statement (large join, + * `ORDER BY random()`, a recursive CTE, ...) cannot hold the deadline hostage. + * + * 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 — + * true enforcement genuinely requires the separate worker thread this fallback exists + * because it could not find. The per-row check is therefore strictly weaker, not just a + * smaller budget: it is only evaluated between rows the native iterator has already + * produced, so a statement that is slow to produce its very first row (a full scan + * before any match, an aggregate over a large recursive CTE, ...) blocks for its full + * cost before the deadline is ever checked. This fallback exists to keep the common + * case usable in a degraded install, not as a substitute for the worker deadline; a + * warning is logged once per process when it activates so a degraded install is + * observable rather than silently under-enforcing its documented time budget. + */ export async function queryGraphSqliteRaw( outputPath: string, sql: string, @@ -40,6 +68,14 @@ export async function queryGraphSqliteRaw( 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, @@ -51,6 +87,8 @@ export async function queryGraphSqliteRaw( return await runRawSqlQueryInWorker({ outputPath, sql, params, maxRows, maxBytes, maxCellBytes }, deadlineMs); } +/** Degraded fallback for `queryGraphSqliteRaw` — see its doc comment for the enforcement + * caveat this path cannot avoid. */ async function queryGraphSqliteRawInProcessBounded( outputPath: string, sql: string, @@ -82,6 +120,9 @@ async function queryGraphSqliteRawInProcessBounded( }); } +/** Throws once the wall-clock deadline has passed between two already-produced rows. + * See the fallback caveat on `queryGraphSqliteRaw`: a statement slow to produce its + * first row is not bounded here — only slow-*between*-rows iteration is caught. */ function* withPerRowDeadline(rows: Iterable, deadlineAt: number, deadlineMs: number): Generator { for (const row of rows) { if (Date.now() > deadlineAt) throw new SqliteQueryDeadlineExceededError(deadlineMs); diff --git a/src/sqlite/rawQueryWorkerPool.ts b/src/sqlite/rawQueryWorkerPool.ts index e40432f4..6cf01c12 100644 --- a/src/sqlite/rawQueryWorkerPool.ts +++ b/src/sqlite/rawQueryWorkerPool.ts @@ -13,6 +13,11 @@ export class SqliteQueryDeadlineExceededError extends Error { } } +/** Resolves the compiled worker entry: 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 (running this module from `src/`, where only + * `dist/` is built). Throwing here is the trigger `queryGraphSqliteRaw` uses to fall + * back to the strictly weaker in-process per-row deadline check. */ export function resolveRawSqlQueryWorkerPath(): string { const selfDirectory = path.dirname(fileURLToPath(import.meta.url)); const sibling = path.resolve(selfDirectory, "rawQueryWorker.js"); @@ -23,6 +28,16 @@ export function resolveRawSqlQueryWorkerPath(): string { throw new Error(`Raw SQLite query worker file not found: ${compiled}`); } +/** + * Runs a single bounded raw SQL read in a dedicated worker thread with a hard execution + * deadline. On expiry, Piscina's `signal` option force-terminates the worker thread and + * rejects immediately — the caller never waits longer than `deadlineMs`, regardless of + * how long the underlying query actually takes, because termination does not need the + * blocked thread's cooperation. `pool.destroy()` is fire-and-forget rather than awaited, + * so an orphaned worker still finishing one already-in-flight synchronous native call + * never delays this call's rejection or a subsequent query against the same file + * (concurrent read-only SQLite connections do not block each other). + */ export async function runRawSqlQueryInWorker(task: RawQueryWorkerTask, deadlineMs: number): Promise { const pool = new Piscina({ filename: resolveRawSqlQueryWorkerPath(), diff --git a/tests/mcp-server.test.ts b/tests/mcp-server.test.ts index da070078..23345aef 100644 --- a/tests/mcp-server.test.ts +++ b/tests/mcp-server.test.ts @@ -2968,3 +2968,85 @@ describe("MCP transport isolation regressions (S8)", () => { } }); }); + +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/sqlite-query-deadline-fallback.test.ts b/tests/sqlite-query-deadline-fallback.test.ts new file mode 100644 index 00000000..917bb8a2 --- /dev/null +++ b/tests/sqlite-query-deadline-fallback.test.ts @@ -0,0 +1,102 @@ +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 { queryGraphSqliteRaw, SqliteQueryDeadlineExceededError } 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("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) " + + " 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("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 a named deadline error for callers regardless of which path produced it", () => { + const error = new SqliteQueryDeadlineExceededError(250); + expect(error).toBeInstanceOf(Error); + expect(error.name).toBe("SqliteQueryDeadlineExceededError"); + expect(error.message).toBe("SQLite query exceeded its 250ms execution budget and was terminated."); + }); +});