diff --git a/src/bedrock_agentcore/_utils/user_agent.py b/src/bedrock_agentcore/_utils/user_agent.py index 25dad669..7146a871 100644 --- a/src/bedrock_agentcore/_utils/user_agent.py +++ b/src/bedrock_agentcore/_utils/user_agent.py @@ -12,29 +12,42 @@ SDK_VERSION = "unknown" -def build_user_agent_suffix(integration_source: Optional[str] = None) -> str: +def _sanitize_token(value: str) -> str: + """Sanitize a User-Agent token to prevent header injection.""" + return "".join(c for c in value.lower() if c.isalnum() or c in "-_") + + +def build_user_agent_suffix(integration_source: Optional[str] = None, feature: Optional[str] = None) -> str: """Build the suffix string to append to boto3 User-Agent header. This value is passed to botocore's Config(user_agent_extra=...) parameter. Args: integration_source: Optional integration framework identifier - (e.g., 'langchain', 'crewai', 'strands') + (e.g., 'langgraph', 'crewai', 'strands', 'raw-sdk') + feature: Optional feature identifier used to separate a capability's + calls from other calls (e.g., 'payments') Returns: String to append to User-Agent header Example: - >>> build_user_agent_suffix("langchain") - 'bedrock-agentcore/1.0.0 (integration_source=langchain)' + >>> build_user_agent_suffix("langgraph") + 'bedrock-agentcore/1.0.0 (integration_source=langgraph)' + >>> build_user_agent_suffix("strands", feature="payments") + 'bedrock-agentcore/1.0.0 (integration_source=strands; feature=payments)' >>> build_user_agent_suffix() 'bedrock-agentcore/1.0.0' """ base = f"bedrock-agentcore/{SDK_VERSION}" + tokens = [] if integration_source: - # Sanitize to prevent header injection - sanitized = "".join(c for c in integration_source.lower() if c.isalnum() or c in "-_") - return f"{base} (integration_source={sanitized})" + tokens.append(f"integration_source={_sanitize_token(integration_source)}") + if feature: + tokens.append(f"feature={_sanitize_token(feature)}") + + if tokens: + return f"{base} ({'; '.join(tokens)})" return base diff --git a/src/bedrock_agentcore/payments/client.py b/src/bedrock_agentcore/payments/client.py index d5c297a9..fb657e77 100644 --- a/src/bedrock_agentcore/payments/client.py +++ b/src/bedrock_agentcore/payments/client.py @@ -250,20 +250,21 @@ def _build_provider_config_input( def __init__( self, region_name: Optional[str] = None, - integration_source: Optional[str] = None, + integration_source: str = "raw-sdk", ) -> None: """Initialize the Payments control plane client. Args: region_name: AWS region name. Defaults to boto3 session region or us-west-2 - integration_source: Optional identifier for tracking integration source in telemetry + integration_source: Identifier of the surface making the calls, propagated via + the boto3 User-Agent header for usage measurement. Defaults to "raw-sdk". """ self.region_name = region_name or boto3.Session().region_name or "us-west-2" - self.integration_source = integration_source + self.integration_source = integration_source or "raw-sdk" # Build config with user-agent for telemetry - user_agent_extra = build_user_agent_suffix(integration_source) + user_agent_extra = build_user_agent_suffix(integration_source=self.integration_source, feature="payments") client_config = Config(user_agent_extra=user_agent_extra) # Control plane operations are available through bedrock-agentcore-control service diff --git a/src/bedrock_agentcore/payments/integrations/langgraph/middleware.py b/src/bedrock_agentcore/payments/integrations/langgraph/middleware.py index e1f9d4e8..dc13a6a9 100644 --- a/src/bedrock_agentcore/payments/integrations/langgraph/middleware.py +++ b/src/bedrock_agentcore/payments/integrations/langgraph/middleware.py @@ -99,6 +99,7 @@ def __init__(self, config: AgentCorePaymentsConfig) -> None: agent_name=config.agent_name, bearer_token=config.bearer_token, token_provider=config.token_provider, + integration_source="langgraph", ) except Exception as e: raise RuntimeError(f"Failed to initialize PaymentManager: {e}") from e diff --git a/src/bedrock_agentcore/payments/integrations/strands/plugin.py b/src/bedrock_agentcore/payments/integrations/strands/plugin.py index 49d0e183..775a5d98 100644 --- a/src/bedrock_agentcore/payments/integrations/strands/plugin.py +++ b/src/bedrock_agentcore/payments/integrations/strands/plugin.py @@ -104,6 +104,7 @@ def init_agent(self, agent) -> None: agent_name=self.config.agent_name, bearer_token=self.config.bearer_token, token_provider=self.config.token_provider, + integration_source="strands", ) logger.info("PaymentManager initialized successfully") except Exception as e: diff --git a/src/bedrock_agentcore/payments/manager.py b/src/bedrock_agentcore/payments/manager.py index 3a56bcde..5a087b21 100644 --- a/src/bedrock_agentcore/payments/manager.py +++ b/src/bedrock_agentcore/payments/manager.py @@ -173,6 +173,7 @@ def __init__( agent_name: Optional[str] = None, bearer_token: Optional[str] = None, token_provider: Optional[Callable[[], str]] = None, + integration_source: str = "raw-sdk", ): """Initialize a PaymentManager instance. @@ -193,6 +194,9 @@ def __init__( token_provider: Optional callable that returns a fresh JWT bearer token string. Called before each request to support token refresh. Mutually exclusive with bearer_token. + integration_source: Identifier of the surface making payment calls, propagated + via the boto3 User-Agent header for usage measurement. Defaults + to "raw-sdk"; integrations set "strands", "langgraph", etc. Raises: ValueError: If payment_manager_arn is invalid, region_name conflicts with boto3_session region, @@ -219,6 +223,7 @@ def __init__( # Store payment manager ARN self._payment_manager_arn: str = payment_manager_arn self._agent_name: Optional[str] = agent_name + self._integration_source: str = integration_source or "raw-sdk" self._bearer_token: Optional[str] = bearer_token self._token_provider: Optional[Callable[[], str]] = token_provider @@ -334,7 +339,7 @@ def _build_client_config(self, boto_client_config: Optional[BotocoreConfig]) -> Returns: Final client configuration with SDK user agent """ - user_agent_extra = build_user_agent_suffix() + user_agent_extra = build_user_agent_suffix(integration_source=self._integration_source, feature="payments") if boto_client_config: existing_user_agent = getattr(boto_client_config, "user_agent_extra", None) diff --git a/tests/bedrock_agentcore/_utils/__init__.py b/tests/bedrock_agentcore/_utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/bedrock_agentcore/_utils/test_user_agent.py b/tests/bedrock_agentcore/_utils/test_user_agent.py new file mode 100644 index 00000000..0e4f9bdd --- /dev/null +++ b/tests/bedrock_agentcore/_utils/test_user_agent.py @@ -0,0 +1,33 @@ +"""Tests for build_user_agent_suffix.""" + +from bedrock_agentcore._utils import user_agent +from bedrock_agentcore._utils.user_agent import build_user_agent_suffix + + +def _base() -> str: + return f"bedrock-agentcore/{user_agent.SDK_VERSION}" + + +def test_no_arguments_returns_base(): + assert build_user_agent_suffix() == _base() + + +def test_integration_source_only(): + assert build_user_agent_suffix("strands") == f"{_base()} (integration_source=strands)" + + +def test_integration_source_and_feature(): + assert build_user_agent_suffix("strands", feature="payments") == ( + f"{_base()} (integration_source=strands; feature=payments)" + ) + + +def test_feature_only(): + assert build_user_agent_suffix(feature="payments") == f"{_base()} (feature=payments)" + + +def test_tokens_are_sanitized_and_lowercased(): + # Injection characters and spaces are stripped; value is lowercased. + assert build_user_agent_suffix("Ra w-SDK) evil", feature="Pay;ments") == ( + f"{_base()} (integration_source=raw-sdkevil; feature=payments)" + ) diff --git a/tests/bedrock_agentcore/payments/integrations/langgraph/test_stage1.py b/tests/bedrock_agentcore/payments/integrations/langgraph/test_stage1.py index 4ae6c2ff..597def53 100644 --- a/tests/bedrock_agentcore/payments/integrations/langgraph/test_stage1.py +++ b/tests/bedrock_agentcore/payments/integrations/langgraph/test_stage1.py @@ -225,6 +225,7 @@ def test_middleware_creates_payment_manager(self, mock_pm_cls): agent_name="test-agent", bearer_token=None, token_provider=None, + integration_source="langgraph", ) assert mw.config is config assert mw.payment_manager is mock_pm_cls.return_value @@ -242,6 +243,7 @@ def test_middleware_passes_bearer_token(self, mock_pm_cls): agent_name=None, bearer_token="my-jwt", token_provider=None, + integration_source="langgraph", ) @patch("bedrock_agentcore.payments.integrations.langgraph.middleware.PaymentManager") diff --git a/tests/bedrock_agentcore/payments/integrations/strands/test_plugin.py b/tests/bedrock_agentcore/payments/integrations/strands/test_plugin.py index c77c7209..181151d6 100644 --- a/tests/bedrock_agentcore/payments/integrations/strands/test_plugin.py +++ b/tests/bedrock_agentcore/payments/integrations/strands/test_plugin.py @@ -127,6 +127,7 @@ def test_init_agent_success(self, mock_payment_manager_class): agent_name=None, bearer_token=None, token_provider=None, + integration_source="strands", ) assert plugin.payment_manager == mock_pm_instance @@ -1282,6 +1283,7 @@ def test_init_agent_passes_agent_name_to_payment_manager(self, mock_payment_mana agent_name="my-agent", bearer_token=None, token_provider=None, + integration_source="strands", ) @patch("bedrock_agentcore.payments.integrations.strands.plugin.PaymentManager") @@ -1307,6 +1309,7 @@ def test_init_agent_passes_none_agent_name_when_not_set(self, mock_payment_manag agent_name=None, bearer_token=None, token_provider=None, + integration_source="strands", )