fix(auth): harden delayed oauth token registration

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<voi
         return;
       }
 
+      const now = Date.now();
       const tokenSnapshot = findNewTokenSnapshotForPendingAuth(localProvider, pendingAuth);
       if (!tokenSnapshot) {
-        const upstreamCompletedAt =
-          pendingAuth.upstreamCompletedAt ?? markManualAuthUpstreamCompleted(state, Date.now());
-        if (
-          upstreamCompletedAt !== null &&
-          Date.now() - upstreamCompletedAt < POLLED_AUTH_LOCAL_TOKEN_GRACE_MS
-        ) {
+        if (shouldKeepWaitingForLocalToken(state, pendingAuth, now)) {
           res.json({ status: 'wait' });
           return;
         }
 
+        pendingManualAuthState.delete(state);
         res.status(409).json({
           status: 'error',
           error:
@@ -965,6 +988,7 @@ router.get('/:provider/status', async (req: Request, res: Response): Promise<voi
       );
 
       if (!account) {
+        pendingManualAuthState.delete(state);
         res.status(409).json({
           status: 'error',
           error: getManualCallbackRegistrationError(localProvider),
@@ -1078,20 +1102,72 @@ router.post('/:provider/submit-callback', async (req: Request, res: Response): P
       return;
     }
 
-    const account = registerAccountFromToken(
-      provider as CLIProxyProvider,
-      getProviderTokenDir(provider as CLIProxyProvider),
-      pendingAuth?.nickname,
-      false,
-      pendingAuth?.expectedAccountId
-    );
-    if (parsed.state) {
-      pendingManualAuthState.delete(parsed.state);
+    const localProvider = provider as CLIProxyProvider;
+    const now = Date.now();
+
+    if (pendingAuth) {
+      const tokenSnapshot = findNewTokenSnapshotForPendingAuth(localProvider, pendingAuth);
+      if (!tokenSnapshot) {
+        if (parsed.state && shouldKeepWaitingForLocalToken(parsed.state, pendingAuth, now)) {
+          res.json({ status: 'wait' });
+          return;
+        }
+
+        if (parsed.state) {
+          pendingManualAuthState.delete(parsed.state);
+        }
+        res.status(409).json({
+          error: getManualCallbackRegistrationError(localProvider),
+        });
+        return;
+      }
+
+      const account = registerAccountFromToken(
+        localProvider,
+        getProviderTokenDir(localProvider),
+        pendingAuth.nickname,
+        false,
+        tokenSnapshot.file
+      );
+
+      if (!account) {
+        if (parsed.state) {
+          pendingManualAuthState.delete(parsed.state);
+        }
+        res.status(409).json({
+          error: getManualCallbackRegistrationError(localProvider),
+        });
+        return;
+      }
+
+      if (parsed.state) {
+        pendingManualAuthState.delete(parsed.state);
+      }
+
+      res.json({
+        success: true,
+        account: {
+          id: account.id,
+          email: account.email,
+          nickname: account.nickname,
+          provider: account.provider,
+          isDefault: account.isDefault,
+        },
+      });
+      return;
     }
 
+    const account = registerAccountFromToken(
+      localProvider,
+      getProviderTokenDir(localProvider),
+      undefined,
+      false,
+      undefined
+    );
+
     if (!account) {
       res.status(409).json({
-        error: getManualCallbackRegistrationError(provider as CLIProxyProvider),
+        error: getManualCallbackRegistrationError(localProvider),
       });
       return;
     }

tests/unit/web-server/cliproxy-auth-routes-manual-callback.test.ts

@@ -173,13 +173,13 @@ describe('cliproxy-auth-routes manual callback nickname persistence', () => {
     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([
       {

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'] });

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',