diff --git a/plugins/account-pool/src/hub.ts b/plugins/account-pool/src/hub.ts index 7b9c83cb27..44aa8c3d04 100644 --- a/plugins/account-pool/src/hub.ts +++ b/plugins/account-pool/src/hub.ts @@ -423,16 +423,18 @@ export class AccountPoolHub { ): Promise { const existing = this.refreshes.get(account.id); if (existing !== undefined) return existing; - const secret = await this.options.accounts.readSecret(account.id); - const refresh = adapter - .refreshSecret({ - account, - secret, - accounts: this.options.accounts, - quotas: this.options.quotas, - fetch: this.options.fetch, - now: this.options.now, - }) + const refresh = this.options.accounts + .readSecret(account.id) + .then((secret) => + adapter.refreshSecret({ + account, + secret, + accounts: this.options.accounts, + quotas: this.options.quotas, + fetch: this.options.fetch, + now: this.options.now, + }), + ) .then((result) => { if (result.refreshed) { const quota = this.options.quotas.get(account.id); diff --git a/plugins/account-pool/src/server.test.ts b/plugins/account-pool/src/server.test.ts index 7047d74b8c..04e873244e 100644 --- a/plugins/account-pool/src/server.test.ts +++ b/plugins/account-pool/src/server.test.ts @@ -21,7 +21,7 @@ import type { ImportedClaudeCredentials, ImportedCodexCredentials, } from "./credentials.js"; -import { HubTokenStore } from "./store.js"; +import { AccountStore, HubTokenStore } from "./store.js"; import { createAccountPoolPlugin, helloResponse, @@ -2405,12 +2405,12 @@ describe("Account Pool plugin", () => { }); it("serializes refresh, writes new tokens with 0600 mode, and uses them", async () => { + let now = 1_800_000_000_000; let refreshCalls = 0; const authorizations: Array = []; const upstream = await startUpstream(async (request, response) => { if (request.url === "/oauth/token") { refreshCalls += 1; - await new Promise((resolve) => setTimeout(resolve, 25)); response.writeHead(200, { "content-type": "application/json" }); response.end( JSON.stringify({ @@ -2431,19 +2431,53 @@ describe("Account Pool plugin", () => { upstreamUrl: upstream.url, source: "import", options: { + now: () => now, importCredentials: async () => - importedCredentials({ expiresAt: Date.now() + 1_000 }), + importedCredentials({ expiresAt: now + 10 * 60 * 1_000 }), refreshUrl: `${upstream.url}/oauth/token`, }, }); + expect(refreshCalls).toBe(0); + now += 6 * 60 * 1_000; + let releaseSecretReads = () => {}; + const secretReadsReleased = new Promise((resolve) => { + releaseSecretReads = resolve; + }); + const readSecret = AccountStore.prototype.readSecret; + const pendingSecretReads: ReturnType[] = []; + const readSecretSpy = vi + .spyOn(AccountStore.prototype, "readSecret") + .mockImplementation(async function (this: AccountStore, accountId) { + const reading = readSecret.call(this, accountId); + pendingSecretReads.push(reading); + const secret = await reading; + await secretReadsReleased; + return secret; + }); + const recordUsed = vi.spyOn(AccountStore.prototype, "recordUsed"); const requests = [1, 2].map(() => fixture.host.harness.behavior.fetchHttp("POST", "/v1/messages", { headers: authHeaders(fixture.key), body: "{}", }), ); - const responses = await Promise.all(requests); - await Promise.all(responses.map((response) => response.text())); + try { + await vi.waitFor(() => { + expect(recordUsed).toHaveBeenCalledTimes(2); + }); + await Promise.all(recordUsed.mock.results.map((result) => result.value)); + await new Promise((resolve) => setImmediate(resolve)); + await Promise.all(pendingSecretReads); + releaseSecretReads(); + const responses = await Promise.all(requests); + expect(responses.map((response) => response.status)).toEqual([200, 200]); + await Promise.all(responses.map((response) => response.text())); + } finally { + releaseSecretReads(); + await Promise.allSettled(requests); + readSecretSpy.mockRestore(); + recordUsed.mockRestore(); + } expect(refreshCalls).toBe(1); expect(authorizations).toEqual(["Bearer oauth-new", "Bearer oauth-new"]); const secretPath = path.join( @@ -2465,6 +2499,103 @@ describe("Account Pool plugin", () => { expect((await fs.stat(secretPath)).mode & 0o777).toBe(0o600); }); + it("refreshes unrelated accounts independently", async () => { + let now = 1_800_000_000_000; + let releaseFirstRefresh = () => {}; + const firstRefreshReleased = new Promise((resolve) => { + releaseFirstRefresh = resolve; + }); + const refreshes: string[] = []; + const authorizations: Array = []; + const upstream = await startUpstream(async (request, response) => { + if (request.url === "/oauth/token") { + const { refresh_token } = z + .object({ refresh_token: z.string() }) + .parse(JSON.parse((await readRequestBody(request)).toString())); + refreshes.push(refresh_token); + if (refresh_token === "refresh-1") await firstRefreshReleased; + response.writeHead(200, { "content-type": "application/json" }); + response.end( + JSON.stringify({ + access_token: `new-${refresh_token}`, + refresh_token: `next-${refresh_token}`, + expires_in: 3600, + }), + ); + return; + } + authorizations.push(request.headers.authorization); + await readRequestBody(request); + response.writeHead(200, { "content-type": "application/json" }); + response.end("{}"); + }); + cleanups.push(upstream.close); + let imported = 0; + const fixture = await createFixture({ + upstreamUrl: upstream.url, + source: "import", + options: { + now: () => now, + importCredentials: async () => { + imported += 1; + return importedCredentials({ + accessToken: `access-${imported}`, + refreshToken: `refresh-${imported}`, + email: `account-${imported}@example.com`, + expiresAt: now + 10 * 60 * 1_000, + }); + }, + refreshUrl: `${upstream.url}/oauth/token`, + }, + }); + const second = accountSchema.parse( + await fixture.host.harness.behavior.callRpc("account.add", { + provider: "claude", + source: { kind: "import" }, + label: "second", + priority: 200, + }), + ); + expect(refreshes).toEqual([]); + now += 6 * 60 * 1_000; + const requests: Promise[] = []; + try { + requests.push( + fixture.host.harness.behavior.fetchHttp("POST", "/v1/messages", { + headers: authHeaders(fixture.key), + body: "{}", + }), + ); + await vi.waitFor(() => { + expect(refreshes).toEqual(["refresh-1"]); + }); + await fixture.host.harness.behavior.callRpc("account.setPriority", { + accountId: second.id, + priority: 0, + }); + requests.push( + fixture.host.harness.behavior.fetchHttp("POST", "/v1/messages", { + headers: authHeaders(fixture.key), + body: "{}", + }), + ); + await vi.waitFor(() => { + expect(authorizations).toEqual(["Bearer new-refresh-2"]); + }); + } finally { + releaseFirstRefresh(); + await Promise.allSettled(requests); + } + const responses = await Promise.all(requests); + expect(responses.map((response) => response.status)).toEqual([200, 200]); + await Promise.all(responses.map((response) => response.text())); + expect(refreshes).toEqual(["refresh-1", "refresh-2"]); + expect(authorizations).toEqual([ + "Bearer new-refresh-2", + "Bearer new-refresh-1", + ]); + }); + it("marks refresh and upstream authorization failures as account errors", async () => { const upstream = await startUpstream((request, response) => { if (request.url === "/oauth/token") {