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 00000000000..ca1fbfaeadb --- /dev/null +++ b/server/src/__tests__/tool-gateway-listing-memory.test.ts @@ -0,0 +1,518 @@ +import { randomUUID } from "node:crypto"; +import { drizzle } from "drizzle-orm/postgres-js"; +import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import { + agents, + companies, + createDb, + heartbeatRuns, + issues, + principalPermissionGrants, + projects, + toolApplications, + toolCatalogEntries, + toolConnections, + toolPolicies, + toolProfileBindings, + toolProfileEntries, + toolProfiles, +} from "@paperclipai/db"; +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"; + +const embeddedPostgresSupport = await getEmbeddedPostgresTestSupport(); +const describeEmbeddedPostgres = embeddedPostgresSupport.supported ? describe : describe.skip; + +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. */ +function recordingDb(db: Db) { + const statements: string[] = []; + const recorded = drizzle(db.$client, { + schema: db._.fullSchema, + logger: { + logQuery: (query: string) => { + statements.push(query); + }, + }, + }) as unknown as Db; + return { db: recorded, statements }; +} + +async function createListingFixture( + db: Db, + toolCount: number, + options: { broadBindings?: boolean } = {}, +) { + 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(); + + // 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 }; +} + +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 }; + } + + 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); + const { statements } = 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"'); + } + }); + + 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]!); + + 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, + // 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/services/tool-access-policy.ts b/server/src/services/tool-access-policy.ts index e09d1b28722..f173f2c6c45 100644 --- a/server/src/services/tool-access-policy.ts +++ b/server/src/services/tool-access-policy.ts @@ -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. + */ +function runContextSnapshotString(key: "issueId" | "projectId" | "routineId") { + 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,24 @@ export function toolAccessPolicyService(db: Db) { return deleted; } - async function loadContext(input: ToolAccessDecisionInput): Promise< + async function loadCatalogEntry(companyId: string, catalogEntryId: string, cache?: ToolAccessDecisionCache) { + // 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)) { + const [entry] = await db.select(catalogEntryContextColumns).from(toolCatalogEntries).where(eq(toolCatalogEntries.id, catalogEntryId)); + return entry && entry.companyId === companyId ? entry : null; + } + // A listing decides most catalog tools of the company, so a batch reads the + // company's entries once and looks each one up in memory. + const entriesById = await cachedRead(cache, `catalog-entries:${companyId}`, async () => { + const entries = await db.select(catalogEntryContextColumns).from(toolCatalogEntries).where(eq(toolCatalogEntries.companyId, companyId)); + return new Map(entries.map((entry) => [entry.id, entry])); + }); + const entry = entriesById.get(catalogEntryId); + return entry && 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 +926,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 +936,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 +966,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 +978,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 +1008,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 +1020,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 +1038,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 +1060,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 +1120,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 +1155,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 +1301,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 +1426,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-gateway.ts b/server/src/services/tool-gateway.ts index 38ae316b313..b8192e52e62 100644 --- a/server/src/services/tool-gateway.ts +++ b/server/src/services/tool-gateway.ts @@ -127,7 +127,11 @@ import { REMOTE_URL_SECRET_CONFIG_PATH, remoteUrlCredentialMatchesPublicUrl, } from "./remote-url-credentials.js"; -import { toolAccessPolicyService } from "./tool-access-policy.js"; +import { + createToolAccessDecisionCache, + toolAccessPolicyService, + type ToolAccessDecisionCache, +} from "./tool-access-policy.js"; import { commitToolActionReview } from "./tool-action-review.js"; import { issueThreadInteractionService } from "./issue-thread-interactions.js"; import { @@ -217,6 +221,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 +466,35 @@ function asRecord(value: unknown): Record | null { return value as Record; } +async function mapWithConcurrency( + items: readonly T[], + concurrency: number, + mapper: (item: T) => Promise, +): 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 { + results[index] = await mapper(items[index]!); + } catch (error) { + // The batch fails with this error, so the other workers stop. + failed = true; + throw error; + } + } + } + + await Promise.all( + Array.from({ length: Math.min(concurrency, items.length) }, () => worker()), + ); + return results; +} + const sensitivePassthroughHeaderPattern = /(^|[-_])(auth|authorization|cookie|secret|session|token)([-_]|$)|(^|[-_])api[-_]?key([-_]|$)/i; const sensitivePassthroughHeaderNames = new Set([ @@ -1216,49 +1253,81 @@ 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"), - ), - eq(toolApplications.companyId, companyId), - inArray(toolApplications.type, ["mcp_http", "mcp_stdio", "rest_api"]), - eq(toolApplications.status, "active"), - ), - ) - .orderBy(toolConnections.name, toolCatalogEntries.name); + // 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. + const { connectionRows, catalogRows } = await db.transaction( + async (tx) => { + const connectionRows = await tx + .select({ + connection: toolConnections, + application: toolApplications, + }) + .from(toolConnections) + .innerJoin( + toolApplications, + eq(toolConnections.applicationId, toolApplications.id), + ) + .where(connectedMcpConnectionFilter(companyId)); + if (connectionRows.length === 0) + return { connectionRows, catalogRows: [] }; + 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); + 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 }) => @@ -2633,6 +2702,22 @@ 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(), + ): Promise> { + return mapWithConcurrency(tools, LISTING_DECISION_CONCURRENCY, async (tool) => ({ + tool, + decision: await policyService.decide(inputForTool(tool), { cache }), + })); + } + function policyErrorStatus(decision: ToolAccessDecision) { if (decision.decision === "rate_limited") return 429; return 403; @@ -2809,13 +2894,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( @@ -2887,13 +2967,11 @@ 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, ); const visibleTools = decisions .filter( @@ -2914,13 +2992,10 @@ 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, ); if ( targetDecisions.some( @@ -8821,16 +8896,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 +8957,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,