Jbbrown/bedrock caching

locales/zh-TW/CONTRIBUTING.md

@@ -170,4 +170,4 @@ code --install-extension bin/roo-cline-<version>.vsix
 
 ## 貢獻協議
 
-透過提交 Pull Request，您同意您的貢獻將依照與專案相同的授權條款（[Apache 2.0](../LICENSE)）進行授權。
+透過提交 Pull Request，您同意您的貢獻將依照與專案相同的授權條款（[Apache 2.0](../LICENSE)）進行授權。

package.json

@@ -356,7 +356,7 @@
 		"@anthropic-ai/bedrock-sdk": "^0.10.2",
 		"@anthropic-ai/sdk": "^0.37.0",
 		"@anthropic-ai/vertex-sdk": "^0.7.0",
-		"@aws-sdk/client-bedrock-runtime": "^3.706.0",
+		"@aws-sdk/client-bedrock-runtime": "^3.779.0",
 		"@google-cloud/vertexai": "^1.9.3",
 		"@google/generative-ai": "^0.18.0",
 		"@mistralai/mistralai": "^1.3.6",

src/api/providers/__tests__/bedrock-custom-arn.test.ts

@@ -1,85 +1,254 @@
 import { AwsBedrockHandler } from "../bedrock"
 import { ApiHandlerOptions } from "../../../shared/api"
+import { logger } from "../../../utils/logging"
 
-// Mock the AWS SDK
+// Mock the logger
+jest.mock("../../../utils/logging", () => ({
+	logger: {
+		debug: jest.fn(),
+		info: jest.fn(),
+		warn: jest.fn(),
+		error: jest.fn(),
+		fatal: jest.fn(),
+		child: jest.fn().mockReturnValue({
+			debug: jest.fn(),
+			info: jest.fn(),
+			warn: jest.fn(),
+			error: jest.fn(),
+			fatal: jest.fn(),
+		}),
+	},
+}))
+
+// Mock AWS SDK
 jest.mock("@aws-sdk/client-bedrock-runtime", () => {
-	const mockResponse = {
-		output: {
-			message: {
-				content: [
-					{
-						text: "Test response",
-					},
-				],
-			},
+	const mockModule = {
+		lastCommandInput: null as Record<string, any> | null,
+		mockSend: jest.fn().mockImplementation(async function () {
+			return {
+				output: new TextEncoder().encode(JSON.stringify({ content: "Test response" })),
+			}
+		}),
+		mockConverseCommand: jest.fn(function (input) {
+			mockModule.lastCommandInput = input
+			return { input }
+		}),
+		MockBedrockRuntimeClient: class {
+			public config: any
+			public send: jest.Mock
+
+			constructor(config: { region?: string }) {
+				this.config = config
+				this.send = mockModule.mockSend
+			}
 		},
 	}
 
-	const mockSend = jest.fn().mockImplementation(() => {
-		return Promise.resolve(mockResponse)
-	})
-
 	return {
-		BedrockRuntimeClient: jest.fn().mockImplementation(() => ({
-			send: mockSend,
-			config: {
-				region: "us-east-1",
-			},
-		})),
-		ConverseCommand: jest.fn(),
+		BedrockRuntimeClient: mockModule.MockBedrockRuntimeClient,
+		ConverseCommand: mockModule.mockConverseCommand,
 		ConverseStreamCommand: jest.fn(),
+		__mock: mockModule, // Expose mock internals for testing
 	}
 })
 
-describe("AwsBedrockHandler with custom ARN", () => {
-	const mockOptions: ApiHandlerOptions = {
-		apiModelId: "custom-arn",
-		awsCustomArn: "arn:aws:bedrock:us-east-1:123456789012:foundation-model/anthropic.claude-3-sonnet-20240229-v1:0",
-		awsRegion: "us-east-1",
+// Get mock module for testing
+const bedrockMock = jest.requireMock("@aws-sdk/client-bedrock-runtime").__mock
+
+describe("Bedrock ARN Handling", () => {
+	// Helper function to create a handler with specific options
+	const createHandler = (options: Partial<ApiHandlerOptions> = {}) => {
+		const defaultOptions: ApiHandlerOptions = {
+			apiModelId: "anthropic.claude-3-sonnet-20240229-v1:0",
+			awsRegion: "us-east-1",
+			...options,
+		}
+		return new AwsBedrockHandler(defaultOptions)
 	}
 
-	it("should use the custom ARN as the model ID", async () => {
-		const handler = new AwsBedrockHandler(mockOptions)
-		const model = handler.getModel()
+	// Direct tests for parseArn function
+	describe("parseArn direct tests", () => {
+		it("should correctly extract modelType and modelId from foundation-model ARN", () => {
+			const handler = createHandler()
+			//note: properly formated foundation-model ARNs don't have an account id.
+			const arn = "arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-3-sonnet-20240229-v1:0"
 
-		expect(model.id).toBe(mockOptions.awsCustomArn)
-		expect(model.info).toHaveProperty("maxTokens")
-		expect(model.info).toHaveProperty("contextWindow")
-		expect(model.info).toHaveProperty("supportsPromptCache")
-	})
+			// Access the private method using type casting
+			const result = (handler as any).parseArn(arn, "us-east-1")
 
-	it("should extract region from ARN and use it for client configuration", () => {
-		// Test with matching region
-		const handler1 = new AwsBedrockHandler(mockOptions)
-		expect((handler1 as any).client.config.region).toBe("us-east-1")
+			// Verify the result contains the expected values
+			expect(result.isValid).toBe(true)
+			expect(result.modelType).toBe("foundation-model")
 
-		// Test with mismatched region
-		const mismatchOptions = {
-			...mockOptions,
-			awsRegion: "us-west-2",
-		}
-		const handler2 = new AwsBedrockHandler(mismatchOptions)
-		// Should use the ARN region, not the provided region
-		expect((handler2 as any).client.config.region).toBe("us-east-1")
-	})
+			//verify the id is not the ARN for foudation models, but the ID
+			expect(result.modelId).toBe("anthropic.claude-3-sonnet-20240229-v1:0")
+			expect(result.crossRegionInference).toBe(false)
+		})
 
-	it("should validate ARN format", async () => {
-		// Invalid ARN format
-		const invalidOptions = {
-			...mockOptions,
-			awsCustomArn: "invalid-arn-format",
-		}
+		it("should correctly extract modelType and modelId from inference-profile ARN", () => {
+			const handler = createHandler()
+			const arn = "arn:aws:bedrock:us-west-2:123456789012:inference-profile/custom-profile"
+
+			// Access the private method using type casting
+			const result = (handler as any).parseArn(arn, "us-west-2")
+
+			// Verify the result contains the expected values
+			expect(result.isValid).toBe(true)
+			// The region is not set in the result for normal ARNs
+			expect(result.modelType).toBe("inference-profile")
+			expect(result.modelId).toBe("custom-profile")
+			expect(result.crossRegionInference).toBe(false)
+		})
+
+		it("should correctly extract modelType and modelId from prompt-router ARN", () => {
+			const handler = createHandler()
+			const arn = "arn:aws:bedrock:eu-west-1:123456789012:prompt-router/custom-router-name"
+
+			// Access the private method using type casting
+			const result = (handler as any).parseArn(arn, "eu-west-1")
+
+			// Verify the result contains the expected values
+			expect(result.isValid).toBe(true)
+			// The region is not set in the result for normal ARNs
+			expect(result.modelType).toBe("prompt-router")
+			expect(result.modelId).toBe("custom-router-name")
+			expect(result.crossRegionInference).toBe(false)
+		})
+
+		it("should set crossRegionInference to true when a known region prefix is found in the model ID", () => {
+			const handler = createHandler()
+			const arn =
+				"arn:aws:bedrock:us-east-1:123456789012:foundation-model/us.anthropic.claude-3-sonnet-20240229-v1:0"
+
+			// Access the private method using type casting
+			const result = (handler as any).parseArn(arn, "us-east-1")
+
+			// Verify crossRegionInference is true
+			expect(result.crossRegionInference).toBe(true)
+			expect(result.modelId).toBe("anthropic.claude-3-sonnet-20240229-v1:0")
+			expect(result.region).toBe("us-east-1")
+		})
 
-		const handler = new AwsBedrockHandler(invalidOptions)
+		it("should set crossRegionInference to true for model IDs with apac (4 digit) region prefix", () => {
+			// This test uses a model ID with a region prefix in a different way
+			// We'll use a real ARN with a model ID that includes a region prefix
+			const handler = createHandler()
 
-		// completePrompt should throw an error for invalid ARN
-		await expect(handler.completePrompt("test")).rejects.toThrow("Invalid ARN format")
+			// Use a model ID with eu. prefix which should be detected
+			const arn =
+				"arn:aws:bedrock:ap-east-1:123456789012:foundation-model/apac.anthropic.claude-3-sonnet-20240229-v1:0"
+
+			// Access the private method using type casting
+			const result = (handler as any).parseArn(arn, "us-east-1")
+
+			// Verify crossRegionInference is true
+			expect(result.crossRegionInference).toBe(true)
+			// The eu. prefix should be removed from the model ID
+			expect(result.modelId).toBe("anthropic.claude-3-sonnet-20240229-v1:0")
+		})
+
+		it("should include region mismatch warning but still extract modelType and modelId", () => {
+			const handler = createHandler()
+			const arn = "arn:aws:bedrock:eu-west-1::foundation-model/anthropic.claude-3-sonnet-20240229-v1:0"
+
+			// Access the private method using type casting with mismatched region
+			const result = (handler as any).parseArn(arn, "us-east-1")
+
+			// Verify the result contains the expected values including error message
+			expect(result.isValid).toBe(true)
+			// In case of region mismatch, the region is set to the ARN region
+			expect(result.region).toBe("eu-west-1")
+			expect(result.modelType).toBe("foundation-model")
+			expect(result.modelId).toBe("anthropic.claude-3-sonnet-20240229-v1:0")
+			expect(result.errorMessage).toContain("Region mismatch")
+			expect(result.crossRegionInference).toBe(false)
+		})
+
+		it("should return isValid: false for simple ARN format", () => {
+			const handler = createHandler()
+			const arn = "arn:aws:bedrock:us-east-1:123456789012:some-other-resource"
+
+			// Access the private method using type casting
+			const result = (handler as any).parseArn(arn, "us-east-1")
+
+			// Verify the result indicates invalid ARN
+			expect(result.isValid).toBe(false)
+			expect(result.region).toBeUndefined()
+			expect(result.errorMessage).toContain("Invalid ARN format")
+			expect(result.crossRegionInference).toBe(false)
+			expect(result.modelType).toBeUndefined()
+			expect(result.modelId).toBeUndefined()
+		})
+
+		it("should return isValid: false for invalid ARN format", () => {
+			const handler = createHandler()
+			const arn = "invalid-arn-format"
+
+			// Access the private method using type casting
+			const result = (handler as any).parseArn(arn)
+
+			// Verify the result indicates invalid ARN
+			expect(result.isValid).toBe(false)
+			expect(result.region).toBeUndefined()
+			expect(result.errorMessage).toContain("Invalid ARN format")
+			expect(result.crossRegionInference).toBe(false)
+			expect(result.modelType).toBeUndefined()
+			expect(result.modelId).toBeUndefined()
+		})
 	})
 
-	it("should complete a prompt successfully with valid ARN", async () => {
-		const handler = new AwsBedrockHandler(mockOptions)
-		const response = await handler.completePrompt("test prompt")
+	// Integration tests for ARN handling in the constructor and other methods
+	describe("ARN handling in constructor and methods", () => {
+		it("should extract model ID from the custom ARN for foundation-model ARNs", async () => {
+			const mockOptions: ApiHandlerOptions = {
+				apiModelId: "custom-arn",
+				//properly formatted foundation-model ARNs don't have an account id
+				awsCustomArn: "arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-3-sonnet-20240229-v1:0",
+				awsRegion: "us-east-1",
+			}
+
+			const handler = new AwsBedrockHandler(mockOptions)
+			const model = handler.getModel()
+
+			// For foundation-model ARNs, the model ID is extracted from the ARN
+			expect(model.id).toBe("anthropic.claude-3-sonnet-20240229-v1:0")
+			expect(model.info).toHaveProperty("maxTokens")
+			expect(model.info).toHaveProperty("contextWindow")
+			expect(model.info).toHaveProperty("supportsPromptCache")
+		})
+
+		it("should extract region from ARN and use it for client configuration", () => {
+			// Test with ARN in eu-west-1 but config region in us-east-1
+			const handler = createHandler({
+				awsRegion: "us-east-1",
+				awsCustomArn:
+					"arn:aws:bedrock:eu-west-1:123456789012:inference-profile/anthropic.claude-3-sonnet-20240229-v1:0",
+			})
+
+			// Verify the client was created with the ARN region, not the provided region
+			expect((handler as any).client.config.region).toBe("eu-west-1")
+		})
+
+		it("should log region mismatch warning when ARN region differs from provided region", () => {
+			// Spy on logger.info which is called when there's a region mismatch
+			const infoSpy = jest.spyOn(logger, "info")
+
+			// Create handler with ARN region different from provided region
+			const arn =
+				"arn:aws:bedrock:eu-west-1:123456789012:inference-profile/anthropic.claude-3-sonnet-20240229-v1:0"
+			const handler = createHandler({
+				awsCustomArn: arn,
+				awsRegion: "us-east-1", // Different from ARN region
+			})
 
-		expect(response).toBe("Test response")
+			// Verify logger was called with region mismatch warning
+			expect(infoSpy).toHaveBeenCalledWith(
+				expect.stringContaining("Region mismatch"),
+				expect.objectContaining({
+					selectedRegion: "us-east-1",
+					arnRegion: "eu-west-1",
+				}),
+			)
+		})
 	})
 })