Skip to content

Commit 6ccd665

Browse files
ndycodendycode
authored
Merge pull request #94 from ndycode/refactor/login-finalization
refactor: unify login finalization across auth flows
2 parents 77edc30 + 3c09c9e commit 6ccd665

3 files changed

Lines changed: 307 additions & 58 deletions

File tree

‎index.ts‎

Lines changed: 41 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,10 @@ import {
3838
createDeviceCodeSession,
3939
} from "./lib/auth/device-code.js";
4040
import {
41+
applyAccountSelectionFallbacks,
42+
persistResolvedAccountSelection,
4143
persistAccountPool,
44+
resolveAndPersistAccountSelection,
4245
resolveAccountSelection,
4346
type AccountSelectionResult,
4447
type TokenSuccessWithAccount,
@@ -617,7 +620,7 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
617620
pkce: { verifier: string },
618621
url: string,
619622
expectedState: string,
620-
onSuccess?: (selection: AccountSelectionResult) => Promise<void>,
623+
replaceAll: boolean,
621624
) => ({
622625
url,
623626
method: "code" as const,
@@ -657,10 +660,10 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => {
657660
REDIRECT_URI,
658661
);
659662
if (tokens?.type === "success") {
660-
const resolved = resolveAccountSelection(tokens);
661-
if (onSuccess) {
662-
await onSuccess(resolved);
663-
}
663+
const resolved = await resolveAndPersistAccountSelection(tokens, {
664+
persistSelections: persistAuthenticatedSelections,
665+
replaceAll,
666+
});
664667
return resolved.primary;
665668
}
666669
return tokens?.type === "failed"
@@ -3095,24 +3098,21 @@ while (attempted.size < Math.max(1, accountCount)) {
30953098
typeof cached.refreshToken === "string" && cached.refreshToken.trim()
30963099
? cached.refreshToken.trim()
30973100
: flagged.refreshToken;
3098-
const resolved = resolveAccountSelection({
3101+
const resolved = applyAccountSelectionFallbacks(
3102+
resolveAccountSelection({
30993103
type: "success",
31003104
access: cached.accessToken,
31013105
refresh: refreshToken,
31023106
expires: cached.expiresAt,
3103-
multiAccount: true,
3104-
});
3105-
if (!resolved.primary.accountIdOverride && flagged.accountId) {
3106-
resolved.primary.accountIdOverride = flagged.accountId;
3107-
resolved.primary.accountIdSource = flagged.accountIdSource ?? "manual";
3108-
resolved.variantsForPersistence = [resolved.primary];
3109-
}
3110-
if (!resolved.primary.organizationIdOverride && flagged.organizationId) {
3111-
resolved.primary.organizationIdOverride = flagged.organizationId;
3112-
}
3113-
if (!resolved.primary.accountLabel && flagged.accountLabel) {
3114-
resolved.primary.accountLabel = flagged.accountLabel;
3115-
}
3107+
multiAccount: true,
3108+
}),
3109+
{
3110+
accountIdOverride: flagged.accountId,
3111+
accountIdSource: flagged.accountIdSource ?? "manual",
3112+
organizationIdOverride: flagged.organizationId,
3113+
accountLabel: flagged.accountLabel,
3114+
},
3115+
);
31163116
restored.push(...resolved.variantsForPersistence);
31173117
console.log(
31183118
`[${i + 1}/${flaggedStorage.accounts.length}] ${label}: RESTORED (Codex CLI cache)`,
@@ -3129,18 +3129,15 @@ while (attempted.size < Math.max(1, accountCount)) {
31293129
continue;
31303130
}
31313131

3132-
const resolved = resolveAccountSelection(refreshResult);
3133-
if (!resolved.primary.accountIdOverride && flagged.accountId) {
3134-
resolved.primary.accountIdOverride = flagged.accountId;
3135-
resolved.primary.accountIdSource = flagged.accountIdSource ?? "manual";
3136-
resolved.variantsForPersistence = [resolved.primary];
3137-
}
3138-
if (!resolved.primary.organizationIdOverride && flagged.organizationId) {
3139-
resolved.primary.organizationIdOverride = flagged.organizationId;
3140-
}
3141-
if (!resolved.primary.accountLabel && flagged.accountLabel) {
3142-
resolved.primary.accountLabel = flagged.accountLabel;
3143-
}
3132+
const resolved = applyAccountSelectionFallbacks(
3133+
resolveAccountSelection(refreshResult),
3134+
{
3135+
accountIdOverride: flagged.accountId,
3136+
accountIdSource: flagged.accountIdSource ?? "manual",
3137+
organizationIdOverride: flagged.organizationId,
3138+
accountLabel: flagged.accountLabel,
3139+
},
3140+
);
31443141
restored.push(...resolved.variantsForPersistence);
31453142
console.log(`[${i + 1}/${flaggedStorage.accounts.length}] ${label}: RESTORED`);
31463143
} catch (error) {
@@ -3342,12 +3339,7 @@ while (attempted.size < Math.max(1, accountCount)) {
33423339

33433340
if (useManualMode) {
33443341
const { pkce, state, url } = await createAuthorizationFlow();
3345-
return buildManualOAuthFlow(pkce, url, state, async (selection) => {
3346-
await persistAuthenticatedSelections(
3347-
selection.variantsForPersistence,
3348-
startFresh,
3349-
);
3350-
});
3342+
return buildManualOAuthFlow(pkce, url, state, startFresh);
33513343
}
33523344

33533345
const explicitCountProvided =
@@ -3358,12 +3350,11 @@ while (attempted.size < Math.max(1, accountCount)) {
33583350
const forceNewLogin = accounts.length > 0 || refreshAccountIndex !== undefined;
33593351
const result = await runOAuthFlow(forceNewLogin);
33603352

3353+
let selection: AccountSelectionResult | null = null;
33613354
let resolved: TokenSuccessWithAccount | null = null;
3362-
let variantsForPersistence: TokenSuccessWithAccount[] = [];
33633355
if (result.type === "success") {
3364-
const selection = resolveAccountSelection(result);
3356+
selection = resolveAccountSelection(result);
33653357
resolved = selection.primary;
3366-
variantsForPersistence = selection.variantsForPersistence;
33673358
const email = extractAccountEmail(resolved.access, resolved.idToken);
33683359
const accountId = resolved.accountIdOverride ?? extractAccountId(resolved.access);
33693360
const label = resolved.accountLabel ?? email ?? accountId ?? "Unknown account";
@@ -3394,20 +3385,18 @@ while (attempted.size < Math.max(1, accountCount)) {
33943385
break;
33953386
}
33963387

3397-
if (!resolved) {
3388+
if (!selection || !resolved) {
33983389
continue;
33993390
}
34003391

34013392
accounts.push(resolved);
34023393
await showToast(`Account ${accounts.length} authenticated`, "success");
34033394

34043395
const isFirstAccount = accounts.length === 1;
3405-
const entriesToPersist =
3406-
variantsForPersistence.length > 0 ? variantsForPersistence : [resolved];
3407-
await persistAuthenticatedSelections(
3408-
entriesToPersist,
3409-
isFirstAccount && startFresh,
3410-
);
3396+
await persistResolvedAccountSelection(selection, {
3397+
persistSelections: persistAuthenticatedSelections,
3398+
replaceAll: isFirstAccount && startFresh,
3399+
});
34113400

34123401
if (accounts.length >= ACCOUNT_LIMITS.MAX_ACCOUNTS) {
34133402
break;
@@ -3490,11 +3479,10 @@ while (attempted.size < Math.max(1, accountCount)) {
34903479
return result;
34913480
}
34923481

3493-
const selection = resolveAccountSelection(result);
3494-
await persistAuthenticatedSelections(
3495-
selection.variantsForPersistence,
3496-
false,
3497-
);
3482+
const selection = await resolveAndPersistAccountSelection(result, {
3483+
persistSelections: persistAuthenticatedSelections,
3484+
replaceAll: false,
3485+
});
34983486
return selection.primary;
34993487
},
35003488
};
@@ -3513,12 +3501,7 @@ while (attempted.size < Math.max(1, accountCount)) {
35133501
setStoragePath(manualPerProjectAccounts ? process.cwd() : null);
35143502

35153503
const { pkce, state, url } = await createAuthorizationFlow();
3516-
return buildManualOAuthFlow(pkce, url, state, async (selection) => {
3517-
await persistAuthenticatedSelections(
3518-
selection.variantsForPersistence,
3519-
false,
3520-
);
3521-
});
3504+
return buildManualOAuthFlow(pkce, url, state, false);
35223505
},
35233506
},
35243507
],

‎lib/auth/login-runner.ts‎

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,19 @@ export type AccountSelectionResult = {
2424
variantsForPersistence: TokenSuccessWithAccount[];
2525
};
2626

27+
export type PersistAccountSelections = (
28+
results: TokenSuccessWithAccount[],
29+
replaceAll: boolean,
30+
) => Promise<void>;
31+
32+
export type AccountSelectionFallbacks = Pick<
33+
TokenSuccessWithAccount,
34+
"accountIdOverride" | "accountIdSource" | "organizationIdOverride" | "accountLabel"
35+
>;
36+
37+
const PERSIST_AUTHENTICATED_SELECTIONS_ERROR =
38+
"Failed to persist authenticated account selections.";
39+
2740
const createSelectionVariant = (
2841
tokens: TokenSuccess,
2942
candidate: {
@@ -104,6 +117,106 @@ export function resolveAccountSelection(tokens: TokenSuccess): AccountSelectionR
104117
};
105118
}
106119

120+
export function applyAccountSelectionFallbacks(
121+
selection: AccountSelectionResult,
122+
fallbacks: AccountSelectionFallbacks,
123+
): AccountSelectionResult {
124+
const primary = { ...selection.primary };
125+
const primaryAccountId = selection.primary.accountIdOverride?.trim();
126+
const primaryOrganizationId = selection.primary.organizationIdOverride?.trim();
127+
const shouldReusePrimaryVariant =
128+
(primaryAccountId?.length ?? 0) > 0 || (primaryOrganizationId?.length ?? 0) > 0;
129+
// Callers may deep-clone `selection.primary`, so reuse the updated primary by
130+
// persisted account identity instead of relying on object aliasing.
131+
let variantsForPersistence = selection.variantsForPersistence.map((variant) =>
132+
variant === selection.primary ||
133+
(shouldReusePrimaryVariant &&
134+
(variant.accountIdOverride?.trim() ?? "") === (primaryAccountId ?? "") &&
135+
(variant.organizationIdOverride?.trim() ?? "") === (primaryOrganizationId ?? ""))
136+
? primary
137+
: { ...variant },
138+
);
139+
140+
const accountIdOverride = fallbacks.accountIdOverride?.trim();
141+
if (!primary.accountIdOverride && accountIdOverride) {
142+
primary.accountIdOverride = accountIdOverride;
143+
primary.accountIdSource = fallbacks.accountIdSource ?? "manual";
144+
variantsForPersistence = [primary];
145+
}
146+
147+
const organizationIdOverride = fallbacks.organizationIdOverride?.trim();
148+
if (!primary.organizationIdOverride && organizationIdOverride) {
149+
primary.organizationIdOverride = organizationIdOverride;
150+
}
151+
152+
const accountLabel = fallbacks.accountLabel?.trim();
153+
if (!primary.accountLabel && accountLabel) {
154+
primary.accountLabel = accountLabel;
155+
}
156+
157+
return {
158+
primary,
159+
variantsForPersistence,
160+
};
161+
}
162+
163+
/**
164+
* Persists the already-resolved selection through the caller's
165+
* `persistSelections` callback. Windows filesystem safety remains delegated to
166+
* that callback, so callers should use `persistAccountPool` or
167+
* `withAccountStorageTransaction` to keep the rename retry and serialized
168+
* read-modify-write behavior covered by `test/login-runner.test.ts`.
169+
* Callback failures are rethrown with a redacted message so callers can log the
170+
* wrapper safely without leaking token-file paths or account identifiers.
171+
*/
172+
export async function persistResolvedAccountSelection(
173+
selection: AccountSelectionResult,
174+
options?: {
175+
persistSelections?: PersistAccountSelections;
176+
replaceAll?: boolean;
177+
},
178+
): Promise<AccountSelectionResult> {
179+
if (!options?.persistSelections) {
180+
return selection;
181+
}
182+
183+
try {
184+
await options.persistSelections(
185+
selection.variantsForPersistence,
186+
options.replaceAll ?? false,
187+
);
188+
} catch (error) {
189+
throw new Error(PERSIST_AUTHENTICATED_SELECTIONS_ERROR, {
190+
cause: error,
191+
});
192+
}
193+
return selection;
194+
}
195+
196+
/**
197+
* Finalizes a resolved selection by delegating persistence to the caller's
198+
* `persistSelections` callback. Windows lock-retry and read-modify-write
199+
* serialization remain the callback's responsibility, so callers should route
200+
* through `persistAccountPool` or `withAccountStorageTransaction` to preserve
201+
* the guarantees covered by `test/login-runner.test.ts`.
202+
* Persistence callback failures are redacted inside
203+
* `persistResolvedAccountSelection()` before they propagate back to callers.
204+
*/
205+
export async function resolveAndPersistAccountSelection(
206+
tokens: TokenSuccess,
207+
options?: {
208+
fallbacks?: AccountSelectionFallbacks;
209+
persistSelections?: PersistAccountSelections;
210+
replaceAll?: boolean;
211+
},
212+
): Promise<AccountSelectionResult> {
213+
let selection = resolveAccountSelection(tokens);
214+
if (options?.fallbacks) {
215+
selection = applyAccountSelectionFallbacks(selection, options.fallbacks);
216+
}
217+
return persistResolvedAccountSelection(selection, options);
218+
}
219+
107220
/**
108221
* Persists login results through the shared storage transaction so overlapping
109222
* login retries serialize their read-modify-write cycle instead of racing stale

0 commit comments

Comments
 (0)