diff --git a/doc/MCP-RUNTIME-OPERATIONS.md b/doc/MCP-RUNTIME-OPERATIONS.md index 011c5c6fbbc..01cd4c14709 100644 --- a/doc/MCP-RUNTIME-OPERATIONS.md +++ b/doc/MCP-RUNTIME-OPERATIONS.md @@ -175,6 +175,19 @@ curl -fsS -X PATCH \ Recovery is complete when stuck-slot alerts clear, timeout/error rates return below threshold, the connection is healthy or intentionally disabled, and audit events show no new restart suppression or capacity deferrals. +### Catalog publication + +Discovery keeps the last descriptor when pages repeat a tool name. Catalog rows +are written in transactions of at most 500 tools. Refresh preserves existing row +IDs, first-seen timestamps, operator-disabled actions, and unchanged pending +reviews. Publication locks the application before the connection, as setup +finalization and removal do. Provider discovery runs before these locks. + +Each completed chunk remains committed if a later chunk fails. Default-profile +synchronization runs in a separate transaction after catalog publication. A +failure there rolls back the profile changes without removing discovered tools. +Retry the catalog refresh after correcting the failure. + ## Verification Coverage Automated coverage includes: @@ -182,3 +195,4 @@ Automated coverage includes: - A synthetic degraded runtime-health scenario in `server/src/__tests__/tool-access-service.test.ts` that creates a stale running slot, degraded connection, timeout event, capacity deferral, and restart suppression. - A durable audit-write failure scenario in `server/src/__tests__/tool-access-service.test.ts` that verifies `mcp_runtime_audit_write_failures` fires from the counter path. - A gateway runtime recovery scenario in `server/src/__tests__/tool-gateway.test.ts` that recovers a stuck local stdio slot before reuse. +- PostgreSQL catalog publication tests in `server/src/__tests__/tool-catalog-publication.test.ts` cover 250 and 2,500 tools, overlapping pages, preserved choices, partial failure, and publication lock contention. diff --git a/maps/feature-map/connection-setup.md b/maps/feature-map/connection-setup.md index ffe815d8fd7..b98314b0956 100644 --- a/maps/feature-map/connection-setup.md +++ b/maps/feature-map/connection-setup.md @@ -105,6 +105,11 @@ Automated: AppsConnect's generic MCP cases cover probe-driven auth, private or unreachable endpoints, preregistered clients, secret-backed headers, and unsafe authorization URLs. These are component/API mocks, not remote provider proof. +[Catalog publication](../../server/src/__tests__/tool-catalog-publication.test.ts) +uses PostgreSQL and a controlled MCP response to check bounded writes, repeated +page entries, preserved action choices, rollback, and concurrent publication. +It does not verify a live provider's discovery contract. + Manual: connect a known test MCP URL through the link flow, verify its saved endpoint and a read-only action, then repeat with a bad URL and wrong credential. The error should explain the next action and preserve non-secret setup context. diff --git a/server/src/__tests__/tool-access-service.test.ts b/server/src/__tests__/tool-access-service.test.ts index 8b39407897a..f4cfcaf4cb4 100644 --- a/server/src/__tests__/tool-access-service.test.ts +++ b/server/src/__tests__/tool-access-service.test.ts @@ -8632,15 +8632,17 @@ describeEmbeddedPostgres("tool access service", () => { "state", )!; const originalTransaction = db.transaction.bind(db); - let transactionCount = 0; const transactionSpy = vi .spyOn(db, "transaction") - .mockImplementation((async (operation, config) => { - transactionCount += 1; - if (transactionCount === 2) - throw new Error("recommended defaults failed"); - return originalTransaction(operation, config); - }) as typeof db.transaction); + .mockImplementation((async (operation, config) => originalTransaction(async tx => { + const insert = tx.insert.bind(tx); + const insertSpy = vi.spyOn(tx, "insert").mockImplementation(((table: Parameters[0]) => { + if (table === toolProfiles) throw new Error("recommended defaults failed"); + return insert(table); + }) as typeof tx.insert); + try { return await operation(tx); } + finally { insertSpy.mockRestore(); } + }, config)) as typeof db.transaction); await expect( service.completePaperclipCloudConnectorCallback({ @@ -10153,6 +10155,91 @@ describeEmbeddedPostgres("tool access service", () => { }, ); + it.each((["personal", "organization", "managed"] as const).flatMap(mode => + (["app finishing", "new tool discovery"] as const).map(change => ({ mode, change }))))("orders $mode OAuth persistence with concurrent $change", async ({ mode, change }) => { + vi.stubEnv("PAPERCLIP_TOOL_OAUTH_SLACK_CLIENT_ID", "slack-client-id"); + vi.stubEnv("PAPERCLIP_TOOL_OAUTH_SLACK_CLIENT_SECRET", "slack-client-secret"); + const company = await createCompany(db); + const userId = `catalog-oauth-${randomUUID()}`; + await grantBoardUser(db, company.id, userId, [], "owner"); + const actor = { actorType: "user" as const, actorId: userId }; + const connector = mode === "managed" ? fakeGmailConnector(company.id, userId) : undefined; + const service = createTestToolAccessService(db, { paperclipCloudConnector: connector }); + const definition = getConnectableAppDefinition(mode === "managed" ? "gmail" : "slack")!; + const availability = definition.ownershipAvailability; + if (mode === "managed") definition.ownershipAvailability = { ...availability, platform_shared: true }; + let discoverNewTool = false; + vi.spyOn(globalThis, "fetch").mockImplementation(async (url, init) => { + if (String(url) === "https://slack.com/api/oauth.v2.access") return mcpHttpResponse(oauthToken("fixture-access-token", "fixture-refresh-token")); + const body = JSON.parse(String(init?.body)); + if (body.method === "notifications/initialized") return new Response(null, { status: 202 }); + return mcpHttpResponse({ jsonrpc: "2.0", id: body.id, result: body.method === "initialize" ? { protocolVersion: "2025-06-18" } + : { tools: [{ name: "search_messages", annotations: { readOnlyHint: true } }, ...(discoverNewTool ? [{ name: "new_search", annotations: { readOnlyHint: true } }] : [])] } }); + }); + let release!: () => void; + let transactionSpy: ReturnType | undefined; + let completed: Promise[]> | undefined; + let blocker: Promise | undefined; + try { + const connected = await service.connectGalleryApp(company.id, { galleryKey: mode === "managed" ? "gmail" : "slack", ...(mode === "managed" ? { connectionMethodKey: "paperclip-draft" } : {}), name: "Concurrent catalog OAuth", grantKind: mode === "organization" ? "organization" : "user" }, actor); + const redirectUri = mode === "managed" ? "https://paperclip.example/api/tools/oauth/cloud-connector/callback" : "https://paperclip.example/api/tools/oauth/callback"; + const start = () => service.startOAuth(company.id, connected.connectionId, { redirectUri, actor, ...(mode === "organization" ? {} : { subjectUserId: userId }) }); + const complete = (state: string, claim: string) => mode === "managed" + ? service.completePaperclipCloudConnectorCallback({ state, claimId: claim, actor }) + : service.completeOAuthCallback({ state, code: claim, redirectUri, actor }); + const initial = await start(); + await complete(new URL(initial.authorizationUrl).searchParams.get("state")!, "initial"); + const reconnect = await start(); + const state = new URL(reconnect.authorizationUrl).searchParams.get("state")!; + const [stateRow] = await db.select().from(toolOauthStates).where(eq(toolOauthStates.state, state)); + expect(stateRow.subjectUserId).toBe(mode === "organization" ? null : userId); + discoverNewTool = change === "new tool discovery"; + await db.execute(sql`CREATE OR REPLACE FUNCTION catalog_oauth_fixture_pause() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN + PERFORM pg_advisory_xact_lock(hashtextextended(NEW.id::text, 0)); RETURN NEW; END $$`); + await db.execute(sql.raw(`CREATE TRIGGER catalog_oauth_fixture_pause AFTER UPDATE ON tool_connections FOR EACH ROW WHEN (NEW.id = '${connected.connectionId}'::uuid) EXECUTE FUNCTION catalog_oauth_fixture_pause()`)); + let locked!: () => void; + const held = new Promise(resolve => { locked = resolve; }); + const resume = new Promise(resolve => { release = resolve; }); + blocker = db.transaction(async tx => { + await tx.execute(sql`SELECT pg_advisory_xact_lock(hashtextextended(${connected.connectionId}, 0))`); + locked(); await resume; + }); + await held; + const pids: number[] = []; + const transaction = db.transaction.bind(db); + transactionSpy = vi.spyOn(db, "transaction").mockImplementation((async (operation, config) => transaction(async tx => { + await tx.execute(sql`SET LOCAL lock_timeout = '8s'`); + const [backend] = await tx.execute(sql`SELECT pg_backend_pid() AS pid`); + pids.push(Number(backend.pid)); + return operation(tx); + }, config)) as typeof db.transaction); + const waitForLock = (index: number) => vi.waitFor(async () => { + expect(pids.length).toBeGreaterThan(index); + const [activity] = await db.execute(sql`SELECT wait_event_type FROM pg_stat_activity WHERE pid = ${pids[index]}`); + expect(activity.wait_event_type).toBe("Lock"); + }, { timeout: 5000 }); + const callback = complete(state, "reconnect"); + const callbackResult = Promise.allSettled([callback]); + await waitForLock(0); + const publication = change === "new tool discovery" ? service.refreshCatalog(connected.connectionId, actor, { credentialHeaders: { Authorization: "Bearer fixture-catalog-token" } }) + : service.finishGalleryAppConnection(company.id, connected.connectionId, { enabledCatalogEntryIds: [], askFirstCatalogEntryIds: [], access: "all_agents" }, actor); + completed = Promise.allSettled([callback, publication]); + await waitForLock(1); + release(); await blocker; + const results = await completed; + for (const result of results) if (result.status === "rejected") throw result.reason; + expect((await callbackResult)[0].status).toBe("fulfilled"); + expect(await service.getConnection(connected.connectionId, company.id)).toMatchObject({ status: "active" }); + if (change === "new tool discovery") expect(await db.select().from(toolCatalogEntries).where(eq(toolCatalogEntries.connectionId, connected.connectionId))).toEqual(expect.arrayContaining([expect.objectContaining({ toolName: "new_search" })])); + else expect(await db.select().from(toolProfiles).where(and(eq(toolProfiles.companyId, company.id), eq(toolProfiles.profileKey, `app:${connected.connectionId}`)))).toHaveLength(1); + } finally { + release?.(); await blocker; await completed; + transactionSpy?.mockRestore(); + await db.execute(sql`DROP TRIGGER IF EXISTS catalog_oauth_fixture_pause ON tool_connections; DROP FUNCTION IF EXISTS catalog_oauth_fixture_pause()`); + definition.ownershipAvailability = availability; + } + }, 20000); + it("promotes a personal OAuth identity only after Everyone in the company is chosen", async () => { vi.stubEnv("PAPERCLIP_TOOL_OAUTH_SLACK_CLIENT_ID", "slack-client-id"); vi.stubEnv( diff --git a/server/src/__tests__/tool-catalog-publication.test.ts b/server/src/__tests__/tool-catalog-publication.test.ts new file mode 100644 index 00000000000..f3b8b79fcb1 --- /dev/null +++ b/server/src/__tests__/tool-catalog-publication.test.ts @@ -0,0 +1,205 @@ +import { randomUUID } from "node:crypto"; +import { and, eq, sql } from "drizzle-orm"; +import { agents, companies, createDb, toolApplications, toolConnections, toolCatalogEntries, toolProfiles, toolProfileBindings, toolProfileEntries, type Db } from "@paperclipai/db"; +import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; +import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "./helpers/embedded-postgres.js"; +import { toolAccessService } from "../services/tool-access.js"; +import { toolAccessPolicyService } from "../services/tool-access-policy.js"; + +const support = await getEmbeddedPostgresTestSupport(); +const describeDatabase = support.supported ? describe : describe.skip; + +describeDatabase("catalog publication", () => { + let database: Awaited>; + let db: Db; + let companyId: string; + let applicationId: string; + let connectionId: string; + beforeAll(async () => { database = await startEmbeddedPostgresTestDatabase("paperclip-catalog-publication-"); db = createDb(database.connectionString); }, 90_000); + afterAll(async () => { await database?.cleanup(); }); + afterEach(() => vi.restoreAllMocks()); + beforeEach(async () => { + companyId = randomUUID(); + await db.insert(companies).values({ id: companyId, name: "Catalog fixture", issuePrefix: `C${companyId.slice(0, 6)}` }); + const [application] = await db.insert(toolApplications).values({ companyId, name: "Fixture MCP", type: "mcp_http" }).returning(); applicationId = application.id; + const [connection] = await db.insert(toolConnections).values({ companyId, applicationId, name: "Fixture connection", uid: randomUUID(), transport: "mcp_remote", authKind: "none", enabled: true, status: "active", config: { url: "https://mcp.example.test/mcp" } }).returning(); connectionId = connection.id; + }); + + const serviceWithTools = (tools: unknown[]) => toolAccessService(db, { + remoteHttpEndpointLookup: async () => [{ address: "8.8.8.8", family: 4 }], + remoteHttpRequest: async (_url, init) => { + const body = JSON.parse(String(init?.body)); + if (body.method === "notifications/initialized") return new Response(null, { status: 202 }); + return new Response(JSON.stringify({ jsonrpc: "2.0", id: body.id, result: body.method === "initialize" ? { protocolVersion: "2025-06-18" } : { tools } }), { headers: { "content-type": "application/json" } }); + }, + }); + it.each([250, 2500])("publishes a real %i-tool catalog in bounded chunks and preserves row identities", async count => { + const tools = Array.from({ length: count }, (_, index) => ({ name: `search_${index}`, description: "Search fixture records", inputSchema: { type: "object", properties: {} }, annotations: { readOnlyHint: true } })); + await db.update(toolConnections).set({ config: { url: "https://mcp.example.test/mcp" } }).where(eq(toolConnections.id, connectionId)); + const service = toolAccessService(db, { + remoteHttpEndpointLookup: async () => [{ address: "8.8.8.8", family: 4 }], + remoteHttpRequest: async (_url, init) => { + const body = JSON.parse(String(init?.body)); + if (body.method === "notifications/initialized") return new Response(null, { status: 202 }); + return new Response(JSON.stringify({ jsonrpc: "2.0", id: body.id, result: body.method === "initialize" ? { protocolVersion: "2025-06-18" } : { tools } }), { headers: { "content-type": "application/json" } }); + }, + }); + let catalogWrites = 0; + const transaction = db.transaction.bind(db); + const transactionSpy = vi.spyOn(db, "transaction").mockImplementation((async (operation, config) => transaction(async tx => { + const insert = tx.insert.bind(tx); + const insertSpy = vi.spyOn(tx, "insert").mockImplementation(((table: Parameters[0]) => { + if (table === toolCatalogEntries) catalogWrites += 1; + return insert(table); + }) as typeof tx.insert); + try { return await operation(tx); } + finally { insertSpy.mockRestore(); } + }, config)) as typeof db.transaction); + expect((await service.refreshCatalog(connectionId)).catalog).toHaveLength(count); + expect(catalogWrites).toBe(Math.ceil(count / 500)); + const first = await db.select().from(toolCatalogEntries).where(eq(toolCatalogEntries.connectionId, connectionId)).orderBy(toolCatalogEntries.toolName); + expect((await service.refreshCatalog(connectionId)).catalog).toHaveLength(count); + expect(catalogWrites).toBe(2 * Math.ceil(count / 500)); + transactionSpy.mockRestore(); + const second = await db.select().from(toolCatalogEntries).where(eq(toolCatalogEntries.connectionId, connectionId)).orderBy(toolCatalogEntries.toolName); + expect(second.map(row => [row.id, row.firstSeenAt, row.status])).toEqual(first.map(row => [row.id, row.firstSeenAt, row.status])); + }); + it("refreshes known tools from overlapping MCP pages using the last descriptor", async () => { + await db.update(toolConnections).set({ config: { url: "https://mcp.example.test/mcp" } }).where(eq(toolConnections.id, connectionId)); + let overlap = false; + const descriptor = (description: string) => ({ name: "search", description, inputSchema: { type: "object" }, annotations: { readOnlyHint: true } }); + const service = toolAccessService(db, { + remoteHttpEndpointLookup: async () => [{ address: "8.8.8.8", family: 4 }], + remoteHttpRequest: async (_url, init) => { + const body = JSON.parse(String(init?.body)); + if (body.method === "notifications/initialized") return new Response(null, { status: 202 }); + const result = body.method === "initialize" ? { protocolVersion: "2025-06-18" } + : body.params?.cursor ? { tools: [descriptor("Last page")] } + : { tools: [descriptor("First page")], ...(overlap ? { nextCursor: "overlap" } : {}) }; + return new Response(JSON.stringify({ jsonrpc: "2.0", id: body.id, result }), { headers: { "content-type": "application/json" } }); + }, + }); + await service.refreshCatalog(connectionId); + const [first] = await db.select().from(toolCatalogEntries).where(eq(toolCatalogEntries.connectionId, connectionId)); + overlap = true; + expect((await service.refreshCatalog(connectionId)).catalog).toHaveLength(1); + const rows = await db.select().from(toolCatalogEntries).where(eq(toolCatalogEntries.connectionId, connectionId)); + expect(rows).toHaveLength(1); + expect(rows[0]).toMatchObject({ id: first.id, firstSeenAt: first.firstSeenAt, description: "Last page" }); + }); + + it("preserves operator-disabled tools and unchanged pending reviews", async () => { + const tools = ["off", "review"].map(name => ({ name, description: "Initial tool", inputSchema: { type: "object" }, annotations: { readOnlyHint: true } })); + const service = serviceWithTools(tools); + await service.refreshCatalog(connectionId); + await db.update(toolConnections).set({ config: { url: "https://mcp.example.test/mcp", quarantineNewEntries: true } }).where(eq(toolConnections.id, connectionId)); + await db.update(toolCatalogEntries).set({ status: "disabled", quarantineReason: null }).where(and(eq(toolCatalogEntries.connectionId, connectionId), eq(toolCatalogEntries.toolName, "off"))); + const reviewedAt = new Date("2020-01-01T00:00:00Z"); + await db.update(toolCatalogEntries).set({ status: "quarantined", quarantineReason: "pending_review", quarantinedAt: reviewedAt }).where(and(eq(toolCatalogEntries.connectionId, connectionId), eq(toolCatalogEntries.toolName, "review"))); + const before = await db.select().from(toolCatalogEntries).where(eq(toolCatalogEntries.connectionId, connectionId)).orderBy(toolCatalogEntries.toolName); + const result = await service.refreshCatalog(connectionId); + expect(result.catalog.map(entry => [entry.toolName, entry.status])).toEqual([["off", "disabled"], ["review", "quarantined"]]); + const after = await db.select().from(toolCatalogEntries).where(eq(toolCatalogEntries.connectionId, connectionId)).orderBy(toolCatalogEntries.toolName); + expect(after.map(row => [row.id, row.firstSeenAt, row.status, row.quarantinedAt, row.quarantineReason])).toEqual(before.map(row => [row.id, row.firstSeenAt, row.status, row.quarantinedAt, row.quarantineReason])); + }); + + it.each(["new", "quarantined"] as const)("does not restore an Off action after overlapping refreshes of a %s tool", async state => { + await serviceWithTools([]).refreshCatalog(connectionId); + if (state === "quarantined") await db.insert(toolCatalogEntries).values({ + companyId, applicationId, connectionId, name: "search", toolName: "search", riskLevel: "read", + status: "quarantined", versionHash: "fixture-version", schemaHash: "fixture-schema", quarantineReason: "pending_review", quarantinedAt: new Date("2020-01-01T00:00:00Z"), + }); + const [agent] = await db.insert(agents).values({ companyId, name: "Catalog permission fixture", role: "engineer", adapterType: "process", adapterConfig: {}, runtimeConfig: {} }).returning(); + const service = serviceWithTools([{ name: "search", annotations: { readOnlyHint: true } }]); + const transaction = db.transaction.bind(db); + let paused = false; + let release!: () => void; + const resume = new Promise(resolve => { release = resolve; }); + const transactionSpy = vi.spyOn(db, "transaction").mockImplementation((async (operation, config) => { + if (!paused) { paused = true; await resume; } + return transaction(operation, config); + }) as typeof db.transaction); + const later = Promise.allSettled([service.refreshCatalog(connectionId, undefined, { enableAllByDefault: true })]); + try { + await vi.waitFor(() => expect(paused).toBe(true), { timeout: 5000 }); + const first = await service.refreshCatalog(connectionId, undefined, { enableAllByDefault: true }); + const entry = first.catalog[0]; + const decide = () => toolAccessPolicyService(db).decide({ companyId, actor: { actorType: "agent", actorId: agent.id }, request: { connectionId, catalogEntryId: entry.id, toolName: entry.toolName, arguments: {} } }); + await expect(decide()).resolves.toMatchObject({ decision: "allow", reasonCode: "allow_profile" }); + const [selection] = await db.select().from(toolProfileEntries).where(eq(toolProfileEntries.catalogEntryId, entry.id)); + await service.deleteProfileEntry(selection.id); + await expect(decide()).resolves.toMatchObject({ decision: "deny" }); + release(); + const [result] = await later; + if (result.status === "rejected") throw result.reason; + await expect(decide()).resolves.toMatchObject({ decision: "deny" }); + expect(await db.select().from(toolCatalogEntries).where(eq(toolCatalogEntries.connectionId, connectionId))).toHaveLength(1); + expect(await db.select().from(toolProfileEntries).where(eq(toolProfileEntries.catalogEntryId, entry.id))).toHaveLength(0); + } finally { release(); await later; transactionSpy.mockRestore(); } + }, 15000); + + it("keeps committed chunks and rolls back a failed default-profile publication", async () => { + const service = serviceWithTools(Array.from({ length: 501 }, (_, index) => ({ name: `search_${index}`, annotations: { readOnlyHint: true } }))); + const transaction = db.transaction.bind(db); + let catalogWrites = 0; + let failProfile = false; + const transactionSpy = vi.spyOn(db, "transaction").mockImplementation((async (operation, config) => transaction(async tx => { + const insert = tx.insert.bind(tx); + const insertSpy = vi.spyOn(tx, "insert").mockImplementation(((table: Parameters[0]) => { + if (table === toolCatalogEntries && ++catalogWrites === 2) throw new Error("second chunk failed"); + if (table === toolProfileBindings && failProfile) throw new Error("profile binding failed"); + return insert(table); + }) as typeof tx.insert); + try { return await operation(tx); } + finally { insertSpy.mockRestore(); } + }, config)) as typeof db.transaction); + await expect(service.refreshCatalog(connectionId)).rejects.toThrow("second chunk failed"); + expect(await db.select().from(toolCatalogEntries).where(eq(toolCatalogEntries.connectionId, connectionId))).toHaveLength(500); + failProfile = true; + await expect(service.refreshCatalog(connectionId)).rejects.toThrow("profile binding failed"); + expect(await db.select().from(toolCatalogEntries).where(eq(toolCatalogEntries.connectionId, connectionId))).toHaveLength(501); + expect(await db.select().from(toolProfiles).where(eq(toolProfiles.companyId, companyId))).toHaveLength(0); + expect(await db.select().from(toolProfileBindings).where(eq(toolProfileBindings.companyId, companyId))).toHaveLength(0); + transactionSpy.mockRestore(); + await service.refreshCatalog(connectionId); + expect(await db.select().from(toolProfiles).where(eq(toolProfiles.companyId, companyId))).toHaveLength(1); + }); + + it("waits for application publication before changing catalog rows", async () => { + const tools = [{ name: "search", description: "Initial search", annotations: { readOnlyHint: true } }]; + const service = serviceWithTools(tools); + await service.refreshCatalog(connectionId); + tools[0].description = "Changed search"; + let locked!: () => void; + let release!: () => void; + const held = new Promise(resolve => { locked = resolve; }); + const resume = new Promise(resolve => { release = resolve; }); + const blocker = db.transaction(async tx => { + await tx.select().from(toolApplications).where(eq(toolApplications.id, applicationId)).for("no key update"); + locked(); await resume; + }); + await held; + const transaction = db.transaction.bind(db); + let started!: (pid: number) => void; + const begun = new Promise(resolve => { started = resolve; }); + const transactionSpy = vi.spyOn(db, "transaction").mockImplementation((async (operation, config) => transaction(async tx => { + await tx.execute(sql`SET LOCAL lock_timeout = '8s'`); + const [backend] = await tx.execute(sql`SELECT pg_backend_pid() AS pid`); + started(Number(backend.pid)); + return operation(tx); + }, config)) as typeof db.transaction); + const refresh = Promise.allSettled([service.refreshCatalog(connectionId)]); + try { + const pid = await begun; + await vi.waitFor(async () => { + const [activity] = await db.execute(sql`SELECT wait_event_type FROM pg_stat_activity WHERE pid = ${pid}`); + expect(activity.wait_event_type).toBe("Lock"); + }, { timeout: 5000 }); + expect((await db.select().from(toolCatalogEntries).where(eq(toolCatalogEntries.connectionId, connectionId)))[0].description).toBe("Initial search"); + } finally { release(); await blocker; } + const [result] = await refresh; + if (result.status === "rejected") throw result.reason; + transactionSpy.mockRestore(); + expect((await db.select().from(toolCatalogEntries).where(eq(toolCatalogEntries.connectionId, connectionId)))[0].description).toBe("Changed search"); + }, 15000); +}); diff --git a/server/src/services/tool-access.ts b/server/src/services/tool-access.ts index 516aa8fba62..852860cc1f6 100644 --- a/server/src/services/tool-access.ts +++ b/server/src/services/tool-access.ts @@ -3065,6 +3065,14 @@ async function gitHubReadGrantAccess(db: Db, companyId: string, connectionId: st return { connection, allowed, legacyShared }; } +async function lockManagedAppPolicies(client: ToolAccessMutationDb, companyId: string, connectionId: string) { + await client.select({ id: toolPolicies.id }).from(toolPolicies).where(and( + eq(toolPolicies.companyId, companyId), eq(toolPolicies.policyType, "require_approval"), + sql`${toolPolicies.config}->>'source' = 'app_gallery_finish'`, + sql`${toolPolicies.config}->>'connectionId' = ${connectionId}`, + )).orderBy(toolPolicies.id).for("no key update"); +} + export function toolAccessService( db: Db, options: ToolAccessServiceOptions = {}, @@ -5629,7 +5637,7 @@ export function toolAccessService( activeCatalogEntryIds: string[]; restoreDraftDefaults?: boolean; actor?: ActorInfo; - }) { + }, dbClient: ToolAccessMutationDb = db) { // Catalog discovery also runs while the setup wizard is still a draft. // Access is not granted until the operator finishes that wizard, so a // draft refresh must never manufacture a profile or company-wide binding. @@ -5640,7 +5648,7 @@ export function toolAccessService( if (input.connection.status !== "active" && !input.restoreDraftDefaults) return; const profileKey = `app:${input.connection.id}`; - let [profile] = await db + let [profile] = await dbClient .select() .from(toolProfiles) .where( @@ -5649,11 +5657,12 @@ export function toolAccessService( eq(toolProfiles.profileKey, profileKey), ), ) - .limit(1); + .limit(1).for("no key update"); + await lockManagedAppPolicies(dbClient, input.connection.companyId, input.connection.id); const createdProfile = !profile; const resetToRecommendedDefaults = !profile || profile.status !== "active"; if (!profile) { - const [sameName] = await db + const [sameName] = await dbClient .select({ id: toolProfiles.id }) .from(toolProfiles) .where( @@ -5666,7 +5675,7 @@ export function toolAccessService( const profileName = sameName ? `${input.connection.name} (${input.connection.id.replace(/-/g, "").slice(0, 8)})` : input.connection.name; - [profile] = await db + [profile] = await dbClient .insert(toolProfiles) .values({ companyId: input.connection.companyId, @@ -5681,7 +5690,7 @@ export function toolAccessService( }, }) .returning(); - await db.insert(toolProfileBindings).values({ + await dbClient.insert(toolProfileBindings).values({ companyId: input.connection.companyId, profileId: profile.id, targetType: "company", @@ -5702,7 +5711,7 @@ export function toolAccessService( // references it (for example, an MCP gateway). Reconnecting revives the // same connection id, so restore that retained profile as a fresh // all-agents assignment instead of leaving every read action Off. - [profile] = await db + [profile] = await dbClient .update(toolProfiles) .set({ name: input.connection.name, @@ -5717,7 +5726,7 @@ export function toolAccessService( }) .where(eq(toolProfiles.id, profile.id)) .returning(); - await db + await dbClient .delete(toolProfileBindings) .where( and( @@ -5725,7 +5734,7 @@ export function toolAccessService( eq(toolProfileBindings.profileId, profile.id), ), ); - await db + await dbClient .delete(toolProfileEntries) .where( and( @@ -5733,7 +5742,7 @@ export function toolAccessService( eq(toolProfileEntries.profileId, profile.id), ), ); - await db.insert(toolProfileBindings).values({ + await dbClient.insert(toolProfileBindings).values({ companyId: input.connection.companyId, profileId: profile.id, targetType: "company", @@ -5761,7 +5770,7 @@ export function toolAccessService( connection: input.connection, askFirstEntries: [], actor: input.actor, - }); + }, dbClient); } // A new or revived connection starts with every discovered action enabled. @@ -5775,7 +5784,7 @@ export function toolAccessService( ), ]; if (candidateIds.length === 0) return; - const existingEntries = await db + const existingEntries = await dbClient .select({ catalogEntryId: toolProfileEntries.catalogEntryId }) .from(toolProfileEntries) .where( @@ -5792,7 +5801,7 @@ export function toolAccessService( ); const entryIds = candidateIds.filter((id) => !configuredIds.has(id)); if (entryIds.length === 0) return; - await db.insert(toolProfileEntries).values( + await dbClient.insert(toolProfileEntries).values( entryIds.map((catalogEntryId) => ({ companyId: input.connection.companyId, profileId: profile.id, @@ -6732,6 +6741,9 @@ export function toolAccessService( .returning({ id: connectionTokenIssuances.id }); const archived = await db.transaction(async (tx) => { + await tx.select({ id: toolApplications.id }).from(toolApplications) + .where(and(eq(toolApplications.id, connection.applicationId), eq(toolApplications.companyId, connection.companyId))) + .for("update"); const [updatedConnection] = await tx .update(toolConnections) .set({ status: "archived", enabled: false, updatedAt: now }) @@ -7755,6 +7767,16 @@ export function toolAccessService( })); } + const publish = (write: (client: ToolAccessMutationDb) => Promise): Promise => db.transaction(async tx => { + await tx.select({ id: toolApplications.id }).from(toolApplications) + .where(and(eq(toolApplications.id, connection.applicationId), eq(toolApplications.companyId, connection.companyId))) + .for("update"); + await tx.select({ id: toolConnections.id }).from(toolConnections) + .where(and(eq(toolConnections.id, connection.id), eq(toolConnections.companyId, connection.companyId))) + .for("update"); + return write(tx); + }); + descriptors = [...new Map(descriptors.map(descriptor => [descriptor.name, descriptor])).values()]; const existingRows = await db .select() .from(toolCatalogEntries) @@ -7767,9 +7789,9 @@ export function toolAccessService( const discoveredNames = new Set(descriptors.map((descriptor) => descriptor.name)); const removedStatus = nativeDaytona ? "removed" : "disabled"; const removedIds = existingRows.filter((entry) => !discoveredNames.has(entry.toolName) && entry.status !== removedStatus).map((entry) => entry.id); - if (removedIds.length) await db.update(toolCatalogEntries) + if (removedIds.length) await publish(tx => tx.update(toolCatalogEntries) .set({ status: removedStatus, quarantineReason: nativeDaytona ? "provider_tool_removed" : "mcp_tool_removed", updatedAt: refreshedAt }) - .where(and(eq(toolCatalogEntries.companyId, connection.companyId), eq(toolCatalogEntries.connectionId, connection.id), inArray(toolCatalogEntries.id, removedIds))); + .where(and(eq(toolCatalogEntries.companyId, connection.companyId), eq(toolCatalogEntries.connectionId, connection.id), inArray(toolCatalogEntries.id, removedIds)))); } // Retired native actions are absent from discovery, but old catalog rows // still need to show as disabled. Gateway denial also applies before refresh. @@ -7777,10 +7799,10 @@ export function toolAccessService( ? existingRows.filter((entry) => isRailwayToolBlocked(entry.toolName)).map((entry) => entry.id) : []; if (blockedRailwayEntryIds.length > 0) { - await db.update(toolCatalogEntries).set({ status: "disabled", updatedAt: refreshedAt }) - .where(and(eq(toolCatalogEntries.connectionId, connection.id), inArray(toolCatalogEntries.id, blockedRailwayEntryIds))); + await publish(tx => tx.update(toolCatalogEntries).set({ status: "disabled", updatedAt: refreshedAt }) + .where(and(eq(toolCatalogEntries.connectionId, connection.id), inArray(toolCatalogEntries.id, blockedRailwayEntryIds)))); } - const updatedEntries: ToolCatalogEntry[] = []; + const catalogValues: Array = []; let quarantinedCount = 0; const sourceTemplateKey = typeof asRecord(connection.config).sourceTemplateKey === "string" @@ -7844,69 +7866,45 @@ export function toolAccessService( : "active"; if (shouldQuarantine && !providerPermanentlyBlocked) quarantinedCount += 1; - if (existing) { - const [updated] = await db - .update(toolCatalogEntries) - .set({ - title: descriptor.title ?? null, - description: descriptor.description ?? null, - inputSchema: descriptor.inputSchema ?? {}, - annotations: descriptor.annotations ?? {}, - riskLevel, - isReadOnly: riskLevel === "read", - isWrite: riskLevel === "write", - isDestructive: riskLevel === "destructive", - status, - versionHash: hash, - schemaHash, - lastSeenAt: refreshedAt, - quarantinedAt: - status === "quarantined" - ? shouldQuarantine - ? refreshedAt - : existing.quarantinedAt - : null, - quarantineReason: - status === "quarantined" - ? shouldQuarantine - ? "pending_review" - : existing.quarantineReason - : null, - updatedAt: refreshedAt, - }) - .where(eq(toolCatalogEntries.id, existing.id)) - .returning(); - updatedEntries.push(toCatalogEntry(updated)); - } else { - const [created] = await db - .insert(toolCatalogEntries) - .values({ - companyId: connection.companyId, - applicationId: connection.applicationId, - connectionId: connection.id, - name: descriptor.name, - toolName: descriptor.name, - entryKind: "tool", - title: descriptor.title ?? null, - description: descriptor.description ?? null, - inputSchema: descriptor.inputSchema ?? {}, - annotations: descriptor.annotations ?? {}, - riskLevel, - isReadOnly: riskLevel === "read", - isWrite: riskLevel === "write", - isDestructive: riskLevel === "destructive", - status, - versionHash: hash, - schemaHash, - firstSeenAt: refreshedAt, - lastSeenAt: refreshedAt, - quarantinedAt: shouldQuarantine ? refreshedAt : null, - quarantineReason: shouldQuarantine ? "pending_review" : null, - }) - .returning(); - updatedEntries.push(toCatalogEntry(created)); - } + catalogValues.push({ + id: existing?.id, companyId: connection.companyId, applicationId: connection.applicationId, + connectionId: connection.id, name: existing?.name ?? descriptor.name, toolName: descriptor.name, entryKind: "tool", + title: descriptor.title ?? null, description: descriptor.description ?? null, + inputSchema: descriptor.inputSchema ?? {}, annotations: descriptor.annotations ?? {}, riskLevel, + isReadOnly: riskLevel === "read", isWrite: riskLevel === "write", isDestructive: riskLevel === "destructive", + status, versionHash: hash, schemaHash, firstSeenAt: existing?.firstSeenAt ?? refreshedAt, lastSeenAt: refreshedAt, + quarantinedAt: existing + ? status === "quarantined" ? shouldQuarantine ? refreshedAt : existing.quarantinedAt : null + : shouldQuarantine ? refreshedAt : null, + quarantineReason: existing + ? status === "quarantined" ? shouldQuarantine ? "pending_review" : existing.quarantineReason : null + : shouldQuarantine ? "pending_review" : null, + updatedAt: existing ? refreshedAt : new Date(), + }); } + const catalogRows: Array = []; + const newCatalogEntryIds = new Set(); + for (let index = 0; index < catalogValues.length; index += 500) { + const chunk = catalogValues.slice(index, index + 500); + catalogRows.push(...await publish(async tx => { + const current = await tx.select({ name: toolCatalogEntries.name, status: toolCatalogEntries.status }) + .from(toolCatalogEntries).where(and(eq(toolCatalogEntries.connectionId, connection.id), inArray(toolCatalogEntries.name, chunk.map(entry => entry.name)))); + const existingNonQuarantinedNames = new Set(current.filter(entry => entry.status !== "quarantined").map(entry => entry.name)); + const rows = await tx.insert(toolCatalogEntries).values(chunk) + .onConflictDoUpdate({ target: [toolCatalogEntries.connectionId, toolCatalogEntries.name], set: { + title: sql`excluded.title`, description: sql`excluded.description`, inputSchema: sql`excluded.input_schema`, + annotations: sql`excluded.annotations`, riskLevel: sql`excluded.risk_level`, isReadOnly: sql`excluded.is_read_only`, + isWrite: sql`excluded.is_write`, isDestructive: sql`excluded.is_destructive`, + status: sql`CASE WHEN ${toolCatalogEntries.status} = 'disabled' AND ${toolCatalogEntries.quarantineReason} IS DISTINCT FROM 'mcp_tool_removed' THEN 'disabled' ELSE excluded.status END`, + versionHash: sql`excluded.version_hash`, schemaHash: sql`excluded.schema_hash`, lastSeenAt: sql`excluded.last_seen_at`, + quarantinedAt: sql`excluded.quarantined_at`, quarantineReason: sql`excluded.quarantine_reason`, updatedAt: sql`excluded.updated_at`, + } }).returning(); + for (const row of rows) if (!existingNonQuarantinedNames.has(row.name)) newCatalogEntryIds.add(row.id); + return rows; + })); + } + const catalogByName = new Map(catalogRows.map(row => [row.toolName, row])); + const updatedEntries = descriptors.map(descriptor => toCatalogEntry(catalogByName.get(descriptor.name)!)); const preserveQuarantine = isRailwayEndpoint(connection.config.url) || sourceTemplateKey === "enterpret" || isDaytonaConnection(connection); @@ -7957,20 +7955,17 @@ export function toolAccessService( const preserveMcpAccess = connection.config.mcpPreserveAccess === true && isRemoteMcpConnectorMethod(connection.config.sourceTemplateKey, connection.config.connectionMethodKey); if (!refreshOptions.skipDefaultProfileSync || preserveMcpAccess) { - await enableCatalogEntriesByDefault({ + await publish(tx => enableCatalogEntriesByDefault({ connection: updatedConnection, // Discovery must not re-enable actions the operator turned Off, // including curated MCP connections during API-key replacement. newCatalogEntryIds: activeEntries - .filter((entry) => { - const previous = existingByName.get(entry.toolName); - return !previous || previous.status === "quarantined"; - }) + .filter((entry) => newCatalogEntryIds.has(entry.id)) .map((entry) => entry.id), activeCatalogEntryIds: activeEntries.map((entry) => entry.id), restoreDraftDefaults: refreshOptions.restoreDraftDefaults || preserveMcpAccess, actor, - }); + }, tx)); } await audit({ companyId: connection.companyId, @@ -13930,6 +13925,9 @@ export function toolAccessService( const config = { ...connection.config }; delete config.mcpPreserveAccess; const updated = await db.transaction(async (tx) => { + await tx.select({ id: toolApplications.id }).from(toolApplications) + .where(and(eq(toolApplications.id, connection.applicationId), eq(toolApplications.companyId, companyId))) + .for("update"); const [row] = await tx.update(toolConnections).set({ config, transportConfig: config, status: "active", enabled: true, ...(input.agentInstructions !== undefined ? { agentInstructions: input.agentInstructions } : {}), updatedAt: new Date() }) .where(and(eq(toolConnections.id, connection.id), eq(toolConnections.companyId, companyId))).returning(); await tx.update(toolApplications).set({ status: "active", updatedAt: new Date() }).where(and(eq(toolApplications.id, connection.applicationId), eq(toolApplications.companyId, companyId))); @@ -14021,6 +14019,9 @@ export function toolAccessService( metadata: { source: "app_gallery_finish" }, })); const transactionResult = await db.transaction(async (tx) => { + await tx.select({ id: toolApplications.id }).from(toolApplications) + .where(and(eq(toolApplications.id, connection.applicationId), eq(toolApplications.companyId, companyId))) + .for("update"); await tx .select({ id: toolConnections.id }) .from(toolConnections) @@ -14040,7 +14041,8 @@ export function toolAccessService( eq(toolProfiles.profileKey, profileKey), ), ) - .limit(1); + .limit(1).for("no key update"); + await lockManagedAppPolicies(tx, companyId, connection.id); // A renamed connection can leave its old profile name in use. Resolve // names against profiles too, including when finalizing an installed draft. const otherProfiles = await tx @@ -15625,6 +15627,12 @@ export function toolAccessService( ), ) .limit(1); + await tx.select({ id: toolApplications.id }).from(toolApplications) + .where(and(eq(toolApplications.id, connection.applicationId), eq(toolApplications.companyId, connection.companyId))) + .for("update"); + await tx.select({ id: toolConnections.id }).from(toolConnections) + .where(and(eq(toolConnections.id, connection.id), eq(toolConnections.companyId, connection.companyId))) + .for("update"); const existingRefs = existingCredentialGrant?.credentialSecretRefs ?? (personalCredential || agentCredential @@ -16245,6 +16253,12 @@ export function toolAccessService( ), ) .limit(1); + await tx.select({ id: toolApplications.id }).from(toolApplications) + .where(and(eq(toolApplications.id, connection.applicationId), eq(toolApplications.companyId, connection.companyId))) + .for("update"); + await tx.select({ id: toolConnections.id }).from(toolConnections) + .where(and(eq(toolConnections.id, connection.id), eq(toolConnections.companyId, connection.companyId))) + .for("update"); const subjectCredentialSecretRefs = existingUserGrant?.credentialSecretRefs ?? []; const accessRef = await createOrRotateOAuthSecret( @@ -16558,6 +16572,12 @@ export function toolAccessService( "Only a company owner, admin, or member with connection-manager permission can share credentials with the organization.", ); } + await tx.select({ id: toolApplications.id }).from(toolApplications) + .where(and(eq(toolApplications.id, connection.applicationId), eq(toolApplications.companyId, connection.companyId))) + .for("update"); + await tx.select({ id: toolConnections.id }).from(toolConnections) + .where(and(eq(toolConnections.id, connection.id), eq(toolConnections.companyId, connection.companyId))) + .for("update"); const txSecrets = secretService(tx); const txSecretContext = { dbClient: tx, secretClient: txSecrets };