From 33da207e7e2c5619fade26352a97c6d2646c9dd4 Mon Sep 17 00:00:00 2001 From: Roo Code Date: Wed, 19 Nov 2025 21:45:47 +0000 Subject: [PATCH] fix: refresh AWS Bedrock credentials automatically on expiration - Add credential refresh logic to handle expired token errors - Implement retry mechanism with automatic client recreation - Support credential refresh for profile-based authentication - Add comprehensive tests for credential refresh functionality - Fixes issue where long-running orchestrator sessions would fail after AWS temporary credentials expired Closes #9408 --- .../bedrock-credential-refresh.spec.ts | 227 ++++++++++++++++++ src/api/providers/bedrock.ts | 132 +++++++++- 2 files changed, 353 insertions(+), 6 deletions(-) create mode 100644 src/api/providers/__tests__/bedrock-credential-refresh.spec.ts diff --git a/src/api/providers/__tests__/bedrock-credential-refresh.spec.ts b/src/api/providers/__tests__/bedrock-credential-refresh.spec.ts new file mode 100644 index 0000000000..b0bb4fe70b --- /dev/null +++ b/src/api/providers/__tests__/bedrock-credential-refresh.spec.ts @@ -0,0 +1,227 @@ +import { vi, describe, it, expect, beforeEach } from "vitest" +import { AwsBedrockHandler } from "../bedrock" +import { BedrockRuntimeClient, ConverseStreamCommand, ConverseCommand } from "@aws-sdk/client-bedrock-runtime" +import { fromIni } from "@aws-sdk/credential-providers" +import { ProviderSettings } from "@roo-code/types" + +// Mock AWS SDK credential providers +vi.mock("@aws-sdk/credential-providers", () => { + const mockFromIni = vi.fn() + return { fromIni: mockFromIni } +}) + +// Mock BedrockRuntimeClient +vi.mock("@aws-sdk/client-bedrock-runtime", () => { + const mockSend = vi.fn() + const BedrockRuntimeClient = vi.fn().mockImplementation(() => ({ + send: mockSend, + })) + const ConverseStreamCommand = vi.fn() + const ConverseCommand = vi.fn() + + return { + BedrockRuntimeClient, + ConverseStreamCommand, + ConverseCommand, + } +}) + +// Mock logger to suppress log output during tests +vi.mock("../../../utils/logging", () => ({ + logger: { + info: vi.fn(), + error: vi.fn(), + warn: vi.fn(), + debug: vi.fn(), + }, +})) + +describe("AwsBedrockHandler - Credential Refresh", () => { + let handler: AwsBedrockHandler + let mockSend: any + let mockFromIni: any + + beforeEach(() => { + vi.clearAllMocks() + + // Get the mocked functions + mockFromIni = vi.mocked(fromIni) + mockSend = vi.fn() + + // Setup BedrockRuntimeClient mock + vi.mocked(BedrockRuntimeClient).mockImplementation( + () => + ({ + send: mockSend, + config: { region: "us-east-1" }, + }) as any, + ) + + // Setup fromIni mock to return fresh credentials + let credentialCallCount = 0 + mockFromIni.mockImplementation(() => { + credentialCallCount++ + return { + accessKeyId: `profile-access-key-${credentialCallCount}`, + secretAccessKey: `profile-secret-key-${credentialCallCount}`, + } + }) + }) + + it("should refresh credentials when receiving expired token error on streaming", async () => { + // Setup handler with profile-based auth + const options: ProviderSettings = { + apiModelId: "anthropic.claude-3-sonnet-20240229-v1:0", + awsRegion: "us-east-1", + awsUseProfile: true, + awsProfile: "test-profile", + } + + handler = new AwsBedrockHandler(options) + + // First call fails with expired token error + const expiredError = new Error("The security token included in the request is expired") + mockSend.mockRejectedValueOnce(expiredError) + + // Second call succeeds with valid stream + const mockStream = { + stream: (async function* () { + yield { messageStart: { role: "assistant" } } + yield { contentBlockStart: { start: { text: "Hello" } } } + yield { contentBlockDelta: { delta: { text: " world" } } } + yield { metadata: { usage: { inputTokens: 10, outputTokens: 5 } } } + yield { messageStop: { stopReason: "end_turn" } } + })(), + } + mockSend.mockResolvedValueOnce(mockStream) + + // Execute createMessage + const systemPrompt = "You are a helpful assistant" + const messages = [{ role: "user" as const, content: "Hello" }] + + const chunks: any[] = [] + for await (const chunk of handler.createMessage(systemPrompt, messages)) { + chunks.push(chunk) + } + + // Verify that the client was recreated with fresh credentials + expect(mockFromIni).toHaveBeenCalledTimes(2) // Initial creation + refresh + expect(mockSend).toHaveBeenCalledTimes(2) // First failed attempt + successful retry + expect(chunks).toContainEqual(expect.objectContaining({ type: "text", text: "Hello" })) + expect(chunks).toContainEqual(expect.objectContaining({ type: "text", text: " world" })) + expect(chunks).toContainEqual(expect.objectContaining({ type: "usage", inputTokens: 10, outputTokens: 5 })) + }) + + it("should refresh credentials when receiving expired token error on completePrompt", async () => { + // Setup handler with profile-based auth + const options: ProviderSettings = { + apiModelId: "anthropic.claude-3-sonnet-20240229-v1:0", + awsRegion: "us-east-1", + awsUseProfile: true, + awsProfile: "test-profile", + } + + handler = new AwsBedrockHandler(options) + + // First call fails with expired token error + const expiredError = new Error("Token has expired") + mockSend.mockRejectedValueOnce(expiredError) + + // Second call succeeds + mockSend.mockResolvedValueOnce({ + output: { + message: { + content: [{ text: "Test response" }], + }, + }, + }) + + // Execute completePrompt + const result = await handler.completePrompt("Test prompt") + + // Verify that the client was recreated with fresh credentials + expect(mockFromIni).toHaveBeenCalledTimes(2) // Initial creation + refresh + expect(mockSend).toHaveBeenCalledTimes(2) // First failed attempt + successful retry + expect(result).toBe("Test response") + }) + + it("should not refresh credentials for non-credential errors", async () => { + // Setup handler with profile-based auth + const options: ProviderSettings = { + apiModelId: "anthropic.claude-3-sonnet-20240229-v1:0", + awsRegion: "us-east-1", + awsUseProfile: true, + awsProfile: "test-profile", + } + + handler = new AwsBedrockHandler(options) + + // Call fails with a different error + const otherError = new Error("Service unavailable") + mockSend.mockRejectedValueOnce(otherError) + + // Execute completePrompt and expect it to throw + await expect(handler.completePrompt("Test prompt")).rejects.toThrow("Request was throttled") + + // Verify that the client was not recreated + expect(mockFromIni).toHaveBeenCalledTimes(1) // Only initial creation + expect(mockSend).toHaveBeenCalledTimes(1) // Only one attempt + }) + + it("should not refresh credentials when using direct credentials", async () => { + // Setup handler with direct credentials (not profile-based) + const options: ProviderSettings = { + apiModelId: "anthropic.claude-3-sonnet-20240229-v1:0", + awsRegion: "us-east-1", + awsAccessKey: "direct-access-key", + awsSecretKey: "direct-secret-key", + awsSessionToken: "direct-session-token", + } + + handler = new AwsBedrockHandler(options) + + // Call fails with expired token error + const expiredError = new Error("The security token included in the request is expired") + mockSend.mockRejectedValueOnce(expiredError) + + // Execute completePrompt and expect it to throw (no retry for direct credentials) + await expect(handler.completePrompt("Test prompt")).rejects.toThrow("AWS credentials have expired") + + // Verify that the client was not recreated (fromIni not called at all since using direct creds) + expect(mockFromIni).toHaveBeenCalledTimes(0) + expect(mockSend).toHaveBeenCalledTimes(1) // Only one attempt + }) + + it("should handle multiple consecutive expired token errors", async () => { + // Setup handler with profile-based auth + const options: ProviderSettings = { + apiModelId: "anthropic.claude-3-sonnet-20240229-v1:0", + awsRegion: "us-east-1", + awsUseProfile: true, + awsProfile: "test-profile", + } + + handler = new AwsBedrockHandler(options) + + // First two calls fail with expired token error + const expiredError = new Error("The security token included in the request is expired") + mockSend.mockRejectedValueOnce(expiredError) + mockSend.mockRejectedValueOnce(expiredError) + + // Third call would succeed, but we shouldn't get there (max retries = 1) + mockSend.mockResolvedValueOnce({ + output: { + message: { + content: [{ text: "Test response" }], + }, + }, + }) + + // Execute completePrompt and expect it to throw after max retries + await expect(handler.completePrompt("Test prompt")).rejects.toThrow("AWS credentials have expired") + + // Verify that we only tried once to refresh + expect(mockFromIni).toHaveBeenCalledTimes(2) // Initial creation + one refresh + expect(mockSend).toHaveBeenCalledTimes(2) // Initial attempt + one retry + }) +}) diff --git a/src/api/providers/bedrock.ts b/src/api/providers/bedrock.ts index faaee0360f..ede6ca7586 100644 --- a/src/api/providers/bedrock.ts +++ b/src/api/providers/bedrock.ts @@ -219,6 +219,14 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH this.costModelConfig = this.getModel() + // Initialize the client + this.client = this.createClient() + } + + /** + * Creates a new BedrockRuntimeClient with fresh credentials + */ + private createClient(): BedrockRuntimeClient { const clientConfig: BedrockRuntimeClientConfig = { userAgentAppId: `RooCode#${Package.version}`, region: this.options.awsRegion, @@ -238,6 +246,7 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH } } else if (this.options.awsUseProfile && this.options.awsProfile) { // Use profile-based credentials if enabled and profile is set + // Force fresh credentials every time to handle credential rotation clientConfig.credentials = fromIni({ profile: this.options.awsProfile, ignoreCache: true, @@ -251,7 +260,19 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH } } - this.client = new BedrockRuntimeClient(clientConfig) + return new BedrockRuntimeClient(clientConfig) + } + + /** + * Refreshes the client with new credentials + */ + private refreshClient(): void { + logger.info("Refreshing BedrockRuntimeClient with new credentials", { + ctx: "bedrock", + useProfile: this.options.awsUseProfile, + profile: this.options.awsProfile, + }) + this.client = this.createClient() } // Helper to guess model info from custom modelId string if not in bedrockModels @@ -413,11 +434,36 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH ) const command = new ConverseStreamCommand(payload) - const response = await this.client.send(command, { - abortSignal: controller.signal, - }) + let response + let retryCount = 0 + const maxRetries = 1 - if (!response.stream) { + // Try to send command with retry on credential expiration + while (retryCount <= maxRetries) { + try { + response = await this.client.send(command, { + abortSignal: controller.signal, + }) + break // Success, exit retry loop + } catch (sendError: unknown) { + // Check if this is an expired credential error + if (this.isExpiredCredentialError(sendError)) { + if (retryCount < maxRetries && (this.options.awsUseProfile || this.options.awsUseApiKey)) { + logger.info("Detected expired credentials, refreshing and retrying", { + ctx: "bedrock", + retryCount: retryCount + 1, + }) + this.refreshClient() + retryCount++ + continue // Retry with fresh credentials + } + } + // Not a credential error or max retries reached + throw sendError + } + } + + if (!response || !response.stream) { clearTimeout(timeoutId) throw new Error("No stream available in the response") } @@ -669,7 +715,32 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH } const command = new ConverseCommand(payload) - const response = await this.client.send(command) + let response + let retryCount = 0 + const maxRetries = 1 + + // Try to send command with retry on credential expiration + while (retryCount <= maxRetries) { + try { + response = await this.client.send(command) + break // Success, exit retry loop + } catch (sendError: unknown) { + // Check if this is an expired credential error + if (this.isExpiredCredentialError(sendError)) { + if (retryCount < maxRetries && (this.options.awsUseProfile || this.options.awsUseApiKey)) { + logger.info("Detected expired credentials in completePrompt, refreshing and retrying", { + ctx: "bedrock", + retryCount: retryCount + 1, + }) + this.refreshClient() + retryCount++ + continue // Retry with fresh credentials + } + } + // Not a credential error or max retries reached + throw sendError + } + } if ( response?.output?.message?.content && @@ -1088,6 +1159,33 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH * *************************************************************************************/ + /** + * Checks if an error is due to expired credentials + */ + private isExpiredCredentialError(error: unknown): boolean { + if (!(error instanceof Error)) { + return false + } + + const errorMessage = error.message.toLowerCase() + const errorName = error.name.toLowerCase() + + // Common patterns for expired credential errors + const expiredPatterns = [ + "the security token included in the request is expired", + "expired token", + "token has expired", + "expiredsignature", + "expiredtoken", + "security token expired", + "invalid security token", + "credential expired", + ] + + // Check for expired credential patterns + return expiredPatterns.some((pattern) => errorMessage.includes(pattern) || errorName.includes(pattern)) + } + /** * Error type definitions for Bedrock API errors */ @@ -1120,6 +1218,27 @@ Please verify: 3. The account ID in the ARN is correct`, logLevel: "error", }, + EXPIRED_CREDENTIALS: { + patterns: [ + "the security token included in the request is expired", + "expired token", + "token has expired", + "expiredsignature", + "expiredtoken", + "security token expired", + "invalid security token", + "credential expired", + ], + messageTemplate: `AWS credentials have expired. + +For profile-based authentication: +- The system will attempt to refresh credentials automatically +- If using temporary credentials, ensure they are being refreshed externally + +For direct credentials: +- Update your AWS access key, secret key, and session token in settings`, + logLevel: "error", + }, THROTTLING: { patterns: [ "throttl", @@ -1275,6 +1394,7 @@ Please check: // Check each error type's patterns in order of specificity (most specific first) const errorTypeOrder = [ + "EXPIRED_CREDENTIALS", // Check first to handle and retry appropriately "SERVICE_QUOTA_EXCEEDED", // Most specific - check before THROTTLING "MODEL_NOT_READY", "TOO_MANY_TOKENS",