Skip to content

Commit 9bace87

Browse files
Abishek10Abishek Kumarclaude
authored and
Sundar Raghavan
committed
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 <[email protected]> * 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 <[email protected]> Co-authored-by: Claude <[email protected]>
1 parent 9317cfb commit 9bace87

File tree

7 files changed

+1975
-855
lines changed

7 files changed

+1975
-855
lines changed

‎docs/examples/agent_runtime_client_examples.md‎

Lines changed: 206 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,206 @@
1+
# AgentCoreRuntimeClient Examples
2+
3+
This document provides practical examples for using the `AgentCoreRuntimeClient` to authenticate WebSocket connections to AgentCore Runtime.
4+
5+
## Basic Usage
6+
7+
### Backend Service (SigV4 Headers)
8+
9+
```python
10+
from bedrock_agentcore.runtime import AgentCoreRuntimeClient
11+
import websockets
12+
import asyncio
13+
14+
async def main():
15+
# Initialize client
16+
client = AgentCoreRuntimeClient(region="us-west-2")
17+
18+
# Generate WebSocket connection with authentication
19+
ws_url, headers = client.generate_ws_connection(
20+
runtime_arn="arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/my-runtime"
21+
)
22+
23+
# Connect using any WebSocket library
24+
async with websockets.connect(ws_url, extra_headers=headers) as ws:
25+
# Send message
26+
await ws.send('{"inputText": "Hello!"}')
27+
28+
# Receive response
29+
response = await ws.recv()
30+
print(f"Received: {response}")
31+
32+
asyncio.run(main())
33+
```
34+
35+
### Frontend Client (Presigned URL)
36+
37+
```python
38+
from bedrock_agentcore.runtime import AgentCoreRuntimeClient
39+
40+
# Backend: Generate presigned URL
41+
client = AgentCoreRuntimeClient(region="us-west-2")
42+
43+
presigned_url = client.generate_presigned_url(
44+
runtime_arn="arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/my-runtime",
45+
expires=300 # 5 minutes
46+
)
47+
48+
# Share presigned_url with frontend
49+
# Frontend JavaScript: new WebSocket(presigned_url)
50+
```
51+
52+
## Advanced Usage
53+
54+
### With Endpoint Qualifier
55+
56+
```python
57+
client = AgentCoreRuntimeClient(region="us-west-2")
58+
59+
# For generate_ws_connection (header-based auth)
60+
ws_url, headers = client.generate_ws_connection(
61+
runtime_arn="arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/my-runtime",
62+
endpoint_name="DEFAULT"
63+
)
64+
# URL will include: ?qualifier=DEFAULT
65+
66+
# For generate_presigned_url (query-based auth)
67+
presigned_url = client.generate_presigned_url(
68+
runtime_arn="arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/my-runtime",
69+
endpoint_name="DEFAULT"
70+
)
71+
# URL will include: ?qualifier=DEFAULT&X-Amz-Algorithm=...
72+
```
73+
74+
### With Custom Query Parameters (Presigned URL only)
75+
76+
```python
77+
client = AgentCoreRuntimeClient(region="us-west-2")
78+
79+
# custom_headers parameter is only available for presigned URLs
80+
presigned_url = client.generate_presigned_url(
81+
runtime_arn="arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/my-runtime",
82+
custom_headers={"custom_param": "value", "another": "param"}
83+
)
84+
85+
# URL will include: ?custom_param=value&another=param&X-Amz-Algorithm=...
86+
```
87+
88+
### With Explicit Session ID
89+
90+
```python
91+
client = AgentCoreRuntimeClient(region="us-west-2")
92+
93+
ws_url, headers = client.generate_ws_connection(
94+
runtime_arn="arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/my-runtime",
95+
session_id="my-custom-session-id"
96+
)
97+
```
98+
99+
## Error Handling
100+
101+
```python
102+
from bedrock_agentcore.runtime import AgentCoreRuntimeClient
103+
104+
client = AgentCoreRuntimeClient(region="us-west-2")
105+
106+
try:
107+
ws_url, headers = client.generate_ws_connection(
108+
runtime_arn="invalid-arn"
109+
)
110+
except ValueError as e:
111+
print(f"Invalid ARN format: {e}")
112+
except RuntimeError as e:
113+
print(f"AWS credentials error: {e}")
114+
```
115+
116+
## Custom Boto3 Session
117+
118+
You can provide your own boto3 session for custom credential handling:
119+
120+
```python
121+
import boto3
122+
from bedrock_agentcore.runtime import AgentCoreRuntimeClient
123+
124+
# Create a custom session with specific profile
125+
session = boto3.Session(profile_name="my-profile")
126+
127+
# Or with specific credentials
128+
session = boto3.Session(
129+
aws_access_key_id="YOUR_ACCESS_KEY",
130+
aws_secret_access_key="YOUR_SECRET_KEY",
131+
aws_session_token="YOUR_SESSION_TOKEN"
132+
)
133+
134+
# Initialize client with custom session
135+
client = AgentCoreRuntimeClient(region="us-west-2", session=session)
136+
137+
# Use the client normally
138+
ws_url, headers = client.generate_ws_connection(runtime_arn)
139+
```
140+
141+
## OAuth Authentication
142+
143+
For scenarios using OAuth bearer tokens instead of AWS credentials:
144+
145+
```python
146+
from bedrock_agentcore.runtime import AgentCoreRuntimeClient
147+
import websockets
148+
import asyncio
149+
150+
async def main():
151+
# Initialize client
152+
client = AgentCoreRuntimeClient(region="us-west-2")
153+
154+
# Your OAuth bearer token (e.g., from JWT authentication)
155+
bearer_token = "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9..."
156+
157+
# Generate WebSocket connection with OAuth
158+
ws_url, headers = client.generate_ws_connection_oauth(
159+
runtime_arn="arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/my-runtime",
160+
bearer_token=bearer_token,
161+
endpoint_name="DEFAULT" # Optional
162+
)
163+
164+
# Connect using OAuth authentication
165+
async with websockets.connect(ws_url, extra_headers=headers) as ws:
166+
await ws.send('{"inputText": "Hello!"}')
167+
response = await ws.recv()
168+
print(f"Received: {response}")
169+
170+
asyncio.run(main())
171+
```
172+
173+
### OAuth with Custom Session ID
174+
175+
```python
176+
client = AgentCoreRuntimeClient(region="us-west-2")
177+
178+
ws_url, headers = client.generate_ws_connection_oauth(
179+
runtime_arn="arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/my-runtime",
180+
bearer_token="your-oauth-token",
181+
session_id="custom-oauth-session-id"
182+
)
183+
```
184+
185+
## Using Different WebSocket Libraries
186+
187+
### With websockets library
188+
189+
```python
190+
import websockets
191+
192+
ws_url, headers = client.generate_ws_connection(runtime_arn)
193+
async with websockets.connect(ws_url, extra_headers=headers) as ws:
194+
await ws.send(message)
195+
```
196+
197+
### With aiohttp library
198+
199+
```python
200+
import aiohttp
201+
202+
ws_url, headers = client.generate_ws_connection(runtime_arn)
203+
async with aiohttp.ClientSession() as session:
204+
async with session.ws_connect(ws_url, headers=headers) as ws:
205+
await ws.send_str(message)
206+
```

‎pyproject.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ dependencies = [
3333
"starlette>=0.46.2",
3434
"typing-extensions>=4.13.2,<5.0.0",
3535
"uvicorn>=0.34.2",
36+
"pre-commit>=4.2.0",
3637
]
3738

3839
[project.scripts]

‎src/bedrock_agentcore/runtime/__init__.py‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,15 @@
66
- BedrockAgentCoreContext: Agent identity context
77
"""
88

9+
from .agent_core_runtime_client import AgentCoreRuntimeClient
910
from .app import BedrockAgentCoreApp
1011
from .context import BedrockAgentCoreContext, RequestContext
1112
from .models import PingStatus
1213

13-
__all__ = ["BedrockAgentCoreApp", "RequestContext", "BedrockAgentCoreContext", "PingStatus"]
14+
__all__ = [
15+
"AgentCoreRuntimeClient",
16+
"BedrockAgentCoreApp",
17+
"RequestContext",
18+
"BedrockAgentCoreContext",
19+
"PingStatus",
20+
]

0 commit comments

Comments
 (0)