From cad26c6bfb736039c8ed5743da650a44792a083c Mon Sep 17 00:00:00 2001 From: Michael Nguyen <13559011+nguyenm7@users.noreply.github.com> Date: Fri, 2 Oct 2026 05:45:27 -0700 Subject: [PATCH] fix(tool-gateway): bound MCP discovery memory and concurrency (#14864) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Thinking Path > - Paperclip is the open source app people use to manage AI agents for work. > - Agents discover governed tools through the MCP gateway. > - A listing repeated policy and full run-row reads for each catalog tool. > - Parallel listings multiplied those allocations during run startup. > - One 900-tool baseline listing used 11,489 queries and about 1.7 GiB of extra heap in a fixture. > - This pull request shares reads within a listing and bounds whole listings across the process. > - The benefit is lower discovery memory use while execution still checks current policy. ## Linked Issues or Issue Description Refs #13115. Its on-demand target change affects the same listing loop. **What happened?** MCP discovery repeated roughly 13 reads per tool. Full run snapshots and repeated connection configurations caused large allocations. Per-listing bounds alone did not limit concurrent listings across gateways. **Expected behavior** Discovery reads shared inputs once per listing. The process bounds active and queued listings. Catalog payload size and policy evaluation still grow with the catalog. Tool execution checks current access rules. **Steps to reproduce** Create a company with a remote MCP connection, 900 catalog tools, large schemas, and a large run snapshot. Send concurrent tools/list requests using a run-bound gateway token. Run the committed benchmark for a deterministic reproduction. **Paperclip version or commit** Baseline: f2e0f196308629ef05c7f65782243e1713b208f0. Measured on Node 26.4.0 on macOS. ## What Changed - Keep Michael Nguyen's per-listing policy cache, scalar policy reads, 16-decision bound, and catalog/connection split under repeatable read. - Admit two whole listings per process and queue at most 32. Return 503 with tool_discovery_busy when full. - Propagate disconnects to discovery. Stop scheduling new reads and drain started reads before freeing the slot. - Project context IDs in gateway authentication and GitHub runtime discovery. Avoid loading task descriptions and results. - Keep discovery audit counts and a SHA-256 digest instead of the full name list. Retain access and activity audit records. - Sweep expired named gateway tokens at startup and on the existing scheduler. Delete at most 500 per pass. Add an idempotent expiry-index migration. Resolve audit token references atomically so cleanup cannot break admitted requests. - Return GET 405 with Allow: POST on stateless MCP gateway and runtime-tools endpoints. Keep runtime-tools authentication. - Add red-green regressions, HTTP concurrency and policy-revocation coverage, and browser approval assertions. - Commit the benchmark harness, raw results, and resource-bound documentation. ## Verification - Baseline listing regressions failed with 722 queries for 50 tools and 6,422 for 500. The connection-row duplication regression also failed. - Follow-up regressions failed before the fixes for full snapshot reads, abandoned listings, full name-list audits, and expired tokens. - Listing and scheduler tests: 14 pass. Existing gateway/policy suites passed after preserving the cleanup return contract. - Two HTTP journeys pass: initialize, GET/SSE rejection, 16 concurrent 500-tool listings, provider call, policy revocation, and denied retry. Token cleanup during provider dispatch also completes successfully and blocks subsequent requests. - pnpm test:e2e tests/e2e/mcp-user-stories.spec.ts --grep '@mcp-runnable': 8 pass. The approval journey clicks Allow once in the browser. Review screenshots wait for loaded content. - pnpm -r typecheck: passes. pnpm build: passes. - The general server group completed with 14,755 passes and 18 failures. The 17 startup mock failures were fixed; all 21 startup tests pass on rerun. The one Discord timing failure passed on the unchanged baseline and on rerun (74 tests). UI and CLI groups pass 7,632 tests. Shared and skill groups pass 853 tests. The remaining database and adapter groups pass 3,010 tests with one worker after a macOS shared-memory limit interrupted a parallel run. All 149 serialized route files pass (2,762 tests). - At 900 tools, one listing falls from 11,489 to 36 queries and from about 1.7 GiB to 37 MiB of extra heap. Sixteen concurrent listings used 169–187 MiB of extra heap. Four connections used 39 queries per listing. - Latest-head verification: 54 successful checks and two expected skips on 5c0793090cec561a00c2bbc7ea5b19466db91018. Fresh Greptile review: 5/5 with no unresolved threads. - Reproduce with server/scripts/benchmark-tool-gateway-listing.ts. See doc/mcp-discovery-performance.md and doc/benchmarks/2026-10-01-mcp-discovery.json. ## Risks - The process-wide FIFO queue can increase discovery latency. Excess callers must retry 503 responses. - Cancellation applies to discovery. Started database reads finish before their slot is released. - Audit consumers must use visibleToolCount and visibleToolsHash instead of visibleTools. - The expiry index can briefly lock the token table during migration. Sweeps preserve unexpired and non-expiring tokens. - Measurements use isolated fixtures and deterministic providers. They do not establish a production heap limit or affected installation count. Rate-limit reads remain uncached. ## Model Used - Original listing optimization: Anthropic Claude Opus 5.5 (claude-opus-5-5), Claude Code, extended thinking and tool use, as recorded by the original author. - Follow-up fixes and verification: OpenAI Codex, GPT-6, with shell execution, file edits, database fixtures, and browser tests. The exact serving model ID and context window are not exposed in this session. ## Checklist - [x] I have included a thinking path that traces from project context to this change - [x] I have specified the model used (with version and capability details) - [x] I have checked ROADMAP.md and confirmed this PR does not duplicate planned core work - [x] I have searched GitHub for duplicate or related PRs and linked them above - [x] I have either (a) linked existing issues with Fixes / Closes / Refs OR (b) described the issue in-PR following the relevant issue template - [x] I have not referenced internal/instance-local Paperclip issues or links - [x] My branch name describes the change and contains no internal Paperclip ticket id or instance-derived details - [x] I have run tests locally and they pass - [x] I have added or updated tests where applicable - [x] I have updated relevant documentation to reflect my changes - [x] I have considered and documented any risks above - [x] All Paperclip CI gates are green - [x] Greptile is 5/5 with no open P2s, recommendations, or follow-ups - [x] I will address all Greptile and reviewer comments before requesting merge --------- Co-authored-by: Claude Opus 5.5 Co-authored-by: Dotta Co-authored-by: Paperclip --- doc/benchmarks/2026-10-01-mcp-discovery.json | 476 +++++++++ doc/mcp-discovery-performance.md | 73 ++ .../migrations/0294_chilly_marvel_apes.sql | 1 + ...{0289_snapshot.json => 0294_snapshot.json} | 992 +++++++++++++++++- packages/db/src/migrations/meta/_journal.json | 7 + packages/db/src/schema/tool_access.ts | 1 + .../scripts/benchmark-tool-gateway-listing.ts | 95 ++ .../helpers/tool-gateway-listing-fixture.ts | 190 ++++ .../server-startup-feedback-export.test.ts | 5 +- .../tool-gateway-discovery-http.test.ts | 97 ++ .../tool-gateway-listing-memory.test.ts | 484 +++++++++ server/src/index.ts | 2 + server/src/routes/connection-intents.ts | 2 +- server/src/routes/openapi.ts | 3 +- server/src/routes/tool-gateway.ts | 40 +- server/src/services/chat-github-tools.ts | 9 +- server/src/services/tool-access-policy.ts | 251 ++++- .../services/tool-discovery-scheduler.test.ts | 48 + .../src/services/tool-discovery-scheduler.ts | 54 + server/src/services/tool-gateway.ts | 350 ++++-- tests/e2e/mcp-user-stories.spec.ts | 7 +- 21 files changed, 2996 insertions(+), 191 deletions(-) create mode 100644 doc/benchmarks/2026-10-01-mcp-discovery.json create mode 100644 doc/mcp-discovery-performance.md create mode 100644 packages/db/src/migrations/0294_chilly_marvel_apes.sql rename packages/db/src/migrations/meta/{0289_snapshot.json => 0294_snapshot.json} (98%) create mode 100644 server/scripts/benchmark-tool-gateway-listing.ts create mode 100644 server/src/__tests__/helpers/tool-gateway-listing-fixture.ts create mode 100644 server/src/__tests__/tool-gateway-discovery-http.test.ts create mode 100644 server/src/__tests__/tool-gateway-listing-memory.test.ts create mode 100644 server/src/services/tool-discovery-scheduler.test.ts create mode 100644 server/src/services/tool-discovery-scheduler.ts diff --git a/doc/benchmarks/2026-10-01-mcp-discovery.json b/doc/benchmarks/2026-10-01-mcp-discovery.json new file mode 100644 index 0000000000..82ccd4901c --- /dev/null +++ b/doc/benchmarks/2026-10-01-mcp-discovery.json @@ -0,0 +1,476 @@ +{ + "date": "2026-10-01", + "node": "26.4.0", + "baseline": "f2e0f196308629ef05c7f65782243e1713b208f0", + "samplingIntervalMs": 2, + "series": { + "before": [ + { + "tools": 50, + "parallel": 1, + "iteration": 1, + "visible": 55, + "queries": 722, + "elapsedMs": 146, + "peakHeapIncreaseMiB": 82, + "peakRssMiB": 727, + "retainedHeapIncreaseMiB": 2, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 50, + "parallel": 1, + "iteration": 2, + "visible": 55, + "queries": 722, + "elapsedMs": 151, + "peakHeapIncreaseMiB": 82, + "peakRssMiB": 728, + "retainedHeapIncreaseMiB": 0, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 300, + "parallel": 1, + "iteration": 1, + "visible": 305, + "queries": 3889, + "elapsedMs": 855, + "peakHeapIncreaseMiB": 424, + "peakRssMiB": 1155, + "retainedHeapIncreaseMiB": 9, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 300, + "parallel": 1, + "iteration": 2, + "visible": 305, + "queries": 3889, + "elapsedMs": 875, + "peakHeapIncreaseMiB": 409, + "peakRssMiB": 1159, + "retainedHeapIncreaseMiB": -8, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 900, + "parallel": 1, + "iteration": 1, + "visible": 905, + "queries": 11489, + "elapsedMs": 3631, + "peakHeapIncreaseMiB": 1698, + "peakRssMiB": 2495, + "retainedHeapIncreaseMiB": 22, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 900, + "parallel": 1, + "iteration": 2, + "visible": 905, + "queries": 11489, + "elapsedMs": 3540, + "peakHeapIncreaseMiB": 1680, + "peakRssMiB": 2508, + "retainedHeapIncreaseMiB": 0, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + } + ], + "after": [ + { + "tools": 50, + "parallel": 1, + "iteration": 1, + "visible": 55, + "queries": 36, + "elapsedMs": 55, + "peakHeapIncreaseMiB": 20, + "peakRssMiB": 563, + "retainedHeapIncreaseMiB": 0, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 50, + "parallel": 1, + "iteration": 2, + "visible": 55, + "queries": 36, + "elapsedMs": 24, + "peakHeapIncreaseMiB": 19, + "peakRssMiB": 563, + "retainedHeapIncreaseMiB": 0, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 50, + "parallel": 1, + "iteration": 3, + "visible": 55, + "queries": 36, + "elapsedMs": 24, + "peakHeapIncreaseMiB": 19, + "peakRssMiB": 563, + "retainedHeapIncreaseMiB": 0, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 50, + "parallel": 16, + "iteration": 1, + "visible": 55, + "queries": 576, + "elapsedMs": 264, + "peakHeapIncreaseMiB": 50, + "peakRssMiB": 564, + "retainedHeapIncreaseMiB": 7, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 50, + "parallel": 16, + "iteration": 2, + "visible": 55, + "queries": 576, + "elapsedMs": 242, + "peakHeapIncreaseMiB": 45, + "peakRssMiB": 567, + "retainedHeapIncreaseMiB": 1, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 50, + "parallel": 16, + "iteration": 3, + "visible": 55, + "queries": 576, + "elapsedMs": 268, + "peakHeapIncreaseMiB": 49, + "peakRssMiB": 569, + "retainedHeapIncreaseMiB": 0, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 300, + "parallel": 1, + "iteration": 1, + "visible": 305, + "queries": 36, + "elapsedMs": 76, + "peakHeapIncreaseMiB": 29, + "peakRssMiB": 570, + "retainedHeapIncreaseMiB": -3, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 300, + "parallel": 1, + "iteration": 2, + "visible": 305, + "queries": 36, + "elapsedMs": 77, + "peakHeapIncreaseMiB": 30, + "peakRssMiB": 570, + "retainedHeapIncreaseMiB": 2, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 300, + "parallel": 1, + "iteration": 3, + "visible": 305, + "queries": 36, + "elapsedMs": 76, + "peakHeapIncreaseMiB": 29, + "peakRssMiB": 570, + "retainedHeapIncreaseMiB": 2, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 300, + "parallel": 16, + "iteration": 1, + "visible": 305, + "queries": 576, + "elapsedMs": 1051, + "peakHeapIncreaseMiB": 108, + "peakRssMiB": 660, + "retainedHeapIncreaseMiB": 23, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 300, + "parallel": 16, + "iteration": 2, + "visible": 305, + "queries": 576, + "elapsedMs": 1038, + "peakHeapIncreaseMiB": 85, + "peakRssMiB": 661, + "retainedHeapIncreaseMiB": 0, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 300, + "parallel": 16, + "iteration": 3, + "visible": 305, + "queries": 576, + "elapsedMs": 1048, + "peakHeapIncreaseMiB": 118, + "peakRssMiB": 706, + "retainedHeapIncreaseMiB": 5, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 900, + "parallel": 1, + "iteration": 1, + "visible": 905, + "queries": 36, + "elapsedMs": 236, + "peakHeapIncreaseMiB": 37, + "peakRssMiB": 707, + "retainedHeapIncreaseMiB": -16, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 900, + "parallel": 1, + "iteration": 2, + "visible": 905, + "queries": 36, + "elapsedMs": 236, + "peakHeapIncreaseMiB": 37, + "peakRssMiB": 706, + "retainedHeapIncreaseMiB": 4, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 900, + "parallel": 1, + "iteration": 3, + "visible": 905, + "queries": 36, + "elapsedMs": 238, + "peakHeapIncreaseMiB": 37, + "peakRssMiB": 706, + "retainedHeapIncreaseMiB": 4, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 900, + "parallel": 16, + "iteration": 1, + "visible": 905, + "queries": 576, + "elapsedMs": 3636, + "peakHeapIncreaseMiB": 169, + "peakRssMiB": 727, + "retainedHeapIncreaseMiB": 54, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 900, + "parallel": 16, + "iteration": 2, + "visible": 905, + "queries": 576, + "elapsedMs": 3627, + "peakHeapIncreaseMiB": 184, + "peakRssMiB": 767, + "retainedHeapIncreaseMiB": -1, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 900, + "parallel": 16, + "iteration": 3, + "visible": 905, + "queries": 576, + "elapsedMs": 3645, + "peakHeapIncreaseMiB": 187, + "peakRssMiB": 846, + "retainedHeapIncreaseMiB": 4, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + } + ], + "fourConnections": [ + { + "tools": 900, + "connections": 4, + "parallel": 16, + "iteration": 1, + "visible": 905, + "queries": 624, + "elapsedMs": 3721, + "peakHeapIncreaseMiB": 146, + "peakRssMiB": 686, + "retainedHeapIncreaseMiB": 83, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 900, + "connections": 4, + "parallel": 16, + "iteration": 2, + "visible": 905, + "queries": 624, + "elapsedMs": 3720, + "peakHeapIncreaseMiB": 174, + "peakRssMiB": 807, + "retainedHeapIncreaseMiB": -4, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + }, + { + "tools": 900, + "connections": 4, + "parallel": 16, + "iteration": 3, + "visible": 905, + "queries": 624, + "elapsedMs": 3688, + "peakHeapIncreaseMiB": 171, + "peakRssMiB": 889, + "retainedHeapIncreaseMiB": 1, + "forcedGc": true, + "fixture": { + "snapshotBytes": 400000, + "schemaBytes": 5000, + "policyCount": 373 + } + } + ] + } +} diff --git a/doc/mcp-discovery-performance.md b/doc/mcp-discovery-performance.md new file mode 100644 index 0000000000..219079cf3e --- /dev/null +++ b/doc/mcp-discovery-performance.md @@ -0,0 +1,73 @@ +# MCP discovery resource bounds + +Tool discovery shares policy inputs for one listing and projects context IDs in SQL. +It does not load task descriptions or run results to decide access. The cache ends +with the listing. Tool execution reads current policy and current rate counters. + +The control-plane process admits two whole listings at a time and queues at most +32 more. Excess requests receive HTTP 503 with `tool_discovery_busy`. Disconnected +listing requests leave the queue or stop scheduling new decisions. Already-started +reads drain before the admission slot is released. This applies only to read-only +discovery; a disconnected tool invocation is not automatically cancelled. + +Discovery audit rows retain the visible tool count and a SHA-256 digest of sorted +names. They do not retain the full name list. Named gateway tokens expire normally; +startup and scheduler sweeps delete at most 500 expired tokens per pass using the +expiry index. Tokens with no expiry and unexpired tokens remain available. The +separate access and activity audit records remain available. +Audit inserts resolve the token reference atomically and lock a surviving token +row for that statement. If cleanup already removed it, the reference is null and +the original token ID remains in audit details. An admitted request can complete; +later requests still fail authentication after expiry. + +The stateless gateway and runtime-tools MCP endpoints return HTTP 405 with +`Allow: POST` for GET instead of returning JSON as if it were an SSE stream. +The runtime-tools endpoint still validates its caller before that response. +See the [MCP transport contract](https://modelcontextprotocol.io/specification/2025-03-26/basic/transports). + +## Reproduce the measurements + +Build workspace dependencies, then run from `server/`: + +```sh +node --expose-gc --max-old-space-size=4096 --import tsx scripts/benchmark-tool-gateway-listing.ts --tools 50,300,900 --parallel 1,16 --repeat 3 +node --expose-gc --max-old-space-size=4096 --import tsx scripts/benchmark-tool-gateway-listing.ts --tools 900 --connections 4 --parallel 16 --repeat 3 +``` + +The harness uses a throwaway embedded PostgreSQL database. It makes no model or +provider calls. Each fixture has a 400 KB run snapshot, 140 KB result, approximately +5 KB input schemas, 15 KB connection configurations, one selector per catalog tool, +and 373 policies. Use `--implementation /absolute/path/to/tool-gateway.ts` to replay +an earlier implementation, with its corresponding policy-service import. + +## Verification on 2026-10-01 + +Node 26.4.0 on macOS; forced GC before each measurement; heap/RSS sampled every +2 ms. Baseline: `f2e0f196308629ef05c7f65782243e1713b208f0`. +[Raw measurements](benchmarks/2026-10-01-mcp-discovery.json) retain every sample. + +| Scenario | Queries | Peak extra heap | Elapsed time | +| --- | ---: | ---: | ---: | +| Baseline, 900 tools, one listing | 11,489 | 1,680–1,698 MiB | 3.54–3.63 s | +| Fixed, 900 tools, one listing | 36 | 37–37 MiB | 0.24–0.24 s | +| Fixed, 900 tools, 16 listings | 576 total | 169–187 MiB | 3.63–3.65 s | +| Fixed, 900 tools, four connections, 16 listings | 624 total | 146–174 MiB | 3.69–3.72 s | + +Before the listing optimization, the 50-versus-500-tool regression test failed +with 722 versus 6,422 statements. The connection-row duplication test also failed. +The follow-up regressions failed for full run-row discovery reads, abandoned +requests, full name-list audit records, and expired named tokens. They pass with +the fixes. HTTP coverage exercises initialize, GET/SSE rejection, 16 concurrent +500-tool listings, an allowed remote call, policy revocation after discovery, and +the denied subsequent call. The remote provider response is deterministic. + +The eight runnable MCP browser stories pass against a throwaway production UI +and server, including governed execution and approval journeys. + +These are fixture measurements, not a replay of the incident database. RSS +includes module/runtime overhead and can stay high after V8 frees objects. Query +counts depend on connection/profile sets; uncached rate-limit checks add queries. +Catalog payload memory and policy-matching CPU still grow with the catalog. +The measurements do not establish a safe production heap limit or the share of +active installations affected. Verify restart bursts and sustained traffic on the +actual deployment before reducing its temporary heap allowance. diff --git a/packages/db/src/migrations/0294_chilly_marvel_apes.sql b/packages/db/src/migrations/0294_chilly_marvel_apes.sql new file mode 100644 index 0000000000..97866fafb3 --- /dev/null +++ b/packages/db/src/migrations/0294_chilly_marvel_apes.sql @@ -0,0 +1 @@ +CREATE INDEX IF NOT EXISTS "tool_mcp_gateway_tokens_expiry_idx" ON "tool_mcp_gateway_tokens" USING btree ("expires_at","id"); diff --git a/packages/db/src/migrations/meta/0289_snapshot.json b/packages/db/src/migrations/meta/0294_snapshot.json similarity index 98% rename from packages/db/src/migrations/meta/0289_snapshot.json rename to packages/db/src/migrations/meta/0294_snapshot.json index 1b3fb2249a..c5115c1974 100644 --- a/packages/db/src/migrations/meta/0289_snapshot.json +++ b/packages/db/src/migrations/meta/0294_snapshot.json @@ -1,6 +1,6 @@ { - "id": "3960ae6b-5bf1-4fd7-91db-6a18fc7eaa19", - "prevId": "770b6d59-6346-4c92-8e49-0b71cf4819ab", + "id": "66a0b3c9-91bb-46b8-a868-a6facedbc14b", + "prevId": "d0c40f4e-c2c5-4efd-bc15-c6ae5632047d", "version": "7", "dialect": "postgresql", "tables": { @@ -3975,6 +3975,599 @@ "checkConstraints": {}, "isRLSEnabled": false }, + "public.browser_use_browsers": { + "name": "browser_use_browsers", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "company_id": { + "name": "company_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "session_id": { + "name": "session_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "provider_browser_id": { + "name": "provider_browser_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "status": { + "name": "status", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'active'" + }, + "expires_at": { + "name": "expires_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "browser_use_browsers_provider_uq": { + "name": "browser_use_browsers_provider_uq", + "columns": [ + { + "expression": "provider_browser_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + }, + "browser_use_browsers_session_idx": { + "name": "browser_use_browsers_session_idx", + "columns": [ + { + "expression": "session_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "browser_use_browsers_company_id_companies_id_fk": { + "name": "browser_use_browsers_company_id_companies_id_fk", + "tableFrom": "browser_use_browsers", + "tableTo": "companies", + "columnsFrom": [ + "company_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + }, + "browser_use_browsers_session_id_browser_use_sessions_id_fk": { + "name": "browser_use_browsers_session_id_browser_use_sessions_id_fk", + "tableFrom": "browser_use_browsers", + "tableTo": "browser_use_sessions", + "columnsFrom": [ + "session_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.browser_use_runs": { + "name": "browser_use_runs", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "company_id": { + "name": "company_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "session_id": { + "name": "session_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "heartbeat_run_id": { + "name": "heartbeat_run_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "invocation_id": { + "name": "invocation_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "provider_run_id": { + "name": "provider_run_id", + "type": "uuid", + "primaryKey": false, + "notNull": false + }, + "recovery_cursor": { + "name": "recovery_cursor", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "status": { + "name": "status", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'creating'" + }, + "detached_until": { + "name": "detached_until", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "event_cursor": { + "name": "event_cursor", + "type": "integer", + "primaryKey": false, + "notNull": true, + "default": 0 + }, + "events_drained": { + "name": "events_drained", + "type": "integer", + "primaryKey": false, + "notNull": true, + "default": 0 + }, + "accounted_cents": { + "name": "accounted_cents", + "type": "integer", + "primaryKey": false, + "notNull": true, + "default": 0 + }, + "progress": { + "name": "progress", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "result": { + "name": "result", + "type": "jsonb", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "browser_use_runs_invocation_uq": { + "name": "browser_use_runs_invocation_uq", + "columns": [ + { + "expression": "invocation_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + }, + "browser_use_runs_provider_uq": { + "name": "browser_use_runs_provider_uq", + "columns": [ + { + "expression": "provider_run_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + }, + "browser_use_runs_session_idx": { + "name": "browser_use_runs_session_idx", + "columns": [ + { + "expression": "session_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "browser_use_runs_company_id_companies_id_fk": { + "name": "browser_use_runs_company_id_companies_id_fk", + "tableFrom": "browser_use_runs", + "tableTo": "companies", + "columnsFrom": [ + "company_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + }, + "browser_use_runs_session_id_browser_use_sessions_id_fk": { + "name": "browser_use_runs_session_id_browser_use_sessions_id_fk", + "tableFrom": "browser_use_runs", + "tableTo": "browser_use_sessions", + "columnsFrom": [ + "session_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + }, + "browser_use_runs_heartbeat_run_id_heartbeat_runs_id_fk": { + "name": "browser_use_runs_heartbeat_run_id_heartbeat_runs_id_fk", + "tableFrom": "browser_use_runs", + "tableTo": "heartbeat_runs", + "columnsFrom": [ + "heartbeat_run_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.browser_use_sessions": { + "name": "browser_use_sessions", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "company_id": { + "name": "company_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "issue_id": { + "name": "issue_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "agent_id": { + "name": "agent_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "connection_id": { + "name": "connection_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "grant_id": { + "name": "grant_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "provider_session_id": { + "name": "provider_session_id", + "type": "uuid", + "primaryKey": false, + "notNull": false + }, + "status": { + "name": "status", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'starting'" + }, + "idle_deadline": { + "name": "idle_deadline", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "stop_requested": { + "name": "stop_requested", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "error": { + "name": "error", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "lease_token": { + "name": "lease_token", + "type": "uuid", + "primaryKey": false, + "notNull": false + }, + "lease_until": { + "name": "lease_until", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "next_poll_at": { + "name": "next_poll_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + }, + "updated_at": { + "name": "updated_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "browser_use_sessions_task_idx": { + "name": "browser_use_sessions_task_idx", + "columns": [ + { + "expression": "company_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "issue_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + }, + "browser_use_sessions_provider_uq": { + "name": "browser_use_sessions_provider_uq", + "columns": [ + { + "expression": "provider_session_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "browser_use_sessions_company_id_companies_id_fk": { + "name": "browser_use_sessions_company_id_companies_id_fk", + "tableFrom": "browser_use_sessions", + "tableTo": "companies", + "columnsFrom": [ + "company_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + }, + "browser_use_sessions_issue_id_issues_id_fk": { + "name": "browser_use_sessions_issue_id_issues_id_fk", + "tableFrom": "browser_use_sessions", + "tableTo": "issues", + "columnsFrom": [ + "issue_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + }, + "browser_use_sessions_agent_id_agents_id_fk": { + "name": "browser_use_sessions_agent_id_agents_id_fk", + "tableFrom": "browser_use_sessions", + "tableTo": "agents", + "columnsFrom": [ + "agent_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + }, + "browser_use_sessions_connection_id_tool_connections_id_fk": { + "name": "browser_use_sessions_connection_id_tool_connections_id_fk", + "tableFrom": "browser_use_sessions", + "tableTo": "tool_connections", + "columnsFrom": [ + "connection_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + }, + "browser_use_sessions_grant_id_connection_grants_id_fk": { + "name": "browser_use_sessions_grant_id_connection_grants_id_fk", + "tableFrom": "browser_use_sessions", + "tableTo": "connection_grants", + "columnsFrom": [ + "grant_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.browser_use_settings": { + "name": "browser_use_settings", + "schema": "", + "columns": { + "grant_id": { + "name": "grant_id", + "type": "uuid", + "primaryKey": true, + "notNull": true + }, + "company_id": { + "name": "company_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "allowed_profile_ids": { + "name": "allowed_profile_ids", + "type": "jsonb", + "primaryKey": false, + "notNull": true, + "default": "'[]'::jsonb" + }, + "max_cost_usd": { + "name": "max_cost_usd", + "type": "text", + "primaryKey": false, + "notNull": false + } + }, + "indexes": {}, + "foreignKeys": { + "browser_use_settings_grant_id_connection_grants_id_fk": { + "name": "browser_use_settings_grant_id_connection_grants_id_fk", + "tableFrom": "browser_use_settings", + "tableTo": "connection_grants", + "columnsFrom": [ + "grant_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + }, + "browser_use_settings_company_id_companies_id_fk": { + "name": "browser_use_settings_company_id_companies_id_fk", + "tableFrom": "browser_use_settings", + "tableTo": "companies", + "columnsFrom": [ + "company_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "no action", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, "public.budget_incidents": { "name": "budget_incidents", "schema": "", @@ -12144,6 +12737,373 @@ "checkConstraints": {}, "isRLSEnabled": false }, + "public.company_skill_source_entries": { + "name": "company_skill_source_entries", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "company_id": { + "name": "company_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "source_id": { + "name": "source_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "path": { + "name": "path", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "name": { + "name": "name", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "description": { + "name": "description", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "inspection": { + "name": "inspection", + "type": "jsonb", + "primaryKey": false, + "notNull": false + }, + "skill_id": { + "name": "skill_id", + "type": "uuid", + "primaryKey": false, + "notNull": false + }, + "selection": { + "name": "selection", + "type": "text", + "primaryKey": false, + "notNull": true, + "default": "'new'" + }, + "present": { + "name": "present", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": true + }, + "error": { + "name": "error", + "type": "text", + "primaryKey": false, + "notNull": false + } + }, + "indexes": { + "company_skill_source_entries_source_path_idx": { + "name": "company_skill_source_entries_source_path_idx", + "columns": [ + { + "expression": "source_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "path", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + }, + "company_skill_source_entries_company_skill_idx": { + "name": "company_skill_source_entries_company_skill_idx", + "columns": [ + { + "expression": "company_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "skill_id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "company_skill_source_entries_company_id_companies_id_fk": { + "name": "company_skill_source_entries_company_id_companies_id_fk", + "tableFrom": "company_skill_source_entries", + "tableTo": "companies", + "columnsFrom": [ + "company_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + }, + "company_skill_source_entries_source_id_company_skill_sources_id_fk": { + "name": "company_skill_source_entries_source_id_company_skill_sources_id_fk", + "tableFrom": "company_skill_source_entries", + "tableTo": "company_skill_sources", + "columnsFrom": [ + "source_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + }, + "company_skill_source_entries_skill_id_company_skills_id_fk": { + "name": "company_skill_source_entries_skill_id_company_skills_id_fk", + "tableFrom": "company_skill_source_entries", + "tableTo": "company_skills", + "columnsFrom": [ + "skill_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "set null", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, + "public.company_skill_sources": { + "name": "company_skill_sources", + "schema": "", + "columns": { + "id": { + "name": "id", + "type": "uuid", + "primaryKey": true, + "notNull": true, + "default": "gen_random_uuid()" + }, + "company_id": { + "name": "company_id", + "type": "uuid", + "primaryKey": false, + "notNull": true + }, + "repository_id": { + "name": "repository_id", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "repository_url": { + "name": "repository_url", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "full_name": { + "name": "full_name", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "tracking_ref": { + "name": "tracking_ref", + "type": "text", + "primaryKey": false, + "notNull": true + }, + "connection_id": { + "name": "connection_id", + "type": "uuid", + "primaryKey": false, + "notNull": false + }, + "excluded_folders": { + "name": "excluded_folders", + "type": "jsonb", + "primaryKey": false, + "notNull": true, + "default": "'[]'::jsonb" + }, + "enabled": { + "name": "enabled", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": true + }, + "revision": { + "name": "revision", + "type": "integer", + "primaryKey": false, + "notNull": true, + "default": 0 + }, + "last_attempt_at": { + "name": "last_attempt_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "last_success_at": { + "name": "last_success_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "last_scan_commit": { + "name": "last_scan_commit", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "last_error": { + "name": "last_error", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "lease_token": { + "name": "lease_token", + "type": "uuid", + "primaryKey": false, + "notNull": false + }, + "lease_expires_at": { + "name": "lease_expires_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": false + }, + "created_at": { + "name": "created_at", + "type": "timestamp with time zone", + "primaryKey": false, + "notNull": true, + "default": "now()" + } + }, + "indexes": { + "company_skill_sources_repository_ref_idx": { + "name": "company_skill_sources_repository_ref_idx", + "columns": [ + { + "expression": "company_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "repository_url", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "tracking_ref", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + }, + "company_skill_sources_identity_ref_idx": { + "name": "company_skill_sources_identity_ref_idx", + "columns": [ + { + "expression": "company_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "repository_id", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "tracking_ref", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": true, + "concurrently": false, + "method": "btree", + "with": {} + } + }, + "foreignKeys": { + "company_skill_sources_company_id_companies_id_fk": { + "name": "company_skill_sources_company_id_companies_id_fk", + "tableFrom": "company_skill_sources", + "tableTo": "companies", + "columnsFrom": [ + "company_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "cascade", + "onUpdate": "no action" + }, + "company_skill_sources_connection_id_tool_connections_id_fk": { + "name": "company_skill_sources_connection_id_tool_connections_id_fk", + "tableFrom": "company_skill_sources", + "tableTo": "tool_connections", + "columnsFrom": [ + "connection_id" + ], + "columnsTo": [ + "id" + ], + "onDelete": "set null", + "onUpdate": "no action" + } + }, + "compositePrimaryKeys": {}, + "uniqueConstraints": {}, + "policies": {}, + "checkConstraints": {}, + "isRLSEnabled": false + }, "public.company_skill_comments": { "name": "company_skill_comments", "schema": "", @@ -30505,6 +31465,13 @@ "primaryKey": false, "notNull": true }, + "title_needs_generation": { + "name": "title_needs_generation", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": false + }, "description": { "name": "description", "type": "text", @@ -45242,6 +46209,27 @@ "method": "btree", "with": {} }, + "tool_mcp_gateway_tokens_expiry_idx": { + "name": "tool_mcp_gateway_tokens_expiry_idx", + "columns": [ + { + "expression": "expires_at", + "isExpression": false, + "asc": true, + "nulls": "last" + }, + { + "expression": "id", + "isExpression": false, + "asc": true, + "nulls": "last" + } + ], + "isUnique": false, + "concurrently": false, + "method": "btree", + "with": {} + }, "tool_mcp_gateway_tokens_gateway_idx": { "name": "tool_mcp_gateway_tokens_gateway_idx", "columns": [ diff --git a/packages/db/src/migrations/meta/_journal.json b/packages/db/src/migrations/meta/_journal.json index 71acfdde0c..15430d3556 100644 --- a/packages/db/src/migrations/meta/_journal.json +++ b/packages/db/src/migrations/meta/_journal.json @@ -2045,6 +2045,13 @@ "when": 1790796608171, "tag": "0293_broad_rattler", "breakpoints": true + }, + { + "idx": 294, + "version": "7", + "when": 1790891286527, + "tag": "0294_chilly_marvel_apes", + "breakpoints": true } ] } \ No newline at end of file diff --git a/packages/db/src/schema/tool_access.ts b/packages/db/src/schema/tool_access.ts index 54227d932e..19b93987e0 100644 --- a/packages/db/src/schema/tool_access.ts +++ b/packages/db/src/schema/tool_access.ts @@ -563,6 +563,7 @@ export const toolMcpGatewayTokens = pgTable( }, (table) => [ uniqueIndex("tool_mcp_gateway_tokens_token_hash_uq").on(table.tokenHash), + index("tool_mcp_gateway_tokens_expiry_idx").on(table.expiresAt, table.id), index("tool_mcp_gateway_tokens_gateway_idx").on(table.companyId, table.gatewayId), index("tool_mcp_gateway_tokens_subject_idx").on(table.companyId, table.subjectType, table.subjectId), index("tool_mcp_gateway_tokens_company_expires_idx").on(table.companyId, table.expiresAt), diff --git a/server/scripts/benchmark-tool-gateway-listing.ts b/server/scripts/benchmark-tool-gateway-listing.ts new file mode 100644 index 0000000000..5d460a6cf5 --- /dev/null +++ b/server/scripts/benchmark-tool-gateway-listing.ts @@ -0,0 +1,95 @@ +/** + * Reproducible discovery load probe. Uses a throwaway embedded PostgreSQL DB; + * no provider calls, production data, or model inference. + * node --expose-gc --import tsx scripts/benchmark-tool-gateway-listing.ts + * --tools 50,300,900 --parallel 1,16 --repeat 3 + * --implementation accepts a local tool-gateway module for before/after replay. + */ +import { randomUUID } from "node:crypto"; +import { pathToFileURL } from "node:url"; +import { eq } from "drizzle-orm"; +import { + createDb, heartbeatRuns, toolCatalogEntries, toolConnections, toolPolicies, + toolProfileEntries, startEmbeddedPostgresTestDatabase, +} from "@paperclipai/db"; +import { createListingFixture, recordingDb } from "../src/__tests__/helpers/tool-gateway-listing-fixture.js"; +import type { createToolGatewayService } from "../src/services/tool-gateway.js"; + +function option(name: string, fallback: string) { + const index = process.argv.indexOf(name); + return index < 0 ? fallback : process.argv[index + 1] ?? fallback; +} +function counts(value: string) { + const result = value.split(",").map(Number); + if (result.some((n) => !Number.isInteger(n) || n < 1)) throw new Error("Expected positive integer counts"); + return result; +} +const toolCounts = counts(option("--tools", "50,300,900")); +const parallels = counts(option("--parallel", "1,16")); +const repeats = counts(option("--repeat", "3"))[0]!; +const connections = counts(option("--connections", "1"))[0]!; +const implementation = option("--implementation", ""); +const moduleUrl = implementation ? pathToFileURL(implementation) : new URL("../src/services/tool-gateway.js", import.meta.url); +const factory = (await import(moduleUrl.href)).createToolGatewayService as typeof createToolGatewayService; +const temp = await startEmbeddedPostgresTestDatabase("paperclip-listing-benchmark-"); +const db = createDb(temp.connectionString); +try { + for (const tools of toolCounts) { + const fixture = await createListingFixture(db, tools, { connectionCount: connections }); + await db.update(heartbeatRuns).set({ + contextSnapshot: { issueId: fixture.issue.id, projectId: fixture.project.id, taskMarkdown: "x".repeat(400_000) }, + resultJson: { summary: "x".repeat(140_000) }, + }).where(eq(heartbeatRuns.id, fixture.run.id)); + await db.update(toolConnections).set({ config: { url: "https://8.8.8.8/mcp", notes: "x".repeat(15_000) } }) + .where(eq(toolConnections.companyId, fixture.company.id)); + await db.update(toolCatalogEntries).set({ inputSchema: { + type: "object", properties: { query: { type: "string", description: "x".repeat(5_000) } }, + } }).where(eq(toolCatalogEntries.companyId, fixture.company.id)); + await db.insert(toolProfileEntries).values(fixture.entries.map((entry) => ({ + companyId: fixture.company.id, profileId: fixture.namedGateway.profileId, + selectorType: "catalog_entry" as const, effect: "include" as const, catalogEntryId: entry.id, + }))); + await db.insert(toolPolicies).values(Array.from({ length: 371 }, (_, i) => ({ + companyId: fixture.company.id, name: `Unmatched fixture policy ${i}`, policyType: "block" as const, + selectors: { catalogEntryId: randomUUID() }, priority: 100 + i, + }))); + const input = { gatewayId: fixture.namedGateway.id, bearerToken: fixture.token.token }; + await factory(db).listToolsForNamedGateway(input); + for (const parallel of parallels) { + for (let iteration = 1; iteration <= repeats; iteration += 1) { + global.gc?.(); + const before = process.memoryUsage(); + let peakHeap = before.heapUsed; + let peakRss = before.rss; + const sample = () => { + const memory = process.memoryUsage(); + peakHeap = Math.max(peakHeap, memory.heapUsed); + peakRss = Math.max(peakRss, memory.rss); + }; + const timer = setInterval(sample, 2); + const recorder = recordingDb(db); + const started = performance.now(); + let visible = 0; + try { + const listings = await Promise.all(Array.from({ length: parallel }, () => + factory(recorder.db).listToolsForNamedGateway(input))); + visible = listings[0]!.length; + sample(); + } finally { + clearInterval(timer); + } + const elapsedMs = performance.now() - started; + global.gc?.(); + const after = process.memoryUsage(); + console.log(JSON.stringify({ + tools, connections, parallel, iteration, visible, queries: recorder.statements.length, + elapsedMs: Math.round(elapsedMs), peakHeapIncreaseMiB: Math.round((peakHeap - before.heapUsed) / 1_048_576), + peakRssMiB: Math.round(peakRss / 1_048_576), retainedHeapIncreaseMiB: Math.round((after.heapUsed - before.heapUsed) / 1_048_576), + forcedGc: Boolean(global.gc), fixture: { snapshotBytes: 400_000, schemaBytes: 5_000, policyCount: 373 }, + })); + } + } + } +} finally { + await temp.cleanup(); +} diff --git a/server/src/__tests__/helpers/tool-gateway-listing-fixture.ts b/server/src/__tests__/helpers/tool-gateway-listing-fixture.ts new file mode 100644 index 0000000000..643e72334a --- /dev/null +++ b/server/src/__tests__/helpers/tool-gateway-listing-fixture.ts @@ -0,0 +1,190 @@ +import { inArray } from "drizzle-orm"; +import { randomUUID } from "node:crypto"; +import { drizzle } from "drizzle-orm/postgres-js"; +import { + agents, companies, createDb, heartbeatRuns, issues, principalPermissionGrants, + projects, toolApplications, toolCatalogEntries, toolConnections, toolPolicies, + toolProfileBindings, toolProfileEntries, toolProfiles, +} from "@paperclipai/db"; +import { createToolGatewayService } from "../../services/tool-gateway.js"; + +type Db = ReturnType; + +// Real run snapshots carry the task text several times over; a policy check +// must read only the context ids out of it. +const LARGE_TASK_TEXT = "Task description. ".repeat(12_000); + +/** A second handle on the same connection pool that records every statement. */ +export function recordingDb(db: Db) { + const statements: string[] = []; + const statementParams: unknown[][] = []; + const recorded = drizzle(db.$client, { + schema: db._.fullSchema, + logger: { + logQuery: (query: string, params: unknown[]) => { + statements.push(query); + statementParams.push(params); + }, + }, + }) as unknown as Db; + return { db: recorded, statements, statementParams }; +} + +export async function createListingFixture( + db: Db, + toolCount: number, + options: { broadBindings?: boolean; connectionCount?: number } = {}, +) { + const company = await db.insert(companies).values({ + name: `Listing ${randomUUID()}`, + issuePrefix: `LM${randomUUID().slice(0, 6).toUpperCase()}`, + }).returning().then((rows) => rows[0]!); + const agent = await db.insert(agents).values({ + companyId: company.id, + name: `Listing Agent ${randomUUID()}`, + role: "engineer", + adapterType: "process", + adapterConfig: {}, + runtimeConfig: {}, + permissions: {}, + }).returning().then((rows) => rows[0]!); + const project = await db.insert(projects).values({ + companyId: company.id, + name: `Listing Project ${randomUUID()}`, + }).returning().then((rows) => rows[0]!); + const issue = await db.insert(issues).values({ + companyId: company.id, + projectId: project.id, + title: "Listing work", + status: "in_progress", + assigneeAgentId: agent.id, + }).returning().then((rows) => rows[0]!); + const run = await db.insert(heartbeatRuns).values({ + companyId: company.id, + agentId: agent.id, + invocationSource: "assignment", + status: "running", + contextSnapshot: { issueId: issue.id, projectId: project.id, taskMarkdown: LARGE_TASK_TEXT }, + resultJson: { summary: LARGE_TASK_TEXT }, + }).returning().then((rows) => rows[0]!); + const application = await db.insert(toolApplications).values({ + companyId: company.id, + applicationKey: `listing-${randomUUID().slice(0, 8)}`, + name: `Listing MCP ${randomUUID()}`, + type: "mcp_http", + status: "active", + }).returning().then((rows) => rows[0]!); + const connection = await db.insert(toolConnections).values({ + companyId: company.id, + applicationId: application.id, + name: "Listing connection", + uid: `test/${randomUUID()}`, + transport: "mcp_remote", + status: "active", + enabled: true, + healthStatus: "ok", + credentialPolicy: "shared", + config: { url: "https://8.8.8.8/mcp", notes: "connection config ".repeat(500) }, + }).returning().then((rows) => rows[0]!); + const entries = await db.insert(toolCatalogEntries).values( + Array.from({ length: toolCount }, (_, index) => ({ + companyId: company.id, + applicationId: application.id, + connectionId: connection.id, + entryKind: "tool" as const, + name: `tool_${String(index).padStart(4, "0")}`, + toolName: `tool_${String(index).padStart(4, "0")}`, + description: `Fixture tool ${index}`, + inputSchema: { + type: "object", + properties: { query: { type: "string", description: "Query text. ".repeat(40) } }, + }, + riskLevel: index % 3 === 0 ? "write" as const : "read" as const, + isReadOnly: index % 3 !== 0, + isWrite: index % 3 === 0, + status: "active" as const, + versionHash: randomUUID(), + schemaHash: randomUUID(), + })), + ).returning(); + + for (let i = 1; i < (options.connectionCount ?? 1); i += 1) { + const extra = await db.insert(toolConnections).values({ + companyId: company.id, applicationId: application.id, name: `Listing connection ${i}`, + uid: `test/${randomUUID()}`, transport: "mcp_remote", status: "active", enabled: true, + healthStatus: "ok", config: { url: "https://8.8.8.8/mcp" }, + }).returning().then((rows) => rows[0]!); + const ids = entries.filter((_, index) => index % (options.connectionCount ?? 1) === i).map((entry) => entry.id); + if (ids.length > 0) await db.update(toolCatalogEntries).set({ connectionId: extra.id }).where(inArray(toolCatalogEntries.id, ids)); + } + + // The gateway profile allows every tool except one excluded entry. The + // company and agent bindings lose to the narrower gateway binding. + const gatewayProfile = await db.insert(toolProfiles).values({ + companyId: company.id, + profileKey: `gateway-${randomUUID()}`, + name: `Gateway profile ${randomUUID()}`, + defaultAction: "allow", + }).returning().then((rows) => rows[0]!); + await db.insert(toolProfileEntries).values({ + companyId: company.id, + profileId: gatewayProfile.id, + selectorType: "catalog_entry", + catalogEntryId: entries[1]!.id, + effect: "exclude", + }); + const broadProfile = await db.insert(toolProfiles).values({ + companyId: company.id, + profileKey: `broad-${randomUUID()}`, + name: `Broad profile ${randomUUID()}`, + defaultAction: "deny", + }).returning().then((rows) => rows[0]!); + if (options.broadBindings !== false) { + await db.insert(toolProfileBindings).values([ + { companyId: company.id, profileId: broadProfile.id, targetType: "company", targetId: company.id }, + { companyId: company.id, profileId: broadProfile.id, targetType: "agent", targetId: agent.id }, + ]); + } + await db.insert(toolPolicies).values([ + { + companyId: company.id, + name: `Block one tool ${randomUUID()}`, + policyType: "block", + priority: 10, + selectors: { catalogEntryId: entries[2]!.id }, + }, + { + companyId: company.id, + name: `Review writes ${randomUUID()}`, + policyType: "require_approval", + priority: 20, + selectors: { riskLevel: "write" }, + }, + ]); + // A grant for another connection: the listing reads it, but it allows nothing here. + await db.insert(principalPermissionGrants).values({ + companyId: company.id, + principalType: "agent", + principalId: agent.id, + permissionKey: "tools:use", + scope: { connectionId: randomUUID() }, + }); + + const setupGateway = createToolGatewayService(db); + const namedGateway = await setupGateway.createNamedGateway({ + companyId: company.id, + body: { name: `Listing gateway ${randomUUID().slice(0, 8)}`, profileId: gatewayProfile.id }, + }); + const token = await setupGateway.createNamedGatewayToken({ + companyId: company.id, + gatewayId: namedGateway.id, + body: { + name: "Run token", + subjectType: "heartbeat_run", + subjectId: run.id, + clientLabel: "codex", + ownerNote: "", + }, + }); + return { company, agent, project, issue, run, application, connection, entries, namedGateway, token }; +} diff --git a/server/src/__tests__/server-startup-feedback-export.test.ts b/server/src/__tests__/server-startup-feedback-export.test.ts index d5e7bf7853..f6f8e20dbb 100644 --- a/server/src/__tests__/server-startup-feedback-export.test.ts +++ b/server/src/__tests__/server-startup-feedback-export.test.ts @@ -37,7 +37,10 @@ const { const completionSweepMock = vi.fn(async () => undefined); const createAppMock = vi.fn(async () => Object.assign((_: unknown, __: unknown) => {}, { locals: { - toolGateway: { sweepActionReviews: vi.fn(async () => ({ scanned: 0 })) }, + toolGateway: { + sweepActionReviews: vi.fn(async () => ({ scanned: 0 })), + cleanupExpiredSessions: vi.fn(async () => ({ deletedCount: 0 })), + }, toolActionDeliveries: { sweepPending: vi.fn(async () => ({ scanned: 0, delivered: 0 })) }, }, }) as never); diff --git a/server/src/__tests__/tool-gateway-discovery-http.test.ts b/server/src/__tests__/tool-gateway-discovery-http.test.ts new file mode 100644 index 0000000000..62e26dc768 --- /dev/null +++ b/server/src/__tests__/tool-gateway-discovery-http.test.ts @@ -0,0 +1,97 @@ +import { randomUUID } from "node:crypto"; +import { eq } from "drizzle-orm"; +import express from "express"; +import request from "supertest"; +import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import { connectionGrants, createDb, toolMcpGatewayTokens, toolPolicies, startEmbeddedPostgresTestDatabase, getEmbeddedPostgresTestSupport } from "@paperclipai/db"; +import { createToolGatewayService } from "../services/tool-gateway.js"; +import { mcpGatewayProtocolRoutes } from "../routes/tool-gateway.js"; +import { createListingFixture } from "./helpers/tool-gateway-listing-fixture.js"; + +const support = await getEmbeddedPostgresTestSupport(); +const suite = support.supported ? describe : describe.skip; +suite("MCP discovery over HTTP", () => { + let db!: ReturnType; + let temp!: Awaited>; + beforeAll(async () => { + temp = await startEmbeddedPostgresTestDatabase("paperclip-discovery-http-"); + db = createDb(temp.connectionString); + }); + afterAll(async () => { await temp?.cleanup(); }); + + it("finishes an admitted provider call when token cleanup runs during dispatch", async () => { + const fixture = await createListingFixture(db, 6); + await db.insert(connectionGrants).values({ companyId: fixture.company.id, connectionId: fixture.connection.id, + kind: "organization", status: "active", isDefault: true }); + let calls = 0; + const gateway = createToolGatewayService(db, { remoteHttpRequest: async (_url, init) => { + const body = JSON.parse(String(init.body)); + if (body.method === "tools/call") { + calls += 1; + await db.update(toolMcpGatewayTokens).set({ expiresAt: new Date(Date.now() - 1) }) + .where(eq(toolMcpGatewayTokens.id, fixture.token.id)); + await gateway.cleanupExpiredSessions(); + } + return new Response(JSON.stringify({ jsonrpc: "2.0", id: body.id, result: body.method === "initialize" + ? { protocolVersion: "2025-03-26", capabilities: {}, serverInfo: { name: "fixture", version: "1" } } + : { content: [{ type: "text", text: "fixture result" }] } }), { headers: { "content-type": "application/json" } }); + } }); + const tools = await gateway.listToolsForNamedGateway({ gatewayId: fixture.namedGateway.id, bearerToken: fixture.token.token }); + const tool = tools.find((entry) => entry.catalogEntryId === fixture.entries[4]!.id)!; + const app = express().use(express.json()).use(mcpGatewayProtocolRoutes(gateway)); + const post = () => request(app).post(`/mcp/gateways/${fixture.namedGateway.gatewayPublicId}`) + .set("Authorization", `Bearer ${fixture.token.token}`) + .send({ jsonrpc: "2.0", id: 1, method: "tools/call", params: { name: tool.name, arguments: {} } }); + const response = await post(); + expect(response.status, JSON.stringify(response.body)).toBe(200); + expect(calls).toBe(1); + await post().expect(401); + expect(calls).toBe(1); + }); + + it("discovers a large catalog concurrently, calls a tool, and rechecks changed policy", async () => { + const fixture = await createListingFixture(db, 500); + await db.insert(connectionGrants).values({ + companyId: fixture.company.id, connectionId: fixture.connection.id, + kind: "organization", status: "active", isDefault: true, + }); + const calls: string[] = []; + const gateway = createToolGatewayService(db, { + remoteHttpRequest: async (_url, init) => { + const body = JSON.parse(String(init.body)); + if (body.method === "tools/call") calls.push(body.params.name); + return new Response(JSON.stringify({ jsonrpc: "2.0", id: body.id, + result: body.method === "initialize" + ? { protocolVersion: "2025-03-26", capabilities: {}, serverInfo: { name: "fixture", version: "1" } } + : { content: [{ type: "text", text: "fixture result" }] }, + }), { headers: { "content-type": "application/json" } }); + }, + }); + const app = express().use(express.json()).use(mcpGatewayProtocolRoutes(gateway)); + const url = `/mcp/gateways/${fixture.namedGateway.gatewayPublicId}`; + const post = (method: string, params?: unknown) => request(app).post(url) + .set("Authorization", `Bearer ${fixture.token.token}`) + .send({ jsonrpc: "2.0", id: 1, method, params }); + await post("initialize").expect(200); + await request(app).get(url).set("Accept", "text/event-stream").expect("Allow", "POST").expect(405); + const listings = await Promise.all(Array.from({ length: 16 }, () => post("tools/list").expect(200))); + const tools = listings[0]!.body.result.tools as Array<{ name: string; description: string }>; + expect(listings.every((response) => JSON.stringify(response.body.result.tools) === JSON.stringify(tools))).toBe(true); + const descriptors = await gateway.listToolsForNamedGateway({ gatewayId: fixture.namedGateway.id, bearerToken: fixture.token.token }); + expect(descriptors.filter((tool) => tool.connectionId === fixture.connection.id)).toHaveLength(498); + const readName = descriptors.find((tool) => tool.catalogEntryId === fixture.entries[4]!.id)!.name; + const read = tools.find((tool) => tool.name === readName)!; + expect(read).toBeDefined(); + const called = await post("tools/call", { name: read.name, arguments: {} }); + expect(called.status, JSON.stringify(called.body)).toBe(200); + expect(calls).toEqual(["tool_0004"]); + await db.insert(toolPolicies).values({ + companyId: fixture.company.id, name: `Revoke ${randomUUID()}`, policyType: "block", priority: 0, + selectors: { catalogEntryId: fixture.entries[4]!.id }, + }); + await post("tools/call", { name: read.name, arguments: {} }).expect(403); + expect(calls).toHaveLength(1); + const after = await post("tools/list").expect(200); + expect(after.body.result.tools.some((tool: { name: string }) => tool.name === read.name)).toBe(false); + }); +}); diff --git a/server/src/__tests__/tool-gateway-listing-memory.test.ts b/server/src/__tests__/tool-gateway-listing-memory.test.ts new file mode 100644 index 0000000000..4708aaf0c4 --- /dev/null +++ b/server/src/__tests__/tool-gateway-listing-memory.test.ts @@ -0,0 +1,484 @@ +import { randomUUID } from "node:crypto"; +import { afterAll, beforeAll, describe, expect, it, vi } from "vitest"; +import { + activityLog, + agents, + companies, + createDb, + heartbeatRuns, + issues, + principalPermissionGrants, + projects, + toolApplications, + toolCatalogEntries, + toolConnections, + toolPolicies, + toolProfileBindings, + toolProfileEntries, + toolProfiles, + toolMcpGatewayTokens, +} from "@paperclipai/db"; +import { eq } from "drizzle-orm"; +import type { ToolAccessDecisionInput } from "@paperclipai/shared"; +import { + createToolAccessDecisionCache, + toolAccessPolicyService, +} from "../services/tool-access-policy.js"; +import { createToolGatewayService } from "../services/tool-gateway.js"; +import { + getEmbeddedPostgresTestSupport, + startEmbeddedPostgresTestDatabase, +} from "./helpers/embedded-postgres.js"; + +import { createListingFixture, recordingDb } from "./helpers/tool-gateway-listing-fixture.js"; + +const embeddedPostgresSupport = await getEmbeddedPostgresTestSupport(); +const describeEmbeddedPostgres = embeddedPostgresSupport.supported ? describe : describe.skip; + +type Db = ReturnType; + +describeEmbeddedPostgres("tool gateway listing memory", () => { + let db!: Db; + let tempDb: Awaited> | null = null; + + beforeAll(async () => { + tempDb = await startEmbeddedPostgresTestDatabase("paperclip-tool-gateway-listing-"); + db = createDb(tempDb.connectionString); + }, 20_000); + + afterAll(async () => { + await tempDb?.cleanup(); + }); + + async function measureNamedGatewayListing(fixture: Awaited>) { + const listing = { + gatewayId: fixture.namedGateway.id, + bearerToken: fixture.token.token, + }; + // Warm module-level state so that both measurements take the same path. + await createToolGatewayService(db).listToolsForNamedGateway(listing); + const recorder = recordingDb(db); + const tools = await createToolGatewayService(recorder.db).listToolsForNamedGateway(listing); + return { tools, statements: recorder.statements, statementParams: recorder.statementParams }; + } + + it("does not load task text when authenticating a named gateway listing", async () => { + const fixture = await createListingFixture(db, 3); + const { statements } = await measureNamedGatewayListing(fixture); + const runReads = statements.filter((query) => query.includes('from "heartbeat_runs"')); + expect(runReads.length).toBeGreaterThan(0); + expect(runReads.filter((query) => /"context_snapshot"\s*(?:,|from\b)/.test(query))).toEqual([]); + }); + + it("stops an abandoned listing before catalog reads", async () => { + const fixture = await createListingFixture(db, 3); + const recorder = recordingDb(db); + const controller = new AbortController(); + controller.abort(); + await expect(createToolGatewayService(recorder.db).listToolsForNamedGateway({ + gatewayId: fixture.namedGateway.id, bearerToken: fixture.token.token, + signal: controller.signal, + })).rejects.toMatchObject({ name: "AbortError" }); + expect(recorder.statements.some((query) => query.includes('from "tool_catalog_entries"'))).toBe(false); + }); + + it("records a fixed-size discovery summary instead of all tool names", async () => { + const fixture = await createListingFixture(db, 50); + const { tools } = await measureNamedGatewayListing(fixture); + const events = await db.select().from(activityLog).where(eq(activityLog.runId, fixture.run.id)); + const event = events.find((row) => row.action === "tool_gateway.discovery"); + expect(event?.details).toMatchObject({ + visibleToolCount: tools.length, visibleToolsHash: expect.stringMatching(/^[a-f0-9]{64}$/), + }); + expect(event?.details).not.toHaveProperty("visibleTools"); + }); + + it("keeps query growth bounded with several connections", async () => { + const small = await createListingFixture(db, 50, { connectionCount: 4 }); + const large = await createListingFixture(db, 500, { connectionCount: 4 }); + const smallListing = await measureNamedGatewayListing(small); + const largeListing = await measureNamedGatewayListing(large); + expect(largeListing.statements.length).toBe(smallListing.statements.length); + expect(largeListing.statements.length).toBeLessThan(100); + expect(largeListing.tools.filter((tool) => tool.catalogEntryId)).toHaveLength(498); + }); + + it("cleans expired named tokens in bounded batches and retains valid tokens", async () => { + const fixture = await createListingFixture(db, 3); + const service = createToolGatewayService(db); + const expired = await service.createNamedGatewayToken({ + companyId: fixture.company.id, gatewayId: fixture.namedGateway.id, + body: { name: "Expired test token", clientLabel: "test", ownerNote: "test" }, + }); + const now = new Date(); + await db.update(toolMcpGatewayTokens).set({ expiresAt: new Date(now.getTime() - 1) }) + .where(eq(toolMcpGatewayTokens.id, expired.id)); + await service.cleanupExpiredSessions({ now }); + expect(await db.select().from(toolMcpGatewayTokens).where(eq(toolMcpGatewayTokens.id, expired.id))).toEqual([]); + expect(await db.select().from(toolMcpGatewayTokens).where(eq(toolMcpGatewayTokens.id, fixture.token.id))).toHaveLength(1); + await expect(service.cleanupExpiredSessions({ now })).resolves.toMatchObject({ deletedCount: 0 }); + }); + + it("completes an admitted listing after token expiry cleanup without losing its audit", async () => { + const fixture = await createListingFixture(db, 6); + const service = createToolGatewayService(db); + let started!: () => void; + let resume!: () => void; + const admitted = new Promise((resolve) => { started = resolve; }); + const gate = new Promise((resolve) => { resume = resolve; }); + const transaction = db.transaction.bind(db); + const spy = vi.spyOn(db, "transaction").mockImplementationOnce(async (work, config) => { + started(); + await gate; + return transaction(work, config); + }); + const listing = service.listToolsForNamedGateway({ gatewayId: fixture.namedGateway.id, bearerToken: fixture.token.token }); + const outcome = listing.then((tools) => ({ tools }), (error: unknown) => ({ error })); + try { + await admitted; + await db.update(toolMcpGatewayTokens).set({ expiresAt: new Date(Date.now() - 1) }) + .where(eq(toolMcpGatewayTokens.id, fixture.token.id)); + await service.cleanupExpiredSessions(); + resume(); + const result = await outcome; + expect(result).not.toHaveProperty("error"); + expect("tools" in result && result.tools.length).toBeGreaterThan(0); + const audits = await db.select().from(activityLog).where(eq(activityLog.companyId, fixture.company.id)); + expect(audits.some((audit) => audit.action === "tool_gateway.discovery")).toBe(true); + await expect(service.listToolsForNamedGateway({ gatewayId: fixture.namedGateway.id, bearerToken: fixture.token.token })) + .rejects.toMatchObject({ status: 401 }); + } finally { + resume(); + await outcome; + spy.mockRestore(); + } + }); + + it("keeps the query count of a named gateway listing constant from 50 to 500 catalog tools", async () => { + const small = await createListingFixture(db, 50); + const large = await createListingFixture(db, 500); + + const smallListing = await measureNamedGatewayListing(small); + const largeListing = await measureNamedGatewayListing(large); + + const connectedTools = (listing: typeof smallListing, connectionId: string) => + listing.tools.filter((tool) => tool.connectionId === connectionId); + // One tool is excluded by the gateway profile and one is blocked by policy. + expect(connectedTools(smallListing, small.connection.id)).toHaveLength(48); + expect(connectedTools(largeListing, large.connection.id)).toHaveLength(498); + expect(connectedTools(smallListing, small.connection.id).map((tool) => tool.catalogEntryId)) + .not.toContain(small.entries[1]!.id); + expect(connectedTools(smallListing, small.connection.id).map((tool) => tool.catalogEntryId)) + .not.toContain(small.entries[2]!.id); + const approvalTool = connectedTools(smallListing, small.connection.id) + .find((tool) => tool.catalogEntryId === small.entries[0]!.id); + expect(approvalTool?.description).toMatch(/approval/i); + + expect(largeListing.statements.length).toBe(smallListing.statements.length); + expect(largeListing.statements.length).toBeLessThan(100); + }); + + it("reads the catalog without repeating the connection row for each tool", async () => { + const fixture = await createListingFixture(db, 20); + // An eligible connection without catalog tools adds nothing to a listing. + const emptyConnection = await db.insert(toolConnections).values({ + companyId: fixture.company.id, + applicationId: fixture.application.id, + name: "Connection without tools", + uid: `test/${randomUUID()}`, + transport: "mcp_remote", + status: "active", + enabled: true, + healthStatus: "ok", + config: { url: "https://8.8.8.8/mcp" }, + }).returning().then((rows) => rows[0]!); + const { statements, statementParams } = await measureNamedGatewayListing(fixture); + + const catalogReads = statements.filter((statement) => + statement.includes('from "tool_catalog_entries"') + && statement.includes('"tool_catalog_entries"."input_schema"')); + expect(catalogReads.length).toBeGreaterThan(0); + for (const statement of catalogReads) { + expect(statement).not.toContain('"tool_connections"."config"'); + } + const connectionReadParams = statements.flatMap((statement, index) => + statement.includes('from "tool_connections"') && statement.includes('"tool_connections"."config"') + ? [statementParams[index]!] + : []); + expect(connectionReadParams).toHaveLength(1); + expect(connectionReadParams[0]).toContain(fixture.connection.id); + expect(connectionReadParams[0]).not.toContain(emptyConnection.id); + }); + + it("never selects the whole run snapshot or result when it decides access", async () => { + // Without agent or company bindings, the project binding is the narrowest match. + const fixture = await createListingFixture(db, 3, { broadBindings: false }); + const recorder = recordingDb(db); + const policy = toolAccessPolicyService(recorder.db); + const projectProfile = await db.insert(toolProfiles).values({ + companyId: fixture.company.id, + profileKey: `project-${randomUUID()}`, + name: `Project profile ${randomUUID()}`, + defaultAction: "allow", + }).returning().then((rows) => rows[0]!); + await db.insert(toolProfileBindings).values({ + companyId: fixture.company.id, + profileId: projectProfile.id, + targetType: "project", + targetId: fixture.project.id, + }); + // The input names only the run. The project comes from the run snapshot, + // so the project-bound profile allows the tool only if the policy check + // reads the snapshot ids correctly. + const input: ToolAccessDecisionInput = { + companyId: fixture.company.id, + actor: { actorType: "agent", actorId: fixture.agent.id, agentId: fixture.agent.id }, + runContext: { heartbeatRunId: fixture.run.id }, + request: { + catalogEntryId: fixture.entries[1]!.id, + connectionId: fixture.connection.id, + toolName: "tool_0001", + arguments: {}, + }, + }; + + for (const cache of [undefined, createToolAccessDecisionCache()]) { + recorder.statements.length = 0; + const decision = await policy.decide(input, { cache }); + expect(decision).toMatchObject({ + allowed: true, + reasonCode: "allow_profile", + effectiveProfileIds: [projectProfile.id], + }); + await policy.writeAudit(input, decision); + const recorded = await policy.recordInvocation(input, decision); + expect(recorded.invocation).toMatchObject({ + runId: fixture.run.id, + issueId: fixture.issue.id, + }); + + expect(recorder.statements.some((statement) => statement.includes('"heartbeat_runs"'))).toBe(true); + for (const statement of recorder.statements) { + expect(statement).not.toMatch(/"context_snapshot"(?!\s*->)/); + expect(statement).not.toContain("result_json"); + expect(statement).not.toContain('"input_schema"'); + } + } + }); + + it("reads run context ids from the snapshot with the same rules as before", async () => { + const fixture = await createListingFixture(db, 3, { broadBindings: false }); + const policy = toolAccessPolicyService(db); + const otherIssue = await db.insert(issues).values({ + companyId: fixture.company.id, + projectId: fixture.project.id, + title: "Other work", + status: "in_progress", + }).returning().then((rows) => rows[0]!); + const projectProfile = await db.insert(toolProfiles).values({ + companyId: fixture.company.id, + profileKey: `project-${randomUUID()}`, + name: `Project profile ${randomUUID()}`, + defaultAction: "allow", + }).returning().then((rows) => rows[0]!); + await db.insert(toolProfileBindings).values({ + companyId: fixture.company.id, + profileId: projectProfile.id, + targetType: "project", + targetId: fixture.project.id, + }); + const runWithSnapshot = (contextSnapshot: Record) => + db.insert(heartbeatRuns).values({ + companyId: fixture.company.id, + agentId: fixture.agent.id, + invocationSource: "assignment", + status: "running", + contextSnapshot, + }).returning().then((rows) => rows[0]!); + const decideForRun = (runId: string, runContext: Partial> = {}) => + policy.decide({ + companyId: fixture.company.id, + actor: { actorType: "agent", actorId: fixture.agent.id, agentId: fixture.agent.id }, + runContext: { heartbeatRunId: runId, ...runContext }, + request: { + catalogEntryId: fixture.entries[1]!.id, + connectionId: fixture.connection.id, + toolName: "tool_0001", + arguments: {}, + }, + }); + + // The issue id alone resolves the project through the issue row. + const issueOnly = await runWithSnapshot({ issueId: fixture.issue.id }); + expect(await decideForRun(issueOnly.id)).toMatchObject({ reasonCode: "allow_profile" }); + + // Non-string and blank ids are ignored, as before. + const malformed = await runWithSnapshot({ issueId: 42, projectId: " ", routineId: { id: "x" } }); + expect(await decideForRun(malformed.id)).toMatchObject({ reasonCode: "deny_default" }); + expect(await decideForRun(malformed.id, { projectId: fixture.project.id })) + .toMatchObject({ reasonCode: "allow_profile" }); + + // A supplied context that disagrees with the stored snapshot is denied. + expect(await decideForRun(issueOnly.id, { issueId: otherIssue.id })) + .toMatchObject({ reasonCode: "deny_run_context_mismatch" }); + }); + + it("returns the same decisions with a shared cache as without one", async () => { + const fixture = await createListingFixture(db, 6); + const other = await createListingFixture(db, 3); + const secondAgent = await db.insert(agents).values({ + companyId: fixture.company.id, + name: `Second Agent ${randomUUID()}`, + role: "engineer", + adapterType: "process", + adapterConfig: {}, + runtimeConfig: {}, + permissions: {}, + }).returning().then((rows) => rows[0]!); + const agentProfile = await db.insert(toolProfiles).values({ + companyId: fixture.company.id, + profileKey: `second-agent-${randomUUID()}`, + name: `Second agent profile ${randomUUID()}`, + defaultAction: "allow", + }).returning().then((rows) => rows[0]!); + await db.insert(toolProfileBindings).values({ + companyId: fixture.company.id, + profileId: agentProfile.id, + targetType: "agent", + targetId: secondAgent.id, + }); + const disabledConnection = await db.insert(toolConnections).values({ + companyId: fixture.company.id, + applicationId: fixture.application.id, + name: "Disabled connection", + uid: `test/${randomUUID()}`, + transport: "mcp_remote", + status: "disabled", + enabled: false, + healthStatus: "ok", + config: { url: "https://8.8.8.8/mcp" }, + }).returning().then((rows) => rows[0]!); + const disabledEntry = await db.insert(toolCatalogEntries).values({ + companyId: fixture.company.id, + applicationId: fixture.application.id, + connectionId: disabledConnection.id, + name: "disabled_tool", + toolName: "disabled_tool", + riskLevel: "read", + versionHash: randomUUID(), + }).returning().then((rows) => rows[0]!); + // Listings never show a quarantined entry, so a cached decision for it + // takes the direct query. + const quarantinedEntry = await db.insert(toolCatalogEntries).values({ + companyId: fixture.company.id, + applicationId: fixture.application.id, + connectionId: fixture.connection.id, + name: "quarantined_tool", + toolName: "quarantined_tool", + riskLevel: "read", + status: "quarantined", + quarantinedAt: new Date(), + versionHash: randomUUID(), + }).returning().then((rows) => rows[0]!); + + const requestFor = (entry: { id: string; connectionId: string; toolName: string }) => ({ + catalogEntryId: entry.id, + connectionId: entry.connectionId, + toolName: entry.toolName, + arguments: {}, + }); + const firstAgentActor = { actorType: "agent" as const, actorId: fixture.agent.id, agentId: fixture.agent.id }; + const secondAgentActor = { actorType: "agent" as const, actorId: secondAgent.id, agentId: secondAgent.id }; + const inputs: ToolAccessDecisionInput[] = [ + ...fixture.entries.map((entry) => ({ + companyId: fixture.company.id, + actor: firstAgentActor, + runContext: { heartbeatRunId: fixture.run.id, gatewayId: fixture.namedGateway.id }, + request: requestFor(entry), + })), + ...fixture.entries.map((entry) => ({ + companyId: fixture.company.id, + actor: secondAgentActor, + request: requestFor(entry), + })), + { + companyId: fixture.company.id, + actor: secondAgentActor, + // The run belongs to the first agent. + runContext: { heartbeatRunId: fixture.run.id }, + request: requestFor(fixture.entries[3]!), + }, + { + companyId: fixture.company.id, + actor: secondAgentActor, + request: requestFor(disabledEntry), + }, + { + companyId: fixture.company.id, + actor: secondAgentActor, + request: requestFor(quarantinedEntry), + }, + { + companyId: fixture.company.id, + actor: secondAgentActor, + // A catalog entry of another company is not in this company's catalog. + request: requestFor(other.entries[0]!), + }, + { + companyId: fixture.company.id, + actor: secondAgentActor, + request: { connectionId: fixture.connection.id, toolName: "tool_0004", arguments: {} }, + }, + { + companyId: fixture.company.id, + actor: secondAgentActor, + // The database matches a uuid in any letter case. + request: { ...requestFor(fixture.entries[4]!), catalogEntryId: fixture.entries[4]!.id.toUpperCase() }, + }, + { + companyId: fixture.company.id, + actor: { actorType: "system", actorId: fixture.company.id }, + request: requestFor(fixture.entries[4]!), + }, + ]; + + const policy = toolAccessPolicyService(db); + const uncached = []; + for (const input of inputs) uncached.push(await policy.decide(input)); + // One cache across mixed actors and runs: the keys must keep them apart. + const cache = createToolAccessDecisionCache(); + const cached = await Promise.all(inputs.map((input) => policy.decide(input, { cache }))); + + expect(cached).toEqual(uncached); + expect(new Set(uncached.map((decision) => decision.reasonCode))).toEqual(new Set([ + "allow_profile", + "deny_default", + "deny_policy_block", + "requires_approval_policy", + "deny_run_context_mismatch", + "deny_disabled_connection", + "deny_missing_tool", + ])); + }); + + it("reads fresh rows for a decision that consumes a rate limit", async () => { + const fixture = await createListingFixture(db, 3); + const policy = toolAccessPolicyService(db); + const cache = createToolAccessDecisionCache(); + + await policy.decide({ + companyId: fixture.company.id, + actor: { actorType: "agent", actorId: fixture.agent.id, agentId: fixture.agent.id }, + runContext: { heartbeatRunId: fixture.run.id }, + request: { + catalogEntryId: fixture.entries[0]!.id, + connectionId: fixture.connection.id, + toolName: "tool_0000", + arguments: {}, + }, + consumeRateLimit: true, + }, { cache }); + + expect(cache.size).toBe(0); + }); +}); diff --git a/server/src/index.ts b/server/src/index.ts index 4b917f08a2..92c57e4b95 100644 --- a/server/src/index.ts +++ b/server/src/index.ts @@ -1293,6 +1293,7 @@ async function startServerWithDatabaseTeardown( await chatCompletionDeliveries.sweepPending().catch((err) => logger.error({ err }, "startup chat completion delivery recovery failed")); await connectionDeliveries.sweepPending(); await app.locals.toolGateway.sweepActionReviews().catch((err: unknown) => logger.error({ err }, "startup tool review recovery failed")); + await app.locals.toolGateway.cleanupExpiredSessions().catch((err: unknown) => logger.error({ err }, "startup gateway token cleanup failed")); await app.locals.toolActionDeliveries.sweepPending().catch((err: unknown) => logger.error({ err }, "startup tool review delivery sweep failed")); await questionResponseDeliveries.sweepPending().then((result) => { if (result.scanned > 0) { @@ -1757,6 +1758,7 @@ async function startServerWithDatabaseTeardown( trackHeartbeatSchedulerWork(chatCompletionDeliveries.sweepPending().catch((err) => logger.error({ err }, "chat completion delivery failed"))); trackHeartbeatSchedulerWork(connectionDeliveries.sweepPending().catch((err) => logger.error({ err }, "connection continuation delivery failed"))); trackHeartbeatSchedulerWork(app.locals.toolGateway.sweepActionReviews().catch((err: unknown) => logger.error({ err }, "tool review recovery failed"))); + trackHeartbeatSchedulerWork(app.locals.toolGateway.cleanupExpiredSessions().catch((err: unknown) => logger.error({ err }, "gateway token cleanup failed"))); trackHeartbeatSchedulerWork(app.locals.toolActionDeliveries.sweepPending().catch((err: unknown) => logger.error({ err }, "tool review delivery sweep failed"))); trackHeartbeatSchedulerWork(questionResponseDeliveries.sweepPending() .then((result) => { diff --git a/server/src/routes/connection-intents.ts b/server/src/routes/connection-intents.ts index ddfaac1470..9d53c87ec9 100644 --- a/server/src/routes/connection-intents.ts +++ b/server/src/routes/connection-intents.ts @@ -59,7 +59,7 @@ export function runtimeConnectionIntentRoutes(db: Db) { router.get("/mcp/runtime-tools", async (req, res) => { await service.validate(runtimeClaims(req)); - res.json({ name: "paperclip-runtime-tools", protocolVersion: "2025-03-26" }); + res.set("Allow", "POST").status(405).end(); }); router.post("/mcp/runtime-tools", async (req, res) => { diff --git a/server/src/routes/openapi.ts b/server/src/routes/openapi.ts index 829064f446..4895b3e2d2 100644 --- a/server/src/routes/openapi.ts +++ b/server/src/routes/openapi.ts @@ -10255,7 +10255,8 @@ registerCurrentRoute({ method: "get", path: "/mcp/runtime-tools", tags: ["connection-intents"], - summary: "Inspect the heartbeat-bound runtime tools MCP endpoint", + summary: "Reject SSE discovery because the runtime tools endpoint supports POST only", + responses: { 405: { description: "SSE stream is not supported" }, 401: r.unauthorized, 403: r.forbidden }, }); registerCurrentRoute({ diff --git a/server/src/routes/tool-gateway.ts b/server/src/routes/tool-gateway.ts index fc73386349..55818267ac 100644 --- a/server/src/routes/tool-gateway.ts +++ b/server/src/routes/tool-gateway.ts @@ -55,6 +55,21 @@ function callerHeaders(req: { headers: Record(req: Request, res: Response, work: (signal: AbortSignal) => Promise) { + const controller = new AbortController(); + const aborted = () => controller.abort(); + const closed = () => { if (!res.writableEnded) controller.abort(); }; + req.once("aborted", aborted); + res.once("close", closed); + if (req.aborted || res.destroyed) controller.abort(); + try { + return await work(controller.signal); + } finally { + req.removeListener("aborted", aborted); + res.removeListener("close", closed); + } +} + async function handleMcpGatewayProtocol( req: Request, res: Response, @@ -96,11 +111,12 @@ async function handleMcpGatewayProtocol( return; } if (body.method === "tools/list") { - const tools = await toolGateway.listToolsForNamedGateway({ + const tools = await discoveryRequest(req, res, (signal) => toolGateway.listToolsForNamedGateway({ ...locator, bearerToken: token, callerHeaders: headers, - }); + signal, + })); res.json({ jsonrpc: "2.0", id, @@ -202,6 +218,7 @@ async function handleMcpGatewayProtocol( } res.status(404).json({ jsonrpc: "2.0", id, error: { code: -32601, message: "Method not found" } }); } catch (err) { + if (res.destroyed) return; if (err instanceof ToolGatewayHttpError) { const id = (req.body as { id?: unknown } | undefined)?.id ?? null; // Provider tool failures are MCP tool results, not successful calls or @@ -227,12 +244,8 @@ async function handleMcpGatewayProtocol( export function mcpGatewayProtocolRoutes(toolGateway: ToolGatewayService) { const router = Router(); - router.get("/mcp/gateways/:gatewayPublicId", async (req, res) => { - res.json({ - transport: "streamable_http", - endpoint: `/mcp/gateways/${req.params.gatewayPublicId}`, - authentication: "bearer", - }); + router.get("/mcp/gateways/:gatewayPublicId", async (_req, res) => { + res.set("Allow", "POST").status(405).end(); }); router.post("/mcp/gateways/:gatewayPublicId", async (req, res) => { await handleMcpGatewayProtocol(req, res, toolGateway, { gatewayPublicId: req.params.gatewayPublicId }); @@ -312,6 +325,7 @@ function outcomeCondition(outcome: string) { } function sendGatewayError(res: import("express").Response, err: unknown) { + if (res.destroyed) return; if (err instanceof ToolGatewayHttpError) { res.status(err.status).json({ error: err.message, @@ -450,12 +464,8 @@ export function toolGatewayRoutes(db: Db, toolGateway: ToolGatewayService) { } }); - router.get("/tool-gateway/gateways/:gatewayId/mcp", async (req, res) => { - res.json({ - transport: "streamable_http", - endpoint: `/api/tool-gateway/gateways/${req.params.gatewayId}/mcp`, - authentication: "bearer", - }); + router.get("/tool-gateway/gateways/:gatewayId/mcp", async (_req, res) => { + res.set("Allow", "POST").status(405).end(); }); router.post("/tool-gateway/gateways/:gatewayId/mcp", async (req, res) => { @@ -551,7 +561,7 @@ export function toolGatewayRoutes(db: Db, toolGateway: ToolGatewayService) { res.status(401).json({ error: "Tool gateway session token is required" }); return; } - const tools = await toolGateway.listToolsForSession(token); + const tools = await discoveryRequest(req, res, (signal) => toolGateway.listToolsForSession(token, { signal })); res.json(tools); } catch (err) { sendGatewayError(res, err); diff --git a/server/src/services/chat-github-tools.ts b/server/src/services/chat-github-tools.ts index 572a78a0bf..f8bcc159ff 100644 --- a/server/src/services/chat-github-tools.ts +++ b/server/src/services/chat-github-tools.ts @@ -21,6 +21,7 @@ import type { ToolGatewayDescriptor, ToolGatewaySession, } from "./tool-gateway.js"; +import { runContextSnapshotString } from "./tool-access-policy.js"; const objectSchema = ( properties: Record, @@ -304,7 +305,7 @@ export async function githubBotConnectionIdsForRun( const [run] = await db .select({ issueId: heartbeatRuns.nativeIssueId, - context: heartbeatRuns.contextSnapshot, + snapshotIssueId: runContextSnapshotString("issueId"), }) .from(heartbeatRuns) .where( @@ -314,7 +315,7 @@ export async function githubBotConnectionIdsForRun( eq(heartbeatRuns.id, runId), ), ); - const issueId = run?.issueId ?? run?.context?.issueId; + const issueId = run?.issueId ?? run?.snapshotIssueId; if (typeof issueId !== "string") return new Set(); const rows = await db .select({ @@ -356,7 +357,7 @@ export async function githubGuestBotConnectionForSession( ): Promise { if (!session.runId || !session.agentId || !session.issueId) return null; const [run] = await db - .select() + .select({ wakeCommentId: runContextSnapshotString("wakeCommentId") }) .from(heartbeatRuns) .where( and( @@ -365,7 +366,7 @@ export async function githubGuestBotConnectionForSession( eq(heartbeatRuns.id, session.runId), ), ); - const commentId = run?.contextSnapshot?.wakeCommentId; + const commentId = run?.wakeCommentId; if (typeof commentId !== "string") return null; const [source] = await db .select({ diff --git a/server/src/services/tool-access-policy.ts b/server/src/services/tool-access-policy.ts index e09d1b2872..827203f82d 100644 --- a/server/src/services/tool-access-policy.ts +++ b/server/src/services/tool-access-policy.ts @@ -1,5 +1,5 @@ import { createHash } from "node:crypto"; -import { and, asc, desc, eq, gt, inArray, ne, sql } from "drizzle-orm"; +import { and, asc, desc, eq, gt, inArray, isNull, ne, sql } from "drizzle-orm"; import type { Db } from "@paperclipai/db"; import { agents, @@ -114,11 +114,69 @@ function isRecord(value: unknown): value is Record { return Boolean(value && typeof value === "object" && !Array.isArray(value)); } -function snapshotString(snapshot: Record, key: string): string | null { - const value = snapshot[key]; +function nonBlankString(value: string | null | undefined): string | null { return typeof value === "string" && value.trim().length > 0 ? value : null; } +/** + * Reads one string field of a run's context snapshot in SQL. The snapshot can + * hold hundreds of KB of task text, so a policy check never loads the whole + * document. A non-string value reads as null. Runs store the snapshot as a + * JSON object, which the run indexes on `context_snapshot ->> 'issueId'` also + * expect. + */ +export function runContextSnapshotString(key: "issueId" | "projectId" | "routineId" | "wakeCommentId") { + const field = sql.raw(`'${key}'`); + return sql`case when jsonb_typeof(${heartbeatRuns.contextSnapshot} -> ${field}) = 'string' + then ${heartbeatRuns.contextSnapshot} ->> ${field} end`; +} + +/** + * Memo for one batch of `decide` calls that share a company, actor and run + * context, such as one tools/list response. A listing decides every tool + * against the same agent, run, profiles and policies, so the batch reads each + * of those rows once instead of once per tool. The cache is a point-in-time + * view: create one per batch and drop it with the request. Rate-limit counters + * are never cached. + */ +export type ToolAccessDecisionCache = Map>; + +export function createToolAccessDecisionCache(): ToolAccessDecisionCache { + return new Map(); +} + +function cachedRead( + cache: ToolAccessDecisionCache | undefined, + key: string, + read: () => Promise, +): Promise { + if (!cache) return read(); + const cached = cache.get(key); + if (cached) return cached as Promise; + const pending = read(); + cache.set(key, pending); + // Do not keep a failed read, so a later call in the batch tries again. + pending.catch(() => { + if (cache.get(key) === pending) cache.delete(key); + }); + return pending; +} + +// The database returns ids in this canonical form. +const CANONICAL_UUID_RE = /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/; + +const catalogEntryContextColumns = { + id: toolCatalogEntries.id, + companyId: toolCatalogEntries.companyId, + connectionId: toolCatalogEntries.connectionId, + applicationId: toolCatalogEntries.applicationId, + riskLevel: toolCatalogEntries.riskLevel, + toolName: toolCatalogEntries.toolName, + status: toolCatalogEntries.status, + versionHash: toolCatalogEntries.versionHash, + schemaHash: toolCatalogEntries.schemaHash, +}; + function stableStringify(value: unknown): string { if (!isRecord(value) && !Array.isArray(value)) return JSON.stringify(value); if (Array.isArray(value)) return `[${value.map(stableStringify).join(",")}]`; @@ -838,7 +896,35 @@ export function toolAccessPolicyService(db: Db) { return deleted; } - async function loadContext(input: ToolAccessDecisionInput): Promise< + async function loadCatalogEntry(companyId: string, catalogEntryId: string, cache?: ToolAccessDecisionCache) { + const readEntry = async () => { + const [entry] = await db.select(catalogEntryContextColumns).from(toolCatalogEntries).where(eq(toolCatalogEntries.id, catalogEntryId)); + return entry && entry.companyId === companyId ? entry : null; + }; + // Only a canonical id can use the in-memory lookup. Any other spelling + // takes the direct query, so a cached and an uncached decision agree. + if (!cache || !CANONICAL_UUID_RE.test(catalogEntryId)) return readEntry(); + // A listing decides the active catalog tools of the company, so a batch + // reads those entries once and looks each one up in memory. Any other + // entry takes the direct query. + const entriesById = await cachedRead(cache, `active-catalog-tools:${companyId}`, async () => { + const entries = await db + .select(catalogEntryContextColumns) + .from(toolCatalogEntries) + .where(and( + eq(toolCatalogEntries.companyId, companyId), + eq(toolCatalogEntries.entryKind, "tool"), + eq(toolCatalogEntries.status, "active"), + isNull(toolCatalogEntries.quarantinedAt), + )); + return new Map(entries.map((entry) => [entry.id, entry])); + }); + const entry = entriesById.get(catalogEntryId); + if (!entry) return readEntry(); + return entry.companyId === companyId ? entry : null; + } + + async function loadContext(input: ToolAccessDecisionInput, cache?: ToolAccessDecisionCache): Promise< | { ok: true; ctx: ToolAccessContext; redaction: RedactionResult } | { ok: false; decision: ToolAccessDecision; redaction: RedactionResult } > { @@ -851,7 +937,9 @@ export function toolAccessPolicyService(db: Db) { const gatewayId = input.runContext?.gatewayId ?? null; if (input.actor.actorType === "agent") { - const [agent] = await db.select().from(agents).where(and(eq(agents.id, agentId ?? ""), eq(agents.companyId, input.companyId))); + const lookupAgentId = agentId ?? ""; + const [agent] = await cachedRead(cache, `agent:${input.companyId}:${lookupAgentId}`, async () => + db.select({ id: agents.id }).from(agents).where(and(eq(agents.id, lookupAgentId), eq(agents.companyId, input.companyId)))); if (!agent) { return { ok: false, redaction, decision: decision("deny", "deny_missing_agent", "Authenticated agent was not found in the company.", [], [], { redactionPlan: redaction.redactionPlan }) }; } @@ -859,15 +947,25 @@ export function toolAccessPolicyService(db: Db) { } if (heartbeatRunId) { - const [run] = await db.select().from(heartbeatRuns).where(eq(heartbeatRuns.id, heartbeatRunId)); + const runId = heartbeatRunId; + const [run] = await cachedRead(cache, `run:${runId}`, async () => + db + .select({ + companyId: heartbeatRuns.companyId, + agentId: heartbeatRuns.agentId, + issueId: runContextSnapshotString("issueId"), + projectId: runContextSnapshotString("projectId"), + routineId: runContextSnapshotString("routineId"), + }) + .from(heartbeatRuns) + .where(eq(heartbeatRuns.id, runId))); if (!run || run.companyId !== input.companyId || (input.actor.actorType === "agent" && run.agentId !== agentId)) { return { ok: false, redaction, decision: decision("deny", "deny_run_context_mismatch", "Supplied run context does not match the authenticated actor.", [], [], { redactionPlan: redaction.redactionPlan }) }; } agentId = run.agentId; - const snapshot = isRecord(run.contextSnapshot) ? run.contextSnapshot : {}; - const runIssueId = snapshotString(snapshot, "issueId"); - const runProjectId = snapshotString(snapshot, "projectId"); - const runRoutineId = snapshotString(snapshot, "routineId"); + const runIssueId = nonBlankString(run.issueId); + const runProjectId = nonBlankString(run.projectId); + const runRoutineId = nonBlankString(run.routineId); if ((issueId && runIssueId && issueId !== runIssueId) || (projectId && runProjectId && projectId !== runProjectId) || (routineId && runRoutineId && routineId !== runRoutineId)) { @@ -879,7 +977,9 @@ export function toolAccessPolicyService(db: Db) { } if (issueId) { - const [issue] = await db.select().from(issues).where(eq(issues.id, issueId)); + const contextIssueId = issueId; + const [issue] = await cachedRead(cache, `issue:${contextIssueId}`, async () => + db.select({ companyId: issues.companyId, projectId: issues.projectId }).from(issues).where(eq(issues.id, contextIssueId))); if (!issue || issue.companyId !== input.companyId) { return { ok: false, redaction, decision: decision("deny", "deny_company_boundary", "Issue context is outside the company.", [], [], { redactionPlan: redaction.redactionPlan }) }; } @@ -889,13 +989,17 @@ export function toolAccessPolicyService(db: Db) { projectId = projectId ?? issue.projectId; } if (projectId) { - const [project] = await db.select().from(projects).where(eq(projects.id, projectId)); + const contextProjectId = projectId; + const [project] = await cachedRead(cache, `project:${contextProjectId}`, async () => + db.select({ companyId: projects.companyId }).from(projects).where(eq(projects.id, contextProjectId))); if (!project || project.companyId !== input.companyId) { return { ok: false, redaction, decision: decision("deny", "deny_company_boundary", "Project context is outside the company.", [], [], { redactionPlan: redaction.redactionPlan }) }; } } if (routineId) { - const [routine] = await db.select().from(routines).where(eq(routines.id, routineId)); + const contextRoutineId = routineId; + const [routine] = await cachedRead(cache, `routine:${contextRoutineId}`, async () => + db.select({ companyId: routines.companyId }).from(routines).where(eq(routines.id, contextRoutineId))); if (!routine || routine.companyId !== input.companyId) { return { ok: false, redaction, decision: decision("deny", "deny_company_boundary", "Routine context is outside the company.", [], [], { redactionPlan: redaction.redactionPlan }) }; } @@ -915,8 +1019,8 @@ export function toolAccessPolicyService(db: Db) { let applicationType: string | null = null; if (catalogEntryId) { - const [entry] = await db.select().from(toolCatalogEntries).where(eq(toolCatalogEntries.id, catalogEntryId)); - if (!entry || entry.companyId !== input.companyId) { + const entry = await loadCatalogEntry(input.companyId, catalogEntryId, cache); + if (!entry) { return { ok: false, redaction, decision: decision("deny", "deny_missing_tool", "Requested tool is not in the company catalog.", [], [], { redactionPlan: redaction.redactionPlan }) }; } connectionId = entry.connectionId; @@ -927,10 +1031,12 @@ export function toolAccessPolicyService(db: Db) { catalogVersionHash = entry.versionHash; catalogSchemaHash = entry.schemaHash; } else if (connectionId) { - const [entry] = await db - .select() - .from(toolCatalogEntries) - .where(and(eq(toolCatalogEntries.companyId, input.companyId), eq(toolCatalogEntries.connectionId, connectionId), eq(toolCatalogEntries.name, input.request.toolName))); + const entryConnectionId = connectionId; + const [entry] = await cachedRead(cache, `catalog-entry-name:${input.companyId}:${entryConnectionId}:${input.request.toolName}`, async () => + db + .select(catalogEntryContextColumns) + .from(toolCatalogEntries) + .where(and(eq(toolCatalogEntries.companyId, input.companyId), eq(toolCatalogEntries.connectionId, entryConnectionId), eq(toolCatalogEntries.name, input.request.toolName)))); if (entry) { catalogEntryId = entry.id; applicationId = entry.applicationId ?? applicationId; @@ -943,7 +1049,18 @@ export function toolAccessPolicyService(db: Db) { } if (connectionId) { - const [connection] = await db.select().from(toolConnections).where(eq(toolConnections.id, connectionId)); + const contextConnectionId = connectionId; + const [connection] = await cachedRead(cache, `connection:${contextConnectionId}`, async () => + db + .select({ + companyId: toolConnections.companyId, + applicationId: toolConnections.applicationId, + enabled: toolConnections.enabled, + status: toolConnections.status, + transport: toolConnections.transport, + }) + .from(toolConnections) + .where(eq(toolConnections.id, contextConnectionId))); if (!connection || connection.companyId !== input.companyId) { return { ok: false, redaction, decision: decision("deny", "deny_company_boundary", "Connection is outside the company.", [], [], { redactionPlan: redaction.redactionPlan }) }; } @@ -954,7 +1071,17 @@ export function toolAccessPolicyService(db: Db) { connectionTransport = connection.transport; } if (applicationId) { - const [application] = await db.select().from(toolApplications).where(eq(toolApplications.id, applicationId)); + const contextApplicationId = applicationId; + const [application] = await cachedRead(cache, `application:${contextApplicationId}`, async () => + db + .select({ + companyId: toolApplications.companyId, + applicationKey: toolApplications.applicationKey, + status: toolApplications.status, + type: toolApplications.type, + }) + .from(toolApplications) + .where(eq(toolApplications.id, contextApplicationId))); if (!application || application.companyId !== input.companyId) { return { ok: false, redaction, decision: decision("deny", "deny_company_boundary", "Application is outside the company.", [], [], { redactionPlan: redaction.redactionPlan }) }; } @@ -1004,24 +1131,28 @@ export function toolAccessPolicyService(db: Db) { }; } - async function effectiveProfiles(ctx: ToolAccessContext) { - const bindings = await db.select().from(toolProfileBindings).where(eq(toolProfileBindings.companyId, ctx.companyId)); + async function effectiveProfiles(ctx: ToolAccessContext, cache?: ToolAccessDecisionCache) { + const bindings = await cachedRead(cache, `profile-bindings:${ctx.companyId}`, async () => + db.select().from(toolProfileBindings).where(eq(toolProfileBindings.companyId, ctx.companyId))); const matchingBindings = bindings.filter((binding) => targetMatches(binding, ctx)); if (matchingBindings.length === 0) return { profiles: [], entries: [] as Array }; const candidateProfileIds = profileIdsInBindingOrder(matchingBindings); - const candidateProfiles = await db.select().from(toolProfiles).where(and( - eq(toolProfiles.companyId, ctx.companyId), - inArray(toolProfiles.id, candidateProfileIds), - )); - const [gateway] = ctx.gatewayId - ? await db - .select({ defaultProfileMode: toolMcpGateways.defaultProfileMode }) - .from(toolMcpGateways) - .where(and( - eq(toolMcpGateways.companyId, ctx.companyId), - eq(toolMcpGateways.id, ctx.gatewayId), - )) - .limit(1) + const candidateProfiles = await cachedRead(cache, `profiles:${ctx.companyId}:${[...candidateProfileIds].sort().join(",")}`, async () => + db.select().from(toolProfiles).where(and( + eq(toolProfiles.companyId, ctx.companyId), + inArray(toolProfiles.id, candidateProfileIds), + ))); + const gatewayId = ctx.gatewayId; + const [gateway] = gatewayId + ? await cachedRead(cache, `gateway-profile-mode:${ctx.companyId}:${gatewayId}`, async () => + db + .select({ defaultProfileMode: toolMcpGateways.defaultProfileMode }) + .from(toolMcpGateways) + .where(and( + eq(toolMcpGateways.companyId, ctx.companyId), + eq(toolMcpGateways.id, gatewayId), + )) + .limit(1)) : []; const activeBindings = effectiveToolProfileBindings( matchingBindings, @@ -1035,25 +1166,29 @@ export function toolAccessPolicyService(db: Db) { .map((profileId) => profilesById.get(profileId) ?? null) .filter((profile): profile is typeof toolProfiles.$inferSelect => Boolean(profile && profile.status === "active")); const activeProfileIds = activeProfiles.map((profile) => profile.id); + // The active profile set depends on the connection (additive app profiles), + // so a batch reads the entries once per distinct set. const entries = activeProfileIds.length > 0 - ? await db.select().from(toolProfileEntries).where(and(eq(toolProfileEntries.companyId, ctx.companyId), inArray(toolProfileEntries.profileId, activeProfileIds))) + ? await cachedRead(cache, `profile-entries:${ctx.companyId}:${[...activeProfileIds].sort().join(",")}`, async () => + db.select().from(toolProfileEntries).where(and(eq(toolProfileEntries.companyId, ctx.companyId), inArray(toolProfileEntries.profileId, activeProfileIds)))) : []; return { profiles: activeProfiles, entries }; } - async function explicitGrant(ctx: ToolAccessContext): Promise { + async function explicitGrant(ctx: ToolAccessContext, cache?: ToolAccessDecisionCache): Promise { const principalType = ctx.actorType === "agent" ? "agent" : ctx.actorType === "user" ? "user" : null; const principalId = ctx.actorType === "agent" ? ctx.agentId : ctx.actorId; if (!principalType || !principalId) return false; - const grants = await db - .select() - .from(principalPermissionGrants) - .where(and( - eq(principalPermissionGrants.companyId, ctx.companyId), - eq(principalPermissionGrants.principalType, principalType), - eq(principalPermissionGrants.principalId, principalId), - eq(principalPermissionGrants.permissionKey, "tools:use"), - )); + const grants = await cachedRead(cache, `tool-grants:${ctx.companyId}:${principalType}:${principalId}`, async () => + db + .select() + .from(principalPermissionGrants) + .where(and( + eq(principalPermissionGrants.companyId, ctx.companyId), + eq(principalPermissionGrants.principalType, principalType), + eq(principalPermissionGrants.principalId, principalId), + eq(principalPermissionGrants.permissionKey, "tools:use"), + ))); return grants.some((grant) => scopeAllowsTool(grant.scope, ctx)); } @@ -1177,18 +1312,30 @@ export function toolAccessPolicyService(db: Db) { }); } - async function decide(input: ToolAccessDecisionInput): Promise { - const loaded = await loadContext(input); + /** + * Decides one tool access request. Pass the same `cache` to every call of a + * batch that shares one actor and run context (for example, one tools/list + * response) so the batch reads the shared policy inputs once. + */ + async function decide( + input: ToolAccessDecisionInput, + options: { cache?: ToolAccessDecisionCache } = {}, + ): Promise { + // A decision that consumes a rate limit can also record a trust-rule hit + // from the policy row, so it always reads fresh rows. + const cache = input.consumeRateLimit === true ? undefined : options.cache; + const loaded = await loadContext(input, cache); if (!loaded.ok) return loaded.decision; const { ctx, redaction } = loaded; - const profileState = await effectiveProfiles(ctx); + const profileState = await effectiveProfiles(ctx, cache); const effectiveProfileIds = profileState.profiles.map((profile) => profile.id); const permittedByProfile = profileState.profiles.some((profile) => { const matchingEntries = profileState.entries.filter((entry) => entry.profileId === profile.id && profileEntryMatches(entry, ctx)); return !matchingEntries.some((entry) => entry.effect === "exclude") && (profile.defaultAction === "allow" || matchingEntries.some((entry) => entry.effect === "include")); }); - const policies = await db.select().from(toolPolicies).where(and(eq(toolPolicies.companyId, ctx.companyId), eq(toolPolicies.enabled, true))).orderBy(asc(toolPolicies.priority), asc(toolPolicies.createdAt)); + const policies = await cachedRead(cache, `enabled-policies:${ctx.companyId}`, async () => + db.select().from(toolPolicies).where(and(eq(toolPolicies.companyId, ctx.companyId), eq(toolPolicies.enabled, true))).orderBy(asc(toolPolicies.priority), asc(toolPolicies.createdAt))); for (const policy of policies) { const conditions = policyConditions(policy); if (conditions && selectorMatches(policy.selectors, ctx)) { @@ -1290,7 +1437,7 @@ export function toolAccessPolicyService(db: Db) { return decision("allow", "allow_policy", "Tool access allowed by policy.", effectiveProfileIds, [policy.id], { redactionPlan: redaction.redactionPlan, policyExplanation }); } } - if (await explicitGrant(ctx)) { + if (await explicitGrant(ctx, cache)) { return decision("allow", "allow_explicit_grant", "Tool access allowed by explicit grant.", effectiveProfileIds, [], { redactionPlan: redaction.redactionPlan }); } diff --git a/server/src/services/tool-discovery-scheduler.test.ts b/server/src/services/tool-discovery-scheduler.test.ts new file mode 100644 index 0000000000..b50f37079c --- /dev/null +++ b/server/src/services/tool-discovery-scheduler.test.ts @@ -0,0 +1,48 @@ +import { describe, expect, it } from "vitest"; +import { createToolDiscoveryScheduler } from "./tool-discovery-scheduler.js"; + +function deferred() { + let resolve!: () => void; + const promise = new Promise((done) => { resolve = done; }); + return { promise, resolve }; +} + +describe("tool discovery admission", () => { + it("bounds whole listings and rejects excess queued requests", async () => { + const scheduler = createToolDiscoveryScheduler(2, 2); + const gate = deferred(); + let active = 0; + let peak = 0; + const work = async () => { + active += 1; + peak = Math.max(peak, active); + await gate.promise; + active -= 1; + }; + const accepted = Array.from({ length: 4 }, () => scheduler.run(work)); + await expect(scheduler.run(work)).rejects.toMatchObject({ name: "ToolDiscoveryBusyError" }); + expect(active).toBe(2); + gate.resolve(); + await Promise.all(accepted); + expect(peak).toBe(2); + await scheduler.run(async () => {}); + }); + + it("removes abandoned queued requests and releases slots after failures", async () => { + const scheduler = createToolDiscoveryScheduler(1, 1); + const gate = deferred(); + const first = scheduler.run(() => gate.promise); + const controller = new AbortController(); + let abandonedStarted = false; + const abandoned = scheduler.run(async () => { abandonedStarted = true; }, controller.signal); + controller.abort(); + await expect(abandoned).rejects.toMatchObject({ name: "AbortError" }); + const next = scheduler.run(async () => { throw new Error("Read failed"); }); + const failure = expect(next).rejects.toThrow("Read failed"); + gate.resolve(); + await first; + await failure; + expect(abandonedStarted).toBe(false); + await scheduler.run(async () => {}); + }); +}); diff --git a/server/src/services/tool-discovery-scheduler.ts b/server/src/services/tool-discovery-scheduler.ts new file mode 100644 index 0000000000..b680729bdd --- /dev/null +++ b/server/src/services/tool-discovery-scheduler.ts @@ -0,0 +1,54 @@ +export class ToolDiscoveryBusyError extends Error { + constructor() { + super("Tool discovery is busy. Retry shortly."); + this.name = "ToolDiscoveryBusyError"; + } +} + +/** Limits whole listings, including catalog reads, across gateway service instances. */ +export function createToolDiscoveryScheduler(concurrency = 2, queueLimit = 32) { + if (!Number.isInteger(concurrency) || concurrency < 1 || !Number.isInteger(queueLimit) || queueLimit < 0) { + throw new Error("Invalid tool discovery scheduler limits"); + } + let active = 0; + const queue: Array<{ start: () => void }> = []; + function drain() { + while (active < concurrency && queue.length > 0) { + active += 1; + queue.shift()!.start(); + } + } + function run(work: () => Promise, signal?: AbortSignal): Promise { + if (signal?.aborted) return Promise.reject(signal.reason); + if (active >= concurrency && queue.length >= queueLimit) { + return Promise.reject(new ToolDiscoveryBusyError()); + } + return new Promise((resolve, reject) => { + const job = { + start() { + signal?.removeEventListener("abort", abort); + // Do not free a slot on abort until already-started reads settle. + Promise.resolve().then(() => { + signal?.throwIfAborted(); + return work(); + }).then(resolve, reject).finally(() => { + active -= 1; + drain(); + }); + }, + }; + function abort() { + const index = queue.indexOf(job); + if (index < 0) return; + queue.splice(index, 1); + reject(signal!.reason); + } + queue.push(job); + signal?.addEventListener("abort", abort, { once: true }); + drain(); + }); + } + return { run }; +} + +export const toolDiscoveryScheduler = createToolDiscoveryScheduler(); diff --git a/server/src/services/tool-gateway.ts b/server/src/services/tool-gateway.ts index 38ae316b31..c601c6bc8b 100644 --- a/server/src/services/tool-gateway.ts +++ b/server/src/services/tool-gateway.ts @@ -127,7 +127,13 @@ import { REMOTE_URL_SECRET_CONFIG_PATH, remoteUrlCredentialMatchesPublicUrl, } from "./remote-url-credentials.js"; -import { toolAccessPolicyService } from "./tool-access-policy.js"; +import { + createToolAccessDecisionCache, + runContextSnapshotString, + toolAccessPolicyService, + type ToolAccessDecisionCache, +} from "./tool-access-policy.js"; +import { toolDiscoveryScheduler, ToolDiscoveryBusyError } from "./tool-discovery-scheduler.js"; import { commitToolActionReview } from "./tool-action-review.js"; import { issueThreadInteractionService } from "./issue-thread-interactions.js"; import { @@ -217,6 +223,10 @@ const ACTION_REQUEST_EXECUTION_WAIT_MS = APPROVED_EXECUTION_TIMEOUT_MS + 5_000; // create) so a live create keeps its own row. const MAX_REMOTE_MCP_RESPONSE_BYTES = 1_000_000; const ACTIVE_GATEWAY_RUN_STATUSES = new Set(["running"]); +// A tool listing decides access for every tool in the company. The decisions +// of one listing share a read cache, so this cap only limits how many of them +// (and their uncached reads, such as rate-limit counters) run at the same time. +const LISTING_DECISION_CONCURRENCY = 16; const uuidPattern = /^[0-9a-f]{8}-[0-9a-f]{4}-[1-5][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/i; @@ -458,6 +468,40 @@ function asRecord(value: unknown): Record | null { return value as Record; } +async function mapWithConcurrency( + items: readonly T[], + concurrency: number, + mapper: (item: T) => Promise, + signal?: AbortSignal, +): Promise { + const results = new Array(items.length); + let nextIndex = 0; + let failed = false; + + async function worker() { + while (!failed && nextIndex < items.length) { + const index = nextIndex; + nextIndex += 1; + try { + signal?.throwIfAborted(); + results[index] = await mapper(items[index]!); + } catch (error) { + // The batch fails with this error, so the other workers stop. + failed = true; + throw error; + } + } + } + + const workers = await Promise.allSettled( + Array.from({ length: Math.min(concurrency, items.length) }, () => worker()), + ); + const failure = workers.find((result) => result.status === "rejected"); + if (failure?.status === "rejected") throw failure.reason; + signal?.throwIfAborted(); + return results; +} + const sensitivePassthroughHeaderPattern = /(^|[-_])(auth|authorization|cookie|secret|session|token)([-_]|$)|(^|[-_])api[-_]?key([-_]|$)/i; const sensitivePassthroughHeaderNames = new Set([ @@ -1216,49 +1260,92 @@ export function createToolGatewayService( return [...BUILTIN_TOOLS, ...pluginTools()]; } + function connectedMcpConnectionFilter(companyId: string) { + return and( + eq(toolConnections.companyId, companyId), + inArray(toolConnections.transport, ["mcp_remote", "local_stdio", "rest_api"]), + eq(toolConnections.status, "active"), + eq(toolConnections.enabled, true), + // A personal connection has no company-level credential to probe. A + // credential-less health sweep can therefore mark it as errored even + // while the responsible user's grant is valid. Keep its cached active + // catalog discoverable; execution resolves and validates that user's + // grant, and a successful call restores the shared health indicator. + or( + inArray(toolConnections.healthStatus, ["ok", "healthy"]), + eq(toolConnections.credentialPolicy, "per_user"), + ), + eq(toolApplications.companyId, companyId), + inArray(toolApplications.type, ["mcp_http", "mcp_stdio", "rest_api"]), + eq(toolApplications.status, "active"), + ); + } + async function connectedMcpToolsForCompany( companyId: string, ): Promise { - const rows = await db - .select({ - catalogEntry: toolCatalogEntries, - connection: toolConnections, - application: toolApplications, - }) - .from(toolCatalogEntries) - .innerJoin( - toolConnections, - eq(toolCatalogEntries.connectionId, toolConnections.id), - ) - .innerJoin( - toolApplications, - eq(toolConnections.applicationId, toolApplications.id), - ) - .where( - and( - eq(toolCatalogEntries.companyId, companyId), - eq(toolCatalogEntries.entryKind, "tool"), - eq(toolCatalogEntries.status, "active"), - isNull(toolCatalogEntries.quarantinedAt), - eq(toolConnections.companyId, companyId), - inArray(toolConnections.transport, ["mcp_remote", "local_stdio", "rest_api"]), - eq(toolConnections.status, "active"), - eq(toolConnections.enabled, true), - // A personal connection has no company-level credential to probe. A - // credential-less health sweep can therefore mark it as errored even - // while the responsible user's grant is valid. Keep its cached active - // catalog discoverable; execution resolves and validates that user's - // grant, and a successful call restores the shared health indicator. - or( - inArray(toolConnections.healthStatus, ["ok", "healthy"]), - eq(toolConnections.credentialPolicy, "per_user"), + // Read each connection and application row once. Selecting them in the + // catalog join repeats the full connection row (config included) for + // every tool, and one connection can expose hundreds of tools. Both reads + // share one read-only snapshot, so they agree like a single join, and the + // second read gets only the connections that have listed tools. + const { connectionRows, catalogRows } = await db.transaction( + async (tx) => { + const catalogRows = await tx + .select({ catalogEntry: toolCatalogEntries }) + .from(toolCatalogEntries) + .innerJoin( + toolConnections, + eq(toolCatalogEntries.connectionId, toolConnections.id), + ) + .innerJoin( + toolApplications, + eq(toolConnections.applicationId, toolApplications.id), + ) + .where( + and( + eq(toolCatalogEntries.companyId, companyId), + eq(toolCatalogEntries.entryKind, "tool"), + eq(toolCatalogEntries.status, "active"), + isNull(toolCatalogEntries.quarantinedAt), + connectedMcpConnectionFilter(companyId), + ), + ) + .orderBy(toolConnections.name, toolCatalogEntries.name); + const connectionIds = [ + ...new Set( + catalogRows.map(({ catalogEntry }) => catalogEntry.connectionId), ), - eq(toolApplications.companyId, companyId), - inArray(toolApplications.type, ["mcp_http", "mcp_stdio", "rest_api"]), - eq(toolApplications.status, "active"), - ), - ) - .orderBy(toolConnections.name, toolCatalogEntries.name); + ]; + if (connectionIds.length === 0) + return { connectionRows: [], catalogRows }; + const connectionRows = await tx + .select({ + connection: toolConnections, + application: toolApplications, + }) + .from(toolConnections) + .innerJoin( + toolApplications, + eq(toolConnections.applicationId, toolApplications.id), + ) + .where( + and( + inArray(toolConnections.id, connectionIds), + connectedMcpConnectionFilter(companyId), + ), + ); + return { connectionRows, catalogRows }; + }, + { isolationLevel: "repeatable read", accessMode: "read only" }, + ); + const connectionRowsById = new Map( + connectionRows.map((row) => [row.connection.id, row]), + ); + const rows = catalogRows.flatMap(({ catalogEntry }) => { + const connectionRow = connectionRowsById.get(catalogEntry.connectionId); + return connectionRow ? [{ catalogEntry, ...connectionRow }] : []; + }); const eligibleRows = rows.filter( ({ catalogEntry, connection, application }) => @@ -1519,7 +1606,8 @@ export function createToolGatewayService( companyId: heartbeatRuns.companyId, agentId: heartbeatRuns.agentId, status: heartbeatRuns.status, - contextSnapshot: heartbeatRuns.contextSnapshot, + issueId: runContextSnapshotString("issueId"), + projectId: runContextSnapshotString("projectId"), }) .from(heartbeatRuns) .where(eq(heartbeatRuns.id, input.runId)) @@ -1543,9 +1631,8 @@ export function createToolGatewayService( throw new ToolGatewayHttpError(403, "Run is not active", "run_inactive"); } - const snapshot = asRecord(run.contextSnapshot); - const snapshotIssueId = stringValue(snapshot?.issueId); - const snapshotProjectId = stringValue(snapshot?.projectId); + const snapshotIssueId = stringValue(run.issueId); + const snapshotProjectId = stringValue(run.projectId); if ( (input.issueId && snapshotIssueId && input.issueId !== snapshotIssueId) || (input.projectId && @@ -1641,6 +1728,10 @@ export function createToolGatewayService( ? "failure" : "success"; try { + const tokenId = input.session?.gatewayTokenId && uuidPattern.test(input.session.gatewayTokenId) + ? input.session.gatewayTokenId + : typeof input.details.gatewayTokenId === "string" && uuidPattern.test(input.details.gatewayTokenId) + ? input.details.gatewayTokenId : null; await db.insert(toolAccessAuditEvents).values({ companyId: input.companyId, gatewayId: @@ -1649,14 +1740,12 @@ export function createToolGatewayService( uuidPattern.test(input.details.gatewayId) ? input.details.gatewayId : null), - gatewayTokenId: - input.session?.gatewayTokenId && - uuidPattern.test(input.session.gatewayTokenId) - ? input.session.gatewayTokenId - : typeof input.details.gatewayTokenId === "string" && - uuidPattern.test(input.details.gatewayTokenId) - ? input.details.gatewayTokenId - : null, + // Cleanup can remove a token after request admission. Resolve the FK + // inside this INSERT and lock a surviving row until the statement ends. + // A plain existence read before insertion still races with deletion. + gatewayTokenId: tokenId ? sql`(select ${toolMcpGatewayTokens.id} from ${toolMcpGatewayTokens} + where ${toolMcpGatewayTokens.id} = ${tokenId} and ${toolMcpGatewayTokens.companyId} = ${input.companyId} + for key share)` : null, gatewayPublicId: typeof input.details.gatewayPublicId === "string" ? input.details.gatewayPublicId @@ -2633,6 +2722,23 @@ export function createToolGatewayService( }; } + /** + * Decides access for each tool of one listing. All decisions of the listing + * share one actor and run context, so they share one read cache: a listing + * reads the agent, run, profiles and policies once, not once per tool. + */ + async function decideToolsForListing( + tools: readonly T[], + inputForTool: (tool: T) => ToolAccessDecisionInput, + cache: ToolAccessDecisionCache = createToolAccessDecisionCache(), + signal?: AbortSignal, + ): Promise> { + return mapWithConcurrency(tools, LISTING_DECISION_CONCURRENCY, async (tool) => ({ + tool, + decision: await policyService.decide(inputForTool(tool), { cache }), + }), signal); + } + function policyErrorStatus(decision: ToolAccessDecision) { if (decision.decision === "rate_limited") return 429; return 403; @@ -2809,13 +2915,8 @@ export function createToolGatewayService( const tools = (await connectedMcpToolsForCompany(session.companyId)).filter( isOnDemandRemoteTool, ); - const decisions = await Promise.all( - tools.map(async (tool) => ({ - tool, - decision: await policyService.decide( - policyInputForTool({ session, tool }), - ), - })), + const decisions = await decideToolsForListing(tools, (tool) => + policyInputForTool({ session, tool }), ); return decisions .filter( @@ -2867,7 +2968,23 @@ export function createToolGatewayService( async function listToolsForContext( session: ToolGatewaySession, + signal?: AbortSignal, ): Promise { + try { + return await toolDiscoveryScheduler.run(() => buildToolsForContext(session, signal), signal); + } catch (error) { + if (error instanceof ToolDiscoveryBusyError) { + throw new ToolGatewayHttpError(503, error.message, "tool_discovery_busy"); + } + throw error; + } + } + + async function buildToolsForContext( + session: ToolGatewaySession, + signal?: AbortSignal, + ): Promise { + signal?.throwIfAborted(); if (session.agentId) { await assertAgentInCompany(session.companyId, session.agentId); } @@ -2887,13 +3004,12 @@ export function createToolGatewayService( (tool.providerType !== "paperclip_self" && tool.providerType !== "paperclip_plugin"), ); - const decisions = await Promise.all( - tools.map(async (tool) => { - const decision = await policyService.decide( - policyInputForTool({ session, tool }), - ); - return { tool, decision }; - }), + const decisionCache = createToolAccessDecisionCache(); + const decisions = await decideToolsForListing( + tools, + (tool) => policyInputForTool({ session, tool }), + decisionCache, + signal, ); const visibleTools = decisions .filter( @@ -2914,13 +3030,11 @@ export function createToolGatewayService( : tool, ); if (onDemandTargets.length > 0) { - const targetDecisions = await Promise.all( - onDemandTargets.map(async (tool) => { - const decision = await policyService.decide( - policyInputForTool({ session, tool }), - ); - return { tool, decision }; - }), + const targetDecisions = await decideToolsForListing( + onDemandTargets, + (tool) => policyInputForTool({ session, tool }), + decisionCache, + signal, ); if ( targetDecisions.some( @@ -8696,7 +8810,9 @@ export function createToolGatewayService( gatewayPublicId?: string | null; bearerToken: string; callerHeaders?: Record; + signal?: AbortSignal; }): Promise { + input.signal?.throwIfAborted(); const session = await namedGatewaySessionFromBearer({ gatewayId: input.gatewayId ?? null, gatewayPublicId: input.gatewayPublicId ?? null, @@ -8705,7 +8821,7 @@ export function createToolGatewayService( callerHeaders: input.callerHeaders, }); await assertGatewayTokenAction(session, "tools/list"); - const tools = await listToolsForContext(session); + const tools = await listToolsForContext(session, input.signal); await writeAudit({ session, companyId: session.companyId, @@ -8717,7 +8833,7 @@ export function createToolGatewayService( decision: "allow", reasonCode: "named_gateway_discovery_filtered", visibleToolCount: tools.length, - visibleTools: tools.map((tool) => tool.name), + visibleToolsHash: createHash("sha256").update(JSON.stringify(tools.map((tool) => tool.name).sort())).digest("hex"), }, }); return tools; @@ -8796,9 +8912,11 @@ export function createToolGatewayService( async listToolsForSession( sessionToken: string, + options: { signal?: AbortSignal } = {}, ): Promise { + options.signal?.throwIfAborted(); const session = await getActiveSession(sessionToken); - const tools = await listToolsForContext(session); + const tools = await listToolsForContext(session, options.signal); await writeAudit({ session, companyId: session.companyId, @@ -8810,7 +8928,7 @@ export function createToolGatewayService( decision: "allow", reasonCode: "discovery_filtered", visibleToolCount: tools.length, - visibleTools: tools.map((tool) => tool.name), + visibleToolsHash: createHash("sha256").update(JSON.stringify(tools.map((tool) => tool.name).sort())).digest("hex"), }, }); return tools; @@ -8821,16 +8939,11 @@ export function createToolGatewayService( agentId: string; }): Promise { await assertAgentInCompany(input.companyId, input.agentId); - const decisions = await Promise.all( - pluginTools().map(async (tool) => { - const decision = await policyService.decide( - policyInputForAgentTool({ - companyId: input.companyId, - agentId: input.agentId, - tool, - }), - ); - return { tool, decision }; + const decisions = await decideToolsForListing(pluginTools(), (tool) => + policyInputForAgentTool({ + companyId: input.companyId, + agentId: input.agentId, + tool, }), ); return decisions @@ -8887,33 +9000,32 @@ export function createToolGatewayService( input.companyId, input.connectionId, ); - const decisions = await Promise.all( - tools.map(async (tool) => { - const decision = await policyService.decide( - policyInputForAgentTool({ - companyId: input.companyId, - agentId: input.agentId, - tool, - }), - ); - const testDecision = - decision.decision === "require_approval" - ? "ask_first" - : decision.allowed - ? "allowed" - : "off"; - return { - toolName: tool.upstreamToolName ?? tool.name, - gatewayToolName: tool.name, - displayName: tool.displayName, - risk: tool.risk, - decision: testDecision, - reasonCode: decision.reasonCode, - matchedPolicyIds: decision.matchedPolicyIds, - effectiveProfileIds: decision.effectiveProfileIds, - }; - }), - ); + const decisions = ( + await decideToolsForListing(tools, (tool) => + policyInputForAgentTool({ + companyId: input.companyId, + agentId: input.agentId, + tool, + }), + ) + ).map(({ tool, decision }) => { + const testDecision = + decision.decision === "require_approval" + ? "ask_first" + : decision.allowed + ? "allowed" + : "off"; + return { + toolName: tool.upstreamToolName ?? tool.name, + gatewayToolName: tool.name, + displayName: tool.displayName, + risk: tool.risk, + decision: testDecision, + reasonCode: decision.reasonCode, + matchedPolicyIds: decision.matchedPolicyIds, + effectiveProfileIds: decision.effectiveProfileIds, + }; + }); const lastChange = await summarizeAccessLastChange({ companyId: input.companyId, connectionId: input.connectionId, @@ -11032,6 +11144,16 @@ export function createToolGatewayService( async cleanupExpiredSessions(input: { now?: Date } = {}) { const now = input.now ?? new Date(); + const expiredTokens = await db.select({ id: toolMcpGatewayTokens.id }) + .from(toolMcpGatewayTokens) + .where(lte(toolMcpGatewayTokens.expiresAt, now)) + .orderBy(asc(toolMcpGatewayTokens.expiresAt), asc(toolMcpGatewayTokens.id)) + .limit(500); + if (expiredTokens.length > 0) { + await db.delete(toolMcpGatewayTokens) + .where(and(inArray(toolMcpGatewayTokens.id, expiredTokens.map((token) => token.id)), lte(toolMcpGatewayTokens.expiresAt, now))) + .returning({ id: toolMcpGatewayTokens.id }); + } const rows = await db .delete(toolGatewaySessions) .where(lte(toolGatewaySessions.expiresAt, now)) diff --git a/tests/e2e/mcp-user-stories.spec.ts b/tests/e2e/mcp-user-stories.spec.ts index 93acd041bb..2585bc1a14 100644 --- a/tests/e2e/mcp-user-stories.spec.ts +++ b/tests/e2e/mcp-user-stories.spec.ts @@ -195,6 +195,10 @@ async function startMockMcp(): Promise { } async function screenshot(page: Page, storyId: string, step: string) { + if (new URL(page.url()).pathname.endsWith("/review")) { + await expect(page.getByRole("heading", { name: "Review", exact: true })).toBeVisible({ timeout: 30_000 }); + await expect(page.getByText("Loading tools…", { exact: true })).toBeHidden(); + } await page.screenshot({ path: `${SCREENSHOT_DIR}/${storyId.toLowerCase()}-${step}.png`, fullPage: true }); } @@ -318,9 +322,10 @@ test.describe.serial("MCP prod Phase 5a user-story harness", () => { expect(pending.actionRequestId).toBeTruthy(); await page.goto(`/${seed.prefix}/apps/${connectionId}/review`); + await expect(page.getByRole("heading", { name: "Waiting for your OK" })).toBeVisible({ timeout: 30_000 }); await screenshot(page, "US-2", "01-review-pending"); - await approveActionRequest(request, seed.companyId, pending.actionRequestId!); + await page.getByRole("button", { name: "Allow once", exact: true }).click(); await pollTestCall(request, connectionId, pending.actionRequestId!, "done"); expect(mock.captures.filter((capture) => capture.method === "tools/call" && capture.toolName === "sheets:update_cell")).toHaveLength(1); await expectAuditEvent(request, seed.companyId, { connectionId, agentId: scout.id, search: "sheets:update_cell" });