From c7e72f1261fdfe4a051f4a8fbeaa4e19d950e088 Mon Sep 17 00:00:00 2001 From: ubaskota <19787410+ubaskota@users.noreply.github.com> Date: Sun, 1 Nov 2026 16:47:42 -0500 Subject: [PATCH 1/3] Add feature usage tracking to the User-Agent header --- ...ture-e45dc6c320c64092a848bccb10a01ce3.json | 4 + .../aio/endpoints/standard_regional.py | 2 + .../src/smithy_aws_core/business_metrics.py | 85 ++++++++ .../interceptors/user_agent.py | 57 ++++++ .../test_user_agent_business_metrics.py | 124 ++++++++++++ .../unit/endpoints/test_standard_regional.py | 19 ++ .../unit/interceptors/test_user_agent.py | 187 +++++++++++++++++- .../tests/unit/test_business_metrics.py | 61 ++++++ ...ment-57cae49903c8488787e87c5d0ab8ceab.json | 4 + .../smithy-core/src/smithy_core/aio/client.py | 10 +- .../src/smithy_core/aio/retries.py | 5 + .../smithy-core/tests/unit/aio/test_client.py | 22 ++- scripts/api-snapshot/api.json | 26 +++ 13 files changed, 602 insertions(+), 4 deletions(-) create mode 100644 packages/smithy-aws-core/.changes/next-release/smithy-aws-core-feature-e45dc6c320c64092a848bccb10a01ce3.json create mode 100644 packages/smithy-aws-core/src/smithy_aws_core/business_metrics.py create mode 100644 packages/smithy-aws-core/tests/integration/test_user_agent_business_metrics.py create mode 100644 packages/smithy-aws-core/tests/unit/test_business_metrics.py create mode 100644 packages/smithy-core/.changes/next-release/smithy-core-enhancement-57cae49903c8488787e87c5d0ab8ceab.json diff --git a/packages/smithy-aws-core/.changes/next-release/smithy-aws-core-feature-e45dc6c320c64092a848bccb10a01ce3.json b/packages/smithy-aws-core/.changes/next-release/smithy-aws-core-feature-e45dc6c320c64092a848bccb10a01ce3.json new file mode 100644 index 000000000..e76d3f56c --- /dev/null +++ b/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." +} \ No newline at end of file diff --git a/packages/smithy-aws-core/src/smithy_aws_core/aio/endpoints/standard_regional.py b/packages/smithy-aws-core/src/smithy_aws_core/aio/endpoints/standard_regional.py index f575810f2..69cb20456 100644 --- a/packages/smithy-aws-core/src/smithy_aws_core/aio/endpoints/standard_regional.py +++ b/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) diff --git a/packages/smithy-aws-core/src/smithy_aws_core/business_metrics.py b/packages/smithy-aws-core/src/smithy_aws_core/business_metrics.py new file mode 100644 index 000000000..a88c8b1cc --- /dev/null +++ b/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 diff --git a/packages/smithy-aws-core/src/smithy_aws_core/interceptors/user_agent.py b/packages/smithy-aws-core/src/smithy_aws_core/interceptors/user_agent.py index 6064e65cc..42b586af8 100644 --- a/packages/smithy-aws-core/src/smithy_aws_core/interceptors/user_agent.py +++ b/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__), diff --git a/packages/smithy-aws-core/tests/integration/test_user_agent_business_metrics.py b/packages/smithy-aws-core/tests/integration/test_user_agent_business_metrics.py new file mode 100644 index 000000000..45cb2fa4d --- /dev/null +++ b/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_ 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"]] diff --git a/packages/smithy-aws-core/tests/unit/endpoints/test_standard_regional.py b/packages/smithy-aws-core/tests/unit/endpoints/test_standard_regional.py index 1da0c2712..a12257958 100644 --- a/packages/smithy-aws-core/tests/unit/endpoints/test_standard_regional.py +++ b/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, {}) diff --git a/packages/smithy-aws-core/tests/unit/interceptors/test_user_agent.py b/packages/smithy-aws-core/tests/unit/interceptors/test_user_agent.py index a9d1f4acb..e05fa5a73 100644 --- a/packages/smithy-aws-core/tests/unit/interceptors/test_user_agent.py +++ b/packages/smithy-aws-core/tests/unit/interceptors/test_user_agent.py @@ -1,13 +1,31 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 +from copy import copy from typing import Any from unittest.mock import Mock import pytest +from smithy_aws_core.business_metrics import ( + FEATURE_IDS, + AwsFeatureId, + register_feature_id, +) from smithy_aws_core.interceptors.user_agent import UserAgentInterceptor -from smithy_core.interceptors import RequestContext +from smithy_core import URI +from smithy_core.aio.retries import ( + RETRY_STRATEGY, + AsyncAdaptiveRetryStrategy, + AsyncSimpleRetryStrategy, + AsyncStandardRetryStrategy, +) +from smithy_core.interceptors import InputContext, RequestContext from smithy_core.types import TypedProperties +from smithy_http import Field, Fields +from smithy_http.aio import HTTPRequest from smithy_http.interceptors.user_agent import USER_AGENT +from smithy_http.interceptors.user_agent import ( + UserAgentInterceptor as GenericUserAgentInterceptor, +) from smithy_http.user_agent import UserAgent, UserAgentComponent @@ -34,3 +52,170 @@ def test_adds_execution_env(monkeypatch: pytest.MonkeyPatch) -> None: def test_omits_execution_env_when_unset(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv("AWS_EXECUTION_ENV", raising=False) assert "exec-env/" not in _user_agent_string() + + +_BASE_USER_AGENT = "aws-sdk-python/1.0.0 ua/2.1 lang/python#3.12" + + +def _interceptor() -> UserAgentInterceptor: + return UserAgentInterceptor( + ua_suffix=None, ua_app_id=None, sdk_version="1.0.0", service_id="svc" + ) + + +def _request(user_agent: str = _BASE_USER_AGENT) -> HTTPRequest: + # The generic user agent interceptor writes the header before signing. + return HTTPRequest( + method="POST", + destination=URI(host="example.com"), + fields=Fields([Field(name="User-Agent", values=[user_agent])]), + body=b"", + ) + + +def _transmit( + interceptor: UserAgentInterceptor, + properties: TypedProperties, + request: HTTPRequest | None = None, +) -> str: + context = RequestContext[Any, HTTPRequest]( + request=Mock(), + transport_request=request if request is not None else _request(), + properties=properties, + ) + return interceptor.modify_before_transmit(context).fields["User-Agent"].as_string() + + +def test_appends_business_metrics() -> None: + properties = TypedProperties() + register_feature_id(properties, AwsFeatureId.RETRY_MODE_STANDARD) + register_feature_id(properties, AwsFeatureId.ENDPOINT_OVERRIDE) + assert _transmit(_interceptor(), properties) == f"{_BASE_USER_AGENT} m/E,N" + + +def test_skips_business_metrics_without_user_agent_header() -> None: + properties = TypedProperties() + register_feature_id(properties, AwsFeatureId.ENDPOINT_OVERRIDE) + request = HTTPRequest( + method="POST", destination=URI(host="example.com"), fields=Fields(), body=b"" + ) + context = RequestContext[Any, HTTPRequest]( + request=Mock(), transport_request=request, properties=properties + ) + + result = _interceptor().modify_before_transmit(context) + + assert "User-Agent" not in result.fields + + +def test_leaves_header_unchanged_without_features() -> None: + assert _transmit(_interceptor(), TypedProperties()) == _BASE_USER_AGENT + + +def test_retry_does_not_duplicate_business_metrics() -> None: + # A retry shallow-copies the request, so it starts with the previous attempt's + # User-Agent, m/ included. The header must not have duplicate m/. + interceptor = _interceptor() + properties = TypedProperties() + register_feature_id(properties, AwsFeatureId.ENDPOINT_OVERRIDE) + request = _request() + _transmit(interceptor, properties, request) + retry = copy(request) + + user_agent = _transmit(interceptor, properties, retry) + + assert user_agent == f"{_BASE_USER_AGENT} m/N" + + +def test_retry_with_more_features_replaces_previous_business_metrics() -> None: + interceptor = _interceptor() + properties = TypedProperties() + register_feature_id(properties, AwsFeatureId.ENDPOINT_OVERRIDE) + request = _request() + _transmit(interceptor, properties, request) + register_feature_id(properties, AwsFeatureId.RETRY_MODE_STANDARD) + + user_agent = _transmit(interceptor, properties, copy(request)) + + assert user_agent == f"{_BASE_USER_AGENT} m/N,E" + + +def test_keeps_m_prefixed_text_from_the_user() -> None: + interceptor = _interceptor() + properties = TypedProperties() + register_feature_id(properties, AwsFeatureId.ENDPOINT_OVERRIDE) + request = _request(f"{_BASE_USER_AGENT} myframework m/legacy-tag") + + user_agent = _transmit(interceptor, properties, request) + + assert user_agent == f"{_BASE_USER_AGENT} myframework m/legacy-tag m/N" + + +def test_retry_with_generic_interceptor_sends_business_metrics_once() -> None: + # Mirrors a generated client, which runs both interceptors on every attempt. + # A retried call must send the same User-Agent each time, with m/ only once, + # and without removing anything the user added, such as the m/legacy-tag in + # this test's suffix. + generic = GenericUserAgentInterceptor() + interceptor = UserAgentInterceptor( + ua_suffix="myframework m/legacy-tag", + ua_app_id=None, + sdk_version="1.0.0", + service_id="svc", + ) + properties = TypedProperties() + generic.read_before_execution(InputContext(request=Mock(), properties=properties)) + request: Any = _request() + interceptor.read_after_serialization( + RequestContext[Any, Any]( + request=Mock(), transport_request=request, properties=properties + ) + ) + register_feature_id(properties, AwsFeatureId.ENDPOINT_OVERRIDE) + + user_agents: list[str] = [] + for _ in range(2): + # Each attempt starts from a shallow copy, as retries do. + context = RequestContext[Any, Any]( + request=Mock(), transport_request=copy(request), properties=properties + ) + generic.modify_before_signing(context) + request = interceptor.modify_before_transmit(context) + user_agents.append(request.fields["User-Agent"].as_string()) + + assert user_agents[0] == user_agents[1] + assert user_agents[0].endswith(" myframework m/legacy-tag m/N") + assert user_agents[0].count(" m/N") == 1 + + +@pytest.mark.parametrize( + ("retry_strategy", "expected"), + [ + (AsyncStandardRetryStrategy(), [AwsFeatureId.RETRY_MODE_STANDARD]), + (AsyncAdaptiveRetryStrategy(), [AwsFeatureId.RETRY_MODE_ADAPTIVE]), + (AsyncSimpleRetryStrategy(), []), + (None, []), + ], +) +def test_registers_retry_mode( + retry_strategy: Any, expected: list[AwsFeatureId] +) -> None: + properties = TypedProperties() + properties[USER_AGENT] = UserAgent() + if retry_strategy is not None: + properties[RETRY_STRATEGY] = retry_strategy + context = RequestContext[Any, Any]( + request=Mock(), transport_request=Mock(), properties=properties + ) + _interceptor().read_after_serialization(context) + assert list(properties.get(FEATURE_IDS, {})) == expected + + +def test_registers_retry_mode_without_user_agent() -> None: + properties = TypedProperties() + properties[RETRY_STRATEGY] = AsyncStandardRetryStrategy() + context = RequestContext[Any, Any]( + request=Mock(), transport_request=Mock(), properties=properties + ) + _interceptor().read_after_serialization(context) + assert list(properties[FEATURE_IDS]) == [AwsFeatureId.RETRY_MODE_STANDARD] diff --git a/packages/smithy-aws-core/tests/unit/test_business_metrics.py b/packages/smithy-aws-core/tests/unit/test_business_metrics.py new file mode 100644 index 000000000..13dbc1742 --- /dev/null +++ b/packages/smithy-aws-core/tests/unit/test_business_metrics.py @@ -0,0 +1,61 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +from typing import Any, cast + +import pytest +from smithy_aws_core import business_metrics +from smithy_aws_core.business_metrics import ( + FEATURE_IDS, + MAX_ENCODED_FEATURE_IDS_BYTES, + AwsFeatureId, + encode_feature_ids, + register_feature_id, +) +from smithy_core.types import TypedProperties + + +def test_register_feature_id_keeps_order_without_duplicates() -> None: + properties = TypedProperties() + register_feature_id(properties, AwsFeatureId.RETRY_MODE_STANDARD) + register_feature_id(properties, AwsFeatureId.ENDPOINT_OVERRIDE) + register_feature_id(properties, AwsFeatureId.RETRY_MODE_STANDARD) + assert list(properties[FEATURE_IDS]) == [ + AwsFeatureId.RETRY_MODE_STANDARD, + AwsFeatureId.ENDPOINT_OVERRIDE, + ] + + +def test_encode_feature_ids_uses_metric_values() -> None: + encoded = encode_feature_ids( + [AwsFeatureId.RETRY_MODE_STANDARD, AwsFeatureId.ENDPOINT_OVERRIDE] + ) + assert encoded == "E,N" + + +def test_encode_feature_ids_empty() -> None: + assert encode_feature_ids([]) == "" + + +def test_encode_feature_ids_drops_whole_values_past_limit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # A 4-byte limit fits "E,F" but not "E,F,N", so N is dropped whole. + monkeypatch.setattr(business_metrics, "MAX_ENCODED_FEATURE_IDS_BYTES", 4) + encoded = encode_feature_ids( + [ + AwsFeatureId.RETRY_MODE_STANDARD, + AwsFeatureId.RETRY_MODE_ADAPTIVE, + AwsFeatureId.ENDPOINT_OVERRIDE, + ] + ) + assert encoded == "E,F" + + +def test_max_encoded_feature_ids_bytes_is_1024() -> None: + assert MAX_ENCODED_FEATURE_IDS_BYTES == 1024 + + +def test_register_feature_id_ignores_plain_strings() -> None: + properties = TypedProperties() + register_feature_id(properties, cast(Any, "N")) + assert FEATURE_IDS not in properties diff --git a/packages/smithy-core/.changes/next-release/smithy-core-enhancement-57cae49903c8488787e87c5d0ab8ceab.json b/packages/smithy-core/.changes/next-release/smithy-core-enhancement-57cae49903c8488787e87c5d0ab8ceab.json new file mode 100644 index 000000000..0da8f2ece --- /dev/null +++ b/packages/smithy-core/.changes/next-release/smithy-core-enhancement-57cae49903c8488787e87c5d0ab8ceab.json @@ -0,0 +1,4 @@ +{ + "type": "enhancement", + "description": "Add the retry strategy chosen for a call to its request properties under `RETRY_STRATEGY`." +} \ No newline at end of file diff --git a/packages/smithy-core/src/smithy_core/aio/client.py b/packages/smithy-core/src/smithy_core/aio/client.py index 257c99ec7..6cf051c97 100644 --- a/packages/smithy-core/src/smithy_core/aio/client.py +++ b/packages/smithy-core/src/smithy_core/aio/client.py @@ -49,7 +49,11 @@ from .interfaces.auth import AsyncAuthScheme from .interfaces.eventstream import AsyncEventReceiver from .interfaces.retries import AsyncRetryStrategy -from .retries import AsyncAdaptiveRetryStrategy, AsyncRetryStrategyResolver +from .retries import ( + RETRY_STRATEGY, + AsyncAdaptiveRetryStrategy, + AsyncRetryStrategyResolver, +) from .utils import close, seek if TYPE_CHECKING: @@ -806,10 +810,12 @@ def _build_call[I: SerializeableShape, O: DeserializeableShape]( pipeline = AsyncRequestPipeline( protocol=config.protocol, transport=config.transport ) + context = _TypedProperties({"config": config}) + context[RETRY_STRATEGY] = retry_strategy call = AsyncClientCall( input=input, operation=operation, - context=_TypedProperties({"config": config}), + context=context, interceptor=InterceptorChain(config.interceptors), auth_scheme_resolver=config.auth_scheme_resolver, supported_auth_schemes=config.auth_schemes, diff --git a/packages/smithy-core/src/smithy_core/aio/retries.py b/packages/smithy-core/src/smithy_core/aio/retries.py index 7e21a7767..462b0b12a 100644 --- a/packages/smithy-core/src/smithy_core/aio/retries.py +++ b/packages/smithy-core/src/smithy_core/aio/retries.py @@ -19,9 +19,11 @@ StandardRetryQuota, StandardRetryToken, ) +from ..types import PropertyKey from .interfaces.retries import AsyncRetryStrategy __all__ = [ + "RETRY_STRATEGY", "AsyncAdaptiveRetryStrategy", "AsyncClientRateLimiter", "AsyncRequestRateTracker", @@ -31,6 +33,9 @@ "AsyncTokenBucket", ] +RETRY_STRATEGY = PropertyKey(key="retry_strategy", value_type=AsyncRetryStrategy) +"""The retry strategy the client chose for the call.""" + class AsyncRetryStrategyResolver: """Retry strategy resolver that caches retry strategies based on configuration options. diff --git a/packages/smithy-core/tests/unit/aio/test_client.py b/packages/smithy-core/tests/unit/aio/test_client.py index 42d99bdfc..bafc15fc1 100644 --- a/packages/smithy-core/tests/unit/aio/test_client.py +++ b/packages/smithy-core/tests/unit/aio/test_client.py @@ -11,7 +11,12 @@ from smithy_core.aio.client import AsyncClient from smithy_core.aio.eventstream import AsyncDuplexEventStream, AsyncInputEventStream from smithy_core.aio.interfaces import AsyncClientTransport -from smithy_core.aio.retries import AsyncAdaptiveRetryStrategy, AsyncSimpleRetryStrategy +from smithy_core.aio.retries import ( + RETRY_STRATEGY, + AsyncAdaptiveRetryStrategy, + AsyncSimpleRetryStrategy, + AsyncStandardRetryStrategy, +) from smithy_core.aio.types import AsyncBytesReader from smithy_core.auth import AuthOption from smithy_core.exceptions import CallError, SmithyError, UnsupportedTransportError @@ -291,6 +296,21 @@ def closed(self) -> bool: return self._closed +def test_build_call_adds_retry_strategy_to_request_properties() -> None: + retry_strategy = AsyncStandardRetryStrategy() + config = Mock(interceptors=[]) + client = _StubClient() + _, call = client._build_call( # pyright: ignore[reportPrivateUsage] + cast(Any, Mock()), + cast(Any, Mock()), + config=config, + retry_strategy=retry_strategy, + ) + + assert call.context[RETRY_STRATEGY] is retry_strategy + assert call.retry_strategy is retry_strategy + + async def test_close_is_idempotent() -> None: client = _StubClient() diff --git a/scripts/api-snapshot/api.json b/scripts/api-snapshot/api.json index 28f12f110..66dabb7ac 100644 --- a/scripts/api-snapshot/api.json +++ b/scripts/api-snapshot/api.json @@ -169,11 +169,13 @@ "smithy_aws_core.aio.endpoints.standard_regional": { "AsyncEndpointResolver": "rex-protocol[async]", "AsyncStandardRegionalEndpointsResolver": "class[async]", + "AwsFeatureId": "rex-class", "Endpoint": "rex-class", "EndpointResolutionError": "rex-exception", "EndpointResolverParams": "rex-class", "REGIONAL_ENDPOINT_CONFIG": "rex-PropertyKey", "URI": "rex-class", + "register_feature_id": "rex-function", "resolve_static_uri": "rex-function" }, "smithy_aws_core.aio.identity": { @@ -406,6 +408,16 @@ "smithy_aws_core.auth.sigv4": { "SigV4AuthScheme": "class" }, + "smithy_aws_core.business_metrics": { + "AwsFeatureId": "class", + "FEATURE_IDS": "PropertyKey", + "MAX_ENCODED_FEATURE_IDS_BYTES": "int", + "PropertyKey": "rex-class", + "TypedProperties": "rex-protocol", + "encode_feature_ids": "function", + "logger": "Logger", + "register_feature_id": "function" + }, "smithy_aws_core.config": { "AsyncAwsConfig": "rex-class[async]", "AsyncDefaultFileSystem": "rex-class[async]", @@ -563,12 +575,23 @@ }, "smithy_aws_core.interceptors": {}, "smithy_aws_core.interceptors.user_agent": { + "AsyncAdaptiveRetryStrategy": "rex-class[async]", + "AsyncStandardRetryStrategy": "rex-class[async]", + "AwsFeatureId": "rex-class", + "FEATURE_IDS": "rex-PropertyKey", + "Field": "rex-class", + "HTTPRequest": "rex-protocol[async]", "Interceptor": "rex-class", + "PropertyKey": "rex-class", + "RETRY_STRATEGY": "rex-PropertyKey", "RawStringUserAgentComponent": "rex-class", "RequestContext": "rex-class", + "TypedProperties": "rex-protocol", "USER_AGENT": "rex-PropertyKey", "UserAgentComponent": "rex-class", "UserAgentInterceptor": "class", + "encode_feature_ids": "rex-function", + "register_feature_id": "rex-function", "smithy_core": "rex-module" }, "smithy_aws_core.interfaces": {}, @@ -749,6 +772,7 @@ "InterceptorChain": "rex-class", "OutputContext": "rex-class", "PropertyKey": "rex-class", + "RETRY_STRATEGY": "rex-PropertyKey", "Request": "rex-protocol[async]", "RequestContext": "rex-class", "Response": "rex-protocol[async]", @@ -855,6 +879,8 @@ "CubicCalculator": "rex-class", "ExponentialBackoffJitterType": "rex-class", "ExponentialRetryBackoffStrategy": "rex-class", + "PropertyKey": "rex-class", + "RETRY_STRATEGY": "PropertyKey", "RetryError": "rex-exception", "RetryStrategyOptions": "rex-class", "SimpleRetryToken": "rex-class", From cb7b751da9b5b932b58d025ed470dacae0a9b25a Mon Sep 17 00:00:00 2001 From: ubaskota <19787410+ubaskota@users.noreply.github.com> Date: Mon, 2 Nov 2026 11:40:31 -0500 Subject: [PATCH 2/3] Update changelog --- ...smithy-core-feature-57cae49903c8488787e87c5d0ab8ceab.json} | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) rename packages/smithy-core/.changes/next-release/{smithy-core-enhancement-57cae49903c8488787e87c5d0ab8ceab.json => smithy-core-feature-57cae49903c8488787e87c5d0ab8ceab.json} (81%) diff --git a/packages/smithy-core/.changes/next-release/smithy-core-enhancement-57cae49903c8488787e87c5d0ab8ceab.json b/packages/smithy-core/.changes/next-release/smithy-core-feature-57cae49903c8488787e87c5d0ab8ceab.json similarity index 81% rename from packages/smithy-core/.changes/next-release/smithy-core-enhancement-57cae49903c8488787e87c5d0ab8ceab.json rename to packages/smithy-core/.changes/next-release/smithy-core-feature-57cae49903c8488787e87c5d0ab8ceab.json index 0da8f2ece..836404d27 100644 --- a/packages/smithy-core/.changes/next-release/smithy-core-enhancement-57cae49903c8488787e87c5d0ab8ceab.json +++ b/packages/smithy-core/.changes/next-release/smithy-core-feature-57cae49903c8488787e87c5d0ab8ceab.json @@ -1,4 +1,4 @@ { - "type": "enhancement", + "type": "feature", "description": "Add the retry strategy chosen for a call to its request properties under `RETRY_STRATEGY`." -} \ No newline at end of file +} From 0b2aacd8cc4f0fb9c491d93dd37e88523dd50d7f Mon Sep 17 00:00:00 2001 From: ubaskota <19787410+ubaskota@users.noreply.github.com> Date: Mon, 2 Nov 2026 15:46:34 -0500 Subject: [PATCH 3/3] Address reviews --- .../interceptors/user_agent.py | 25 +++-------------- .../unit/interceptors/test_user_agent.py | 28 ------------------- scripts/api-snapshot/api.json | 1 - 3 files changed, 4 insertions(+), 50 deletions(-) diff --git a/packages/smithy-aws-core/src/smithy_aws_core/interceptors/user_agent.py b/packages/smithy-aws-core/src/smithy_aws_core/interceptors/user_agent.py index 42b586af8..9b82f4f21 100644 --- a/packages/smithy-aws-core/src/smithy_aws_core/interceptors/user_agent.py +++ b/packages/smithy-aws-core/src/smithy_aws_core/interceptors/user_agent.py @@ -12,7 +12,6 @@ ) 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 @@ -30,15 +29,6 @@ _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.""" @@ -100,17 +90,10 @@ def modify_before_transmit( 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) + user_agent = request.fields["User-Agent"].as_string() + request.fields.set_field( + Field(name="User-Agent", values=[f"{user_agent} m/{encoded}"]) + ) return request def _register_retry_mode(self, properties: TypedProperties) -> None: diff --git a/packages/smithy-aws-core/tests/unit/interceptors/test_user_agent.py b/packages/smithy-aws-core/tests/unit/interceptors/test_user_agent.py index e05fa5a73..4aa4bad21 100644 --- a/packages/smithy-aws-core/tests/unit/interceptors/test_user_agent.py +++ b/packages/smithy-aws-core/tests/unit/interceptors/test_user_agent.py @@ -112,34 +112,6 @@ def test_leaves_header_unchanged_without_features() -> None: assert _transmit(_interceptor(), TypedProperties()) == _BASE_USER_AGENT -def test_retry_does_not_duplicate_business_metrics() -> None: - # A retry shallow-copies the request, so it starts with the previous attempt's - # User-Agent, m/ included. The header must not have duplicate m/. - interceptor = _interceptor() - properties = TypedProperties() - register_feature_id(properties, AwsFeatureId.ENDPOINT_OVERRIDE) - request = _request() - _transmit(interceptor, properties, request) - retry = copy(request) - - user_agent = _transmit(interceptor, properties, retry) - - assert user_agent == f"{_BASE_USER_AGENT} m/N" - - -def test_retry_with_more_features_replaces_previous_business_metrics() -> None: - interceptor = _interceptor() - properties = TypedProperties() - register_feature_id(properties, AwsFeatureId.ENDPOINT_OVERRIDE) - request = _request() - _transmit(interceptor, properties, request) - register_feature_id(properties, AwsFeatureId.RETRY_MODE_STANDARD) - - user_agent = _transmit(interceptor, properties, copy(request)) - - assert user_agent == f"{_BASE_USER_AGENT} m/N,E" - - def test_keeps_m_prefixed_text_from_the_user() -> None: interceptor = _interceptor() properties = TypedProperties() diff --git a/scripts/api-snapshot/api.json b/scripts/api-snapshot/api.json index 66dabb7ac..70d5733b0 100644 --- a/scripts/api-snapshot/api.json +++ b/scripts/api-snapshot/api.json @@ -582,7 +582,6 @@ "Field": "rex-class", "HTTPRequest": "rex-protocol[async]", "Interceptor": "rex-class", - "PropertyKey": "rex-class", "RETRY_STRATEGY": "rex-PropertyKey", "RawStringUserAgentComponent": "rex-class", "RequestContext": "rex-class",