Skip to content

Commit 9e0e10b

Browse files
ndycodendycode
committed
refactor: unify login finalization across auth flows
1 parent 8336f16 commit 9e0e10b

3 files changed

Lines changed: 173 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: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,16 @@ 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+
2737
const createSelectionVariant = (
2838
tokens: TokenSuccess,
2939
candidate: {
@@ -104,6 +114,71 @@ export function resolveAccountSelection(tokens: TokenSuccess): AccountSelectionR
104114
};
105115
}
106116

117+
export function applyAccountSelectionFallbacks(
118+
selection: AccountSelectionResult,
119+
fallbacks: AccountSelectionFallbacks,
120+
): AccountSelectionResult {
121+
const primary = { ...selection.primary };
122+
let variantsForPersistence = selection.variantsForPersistence.map((variant) =>
123+
variant === selection.primary ? primary : { ...variant },
124+
);
125+
126+
const accountIdOverride = fallbacks.accountIdOverride?.trim();
127+
if (!primary.accountIdOverride && accountIdOverride) {
128+
primary.accountIdOverride = accountIdOverride;
129+
primary.accountIdSource = fallbacks.accountIdSource ?? "manual";
130+
variantsForPersistence = [primary];
131+
}
132+
133+
const organizationIdOverride = fallbacks.organizationIdOverride?.trim();
134+
if (!primary.organizationIdOverride && organizationIdOverride) {
135+
primary.organizationIdOverride = organizationIdOverride;
136+
}
137+
138+
const accountLabel = fallbacks.accountLabel?.trim();
139+
if (!primary.accountLabel && accountLabel) {
140+
primary.accountLabel = accountLabel;
141+
}
142+
143+
return {
144+
primary,
145+
variantsForPersistence,
146+
};
147+
}
148+
149+
export async function persistResolvedAccountSelection(
150+
selection: AccountSelectionResult,
151+
options?: {
152+
persistSelections?: PersistAccountSelections;
153+
replaceAll?: boolean;
154+
},
155+
): Promise<AccountSelectionResult> {
156+
if (!options?.persistSelections) {
157+
return selection;
158+
}
159+
160+
await options.persistSelections(
161+
selection.variantsForPersistence,
162+
options.replaceAll ?? false,
163+
);
164+
return selection;
165+
}
166+
167+
export async function resolveAndPersistAccountSelection(
168+
tokens: TokenSuccess,
169+
options?: {
170+
fallbacks?: AccountSelectionFallbacks;
171+
persistSelections?: PersistAccountSelections;
172+
replaceAll?: boolean;
173+
},
174+
): Promise<AccountSelectionResult> {
175+
let selection = resolveAccountSelection(tokens);
176+
if (options?.fallbacks) {
177+
selection = applyAccountSelectionFallbacks(selection, options.fallbacks);
178+
}
179+
return persistResolvedAccountSelection(selection, options);
180+
}
181+
107182
/**
108183
* Persists login results through the shared storage transaction so overlapping
109184
* login retries serialize their read-modify-write cycle instead of racing stale

‎test/login-runner.test.ts‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,10 @@ import { tmpdir } from "node:os";
33
import { join } from "node:path";
44
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
55
import {
6+
applyAccountSelectionFallbacks,
67
persistAccountPool,
8+
resolveAccountSelection,
9+
resolveAndPersistAccountSelection,
710
type TokenSuccessWithAccount,
811
} from "../lib/auth/login-runner.js";
912
import { loadAccounts, setStoragePathDirect } from "../lib/storage.js";
@@ -94,3 +97,57 @@ describe("login-runner persistAccountPool", () => {
9497
}
9598
});
9699
});
100+
101+
describe("login-runner selection finalization", () => {
102+
it("applies flagged-account fallbacks without overwriting resolved ids", () => {
103+
const selection = resolveAccountSelection({
104+
type: "success",
105+
access: "access-token",
106+
refresh: "refresh-token",
107+
expires: Date.now() + 60_000,
108+
idToken: "id-token",
109+
accountIdOverride: "resolved-account",
110+
organizationIdOverride: "resolved-org",
111+
accountLabel: "Resolved label",
112+
});
113+
114+
const updated = applyAccountSelectionFallbacks(selection, {
115+
accountIdOverride: "flagged-account",
116+
accountIdSource: "manual",
117+
organizationIdOverride: "flagged-org",
118+
accountLabel: "Flagged label",
119+
});
120+
121+
expect(updated.primary.accountIdOverride).toBe("resolved-account");
122+
expect(updated.primary.organizationIdOverride).toBe("resolved-org");
123+
expect(updated.primary.accountLabel).toBe("Resolved label");
124+
expect(updated.variantsForPersistence).toHaveLength(selection.variantsForPersistence.length);
125+
});
126+
127+
it("resolves and persists the selected variants through the shared callback", async () => {
128+
const persistSelections = vi.fn(async () => {});
129+
const result = await resolveAndPersistAccountSelection(
130+
{
131+
type: "success",
132+
access: "persist-access",
133+
refresh: "persist-refresh",
134+
expires: Date.now() + 60_000,
135+
idToken: "persist-id",
136+
},
137+
{
138+
persistSelections,
139+
replaceAll: true,
140+
fallbacks: {
141+
accountIdOverride: "flagged-account",
142+
accountIdSource: "manual",
143+
accountLabel: "Flagged label",
144+
},
145+
},
146+
);
147+
148+
expect(result.primary.accountIdOverride).toBe("flagged-account");
149+
expect(result.primary.accountLabel).toBe("Flagged label");
150+
expect(persistSelections).toHaveBeenCalledTimes(1);
151+
expect(persistSelections).toHaveBeenCalledWith(result.variantsForPersistence, true);
152+
});
153+
});

0 commit comments

Comments
 (0)