Skip to content

Commit f68e3d9

Browse files
ndycodendycode
authored
Merge pull request #99 from sdip15fa/feature/deactivated-workspace-rotation
Retry deactivated workspaces during rotation
2 parents f613afd + 9c81a2f commit f68e3d9

File tree

2 files changed

+186
-41
lines changed

2 files changed

+186
-41
lines changed

‎index.ts‎

Lines changed: 49 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1782,9 +1782,10 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
17821782
}
17831783

17841784
while (true) {
1785-
let accountCount = accountManager.getAccountCount();
1786-
const attempted = new Set<number>();
1787-
let restartAccountTraversalWithFallback = false;
1785+
let accountCount = accountManager.getAccountCount();
1786+
const attempted = new Set<number>();
1787+
let restartAccountTraversalWithFallback = false;
1788+
let restartAccountTraversalAfterWorkspaceDeactivation = false;
17881789

17891790
while (attempted.size < Math.max(1, accountCount)) {
17901791
const selectionExplainability = accountManager.getSelectionExplainability(
@@ -2082,10 +2083,10 @@ while (attempted.size < Math.max(1, accountCount)) {
20822083
});
20832084

20842085
const workspaceDeactivated = isDeactivatedWorkspaceError(errorBody, response.status);
2085-
if (workspaceDeactivated) {
2086-
const accountLabel = formatAccountLabel(account, account.index);
2087-
accountManager.refundToken(account, modelFamily, model);
2088-
accountManager.recordFailure(account, modelFamily, model);
2086+
if (workspaceDeactivated) {
2087+
const accountLabel = formatAccountLabel(account, account.index);
2088+
accountManager.refundToken(account, modelFamily, model);
2089+
accountManager.recordFailure(account, modelFamily, model);
20892090
account.lastSwitchReason = "rotation";
20902091
runtimeMetrics.failedRequests++;
20912092
runtimeMetrics.accountRotations++;
@@ -2113,27 +2114,38 @@ while (attempted.size < Math.max(1, accountCount)) {
21132114
);
21142115
}
21152116

2116-
if (accountManager.removeAccount(account)) {
2117-
accountManager.saveToDiskDebounced();
2118-
attempted.clear();
2119-
accountCount = accountManager.getAccountCount();
2120-
await showToast(
2121-
`Workspace deactivated. Removed ${accountLabel} from rotation and switching accounts.`,
2122-
"warning",
2123-
{ duration: toastDurationMs },
2117+
// Keep deactivated workspace cleanup aligned with the existing
2118+
// refresh-token removal path. saveToDiskDebounced reuses the
2119+
// storage temp-file + EPERM/EBUSY retry path covered in
2120+
// test/storage.test.ts, and surfaced persistence failures stay
2121+
// sanitized by the redaction checks in test/login-runner.test.ts.
2122+
const removedCount = accountManager.removeAccountsWithSameRefreshToken(account);
2123+
if (removedCount > 0) {
2124+
accountManager.saveToDiskDebounced();
2125+
restartAccountTraversalAfterWorkspaceDeactivation = true;
2126+
const removalMessage = removedCount > 1
2127+
? `Workspace deactivated. Removed ${removedCount} related entries from rotation and switching accounts.`
2128+
: `Workspace deactivated. Removed ${accountLabel} from rotation and switching accounts.`;
2129+
await showToast(
2130+
removalMessage,
2131+
"warning",
2132+
{ duration: toastDurationMs },
2133+
);
2134+
break;
2135+
}
2136+
2137+
logWarn(
2138+
`[${PLUGIN_NAME}] Expected grouped account removal after workspace deactivation, but removed ${removedCount}.`,
2139+
);
2140+
accountManager.markAccountCoolingDown(
2141+
account,
2142+
ACCOUNT_LIMITS.AUTH_FAILURE_COOLDOWN_MS,
2143+
"auth-failure",
21242144
);
2145+
accountManager.saveToDiskDebounced();
21252146
break;
21262147
}
21272148

2128-
accountManager.markAccountCoolingDown(
2129-
account,
2130-
ACCOUNT_LIMITS.AUTH_FAILURE_COOLDOWN_MS,
2131-
"auth-failure",
2132-
);
2133-
accountManager.saveToDiskDebounced();
2134-
break;
2135-
}
2136-
21372149
const unsupportedModelInfo = getUnsupportedCodexModelInfo(errorBody);
21382150
const hasRemainingAccounts = attempted.size < Math.max(1, accountCount);
21392151

@@ -2432,14 +2444,20 @@ while (attempted.size < Math.max(1, accountCount)) {
24322444
runtimeMetrics.lastErrorCategory = null;
24332445
return successResponse;
24342446
}
2435-
if (restartAccountTraversalWithFallback) {
2436-
break;
2437-
}
2438-
}
2447+
if (restartAccountTraversalWithFallback) {
2448+
break;
2449+
}
2450+
if (restartAccountTraversalAfterWorkspaceDeactivation) {
2451+
break;
2452+
}
2453+
}
24392454

2440-
if (restartAccountTraversalWithFallback) {
2441-
continue;
2442-
}
2455+
if (restartAccountTraversalWithFallback) {
2456+
continue;
2457+
}
2458+
if (restartAccountTraversalAfterWorkspaceDeactivation) {
2459+
continue;
2460+
}
24432461

24442462
const waitMs = accountManager.getMinWaitTimeForFamily(modelFamily, model);
24452463
const count = accountManager.getAccountCount();

‎test/index.test.ts‎

Lines changed: 137 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -3089,7 +3089,7 @@ describe("OpenAIOAuthPlugin fetch handler", () => {
30893089
expect(response.status).toBe(200);
30903090
});
30913091

3092-
it("removes only the deactivated workspace and fails over to a healthy sibling workspace", async () => {
3092+
it("removes all entries sharing the deactivated refresh token and fails over to a healthy account", async () => {
30933093
const fetchHelpers = await import("../lib/request/fetch-helpers.js");
30943094
const storageModule = await import("../lib/storage.js");
30953095
const accountsModule = await import("../lib/accounts.js");
@@ -3104,17 +3104,26 @@ describe("OpenAIOAuthPlugin fetch handler", () => {
31043104
email: "same@example.com",
31053105
refreshToken: "shared-refresh",
31063106
};
3107-
const liveWorkspace = {
3107+
const duplicateWorkspace = {
31083108
index: 1,
3109+
accountId: "org-dead-duplicate",
3110+
organizationId: "org-dead-duplicate",
3111+
accountIdSource: "org",
3112+
accountLabel: "Duplicate dead workspace",
3113+
email: "same@example.com",
3114+
refreshToken: "shared-refresh",
3115+
};
3116+
const healthyFallback = {
3117+
index: 2,
31093118
accountId: "org-live",
31103119
organizationId: "org-live",
31113120
accountIdSource: "org",
31123121
accountLabel: "Live workspace",
3113-
email: "same@example.com",
3114-
refreshToken: "shared-refresh",
3122+
email: "live@example.com",
3123+
refreshToken: "healthy-refresh",
31153124
};
31163125

3117-
const accounts = [deadWorkspace, liveWorkspace];
3126+
const accounts = [deadWorkspace, duplicateWorkspace, healthyFallback];
31183127
const removeAccount = vi.fn((target: typeof deadWorkspace) => {
31193128
const idx = accounts.findIndex((account) => account.accountId === target.accountId);
31203129
if (idx < 0) return false;
@@ -3124,7 +3133,15 @@ describe("OpenAIOAuthPlugin fetch handler", () => {
31243133
});
31253134
return true;
31263135
});
3127-
const removeAccountsWithSameRefreshToken = vi.fn(() => 0);
3136+
const removeAccountsWithSameRefreshToken = vi.fn((target: typeof deadWorkspace) => {
3137+
const nextAccounts = accounts.filter((account) => account.refreshToken !== target.refreshToken);
3138+
const removedCount = accounts.length - nextAccounts.length;
3139+
accounts.splice(0, accounts.length, ...nextAccounts);
3140+
accounts.forEach((account, index) => {
3141+
account.index = index;
3142+
});
3143+
return removedCount;
3144+
});
31283145

31293146
const customManager = {
31303147
getAccountCount: () => accounts.length,
@@ -3187,7 +3204,10 @@ describe("OpenAIOAuthPlugin fetch handler", () => {
31873204
const headers = new Headers(init?.headers);
31883205
const accessToken = headers.get("x-test-access-token");
31893206
if (accessToken === "access-org-dead") {
3190-
return new Response(JSON.stringify({ error: { code: "deactivated_workspace", message: "workspace dead" } }), {
3207+
return new Response(JSON.stringify({
3208+
error: { code: "deactivated_workspace", message: "workspace dead" },
3209+
detail: { code: "deactivated_workspace", message: "workspace dead" },
3210+
}), {
31913211
status: 402,
31923212
});
31933213
}
@@ -3202,8 +3222,8 @@ describe("OpenAIOAuthPlugin fetch handler", () => {
32023222

32033223
expect(response.status).toBe(200);
32043224
expect(globalThis.fetch).toHaveBeenCalledTimes(2);
3205-
expect(removeAccount).toHaveBeenCalledTimes(1);
3206-
expect(removeAccountsWithSameRefreshToken).not.toHaveBeenCalled();
3225+
expect(removeAccount).not.toHaveBeenCalled();
3226+
expect(removeAccountsWithSameRefreshToken).toHaveBeenCalledTimes(1);
32073227
expect(accounts.map((account) => account.accountId)).toEqual(["org-live"]);
32083228
expect(vi.mocked(storageModule.withFlaggedAccountStorageTransaction)).toHaveBeenCalledTimes(1);
32093229
expect(mockFlaggedStorage.accounts).toEqual(
@@ -3218,6 +3238,110 @@ describe("OpenAIOAuthPlugin fetch handler", () => {
32183238
);
32193239
});
32203240

3241+
it("cools down the deactivated workspace when grouped removal returns zero", async () => {
3242+
const fetchHelpers = await import("../lib/request/fetch-helpers.js");
3243+
const storageModule = await import("../lib/storage.js");
3244+
const accountsModule = await import("../lib/accounts.js");
3245+
const { AccountManager } = accountsModule;
3246+
3247+
const deadWorkspace = {
3248+
index: 0,
3249+
accountId: "org-dead",
3250+
organizationId: "org-dead",
3251+
accountIdSource: "org",
3252+
accountLabel: "Dead workspace",
3253+
email: "same@example.com",
3254+
refreshToken: "shared-refresh",
3255+
};
3256+
3257+
const markAccountCoolingDown = vi.fn();
3258+
const saveToDiskDebounced = vi.fn();
3259+
const removeAccountsWithSameRefreshToken = vi.fn(() => 0);
3260+
const customManager = {
3261+
getAccountCount: () => 1,
3262+
getCurrentOrNextForFamilyHybrid: () => deadWorkspace,
3263+
getSelectionExplainability: () => [
3264+
{
3265+
index: 0,
3266+
enabled: true,
3267+
isCurrentForFamily: true,
3268+
eligible: true,
3269+
reasons: ["eligible"],
3270+
healthScore: 100,
3271+
tokensAvailable: 50,
3272+
lastUsed: Date.now(),
3273+
},
3274+
],
3275+
toAuthDetails: () => ({
3276+
type: "oauth" as const,
3277+
access: "access-org-dead",
3278+
refresh: deadWorkspace.refreshToken,
3279+
expires: Date.now() + 60_000,
3280+
}),
3281+
hasRefreshToken: () => true,
3282+
saveToDiskDebounced,
3283+
updateFromAuth: vi.fn(),
3284+
clearAuthFailures: vi.fn(),
3285+
incrementAuthFailures: vi.fn(() => 1),
3286+
markAccountCoolingDown,
3287+
markRateLimitedWithReason: vi.fn(),
3288+
recordRateLimit: vi.fn(),
3289+
consumeToken: vi.fn(() => true),
3290+
refundToken: vi.fn(),
3291+
markSwitched: vi.fn(),
3292+
removeAccount: vi.fn(() => false),
3293+
removeAccountsWithSameRefreshToken,
3294+
recordFailure: vi.fn(),
3295+
recordSuccess: vi.fn(),
3296+
getMinWaitTimeForFamily: vi.fn(() => 0),
3297+
shouldShowAccountToast: vi.fn(() => false),
3298+
markToastShown: vi.fn(),
3299+
setActiveIndex: vi.fn(() => deadWorkspace),
3300+
getAccountsSnapshot: vi.fn(() => [deadWorkspace]),
3301+
};
3302+
vi.spyOn(AccountManager, "loadFromDisk").mockResolvedValueOnce(customManager as never);
3303+
vi.mocked(accountsModule.extractAccountId).mockReturnValue("org-dead");
3304+
vi.mocked(fetchHelpers.createCodexHeaders).mockImplementation(
3305+
(_init, _accountId, accessToken) =>
3306+
new Headers({ "x-test-access-token": String(accessToken) }),
3307+
);
3308+
vi.mocked(fetchHelpers.handleErrorResponse).mockImplementation(async (response) => {
3309+
const errorBody = await response.clone().json().catch(() => ({}));
3310+
return { response, rateLimit: undefined, errorBody };
3311+
});
3312+
3313+
globalThis.fetch = vi.fn(async () =>
3314+
new Response(JSON.stringify({
3315+
error: { code: "deactivated_workspace", message: "workspace dead" },
3316+
detail: { code: "deactivated_workspace", message: "workspace dead" },
3317+
}), {
3318+
status: 402,
3319+
}),
3320+
);
3321+
3322+
const { sdk } = await setupPlugin();
3323+
const response = await sdk.fetch!("https://api.openai.com/v1/chat", {
3324+
method: "POST",
3325+
body: JSON.stringify({ model: "gpt-5.1" }),
3326+
});
3327+
const body = await response.json();
3328+
3329+
expect(response.status).toBe(503);
3330+
expect(globalThis.fetch).toHaveBeenCalledTimes(1);
3331+
expect(removeAccountsWithSameRefreshToken).toHaveBeenCalledTimes(1);
3332+
expect(markAccountCoolingDown).toHaveBeenCalledWith(
3333+
deadWorkspace,
3334+
expect.any(Number),
3335+
"auth-failure",
3336+
);
3337+
expect(saveToDiskDebounced).toHaveBeenCalledTimes(1);
3338+
expect(body).toEqual({
3339+
error: {
3340+
message: "All 1 account(s) failed (server errors or auth issues). Check account health with `codex-health`.",
3341+
},
3342+
});
3343+
});
3344+
32213345
it("handles empty body in request", async () => {
32223346
globalThis.fetch = vi.fn().mockResolvedValue(
32233347
new Response(JSON.stringify({ content: "test" }), { status: 200 }),
@@ -4267,7 +4391,10 @@ describe("OpenAIOAuthPlugin persistAccountPool", () => {
42674391
const headers = new Headers(init?.headers);
42684392
const accessToken = headers.get("x-test-access-token");
42694393
if (accessToken === "access-dead") {
4270-
return new Response(JSON.stringify({ error: { code: "deactivated_workspace", message: "workspace dead" } }), {
4394+
return new Response(JSON.stringify({
4395+
error: { code: "deactivated_workspace", message: "workspace dead" },
4396+
detail: { code: "deactivated_workspace", message: "workspace dead" },
4397+
}), {
42714398
status: 402,
42724399
headers: { "content-type": "application/json" },
42734400
});

0 commit comments

Comments
 (0)