Skip to content

add optional session parameter to CodeInterpreter for custom credential management#41

Merged
siwachabhi merged 3 commits intoaws:mainfrom
gonzalo123:feature/inyect_aws_session_to_code_interpreter_client
Sep 9, 2025
Merged

add optional session parameter to CodeInterpreter for custom credential management#41
siwachabhi merged 3 commits intoaws:mainfrom
gonzalo123:feature/inyect_aws_session_to_code_interpreter_client

Conversation

@gonzalo123
Copy link
Contributor

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 #40

…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
@codecov-commenter
Copy link

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 siwachabhi merged commit fa8b028 into aws:main Sep 9, 2025
16 of 17 checks passed
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add optional session parameter to CodeInterpreter for custom credential management

3 participants