-
Notifications
You must be signed in to change notification settings - Fork 2.6k
feat: enable prompt caching detection for AWS Bedrock custom ARNs #6666
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
- Add new message types for requesting Bedrock model capabilities - Implement backend handler to parse ARN and return model capabilities - Create useBedrockModelCapabilities hook to fetch capabilities from backend - Update useSelectedModel to use dynamic capabilities instead of hardcoded values - Add comprehensive tests for the new functionality Fixes #6429
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Reviewing my own code is like debugging in production - technically possible but morally questionable.
| }) | ||
| } | ||
| } catch (error) { | ||
| provider.log(`Error getting Bedrock model capabilities: ${error}`) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Consider enhancing the error logging to include the ARN for better debugging:
| provider.log(`Error getting Bedrock model capabilities: ${error}`) | |
| provider.log(`Error getting Bedrock model capabilities for ARN ${message.values.customArn}: ${error}`) |
| // Listen for response | ||
| const handler = (event: MessageEvent) => { | ||
| const message = event.data | ||
| if (message.type === "bedrockModelCapabilities" && message.values?.customArn === customArn) { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The hook doesn't handle potential race conditions if the ARN changes rapidly. Consider adding a cleanup mechanism to ignore responses from outdated requests. You could use a ref to track the current request or implement request cancellation.
| import type { ModelInfo } from "@roo-code/types" | ||
|
|
||
| export function useBedrockModelCapabilities(customArn?: string): ModelInfo | undefined { | ||
| const [capabilities, setCapabilities] = useState<ModelInfo | undefined>(undefined) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Consider exposing a loading state from this hook so the UI can show appropriate feedback while capabilities are being fetched. This would improve the user experience during the async operation.
|
|
||
| // Listen for response | ||
| const handler = (event: MessageEvent) => { | ||
| const message = event.data |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
For better type safety, consider defining a specific type for the message event data instead of relying on implicit any:
interface BedrockCapabilitiesMessage {
type: string;
values?: {
customArn: string;
modelInfo?: ModelInfo;
error?: string;
};
}|
Closing for now as I don't think this is the proper way to implement this |
Summary
This PR fixes the issue where prompt caching capabilities are not properly detected when using AWS Bedrock Application Inference Profiles (AIPs) with custom ARNs.
Problem
When using a custom ARN for an AIP in Bedrock, the frontend hardcodes the model capabilities to
supportsPromptCache: false, regardless of what the underlying model actually supports. This happens because the frontend doesn't have access to the backend's ARN parsing logic.Solution
useBedrockModelCapabilities) that:Testing
Related Issue
Fixes #6429
Important
Enable dynamic detection and caching of AWS Bedrock model capabilities for custom ARNs, replacing hardcoded values.
requestBedrockModelCapabilitiesmessage handling inwebviewMessageHandler.tsto dynamically fetch model capabilities for custom ARNs.useBedrockModelCapabilitieshook to request and cache capabilities from the backend.useSelectedModelto utilize dynamic capabilities for Bedrock custom ARNs.useBedrockModelCapabilitiesinuseBedrockModelCapabilities.spec.ts.bedrockModelCapabilitiestoExtensionMessageandWebviewMessagetypes.This description was created by
for 626fd45. You can customize this summary. It will automatically update as commits are pushed.