fix: Add optional ARN to AWS Bedrock connection to assume role when u…#4396
fix: Add optional ARN to AWS Bedrock connection to assume role when u…#4396mcjraquel wants to merge 4 commits into
Conversation
dosubot
Bot
added
the
size:M
|
@mcjraquel is attempting to deploy a commit to the agenta projects Team on Vercel. A member of the Team first needs to authorize it. |
|
|
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
📝 WalkthroughSummary by CodeRabbit
WalkthroughThis PR introduces AWS IAM role assumption credential resolution across the Agenta platform. The backend adds a helper function that detects role ARNs and uses boto3 STS to obtain temporary credentials, integrates it into two LLM execution paths, and includes comprehensive unit tests. The frontend extends provider types and configuration UI to capture role ARN input and maps it through form initialization and secret transforms. ChangesAWS Role Assumption Credential Resolution
🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
sdks/python/oss/tests/pytest/unit/test_resolve_aws_credentials.py (1)
195-232: ⚡ Quick winAdd regression coverage for forwarding source
AWS_SESSION_TOKENCurrent forwarding tests cover key/secret but not source session token. Please add a case asserting
aws_session_token(and/orAWS_SESSION_TOKEN) is passed intoboto3.client(...).🧪 Suggested test addition
+def test_base_session_token_forwarded_to_sts(): + sts = _mock_sts() + settings = { + "aws_role_arn": _ROLE_ARN, + "aws_access_key_id": "BASE_KEY", + "aws_secret_access_key": "BASE_SECRET", + "aws_session_token": "BASE_TOKEN", + } + + with patch("boto3.client", return_value=sts) as mock_client: + _resolve_aws_credentials(settings) + + mock_client.assert_called_once_with( + "sts", + aws_access_key_id="BASE_KEY", + aws_secret_access_key="BASE_SECRET", + aws_session_token="BASE_TOKEN", + region_name="us-east-1", + )
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 4fd33654-96e0-4922-9315-d8c180805f5e
⛔ Files ignored due to path filters (3)
api/uv.lockis excluded by!**/*.locksdks/python/uv.lockis excluded by!**/*.lockservices/uv.lockis excluded by!**/*.lock
📒 Files selected for processing (7)
sdks/python/agenta/sdk/engines/running/handlers.pysdks/python/oss/tests/pytest/unit/test_resolve_aws_credentials.pysdks/python/pyproject.tomlweb/oss/src/components/ModelRegistry/Drawers/ConfigureProviderDrawer/assets/ConfigureProviderDrawerContent.tsxweb/oss/src/components/ModelRegistry/Drawers/ConfigureProviderDrawer/assets/constants.tsweb/packages/agenta-entities/src/secret/core/transforms.tsweb/packages/agenta-shared/src/types/llmProvider.ts
|
Hi @mcjraquel, Thank you for your contribution. While we review it, could you please take a look at the QA agent's comments. Cheers, |
|
Hi, @junaway. I can work on the issue you linked. |
3 participants
Summary
Added AWS IAM role assumption support for Bedrock providers, enabling users to supply a Role ARN as an alternative to static credentials.
- SDK: new
_resolve_aws_credentials()inhandlers.py— whenaws_role_arnis present, calls AWS STSassume_roleand replaces any static credentials with the short-lived session token beforeforwarding to litellm. Wired into both
auto_ai_critique_v0and_run_prompt_llm_config_with_retry. Addsboto3as a production dependency.- Frontend: adds an optional Role ARN field to the Configure Provider drawer for Bedrock / Bedrock Converse / SageMaker; fixes a bug where
required: falsefields were forced required when a knownprovider was selected; wires
roleArn :left_right_arrow: aws_role_arnthroughLlmProvidertype and secret transforms.## Testing
### Verified locally
-
uv run pytest sdks/python/oss/tests/pytest/unit/ -v— all 32 unit tests pass-
ruff format+ruff check --fix— clean insdks/python/andservices/-
pnpm lint-fix— no ESLint or TypeScript errors across all 11 web packages### Added or updated tests
New:
sdks/python/oss/tests/pytest/unit/test_resolve_aws_credentials.py— 15 unit tests covering:- No-op path (no ARN → settings returned unchanged)
- Both
aws_role_arnandAWS_ROLE_ARNkey casings trigger STSassume_role- Role ARN and uppercase
AWS_*credential keys removed from result- Session token injected from STS response
- Region defaults to
us-east-1; resolved fromaws_region_name,aws_region, orAWS_REGION- Base credentials forwarded to the STS client constructor
- Original dict not mutated
### QA follow-up
- Open Configure Provider drawer, select Bedrock — verify Role ARN field renders without
*and form submits without a value- Supply a valid Role ARN — verify credentials are assumed correctly end-to-end against a real AWS account
## Demo
N/A — backend credential resolution; no visible UI change beyond the new optional Role ARN input field in the drawer.
## Checklist
- [x] I have included a video or screen recording for UI changes, or marked Demo as N/A
- [x] Relevant tests pass locally
- [x] Relevant linting and formatting pass locally
- [x] I have signed the CLA, or I will sign it when the bot prompts me
Contributor Resources