Skip to content

fix: keep rate-limited sessions on the failover account #5

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Emilvooo wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: keep rate-limited sessions on the failover account #5

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
69 changes: 49 additions & 20 deletions index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -213,14 +221,27 @@ 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);
}

const account = await accountManager.getNextAvailableAccountForNewSession(model);
if (account) {
sessionBindingStore.set(sessionKey, account.index);
bindSessionAccount(sessionKey, account);
}
return account;
};
Expand Down Expand Up @@ -337,12 +358,34 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => {
retryCount = 0,
triedAccountIndices: Set<number> = new Set(),
): Promise<Response> => {
let originalBody: Record<string, unknown> = {};
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);
}
Expand Down Expand Up @@ -372,24 +415,6 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => {
);
}

let originalBody: Record<string, unknown> = {};
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 || "");

Expand Down Expand Up @@ -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);
}
Expand All @@ -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);
}
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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());
}
Expand Down
9 changes: 9 additions & 0 deletions lib/accounts/manager.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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 {
Expand Down
20 changes: 20 additions & 0 deletions test/account-manager-strategy.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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");
Expand Down
Loading