Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
8 changes: 5 additions & 3 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -75,9 +75,11 @@ When you hit a rate limit:

1. Plugin detects 429 (rate limited) response
2. Marks current account as limited for that model
3. Switches to next available account
4. Retries your request automatically
5. Shows toast notification: `Switched to account2@example.com`
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
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.

### Account Selection Strategies

Expand Down
170 changes: 118 additions & 52 deletions index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,19 +3,19 @@ import type { Auth } from "@opencode-ai/sdk";
import {
createAuthorizationFlow,
decodeJWT,
extractAccountIdFromToken,
exchangeAuthorizationCode,
parseAuthorizationInput,
REDIRECT_URI,
validateAuthorizationState,
} from "./lib/auth/auth.js";
import { openBrowserUrl } from "./lib/auth/browser.js";
import { startLocalOAuthServer } from "./lib/auth/server.js";
import { getCodexMode, loadPluginConfig } from "./lib/config.js";
import {
AUTH_LABELS,
CODEX_BASE_URL,
DUMMY_API_KEY,
ERROR_MESSAGES,
JWT_CLAIM_PATH,
LOG_STAGES,
PROVIDER_ID,
HTTP_STATUS,
Expand All @@ -28,13 +28,13 @@ import {
handleErrorResponse,
handleSuccessResponse,
rewriteUrlForCodex,
transformRequestForCodex,
validateCodexBackendUrl,
} from "./lib/request/fetch-helpers.js";
import type { UserConfig } from "./lib/types.js";
import { AccountManager } from "./lib/accounts/index.js";
import type { ManagedAccount } from "./lib/accounts/index.js";
import { codexStatus } from "./lib/codex-status.js";
import { prefetchModels } from "./lib/models.js";
import { SessionBindingStore } from "./lib/session-bindings.js";

function extractModelFromBody(body: string | undefined): string | undefined {
if (!body) return undefined;
Expand All @@ -46,6 +46,20 @@ function extractModelFromBody(body: string | undefined): string | undefined {
}
}

function extractPromptCacheKeyFromBody(body: string | undefined): string | undefined {
if (!body) return undefined;
try {
const parsed = JSON.parse(body);
if (typeof parsed?.prompt_cache_key !== "string") {
return undefined;
}
const key = parsed.prompt_cache_key.trim();
return key.length > 0 ? key : undefined;
} catch {
return undefined;
}
}

let lastToastAccountIndex: number | null = null;
let lastToastTime = 0;
const TOAST_DEBOUNCE_MS = 5000;
Expand Down Expand Up @@ -89,7 +103,7 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => {
try {
await client.tui.showToast({
body: {
message: `Switching ${fromLabel} ${toLabel}${toPlanLabel}`,
message: `Switching ${fromLabel} -> ${toLabel}${toPlanLabel}`,
variant: "info",
},
});
Expand Down Expand Up @@ -180,13 +194,48 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => {
await accountManager.loadFromDisk();
await accountManager.importFromOpenCodeAuth();

const buildManualOAuthFlow = (pkce: { verifier: string }, url: string) => ({
const sessionBindingStore = new SessionBindingStore();
sessionBindingStore.loadFromDisk();

const findAccountByIndex = (index: number): ManagedAccount | null => {
return accountManager.getAllAccounts().find((acc) => acc.index === index) || null;
};

const getSessionBoundAccount = async (
sessionKey: string | undefined,
model?: string,
): Promise<ManagedAccount | null> => {
if (!sessionKey) {
return accountManager.getNextAvailableAccount(model);
}

const boundIndex = sessionBindingStore.get(sessionKey);
if (boundIndex !== undefined) {
const bound = findAccountByIndex(boundIndex);
if (bound) {
return bound;
}
sessionBindingStore.delete(sessionKey);
}

const account = await accountManager.getNextAvailableAccountForNewSession(model);
if (account) {
sessionBindingStore.set(sessionKey, account.index);
}
return account;
};

const buildManualOAuthFlow = (
pkce: { verifier: string },
expectedState: string,
url: string,
) => ({
url,
method: "code" as const,
instructions: AUTH_LABELS.INSTRUCTIONS_MANUAL,
callback: async (input: string) => {
const parsed = parseAuthorizationInput(input);
if (!parsed.code) {
if (!parsed.code || !validateAuthorizationState(parsed.state, expectedState)) {
return { type: "failed" as const };
}
const tokens = await exchangeAuthorizationCode(
Expand Down Expand Up @@ -281,17 +330,6 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => {
);
}

const providerConfig = provider as
| { options?: Record<string, unknown>; models?: UserConfig["models"] }
| undefined;
const userConfig: UserConfig = {
global: providerConfig?.options || {},
models: providerConfig?.models || {},
};

const pluginConfig = loadPluginConfig();
const codexMode = getCodexMode(pluginConfig);

const executeRequest = async (
account: ManagedAccount,
input: Request | string | URL,
Expand All @@ -309,7 +347,10 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => {
return executeRequest(nextAccount, input, init, retryCount, triedAccountIndices);
}
return new Response(
JSON.stringify({ error: "All accounts failed token refresh" }),
JSON.stringify({
error:
"Token refresh failed for the current session account. Start a new session to switch accounts.",
}),
{
status: HTTP_STATUS.UNAUTHORIZED,
headers: { "Content-Type": "application/json" },
Expand All @@ -318,28 +359,39 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => {
}

const originalUrl = extractRequestUrl(input);
const url = rewriteUrlForCodex(originalUrl);
let url: string;
try {
url = validateCodexBackendUrl(rewriteUrlForCodex(originalUrl));
} catch {
return new Response(
JSON.stringify({ error: ERROR_MESSAGES.INVALID_BACKEND_URL }),
{
status: HTTP_STATUS.BAD_REQUEST,
headers: { "Content-Type": "application/json" },
},
);
}

const originalBody = init?.body
? JSON.parse(init.body as string)
: {};
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 = originalBody.model;

const transformation = await transformRequestForCodex(
init,
url,
userConfig,
codexMode,
);
const requestInit = transformation?.updatedInit ?? init;
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 ||
(() => {
const decoded = decodeJWT(account.access || "");
return decoded?.[JWT_CLAIM_PATH]?.chatgpt_account_id;
})();
account.accountId || extractAccountIdFromToken(account.access || "");

if (!accountId) {
logDebug(
Expand All @@ -363,17 +415,17 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => {
}

const headers = createCodexHeaders(
requestInit,
init,
accountId,
account.access || "",
{
model: transformation?.body.model,
promptCacheKey: (transformation?.body as any)?.prompt_cache_key,
model,
promptCacheKey,
},
);

const response = await fetch(url, {
...requestInit,
...init,
headers,
});

Expand Down Expand Up @@ -435,7 +487,6 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => {
);
} catch {}
}

if (retryCount < accountManager.getAccountCount() - 1) {
const nextAccount =
await accountManager.getNextAvailableAccountExcluding(triedAccountIndices, model);
Expand All @@ -446,7 +497,7 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => {
}
}

if (response.status === HTTP_STATUS.UNAUTHORIZED && retryCount < 1) {
if (response.status === HTTP_STATUS.UNAUTHORIZED) {
accountManager.markRefreshFailed(account, "401 Unauthorized");
const nextAccount =
await accountManager.getNextAvailableAccountExcluding(triedAccountIndices, model);
Expand All @@ -471,7 +522,7 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => {
fs.mkdirSync(logDir, { recursive: true });
fs.writeFileSync(path.join(logDir, "last-400-error.json"), JSON.stringify({
timestamp: new Date().toISOString(),
model: transformation?.body.model,
model,
status: response.status,
errorBody,
detail,
Expand All @@ -484,12 +535,15 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => {

// Log the error for debugging
if (debugMode) {
console.log(`[openai-multi-auth] 400 error for model ${transformation?.body.model} on account ${account.email || account.index} [${account.planType}]: ${JSON.stringify(errorBody)}`);
console.log(`[openai-multi-auth] 400 error for model ${model} on account ${account.email || account.index} [${account.planType}]: ${JSON.stringify(errorBody)}`);
}

// Check if it's a "model not supported" error
if (detail.includes("model is not supported") || detail.includes("not supported when using Codex")) {
const requestedModel = transformation?.body.model || model;
const requestedModel = typeof model === "string" ? model : "";
if (!requestedModel) {
return await handleErrorResponse(response);
}

// STEP 1: Try other accounts first (they might be Plus/Pro/Team and support the model)
const nextAccount = await accountManager.getNextAvailableAccountExcluding(triedAccountIndices, requestedModel);
Expand Down Expand Up @@ -552,8 +606,11 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => {
input: Request | string | URL,
init?: RequestInit,
): Promise<Response> {
const model = extractModelFromBody(init?.body as string);
const account = await accountManager.getNextAvailableAccount(model);
const requestBody =
typeof init?.body === "string" ? (init.body as string) : undefined;
const model = extractModelFromBody(requestBody);
const sessionKey = extractPromptCacheKeyFromBody(requestBody);
const account = await getSessionBoundAccount(sessionKey, model);

if (!account) {
return new Response(
Expand Down Expand Up @@ -583,7 +640,7 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => {

if (!serverInfo.ready) {
serverInfo.close();
return buildManualOAuthFlow(pkce, url);
return buildManualOAuthFlow(pkce, state, url);
}

return buildAutoOAuthFlow(pkce, state, url, serverInfo);
Expand All @@ -600,7 +657,7 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => {

if (!serverInfo.ready) {
serverInfo.close();
return buildManualOAuthFlow(pkce, url);
return buildManualOAuthFlow(pkce, state, url);
}

return buildAutoOAuthFlow(pkce, state, url, serverInfo);
Expand All @@ -610,8 +667,8 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => {
label: AUTH_LABELS.OAUTH_MANUAL,
type: "oauth" as const,
authorize: async () => {
const { pkce, url } = await createAuthorizationFlow();
return buildManualOAuthFlow(pkce, url);
const { pkce, state, url } = await createAuthorizationFlow();
return buildManualOAuthFlow(pkce, state, url);
},
},
{
Expand All @@ -620,6 +677,15 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => {
},
],
},
"chat.headers": async (
input: { model: { providerID: string }; sessionID: string },
output: { headers: Record<string, string> },
) => {
if (input.model.providerID !== PROVIDER_ID) return;
output.headers = output.headers || {};
output.headers.originator = "opencode";
output.headers.session_id = input.sessionID;
},
config: async (cfg) => {
cfg.command = cfg.command || {};
cfg.command["codex-status"] = {
Expand Down
Loading