diff --git a/src/main/services/provider-tool-turns/credential-resolver.ts b/src/main/services/provider-tool-turns/credential-resolver.ts index 4cc6e58d..468dfd6c 100644 --- a/src/main/services/provider-tool-turns/credential-resolver.ts +++ b/src/main/services/provider-tool-turns/credential-resolver.ts @@ -101,15 +101,26 @@ export async function resolveTurnCredential(deps: { } const creds = parseStoredOAuth(deps.provider.provider_key, secret); if (Date.now() < creds.expires) return toTurnCredential(creds); - if (!deps.provider.credential_id) throw new Error("Provider credential id is missing; reconnect the provider."); - const refreshed = await refreshOAuth(deps.provider.provider_key, creds, deps.integrationKeys); - deps.authBroker.updateProviderSecret({ - orgId: deps.orgId, - providerKey: deps.provider.provider_key, - credentialId: deps.provider.credential_id, - secret: JSON.stringify(refreshed), - }); - return toTurnCredential(refreshed); + return toTurnCredential(await refreshOnce(deps, creds)); +} + +/** Refreshes in flight, by stored credential: a rotating provider honours a refresh token once, so callers at once share one. */ +const refreshing = new Map>(); + +function refreshOnce(deps: { authBroker: AuthBrokerService; integrationKeys: IntegrationKeyStore; orgId: string; provider: ProviderRow }, creds: StoredOAuth): Promise { + const credentialId = deps.provider.credential_id; + if (!credentialId) throw new Error("Provider credential id is missing; reconnect the provider."); + const running = refreshing.get(credentialId); + if (running) return running; + console.info(`[credentials] refreshing the ${deps.provider.provider_key} sign-in`); + const refresh = refreshOAuth(deps.provider.provider_key, creds, deps.integrationKeys) + .then((refreshed) => { + deps.authBroker.updateProviderSecret({ orgId: deps.orgId, providerKey: deps.provider.provider_key, credentialId, secret: JSON.stringify(refreshed) }); + return refreshed; + }) + .finally(() => refreshing.delete(credentialId)); + refreshing.set(credentialId, refresh); + return refresh; } function requireProjectId(creds: StoredOAuth): string { diff --git a/tests/unit/credential-refresh-race.test.ts b/tests/unit/credential-refresh-race.test.ts new file mode 100644 index 00000000..f47b8b79 --- /dev/null +++ b/tests/unit/credential-refresh-race.test.ts @@ -0,0 +1,68 @@ +import { randomUUID } from "node:crypto"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, describe, expect, test, vi } from "vitest"; +import { DatabaseConnection } from "../../src/main/database/connection"; +import { OrganizationService } from "../../src/main/services/organization-service"; +import { ProviderService } from "../../src/main/services/provider-service"; +import { providerAdapters } from "../../src/main/services/provider-adapters"; +import { AuthBrokerService } from "../../src/main/services/auth-broker-service"; +import { IntegrationKeyStore } from "../../src/main/services/integration-key-store"; +import { loadConfiguredProvider } from "../../src/main/services/provider-internal"; +import { resolveTurnCredential } from "../../src/main/services/provider-tool-turns/credential-resolver"; + +const vault = { encryptSecret: (secret: string) => Buffer.from(secret, "utf8"), decryptSecret: (ciphertext: Buffer) => ciphertext.toString("utf8") }; + +afterEach(() => vi.unstubAllGlobals()); + +/** A token endpoint that rotates refresh tokens: each one works once, as rotating providers do. */ +function rotatingTokenEndpoint() { + let generation = 0; + const valid = new Set(["refresh-0"]); + const calls: string[] = []; + vi.stubGlobal("fetch", async (_url: string, init: RequestInit) => { + const used = (JSON.parse(String(init.body)) as { refresh_token: string }).refresh_token; + calls.push(used); + await new Promise((resolve) => setTimeout(resolve, 20)); + if (!valid.delete(used)) return new Response(JSON.stringify({ error: "invalid_grant" }), { status: 400, headers: { "content-type": "application/json" } }); + generation += 1; + valid.add(`refresh-${generation}`); + return new Response(JSON.stringify({ access_token: `access-${generation}`, refresh_token: `refresh-${generation}`, expires_in: 3600 }), { status: 200, headers: { "content-type": "application/json" } }); + }); + return calls; +} + +function expiredClaudePro() { + const connection = new DatabaseConnection(join(tmpdir(), `ac-refresh-${randomUUID()}.db`)); + const providers = new ProviderService(connection.db, providerAdapters, vault); + let orgId = ""; + let providerId = ""; + connection.transaction(() => { + orgId = new OrganizationService(connection.db).createOrganization({ organizationName: "Acme", employeeName: "Dana", mission: "m", goal: "g" }); + providerId = (connection.db.prepare("select id from providers where org_id = ? and provider_key = 'claude-pro'").get(orgId) as { id: string }).id; + providers.saveCredential({ orgId, providerId, secret: JSON.stringify({ access: "access-0", refresh: "refresh-0", expires: 0 }), model: "claude-sonnet-4-5" }); + }); + const deps = { authBroker: new AuthBrokerService(connection.db, vault), integrationKeys: new IntegrationKeyStore(connection.db, vault), orgId }; + const resolve = () => resolveTurnCredential({ ...deps, provider: loadConfiguredProvider(connection.db, orgId, providerId) }); + return { connection, resolve }; +} + +describe("an expired sign-in is refreshed once, however many runs need it", () => { + test("runs that need it at the same time share one refresh and the rotated token survives", async () => { + const calls = rotatingTokenEndpoint(); + const h = expiredClaudePro(); + const credentials = await Promise.all([h.resolve(), h.resolve(), h.resolve()]); + expect(credentials.map((credential) => credential.accessToken)).toEqual(["access-1", "access-1", "access-1"]); + expect(calls).toEqual(["refresh-0"]); + h.connection.close(); + }); + + test("a run that comes after the refresh uses the stored new sign-in", async () => { + const calls = rotatingTokenEndpoint(); + const h = expiredClaudePro(); + expect((await h.resolve()).accessToken).toBe("access-1"); + expect((await h.resolve()).accessToken).toBe("access-1"); + expect(calls).toEqual(["refresh-0"]); + h.connection.close(); + }); +});