Skip to content

Commit f51f7f0

Browse files
TheGreatAxiosTheGreatAxios
committed
Defer OAuth credential persistence until setup validation
1 parent af0db0f commit f51f7f0

19 files changed

Lines changed: 453 additions & 134 deletions

‎src/auth/callback-page.test.ts‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import {
77
PRODUCT_SITE_URL,
88
} from "../branding.js";
99
import { callbackPageHtml, humanizeIdentifier } from "./callback-page.js";
10+
import { authorizationDoneHtml } from "./oauth/callback-server.js";
1011

1112
describe("humanizeIdentifier", () => {
1213
test("machine identifiers lose their separators and lead with a capital", () => {
@@ -28,6 +29,13 @@ describe("callbackPageHtml", () => {
2829
expect(html).not.toContain("access_denied");
2930
});
3031

32+
test("provider authorization waits for native setup before claiming connection", () => {
33+
const html = authorizationDoneHtml("Codex");
34+
expect(html).toContain("Codex authorization received");
35+
expect(html).toContain("finish setup");
36+
expect(html).not.toContain("connected successfully");
37+
});
38+
3139
test("failure names the server and the humanized reason", () => {
3240
const html = callbackPageHtml({ subject: "granola", error: "access_denied" });
3341
expect(html).toContain("Granola failed to connect");

‎src/auth/callback-page.ts‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -262,6 +262,8 @@ export interface CallbackPage {
262262
readonly subject?: string;
263263
/** Why it failed. Omit for the success page. */
264264
readonly error?: string;
265+
/** Authorization succeeded, but the native setup flow still has work to do. */
266+
readonly pendingSetup?: boolean;
265267
}
266268

267269
/**
@@ -276,19 +278,26 @@ export function callbackPageHtml(page: CallbackPage = {}): string {
276278
const failed = page.error !== undefined;
277279
const subject =
278280
page.subject === undefined ? undefined : escapeHtml(humanizeIdentifier(page.subject));
281+
const pendingSetup = !failed && page.pendingSetup === true;
279282
const tone = failed ? "var(--accent)" : "var(--ok)";
280-
const label = failed ? "not connected" : "connected";
283+
const label = failed ? "not connected" : pendingSetup ? "authorization received" : "connected";
281284
const heading = failed
282285
? subject === undefined
283286
? "Authorization did not complete"
284287
: `${subject} failed to connect`
285-
: subject === undefined
286-
? "Authorization complete"
287-
: `${subject} connected successfully`;
288+
: pendingSetup
289+
? subject === undefined
290+
? "Authorization received"
291+
: `${subject} authorization received`
292+
: subject === undefined
293+
? "Authorization complete"
294+
: `${subject} connected successfully`;
288295
const reason = escapeHtml(humanizeIdentifier(page.error ?? ""));
289296
const body = failed
290297
? `${reason}. Close this tab and try again from ${PRODUCT_NAME}.`
291-
: `You can close this tab and return to ${PRODUCT_NAME}.`;
298+
: pendingSetup
299+
? `Return to ${PRODUCT_NAME} to finish setup.`
300+
: `You can close this tab and return to ${PRODUCT_NAME}.`;
292301
return [
293302
"<!doctype html>",
294303
'<html lang="en">',

‎src/auth/codex/login.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,11 +7,11 @@ import {
77
import { CODEX_BASE_URL, CODEX_DEFAULT_MODELS } from "./constants.js";
88
import { startCodexCallbackServer } from "./callback-server.js";
99
import { buildAuthorizeUrl, exchangeCode } from "./oauth.js";
10-
import { saveCodexProfile } from "./store.js";
10+
import { saveCodexProfile, type CodexTokens } from "./store.js";
1111

1212
export { openInBrowser };
1313

14-
export type CodexLoginHandle = OAuthLoginHandle;
14+
export type CodexLoginHandle = OAuthLoginHandle<CodexTokens>;
1515
export type StartCodexLoginOptions = StartOAuthLoginOptions;
1616

1717
// Drive the loopback PKCE login for a Codex profile.

‎src/auth/codex/session.ts‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,3 +52,13 @@ const session = createTokenSession<CodexTokens, CodexAccess>({
5252

5353
export const isCodexTokenExpired = session.isExpired;
5454
export const getValidCodexToken = session.getValidToken;
55+
56+
export async function refreshStagedCodexTokens(
57+
tokens: CodexTokens,
58+
now: number = Date.now(),
59+
): Promise<CodexTokens> {
60+
if (!isCodexTokenExpired(tokens, now)) return tokens;
61+
const refreshed = await refreshTokens(tokens.refresh, now);
62+
Object.assign(tokens, refreshed);
63+
return tokens;
64+
}

‎src/auth/codex/usage.ts‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -73,17 +73,23 @@ function parseUsage(payload: unknown): CodexUsage {
7373
};
7474
}
7575

76-
export async function codexAuthHeaders(profileName: string): Promise<Record<string, string>> {
77-
const { access, accountId } = await getValidCodexToken(profileName);
76+
export function codexAuthHeadersForToken(token: {
77+
readonly access: string;
78+
readonly accountId?: string | undefined;
79+
}): Record<string, string> {
7880
const headers: Record<string, string> = {
79-
authorization: `Bearer ${access}`,
81+
authorization: `Bearer ${token.access}`,
8082
originator: CODEX_AUTHORIZE_EXTRA_PARAMS["originator"] ?? "codex_cli_rs",
8183
"user-agent": `${COMMAND_NAME} (codex_cli_rs/${CODEX_CLIENT_VERSION})`,
8284
};
83-
if (accountId !== undefined) headers["chatgpt-account-id"] = accountId;
85+
if (token.accountId !== undefined) headers["chatgpt-account-id"] = token.accountId;
8486
return headers;
8587
}
8688

89+
export async function codexAuthHeaders(profileName: string): Promise<Record<string, string>> {
90+
return codexAuthHeadersForToken(await getValidCodexToken(profileName));
91+
}
92+
8793
// Fetch the live usage/quota snapshot for a Codex profile.
8894
export async function fetchCodexUsage(profileName: string): Promise<CodexUsage> {
8995
const res = await fetch(`${CODEX_BASE_URL}${CODEX_USAGE_PATH}`, {

‎src/auth/oauth-scope-check.test.ts‎

Lines changed: 102 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -1,49 +1,79 @@
11
import { afterEach, describe, expect, test } from "bun:test";
2-
import { withMockedModule } from "../../tests/helpers/mock-module.js";
3-
4-
// getValidCodexToken/getValidXaiToken hit the real home-level auth store and
5-
// refresh endpoints; stub the session layer so this test only exercises the
6-
// scope probe's own HTTP call and status classification. Other suites
7-
// (tests/unit/codex-session.test.ts) import the real modules directly, so the
8-
// mocks must be torn down after this file's tests run rather than leaking
9-
// into the rest of the bun test process.
10-
await withMockedModule(
11-
import.meta.resolve("./codex/session.js"),
12-
(real: typeof import("./codex/session.js")) => ({
13-
...real,
14-
getValidCodexToken: async () => ({ access: "codex-token", accountId: "acct-1" }),
15-
}),
16-
);
17-
await withMockedModule(
18-
import.meta.resolve("./xai/session.js"),
19-
(real: typeof import("./xai/session.js")) => ({
20-
...real,
21-
getValidXaiToken: async () => ({ access: "xai-token" }),
22-
}),
23-
);
24-
25-
const { checkOAuthProviderScope } = await import("./oauth-scope-check.js");
2+
3+
import { checkOAuthProviderScope } from "./oauth-scope-check.js";
264

275
const originalFetch = global.fetch;
6+
const codexTokens = {
7+
access: "staged-codex-token",
8+
refresh: "codex-refresh",
9+
expiresAt: Date.now() + 3_600_000,
10+
accountId: "acct-staged",
11+
};
12+
const xaiTokens = {
13+
access: "staged-xai-token",
14+
refresh: "xai-refresh",
15+
expiresAt: Date.now() + 3_600_000,
16+
};
2817

29-
function stubFetch(impl: (url: string) => Response | Promise<Response>): void {
30-
global.fetch = (async (input: RequestInfo | URL) => impl(String(input))) as typeof fetch;
18+
function stubFetch(impl: (url: string, init?: RequestInit) => Response | Promise<Response>): void {
19+
global.fetch = (async (input: RequestInfo | URL, init?: RequestInit) =>
20+
impl(String(input), init)) as typeof fetch;
3121
}
3222

3323
describe("checkOAuthProviderScope", () => {
3424
afterEach(() => {
3525
global.fetch = originalFetch;
3626
});
3727

38-
test("codex: ok when the catalog call succeeds", async () => {
39-
stubFetch(() => new Response(JSON.stringify({ models: ["gpt-5"] }), { status: 200 }));
40-
const result = await checkOAuthProviderScope("codex", "work");
28+
test("codex: builds the probe from staged tokens", async () => {
29+
stubFetch((_url, init) => {
30+
expect(init?.headers).toMatchObject({
31+
authorization: "Bearer staged-codex-token",
32+
"chatgpt-account-id": "acct-staged",
33+
});
34+
return new Response(JSON.stringify({ models: ["gpt-5"] }), { status: 200 });
35+
});
36+
const result = await checkOAuthProviderScope("codex", codexTokens);
4137
expect(result.status).toBe("ok");
4238
});
4339

40+
test("codex: refreshes expired staged tokens before classifying the probe", async () => {
41+
const expired = { ...codexTokens, expiresAt: 0 };
42+
const requests: string[] = [];
43+
stubFetch((url, init) => {
44+
requests.push(url);
45+
if (url.includes("/oauth/token")) {
46+
return new Response(JSON.stringify({ access_token: "refreshed-codex", expires_in: 3600 }), {
47+
status: 200,
48+
headers: { "content-type": "application/json" },
49+
});
50+
}
51+
expect(init?.headers).toMatchObject({
52+
authorization: "Bearer refreshed-codex",
53+
"chatgpt-account-id": "acct-staged",
54+
});
55+
return new Response(JSON.stringify({ models: ["gpt-5"] }), { status: 200 });
56+
});
57+
58+
const result = await checkOAuthProviderScope("codex", expired);
59+
60+
expect(result.status).toBe("ok");
61+
expect(requests).toHaveLength(2);
62+
expect(expired.access).toBe("refreshed-codex");
63+
});
64+
65+
test("codex: reports an expired staged token refresh failure as unavailable", async () => {
66+
const expired = { ...codexTokens, expiresAt: 0 };
67+
stubFetch(() => new Response("refresh rejected", { status: 401 }));
68+
69+
const result = await checkOAuthProviderScope("codex", expired);
70+
71+
expect(result.status).toBe("unavailable");
72+
});
73+
4474
test("codex: insufficient-scope on a definitive 403", async () => {
4575
stubFetch(() => new Response("forbidden", { status: 403 }));
46-
const result = await checkOAuthProviderScope("codex", "work");
76+
const result = await checkOAuthProviderScope("codex", codexTokens);
4777
expect(result.status).toBe("insufficient-scope");
4878
if (result.status === "insufficient-scope") {
4979
expect(result.message).toMatch(/reconnect/i);
@@ -54,41 +84,75 @@ describe("checkOAuthProviderScope", () => {
5484

5585
test("codex: insufficient-scope on a definitive 401", async () => {
5686
stubFetch(() => new Response("nope", { status: 401 }));
57-
const result = await checkOAuthProviderScope("codex", "work");
87+
const result = await checkOAuthProviderScope("codex", codexTokens);
5888
expect(result.status).toBe("insufficient-scope");
5989
});
6090

6191
test("codex: unavailable on a network failure, not blocked", async () => {
6292
stubFetch(() => {
6393
throw new Error("fetch failed");
6494
});
65-
const result = await checkOAuthProviderScope("codex", "work");
95+
const result = await checkOAuthProviderScope("codex", codexTokens);
6696
expect(result.status).toBe("unavailable");
6797
});
6898

6999
test("codex: unavailable (not scope failure) on a 500", async () => {
70100
stubFetch(() => new Response("boom", { status: 500 }));
71-
const result = await checkOAuthProviderScope("codex", "work");
101+
const result = await checkOAuthProviderScope("codex", codexTokens);
72102
expect(result.status).toBe("unavailable");
73103
});
74104

75-
test("xai: ok when the models call succeeds", async () => {
76-
stubFetch(() => new Response(JSON.stringify({ data: [] }), { status: 200 }));
77-
const result = await checkOAuthProviderScope("xai", "personal");
105+
test("xai: builds the probe from staged tokens", async () => {
106+
stubFetch((_url, init) => {
107+
expect(init?.headers).toMatchObject({ authorization: "Bearer staged-xai-token" });
108+
return new Response(JSON.stringify({ data: [] }), { status: 200 });
109+
});
110+
const result = await checkOAuthProviderScope("xai", xaiTokens);
111+
expect(result.status).toBe("ok");
112+
});
113+
114+
test("xai: refreshes expired staged tokens before classifying the probe", async () => {
115+
const expired = { ...xaiTokens, expiresAt: 0 };
116+
const requests: string[] = [];
117+
stubFetch((url, init) => {
118+
requests.push(url);
119+
if (url.includes("/oauth2/token")) {
120+
return new Response(JSON.stringify({ access_token: "refreshed-xai", expires_in: 3600 }), {
121+
status: 200,
122+
headers: { "content-type": "application/json" },
123+
});
124+
}
125+
expect(init?.headers).toMatchObject({ authorization: "Bearer refreshed-xai" });
126+
return new Response(JSON.stringify({ data: [] }), { status: 200 });
127+
});
128+
129+
const result = await checkOAuthProviderScope("xai", expired);
130+
78131
expect(result.status).toBe("ok");
132+
expect(requests).toHaveLength(2);
133+
expect(expired.access).toBe("refreshed-xai");
134+
});
135+
136+
test("xai: reports an expired staged token refresh failure as unavailable", async () => {
137+
const expired = { ...xaiTokens, expiresAt: 0 };
138+
stubFetch(() => new Response("refresh rejected", { status: 401 }));
139+
140+
const result = await checkOAuthProviderScope("xai", expired);
141+
142+
expect(result.status).toBe("unavailable");
79143
});
80144

81145
test("xai: insufficient-scope on a definitive 403", async () => {
82146
stubFetch(() => new Response("forbidden", { status: 403 }));
83-
const result = await checkOAuthProviderScope("xai", "personal");
147+
const result = await checkOAuthProviderScope("xai", xaiTokens);
84148
expect(result.status).toBe("insufficient-scope");
85149
});
86150

87151
test("xai: unavailable on a timeout-style abort", async () => {
88152
stubFetch(() => {
89153
throw new DOMException("The operation timed out.", "TimeoutError");
90154
});
91-
const result = await checkOAuthProviderScope("xai", "personal");
155+
const result = await checkOAuthProviderScope("xai", xaiTokens);
92156
expect(result.status).toBe("unavailable");
93157
});
94158
});

‎src/auth/oauth-scope-check.ts‎

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,13 @@
1111
// status is inspected to classify the result.
1212

1313
import { CODEX_BASE_URL, CODEX_MODELS_PATH, CODEX_CLIENT_VERSION } from "./codex/constants.js";
14-
import { codexAuthHeaders } from "./codex/usage.js";
14+
import { refreshStagedCodexTokens } from "./codex/session.js";
15+
import type { CodexTokens } from "./codex/store.js";
16+
import { codexAuthHeadersForToken } from "./codex/usage.js";
1517
import { XAI_BASE_URL, XAI_TOKEN_TIMEOUT_MS } from "./xai/constants.js";
16-
import { xaiAuthHeaders } from "./xai/usage.js";
18+
import { refreshStagedXaiTokens } from "./xai/session.js";
19+
import type { XaiTokens } from "./xai/store.js";
20+
import { xaiAuthHeadersForToken } from "./xai/usage.js";
1721

1822
export type OAuthScopeCheckKind = "codex" | "xai";
1923

@@ -53,10 +57,10 @@ function classifyStatus(status: number, providerLabel: string): OAuthScopeCheckR
5357
return unavailable(providerLabel);
5458
}
5559

56-
async function checkCodexScope(profile: string): Promise<OAuthScopeCheckResult> {
60+
async function checkCodexScope(tokens: CodexTokens): Promise<OAuthScopeCheckResult> {
5761
const providerLabel = "Codex";
5862
try {
59-
const headers = await codexAuthHeaders(profile);
63+
const headers = codexAuthHeadersForToken(await refreshStagedCodexTokens(tokens));
6064
const url = `${CODEX_BASE_URL}${CODEX_MODELS_PATH}?client_version=${encodeURIComponent(CODEX_CLIENT_VERSION)}`;
6165
const res = await fetch(url, {
6266
headers,
@@ -69,10 +73,10 @@ async function checkCodexScope(profile: string): Promise<OAuthScopeCheckResult>
6973
}
7074
}
7175

72-
async function checkXaiScope(profile: string): Promise<OAuthScopeCheckResult> {
76+
async function checkXaiScope(tokens: XaiTokens): Promise<OAuthScopeCheckResult> {
7377
const providerLabel = "Grok";
7478
try {
75-
const headers = await xaiAuthHeaders(profile);
79+
const headers = xaiAuthHeadersForToken(await refreshStagedXaiTokens(tokens));
7680
const res = await fetch(`${XAI_BASE_URL}/models`, {
7781
headers,
7882
signal: AbortSignal.timeout(XAI_TOKEN_TIMEOUT_MS),
@@ -89,7 +93,7 @@ async function checkXaiScope(profile: string): Promise<OAuthScopeCheckResult> {
8993
// login result alone.
9094
export async function checkOAuthProviderScope(
9195
kind: OAuthScopeCheckKind,
92-
profile: string,
96+
tokens: CodexTokens | XaiTokens,
9397
): Promise<OAuthScopeCheckResult> {
94-
return kind === "codex" ? checkCodexScope(profile) : checkXaiScope(profile);
98+
return kind === "codex" ? checkCodexScope(tokens) : checkXaiScope(tokens);
9599
}

‎src/auth/oauth/callback-server.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -121,5 +121,5 @@ export async function startCallbackServer(
121121
}
122122

123123
export function authorizationDoneHtml(providerName: string): string {
124-
return callbackPageHtml({ subject: providerName });
124+
return callbackPageHtml({ subject: providerName, pendingSetup: true });
125125
}

‎src/auth/oauth/index.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ export {
3131
startOAuthLogin,
3232
type OAuthLoginDeps,
3333
type OAuthLoginHandle,
34+
type StagedOAuthProfile,
3435
type StartOAuthLoginOptions,
3536
} from "./login.js";
3637
export {

0 commit comments

Comments
 (0)