diff --git a/src/__tests__/bearer-retry.test.ts b/src/__tests__/bearer-retry.test.ts new file mode 100644 index 0000000..a6239b3 --- /dev/null +++ b/src/__tests__/bearer-retry.test.ts @@ -0,0 +1,117 @@ +import { describe, expect, mock, test } from 'bun:test' + +let sendCalls = 0 +let apiResponseLogCalls = 0 + +mock.module('../plugin/logger.js', () => ({ + debug: () => {}, + error: () => {}, + getTimestamp: () => '2026-07-22T00:00:00.000Z', + log: () => {}, + logApiError: () => {}, + logApiRequest: () => {}, + logApiResponse: () => { + apiResponseLogCalls++ + }, + warn: () => {} +})) + +mock.module('../plugin/sdk-client.js', () => ({ + createSdkClient: () => ({ + send: async () => { + sendCalls++ + const error: any = new Error('The bearer token included in the request is invalid') + error.name = 'ForbiddenException' + error.$metadata = { httpStatusCode: 403 } + throw error + } + }) +})) + +const { RequestHandler } = await import('../core/request/request-handler.js') + +describe('RequestHandler bearer-invalid recovery', () => { + test('forces one refresh, retries once, then applies permanent failure handling', async () => { + sendCalls = 0 + apiResponseLogCalls = 0 + + const account: any = { + id: 'account-1', + email: 'user@example.com', + authMethod: 'idc', + region: 'us-east-1', + refreshToken: 'refresh-token', + accessToken: 'stale-access-token', + expiresAt: Date.now() + 60_000, + isHealthy: true, + failCount: 0, + usedCount: 0, + limitCount: 0 + } + + const accountManager: any = { + getAccounts: () => [account], + getAccountCount: () => 1, + toAuthDetails: (selected: any) => ({ + access: selected.accessToken, + refresh: selected.refreshToken, + expires: selected.expiresAt, + authMethod: selected.authMethod, + region: selected.region, + email: selected.email + }) + } + + const repository: any = { + batchSave: async () => {}, + save: async () => {}, + invalidateCache: () => {}, + findAll: async () => [account] + } + + const config: any = { + max_request_iterations: 5, + request_timeout_ms: 5_000, + rate_limit_max_retries: 2, + rate_limit_retry_delay_ms: 1, + token_expiry_buffer_ms: 0, + auto_sync_kiro_cli: false, + account_selection_strategy: 'sticky', + enable_log_api_request: true + } + + const handler: any = new RequestHandler(accountManager, config, repository) + let forceRefreshCalls = 0 + + handler.accountSelector = { + selectHealthyAccount: async () => account + } + handler.tokenRefresher = { + refreshIfNeeded: async (selected: any) => ({ account: selected, shouldContinue: false }), + forceRefresh: async () => { + forceRefreshCalls++ + } + } + handler.prepareSdkRequest = () => ({ + region: 'us-east-1', + effort: undefined, + conversationState: {}, + profileArn: undefined, + conversationId: 'conversation-1', + streaming: false, + effectiveModel: 'claude-sonnet-4-5' + }) + + await expect( + handler.handle( + 'https://q.us-east-1.amazonaws.com/models/claude-sonnet-4-5', + { body: '{}' }, + () => {} + ) + ).rejects.toThrow('Kiro Error: 403') + + expect(sendCalls).toBe(2) + expect(forceRefreshCalls).toBe(1) + expect(apiResponseLogCalls).toBe(2) + }) +}) diff --git a/src/core/auth/token-refresher.ts b/src/core/auth/token-refresher.ts index a1b927c..8e47f91 100644 --- a/src/core/auth/token-refresher.ts +++ b/src/core/auth/token-refresher.ts @@ -41,6 +41,34 @@ export class TokenRefresher { } } + async forceRefresh(account: ManagedAccount, auth: KiroAuthDetails): Promise { + if (this.config.auto_sync_kiro_cli) { + await this.syncFromKiroCli() + } + + this.repository.invalidateCache() + const accounts = await this.repository.findAll() + const synced = accounts.find((a: ManagedAccount) => a.id === account.id) + + if (synced && synced.accessToken !== account.accessToken) { + this.accountManager.updateFromAuth(account, this.accountManager.toAuthDetails(synced)) + await this.repository.batchSave(this.accountManager.getAccounts()) + logger.debug('Force refresh: recovered newer token from CLI sync') + return + } + + try { + const newAuth = await refreshAccessToken(auth) + this.accountManager.updateFromAuth(account, newAuth) + await this.repository.batchSave(this.accountManager.getAccounts()) + logger.debug('Force refresh: token refreshed via OIDC') + } catch (e: any) { + logger.warn('Force refresh failed, will retry with current token', { + message: e instanceof Error ? e.message : String(e) + }) + } + } + private async handleRefreshError( error: any, account: ManagedAccount, diff --git a/src/core/request/error-handler.ts b/src/core/request/error-handler.ts index c9f968f..235d431 100644 --- a/src/core/request/error-handler.ts +++ b/src/core/request/error-handler.ts @@ -6,6 +6,7 @@ type ToastFunction = (message: string, variant: 'info' | 'warning' | 'success' | interface RequestContext { retry: number + bearerRetried?: boolean } interface ErrorHandlerConfig { @@ -26,7 +27,12 @@ export class ErrorHandler { account: ManagedAccount, context: RequestContext, showToast: ToastFunction - ): Promise<{ shouldRetry: boolean; newContext?: RequestContext; switchAccount?: boolean }> { + ): Promise<{ + shouldRetry: boolean + newContext?: RequestContext + switchAccount?: boolean + forceRefresh?: boolean + }> { const readBody = async (): Promise => { try { const body = JSON.parse(await response.clone().text()) @@ -114,10 +120,20 @@ export class ErrorHandler { errorReason = 'Account Suspended' isPermanent = true } - if ( + const isBearerInvalid = errorReason.includes('bearer token included in the request is invalid') || errorReason.includes('The bearer token included in the request is invalid') - ) { + + if (isBearerInvalid && !context.bearerRetried) { + showToast('403: Bearer token stale after idle. Refreshing and retrying...', 'warning') + return { + shouldRetry: true, + newContext: { ...context, retry: context.retry + 1, bearerRetried: true }, + forceRefresh: true + } + } + + if (isBearerInvalid) { isPermanent = true } if (isPermanent) { diff --git a/src/core/request/request-handler.ts b/src/core/request/request-handler.ts index 7dcfdf4..727151a 100644 --- a/src/core/request/request-handler.ts +++ b/src/core/request/request-handler.ts @@ -88,6 +88,7 @@ export class RequestHandler { 20000 let retry = 0 + let bearerRetried = false let consecutiveNullAccounts = 0 const retryContext = this.retryStrategy.createContext() @@ -162,11 +163,24 @@ export class RequestHandler { } catch (e: any) { const httpStatus = e?.$metadata?.httpStatusCode - if (httpStatus) { - if (apiTimestamp) { - this.logSdkError(sdkPrep, e, acc, apiTimestamp) + if (httpStatus && apiTimestamp) { + this.logSdkError(sdkPrep, e, acc, apiTimestamp) + } + + if (httpStatus === 403 && !bearerRetried) { + const msg = e?.message || '' + if ( + msg.includes('bearer token included in the request is invalid') || + msg.includes('The bearer token included in the request is invalid') + ) { + bearerRetried = true + logger.warn('403 bearer invalid on first attempt, forcing token refresh and retrying') + await this.tokenRefresher.forceRefresh(acc, this.accountManager.toAuthDetails(acc)) + continue } + } + if (httpStatus) { const mockResponse = new Response( JSON.stringify({ message: e.message, __type: e.name }), { @@ -180,13 +194,17 @@ export class RequestHandler { e, mockResponse, acc, - { retry }, + { retry, bearerRetried }, showToast ) if (errorResult.shouldRetry) { if (errorResult.newContext) { retry = errorResult.newContext.retry + bearerRetried = errorResult.newContext.bearerRetried ?? bearerRetried + } + if (errorResult.forceRefresh) { + await this.tokenRefresher.forceRefresh(acc, this.accountManager.toAuthDetails(acc)) } if (errorResult.switchAccount) { continue