Feature/request header allowlist

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"