Add feature usage tracking to the User-Agent header

packages/smithy-aws-core/.changes/next-release/smithy-aws-core-feature-e45dc6c320c64092a848bccb10a01ce3.json

@@ -0,0 +1,4 @@
+{
+  "type": "feature",
+  "description": "Add User Agent business metrics to the `User-Agent` header (`m/`), tracking the retry mode and endpoint override."
+}

packages/smithy-aws-core/src/smithy_aws_core/aio/endpoints/standard_regional.py

@@ -7,6 +7,7 @@
 from smithy_core.endpoints import Endpoint, EndpointResolverParams, resolve_static_uri
 from smithy_core.exceptions import EndpointResolutionError
 
+from ...business_metrics import AwsFeatureId, register_feature_id
 from ...endpoints import REGIONAL_ENDPOINT_CONFIG
 
 __all__ = ["AsyncStandardRegionalEndpointsResolver"]
@@ -20,6 +21,7 @@ def __init__(self, endpoint_prefix: str = "bedrock-runtime"):
 
     async def resolve_endpoint(self, params: EndpointResolverParams[Any]) -> Endpoint:
         if (static_uri := resolve_static_uri(params)) is not None:
+            register_feature_id(params.context, AwsFeatureId.ENDPOINT_OVERRIDE)
             return Endpoint(uri=static_uri)
 
         region_config = params.context.get(REGIONAL_ENDPOINT_CONFIG)

packages/smithy-aws-core/src/smithy_aws_core/business_metrics.py

@@ -0,0 +1,85 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  SPDX-License-Identifier: Apache-2.0
+"""Business metrics (feature ID) tracking for the User-Agent header.
+
+Features are recorded in the request properties of each call, where they are used.
+When the request is sent, their metric values are added to the ``m/`` section of the
+User-Agent header.
+"""
+
+import logging
+from collections.abc import Iterable
+from enum import StrEnum, unique
+
+from smithy_core.interfaces import TypedProperties
+from smithy_core.types import PropertyKey
+
+__all__ = [
+    "FEATURE_IDS",
+    "MAX_ENCODED_FEATURE_IDS_BYTES",
+    "AwsFeatureId",
+    "encode_feature_ids",
+    "register_feature_id",
+]
+
+logger = logging.getLogger(__name__)
+
+
+@unique
+class AwsFeatureId(StrEnum):
+    """Features that can be tracked.
+
+    Each member is named after the feature ID, and its value is the
+    feature's metric value, which is what is sent in the User-Agent header. Only
+    features this SDK tracks are listed.
+    """
+
+    RETRY_MODE_STANDARD = "E"
+    RETRY_MODE_ADAPTIVE = "F"
+    ENDPOINT_OVERRIDE = "N"
+
+
+FEATURE_IDS = PropertyKey(key="feature_ids", value_type=dict[AwsFeatureId, None])
+"""Features recorded for the current call, in the order they were registered."""
+
+
+def register_feature_id(properties: TypedProperties, feature_id: AwsFeatureId) -> None:
+    """Record that the current call used a feature.
+
+    Registering a feature that is already recorded has no effect. Features must be
+    registered in ``read_after_signing`` or earlier. Features registered later in an
+    attempt aren't sent with that attempt. Values that aren't ``AwsFeatureId``
+    members are ignored, so that tracking never fails a request.
+
+    :param properties: The request properties of the current call.
+    :param feature_id: The feature that was used.
+    """
+    if not isinstance(feature_id, AwsFeatureId):  # pyright: ignore[reportUnnecessaryIsInstance]
+        logger.debug("Ignoring unknown feature ID: %r", feature_id)
+        return
+    features = properties.get(FEATURE_IDS)
+    if features is None:
+        features = {}
+        properties[FEATURE_IDS] = features
+    features[feature_id] = None
+
+
+MAX_ENCODED_FEATURE_IDS_BYTES = 1024
+"""The maximum size of the encoded feature ID list."""
+
+
+def encode_feature_ids(feature_ids: Iterable[AwsFeatureId]) -> str:
+    """Encode features as a comma separated list of metric values of at most 1024 bytes.
+
+    Values are kept in the given order. If the list is too long, values are dropped
+    from the end so that no value is ever partially included.
+
+    :param feature_ids: The features to encode, without duplicates.
+    """
+    encoded = ""
+    for feature_id in feature_ids:
+        candidate = f"{encoded},{feature_id.value}" if encoded else feature_id.value
+        if len(candidate.encode("utf-8")) > MAX_ENCODED_FEATURE_IDS_BYTES:
+            break
+        encoded = candidate
+    return encoded

packages/smithy-aws-core/src/smithy_aws_core/interceptors/user_agent.py

@@ -5,16 +5,40 @@
 from typing import Any
 
 import smithy_core
+from smithy_core.aio.retries import (
+    RETRY_STRATEGY,
+    AsyncAdaptiveRetryStrategy,
+    AsyncStandardRetryStrategy,
+)
 from smithy_core.interceptors import Interceptor, RequestContext
+from smithy_core.interfaces import TypedProperties
+from smithy_core.types import PropertyKey
+from smithy_http import Field
+from smithy_http.aio.interfaces import HTTPRequest
 from smithy_http.interceptors.user_agent import USER_AGENT
 from smithy_http.user_agent import RawStringUserAgentComponent, UserAgentComponent
 
 from .. import __version__
+from ..business_metrics import (
+    FEATURE_IDS,
+    AwsFeatureId,
+    encode_feature_ids,
+    register_feature_id,
+)
 
 __all__ = ["UserAgentInterceptor"]
 
 _USERAGENT_SDK_NAME = "aws-sdk-python"
 
+_WRITTEN_USER_AGENT = PropertyKey(
+    key="aws_written_user_agent", value_type=tuple[str, str]
+)
+"""
+The User-Agent this attempt started with (without m/) and the one it wrote (with m/).
+
+The next attempt uses them, if there is one, so that m/ isn't added twice.
+"""
+
 
 class UserAgentInterceptor(Interceptor[Any, Any, Any, Any]):
     """Adds AWS fields to the UserAgent."""
@@ -65,6 +89,39 @@ def read_after_serialization(self, context: RequestContext[Any, Any]) -> None:
                     RawStringUserAgentComponent(self._ua_suffix)
                 )
 
+        self._register_retry_mode(context.properties)
+
+    def modify_before_transmit(
+        self, context: RequestContext[Any, HTTPRequest]
+    ) -> HTTPRequest:
+        # Runs on every attempt, after identity resolution and signing, so features
+        # registered up to that point are included. User-Agent isn't signed by AWS
+        # auth schemes, so changing it after signing doesn't affect the signature.
+        request = context.transport_request
+        encoded = encode_feature_ids(context.properties.get(FEATURE_IDS, {}))
+        if encoded and "User-Agent" in request.fields:
+            current = request.fields["User-Agent"].as_string()
+            # In generated clients each attempt gets a fresh User-Agent. The signer
+            # copies the headers and the generic interceptor rewrites them. As a
+            # fallback, for example with a signer that doesn't copy them, if a retry
+            # still has the header written here, start from its value before m/ was
+            # added, so m/ is never added twice and nothing else is removed.
+            previous = context.properties.get(_WRITTEN_USER_AGENT)
+            base = previous[0] if previous and current == previous[1] else current
+            written = f"{base} m/{encoded}"
+            request.fields.set_field(Field(name="User-Agent", values=[written]))
+            context.properties[_WRITTEN_USER_AGENT] = (base, written)
+        return request
+
+    def _register_retry_mode(self, properties: TypedProperties) -> None:
+        # Adaptive is a subclass of standard, so it's checked first. Other
+        # strategies have no feature ID and aren't reported.
+        strategy = properties.get(RETRY_STRATEGY)
+        if isinstance(strategy, AsyncAdaptiveRetryStrategy):
+            register_feature_id(properties, AwsFeatureId.RETRY_MODE_ADAPTIVE)
+        elif isinstance(strategy, AsyncStandardRetryStrategy):
+            register_feature_id(properties, AwsFeatureId.RETRY_MODE_STANDARD)
+
     def _build_sdk_metadata(self) -> list[UserAgentComponent]:
         return [
             UserAgentComponent(_USERAGENT_SDK_NAME, __version__),

packages/smithy-aws-core/tests/integration/test_user_agent_business_metrics.py

@@ -0,0 +1,124 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  SPDX-License-Identifier: Apache-2.0
+"""End-to-end User Agent 2.1 business metrics through a generated AWS client.
+
+Drives the generated awsJson1.0 client (``codegen-output/aws-json-1-0``) with a stub
+transport that records each request's User-Agent header, and checks the ``m/``
+section the SDK sends for the features each call uses.
+"""
+
+import os
+from pathlib import Path
+from typing import Any
+
+import pytest
+
+pytest.importorskip("awsjson10")
+
+from awsjson10.client import AsyncJSONRPC10Client
+from awsjson10.config import AsyncJSONRPC10Config
+from awsjson10.models import EmptyInputAndEmptyOutputInput
+from smithy_http import tuples_to_fields
+from smithy_http.aio import HTTPResponse
+
+_JSON_CONTENT_TYPE = ("content-type", "application/x-amz-json-1.0")
+
+
+class _RecordingTransport:
+    """Returns canned responses and records the User-Agent of every request."""
+
+    TIMEOUT_EXCEPTIONS = ()
+
+    def __init__(self, *, throttle_first: bool = False) -> None:
+        self.user_agents: list[str] = []
+        self._throttle_first = throttle_first
+
+    async def send(self, request: Any, *, request_config: Any = None) -> HTTPResponse:
+        self.user_agents.append(request.fields["User-Agent"].as_string())
+        if self._throttle_first and len(self.user_agents) == 1:
+            return HTTPResponse(
+                status=429,
+                fields=tuples_to_fields([_JSON_CONTENT_TYPE]),
+                body=b'{"__type":"ThrottlingException"}',
+            )
+        return HTTPResponse(
+            status=200,
+            fields=tuples_to_fields([_JSON_CONTENT_TYPE]),
+            body=b"{}",
+        )
+
+
+def _business_metrics(user_agent: str) -> list[str]:
+    return [part for part in user_agent.split() if part.startswith("m/")]
+
+
+@pytest.fixture(autouse=True)
+def _clean_environment(monkeypatch: pytest.MonkeyPatch) -> None:
+    # Keep the host's AWS settings from changing what each call reports.
+    for name in ("AWS_RETRY_MODE", "AWS_MAX_ATTEMPTS", "AWS_PROFILE"):
+        monkeypatch.delenv(name, raising=False)
+    # AWS_ENDPOINT_URL and the service-specific AWS_ENDPOINT_URL_<SERVICE> variables.
+    for name in list(os.environ):
+        if name.startswith("AWS_ENDPOINT_URL"):
+            monkeypatch.delenv(name)
+
+
+async def _call(
+    tmp_path: Path, transport: _RecordingTransport, **config: Any
+) -> list[str]:
+    missing = str(tmp_path / "missing")
+    resolved = await AsyncJSONRPC10Config.resolve(
+        config_file_path=missing,
+        credentials_file_path=missing,
+        region="us-west-2",
+        transport=transport,
+        aws_access_key_id="AKID",
+        aws_secret_access_key="secret",
+        **config,
+    )
+    client = AsyncJSONRPC10Client(config=resolved)
+    await client.empty_input_and_empty_output(EmptyInputAndEmptyOutputInput())
+    return transport.user_agents
+
+
+async def test_reports_standard_retry_mode_by_default(tmp_path: Path) -> None:
+    user_agents = await _call(tmp_path, _RecordingTransport())
+    assert _business_metrics(user_agents[0]) == ["m/E"]
+
+
+async def test_reports_adaptive_retry_mode(tmp_path: Path) -> None:
+    user_agents = await _call(tmp_path, _RecordingTransport(), retry_mode="adaptive")
+    assert _business_metrics(user_agents[0]) == ["m/F"]
+
+
+async def test_reports_adaptive_retry_mode_from_environment(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setenv("AWS_RETRY_MODE", "adaptive")
+    user_agents = await _call(tmp_path, _RecordingTransport())
+    assert _business_metrics(user_agents[0]) == ["m/F"]
+
+
+async def test_reports_endpoint_override(tmp_path: Path) -> None:
+    user_agents = await _call(
+        tmp_path, _RecordingTransport(), endpoint_uri="https://localhost:8443"
+    )
+    assert _business_metrics(user_agents[0]) == ["m/E,N"]
+
+
+async def test_reports_endpoint_override_from_environment(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setenv("AWS_ENDPOINT_URL_JSON_RPC_10", "https://localhost:8443")
+    user_agents = await _call(tmp_path, _RecordingTransport())
+    assert _business_metrics(user_agents[0]) == ["m/E,N"]
+
+
+async def test_retried_call_sends_same_business_metrics_once(tmp_path: Path) -> None:
+    user_agents = await _call(
+        tmp_path,
+        _RecordingTransport(throttle_first=True),
+        endpoint_uri="https://localhost:8443",
+    )
+    assert len(user_agents) == 2
+    assert [_business_metrics(ua) for ua in user_agents] == [["m/E,N"], ["m/E,N"]]

packages/smithy-aws-core/tests/unit/endpoints/test_standard_regional.py

@@ -8,6 +8,7 @@
 from smithy_aws_core.aio.endpoints.standard_regional import (
     AsyncStandardRegionalEndpointsResolver,
 )
+from smithy_aws_core.business_metrics import FEATURE_IDS, AwsFeatureId
 from smithy_aws_core.endpoints import REGIONAL_ENDPOINT_CONFIG
 from smithy_core import URI
 from smithy_core.endpoints import EndpointResolverParams
@@ -94,3 +95,21 @@ async def test_resolve_endpoint_with_sdk_endpoint_and_region():
     endpoint = await resolver.resolve_endpoint(params)
 
     assert endpoint.uri.host == "example.com"
+
+
+async def test_resolve_endpoint_registers_endpoint_override():
+    resolver = AsyncStandardRegionalEndpointsResolver(endpoint_prefix="service")
+    params = EndpointConfig.params("https://example.com")
+
+    await resolver.resolve_endpoint(params)
+
+    assert AwsFeatureId.ENDPOINT_OVERRIDE in params.context.get(FEATURE_IDS, {})
+
+
+async def test_resolve_regional_endpoint_does_not_register_endpoint_override():
+    resolver = AsyncStandardRegionalEndpointsResolver(endpoint_prefix="service")
+    params = EndpointConfig.params(region="us-west-2")
+
+    await resolver.resolve_endpoint(params)
+
+    assert AwsFeatureId.ENDPOINT_OVERRIDE not in params.context.get(FEATURE_IDS, {})