diff --git a/server/src/__tests__/heartbeat-runtime-mcp-servers.test.ts b/server/src/__tests__/heartbeat-runtime-mcp-servers.test.ts index 49e3a0bde3f..a883315d6fa 100644 --- a/server/src/__tests__/heartbeat-runtime-mcp-servers.test.ts +++ b/server/src/__tests__/heartbeat-runtime-mcp-servers.test.ts @@ -23,6 +23,7 @@ import { startEmbeddedPostgresTestDatabase, } from "./helpers/embedded-postgres.js"; import { buildPaperclipRuntimeMcpServers, createManagedMcpRunConfig } from "../services/heartbeat.js"; +import { createToolGatewayService } from "../services/tool-gateway.js"; const embeddedPostgresSupport = await getEmbeddedPostgresTestSupport(); const describeEmbeddedPostgres = embeddedPostgresSupport.supported ? describe : describe.skip; @@ -307,6 +308,18 @@ describeEmbeddedPostgres("heartbeat runtime MCP servers", () => { name: "Allow one uninstalled tool", defaultAction: "deny" as const, }, + { + companyId: company!.id, + profileKey: `gateway:mixed-allow:${randomUUID()}`, + name: "Allow both connections by default", + defaultAction: "allow" as const, + }, + { + companyId: company!.id, + profileKey: `gateway:mixed-deny:${randomUUID()}`, + name: "Explicitly allow both connections", + defaultAction: "deny" as const, + }, ]).returning(); const uninstalledGatewayToolName = [ `mcp.${application!.applicationKey}`, @@ -327,10 +340,17 @@ describeEmbeddedPostgres("heartbeat runtime MCP servers", () => { effect: "include", toolName: uninstalledGatewayToolName, }, + ...connections.map((connection) => ({ + companyId: company!.id, + profileId: profiles[3]!.id, + selectorType: "connection" as const, + effect: "include" as const, + connectionId: connection.id, + })), ]); const gateways = await db.insert(toolMcpGateways).values(profiles.map((profile, index) => ({ companyId: company!.id, - name: `${connections[index]!.name} gateway`, + name: `${profile.name} gateway`, slug: `gateway-${index}-${randomUUID().slice(0, 8)}`, profileId: profile.id, status: "active" as const, @@ -365,8 +385,23 @@ describeEmbeddedPostgres("heartbeat runtime MCP servers", () => { issueId: null, }); - expect(config?.gateways).toHaveLength(1); - expect(config?.gateways[0]).toMatchObject({ id: gateways[0]!.id, name: gateways[0]!.name }); + expect(config?.gateways).toHaveLength(3); + expect(config?.gateways.some((gateway) => gateway.id === gateways[0]!.id)).toBe(true); expect(config?.gateways.some((gateway) => gateway.id === gateways[1]!.id)).toBe(false); + expect(config?.gateways.some((gateway) => gateway.id === gateways[2]!.id)).toBe(true); + expect(config?.gateways.some((gateway) => gateway.id === gateways[3]!.id)).toBe(true); + + const mixedGateway = config?.gateways.find((gateway) => gateway.id === gateways[2]!.id); + const visibleTools = await createToolGatewayService(db).listToolsForNamedGateway({ + gatewayId: mixedGateway!.id, + bearerToken: mixedGateway!.bearerToken, + }); + expect(visibleTools.some((tool) => tool.connectionId === connections[0]!.id)).toBe(true); + expect(visibleTools.some((tool) => tool.connectionId === connections[1]!.id)).toBe(false); + await expect(createToolGatewayService(db).executeTool({ + sessionToken: mixedGateway!.bearerToken, + gatewayId: mixedGateway!.id, + tool: uninstalledGatewayToolName, + })).rejects.toMatchObject({ status: 403, reasonCode: "installation_required" }); }); }); diff --git a/server/src/services/heartbeat.ts b/server/src/services/heartbeat.ts index d8faf6c30d4..c07a485fb1c 100644 --- a/server/src/services/heartbeat.ts +++ b/server/src/services/heartbeat.ts @@ -3535,7 +3535,7 @@ export async function createManagedMcpRunConfig(input: { })))) .filter(({ connectionIds }) => connectionIds.size > 0 - && [...connectionIds].every((connectionId) => installedConnectionIds.has(connectionId))) + && [...connectionIds].some((connectionId) => installedConnectionIds.has(connectionId))) .map(({ gateway }) => gateway); if (gateways.length === 0) return null; diff --git a/server/src/services/tool-gateway.ts b/server/src/services/tool-gateway.ts index bbc55eb4945..de037f573f3 100644 --- a/server/src/services/tool-gateway.ts +++ b/server/src/services/tool-gateway.ts @@ -17,6 +17,7 @@ import { toolApplications, toolCallEvents, toolCatalogEntries, + toolConnectionInstalls, toolConnections, toolGatewayRateLimitCounters, toolGatewaySessions, @@ -172,6 +173,7 @@ export interface ToolGatewaySession { gatewayName?: string | null; gatewayTokenId?: string | null; gatewayTokenAllowedActions?: ToolMcpGatewayTokenAction[]; + enforceConnectionInstalls?: boolean; actorType?: "agent" | "user" | "system" | "plugin"; actorId?: string | null; createdAt: Date; @@ -1893,7 +1895,10 @@ export function createToolGatewayService( } async function searchableOnDemandTools(session: ToolGatewaySession): Promise { - const tools = (await connectedMcpToolsForCompany(session.companyId)).filter(isOnDemandRemoteTool); + const tools = (await installedToolsForSession( + session, + await connectedMcpToolsForCompany(session.companyId), + )).filter(isOnDemandRemoteTool); const decisions = await Promise.all(tools.map(async (tool) => ({ tool, decision: await policyService.decide(policyInputForTool({ session, tool })), @@ -1942,7 +1947,10 @@ export function createToolGatewayService( if (session.agentId) { await assertAgentInCompany(session.companyId, session.agentId); } - const allConnectedTools = await connectedMcpToolsForCompany(session.companyId); + const allConnectedTools = await installedToolsForSession( + session, + await connectedMcpToolsForCompany(session.companyId), + ); const onDemandTargets = allConnectedTools.filter(isOnDemandRemoteTool); const tools = [...allTools(), ...allConnectedTools.filter((tool) => !isOnDemandRemoteTool(tool))].filter( (tool) => session.agentId || (tool.providerType !== "paperclip_self" && tool.providerType !== "paperclip_plugin"), @@ -1971,6 +1979,39 @@ export function createToolGatewayService( return visibleTools; } + async function installedConnectionIdsForSession(session: ToolGatewaySession): Promise | null> { + if (!session.enforceConnectionInstalls || !session.agentId) return null; + const installs = await db + .select({ connectionId: toolConnectionInstalls.connectionId }) + .from(toolConnectionInstalls) + .where(and( + eq(toolConnectionInstalls.companyId, session.companyId), + sql`((${toolConnectionInstalls.targetType} = 'company' and ${toolConnectionInstalls.targetId} = ${session.companyId}) or (${toolConnectionInstalls.targetType} = 'agent' and ${toolConnectionInstalls.targetId} = ${session.agentId}))`, + )); + return new Set(installs.map((install) => install.connectionId)); + } + + async function installedToolsForSession( + session: ToolGatewaySession, + tools: ToolGatewayDescriptor[], + ): Promise { + const installedConnectionIds = await installedConnectionIdsForSession(session); + if (!installedConnectionIds) return tools; + return tools.filter((tool) => !tool.connectionId || installedConnectionIds.has(tool.connectionId)); + } + + async function assertToolInstalledForSession(session: ToolGatewaySession, tool: ToolGatewayDescriptor) { + if (!tool.connectionId) return; + const installedConnectionIds = await installedConnectionIdsForSession(session); + if (!installedConnectionIds || installedConnectionIds.has(tool.connectionId)) return; + throw new ToolGatewayHttpError( + 403, + "This connection is not installed for the current agent", + "installation_required", + { connectionId: tool.connectionId }, + ); + } + async function executeBuiltinTool(session: ToolGatewaySession, tool: ToolGatewayDescriptor, parameters: unknown) { const params = asRecord(parameters) ?? {}; @@ -3723,6 +3764,7 @@ export function createToolGatewayService( gatewayName: row.gateway.name, gatewayTokenId: row.token.id || tokenId, gatewayTokenAllowedActions: normalizeGatewayTokenActions(row.token.allowedActions), + enforceConnectionInstalls: row.token.subjectType === "heartbeat_run", actorType: runId ? "agent" : "system", actorId: runId ? agentId : row.token.id, createdAt: row.token.createdAt, @@ -5535,6 +5577,8 @@ export function createToolGatewayService( requestedParameters = targetParameters; } + await assertToolInstalledForSession(session, tool); + const argumentValidation = validateToolContent({ value: requestedParameters, direction: "arguments",