From c9fe1801005b05f80c48fd3e3d4d4d83aaa67787 Mon Sep 17 00:00:00 2001 From: Milos Date: Sat, 14 Mar 2026 22:31:48 +0100 Subject: [PATCH] fix: keep rate-limited sessions on the failover account --- README.md | 4 +- index.ts | 69 ++++++--- lib/accounts/manager.ts | 9 ++ test/account-manager-strategy.test.ts | 20 +++ test/session-failover-binding.test.ts | 209 ++++++++++++++++++++++++++ 5 files changed, 289 insertions(+), 22 deletions(-) create mode 100644 test/session-failover-binding.test.ts diff --git a/README.md b/README.md index d872f22..a216a4a 100644 --- a/README.md +++ b/README.md @@ -75,8 +75,8 @@ When you hit a rate limit: 1. Plugin detects 429 (rate limited) response 2. Marks current account as limited for that model -3. Keeps the current session on the same account (no mid-turn hot-swap) -4. Keeps that session/account binding; start a new session to switch accounts +3. Retries the request on the next available account +4. Rebinds that session to the new account so later prompts keep using it 5. Shows toast notification for account usage and rate limit status Session bindings are persisted locally so the same `prompt_cache_key` stays on the same account even after plugin process restarts. diff --git a/index.ts b/index.ts index abb0c12..a86c7de 100644 --- a/index.ts +++ b/index.ts @@ -201,6 +201,14 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => { return accountManager.getAllAccounts().find((acc) => acc.index === index) || null; }; + const bindSessionAccount = ( + sessionKey: string | undefined, + account: ManagedAccount, + ) => { + if (!sessionKey) return; + sessionBindingStore.set(sessionKey, account.index); + }; + const getSessionBoundAccount = async ( sessionKey: string | undefined, model?: string, @@ -213,6 +221,19 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => { if (boundIndex !== undefined) { const bound = findAccountByIndex(boundIndex); if (bound) { + if (accountManager.isAccountAvailableForModel(bound, model)) { + return bound; + } + + const nextAccount = await accountManager.getNextAvailableAccountExcluding( + new Set([bound.index]), + model, + ); + if (nextAccount) { + bindSessionAccount(sessionKey, nextAccount); + return nextAccount; + } + return bound; } sessionBindingStore.delete(sessionKey); @@ -220,7 +241,7 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => { const account = await accountManager.getNextAvailableAccountForNewSession(model); if (account) { - sessionBindingStore.set(sessionKey, account.index); + bindSessionAccount(sessionKey, account); } return account; }; @@ -337,12 +358,34 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => { retryCount = 0, triedAccountIndices: Set = new Set(), ): Promise => { + let originalBody: Record = {}; + if (typeof init?.body === "string") { + try { + originalBody = JSON.parse(init.body); + } catch { + originalBody = {}; + } + } + const isStreaming = originalBody.stream === true; + const model = + typeof originalBody.model === "string" + ? originalBody.model + : undefined; + const promptCacheKey = + typeof originalBody.prompt_cache_key === "string" + ? originalBody.prompt_cache_key + : undefined; + // Track this account as tried triedAccountIndices.add(account.index); const isTokenValid = await accountManager.ensureValidToken(account); if (!isTokenValid) { - const nextAccount = await accountManager.getNextAvailableAccountExcluding(triedAccountIndices); + const nextAccount = await accountManager.getNextAvailableAccountExcluding( + triedAccountIndices, + model, + ); if (nextAccount && nextAccount.index !== account.index) { + bindSessionAccount(promptCacheKey, nextAccount); await showAccountSwitchToast(account, nextAccount); return executeRequest(nextAccount, input, init, retryCount, triedAccountIndices); } @@ -372,24 +415,6 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => { ); } - let originalBody: Record = {}; - if (typeof init?.body === "string") { - try { - originalBody = JSON.parse(init.body); - } catch { - originalBody = {}; - } - } - const isStreaming = originalBody.stream === true; - const model = - typeof originalBody.model === "string" - ? originalBody.model - : undefined; - const promptCacheKey = - typeof originalBody.prompt_cache_key === "string" - ? originalBody.prompt_cache_key - : undefined; - const accountId = account.accountId || extractAccountIdFromToken(account.access || ""); @@ -491,6 +516,7 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => { const nextAccount = await accountManager.getNextAvailableAccountExcluding(triedAccountIndices, model); if (nextAccount && nextAccount.index !== account.index) { + bindSessionAccount(promptCacheKey, nextAccount); await showAccountSwitchToast(account, nextAccount); return executeRequest(nextAccount, input, init, retryCount + 1, triedAccountIndices); } @@ -502,6 +528,7 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => { const nextAccount = await accountManager.getNextAvailableAccountExcluding(triedAccountIndices, model); if (nextAccount && nextAccount.index !== account.index) { + bindSessionAccount(promptCacheKey, nextAccount); await showAccountSwitchToast(account, nextAccount); return executeRequest(nextAccount, input, init, retryCount + 1, triedAccountIndices); } @@ -551,6 +578,7 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => { if (debugMode) { console.log(`[openai-multi-auth] Model ${requestedModel} not supported on ${account.email || account.index} [${account.planType}], trying ${nextAccount.email || nextAccount.index} [${nextAccount.planType}]`); } + bindSessionAccount(promptCacheKey, nextAccount); await showModelRetryToast( requestedModel, account, @@ -580,6 +608,7 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => { // Get first available account for the fallback model const fallbackAccount = await accountManager.getNextAvailableAccount(fallbackModel); if (fallbackAccount) { + bindSessionAccount(promptCacheKey, fallbackAccount); // Reset tried accounts for the new model return executeRequest(fallbackAccount, input, modifiedInit, 0, new Set()); } diff --git a/lib/accounts/manager.ts b/lib/accounts/manager.ts index a5661e6..e936aec 100644 --- a/lib/accounts/manager.ts +++ b/lib/accounts/manager.ts @@ -311,6 +311,13 @@ export class AccountManager { return true; } + isAccountAvailableForModel( + account: ManagedAccount, + model?: string, + ): boolean { + return this.isAccountAvailable(account, model, Date.now()); + } + private getLeastRateLimitedAccount(model?: string): ManagedAccount | null { if (this.accounts.length === 0) return null; @@ -355,6 +362,8 @@ export class AccountManager { `[openai-multi-auth] ${identifier} rate limited until ${new Date(resetTime).toISOString()}`, ); } + + void this.saveToDisk(); } markRefreshFailed(account: ManagedAccount, error: string): void { diff --git a/test/account-manager-strategy.test.ts b/test/account-manager-strategy.test.ts index 8307ddf..2a2f68b 100644 --- a/test/account-manager-strategy.test.ts +++ b/test/account-manager-strategy.test.ts @@ -101,6 +101,26 @@ describe("AccountManager strategy selection", () => { expect(third?.index).toBe(1); }); + it("persists rate-limited accounts across manager reloads", async () => { + const home = mkdtempSync(join(tmpdir(), "strategy-sticky-persisted-limit-")); + const manager = await createManager(home, "sticky"); + await manager.loadFromDisk(); + + await manager.addAccount("a@example.com", "rt-1"); + await manager.addAccount("b@example.com", "rt-2"); + + const first = await manager.getNextAvailableAccount("gpt-5.2-codex"); + expect(first?.index).toBe(0); + + manager.markRateLimited(first!, 60_000, "gpt-5.2-codex"); + + const reloadedManager = await createManager(home, "sticky"); + await reloadedManager.loadFromDisk(); + + const next = await reloadedManager.getNextAvailableAccount("gpt-5.2-codex"); + expect(next?.index).toBe(1); + }); + it("skips rate-limited accounts and keeps round-robin progression", async () => { const home = mkdtempSync(join(tmpdir(), "strategy-rr-failover-")); const manager = await createManager(home, "round-robin"); diff --git a/test/session-failover-binding.test.ts b/test/session-failover-binding.test.ts new file mode 100644 index 0000000..09db37c --- /dev/null +++ b/test/session-failover-binding.test.ts @@ -0,0 +1,209 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +const mockState = vi.hoisted(() => ({ + accounts: [ + { + index: 0, + email: 'test-1@example.com', + access: 'access-token-1', + expires: Date.now() + 60_000, + accountId: 'acct_1', + }, + ], + activeIndex: 0, + limitedByModel: new Map>(), + sessionBindings: new Map(), +})); + +function createMockAccount(index: number) { + return { + index, + email: `test-${index + 1}@example.com`, + access: `access-token-${index + 1}`, + expires: Date.now() + 60_000, + accountId: `acct_${index + 1}`, + }; +} + +function getLimitedSet(model?: string) { + const key = model || '__global__'; + const existing = mockState.limitedByModel.get(key); + if (existing) return existing; + + const created = new Set(); + mockState.limitedByModel.set(key, created); + return created; +} + +function getNextAvailableAccount(model?: string, exclude: Set = new Set()) { + const limited = getLimitedSet(model); + const account = mockState.accounts.find( + (candidate) => !exclude.has(candidate.index) && !limited.has(candidate.index), + ); + if (!account) return null; + + mockState.activeIndex = account.index; + return account; +} + +vi.mock('@opencode-ai/plugin', () => ({ + tool: (definition: unknown) => definition, +})); + +vi.mock('../lib/accounts/index.js', () => { + class AccountManager { + async loadFromDisk() {} + async importFromOpenCodeAuth() {} + getAllAccounts() { + return mockState.accounts; + } + getAccountCount() { + return mockState.accounts.length; + } + getActiveAccount() { + return mockState.accounts[mockState.activeIndex] || null; + } + async getNextAvailableAccount(model?: string) { + return getNextAvailableAccount(model); + } + async getNextAvailableAccountForNewSession(model?: string) { + return getNextAvailableAccount(model); + } + async getNextAvailableAccountExcluding( + excludeIndices: Set, + model?: string, + ) { + return getNextAvailableAccount(model, excludeIndices); + } + isAccountAvailableForModel(account: { index: number }, model?: string) { + return !getLimitedSet(model).has(account.index); + } + async ensureValidToken() { + return true; + } + markRateLimited(account: { index: number }, _retryAfterMs: number, model?: string) { + getLimitedSet(model).add(account.index); + } + markRefreshFailed() {} + async addAccount() {} + } + + return { AccountManager }; +}); + +vi.mock('../lib/session-bindings.js', () => { + class SessionBindingStore { + loadFromDisk() {} + get(key: string) { + return mockState.sessionBindings.get(key); + } + set(key: string, value: number) { + mockState.sessionBindings.set(key, value); + } + delete(key: string) { + mockState.sessionBindings.delete(key); + } + } + + return { SessionBindingStore }; +}); + +describe('Session failover binding', () => { + beforeEach(() => { + vi.resetModules(); + mockState.accounts = [createMockAccount(0), createMockAccount(1)]; + mockState.activeIndex = 0; + mockState.limitedByModel.clear(); + mockState.sessionBindings.clear(); + }); + + it('rebinds later prompts to the failover account after a 429', async () => { + const authHeaders: string[] = []; + let requestCount = 0; + (globalThis as any).fetch = vi.fn(async (input: Request | string | URL, init?: RequestInit) => { + const url = String(input); + if (url.includes('/models')) { + return new Response(JSON.stringify({ models: [] }), { + status: 200, + headers: { 'content-type': 'application/json' }, + }); + } + + const headers = new Headers(init?.headers as HeadersInit | undefined); + authHeaders.push(headers.get('authorization') || ''); + requestCount++; + + if (requestCount === 1) { + return new Response( + JSON.stringify({ + resets_at: new Date(Date.now() + 60_000).toISOString(), + }), + { + status: 429, + headers: { 'content-type': 'application/json' }, + }, + ); + } + + return new Response('data: {"type":"response.done"}\n\n', { + status: 200, + headers: { 'content-type': 'text/event-stream' }, + }); + }); + + const { OpenAIAuthPlugin } = await import('../index.js'); + const plugin = await OpenAIAuthPlugin({ + client: { + auth: { set: vi.fn() }, + tui: { showToast: vi.fn() }, + }, + } as any); + + const loader = await plugin.auth.loader( + async () => ({ + type: 'oauth', + access: 'access-token-1', + refresh: 'refresh-token', + expires: Date.now() + 60_000, + }) as any, + {} as any, + ); + + const request = { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ + model: 'gpt-5.2-codex', + prompt_cache_key: 'ses_test_key', + input: [{ type: 'message', role: 'user', content: 'hello' }], + }), + }; + + await loader.fetch('https://chatgpt.com/backend-api/responses', request); + + expect(authHeaders).toEqual(['Bearer access-token-1', 'Bearer access-token-2']); + expect(mockState.sessionBindings.get('ses_test_key')).toBe(1); + + const secondPromptAuthHeaders: string[] = []; + (globalThis as any).fetch = vi.fn(async (input: Request | string | URL, init?: RequestInit) => { + const url = String(input); + if (url.includes('/models')) { + return new Response(JSON.stringify({ models: [] }), { + status: 200, + headers: { 'content-type': 'application/json' }, + }); + } + + const headers = new Headers(init?.headers as HeadersInit | undefined); + secondPromptAuthHeaders.push(headers.get('authorization') || ''); + return new Response('data: {"type":"response.done"}\n\n', { + status: 200, + headers: { 'content-type': 'text/event-stream' }, + }); + }); + + await loader.fetch('https://chatgpt.com/backend-api/responses', request); + + expect(secondPromptAuthHeaders).toEqual(['Bearer access-token-2']); + }); +});