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
2026-09-05 08:10:14 +00:00 · 2025-11-19 21:45:47 +00:00 · 2025-11-19 21:45:47 +00:00 · 33da207e7e
commit 33da207e7e
parent e618d881d6
2 changed files with 353 additions and 6 deletions
--- a/src/api/providers/tests/bedrock-credential-refresh.spec.ts
+++ 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
+	})
+})
--- 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",