From c3ba178ec1218bad4ebdd8e536e759626e51fc89 Mon Sep 17 00:00:00 2001 From: Dotta Date: Wed, 22 Jul 2026 15:52:25 +0000 Subject: [PATCH] fix(connections): isolate broker claim recovery state Co-Authored-By: Paperclip --- server/src/services/connection-broker.test.ts | 50 ++++++++++++++++++- server/src/services/connection-broker.ts | 42 +++++++++++----- 2 files changed, 79 insertions(+), 13 deletions(-) diff --git a/server/src/services/connection-broker.test.ts b/server/src/services/connection-broker.test.ts index cc530be0b33..5052ad874e0 100644 --- a/server/src/services/connection-broker.test.ts +++ b/server/src/services/connection-broker.test.ts @@ -1,9 +1,12 @@ import { generateKeyPairSync } from "node:crypto"; -import { describe, expect, it } from "vitest"; +import { afterAll, beforeAll, describe, expect, it } from "vitest"; import { verifyRequestEnvelope, type ReplayStore } from "@paperclip/connect-protocol"; +import { companies, createDb, toolApplications, toolConnections } from "@paperclipai/db"; +import { getEmbeddedPostgresTestSupport, startEmbeddedPostgresTestDatabase } from "../__tests__/helpers/embedded-postgres.js"; import { ConnectClaimInterceptionError, ConnectServiceResponseError, + connectionBrokerStore, connectionBrokerService, createSignedConnectBrokerClient, type ConnectionBrokerStore, @@ -131,3 +134,48 @@ describe("connection broker", () => { expect(f.claims.get("claim-code-at-least-22-chars")?.consumed).toBe(false); }); }); + +const embeddedPostgresSupport = await getEmbeddedPostgresTestSupport(); +const describeEmbeddedPostgres = embeddedPostgresSupport.supported ? describe : describe.skip; + +describeEmbeddedPostgres("connection broker claim recovery persistence", () => { + let db!: ReturnType; + let tempDb: Awaited> | null = null; + + beforeAll(async () => { + tempDb = await startEmbeddedPostgresTestDatabase("paperclip-connection-broker-"); + db = createDb(tempDb.connectionString); + }, 20_000); + + afterAll(async () => { + await tempDb?.cleanup(); + }); + + it("preserves independent concurrent claim recovery records", async () => { + const company = await db.insert(companies).values({ name: "Broker Claims", issuePrefix: "BCL" }).returning().then((rows) => rows[0]!); + const application = await db.insert(toolApplications).values({ companyId: company.id, name: "Broker App", type: "mcp_http" }).returning().then((rows) => rows[0]!); + const connection = await db.insert(toolConnections).values({ + companyId: company.id, + applicationId: application.id, + name: "Broker Connection", + uid: "broker-claim-recovery", + transport: "rest_api", + enabled: true, + config: {}, + }).returning().then((rows) => rows[0]!); + const store = connectionBrokerStore(db); + const claimA = { custodyMode: "B2" as const, grantSealed: "sealed-a", tokenMeta: { scopes: ["read"], expiresAt: null } }; + const claimB = { custodyMode: "B2" as const, grantSealed: "sealed-b", tokenMeta: { scopes: ["write"], expiresAt: null } }; + + await Promise.all([ + store.saveClaimState({ companyId: company.id, connectionId: connection.id, claimCodeHash: "hash-a", claim: claimA }), + store.saveClaimState({ companyId: company.id, connectionId: connection.id, claimCodeHash: "hash-b", claim: claimB }), + ]); + + await expect(store.loadClaimState({ companyId: company.id, connectionId: connection.id, claimCodeHash: "hash-a" })).resolves.toEqual({ claim: claimA }); + await expect(store.loadClaimState({ companyId: company.id, connectionId: connection.id, claimCodeHash: "hash-b" })).resolves.toEqual({ claim: claimB }); + await store.clearClaimState({ companyId: company.id, connectionId: connection.id, claimCodeHash: "hash-a" }); + await expect(store.loadClaimState({ companyId: company.id, connectionId: connection.id, claimCodeHash: "hash-a" })).resolves.toBeNull(); + await expect(store.loadClaimState({ companyId: company.id, connectionId: connection.id, claimCodeHash: "hash-b" })).resolves.toEqual({ claim: claimB }); + }); +}); diff --git a/server/src/services/connection-broker.ts b/server/src/services/connection-broker.ts index d05081b5015..f44971b454a 100644 --- a/server/src/services/connection-broker.ts +++ b/server/src/services/connection-broker.ts @@ -1,5 +1,5 @@ import { createHash, randomBytes, type KeyObject } from "node:crypto"; -import { and, eq } from "drizzle-orm"; +import { and, eq, sql } from "drizzle-orm"; import type { Db } from "@paperclipai/db"; import { companyMemberships, connectionGrants, toolConnections } from "@paperclipai/db"; import { @@ -64,21 +64,39 @@ export function connectionBrokerStore(db: Db): ConnectionBrokerStore { }, async loadClaimState({ companyId, connectionId, claimCodeHash }) { const row = await db.select({ config: toolConnections.config }).from(toolConnections).where(and(eq(toolConnections.companyId, companyId), eq(toolConnections.id, connectionId))).limit(1).then((rows) => rows[0]); - const state = (row?.config as { brokerClaimState?: { claimCodeHash: string; claim: ReturnType } } | undefined)?.brokerClaimState; - return state?.claimCodeHash === claimCodeHash ? { claim: state.claim } : null; + const config = row?.config as { + brokerClaimStates?: Record }>; + brokerClaimState?: { claimCodeHash: string; claim: ReturnType }; + } | undefined; + const keyedState = config?.brokerClaimStates?.[claimCodeHash]; + if (keyedState) return keyedState; + return config?.brokerClaimState?.claimCodeHash === claimCodeHash ? { claim: config.brokerClaimState.claim } : null; }, async saveClaimState({ companyId, connectionId, claimCodeHash, claim }) { - const row = await db.select({ config: toolConnections.config }).from(toolConnections).where(and(eq(toolConnections.companyId, companyId), eq(toolConnections.id, connectionId))).limit(1).then((rows) => rows[0]); - if (!row) throw new Error("Tool connection not found"); - await db.update(toolConnections).set({ config: { ...row.config, brokerClaimState: { claimCodeHash, claim } }, updatedAt: new Date() }).where(and(eq(toolConnections.companyId, companyId), eq(toolConnections.id, connectionId))); + await db.transaction(async (tx) => { + const txDb = tx as unknown as Db; + await tx.execute(sql`select id from ${toolConnections} where ${toolConnections.companyId} = ${companyId} and ${toolConnections.id} = ${connectionId} for update`); + const row = await txDb.select({ config: toolConnections.config }).from(toolConnections).where(and(eq(toolConnections.companyId, companyId), eq(toolConnections.id, connectionId))).limit(1).then((rows) => rows[0]); + if (!row) throw new Error("Tool connection not found"); + const config = { ...(row.config ?? {}) } as Record; + const brokerClaimStates = { ...(config.brokerClaimStates as Record | undefined) }; + brokerClaimStates[claimCodeHash] = { claim }; + await txDb.update(toolConnections).set({ config: { ...config, brokerClaimStates }, updatedAt: new Date() }).where(and(eq(toolConnections.companyId, companyId), eq(toolConnections.id, connectionId))); + }); }, async clearClaimState({ companyId, connectionId, claimCodeHash }) { - const row = await db.select({ config: toolConnections.config }).from(toolConnections).where(and(eq(toolConnections.companyId, companyId), eq(toolConnections.id, connectionId))).limit(1).then((rows) => rows[0]); - const config = { ...(row?.config ?? {}) } as Record; - const state = config.brokerClaimState as { claimCodeHash?: string } | undefined; - if (state?.claimCodeHash !== claimCodeHash) return; - delete config.brokerClaimState; - await db.update(toolConnections).set({ config, updatedAt: new Date() }).where(and(eq(toolConnections.companyId, companyId), eq(toolConnections.id, connectionId))); + await db.transaction(async (tx) => { + const txDb = tx as unknown as Db; + await tx.execute(sql`select id from ${toolConnections} where ${toolConnections.companyId} = ${companyId} and ${toolConnections.id} = ${connectionId} for update`); + const row = await txDb.select({ config: toolConnections.config }).from(toolConnections).where(and(eq(toolConnections.companyId, companyId), eq(toolConnections.id, connectionId))).limit(1).then((rows) => rows[0]); + if (!row) return; + const config = { ...(row.config ?? {}) } as Record; + const brokerClaimStates = { ...(config.brokerClaimStates as Record | undefined) }; + delete brokerClaimStates[claimCodeHash]; + const legacyState = config.brokerClaimState as { claimCodeHash?: string } | undefined; + if (legacyState?.claimCodeHash === claimCodeHash) delete config.brokerClaimState; + await txDb.update(toolConnections).set({ config: { ...config, brokerClaimStates }, updatedAt: new Date() }).where(and(eq(toolConnections.companyId, companyId), eq(toolConnections.id, connectionId))); + }); }, }; }