fix(connections): make broker claims retryable

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

packages/connect-protocol/src/index.test.ts

@@ -114,7 +114,7 @@ describe("connect protocol", () => {
     const { privateKey, publicKey } = generateKeyPairSync("ed25519");
     const store = replayStore();
     const compactJws = signRequestEnvelope({
-      body: { claimCode: "claim-code-with-enough-entropy" },
+      body: { claimCode: "claim-code-with-enough-entropy", idempotencyKey: "a".repeat(64) },
       instanceId: "in_devinstance",
       keyId: "key_1",
       privateKey,
@@ -131,7 +131,7 @@ describe("connect protocol", () => {
       replayStore: store,
       now: new Date("2026-07-21T18:00:01.000Z"),
     };
-    await expect(verifyRequestEnvelope(input)).resolves.toMatchObject({ body: { claimCode: "claim-code-with-enough-entropy" } });
+    await expect(verifyRequestEnvelope(input)).resolves.toMatchObject({ body: { claimCode: "claim-code-with-enough-entropy", idempotencyKey: "a".repeat(64) } });
     await expect(verifyRequestEnvelope(input)).rejects.toMatchObject({ code: "replayed_jti", status: 401 });
   });
 

packages/connect-protocol/src/index.ts

@@ -52,7 +52,10 @@ export const handshakeResponseSchema = z.object({
   expiresAt: z.string().datetime(),
 });
 
-export const claimRequestSchema = z.object({ claimCode: z.string().min(22) });
+export const claimRequestSchema = z.object({
+  claimCode: z.string().min(22),
+  idempotencyKey: z.string().regex(/^[a-f0-9]{64}$/),
+});
 
 export const tokenMetadataSchema = z.object({
   scopes: z.array(z.string()),

server/src/services/connection-broker.test.ts

@@ -17,7 +17,7 @@ const CONNECTION_ID = "44444444-4444-4444-8444-444444444444";
 function fixture(mode: "A" | "B2" = "A") {
   const { privateKey, publicKey } = generateKeyPairSync("ed25519");
   const replay = new Set<string>();
-  const claims = new Map([["claim-code-at-least-22-chars", { owner: INSTANCE_ID, attempts: 0, consumed: false }]]);
+  const claims = new Map([["claim-code-at-least-22-chars", { owner: INSTANCE_ID, attempts: 0, consumed: false, idempotencyKey: "", response: null as Record<string, unknown> | null }]]);
   const saved: Parameters<ConnectionBrokerStore["saveGrant"]>[0][] = [];
   const secrets: Array<{ purpose: string; value: string }> = [];
   const replayStore: ReplayStore = { consume: (_instance, jti) => !replay.has(jti) && Boolean(replay.add(jti)) };
@@ -28,16 +28,22 @@ function fixture(mode: "A" | "B2" = "A") {
       try {
         const envelope = await verifyRequestEnvelope({ compactJws, publicKey, expectedInstanceId: INSTANCE_ID, expectedPath: path, replayStore });
         if (path === "/v1/claims") {
-          const code = (envelope.body as { claimCode: string }).claimCode;
+          const { claimCode: code, idempotencyKey } = envelope.body as { claimCode: string; idempotencyKey: string };
           const claim = claims.get(code);
           if (!claim) return { status: 404, body: { error: "unknown_claim" } };
           if (claim.attempts >= 5) return { status: 404, body: { error: "unknown_claim" } };
           if (claim.owner !== envelope.iss) { claim.attempts += 1; return { status: 403, body: { error: "claim_instance_mismatch" } }; }
-          if (claim.consumed) return { status: 409, body: { error: "claim_already_consumed" } };
+          if (claim.consumed) {
+            return claim.idempotencyKey === idempotencyKey && claim.response
+              ? { status: 200, body: claim.response }
+              : { status: 409, body: { error: "claim_already_consumed" } };
+          }
           claim.consumed = true;
-          return mode === "A"
-            ? { status: 200, body: { custodyMode: "A", grantRef: "gr_testgrant", tokenMeta: { scopes: ["read"], expiresAt: null } } }
-            : { status: 200, body: { custodyMode: "B2", grantSealed: "sealed-refresh", tokenMeta: { scopes: ["read"], expiresAt: null } } };
+          claim.idempotencyKey = idempotencyKey;
+          claim.response = mode === "A"
+            ? { custodyMode: "A", grantRef: "gr_testgrant", tokenMeta: { scopes: ["read"], expiresAt: null } }
+            : { custodyMode: "B2", grantSealed: "sealed-refresh", tokenMeta: { scopes: ["read"], expiresAt: null } };
+          return { status: 200, body: claim.response };
         }
         if (path.endsWith("/token")) return { status: 200, body: { grantSealed: "sealed-access", tokenMeta: { scopes: ["read"], expiresAt: null } } };
         if (path === "/v1/handshakes") return { status: 201, body: { handshakeId: "hs_test", authorizeUrl: "https://connect.example/authorize", expiresAt: new Date(Date.now() + 60_000).toISOString() } };
@@ -74,7 +80,7 @@ describe("connection broker", () => {
   it("surfaces a 409 redemption as an interception alert", async () => {
     const f = fixture();
     await f.service.claimOAuth({ companyId: "company", connectionId: CONNECTION_ID, claimCode: "claim-code-at-least-22-chars", kind: "workspace" });
-    await expect(f.service.claimOAuth({ companyId: "company", connectionId: CONNECTION_ID, claimCode: "claim-code-at-least-22-chars", kind: "workspace" })).rejects.toBeInstanceOf(ConnectClaimInterceptionError);
+    await expect(f.client.claim("claim-code-at-least-22-chars", "b".repeat(64))).rejects.toBeInstanceOf(ConnectClaimInterceptionError);
   });
 
   it("does not consume mismatched claims and locks the service-side double after five attempts", async () => {
@@ -122,6 +128,28 @@ describe("connection broker", () => {
     expect(f.claimStates.size).toBe(0);
   });
 
+  it("retries the remote claim idempotently when local recovery persistence fails", async () => {
+    const f = fixture("B2");
+    let failSave = true;
+    const service = connectionBrokerService({
+      client: f.client,
+      store: {
+        ...f.store,
+        saveClaimState: async (input) => {
+          if (failSave) { failSave = false; throw new Error("database unavailable"); }
+          await f.store.saveClaimState(input);
+        },
+      },
+      enabled: true,
+      vault: { put: async (input) => ({ secretId: "secret-grant", versionSelector: "latest", configPath: "broker.grant", required: true, label: input.purpose }) },
+    });
+    const params = { companyId: "company", connectionId: CONNECTION_ID, claimCode: "claim-code-at-least-22-chars", kind: "workspace" as const };
+
+    await expect(service.claimOAuth(params)).rejects.toThrow("database unavailable");
+    expect(f.claims.get(params.claimCode)?.consumed).toBe(true);
+    await expect(service.claimOAuth(params)).resolves.toMatchObject({ custodyMode: "B2" });
+  });
+
   it("rejects user grants before consuming a claim when membership is inactive", async () => {
     const f = fixture();
     const service = connectionBrokerService({

server/src/services/connection-broker.ts

@@ -28,7 +28,7 @@ export function createSignedConnectBrokerClient(input: { instanceId: string; key
   }
   return {
     createHandshake: async (body: unknown) => handshakeResponseSchema.parse(await request("/v1/handshakes", handshakeRequestSchema.parse(body))),
-    claim: async (claimCode: string) => claimResponseSchema.parse(await request("/v1/claims", claimRequestSchema.parse({ claimCode }))),
+    claim: async (claimCode: string, idempotencyKey: string) => claimResponseSchema.parse(await request("/v1/claims", claimRequestSchema.parse({ claimCode, idempotencyKey }))),
     mintGrantToken: async (grantRef: string, grantWrapKey: string) => grantTokenResponseSchema.parse(await request(`/v1/grants/${encodeURIComponent(grantRef)}/token`, grantTokenRequestSchema.parse({ grantWrapKey }))),
     registerRelay: async (body: unknown) => relayRegistrationResponseSchema.parse(await request("/v1/connections/register", relayRegistrationRequestSchema.parse(body))),
   };
@@ -112,7 +112,7 @@ export function connectionBrokerService(input: { client: ConnectBrokerClient; st
       }
       const claimCodeHash = createHash("sha256").update(params.claimCode).digest("hex");
       const existingState = await input.store.loadClaimState({ companyId: params.companyId, connectionId: params.connectionId, claimCodeHash });
-      const claim = existingState?.claim ?? await input.client.claim(params.claimCode);
+      const claim = existingState?.claim ?? await input.client.claim(params.claimCode, claimCodeHash);
       if (!existingState) await input.store.saveClaimState({ companyId: params.companyId, connectionId: params.connectionId, claimCodeHash, claim });
       if (claim.custodyMode !== "A" && claim.custodyMode !== "B2") throw new ConnectServiceResponseError(403, "mode_not_allowed");
       const serviceGrantRef = claim.custodyMode === "A" ? claim.grantRef : null;