feat(identity): Add X-Amz-Bedrock-AgentCore-Identity-WAT header propa…

src/bedrock_agentcore/_utils/identity_propagation.py

@@ -0,0 +1,72 @@
+"""Utilities for propagating the Identity WAT header on outbound boto3 calls.
+
+The X-Amz-Bedrock-AgentCore-Identity-WAT header carries the Workload Access Token
+across service hops, enabling the workload identity chain to be extended automatically.
+When present in the request context, it must be attached to outbound calls to
+Runtime and Gateway services so downstream services can verify and extend the chain.
+
+WAT is NOT propagated to Identity calls (Identity is the WAT issuer, not a consumer).
+
+Note on Gateway MCP calls:
+    Gateway invocations use the MCP protocol over raw HTTP (no boto3 API exists for
+    invoke_gateway). For these calls, the developer must manually attach the WAT header
+    using ``BedrockAgentCoreContext.get_workload_access_token()``::
+
+        wat = BedrockAgentCoreContext.get_workload_access_token()
+        if wat:
+            headers["X-Amz-Bedrock-AgentCore-Identity-WAT"] = wat
+"""
+
+import logging
+
+from botocore.client import BaseClient
+
+from ..runtime.context import BedrockAgentCoreContext
+from ..runtime.models import IDENTITY_WAT_HEADER
+
+logger = logging.getLogger(__name__)
+
+# Operations that should receive the WAT header (Runtime and Gateway invocations)
+_WAT_PROPAGATION_OPERATIONS = (
+    "InvokeAgentRuntime",
+    "InvokeAgentRuntimeCommand",
+    "InvokeHarness",
+    "InvokeGateway",
+)
+
+
+def _inject_identity_wat_header(request, **kwargs) -> None:
+    """Inject the Identity WAT header into an outgoing boto3 request.
+
+    Reads the WAT from ``BedrockAgentCoreContext.get_workload_access_token()``
+    and attaches it as the ``X-Amz-Bedrock-AgentCore-Identity-WAT`` header.
+    No-ops if no WAT is present in the current context.
+
+    Args:
+        request: The ``AWSPreparedRequest`` about to be sent.
+        **kwargs: Additional event keyword arguments (ignored).
+    """
+    wat = BedrockAgentCoreContext.get_workload_access_token()
+    if wat:
+        request.headers[IDENTITY_WAT_HEADER] = wat
+
+
+def register_identity_wat_propagation(client: BaseClient) -> None:
+    """Register the Identity WAT propagation handler on a boto3 client.
+
+    Registers a ``before-sign`` event handler for each allowlisted operation
+    so that Runtime and Gateway invocation requests automatically include the
+    ``X-Amz-Bedrock-AgentCore-Identity-WAT`` header when a WAT is present
+    in the current request context.
+
+    Identity operations (token minting, credential retrieval) are excluded.
+
+    Args:
+        client: A boto3/botocore client instance (bedrock-agentcore data plane).
+    """
+    for op in _WAT_PROPAGATION_OPERATIONS:
+        client.meta.events.register(
+            f"before-sign.bedrock-agentcore.{op}",
+            _inject_identity_wat_header,
+            unique_id=f"identity-wat-{op}",
+        )

src/bedrock_agentcore/identity/__init__.py

@@ -1,5 +1,5 @@
 """Bedrock AgentCore SDK identity package."""
 
-from .auth import requires_access_token, requires_api_key
+from .auth import requires_access_token, requires_api_key, requires_wat
 
-__all__ = ["requires_access_token", "requires_api_key"]
+__all__ = ["requires_access_token", "requires_api_key", "requires_wat"]

src/bedrock_agentcore/identity/auth.py

@@ -274,6 +274,74 @@ def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
     return decorator
 
 
+def requires_wat(
+    *,
+    into: str = "identity_wat",
+) -> Callable:
+    """Decorator that reads the Workload Access Token (WAT) from context and propagates it.
+
+    The WAT is minted by the Runtime platform before the agent code runs.
+    This decorator reads the WAT from ``BedrockAgentCoreContext``, stores it for
+    automatic propagation via the ``X-Amz-Bedrock-AgentCore-Identity-WAT`` header
+    on outbound Runtime/Gateway calls, and injects it into the decorated function.
+
+    Use this when your agent needs to propagate workload identity across hops
+    (e.g., calling other AgentCore runtimes or gateways).
+
+    The agent code does NOT need to call Identity/ACPS to get the WAT — the Runtime
+    service handles WAT minting (with chain extension) and delivers the WAT to the
+    agent via the request context.
+
+    Args:
+        into: Parameter name to inject the WAT into (default: 'identity_wat').
+
+    Returns:
+        Decorator function
+
+    Example::
+
+        @requires_wat()
+        def call_downstream(query: str, *, identity_wat: str) -> str:
+            '''Call a downstream service with WAT propagation.'''
+            # identity_wat is the WAT minted by the Runtime platform
+            # For boto3 calls (invoke_agent_runtime), it auto-propagates
+            # For raw HTTP calls (gateway MCP), attach it manually:
+            #   headers["X-Amz-Bedrock-AgentCore-Identity-WAT"] = identity_wat
+            ...
+    """
+
+    def decorator(func: Callable) -> Callable:
+
+        def _get_wat() -> str:
+            """Read the WAT from context (already provisioned by Runtime)."""
+            wat = BedrockAgentCoreContext.get_workload_access_token()
+            if not wat:
+                raise ValueError(
+                    "No workload access token available in context. "
+                    "Ensure the agent is running on AgentCore runtime."
+                )
+            return wat
+
+        @wraps(func)
+        async def async_wrapper(*args: Any, **kwargs_func: Any) -> Any:
+            wat = _get_wat()
+            kwargs_func[into] = wat
+            return await func(*args, **kwargs_func)
+
+        @wraps(func)
+        def sync_wrapper(*args: Any, **kwargs_func: Any) -> Any:
+            wat = _get_wat()
+            kwargs_func[into] = wat
+            return func(*args, **kwargs_func)
+
+        if asyncio.iscoroutinefunction(func):
+            return async_wrapper
+        else:
+            return sync_wrapper
+
+    return decorator
+
+
 def _get_oauth2_callback_url(user_provided_oauth2_callback_url: Optional[str]):
     if user_provided_oauth2_callback_url:
         return user_provided_oauth2_callback_url

src/bedrock_agentcore/runtime/agent_core_runtime_client.py

@@ -25,6 +25,7 @@
 
 from .._utils.config import WaitConfig
 from .._utils.endpoints import get_data_plane_endpoint, validate_region
+from .._utils.identity_propagation import register_identity_wat_propagation
 from .._utils.polling import wait_until, wait_until_deleted
 from .._utils.snake_case import accept_snake_case_kwargs, convert_kwargs
 from .._utils.user_agent import build_user_agent_suffix
@@ -102,6 +103,10 @@ def __init__(
             region_name=self.region,
             config=client_config,
         )
+
+        # Register Identity WAT header propagation on outbound data plane calls
+        register_identity_wat_propagation(self.dp_client)
+
         self.logger.info(
             "Initialized AgentCoreRuntimeClient for region: %s",
             self.region,

src/bedrock_agentcore/runtime/app.py

@@ -37,6 +37,7 @@
     AUTHORIZATION_HEADER,
     BAGGAGE_KEY_EXPERIMENT_ARN,
     BAGGAGE_KEY_EXPERIMENT_VARIANT,
+    IDENTITY_WAT_HEADER,
     OAUTH2_CALLBACK_URL_HEADER,
     REQUEST_ID_HEADER,
     SESSION_HEADER,
@@ -409,7 +410,7 @@ def _build_request_context(self, request) -> RequestContext:
             session_id = headers.get(SESSION_HEADER)
             BedrockAgentCoreContext.set_request_context(request_id, session_id)
 
-            agent_identity_token = headers.get(ACCESS_TOKEN_HEADER)
+            agent_identity_token = headers.get(IDENTITY_WAT_HEADER) or headers.get(ACCESS_TOKEN_HEADER)
             if agent_identity_token:
                 BedrockAgentCoreContext.set_workload_access_token(agent_identity_token)
 

src/bedrock_agentcore/runtime/models.py

@@ -18,6 +18,7 @@ class PingStatus(str, Enum):
 SHELL_ID_HEADER = "X-Amzn-Bedrock-AgentCore-Shell-Id"
 REQUEST_ID_HEADER = "X-Amzn-Bedrock-AgentCore-Runtime-Request-Id"
 ACCESS_TOKEN_HEADER = "WorkloadAccessToken"  # nosec
+IDENTITY_WAT_HEADER = "X-Amz-Bedrock-AgentCore-Identity-WAT"  # nosec
 OAUTH2_CALLBACK_URL_HEADER = "OAuth2CallbackUrl"
 AUTHORIZATION_HEADER = "Authorization"
 CUSTOM_HEADER_PREFIX = "X-Amzn-Bedrock-AgentCore-Runtime-Custom-"
@@ -171,6 +172,7 @@ class PingStatus(str, Enum):
 
 _CUSTOM_HEADER_PREFIX_LOWER = CUSTOM_HEADER_PREFIX.lower()
 _AUTHORIZATION_HEADER_LOWER = AUTHORIZATION_HEADER.lower()
+_IDENTITY_WAT_HEADER_LOWER = IDENTITY_WAT_HEADER.lower()
 
 
 def is_forwardable_header(header_name: str) -> bool:
@@ -179,13 +181,17 @@ def is_forwardable_header(header_name: str) -> bool:
     Rules (from the AgentCore runtime header allowlist docs):
     - Not in the restricted headers list
     - Does not start with ``x-amz-`` (reserved for AWS SigV4 signing)
+      unless it is the Identity WAT header used for workload identity chain propagation
     - Does not start with ``x-amzn-`` unless it starts with the legacy
       ``X-Amzn-Bedrock-AgentCore-Runtime-Custom-`` prefix
     """
     lower = header_name.lower()
     if lower in RESTRICTED_HEADERS:
         return False
     if lower.startswith("x-amz-"):
+        # Allow the Identity WAT header through for workload identity chain propagation
+        if lower == _IDENTITY_WAT_HEADER_LOWER:
+            return True
         return False
     if lower.startswith("x-amzn-") and not lower.startswith(_CUSTOM_HEADER_PREFIX_LOWER):
         return False

tests/bedrock_agentcore/runtime/test_identity_wat_propagation.py

@@ -0,0 +1,114 @@
+"""Tests for X-Amz-Bedrock-AgentCore-Identity-WAT header propagation."""
+
+import contextvars
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from bedrock_agentcore._utils.identity_propagation import (
+    _inject_identity_wat_header,
+    register_identity_wat_propagation,
+)
+from bedrock_agentcore.runtime.context import BedrockAgentCoreContext
+from bedrock_agentcore.runtime.models import (
+    IDENTITY_WAT_HEADER,
+    is_forwardable_header,
+)
+
+
+class TestIdentityWatHeader:
+    """Tests for the IDENTITY_WAT_HEADER constant and forwarding rules."""
+
+    def test_identity_wat_header_value(self):
+        """Test that the header constant has the correct value."""
+        assert IDENTITY_WAT_HEADER == "X-Amz-Bedrock-AgentCore-Identity-WAT"
+
+    def test_identity_wat_header_is_forwardable(self):
+        """Test that the Identity WAT header passes the forwardable check."""
+        assert is_forwardable_header(IDENTITY_WAT_HEADER) is True
+
+    def test_identity_wat_header_is_forwardable_case_insensitive(self):
+        """Test that the Identity WAT header passes regardless of casing."""
+        assert is_forwardable_header("x-amz-bedrock-agentcore-identity-wat") is True
+        assert is_forwardable_header("X-AMZ-BEDROCK-AGENTCORE-IDENTITY-WAT") is True
+
+    def test_other_x_amz_headers_still_blocked(self):
+        """Test that other x-amz- headers are still blocked."""
+        assert is_forwardable_header("X-Amz-Date") is False
+        assert is_forwardable_header("X-Amz-Security-Token") is False
+        assert is_forwardable_header("x-amz-content-sha256") is False
+
+
+class TestInjectIdentityWatHeader:
+    """Tests for the _inject_identity_wat_header event handler."""
+
+    def test_injects_header_when_wat_present(self):
+        """Test that the WAT header is injected when present in context."""
+        token = "my-wat-token"
+        BedrockAgentCoreContext.set_workload_access_token(token)
+
+        request = MagicMock()
+        request.headers = {}
+        _inject_identity_wat_header(request)
+
+        assert request.headers[IDENTITY_WAT_HEADER] == token
+
+    def test_does_not_inject_when_wat_absent(self):
+        """Test that no header is injected when no WAT in context."""
+        ctx = contextvars.Context()
+
+        def test_in_new_context():
+            request = MagicMock()
+            request.headers = {}
+            _inject_identity_wat_header(request)
+            return request.headers
+
+        headers = ctx.run(test_in_new_context)
+        assert IDENTITY_WAT_HEADER not in headers
+
+
+class TestRegisterIdentityWatPropagation:
+    """Tests for the register_identity_wat_propagation utility."""
+
+    def test_registers_event_handler_per_operation(self):
+        """Test that event handlers are registered for each operation."""
+        mock_client = MagicMock()
+        register_identity_wat_propagation(mock_client)
+
+        calls = mock_client.meta.events.register.call_args_list
+        registered_events = [call[0][0] for call in calls]
+
+        assert "before-sign.bedrock-agentcore.InvokeAgentRuntime" in registered_events
+        assert "before-sign.bedrock-agentcore.InvokeAgentRuntimeCommand" in registered_events
+        assert "before-sign.bedrock-agentcore.InvokeHarness" in registered_events
+        assert "before-sign.bedrock-agentcore.InvokeGateway" in registered_events
+
+    def test_does_not_register_for_identity_operations(self):
+        """Test that Identity operations are not registered."""
+        mock_client = MagicMock()
+        register_identity_wat_propagation(mock_client)
+
+        calls = mock_client.meta.events.register.call_args_list
+        registered_events = [call[0][0] for call in calls]
+
+        for event in registered_events:
+            assert "GetWorkloadAccessToken" not in event
+            assert "GetResourceOauth2Token" not in event
+
+
+class TestAppExtractsWatHeader:
+    """Tests for WAT header extraction in _build_request_context."""
+
+    def test_workload_access_token_stored_in_context(self):
+        """Test that app extracts WorkloadAccessToken and stores in context."""
+        from bedrock_agentcore.runtime.app import BedrockAgentCoreApp
+
+        app = BedrockAgentCoreApp()
+        mock_request = MagicMock()
+        mock_request.headers = {
+            "WorkloadAccessToken": "wat-from-runtime",
+        }
+
+        app._build_request_context(mock_request)
+
+        assert BedrockAgentCoreContext.get_workload_access_token() == "wat-from-runtime"