From 5bad65e63ec04eb1d89cd298a7c2d883cd7b2e64 Mon Sep 17 00:00:00 2001 From: Sudip Dutta Date: Tue, 28 Jul 2026 00:31:58 -0700 Subject: [PATCH] feat(runtime): Propagate X-Amz-Bedrock-AgentCore-Identity-WAT on outbound calls Enable automatic workload identity chain propagation for agent code running on AgentCore Runtime. The Runtime platform mints a WAT (with chain extension) and delivers it to the agent via the X-Amz-Bedrock-AgentCore-Identity-WAT header. This PR adds SDK support to: 1. Extract the WAT from inbound requests (falls back to WorkloadAccessToken) 2. Auto-attach it as X-Amz-Bedrock-AgentCore-Identity-WAT on outbound boto3 calls to Runtime/Gateway (InvokeAgentRuntime, InvokeHarness, etc.) 3. Expose @requires_wat decorator for reading the WAT from context 4. Allowlist the header in is_forwardable_header() For Gateway MCP calls (raw HTTP), developers read the WAT via BedrockAgentCoreContext.get_workload_access_token() and attach manually. --- .../_utils/identity_propagation.py | 72 +++++++++++ src/bedrock_agentcore/identity/__init__.py | 4 +- src/bedrock_agentcore/identity/auth.py | 68 +++++++++++ .../runtime/agent_core_runtime_client.py | 5 + src/bedrock_agentcore/runtime/app.py | 3 +- src/bedrock_agentcore/runtime/models.py | 6 + .../runtime/test_identity_wat_propagation.py | 114 ++++++++++++++++++ 7 files changed, 269 insertions(+), 3 deletions(-) create mode 100644 src/bedrock_agentcore/_utils/identity_propagation.py create mode 100644 tests/bedrock_agentcore/runtime/test_identity_wat_propagation.py diff --git a/src/bedrock_agentcore/_utils/identity_propagation.py b/src/bedrock_agentcore/_utils/identity_propagation.py new file mode 100644 index 00000000..2cbfbced --- /dev/null +++ b/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}", + ) diff --git a/src/bedrock_agentcore/identity/__init__.py b/src/bedrock_agentcore/identity/__init__.py index e54ab08e..70016267 100644 --- a/src/bedrock_agentcore/identity/__init__.py +++ b/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"] diff --git a/src/bedrock_agentcore/identity/auth.py b/src/bedrock_agentcore/identity/auth.py index 447d9c24..ae7a796b 100644 --- a/src/bedrock_agentcore/identity/auth.py +++ b/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 diff --git a/src/bedrock_agentcore/runtime/agent_core_runtime_client.py b/src/bedrock_agentcore/runtime/agent_core_runtime_client.py index 57190012..63a00051 100644 --- a/src/bedrock_agentcore/runtime/agent_core_runtime_client.py +++ b/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, diff --git a/src/bedrock_agentcore/runtime/app.py b/src/bedrock_agentcore/runtime/app.py index f2e596be..61a6f3bb 100644 --- a/src/bedrock_agentcore/runtime/app.py +++ b/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) diff --git a/src/bedrock_agentcore/runtime/models.py b/src/bedrock_agentcore/runtime/models.py index cad309da..db0cef82 100644 --- a/src/bedrock_agentcore/runtime/models.py +++ b/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,6 +181,7 @@ 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 """ @@ -186,6 +189,9 @@ def is_forwardable_header(header_name: str) -> bool: 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 diff --git a/tests/bedrock_agentcore/runtime/test_identity_wat_propagation.py b/tests/bedrock_agentcore/runtime/test_identity_wat_propagation.py new file mode 100644 index 00000000..d6b413f5 --- /dev/null +++ b/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"