From 7142a8b16c81eb4a03a63c8fbf1823f011c8e1bc Mon Sep 17 00:00:00 2001 From: Sharjeel Yunus Date: Thu, 20 Aug 2026 02:40:32 +0500 Subject: [PATCH 1/2] feat(auth): implement token refresh with retry logic and error handling - Added `TokenRefreshError` class to handle token refresh errors with retryable status. - Enhanced `refreshIdToken` function to retry token refresh on transient errors and specific HTTP status codes (429, 500). - Introduced utility functions for parsing refresh responses and checking token expiry. - Updated tests to cover new retry logic and error scenarios for token refresh. --- src/auth/session.ts | 126 +++++++++++--- src/auth/token.ts | 7 + tests/auth/session.test.ts | 326 ++++++++++++++++++++++++++++++++----- tests/auth/token.test.ts | 37 ++++- 4 files changed, 434 insertions(+), 62 deletions(-) diff --git a/src/auth/session.ts b/src/auth/session.ts index 995d562..71e6bce 100644 --- a/src/auth/session.ts +++ b/src/auth/session.ts @@ -3,10 +3,12 @@ import { writeGlobalConfig, type EnsembleUserConfig, } from '../config/globalConfig.js'; -import { decodeIdTokenClaims, isTokenExpired } from './token.js'; +import { decodeIdTokenClaims, isTokenExpired, isTokenPastExpiry } from './token.js'; import { getEnsembleFirebaseApiKey } from '../config/env.js'; const DEFAULT_REFRESH_API_BASE = 'https://securetoken.googleapis.com/v1/token'; +const REFRESH_MAX_ATTEMPTS = 3; +const REFRESH_RETRY_DELAY_MS = 200; interface RefreshTokenResponse { id_token?: string; @@ -18,6 +20,16 @@ interface RefreshTokenResponse { }; } +class TokenRefreshError extends Error { + readonly retryable: boolean; + + constructor(message: string, retryable: boolean) { + super(message); + this.name = 'TokenRefreshError'; + this.retryable = retryable; + } +} + export type AuthSessionResult = | { ok: true; @@ -33,32 +45,59 @@ export type AuthSessionResult = message: string; }; -async function refreshIdToken(refreshToken: string): Promise<{ +function isRetryableRefreshStatus(status: number): boolean { + return status === 429 || status >= 500; +} + +function parseRefreshResponseBody(text: string): RefreshTokenResponse { + if (!text.trim()) { + throw new TokenRefreshError('Token refresh failed: empty response body', true); + } + + try { + return JSON.parse(text) as RefreshTokenResponse; + } catch { + throw new TokenRefreshError('Token refresh failed: invalid JSON response', true); + } +} + +async function sleep(ms: number): Promise { + await new Promise((resolve) => setTimeout(resolve, ms)); +} + +async function refreshIdTokenOnce( + refreshToken: string, + apiKey: string +): Promise<{ idToken: string; refreshToken: string; userId?: string; }> { - const apiKey = getEnsembleFirebaseApiKey(); - if (!apiKey) { - throw new Error('Missing Firebase API key for token refresh. Set ENSEMBLE_FIREBASE_API_KEY.'); - } - const refreshUrl = `${DEFAULT_REFRESH_API_BASE}?key=${encodeURIComponent(apiKey)}`; const body = new URLSearchParams({ grant_type: 'refresh_token', refresh_token: refreshToken, }); - const res = await fetch(refreshUrl, { - method: 'POST', - headers: { 'Content-Type': 'application/x-www-form-urlencoded' }, - body, - }); + let res: Response; + try { + res = await fetch(refreshUrl, { + method: 'POST', + headers: { 'Content-Type': 'application/x-www-form-urlencoded' }, + body, + }); + } catch (err) { + const message = err instanceof Error ? err.message : 'network error'; + throw new TokenRefreshError(`Token refresh failed: ${message}`, true); + } + + const text = await res.text(); + const data = parseRefreshResponseBody(text); - const data = (await res.json()) as RefreshTokenResponse; if (!res.ok || !data.id_token) { const reason = data?.error?.message ?? `HTTP ${res.status}`; - throw new Error(`Token refresh failed: ${reason}`); + const retryable = isRetryableRefreshStatus(res.status); + throw new TokenRefreshError(`Token refresh failed: ${reason}`, retryable); } return { @@ -68,6 +107,34 @@ async function refreshIdToken(refreshToken: string): Promise<{ }; } +async function refreshIdToken(refreshToken: string): Promise<{ + idToken: string; + refreshToken: string; + userId?: string; +}> { + const apiKey = getEnsembleFirebaseApiKey(); + if (!apiKey) { + throw new Error('Missing Firebase API key for token refresh. Set ENSEMBLE_FIREBASE_API_KEY.'); + } + + let lastError: Error = new Error('Token refresh failed.'); + + for (let attempt = 1; attempt <= REFRESH_MAX_ATTEMPTS; attempt += 1) { + try { + return await refreshIdTokenOnce(refreshToken, apiKey); + } catch (err) { + lastError = err instanceof Error ? err : new Error('Token refresh failed.'); + const retryable = err instanceof TokenRefreshError && err.retryable; + if (!retryable || attempt === REFRESH_MAX_ATTEMPTS) { + break; + } + await sleep(REFRESH_RETRY_DELAY_MS); + } + } + + throw lastError; +} + const ENSEMBLE_TOKEN_ENV = 'ENSEMBLE_TOKEN'; /** @@ -99,6 +166,19 @@ async function sessionFromEnvToken(): Promise { } } +function sessionFromStoredUser( + user: NonNullable +): AuthSessionResult { + return { + ok: true, + idToken: user.idToken, + userId: user.uid, + name: user.name, + email: user.email, + refreshed: false, + }; +} + export async function getValidAuthSession(): Promise { const fromEnv = await sessionFromEnvToken(); if (fromEnv.ok) return fromEnv; @@ -117,14 +197,7 @@ export async function getValidAuthSession(): Promise { } if (!isTokenExpired(user.idToken)) { - return { - ok: true, - idToken: user.idToken, - userId: user.uid, - name: user.name, - email: user.email, - refreshed: false, - }; + return sessionFromStoredUser(user); } if (!user.refreshToken) { @@ -159,11 +232,16 @@ export async function getValidAuthSession(): Promise { email: updatedUser.email, refreshed: true, }; - } catch { + } catch (err) { + if (!isTokenPastExpiry(user.idToken)) { + return sessionFromStoredUser(user); + } + + const reason = err instanceof Error ? err.message : 'Token refresh failed.'; return { ok: false, reason: 'expired', - message: `Session expired and automatic refresh failed. Run \`ensemble login\` again.`, + message: `Session expired and automatic refresh failed: ${reason}. Run \`ensemble login\` again.`, }; } } diff --git a/src/auth/token.ts b/src/auth/token.ts index 5e10746..d843d48 100644 --- a/src/auth/token.ts +++ b/src/auth/token.ts @@ -45,3 +45,10 @@ export function isTokenExpired(idToken: string, bufferSeconds = 60): boolean { if (expiry === undefined) return true; return expiry <= Date.now() + bufferSeconds * 1000; } + +/** True only when JWT exp is present and already past (no proactive buffer). Missing exp => false. */ +export function isTokenPastExpiry(idToken: string): boolean { + const expiry = getIdTokenExpiryMs(idToken); + if (expiry === undefined) return false; + return expiry <= Date.now(); +} diff --git a/tests/auth/session.test.ts b/tests/auth/session.test.ts index c2ab841..3bae95d 100644 --- a/tests/auth/session.test.ts +++ b/tests/auth/session.test.ts @@ -20,17 +20,33 @@ function makeJwt(payload: Record): string { return `${header}.${payloadB64}.${base64url('sig')}`; } +function mockRefreshFetchResponse(params: { + ok: boolean; + status?: number; + body?: unknown; + text?: string; +}) { + const bodyText = params.text ?? JSON.stringify(params.body ?? {}); + return { + ok: params.ok, + status: params.status ?? (params.ok ? 200 : 400), + text: async () => bodyText, + }; +} + describe('getValidAuthSession', () => { const originalEnv = process.env.ENSEMBLE_FIREBASE_API_KEY; const originalToken = process.env.ENSEMBLE_TOKEN; beforeEach(() => { vi.mocked(globalConfig.readGlobalConfig).mockReset(); + vi.mocked(globalConfig.writeGlobalConfig).mockReset(); process.env.ENSEMBLE_FIREBASE_API_KEY = 'test-api-key'; delete process.env.ENSEMBLE_TOKEN; }); afterEach(() => { + vi.useRealTimers(); process.env.ENSEMBLE_FIREBASE_API_KEY = originalEnv; if (originalToken !== undefined) process.env.ENSEMBLE_TOKEN = originalToken; else delete process.env.ENSEMBLE_TOKEN; @@ -44,14 +60,16 @@ describe('getValidAuthSession', () => { exp: Math.floor(Date.now() / 1000) + 3600, }); const originalFetch = globalThis.fetch; - globalThis.fetch = vi.fn().mockResolvedValue({ - ok: true, - json: async () => ({ - id_token: newToken, - refresh_token: 'env-refresh-token', - expires_in: '3600', - }), - }); + globalThis.fetch = vi.fn().mockResolvedValue( + mockRefreshFetchResponse({ + ok: true, + body: { + id_token: newToken, + refresh_token: 'env-refresh-token', + expires_in: '3600', + }, + }) + ); const result = await getValidAuthSession(); @@ -69,11 +87,14 @@ describe('getValidAuthSession', () => { it('returns expired when ENSEMBLE_TOKEN is set but refresh fails', async () => { process.env.ENSEMBLE_TOKEN = 'bad-refresh-token'; const originalFetch = globalThis.fetch; - globalThis.fetch = vi.fn().mockResolvedValue({ - ok: false, - status: 400, - json: async () => ({ error: { message: 'INVALID_GRANT' } }), - }); + const fetchMock = vi.fn().mockResolvedValue( + mockRefreshFetchResponse({ + ok: false, + status: 400, + body: { error: { message: 'INVALID_GRANT' } }, + }) + ); + globalThis.fetch = fetchMock; const result = await getValidAuthSession(); @@ -81,10 +102,12 @@ describe('getValidAuthSession', () => { expect(result.ok).toBe(false); if (!result.ok) { expect(result.reason).toBe('expired'); + expect(result.message).toContain('INVALID_GRANT'); expect(result.message).toContain('ENSEMBLE_TOKEN'); expect(result.message).toContain('ensemble token'); } expect(globalConfig.readGlobalConfig).not.toHaveBeenCalled(); + expect(fetchMock).toHaveBeenCalledTimes(1); }); it('returns not_logged_in when no config', async () => { @@ -171,13 +194,23 @@ describe('getValidAuthSession', () => { refreshToken: 'refresh-token', }, }); - delete process.env.ENSEMBLE_FIREBASE_API_KEY; + const originalFetch = globalThis.fetch; + globalThis.fetch = vi.fn().mockResolvedValue( + mockRefreshFetchResponse({ + ok: false, + status: 400, + body: { error: { message: 'INVALID_GRANT' } }, + }) + ); const result = await getValidAuthSession(); + globalThis.fetch = originalFetch; + expect(result.ok).toBe(false); if (!result.ok) { expect(result.reason).toBe('expired'); + expect(result.message).toContain('INVALID_GRANT'); expect(result.message).toContain('Run `ensemble login` again.'); } }); @@ -224,31 +257,25 @@ describe('getValidAuthSession', () => { email: 'a@b.com', exp: Math.floor(Date.now() / 1000) + 3600, }); - vi.mocked(globalConfig.readGlobalConfig) - .mockResolvedValueOnce({ - user: { - uid: 'u1', - idToken: oldToken, - refreshToken: 'refresh-123', - }, - }) - .mockResolvedValue({ - user: { - uid: 'u1', - idToken: newToken, - refreshToken: 'refresh-123', - }, - }); + vi.mocked(globalConfig.readGlobalConfig).mockResolvedValue({ + user: { + uid: 'u1', + idToken: oldToken, + refreshToken: 'refresh-123', + }, + }); const originalFetch = globalThis.fetch; - globalThis.fetch = vi.fn().mockResolvedValue({ - ok: true, - json: async () => ({ - id_token: newToken, - refresh_token: 'refresh-456', - expires_in: '3600', - }), - }); + globalThis.fetch = vi.fn().mockResolvedValue( + mockRefreshFetchResponse({ + ok: true, + body: { + id_token: newToken, + refresh_token: 'refresh-456', + expires_in: '3600', + }, + }) + ); const result = await getValidAuthSession(); @@ -261,4 +288,229 @@ describe('getValidAuthSession', () => { } expect(globalConfig.writeGlobalConfig).toHaveBeenCalled(); }); + + it('retries refresh when first attempt fails transiently then succeeds', async () => { + vi.useFakeTimers(); + const oldToken = makeJwt({ + userId: 'u1', + exp: Math.floor(Date.now() / 1000) - 3600, + }); + const newToken = makeJwt({ + userId: 'u1', + email: 'a@b.com', + exp: Math.floor(Date.now() / 1000) + 3600, + }); + vi.mocked(globalConfig.readGlobalConfig).mockResolvedValue({ + user: { + uid: 'u1', + idToken: oldToken, + refreshToken: 'refresh-123', + }, + }); + + const originalFetch = globalThis.fetch; + const fetchMock = vi + .fn() + .mockRejectedValueOnce(new Error('fetch failed')) + .mockResolvedValueOnce( + mockRefreshFetchResponse({ + ok: true, + body: { + id_token: newToken, + refresh_token: 'refresh-456', + expires_in: '3600', + }, + }) + ); + globalThis.fetch = fetchMock; + + const resultPromise = getValidAuthSession(); + await vi.advanceTimersByTimeAsync(200); + const result = await resultPromise; + + globalThis.fetch = originalFetch; + + expect(result.ok).toBe(true); + expect(fetchMock).toHaveBeenCalledTimes(2); + expect(globalConfig.writeGlobalConfig).toHaveBeenCalled(); + }); + + it('retries refresh on 503 then succeeds', async () => { + vi.useFakeTimers(); + const oldToken = makeJwt({ + userId: 'u1', + exp: Math.floor(Date.now() / 1000) - 3600, + }); + const newToken = makeJwt({ + userId: 'u1', + exp: Math.floor(Date.now() / 1000) + 3600, + }); + vi.mocked(globalConfig.readGlobalConfig).mockResolvedValue({ + user: { + uid: 'u1', + idToken: oldToken, + refreshToken: 'refresh-123', + }, + }); + + const originalFetch = globalThis.fetch; + const fetchMock = vi + .fn() + .mockResolvedValueOnce( + mockRefreshFetchResponse({ + ok: false, + status: 503, + body: { error: { message: 'UNAVAILABLE' } }, + }) + ) + .mockResolvedValueOnce( + mockRefreshFetchResponse({ + ok: true, + body: { id_token: newToken, refresh_token: 'refresh-456' }, + }) + ); + globalThis.fetch = fetchMock; + + const resultPromise = getValidAuthSession(); + await vi.advanceTimersByTimeAsync(200); + const result = await resultPromise; + + globalThis.fetch = originalFetch; + + expect(result.ok).toBe(true); + expect(fetchMock).toHaveBeenCalledTimes(2); + }); + + it('falls back to stored token when refresh fails within proactive buffer', async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date('2025-01-01T12:59:30Z')); + const bufferToken = makeJwt({ userId: 'u1', email: 'a@b.com', exp: 1735736400 }); + vi.mocked(globalConfig.readGlobalConfig).mockResolvedValue({ + user: { + uid: 'u1', + email: 'a@b.com', + idToken: bufferToken, + refreshToken: 'refresh-123', + }, + }); + + const originalFetch = globalThis.fetch; + const fetchMock = vi.fn().mockRejectedValue(new Error('fetch failed')); + globalThis.fetch = fetchMock; + + const resultPromise = getValidAuthSession(); + await vi.advanceTimersByTimeAsync(500); + const result = await resultPromise; + + globalThis.fetch = originalFetch; + + expect(result.ok).toBe(true); + if (result.ok) { + expect(result.idToken).toBe(bufferToken); + expect(result.refreshed).toBe(false); + } + expect(globalConfig.writeGlobalConfig).not.toHaveBeenCalled(); + expect(fetchMock).toHaveBeenCalledTimes(3); + }); + + it('returns expired with underlying reason when token is past exp and refresh fails', async () => { + vi.useFakeTimers(); + const expiredToken = makeJwt({ + userId: 'u1', + exp: Math.floor(Date.now() / 1000) - 3600, + }); + vi.mocked(globalConfig.readGlobalConfig).mockResolvedValue({ + user: { + uid: 'u1', + idToken: expiredToken, + refreshToken: 'refresh-token', + }, + }); + + const originalFetch = globalThis.fetch; + const fetchMock = vi.fn().mockRejectedValue(new Error('fetch failed')); + globalThis.fetch = fetchMock; + + const resultPromise = getValidAuthSession(); + await vi.advanceTimersByTimeAsync(500); + const result = await resultPromise; + + globalThis.fetch = originalFetch; + + expect(result.ok).toBe(false); + if (!result.ok) { + expect(result.message).toContain('fetch failed'); + expect(result.message).toContain('Run `ensemble login` again.'); + } + expect(fetchMock).toHaveBeenCalledTimes(3); + }); + + it('falls back to stored token when jwt has no exp and refresh fails', async () => { + vi.useFakeTimers(); + const tokenWithoutExp = makeJwt({ userId: 'u1', email: 'a@b.com' }); + vi.mocked(globalConfig.readGlobalConfig).mockResolvedValue({ + user: { + uid: 'u1', + email: 'a@b.com', + idToken: tokenWithoutExp, + refreshToken: 'refresh-123', + }, + }); + + const originalFetch = globalThis.fetch; + globalThis.fetch = vi.fn().mockRejectedValue(new Error('fetch failed')); + + const resultPromise = getValidAuthSession(); + await vi.advanceTimersByTimeAsync(500); + const result = await resultPromise; + + globalThis.fetch = originalFetch; + + expect(result.ok).toBe(true); + if (result.ok) { + expect(result.idToken).toBe(tokenWithoutExp); + expect(result.refreshed).toBe(false); + } + expect(globalConfig.writeGlobalConfig).not.toHaveBeenCalled(); + }); + + it('retries on empty response body then succeeds', async () => { + vi.useFakeTimers(); + const oldToken = makeJwt({ + userId: 'u1', + exp: Math.floor(Date.now() / 1000) - 3600, + }); + const newToken = makeJwt({ + userId: 'u1', + exp: Math.floor(Date.now() / 1000) + 3600, + }); + vi.mocked(globalConfig.readGlobalConfig).mockResolvedValue({ + user: { + uid: 'u1', + idToken: oldToken, + refreshToken: 'refresh-123', + }, + }); + + const originalFetch = globalThis.fetch; + const fetchMock = vi + .fn() + .mockResolvedValueOnce(mockRefreshFetchResponse({ ok: true, text: '' })) + .mockResolvedValueOnce( + mockRefreshFetchResponse({ + ok: true, + body: { id_token: newToken, refresh_token: 'refresh-456' }, + }) + ); + globalThis.fetch = fetchMock; + + const resultPromise = getValidAuthSession(); + await vi.advanceTimersByTimeAsync(200); + const result = await resultPromise; + + globalThis.fetch = originalFetch; + + expect(result.ok).toBe(true); + expect(fetchMock).toHaveBeenCalledTimes(2); + }); }); diff --git a/tests/auth/token.test.ts b/tests/auth/token.test.ts index ce7d01c..e28cfdb 100644 --- a/tests/auth/token.test.ts +++ b/tests/auth/token.test.ts @@ -1,5 +1,5 @@ import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; -import { decodeIdTokenClaims, getIdTokenExpiryMs, isTokenExpired } from '../../src/auth/token.js'; +import { decodeIdTokenClaims, getIdTokenExpiryMs, isTokenExpired, isTokenPastExpiry } from '../../src/auth/token.js'; function base64urlEncode(str: string): string { return Buffer.from(str, 'utf8') @@ -118,3 +118,38 @@ describe('isTokenExpired', () => { expect(isTokenExpired(token)).toBe(true); }); }); + +describe('isTokenPastExpiry', () => { + beforeEach(() => { + vi.useFakeTimers(); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + it('returns false when token is not past expiry', () => { + vi.setSystemTime(new Date('2025-01-01T12:00:00Z')); + const token = makeJwt({ userId: 'u1', exp: 1735736400 }); + expect(isTokenPastExpiry(token)).toBe(false); + }); + + it('returns true when token exp is in the past', () => { + vi.setSystemTime(new Date('2025-01-01T13:00:00Z')); + const token = makeJwt({ userId: 'u1', exp: 1735732800 }); + expect(isTokenPastExpiry(token)).toBe(true); + }); + + it('returns false when jwt has no exp claim', () => { + vi.setSystemTime(new Date('2025-01-01T12:00:00Z')); + const token = makeJwt({ userId: 'u1' }); + expect(isTokenPastExpiry(token)).toBe(false); + }); + + it('returns false when token is within proactive buffer but not past exp', () => { + vi.setSystemTime(new Date('2025-01-01T12:59:30Z')); + const token = makeJwt({ userId: 'u1', exp: 1735736400 }); + expect(isTokenExpired(token, 60)).toBe(true); + expect(isTokenPastExpiry(token)).toBe(false); + }); +}); From f080af810ffa53091494a1325ff569355e945704 Mon Sep 17 00:00:00 2001 From: Sharjeel Yunus Date: Thu, 20 Aug 2026 02:43:27 +0500 Subject: [PATCH 2/2] refactor(auth): simplify sessionFromStoredUser function signature - Updated the function signature of `sessionFromStoredUser` for improved readability by removing unnecessary line breaks. - Enhanced code formatting in the test file for better consistency and clarity. --- src/auth/session.ts | 4 +--- tests/auth/token.test.ts | 7 ++++++- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/src/auth/session.ts b/src/auth/session.ts index 71e6bce..f0cc65a 100644 --- a/src/auth/session.ts +++ b/src/auth/session.ts @@ -166,9 +166,7 @@ async function sessionFromEnvToken(): Promise { } } -function sessionFromStoredUser( - user: NonNullable -): AuthSessionResult { +function sessionFromStoredUser(user: NonNullable): AuthSessionResult { return { ok: true, idToken: user.idToken, diff --git a/tests/auth/token.test.ts b/tests/auth/token.test.ts index e28cfdb..bde3163 100644 --- a/tests/auth/token.test.ts +++ b/tests/auth/token.test.ts @@ -1,5 +1,10 @@ import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; -import { decodeIdTokenClaims, getIdTokenExpiryMs, isTokenExpired, isTokenPastExpiry } from '../../src/auth/token.js'; +import { + decodeIdTokenClaims, + getIdTokenExpiryMs, + isTokenExpired, + isTokenPastExpiry, +} from '../../src/auth/token.js'; function base64urlEncode(str: string): string { return Buffer.from(str, 'utf8')