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" });