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 <noreply@paperclip.ing>

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.

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({

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<void>((resolve) => { started = resolve; });
+    const gate = new Promise<void>((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);

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