Skip to content

Commit de8615b

Browse files
ndycodendycode
committed
fix: preserve workspace identity in flagged cleanup
1 parent f667464 commit de8615b

File tree

5 files changed

+322
-62
lines changed

5 files changed

+322
-62
lines changed

‎index.ts‎

Lines changed: 59 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,7 @@ import {
117117
createTimestampedBackupPath,
118118
loadFlaggedAccounts,
119119
saveFlaggedAccounts,
120+
withFlaggedAccountStorageTransaction,
120121
clearFlaggedAccounts,
121122
StorageError,
122123
formatStorageErrorHint,
@@ -189,9 +190,16 @@ function getWorkspaceIdentityKey(account: {
189190
accountId?: string;
190191
refreshToken: string;
191192
}): string {
192-
if (account.organizationId) return `organizationId:${account.organizationId}`;
193-
if (account.accountId) return `accountId:${account.accountId}`;
194-
return `refreshToken:${account.refreshToken}`;
193+
const organizationId = account.organizationId?.trim();
194+
const accountId = account.accountId?.trim();
195+
const refreshToken = account.refreshToken.trim();
196+
if (organizationId) {
197+
return accountId
198+
? `organizationId:${organizationId}|accountId:${accountId}`
199+
: `organizationId:${organizationId}`;
200+
}
201+
if (accountId) return `accountId:${accountId}`;
202+
return `refreshToken:${refreshToken}`;
195203
}
196204

197205
function matchesWorkspaceIdentity(
@@ -205,6 +213,21 @@ function matchesWorkspaceIdentity(
205213
return getWorkspaceIdentityKey(account) === identityKey;
206214
}
207215

216+
function upsertFlaggedAccountRecord(
217+
accounts: FlaggedAccountMetadataV1[],
218+
record: FlaggedAccountMetadataV1,
219+
): void {
220+
const identityKey = getWorkspaceIdentityKey(record);
221+
const existingIndex = accounts.findIndex((flagged) =>
222+
matchesWorkspaceIdentity(flagged, identityKey),
223+
);
224+
if (existingIndex >= 0) {
225+
accounts[existingIndex] = record;
226+
return;
227+
}
228+
accounts.push(record);
229+
}
230+
208231
/**
209232
* OpenAI Codex OAuth authentication plugin for opencode
210233
*
@@ -2382,7 +2405,6 @@ while (attempted.size < Math.max(1, accountCount)) {
23822405

23832406
const workspaceDeactivated = isDeactivatedWorkspaceError(errorBody, response.status);
23842407
if (workspaceDeactivated) {
2385-
const identityKey = getWorkspaceIdentityKey(account);
23862408
const accountLabel = formatAccountLabel(account, account.index);
23872409
accountManager.refundToken(account, modelFamily, model);
23882410
accountManager.recordFailure(account, modelFamily, model);
@@ -2393,22 +2415,20 @@ while (attempted.size < Math.max(1, accountCount)) {
23932415
runtimeMetrics.lastErrorCategory = "workspace-deactivated";
23942416

23952417
try {
2396-
const flaggedStorage = await loadFlaggedAccounts();
23972418
const flaggedRecord: FlaggedAccountMetadataV1 = {
23982419
...account,
23992420
flaggedAt: Date.now(),
24002421
flaggedReason: "workspace-deactivated",
24012422
lastError: "deactivated_workspace",
24022423
};
2403-
const existingIndex = flaggedStorage.accounts.findIndex((flagged) =>
2404-
matchesWorkspaceIdentity(flagged, identityKey),
2405-
);
2406-
if (existingIndex >= 0) {
2407-
flaggedStorage.accounts[existingIndex] = flaggedRecord;
2408-
} else {
2409-
flaggedStorage.accounts.push(flaggedRecord);
2410-
}
2411-
await saveFlaggedAccounts(flaggedStorage);
2424+
await withFlaggedAccountStorageTransaction(async (current, persist) => {
2425+
const nextStorage: typeof current = {
2426+
...current,
2427+
accounts: current.accounts.map((flagged) => ({ ...flagged })),
2428+
};
2429+
upsertFlaggedAccountRecord(nextStorage.accounts, flaggedRecord);
2430+
await persist(nextStorage);
2431+
});
24122432
} catch (flagError) {
24132433
logWarn(
24142434
`Failed to persist deactivated workspace flag for ${accountLabel}: ${flagError instanceof Error ? flagError.message : String(flagError)}`,
@@ -3074,9 +3094,9 @@ while (attempted.size < Math.max(1, accountCount)) {
30743094
return;
30753095
}
30763096

3077-
const flaggedStorage = await loadFlaggedAccounts();
30783097
let storageChanged = false;
30793098
let flaggedChanged = false;
3099+
const flaggedUpdates = new Map<string, FlaggedAccountMetadataV1>();
30803100
const removeFromActive = new Set<string>();
30813101
const total = workingStorage.accounts.length;
30823102
let ok = 0;
@@ -3182,21 +3202,17 @@ while (attempted.size < Math.max(1, accountCount)) {
31823202
refreshResult.message ?? refreshResult.reason ?? "refresh failed";
31833203
console.log(`[${i + 1}/${total}] ${label}: ERROR (${message})`);
31843204
if (deepProbe && isFlaggableFailure(refreshResult)) {
3185-
const existingIndex = flaggedStorage.accounts.findIndex(
3186-
(flagged) => flagged.refreshToken === account.refreshToken,
3187-
);
31883205
const flaggedRecord: FlaggedAccountMetadataV1 = {
31893206
...account,
31903207
flaggedAt: Date.now(),
31913208
flaggedReason: "token-invalid",
31923209
lastError: message,
31933210
};
3194-
if (existingIndex >= 0) {
3195-
flaggedStorage.accounts[existingIndex] = flaggedRecord;
3196-
} else {
3197-
flaggedStorage.accounts.push(flaggedRecord);
3198-
}
3199-
removeFromActive.add(`refreshToken:${account.refreshToken}`);
3211+
flaggedUpdates.set(
3212+
getWorkspaceIdentityKey(flaggedRecord),
3213+
flaggedRecord,
3214+
);
3215+
removeFromActive.add(getWorkspaceIdentityKey(account));
32003216
flaggedChanged = true;
32013217
}
32023218
continue;
@@ -3279,20 +3295,16 @@ while (attempted.size < Math.max(1, accountCount)) {
32793295
errors += 1;
32803296
const message = error instanceof Error ? error.message : String(error);
32813297
if (message.includes("deactivated_workspace")) {
3282-
const existingIndex = flaggedStorage.accounts.findIndex((flagged) =>
3283-
matchesWorkspaceIdentity(flagged, getWorkspaceIdentityKey(account)),
3284-
);
32853298
const flaggedRecord: FlaggedAccountMetadataV1 = {
32863299
...account,
32873300
flaggedAt: Date.now(),
32883301
flaggedReason: "workspace-deactivated",
32893302
lastError: message,
32903303
};
3291-
if (existingIndex >= 0) {
3292-
flaggedStorage.accounts[existingIndex] = flaggedRecord;
3293-
} else {
3294-
flaggedStorage.accounts.push(flaggedRecord);
3295-
}
3304+
flaggedUpdates.set(
3305+
getWorkspaceIdentityKey(flaggedRecord),
3306+
flaggedRecord,
3307+
);
32963308
removeFromActive.add(getWorkspaceIdentityKey(account));
32973309
flaggedChanged = true;
32983310
}
@@ -3320,7 +3332,16 @@ while (attempted.size < Math.max(1, accountCount)) {
33203332
invalidateAccountManagerCache();
33213333
}
33223334
if (flaggedChanged) {
3323-
await saveFlaggedAccounts(flaggedStorage);
3335+
await withFlaggedAccountStorageTransaction(async (current, persist) => {
3336+
const nextStorage: typeof current = {
3337+
...current,
3338+
accounts: current.accounts.map((flagged) => ({ ...flagged })),
3339+
};
3340+
for (const flaggedRecord of flaggedUpdates.values()) {
3341+
upsertFlaggedAccountRecord(nextStorage.accounts, flaggedRecord);
3342+
}
3343+
await persist(nextStorage);
3344+
});
33243345
}
33253346

33263347
console.log("");
@@ -3524,7 +3545,11 @@ while (attempted.size < Math.max(1, accountCount)) {
35243545
await saveFlaggedAccounts({
35253546
version: 1,
35263547
accounts: flaggedStorage.accounts.filter(
3527-
(flagged) => flagged.refreshToken !== target.refreshToken,
3548+
(flagged) =>
3549+
!matchesWorkspaceIdentity(
3550+
flagged,
3551+
getWorkspaceIdentityKey(target),
3552+
),
35283553
),
35293554
});
35303555
invalidateAccountManagerCache();

‎lib/storage.ts‎

Lines changed: 69 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -921,7 +921,28 @@ function normalizeFlaggedStorage(data: unknown): FlaggedAccountStorageV1 {
921921
return { version: 1, accounts: [] };
922922
}
923923

924-
const byRefreshToken = new Map<string, FlaggedAccountMetadataV1>();
924+
const normalizeFlaggedIdentityPart = (value: unknown): string | undefined =>
925+
typeof value === "string" && value.trim().length > 0 ? value.trim() : undefined;
926+
const getFlaggedIdentityKey = (account: {
927+
organizationId?: string;
928+
accountId?: string;
929+
refreshToken: string;
930+
}): string => {
931+
const organizationId = normalizeFlaggedIdentityPart(account.organizationId);
932+
const accountId = normalizeFlaggedIdentityPart(account.accountId);
933+
const refreshToken = normalizeFlaggedIdentityPart(account.refreshToken) ?? "";
934+
if (organizationId) {
935+
return accountId
936+
? `organizationId:${organizationId}|accountId:${accountId}`
937+
: `organizationId:${organizationId}`;
938+
}
939+
if (accountId) {
940+
return `accountId:${accountId}`;
941+
}
942+
return `refreshToken:${refreshToken}`;
943+
};
944+
945+
const byIdentityKey = new Map<string, FlaggedAccountMetadataV1>();
925946
for (const rawAccount of data.accounts) {
926947
if (!isRecord(rawAccount)) continue;
927948
const refreshToken =
@@ -1000,16 +1021,18 @@ function normalizeFlaggedStorage(data: unknown): FlaggedAccountStorageV1 {
10001021
flaggedReason: typeof rawAccount.flaggedReason === "string" ? rawAccount.flaggedReason : undefined,
10011022
lastError: typeof rawAccount.lastError === "string" ? rawAccount.lastError : undefined,
10021023
};
1003-
byRefreshToken.set(refreshToken, normalized);
1024+
byIdentityKey.set(getFlaggedIdentityKey(normalized), normalized);
10041025
}
10051026

10061027
return {
10071028
version: 1,
1008-
accounts: Array.from(byRefreshToken.values()),
1029+
accounts: Array.from(byIdentityKey.values()),
10091030
};
10101031
}
10111032

1012-
export async function loadFlaggedAccounts(): Promise<FlaggedAccountStorageV1> {
1033+
async function loadFlaggedAccountsUnlocked(
1034+
saveUnlocked: (storage: FlaggedAccountStorageV1) => Promise<void>,
1035+
): Promise<FlaggedAccountStorageV1> {
10131036
const path = getFlaggedAccountsPath();
10141037
const empty: FlaggedAccountStorageV1 = { version: 1, accounts: [] };
10151038

@@ -1035,7 +1058,7 @@ export async function loadFlaggedAccounts(): Promise<FlaggedAccountStorageV1> {
10351058
const legacyData = JSON.parse(legacyContent) as unknown;
10361059
const migrated = normalizeFlaggedStorage(legacyData);
10371060
if (migrated.accounts.length > 0) {
1038-
await saveFlaggedAccounts(migrated);
1061+
await saveUnlocked(migrated);
10391062
}
10401063
try {
10411064
await fs.unlink(legacyPath);
@@ -1058,26 +1081,50 @@ export async function loadFlaggedAccounts(): Promise<FlaggedAccountStorageV1> {
10581081
}
10591082
}
10601083

1061-
export async function saveFlaggedAccounts(storage: FlaggedAccountStorageV1): Promise<void> {
1062-
return withStorageLock(async () => {
1063-
const path = getFlaggedAccountsPath();
1064-
const uniqueSuffix = `${Date.now()}.${Math.random().toString(36).slice(2, 8)}`;
1065-
const tempPath = `${path}.${uniqueSuffix}.tmp`;
1084+
export async function loadFlaggedAccounts(): Promise<FlaggedAccountStorageV1> {
1085+
return withStorageLock(async () => loadFlaggedAccountsUnlocked(saveFlaggedAccountsUnlocked));
1086+
}
1087+
1088+
async function saveFlaggedAccountsUnlocked(storage: FlaggedAccountStorageV1): Promise<void> {
1089+
const path = getFlaggedAccountsPath();
1090+
const uniqueSuffix = `${Date.now()}.${Math.random().toString(36).slice(2, 8)}`;
1091+
const tempPath = `${path}.${uniqueSuffix}.tmp`;
10661092

1093+
try {
1094+
await fs.mkdir(dirname(path), { recursive: true });
1095+
const content = JSON.stringify(normalizeFlaggedStorage(storage), null, 2);
1096+
await fs.writeFile(tempPath, content, { encoding: "utf-8", mode: 0o600 });
1097+
await renameWithWindowsRetry(tempPath, path);
1098+
} catch (error) {
10671099
try {
1068-
await fs.mkdir(dirname(path), { recursive: true });
1069-
const content = JSON.stringify(normalizeFlaggedStorage(storage), null, 2);
1070-
await fs.writeFile(tempPath, content, { encoding: "utf-8", mode: 0o600 });
1071-
await renameWithWindowsRetry(tempPath, path);
1072-
} catch (error) {
1073-
try {
1074-
await fs.unlink(tempPath);
1075-
} catch {
1076-
// Ignore cleanup failures.
1077-
}
1078-
log.error("Failed to save flagged account storage", { path, error: String(error) });
1079-
throw error;
1100+
await fs.unlink(tempPath);
1101+
} catch {
1102+
// Ignore cleanup failures.
10801103
}
1104+
log.error("Failed to save flagged account storage", { path, error: String(error) });
1105+
throw error;
1106+
}
1107+
}
1108+
1109+
/**
1110+
* Executes a read-modify-write transaction for flagged account storage under the
1111+
* shared storage lock so concurrent callers cannot lose updates.
1112+
*/
1113+
export async function withFlaggedAccountStorageTransaction<T>(
1114+
handler: (
1115+
current: FlaggedAccountStorageV1,
1116+
persist: (storage: FlaggedAccountStorageV1) => Promise<void>,
1117+
) => Promise<T>,
1118+
): Promise<T> {
1119+
return withStorageLock(async () => {
1120+
const current = await loadFlaggedAccountsUnlocked(saveFlaggedAccountsUnlocked);
1121+
return handler(current, saveFlaggedAccountsUnlocked);
1122+
});
1123+
}
1124+
1125+
export async function saveFlaggedAccounts(storage: FlaggedAccountStorageV1): Promise<void> {
1126+
return withStorageLock(async () => {
1127+
await saveFlaggedAccountsUnlocked(storage);
10811128
});
10821129
}
10831130

‎test/fetch-helpers.test.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,9 +22,9 @@ import type { Auth } from '../lib/types.js';
2222
import { URL_PATHS, OPENAI_HEADERS, OPENAI_HEADER_VALUES, CODEX_BASE_URL } from '../lib/constants.js';
2323

2424
describe('Fetch Helpers Module', () => {
25-
afterEach(() => {
26-
vi.restoreAllMocks();
27-
});
25+
afterEach(() => {
26+
vi.restoreAllMocks();
27+
});
2828

2929
describe('shouldRefreshToken', () => {
3030
it('should return true for non-oauth auth', () => {

0 commit comments

Comments
 (0)