Added request header allowlist support (#93)

Add support for specifying request header allowlist feature which is supported with agent core runtime with the latest release.

This change will extract any request headers passed as per the request header allowlist. These headers are set in BedrockAgentCoreContext and passed via RequestContext.

Co-authored-by: donmanue <[email protected]>

Author: evolutiontheory

pyproject.toml

@@ -26,8 +26,8 @@ classifiers = [
     "Topic :: Software Development :: Libraries :: Python Modules",
 ]
 dependencies = [
-    "boto3>=1.39.7",
-    "botocore>=1.39.7",
+    "boto3>=1.40.35",
+    "botocore>=1.40.35",
     "pydantic>=2.0.0,<3.0.0",
     "urllib3>=1.26.0",
     "starlette>=0.46.2",

src/bedrock_agentcore/runtime/app.py

@@ -21,6 +21,8 @@
 from .context import BedrockAgentCoreContext, RequestContext
 from .models import (
     ACCESS_TOKEN_HEADER,
+    AUTHORIZATION_HEADER,
+    CUSTOM_HEADER_PREFIX,
     REQUEST_ID_HEADER,
     SESSION_HEADER,
     TASK_ACTION_CLEAR_FORCED_STATUS,
@@ -279,7 +281,27 @@ def _build_request_context(self, request) -> RequestContext:
             if agent_identity_token:
                 BedrockAgentCoreContext.set_workload_access_token(agent_identity_token)
 
-            return RequestContext(session_id=session_id)
+            # Collect relevant request headers (Authorization + Custom headers)
+            request_headers = {}
+
+            # Add Authorization header if present
+            authorization_header = headers.get(AUTHORIZATION_HEADER)
+            if authorization_header is not None:
+                request_headers[AUTHORIZATION_HEADER] = authorization_header
+
+            # Add custom headers with the specified prefix
+            for header_name, header_value in headers.items():
+                if header_name.lower().startswith(CUSTOM_HEADER_PREFIX.lower()):
+                    request_headers[header_name] = header_value
+
+            # Set in context if any headers were found
+            if request_headers:
+                BedrockAgentCoreContext.set_request_headers(request_headers)
+
+            # Get the headers from context to pass to RequestContext
+            req_headers = BedrockAgentCoreContext.get_request_headers()
+
+            return RequestContext(session_id=session_id, request_headers=req_headers)
         except Exception as e:
             self.logger.warning("Failed to build request context: %s: %s", type(e).__name__, e)
             request_id = str(uuid.uuid4())

src/bedrock_agentcore/runtime/context.py

@@ -4,7 +4,7 @@
 """
 
 from contextvars import ContextVar
-from typing import Optional
+from typing import Dict, Optional
 
 from pydantic import BaseModel, Field
 
@@ -13,6 +13,7 @@ class RequestContext(BaseModel):
     """Request context containing metadata from HTTP requests."""
 
     session_id: Optional[str] = Field(None)
+    request_headers: Optional[Dict[str, str]] = Field(None)
 
 
 class BedrockAgentCoreContext:
@@ -21,6 +22,7 @@ class BedrockAgentCoreContext:
     _workload_access_token: ContextVar[Optional[str]] = ContextVar("workload_access_token")
     _request_id: ContextVar[Optional[str]] = ContextVar("request_id")
     _session_id: ContextVar[Optional[str]] = ContextVar("session_id")
+    _request_headers: ContextVar[Optional[Dict[str, str]]] = ContextVar("request_headers")
 
     @classmethod
     def set_workload_access_token(cls, token: str):
@@ -56,3 +58,16 @@ def get_session_id(cls) -> Optional[str]:
             return cls._session_id.get()
         except LookupError:
             return None
+
+    @classmethod
+    def set_request_headers(cls, headers: Dict[str, str]):
+        """Set request headers in the context."""
+        cls._request_headers.set(headers)
+
+    @classmethod
+    def get_request_headers(cls) -> Optional[Dict[str, str]]:
+        """Get request headers from the context."""
+        try:
+            return cls._request_headers.get()
+        except LookupError:
+            return None

src/bedrock_agentcore/runtime/models.py

@@ -17,6 +17,8 @@ class PingStatus(str, Enum):
 SESSION_HEADER = "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id"
 REQUEST_ID_HEADER = "X-Amzn-Bedrock-AgentCore-Runtime-Request-Id"
 ACCESS_TOKEN_HEADER = "WorkloadAccessToken"  # nosec
+AUTHORIZATION_HEADER = "Authorization"
+CUSTOM_HEADER_PREFIX = "X-Amzn-Bedrock-AgentCore-Runtime-Custom-"
 
 # Task action constants
 TASK_ACTION_PING_STATUS = "ping_status"