From 57dc939e1160e38ac0737cca89ef1fef48aac2f0 Mon Sep 17 00:00:00 2001 From: Raju Ansari Date: Thu, 17 Sep 2026 15:21:18 +0000 Subject: [PATCH] feat(payments): attribute SDK usage via user-agent surface tag and created-by resource tag Adds two complementary origin signals for AgentCore Payments: - Per-call surface tag: PaymentManager and PaymentClient now propagate an integration_source (default "raw-sdk") plus a feature=payments token through the boto3 User-Agent. Strands sets "strands", LangGraph sets "langgraph". Lets the payments DP count callers by surface/version/operation from each request. - Durable resource tag: create_payment_manager stamps agentcore:created-by=agentcore-sdk (default-if-absent, caller value preserved) so payment managers created through the SDK are attributable after creation. build_user_agent_suffix gains an optional feature token. Adds unit coverage for the user-agent formatting and the created-by tag injection/override behavior. --- src/bedrock_agentcore/_utils/user_agent.py | 27 +++++++++--- src/bedrock_agentcore/payments/client.py | 21 +++++++-- .../integrations/langgraph/middleware.py | 1 + .../payments/integrations/strands/plugin.py | 1 + src/bedrock_agentcore/payments/manager.py | 7 ++- tests/bedrock_agentcore/_utils/__init__.py | 0 .../_utils/test_user_agent.py | 33 ++++++++++++++ .../integrations/langgraph/test_stage1.py | 2 + .../integrations/strands/test_plugin.py | 3 ++ .../bedrock_agentcore/payments/test_client.py | 43 +++++++++++++++++++ 10 files changed, 126 insertions(+), 12 deletions(-) create mode 100644 tests/bedrock_agentcore/_utils/__init__.py create mode 100644 tests/bedrock_agentcore/_utils/test_user_agent.py 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..b85020d7 100644 --- a/src/bedrock_agentcore/payments/client.py +++ b/src/bedrock_agentcore/payments/client.py @@ -24,6 +24,11 @@ logger = logging.getLogger(__name__) +# Tag stamped on resources created through the SDK, used to attribute resources by origin. +# Applied only when the caller has not already supplied this key. +CREATED_BY_TAG_KEY = "agentcore:created-by" +CREATED_BY_TAG_VALUE = "agentcore-sdk" + class CoinbaseCdpConfigurationInput(TypedDict, total=False): """Configuration for Coinbase CDP credential provider. @@ -250,20 +255,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 @@ -391,6 +397,7 @@ def create_payment_manager( wait_for_ready: bool = False, max_wait: int = 300, poll_interval: int = 10, + tags: Optional[Dict[str, str]] = None, ) -> Dict[str, Any]: """Create a payment manager resource. @@ -404,6 +411,8 @@ def create_payment_manager( wait_for_ready: Whether to wait for manager to reach READY status max_wait: Maximum seconds to wait if wait_for_ready is True poll_interval: Seconds between checks if wait_for_ready is True + tags: Optional resource tags. The SDK stamps ``agentcore:created-by=agentcore-sdk`` + for origin attribution unless the caller already provides that key. Returns: Dictionary with paymentManagerArn, paymentManagerId, and status @@ -430,6 +439,10 @@ def create_payment_manager( if authorizer_configuration: params["authorizerConfiguration"] = authorizer_configuration + resource_tags = dict(tags) if tags else {} + resource_tags.setdefault(CREATED_BY_TAG_KEY, CREATED_BY_TAG_VALUE) + params["tags"] = resource_tags + response = self.payments_cp_client.create_payment_manager(**params) manager_arn = response.get("paymentManagerArn") 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", ) diff --git a/tests/bedrock_agentcore/payments/test_client.py b/tests/bedrock_agentcore/payments/test_client.py index 799e3db9..243dabcd 100644 --- a/tests/bedrock_agentcore/payments/test_client.py +++ b/tests/bedrock_agentcore/payments/test_client.py @@ -120,6 +120,49 @@ def test_create_payment_manager_success(self, mock_session, mock_boto3_client): assert result["paymentManagerId"] == "pm-123" assert result["status"] == "ACTIVE" + @patch("bedrock_agentcore.payments.client.boto3.client") + @patch("bedrock_agentcore.payments.client.boto3.Session") + def test_create_payment_manager_stamps_created_by_tag(self, mock_session, mock_boto3_client): + """SDK stamps agentcore:created-by=agentcore-sdk when caller omits it.""" + mock_session.return_value.region_name = "us-west-2" + mock_cp_client = MagicMock() + mock_boto3_client.return_value = mock_cp_client + mock_cp_client.create_payment_manager.return_value = { + "paymentManagerArn": "arn:aws:bedrock:us-west-2:123456789012:payment-manager/pm-123", + "paymentManagerId": "pm-123", + "status": "ACTIVE", + } + + client = PaymentClient(region_name="us-west-2") + client.create_payment_manager(name="test-manager", role_arn=self.role_arn) + + sent_tags = mock_cp_client.create_payment_manager.call_args[1]["tags"] + assert sent_tags["agentcore:created-by"] == "agentcore-sdk" + + @patch("bedrock_agentcore.payments.client.boto3.client") + @patch("bedrock_agentcore.payments.client.boto3.Session") + def test_create_payment_manager_preserves_caller_created_by_tag(self, mock_session, mock_boto3_client): + """Caller-supplied agentcore:created-by is not overwritten; other tags are kept.""" + mock_session.return_value.region_name = "us-west-2" + mock_cp_client = MagicMock() + mock_boto3_client.return_value = mock_cp_client + mock_cp_client.create_payment_manager.return_value = { + "paymentManagerArn": "arn:aws:bedrock:us-west-2:123456789012:payment-manager/pm-123", + "paymentManagerId": "pm-123", + "status": "ACTIVE", + } + + client = PaymentClient(region_name="us-west-2") + client.create_payment_manager( + name="test-manager", + role_arn=self.role_arn, + tags={"agentcore:created-by": "my-app", "team": "payments"}, + ) + + sent_tags = mock_cp_client.create_payment_manager.call_args[1]["tags"] + assert sent_tags["agentcore:created-by"] == "my-app" + assert sent_tags["team"] == "payments" + @patch("bedrock_agentcore.payments.client.boto3.client") @patch("bedrock_agentcore.payments.client.boto3.Session") def test_get_payment_manager_success(self, mock_session, mock_boto3_client):