Cost display updating for Bedrock custom ARNs that are prompt routers

[Smartsheet-JB-Brown]

Fix AWS Bedrock Provider with Cost Calculations for Prompt Router ARNs

Note - 2nd PR coming in a couple hours to make AWS prices in src/shared/api.ts current. These 2 PRs together will make Roo's cost display more accurate for Bedrock users.

Overview

This PR fixes the AWS Bedrock provider to properly handle and utilize the invokedModelId information returned by AWS Bedrock Prompt Router responses. The code has been updated to use the token cost information from the invokedModelId, because at request time the inference model, and therefore token costs, are not known for prompt router ARNs. This improvement ensures more accurate cost tracking and model configuration when using AWS Bedrock's Prompt Router feature.

Problem

When using AWS Bedrock's Prompt Router feature, the actual model used for inference is not known at request time. The Prompt Router returns the actual model used via the invokedModelId field in the response stream, and the token usage values are located in a different structure than responses from individual models. Previously, our implementation wasn't properly extracting and utilizing this information, which led to:

1. No cost tracking for prompt router ARNs
2. Incorrect model configuration for subsequent requests
3. Missing token usage metrics from Prompt Router responses

Changes

• Added support for extracting and parsing the invokedModelId from AWS Bedrock Prompt Router responses
• Implemented logic to update the model configuration based on the actual model used
• Enhanced region handling for cross-region inference with custom ARNs
• Added proper extraction of model names from ARNs with regional prefixes (e.g., "us." or "eu.")
• Improved error handling for custom ARN scenarios with more descriptive error messages
• Added comprehensive test coverage for the new functionality

Benefits

• Accurate Cost Tracking: Costs are now calculated based on the actual model used by the Prompt Router
• Improved Model Configuration: The provider now correctly updates its configuration based on the model selected by the Prompt Router
• Better Error Handling: More descriptive error messages for common issues with custom ARNs
• Enhanced Region Support: Better handling of cross-region inference scenarios

Testing

Added new test files and expanded existing tests:

• bedrock-invokedModelId.test.ts: Tests for proper handling of the invokedModelId field
• Enhanced bedrock-custom-arn.test.ts: Tests for custom ARN validation and region extraction
• Updated bedrock.test.ts: Added tests for invokedModelId handling in various scenarios

Technical Details

The key technical changes include:

1. Enhanced the StreamEvent interface to include the trace.promptRouter.invokedModelId field
2. Added logic to extract the model name from the invokedModelId ARN
3. Implemented the costModelConfig property to track the actual model used
4. Updated the getModel method to return the correct model configuration
5. Added special handling for region prefixes in model names (us./eu.)
6. Improved error messages for common ARN-related issues

This implementation ensures that when AWS Bedrock's Prompt Router selects a different model than initially requested, our system correctly identifies and uses that model's configuration for accurate cost tracking and proper functionality.

---

Important

Enhances AWS Bedrock provider to accurately handle cost calculations and model configurations for prompt router ARNs, with improved error handling and comprehensive testing.

• Behavior:
  • Updates AwsBedrockHandler in bedrock.ts to handle invokedModelId for accurate cost tracking and model configuration.
  • Adds logic to extract model name from invokedModelId and update costModelConfig.
  • Improves error handling for custom ARNs with descriptive messages.
• Models:
  • Updates AWS Bedrock model definitions in api.ts with latest pricing and new models.
  • Adds supportsComputerUse flag to relevant models.
• Testing:
  • Adds bedrock-invokedModelId.test.ts for invokedModelId handling.
  • Updates bedrock-createMessage.test.ts and bedrock.test.ts for custom ARN and cross-region inference scenarios.
• Misc:
  • Mocks logger in jest.setup.ts for consistent test logging.

^{This description was created by for d8df9a5. It will automatically update as commits are pushed.}

[changeset-bot]

⚠️ No Changeset found

Latest commit: 4ada518

Merging this PR will not cause a version bump for any packages. If these changes should not result in a new version, you're good to go. If these changes should result in a version bump, you need to add a changeset.

This PR includes no changesets

When changesets are added to this PR, you'll see the packages that this PR includes changesets for and the associated semver types

Click here to learn what changesets are, and how to add one.

Click here if you're a maintainer who wants to add a changeset to this PR

[mrubens]

Think we can clean up some of the logger.debugs? And if the tests that rely on them are important, update them to be based on the behavior or data?

[mrubens]

Can we revert the changes to this file and mock debug within specific tests if necessary?

[Smartsheet-JB-Brown]

yes, done

[mrubens]

Think we can revert this one as well?

[mrubens]

Looks good!