From f17bfd73ba513a330ce2aecd3e3927d87230e0ef Mon Sep 17 00:00:00 2001 From: Dotta Date: Thu, 20 Aug 2026 19:51:42 +0000 Subject: [PATCH] fix: align gateway install filtering with policy evaluation Use the canonical tool gateway policy evaluator to identify every connection a managed gateway can expose. Cover allow-default exclusions and namespaced tool selectors so install checks cannot diverge from runtime gateway behavior. Co-Authored-By: Paperclip --- .../heartbeat-runtime-mcp-servers.test.ts | 75 +++++++++++++++---- server/src/services/heartbeat.ts | 71 +++++------------- server/src/services/tool-gateway.ts | 35 +++++++++ 3 files changed, 114 insertions(+), 67 deletions(-) diff --git a/server/src/__tests__/heartbeat-runtime-mcp-servers.test.ts b/server/src/__tests__/heartbeat-runtime-mcp-servers.test.ts index d7b8958f791..49e3a0bde3f 100644 --- a/server/src/__tests__/heartbeat-runtime-mcp-servers.test.ts +++ b/server/src/__tests__/heartbeat-runtime-mcp-servers.test.ts @@ -9,6 +9,7 @@ import { heartbeatRuns, toolAccessAuditEvents, toolApplications, + toolCatalogEntries, toolConnectionInstalls, toolConnections, toolMcpGateways, @@ -241,7 +242,7 @@ describeEmbeddedPostgres("heartbeat runtime MCP servers", () => { }); }); - it("injects only managed gateways whose profile connections are installed for the agent", async () => { + it("uses canonical gateway profile semantics before injecting installed gateways", async () => { const [company] = await db.insert(companies).values({ name: `Managed gateway installs ${randomUUID()}`, issuePrefix: `MG${randomUUID().slice(0, 5).toUpperCase()}`, @@ -269,6 +270,8 @@ describeEmbeddedPostgres("heartbeat runtime MCP servers", () => { transport: "mcp_remote", status: "active", enabled: true, + healthStatus: "healthy", + config: { url: "https://installed-gateway.example.test/mcp" }, }, { companyId: company!.id, @@ -278,21 +281,53 @@ describeEmbeddedPostgres("heartbeat runtime MCP servers", () => { transport: "mcp_remote", status: "active", enabled: true, + healthStatus: "healthy", + config: { url: "https://uninstalled-gateway.example.test/mcp" }, }, ]).returning(); - const profiles = await db.insert(toolProfiles).values(connections.map((connection) => ({ + await db.insert(toolCatalogEntries).values(connections.map((connection, index) => ({ companyId: company!.id, - profileKey: `gateway:${connection.id}`, - name: connection.name, - defaultAction: "deny" as const, - }))).returning(); - await db.insert(toolProfileEntries).values(profiles.map((profile, index) => ({ - companyId: company!.id, - profileId: profile.id, - selectorType: "connection" as const, - effect: "include" as const, - connectionId: connections[index]!.id, + applicationId: application!.id, + connectionId: connection.id, + entryKind: "tool" as const, + name: `gateway-tool-${index}`, + toolName: index === 0 ? "installed_tool" : "uninstalled_tool", + versionHash: randomUUID(), }))); + const profiles = await db.insert(toolProfiles).values([ + { + companyId: company!.id, + profileKey: `gateway:allow-excluding-uninstalled:${randomUUID()}`, + name: "Allow installed tools", + defaultAction: "allow" as const, + }, + { + companyId: company!.id, + profileKey: `gateway:namespaced-uninstalled:${randomUUID()}`, + name: "Allow one uninstalled tool", + defaultAction: "deny" as const, + }, + ]).returning(); + const uninstalledGatewayToolName = [ + `mcp.${application!.applicationKey}`, + `${connections[1]!.id.replaceAll("-", "").slice(0, 8)}:uninstalled-tool`, + ].join("-"); + await db.insert(toolProfileEntries).values([ + { + companyId: company!.id, + profileId: profiles[0]!.id, + selectorType: "connection", + effect: "exclude", + connectionId: connections[1]!.id, + }, + { + companyId: company!.id, + profileId: profiles[1]!.id, + selectorType: "tool_name", + effect: "include", + toolName: uninstalledGatewayToolName, + }, + ]); const gateways = await db.insert(toolMcpGateways).values(profiles.map((profile, index) => ({ companyId: company!.id, name: `${connections[index]!.name} gateway`, @@ -300,6 +335,12 @@ describeEmbeddedPostgres("heartbeat runtime MCP servers", () => { profileId: profile.id, status: "active" as const, }))).returning(); + await db.insert(toolProfileBindings).values(gateways.map((gateway, index) => ({ + companyId: company!.id, + profileId: profiles[index]!.id, + targetType: "gateway" as const, + targetId: gateway.id, + }))); await db.insert(toolConnectionInstalls).values({ companyId: company!.id, connectionId: connections[0]!.id, @@ -307,10 +348,18 @@ describeEmbeddedPostgres("heartbeat runtime MCP servers", () => { targetId: agent!.id, }); + const runId = randomUUID(); + await db.insert(heartbeatRuns).values({ + id: runId, + companyId: company!.id, + agentId: agent!.id, + status: "running", + contextSnapshot: {}, + }); const config = await createManagedMcpRunConfig({ db, agent: agent!, - runId: randomUUID(), + runId, config: {}, projectId: null, issueId: null, diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index e0c0b044d56..d8faf6c30d4 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -65,10 +65,8 @@ import { routines, toolMcpGateways, toolMcpGatewayTokens, - toolCatalogEntries, toolConnectionInstalls, toolConnections, - toolProfileEntries, toolProfiles, workspaceOperations, } from "@paperclipai/db"; @@ -3468,62 +3466,23 @@ function gatewayAppliesToRun(input: { async function gatewayConnectionIds(input: { db: Db; companyId: string; + agentId: string; + runId: string; + issueId: string | null; + projectId: string | null; gateway: typeof toolMcpGateways.$inferSelect; }): Promise> { const managedRuntimeConnectionId = readNonEmptyString(input.gateway.metadata?.managedRuntimeConnectionId); if (managedRuntimeConnectionId) return new Set([managedRuntimeConnectionId]); - - const [profile, entries, catalog, connections] = await Promise.all([ - input.db - .select({ defaultAction: toolProfiles.defaultAction }) - .from(toolProfiles) - .where(and(eq(toolProfiles.companyId, input.companyId), eq(toolProfiles.id, input.gateway.profileId))) - .then((rows) => rows[0] ?? null), - input.db - .select() - .from(toolProfileEntries) - .where(and( - eq(toolProfileEntries.companyId, input.companyId), - eq(toolProfileEntries.profileId, input.gateway.profileId), - )), - input.db - .select({ - id: toolCatalogEntries.id, - connectionId: toolCatalogEntries.connectionId, - applicationId: toolCatalogEntries.applicationId, - toolName: toolCatalogEntries.toolName, - riskLevel: toolCatalogEntries.riskLevel, - }) - .from(toolCatalogEntries) - .where(and(eq(toolCatalogEntries.companyId, input.companyId), eq(toolCatalogEntries.status, "active"))), - input.db - .select({ id: toolConnections.id, applicationId: toolConnections.applicationId }) - .from(toolConnections) - .where(eq(toolConnections.companyId, input.companyId)), - ]); - if (!profile) return new Set(); - if (profile.defaultAction === "allow") return new Set(catalog.map((entry) => entry.connectionId)); - - const connectionIds = new Set(); - for (const entry of entries) { - if (entry.effect !== "include") continue; - if (entry.connectionId) connectionIds.add(entry.connectionId); - if (entry.applicationId) { - for (const connection of connections) { - if (connection.applicationId === entry.applicationId) connectionIds.add(connection.id); - } - } - for (const catalogEntry of catalog) { - if ( - (entry.catalogEntryId && entry.catalogEntryId === catalogEntry.id) - || (entry.toolName && entry.toolName === catalogEntry.toolName) - || (entry.riskLevel && entry.riskLevel === catalogEntry.riskLevel) - ) { - connectionIds.add(catalogEntry.connectionId); - } - } - } - return connectionIds; + const service = createToolGatewayService(input.db); + return new Set(await service.listAccessibleConnectionIdsForGatewayContext({ + companyId: input.companyId, + gatewayId: input.gateway.id, + agentId: input.agentId, + runId: input.runId, + issueId: input.issueId, + projectId: input.projectId, + })); } export async function createManagedMcpRunConfig(input: { @@ -3567,6 +3526,10 @@ export async function createManagedMcpRunConfig(input: { connectionIds: await gatewayConnectionIds({ db: input.db, companyId: input.agent.companyId, + agentId: input.agent.id, + runId: input.runId, + issueId: input.issueId, + projectId: input.projectId, gateway, }), })))) diff --git a/server/src/services/tool-gateway.ts b/server/src/services/tool-gateway.ts index 5c86958d3a4..bbc55eb4945 100644 --- a/server/src/services/tool-gateway.ts +++ b/server/src/services/tool-gateway.ts @@ -4844,6 +4844,41 @@ export function createToolGatewayService( return tools; }, + async listAccessibleConnectionIdsForGatewayContext(input: { + companyId: string; + gatewayId: string; + agentId: string; + runId: string; + issueId?: string | null; + projectId?: string | null; + }): Promise { + await assertAgentInCompany(input.companyId, input.agentId); + const now = new Date(); + const session: ToolGatewaySession = { + id: `gateway:${input.gatewayId}`, + token: "", + companyId: input.companyId, + agentId: input.agentId, + runId: input.runId, + issueId: input.issueId ?? null, + projectId: input.projectId ?? null, + gatewayId: input.gatewayId, + actorType: "agent", + actorId: input.agentId, + createdAt: now, + expiresAt: new Date(now.getTime() + DEFAULT_SESSION_TTL_MS), + }; + const connectedTools = await connectedMcpToolsForCompany(input.companyId); + const decisions = await Promise.all(connectedTools.map(async (tool) => ({ + tool, + decision: await policyService.decide(policyInputForTool({ session, tool })), + }))); + return [...new Set(decisions + .filter(({ decision }) => decision.allowed || decision.decision === "require_approval") + .map(({ tool }) => tool.connectionId) + .filter((connectionId): connectionId is string => Boolean(connectionId)))]; + }, + async createSession(input: { companyId: string; agentId: string;