From 5c0793090cec561a00c2bbc7ea5b19466db91018 Mon Sep 17 00:00:00 2001 From: Dotta Date: Thu, 1 Oct 2026 17:33:39 -0500 Subject: [PATCH] fix(tool-gateway): keep admitted request audits safe during token cleanup Resolve the nullable token reference with a row lock inside the audit INSERT. Deleted tokens retain their original ID in audit details. Red-green listing and HTTP provider-call tests reproduce cleanup during admitted work and verify that later requests still fail authentication. Co-Authored-By: Paperclip --- doc/mcp-discovery-performance.md | 4 ++ .../tool-gateway-discovery-http.test.ts | 33 ++++++++++++++++- .../tool-gateway-listing-memory.test.ts | 37 ++++++++++++++++++- server/src/services/tool-gateway.ts | 18 +++++---- 4 files changed, 82 insertions(+), 10 deletions(-) diff --git a/doc/mcp-discovery-performance.md b/doc/mcp-discovery-performance.md index bd36964b94..219079cf3e 100644 --- a/doc/mcp-discovery-performance.md +++ b/doc/mcp-discovery-performance.md @@ -15,6 +15,10 @@ names. They do not retain the full name list. Named gateway tokens expire normal 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. diff --git a/server/src/__tests__/tool-gateway-discovery-http.test.ts b/server/src/__tests__/tool-gateway-discovery-http.test.ts index 14429cdb94..62e26dc768 100644 --- a/server/src/__tests__/tool-gateway-discovery-http.test.ts +++ b/server/src/__tests__/tool-gateway-discovery-http.test.ts @@ -1,8 +1,9 @@ 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, toolPolicies, startEmbeddedPostgresTestDatabase, getEmbeddedPostgresTestSupport } from "@paperclipai/db"; +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"; @@ -18,6 +19,36 @@ suite("MCP discovery over HTTP", () => { }); 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({ diff --git a/server/src/__tests__/tool-gateway-listing-memory.test.ts b/server/src/__tests__/tool-gateway-listing-memory.test.ts index bc98242e1d..4708aaf0c4 100644 --- a/server/src/__tests__/tool-gateway-listing-memory.test.ts +++ b/server/src/__tests__/tool-gateway-listing-memory.test.ts @@ -1,5 +1,5 @@ import { randomUUID } from "node:crypto"; -import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import { afterAll, beforeAll, describe, expect, it, vi } from "vitest"; import { activityLog, agents, @@ -119,6 +119,41 @@ describeEmbeddedPostgres("tool gateway listing memory", () => { 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); diff --git a/server/src/services/tool-gateway.ts b/server/src/services/tool-gateway.ts index f3e7c30093..c601c6bc8b 100644 --- a/server/src/services/tool-gateway.ts +++ b/server/src/services/tool-gateway.ts @@ -1728,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: @@ -1736,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