Merge pull request #907 from kaitranntt/kai/fix/906-polled-oauth-token-persistence

fix(auth): wait for polled oauth token persistence
This commit is contained in:
Kai (Tam Nhu) Tran
2026-04-04 12:13:40 -04:00
committed by GitHub
4 changed files with 555 additions and 38 deletions
+119 -13
View File
@@ -62,6 +62,7 @@ import { requireLocalAccessWhenAuthDisabled } from '../middleware/auth-middlewar
const router = Router();
const MANUAL_AUTH_STATE_TTL_MS = 10 * 60 * 1000;
const POLLED_AUTH_LOCAL_TOKEN_GRACE_MS = 15 * 1000;
type ProviderTokenSnapshot = {
file: string;
mtimeMs: number;
@@ -74,6 +75,7 @@ const pendingManualAuthState = new Map<
nickname?: string;
expectedAccountId?: string;
createdAt: number;
upstreamCompletedAt?: number;
knownTokenFiles: ProviderTokenSnapshot[];
}
>();
@@ -97,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);
}
}
@@ -122,6 +129,7 @@ function getManualAuthState(state: string | undefined): {
nickname?: string;
expectedAccountId?: string;
createdAt: number;
upstreamCompletedAt?: number;
knownTokenFiles: ProviderTokenSnapshot[];
} | null {
if (!state) {
@@ -138,10 +146,27 @@ function getManualAuthState(state: string | undefined): {
nickname: pending.nickname,
expectedAccountId: pending.expectedAccountId,
createdAt: pending.createdAt,
upstreamCompletedAt: pending.upstreamCompletedAt,
knownTokenFiles: pending.knownTokenFiles,
};
}
function markManualAuthUpstreamCompleted(state: string, now = Date.now()): number | null {
pruneExpiredManualAuthState(now);
const pending = pendingManualAuthState.get(state);
if (!pending) {
return null;
}
if (pending.upstreamCompletedAt !== undefined) {
return pending.upstreamCompletedAt;
}
pending.upstreamCompletedAt = now;
pendingManualAuthState.set(state, pending);
return now;
}
function listProviderTokenSnapshots(provider: CLIProxyProvider): ProviderTokenSnapshot[] {
const tokenDir = getProviderTokenDir(provider);
if (!fs.existsSync(tokenDir)) {
@@ -179,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])
@@ -188,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 };
@@ -916,8 +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) {
if (shouldKeepWaitingForLocalToken(state, pendingAuth, now)) {
res.json({ status: 'wait' });
return;
}
pendingManualAuthState.delete(state);
res.status(409).json({
status: 'error',
error:
@@ -935,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),
@@ -1048,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;
}
@@ -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,26 +187,44 @@ 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('returns 409 when status polling completes upstream but no new local token was written', async () => {
it('keeps polling briefly after upstream completion before surfacing the missing-token error', async () => {
const tokenDir = path.join(tempHome, '.ccs', 'cliproxy', 'auth');
fs.mkdirSync(tokenDir, { recursive: true });
fs.writeFileSync(
@@ -232,18 +250,269 @@ describe('cliproxy-auth-routes manual callback nickname persistence', () => {
const startResponse = await postJson('/api/cliproxy/auth/codex/start-url', {});
expect(startResponse.status).toBe(200);
const statusResponse = await getJson(
'/api/cliproxy/auth/codex/status?state=state-status-missing'
const realDateNow = Date.now;
let now = realDateNow();
Date.now = () => now;
try {
const firstStatusResponse = await getJson(
'/api/cliproxy/auth/codex/status?state=state-status-missing'
);
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-status-missing'
);
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('continues polling until the local token appears after upstream completion', async () => {
mockFetch([
{
url: /\/v0\/management\/codex-auth-url\?is_webui=true$/,
response: {
auth_url: 'https://auth.example.com/authorize?state=state-status-delayed',
state: 'state-status-delayed',
},
},
{
url: /\/v0\/management\/get-auth-status\?state=state-status-delayed$/,
response: { status: 'ok' },
},
]);
const startResponse = await postJson('/api/cliproxy/auth/codex/start-url', {});
expect(startResponse.status).toBe(200);
const firstStatusResponse = await getJson(
'/api/cliproxy/auth/codex/status?state=state-status-delayed'
);
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.',
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-delayed@example.com.json'),
JSON.stringify({ type: 'codex', email: 'delayed@example.com' }),
'utf8'
);
const secondStatusResponse = await getJson(
'/api/cliproxy/auth/codex/status?state=state-status-delayed'
);
expect(secondStatusResponse.status).toBe(200);
expect(secondStatusResponse.body).toEqual({
status: 'ok',
account: {
id: 'delayed@example.com',
email: 'delayed@example.com',
nickname: 'delayed',
provider: 'codex',
isDefault: true,
},
});
});
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([
{
+8 -1
View File
@@ -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'] });
@@ -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',