From 01313a5a8e6dcaba2331865cd5999d35f7be6334 Mon Sep 17 00:00:00 2001 From: Tam Nhu Tran Date: Sat, 4 Apr 2026 00:37:04 -0400 Subject: [PATCH] fix(auth): harden delayed oauth token registration --- src/web-server/routes/cliproxy-auth-routes.ts | 114 +++++++-- ...iproxy-auth-routes-manual-callback.test.ts | 235 ++++++++++++++++-- ui/src/hooks/use-cliproxy-auth-flow.ts | 9 +- .../hooks/use-cliproxy-auth-flow.test.tsx | 135 ++++++++++ 4 files changed, 457 insertions(+), 36 deletions(-) diff --git a/src/web-server/routes/cliproxy-auth-routes.ts b/src/web-server/routes/cliproxy-auth-routes.ts index fca6f37c..7341c4d0 100644 --- a/src/web-server/routes/cliproxy-auth-routes.ts +++ b/src/web-server/routes/cliproxy-auth-routes.ts @@ -99,7 +99,12 @@ router.use((req: Request, res: Response, next) => { function pruneExpiredManualAuthState(now = Date.now()): void { for (const [state, pending] of pendingManualAuthState.entries()) { - if (now - pending.createdAt > MANUAL_AUTH_STATE_TTL_MS) { + const authExpired = now - pending.createdAt > MANUAL_AUTH_STATE_TTL_MS; + const withinLocalTokenGrace = + pending.upstreamCompletedAt !== undefined && + now - pending.upstreamCompletedAt < POLLED_AUTH_LOCAL_TOKEN_GRACE_MS; + + if (authExpired && !withinLocalTokenGrace) { pendingManualAuthState.delete(state); } } @@ -199,7 +204,7 @@ function listProviderTokenSnapshots(provider: CLIProxyProvider): ProviderTokenSn function findNewTokenSnapshotForPendingAuth( provider: CLIProxyProvider, - pending: { knownTokenFiles: ProviderTokenSnapshot[] } + pending: { expectedAccountId?: string; knownTokenFiles: ProviderTokenSnapshot[] } ): ProviderTokenSnapshot | null { const knownTokenMtimes = new Map( pending.knownTokenFiles.map((snapshot) => [snapshot.file, snapshot.mtimeMs]) @@ -208,11 +213,32 @@ function findNewTokenSnapshotForPendingAuth( return ( listProviderTokenSnapshots(provider).find((snapshot) => { const knownMtime = knownTokenMtimes.get(snapshot.file); - return knownMtime === undefined || snapshot.mtimeMs > knownMtime + 1; + if (knownMtime === undefined) { + return true; + } + + if (!pending.expectedAccountId) { + return false; + } + + return snapshot.mtimeMs > knownMtime + 1; }) || null ); } +function shouldKeepWaitingForLocalToken( + state: string, + pending: { upstreamCompletedAt?: number }, + now = Date.now() +): boolean { + const upstreamCompletedAt = + pending.upstreamCompletedAt ?? markManualAuthUpstreamCompleted(state, now); + + return ( + upstreamCompletedAt !== null && now - upstreamCompletedAt < POLLED_AUTH_LOCAL_TOKEN_GRACE_MS + ); +} + function parseKiroMethod(raw: unknown): { method: KiroAuthMethod; invalid: boolean } { if (raw === undefined || raw === null) { return { method: normalizeKiroAuthMethod(), invalid: false }; @@ -936,18 +962,15 @@ router.get('/:provider/status', async (req: Request, res: Response): Promise { expect(registry.providers.kiro.accounts['github-ABC123']?.nickname).toBe('work'); }); - it('returns 409 when callback completes upstream but no account can be registered locally', async () => { + it('returns wait after callback submission when the local token is not yet available', async () => { mockFetch([ { - url: /\/v0\/management\/kiro-auth-url\?is_webui=true&method=google$/, + url: /\/v0\/management\/codex-auth-url\?is_webui=true$/, response: { - auth_url: 'https://auth.example.com/authorize?state=state-409', - state: 'state-409', + auth_url: 'https://auth.example.com/authorize?state=state-callback-wait', + state: 'state-callback-wait', }, }, { @@ -187,23 +187,41 @@ describe('cliproxy-auth-routes manual callback nickname persistence', () => { method: 'POST', response: { status: 'ok' }, }, + { + url: /\/v0\/management\/get-auth-status\?state=state-callback-wait$/, + response: { status: 'ok' }, + }, ]); - const startResponse = await postJson('/api/cliproxy/auth/kiro/start-url', { - nickname: 'work', - kiroMethod: 'google', - }); + const startResponse = await postJson('/api/cliproxy/auth/codex/start-url', {}); expect(startResponse.status).toBe(200); - const callbackResponse = await postJson('/api/cliproxy/auth/kiro/submit-callback', { - redirectUrl: 'http://localhost/callback?code=abc123&state=state-409', - }); + const realDateNow = Date.now; + let now = realDateNow(); + Date.now = () => now; + try { + const callbackResponse = await postJson('/api/cliproxy/auth/codex/submit-callback', { + redirectUrl: 'http://localhost/callback?code=abc123&state=state-callback-wait', + }); - expect(callbackResponse.status).toBe(409); - expect(callbackResponse.body).toEqual({ - error: - 'Authenticated token could not be matched to a new account. Retry the flow and choose a different nickname if needed.', - }); + expect(callbackResponse.status).toBe(200); + expect(callbackResponse.body).toEqual({ status: 'wait' }); + + now += 16_000; + + const statusResponse = await getJson( + '/api/cliproxy/auth/codex/status?state=state-callback-wait' + ); + + expect(statusResponse.status).toBe(409); + expect(statusResponse.body).toEqual({ + status: 'error', + error: + 'Authentication completed upstream, but no new local token was saved for this account. Update CCS/CLIProxy and retry.', + }); + } finally { + Date.now = realDateNow; + } }); it('keeps polling briefly after upstream completion before surfacing the missing-token error', async () => { @@ -310,6 +328,191 @@ describe('cliproxy-auth-routes manual callback nickname persistence', () => { }); }); + it('lets polling finish successfully after callback submission once the token appears', async () => { + mockFetch([ + { + url: /\/v0\/management\/codex-auth-url\?is_webui=true$/, + response: { + auth_url: 'https://auth.example.com/authorize?state=state-callback-delayed', + state: 'state-callback-delayed', + }, + }, + { + url: /\/v0\/management\/oauth-callback$/, + method: 'POST', + response: { status: 'ok' }, + }, + { + url: /\/v0\/management\/get-auth-status\?state=state-callback-delayed$/, + response: { status: 'ok' }, + }, + ]); + + const startResponse = await postJson('/api/cliproxy/auth/codex/start-url', {}); + expect(startResponse.status).toBe(200); + + const callbackResponse = await postJson('/api/cliproxy/auth/codex/submit-callback', { + redirectUrl: 'http://localhost/callback?code=abc123&state=state-callback-delayed', + }); + + expect(callbackResponse.status).toBe(200); + expect(callbackResponse.body).toEqual({ status: 'wait' }); + + const tokenDir = path.join(tempHome, '.ccs', 'cliproxy', 'auth'); + fs.mkdirSync(tokenDir, { recursive: true }); + fs.writeFileSync( + path.join(tokenDir, 'codex-callback@example.com.json'), + JSON.stringify({ type: 'codex', email: 'callback@example.com' }), + 'utf8' + ); + + const statusResponse = await getJson( + '/api/cliproxy/auth/codex/status?state=state-callback-delayed' + ); + + expect(statusResponse.status).toBe(200); + expect(statusResponse.body).toEqual({ + status: 'ok', + account: { + id: 'callback@example.com', + email: 'callback@example.com', + nickname: 'callback', + provider: 'codex', + isDefault: true, + }, + }); + }); + + it('keeps the pending auth state alive through the local-token grace window even near TTL expiry', async () => { + mockFetch([ + { + url: /\/v0\/management\/codex-auth-url\?is_webui=true$/, + response: { + auth_url: 'https://auth.example.com/authorize?state=state-ttl-grace', + state: 'state-ttl-grace', + }, + }, + { + url: /\/v0\/management\/get-auth-status\?state=state-ttl-grace$/, + response: { status: 'ok' }, + }, + { + url: /\/v0\/management\/get-auth-status\?state=state-ttl-grace$/, + response: { status: 'ok' }, + }, + ]); + + const realDateNow = Date.now; + let now = realDateNow(); + Date.now = () => now; + try { + const startResponse = await postJson('/api/cliproxy/auth/codex/start-url', {}); + expect(startResponse.status).toBe(200); + + now += 10 * 60 * 1000 - 1_000; + + const firstStatusResponse = await getJson( + '/api/cliproxy/auth/codex/status?state=state-ttl-grace' + ); + + expect(firstStatusResponse.status).toBe(200); + expect(firstStatusResponse.body).toEqual({ status: 'wait' }); + + const tokenDir = path.join(tempHome, '.ccs', 'cliproxy', 'auth'); + fs.mkdirSync(tokenDir, { recursive: true }); + fs.writeFileSync( + path.join(tokenDir, 'codex-ttl@example.com.json'), + JSON.stringify({ type: 'codex', email: 'ttl@example.com' }), + 'utf8' + ); + + now += 2_000; + + const secondStatusResponse = await getJson( + '/api/cliproxy/auth/codex/status?state=state-ttl-grace' + ); + + expect(secondStatusResponse.status).toBe(200); + expect(secondStatusResponse.body).toEqual({ + status: 'ok', + account: { + id: 'ttl@example.com', + email: 'ttl@example.com', + nickname: 'ttl', + provider: 'codex', + isDefault: true, + }, + }); + } finally { + Date.now = realDateNow; + } + }); + + it('does not treat rewrites of pre-existing token files as a newly added account', async () => { + const tokenDir = path.join(tempHome, '.ccs', 'cliproxy', 'auth'); + fs.mkdirSync(tokenDir, { recursive: true }); + const tokenPath = path.join(tokenDir, 'codex-existing@example.com.json'); + fs.writeFileSync( + tokenPath, + JSON.stringify({ type: 'codex', email: 'existing@example.com', version: 1 }), + 'utf8' + ); + + mockFetch([ + { + url: /\/v0\/management\/codex-auth-url\?is_webui=true$/, + response: { + auth_url: 'https://auth.example.com/authorize?state=state-existing-rewrite', + state: 'state-existing-rewrite', + }, + }, + { + url: /\/v0\/management\/get-auth-status\?state=state-existing-rewrite$/, + response: { status: 'ok' }, + }, + { + url: /\/v0\/management\/get-auth-status\?state=state-existing-rewrite$/, + response: { status: 'ok' }, + }, + ]); + + const startResponse = await postJson('/api/cliproxy/auth/codex/start-url', {}); + expect(startResponse.status).toBe(200); + + fs.writeFileSync( + tokenPath, + JSON.stringify({ type: 'codex', email: 'existing@example.com', version: 2 }), + 'utf8' + ); + + const realDateNow = Date.now; + let now = realDateNow(); + Date.now = () => now; + try { + const firstStatusResponse = await getJson( + '/api/cliproxy/auth/codex/status?state=state-existing-rewrite' + ); + + expect(firstStatusResponse.status).toBe(200); + expect(firstStatusResponse.body).toEqual({ status: 'wait' }); + + now += 16_000; + + const secondStatusResponse = await getJson( + '/api/cliproxy/auth/codex/status?state=state-existing-rewrite' + ); + + expect(secondStatusResponse.status).toBe(409); + expect(secondStatusResponse.body).toEqual({ + status: 'error', + error: + 'Authentication completed upstream, but no new local token was saved for this account. Update CCS/CLIProxy and retry.', + }); + } finally { + Date.now = realDateNow; + } + }); + it('registers the new account before reporting polled auth success', async () => { mockFetch([ { diff --git a/ui/src/hooks/use-cliproxy-auth-flow.ts b/ui/src/hooks/use-cliproxy-auth-flow.ts index 34431809..dcc57e09 100644 --- a/ui/src/hooks/use-cliproxy-auth-flow.ts +++ b/ui/src/hooks/use-cliproxy-auth-flow.ts @@ -432,9 +432,16 @@ export function useCliproxyAuthFlow() { return; } const success = data.success === true; + const waitingForLocalToken = data.status === 'wait'; const hasAccount = typeof data.account === 'object' && data.account !== null; - if (response.ok && success && hasAccount) { + if (response.ok && waitingForLocalToken) { + setState((prev) => ({ + ...prev, + isSubmittingCallback: false, + error: null, + })); + } else if (response.ok && success && hasAccount) { stopPolling(); queryClient.invalidateQueries({ queryKey: ['cliproxy-auth'] }); queryClient.invalidateQueries({ queryKey: ['cliproxy-accounts'] }); diff --git a/ui/tests/unit/hooks/use-cliproxy-auth-flow.test.tsx b/ui/tests/unit/hooks/use-cliproxy-auth-flow.test.tsx index 1ad319b1..d90c4593 100644 --- a/ui/tests/unit/hooks/use-cliproxy-auth-flow.test.tsx +++ b/ui/tests/unit/hooks/use-cliproxy-auth-flow.test.tsx @@ -151,6 +151,76 @@ describe('useCliproxyAuthFlow', () => { expect(toast.error).toHaveBeenCalledWith('poll failed'); }); + it('keeps polling through wait responses until a later ok response includes the account', async () => { + let pollCount = 0; + + vi.stubGlobal( + 'fetch', + vi.fn((input: RequestInfo | URL) => { + const url = String(input); + + if (url.includes('/start-url')) { + return Promise.resolve( + createJsonResponse({ + success: true, + authUrl: 'https://auth.example/wait-ok', + state: 'state-wait-ok', + }) + ); + } + + if (url.includes('/status?state=state-wait-ok')) { + pollCount += 1; + if (pollCount === 1) { + return Promise.resolve(createJsonResponse({ status: 'wait' })); + } + + return Promise.resolve( + createJsonResponse({ + status: 'ok', + account: { + id: 'delayed@example.com', + email: 'delayed@example.com', + nickname: 'delayed', + provider: 'codex', + isDefault: true, + }, + }) + ); + } + + return Promise.reject(new Error(`Unexpected fetch: ${url}`)); + }) + ); + + const { result } = renderHook(() => useCliproxyAuthFlow(), { wrapper }); + + await act(async () => { + await result.current.startAuth('codex', { startEndpoint: 'start-url' }); + }); + + await act(async () => { + vi.advanceTimersByTime(3000); + await Promise.resolve(); + await Promise.resolve(); + }); + + expect(result.current.isAuthenticating).toBe(true); + expect(result.current.oauthState).toBe('state-wait-ok'); + expect(result.current.error).toBeNull(); + expect(toast.success).not.toHaveBeenCalled(); + + await act(async () => { + vi.advanceTimersByTime(3000); + await Promise.resolve(); + await Promise.resolve(); + }); + + expect(result.current.isAuthenticating).toBe(false); + expect(result.current.oauthState).toBeNull(); + expect(toast.success).toHaveBeenCalledWith('codex authentication successful'); + }); + it('treats callback responses without an account as failures', async () => { vi.stubGlobal( 'fetch', @@ -197,6 +267,71 @@ describe('useCliproxyAuthFlow', () => { expect(toast.success).not.toHaveBeenCalled(); }); + it('keeps auth active when callback submission returns wait and only errors on the terminal poll', async () => { + vi.stubGlobal( + 'fetch', + vi.fn((input: RequestInfo | URL) => { + const url = String(input); + + if (url.includes('/start-url')) { + return Promise.resolve( + createJsonResponse({ + success: true, + authUrl: 'https://auth.example/callback-wait', + state: 'state-callback-wait', + }) + ); + } + + if (url.includes('/submit-callback')) { + return Promise.resolve(createJsonResponse({ status: 'wait' })); + } + + if (url.includes('/status?state=state-callback-wait')) { + return Promise.resolve( + createJsonResponse({ + status: 'error', + error: + 'Authentication completed upstream, but no new local token was saved for this account. Update CCS/CLIProxy and retry.', + }) + ); + } + + return Promise.reject(new Error(`Unexpected fetch: ${url}`)); + }) + ); + + const { result } = renderHook(() => useCliproxyAuthFlow(), { wrapper }); + + await act(async () => { + await result.current.startAuth('codex', { startEndpoint: 'start-url' }); + }); + + await act(async () => { + await result.current.submitCallback( + 'http://localhost/callback?code=abc123&state=state-callback-wait' + ); + }); + + expect(result.current.isSubmittingCallback).toBe(false); + expect(result.current.isAuthenticating).toBe(true); + expect(result.current.error).toBeNull(); + + await act(async () => { + vi.advanceTimersByTime(3000); + await Promise.resolve(); + await Promise.resolve(); + }); + + expect(result.current.isAuthenticating).toBe(false); + expect(result.current.error).toBe( + 'Authentication completed upstream, but no new local token was saved for this account. Update CCS/CLIProxy and retry.' + ); + expect(toast.error).toHaveBeenCalledWith( + 'Authentication completed upstream, but no new local token was saved for this account. Update CCS/CLIProxy and retry.' + ); + }); + it('treats status ok responses without an account as failures', async () => { vi.stubGlobal( 'fetch',