add optional session parameter to CodeInterpreter for custom credential management#41
Merged
siwachabhi merged 3 commits intoaws:mainfrom Sep 9, 2025
Conversation
…edential management Add boto3.Session injection support to CodeInterpreter constructor to enable custom credential configurations and role assumption while maintaining full backward compatibility. Features: - Optional session parameter in CodeInterpreter constructor - Support for pre-configured boto3 sessions with custom credentials - Enable AWS role assumption and cross-account access - Session reuse across multiple service clients - Full backward compatibility with existing implementations Use Cases: - Dynamic role switching for multi-tenant applications - Custom credential configurations beyond environment variables - Cross-account resource access with assumed roles - Containerized environments with temporary credentials - Testing scenarios with mock sessions - Fine-grained credential management and security Configuration: - Compatible with all existing CodeInterpreter functionality Implementation: - Updated CodeInterpreter.__init__() to accept optional session parameter - Modified code_session() context manager to support session injection - Session validation and client creation logic preserved - Comprehensive test coverage for both default and custom session scenarios - Updated documentation with usage examples and best practices Benefits: - Enhanced credential flexibility for enterprise applications - Improved security through controlled session management - Better testing capabilities with session mocking - Simplified multi-account and role-based access patterns - Maintains existing API contract and behavior Breaking Changes: None - Default behavior remains identical to previous implementation - All existing code continues to work without modification - New functionality only active when session parameter is provided CLOSES aws#40
Welcome to Codecov 🎉Once you merge this PR into your default branch, you're all set! Codecov will compare coverage reports and display results in all future pull requests. Thanks for integrating Codecov - We've got you covered ☂️ |
siwachabhi
approved these changes
Sep 9, 2025
sundargthb
pushed a commit
that referenced
this pull request
Dec 2, 2025
* feat: add WebSocket support with @app.websocket decorator Add bidirectional streaming support for agent invocations: - Add @app.websocket decorator for registering WebSocket handlers - Register /ws endpoint in BedrockAgentCoreApp constructor - Implement _handle_websocket method with proper error handling - Follow existing decorator patterns (entrypoint, ping) for consistency - Support WebSocketDisconnect for graceful connection handling - Integrate with existing RequestContext for session management test: add unit and integration tests for WebSocket support Add comprehensive test coverage for WebSocket decorator: Unit tests (10 tests): - WebSocket route initialization and registration - Decorator functionality and handler storage - Basic send/receive communication - Context integration with session IDs - Exception handling and error cases - Multiple message handling - Graceful disconnect handling - Custom request headers via context - Streaming data functionality Integration test: - End-to-end WebSocket echo server - Streaming multiple messages - Session ID propagation through headers - Real WebSocket client connection testing All tests pass successfully. * fix: remove exception details from WebSocket close reason Don't send exception messages to clients when closing WebSocket connections with code 1011. This prevents leaking internal error details to clients. Changes: - Remove reason parameter from websocket.close() calls - Exception details are still logged server-side - Tests continue to pass as they only check for disconnect exceptions --------- Co-authored-by: Abishek Kumar <kumabish@amazon.com>
Abishek10
added a commit
that referenced
this pull request
Dec 2, 2025
* feat: add WebSocket support with @app.websocket decorator (#41) * feat: add WebSocket support with @app.websocket decorator Add bidirectional streaming support for agent invocations: - Add @app.websocket decorator for registering WebSocket handlers - Register /ws endpoint in BedrockAgentCoreApp constructor - Implement _handle_websocket method with proper error handling - Follow existing decorator patterns (entrypoint, ping) for consistency - Support WebSocketDisconnect for graceful connection handling - Integrate with existing RequestContext for session management test: add unit and integration tests for WebSocket support Add comprehensive test coverage for WebSocket decorator: Unit tests (10 tests): - WebSocket route initialization and registration - Decorator functionality and handler storage - Basic send/receive communication - Context integration with session IDs - Exception handling and error cases - Multiple message handling - Graceful disconnect handling - Custom request headers via context - Streaming data functionality Integration test: - End-to-end WebSocket echo server - Streaming multiple messages - Session ID propagation through headers - Real WebSocket client connection testing All tests pass successfully. * fix: remove exception details from WebSocket close reason Don't send exception messages to clients when closing WebSocket connections with code 1011. This prevents leaking internal error details to clients. Changes: - Remove reason parameter from websocket.close() calls - Exception details are still logged server-side - Tests continue to pass as they only check for disconnect exceptions --------- Co-authored-by: Abishek Kumar <kumabish@amazon.com> * feat: add AgentRuntimeClient for WebSocket authentication (#42) * feat: add AgentRuntimeClient for WebSocket authentication - Implemented AgentRuntimeClient with methods for generating WebSocket credentials - Added generate_ws_connection() for backend services (returns URL + SigV4 headers) - Added generate_presigned_url() for frontend clients (returns presigned URL) - Support for endpoint qualifiers and custom query parameters (presigned URLs only) - Comprehensive unit tests with 22 test cases covering all functionality - Integration tests for validating credential format - Complete documentation with usage examples 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * fix: move websockets to dev dependencies for CI The websockets package was in optional dependencies which caused CI test failures. Moving it to dev dependencies ensures it's installed during CI test runs. * fix: skip integration tests when AWS credentials unavailable Add skipif decorator to integration tests that require AWS credentials. This prevents CI failures when credentials are not configured. * refactor: use mocked credentials instead of skipping tests Replace skipif decorators with pytest fixture that mocks boto3 credentials using botocore.credentials.Credentials class. This ensures integration tests run in all environments (local and CI) while still validating URL and header generation logic. Benefits: - Tests always run, providing better coverage - Uses real Credentials class for proper SigV4 signing - No need for AWS credentials to be configured * refactor: rename AgentRuntimeClient to AgentCoreRuntimeClient - Rename class from AgentRuntimeClient to AgentCoreRuntimeClient - Rename module from agent_runtime_client.py to agent_core_runtime_client.py - Update all unit and integration tests to use new class name - Update __init__.py exports - Update documentation examples This provides a more accurate name that reflects the AgentCore Runtime service. --------- Co-authored-by: Abishek Kumar <kumabish@amazon.com> Co-authored-by: Claude <noreply@anthropic.com> * chore: exclude integration testing workflow changes * fix: add websockets to integration test --------- Co-authored-by: Abishek K <Abishek10@users.noreply.github.com> Co-authored-by: Abishek Kumar <kumabish@amazon.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Sundar Raghavan <sdraghav@amazon.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add boto3.Session injection support to CodeInterpreter constructor to enable custom credential configurations and role assumption while maintaining full backward compatibility.
Features:
Use Cases:
Configuration:
Implementation:
Benefits:
Breaking Changes: None
CLOSES #40