feat(runtime): Add session_id support to WebSocket connection methods (#186)

* feat(runtime): Add session_id support to WebSocket connection methods

Add session_id parameter to both generate_ws_connection and generate_presigned_url
methods in AgentCoreRuntimeClient to enable proper session tracking in WebSocket
connections.

Changes:
- generate_ws_connection: Add X-Amzn-Bedrock-AgentCore-Runtime-Session-Id header
- generate_ws_connection: Update User-Agent to AgentCoreRuntimeClient/1.0
- generate_presigned_url: Add session_id as query parameter via custom_headers
- Update unit tests to verify session_id in headers and URL params
- Update integration tests to verify session_id end-to-end

All tests passing (32 unit tests, 4 integration tests)

* docs(runtime): Update class name in docstring examples from AgentRuntimeClient to AgentCoreRuntimeClient

---------

Co-authored-by: Abishek Kumar <[email protected]>

Author: Abishek10

src/bedrock_agentcore/runtime/agent_core_runtime_client.py

@@ -159,7 +159,7 @@ def generate_ws_connection(
             ValueError: If runtime_arn format is invalid.
 
         Example:
-            >>> client = AgentRuntimeClient('us-west-2')
+            >>> client = AgentCoreRuntimeClient('us-west-2')
             >>> ws_url, headers = client.generate_ws_connection(
             ...     runtime_arn='arn:aws:bedrock-agentcore:us-west-2:123:runtime/my-runtime',
             ...     endpoint_name='DEFAULT'
@@ -209,11 +209,12 @@ def generate_ws_connection(
             "Host": host,
             "X-Amz-Date": request.headers["x-amz-date"],
             "Authorization": request.headers["Authorization"],
+            "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id": session_id,
             "Upgrade": "websocket",
             "Connection": "Upgrade",
             "Sec-WebSocket-Version": "13",
             "Sec-WebSocket-Key": base64.b64encode(secrets.token_bytes(16)).decode(),
-            "User-Agent": f"AgentRuntime-Client/1.0 (Session: {session_id})",
+            "User-Agent": "AgentCoreRuntimeClient/1.0",
         }
 
         # Add session token if present
@@ -256,7 +257,7 @@ def generate_presigned_url(
             RuntimeError: If URL generation fails or no credentials found.
 
         Example:
-            >>> client = AgentRuntimeClient('us-west-2')
+            >>> client = AgentCoreRuntimeClient('us-west-2')
             >>> presigned_url = client.generate_presigned_url(
             ...     runtime_arn='arn:aws:bedrock-agentcore:us-west-2:123:runtime/my-runtime',
             ...     endpoint_name='DEFAULT',
@@ -278,6 +279,11 @@ def generate_presigned_url(
             session_id = str(uuid.uuid4())
             self.logger.debug("Auto-generated session ID: %s", session_id)
 
+        # Add session_id to custom_headers (which become query params)
+        if custom_headers is None:
+            custom_headers = {}
+        custom_headers["X-Amzn-Bedrock-AgentCore-Runtime-Session-Id"] = session_id
+
         # Build WebSocket URL with query parameters
         ws_url = self._build_websocket_url(runtime_arn, endpoint_name, custom_headers)
 
@@ -344,7 +350,7 @@ def generate_ws_connection_oauth(
             ValueError: If runtime_arn format is invalid or bearer_token is empty.
 
         Example:
-            >>> client = AgentRuntimeClient('us-west-2')
+            >>> client = AgentCoreRuntimeClient('us-west-2')
             >>> ws_url, headers = client.generate_ws_connection_oauth(
             ...     runtime_arn='arn:aws:bedrock-agentcore:us-west-2:123:runtime/my-runtime',
             ...     bearer_token='eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9...',

tests/integration/runtime/test_agent_core_runtime_client_integration.py

@@ -56,6 +56,21 @@ def test_generate_ws_connection_returns_valid_format(self, mock_boto_session):
         assert "Authorization" in headers
         assert "X-Amz-Date" in headers
         assert "Host" in headers
+        assert "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id" in headers
+        assert "User-Agent" in headers
+        assert headers["User-Agent"] == "AgentCoreRuntimeClient/1.0"
+
+    def test_generate_ws_connection_with_session_id(self, mock_boto_session):
+        """Test that generate_ws_connection includes provided session ID in headers."""
+        client = AgentCoreRuntimeClient(region="us-west-2")
+        runtime_arn = "arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/test-runtime"
+        test_session_id = "integration-test-session-789"
+
+        ws_url, headers = client.generate_ws_connection(runtime_arn, session_id=test_session_id)
+
+        # Verify session ID is in headers
+        assert "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id" in headers
+        assert headers["X-Amzn-Bedrock-AgentCore-Runtime-Session-Id"] == test_session_id
 
     def test_generate_presigned_url_returns_valid_format(self, mock_boto_session):
         """Test that generate_presigned_url returns properly formatted URL."""
@@ -69,6 +84,19 @@ def test_generate_presigned_url_returns_valid_format(self, mock_boto_session):
         assert "runtimes" in presigned_url
         assert "X-Amz-Algorithm" in presigned_url
         assert "X-Amz-Signature" in presigned_url
+        # Verify session ID is in query params
+        assert "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id=" in presigned_url
+
+    def test_generate_presigned_url_with_session_id(self, mock_boto_session):
+        """Test that generate_presigned_url includes session ID in query params."""
+        client = AgentCoreRuntimeClient(region="us-west-2")
+        runtime_arn = "arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/test-runtime"
+        test_session_id = "integration-test-presigned-session"
+
+        presigned_url = client.generate_presigned_url(runtime_arn, session_id=test_session_id)
+
+        # Verify session ID is in query params
+        assert f"X-Amzn-Bedrock-AgentCore-Runtime-Session-Id={test_session_id}" in presigned_url
 
     @pytest.mark.skip(reason="Requires actual runtime endpoint")
     async def test_connect_with_generated_headers(self):

tests/unit/runtime/test_agent_core_runtime_client.py

@@ -181,6 +181,26 @@ def test_generate_connection_with_session_id(self, mock_endpoint, mock_session):
 
         assert ws_url is not None
         assert headers is not None
+        # Verify session ID is in headers
+        assert "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id" in headers
+        assert headers["X-Amzn-Bedrock-AgentCore-Runtime-Session-Id"] == "test-session-123"
+
+    @patch("bedrock_agentcore.runtime.agent_core_runtime_client.boto3.Session")
+    @patch("bedrock_agentcore.runtime.agent_core_runtime_client.get_data_plane_endpoint")
+    def test_generate_connection_user_agent(self, mock_endpoint, mock_session):
+        """Test that User-Agent header is set correctly."""
+        mock_endpoint.return_value = "https://example.aws.dev"
+        mock_creds = Mock()
+        mock_creds.get_frozen_credentials.return_value = Mock(access_key="AKIATEST", secret_key="secret", token=None)
+        mock_session.return_value.get_credentials.return_value = mock_creds
+
+        client = AgentCoreRuntimeClient(region="us-west-2")
+        runtime_arn = "arn:aws:bedrock-agentcore:us-west-2:123:runtime/my-runtime"
+
+        ws_url, headers = client.generate_ws_connection(runtime_arn)
+
+        assert "User-Agent" in headers
+        assert headers["User-Agent"] == "AgentCoreRuntimeClient/1.0"
 
     @patch("bedrock_agentcore.runtime.agent_core_runtime_client.boto3.Session")
     @patch("bedrock_agentcore.runtime.agent_core_runtime_client.get_data_plane_endpoint")
@@ -271,6 +291,23 @@ def test_generate_presigned_url_with_custom_headers(self, mock_endpoint, mock_se
 
         assert "abc=pqr" in presigned_url
 
+    @patch("bedrock_agentcore.runtime.agent_core_runtime_client.boto3.Session")
+    @patch("bedrock_agentcore.runtime.agent_core_runtime_client.get_data_plane_endpoint")
+    def test_generate_presigned_url_with_session_id(self, mock_endpoint, mock_session):
+        """Test generating presigned URL with explicit session ID."""
+        mock_endpoint.return_value = "https://example.aws.dev"
+        mock_creds = Mock()
+        mock_creds.get_frozen_credentials.return_value = Mock(access_key="AKIATEST", secret_key="secret", token=None)
+        mock_session.return_value.get_credentials.return_value = mock_creds
+
+        client = AgentCoreRuntimeClient(region="us-west-2")
+        runtime_arn = "arn:aws:bedrock-agentcore:us-west-2:123:runtime/my-runtime"
+
+        presigned_url = client.generate_presigned_url(runtime_arn, session_id="test-session-456")
+
+        # Verify session ID is in query params
+        assert "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id=test-session-456" in presigned_url
+
     @patch("bedrock_agentcore.runtime.agent_core_runtime_client.boto3.Session")
     @patch("bedrock_agentcore.runtime.agent_core_runtime_client.get_data_plane_endpoint")
     def test_generate_presigned_url_with_custom_expires(self, mock_endpoint, mock_session):
@@ -286,6 +323,8 @@ def test_generate_presigned_url_with_custom_expires(self, mock_endpoint, mock_se
         presigned_url = client.generate_presigned_url(runtime_arn, expires=60)
 
         assert "X-Amz-Expires=60" in presigned_url
+        # Verify auto-generated session ID is present
+        assert "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id=" in presigned_url
 
     def test_generate_presigned_url_exceeds_max_expires_raises_error(self):
         """Test that exceeding max expiration raises ValueError."""