fix(connections): isolate broker claim recovery state

Co-Authored-By: Paperclip <noreply@paperclip.ing>

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<typeof createDb>;
+  let tempDb: Awaited<ReturnType<typeof startEmbeddedPostgresTestDatabase>> | 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 });
+  });
+});

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<typeof claimResponseSchema.parse> } } | undefined)?.brokerClaimState;
-      return state?.claimCodeHash === claimCodeHash ? { claim: state.claim } : null;
+      const config = row?.config as {
+        brokerClaimStates?: Record<string, { claim: ReturnType<typeof claimResponseSchema.parse> }>;
+        brokerClaimState?: { claimCodeHash: string; claim: ReturnType<typeof claimResponseSchema.parse> };
+      } | 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<string, unknown>;
+        const brokerClaimStates = { ...(config.brokerClaimStates as Record<string, unknown> | 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<string, unknown>;
-      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<string, unknown>;
+        const brokerClaimStates = { ...(config.brokerClaimStates as Record<string, unknown> | 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)));
+      });
     },
   };
 }