[Bugfix] UserAgent Refactoring and Remove Duplicate Custom Prefix (#44189)

* change workloads based on feedback

* add staging yml file

* add staging yml file

* refactor useragent changes for enabled features and fix bug on init

* update samples

* update changelog

* refactor tests, add suffix tests, react to copilot

* pylint unused imports

* Update sdk/cosmos/azure-cosmos/azure/cosmos/user_agent_policy.py

---------

Co-authored-by: Simon Moreno <[email protected]>

Author: tvaron3

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -11,6 +11,7 @@
 #### Bugs Fixed
 * Fixed bug where client timeout/read_timeout values were not properly enforced[PR 42652](https://github.com/Azure/azure-sdk-for-python/pull/42652).
 * Fixed bug when passing in None for some option in `query_items` would cause unexpected errors. See [PR 44098](https://github.com/Azure/azure-sdk-for-python/pull/44098)
+* Fixed bug where first metadata requests have duplicated custom `user_agent` in headers. See [PR 44189](https://github.com/Azure/azure-sdk-for-python/pull/44189)
 
 #### Other Changes
 * Added cross-regional retries for 503 (Service Unavailable) errors. See [PR 41588](https://github.com/Azure/azure-sdk-for-python/pull/41588).

sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py

@@ -38,7 +38,6 @@
     HTTPPolicy,
     ContentDecodePolicy,
     HeadersPolicy,
-    UserAgentPolicy,
     NetworkTraceLoggingPolicy,
     CustomHookPolicy,
     DistributedTracingPolicy,
@@ -49,6 +48,7 @@
 from azure.core.utils import CaseInsensitiveDict
 
 from . import _base as base
+from .user_agent_policy import CosmosUserAgentPolicy
 from ._global_partition_endpoint_manager_per_partition_automatic_failover import _GlobalPartitionEndpointManagerForPerPartitionAutomaticFailover # pylint: disable=line-too-long
 from . import _query_iterable as query_iterable
 from . import _runtime_constants as runtime_constants
@@ -231,7 +231,7 @@ def __init__( # pylint: disable=too-many-statements
         policies = [
             HeadersPolicy(**kwargs),
             ProxyPolicy(proxies=proxies),
-            UserAgentPolicy(base_user_agent=self._user_agent, **kwargs),
+            CosmosUserAgentPolicy(base_user_agent=self._user_agent, **kwargs),
             ContentDecodePolicy(),
             retry_policy,
             credentials_policy,
@@ -245,6 +245,8 @@ def __init__( # pylint: disable=too-many-statements
                 **kwargs
             ),
         ]
+        # after passing in the user_agent into the user agent policy the user_agent is no longer needed
+        kwargs.pop("user_agent", None)
 
         transport = kwargs.pop("transport", None)
         self.pipeline_client: PipelineClient[HttpRequest, HttpResponse] = PipelineClient(

sdk/cosmos/azure-cosmos/azure/cosmos/_synchronized_request.py

@@ -34,9 +34,7 @@
 from ._availability_strategy_handler import execute_with_hedging
 from ._constants import _Constants
 from ._request_object import RequestObject
-from ._utils import get_user_agent_features
 from .documents import _OperationType
-from .http_constants import ResourceType
 
 # cspell:ignore ppaf
 def _is_readable_stream(obj):
@@ -133,15 +131,6 @@ def _Request(global_endpoint_manager, request_params, connection_policy, pipelin
 
     parse_result = urlparse(request.url)
 
-    # Add relevant enabled features to user agent for debugging
-    if request.headers[http_constants.HttpHeaders.ThinClientProxyResourceType] == http_constants.ResourceType.Document:
-        user_agent_features = get_user_agent_features(global_endpoint_manager)
-        if len(user_agent_features) > 0:
-            user_agent = kwargs.pop("user_agent", global_endpoint_manager.client._user_agent)
-            user_agent = "{} {}".format(user_agent, user_agent_features)
-            kwargs.update({"user_agent": user_agent})
-            kwargs.update({"user_agent_overwrite": True})
-
     # The requests library now expects header values to be strings only starting 2.11,
     # and will raise an error on validation if they are not, so casting all header values to strings.
     request.headers.update({header: str(value) for header, value in request.headers.items()})
@@ -220,7 +209,7 @@ def _is_availability_strategy_applicable(request_params: RequestObject) -> bool:
     """
     return (request_params.availability_strategy_config is not None and
             not request_params.is_hedging_request and
-            request_params.resource_type == ResourceType.Document and
+            request_params.resource_type == http_constants.ResourceType.Document and
             (not _OperationType.IsWriteOperation(request_params.operation_type) or
              request_params.retry_write > 0))
 

sdk/cosmos/azure-cosmos/azure/cosmos/aio/_asynchronous_request.py

@@ -36,9 +36,7 @@
 from .._constants import _Constants
 from .._request_object import RequestObject
 from .._synchronized_request import _request_body_from_data, _replace_url_prefix
-from .._utils import get_user_agent_features
 from ..documents import _OperationType
-from ..http_constants import ResourceType
 
 # cspell:ignore ppaf
 async def _Request(global_endpoint_manager, request_params, connection_policy, pipeline_client, request, **kwargs): # pylint: disable=too-many-statements
@@ -97,15 +95,6 @@ async def _Request(global_endpoint_manager, request_params, connection_policy, p
 
     parse_result = urlparse(request.url)
 
-    # Add relevant enabled features to user agent for debugging
-    if request.headers[http_constants.HttpHeaders.ThinClientProxyResourceType] == http_constants.ResourceType.Document:
-        user_agent_features = get_user_agent_features(global_endpoint_manager)
-        if len(user_agent_features) > 0:
-            user_agent = kwargs.pop("user_agent", global_endpoint_manager.client._user_agent)
-            user_agent = "{} {}".format(user_agent, user_agent_features)
-            kwargs.update({"user_agent": user_agent})
-            kwargs.update({"user_agent_overwrite": True})
-
     # The requests library now expects header values to be strings only starting 2.11,
     # and will raise an error on validation if they are not, so casting all header values to strings.
     request.headers.update({header: str(value) for header, value in request.headers.items()})
@@ -190,7 +179,7 @@ def _is_availability_strategy_applicable(request_params: RequestObject) -> bool:
     """
     return (request_params.availability_strategy_config is not None and
             not request_params.is_hedging_request and
-            request_params.resource_type == ResourceType.Document and
+            request_params.resource_type == http_constants.ResourceType.Document and
             (not _OperationType.IsWriteOperation(request_params.operation_type) or
              request_params.retry_write > 0))
 

sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py

@@ -38,7 +38,6 @@
     AsyncHTTPPolicy,
     ContentDecodePolicy,
     HeadersPolicy,
-    UserAgentPolicy,
     NetworkTraceLoggingPolicy,
     CustomHookPolicy,
     DistributedTracingPolicy,
@@ -79,8 +78,7 @@
 from .._cosmos_http_logging_policy import CosmosHttpLoggingPolicy
 from .._range_partition_resolver import RangePartitionResolver
 from ._read_items_helper_async import ReadItemsHelperAsync
-
-
+from ..user_agent_policy import CosmosUserAgentPolicy
 
 
 class CredentialDict(TypedDict, total=False):
@@ -230,7 +228,7 @@ def __init__( # pylint: disable=too-many-statements
         policies = [
             HeadersPolicy(**kwargs),
             ProxyPolicy(proxies=proxies),
-            UserAgentPolicy(base_user_agent=self._user_agent, **kwargs),
+            CosmosUserAgentPolicy(base_user_agent=self._user_agent, **kwargs),
             ContentDecodePolicy(),
             retry_policy,
             credentials_policy,
@@ -244,6 +242,8 @@ def __init__( # pylint: disable=too-many-statements
                 **kwargs
             ),
         ]
+        # after passing in the user_agent into the user agent policy the user_agent is no longer needed
+        kwargs.pop("user_agent", None)
 
         transport = kwargs.pop("transport", None)
         self.pipeline_client: AsyncPipelineClient[HttpRequest, AsyncHttpResponse] = AsyncPipelineClient(

sdk/cosmos/azure-cosmos/azure/cosmos/user_agent_policy.py

@@ -0,0 +1,37 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See LICENSE.txt in the project root for
+# license information.
+# -------------------------------------------------------------------------
+
+from azure.core.pipeline import PipelineRequest
+from azure.core.pipeline.policies import UserAgentPolicy
+from azure.core.pipeline.policies._universal import HTTPRequestType
+
+from azure.cosmos._utils import get_user_agent_features
+
+
+class CosmosUserAgentPolicy(UserAgentPolicy):
+    """Custom user agent policy for Cosmos DB that appends feature flags to the user agent string.
+
+    This policy extends the standard UserAgentPolicy to include Cosmos-specific feature flags
+    (e.g., circuit breaker, per-partition automatic failover) in the user agent header for
+    debugging and telemetry purposes.
+    """
+
+    def on_request(self, request: PipelineRequest[HTTPRequestType]) -> None:
+        """Modifies the User-Agent header before the request is sent.
+
+        :param request: The PipelineRequest object
+        :type request: ~azure.core.pipeline.PipelineRequest
+        """
+        options_dict = request.context.options
+        # Add relevant enabled features to user agent for debugging
+        if "global_endpoint_manager" in options_dict:
+            global_endpoint_manager = options_dict["global_endpoint_manager"]
+            user_agent_features = get_user_agent_features(global_endpoint_manager)
+            if len(user_agent_features) > 0:
+                user_agent = "{} {}".format(self._user_agent, user_agent_features)
+                options_dict["user_agent"] = user_agent
+                options_dict["user_agent_overwrite"] = True
+        super().on_request(request)

sdk/cosmos/azure-cosmos/samples/user_agent_management.py

@@ -74,7 +74,7 @@ def create_client_with_custom_user_agent():
     client = cosmos_client.CosmosClient(
         HOST,
         {'masterKey': MASTER_KEY},
-        user_agent=custom_user_agent
+        user_agent_suffix=custom_user_agent
     )
 
     print(f'Client created with custom user agent: {custom_user_agent}')
@@ -106,7 +106,7 @@ def create_client_with_detailed_user_agent():
     client = cosmos_client.CosmosClient(
         HOST,
         {'masterKey': MASTER_KEY},
-        user_agent=detailed_user_agent
+        user_agent_suffix=detailed_user_agent
     )
 
     print(f'Client created with detailed user agent: {detailed_user_agent}')
@@ -128,23 +128,23 @@ def create_multiple_clients_with_different_user_agents():
     ingestion_client = cosmos_client.CosmosClient(
         HOST,
         {'masterKey': MASTER_KEY},
-        user_agent="DataIngestionService/1.0.0"
+        user_agent_suffix="DataIngestionService/1.0.0"
     )
     print('Created client for data ingestion: DataIngestionService/1.0.0')
 
     # Client for query service
     query_client = cosmos_client.CosmosClient(
         HOST,
         {'masterKey': MASTER_KEY},
-        user_agent="QueryService/1.0.0"
+        user_agent_suffix="QueryService/1.0.0"
     )
     print('Created client for queries: QueryService/1.0.0')
 
     # Client for analytics service
     analytics_client = cosmos_client.CosmosClient(
         HOST,
         {'masterKey': MASTER_KEY},
-        user_agent="AnalyticsService/1.0.0"
+        user_agent_suffix="AnalyticsService/1.0.0"
     )
     print('Created client for analytics: AnalyticsService/1.0.0')
 
@@ -215,7 +215,7 @@ def operation_with_user_agent(user_agent_suffix, operation_name):
         client = cosmos_client.CosmosClient(
                 HOST,
                 {'masterKey': MASTER_KEY},
-                user_agent=user_agent_suffix
+            user_agent_suffix=user_agent_suffix
         )
         db = client.create_database_if_not_exists(id=DATABASE_ID)
         print(f'{operation_name} completed with user agent: {user_agent_suffix}')

sdk/cosmos/azure-cosmos/samples/user_agent_management_async.py

@@ -73,7 +73,7 @@ async def create_client_with_custom_user_agent():
     async with CosmosClient(
         HOST,
         {'masterKey': MASTER_KEY},
-        user_agent=custom_user_agent
+        user_agent_suffix=custom_user_agent
     ) as client:
         print(f'Async client created with custom user agent: {custom_user_agent}')
         print('This will appear as: [default-user-agent] {}'.format(custom_user_agent))
@@ -103,7 +103,7 @@ async def create_client_with_detailed_user_agent():
     async with CosmosClient(
         HOST,
         {'masterKey': MASTER_KEY},
-        user_agent=detailed_user_agent
+        user_agent_suffix=detailed_user_agent
     ) as client:
         print(f'Async client created with detailed user agent: {detailed_user_agent}')
         print('This helps identify specific instances in diagnostics and logs')
@@ -146,7 +146,7 @@ async def demonstrate_user_agent_in_operations():
     async with CosmosClient(
         HOST,
         {'masterKey': MASTER_KEY},
-        user_agent=custom_suffix
+        user_agent_suffix=custom_suffix
     ) as client:
         try:
             # Create database
@@ -199,7 +199,7 @@ async def operation_with_user_agent(user_agent_suffix, operation_name):
         async with CosmosClient(
             HOST,
             {'masterKey': MASTER_KEY},
-            user_agent=user_agent_suffix
+            user_agent_suffix=user_agent_suffix
         ) as client:
             db = await client.create_database_if_not_exists(id=DATABASE_ID)
             print(f'{operation_name} completed with user agent: {user_agent_suffix}')

sdk/cosmos/azure-cosmos/tests/test_user_agent.py

@@ -0,0 +1,100 @@
+# The MIT License (MIT)
+# Copyright (c) Microsoft Corporation. All rights reserved.
+import os
+import unittest
+import uuid
+import re
+from typing import Optional, Any
+
+import test_config
+import pytest
+from azure.cosmos import CosmosClient
+from test_cosmos_http_logging_policy import create_logger
+
+
+@pytest.mark.cosmosEmulator
+class TestUserAgent(unittest.TestCase):
+    """Python User Agent Prefix/Suffix Tests (sync)."""
+
+    client: Optional[CosmosClient] = None
+    host = test_config.TestConfig.host
+    masterKey = test_config.TestConfig.masterKey
+    connectionPolicy = test_config.TestConfig.connectionPolicy
+
+    @classmethod
+    def setUpClass(cls):
+        if (cls.masterKey == '[YOUR_KEY_HERE]' or
+                cls.host == '[YOUR_ENDPOINT_HERE]'):
+            raise Exception(
+                "You must specify your Azure Cosmos account values for "
+                "'masterKey' and 'host' at the top of this class to run the "
+                "tests.")
+
+    def _check(self, user_agent_kw: dict[str, Any]) -> None:
+        self.client = CosmosClient(self.host, self.masterKey, **user_agent_kw)
+        created_database = self.client.get_database_client(test_config.TestConfig.TEST_DATABASE_ID)
+        read_result = created_database.read()
+        assert read_result['id'] == created_database.id
+
+    def test_user_agent_special_characters(self):
+        # Values cover normal, accented, apostrophe, unicode, and spaces / symbols
+        user_agents = ["TestUserAgent", "TéstUserAgent's", "UserAgent with space$%_^()*&"] # cspell:disable-line
+        for ua in user_agents:
+            # Prefix usage
+            self._check({'user_agent': ua})
+            # Suffix usage
+            self._check({'user_agent_suffix': ua})
+
+        # Explicit negative test for disallowed unicode (same logic as previous unicode test)
+        bad = "UnicodeChar鱀InUserAgent"  # expecting potential encode issues when used as prefix
+        try:
+            # Prefix usage
+            self._check({'user_agent': bad})
+            # Suffix usage
+            self._check({'user_agent_suffix': bad})
+            pytest.fail("Unicode characters should not be allowed.")
+        except UnicodeEncodeError as e:
+            assert "ordinal not in range(256)" in e.reason
+
+    def test_user_agent_with_features(self):
+        def _run_case(use_suffix: bool) -> None:
+            os.environ["AZURE_COSMOS_ENABLE_CIRCUIT_BREAKER"] = "True"
+            user_agent_value = "customString"
+            mock_handler = test_config.MockHandler()
+            logger = create_logger("testloggerdefault_sync", mock_handler)
+            try:
+                kwargs: dict[str, Any] = {"logger": logger}
+                if use_suffix:
+                    kwargs["user_agent_suffix"] = user_agent_value
+                else:
+                    kwargs["user_agent"] = user_agent_value
+                self.client = CosmosClient(self.host, self.masterKey, **kwargs)
+                created_database = self.client.get_database_client(test_config.TestConfig.TEST_DATABASE_ID)
+                container = created_database.get_container_client(test_config.TestConfig.TEST_SINGLE_PARTITION_CONTAINER_ID)
+                container.create_item(body={"id": "testItem" + str(uuid.uuid4()), "content": "testContent"})
+                container.create_item(body={"id": "testItem" + str(uuid.uuid4()), "content": "testContent"})
+            finally:
+                os.environ["AZURE_COSMOS_ENABLE_CIRCUIT_BREAKER"] = "False"
+
+            # Examine log messages for User-Agent header
+            for log_message in mock_handler.messages:
+                msg = log_message.message
+                if msg.startswith("Request URL:"):
+                    m = re.search(r"'User-Agent':\s*'([^']+)'", msg, re.MULTILINE)
+                    if m:
+                        ua = m.group(1)
+                        if use_suffix:
+                            # Suffix case: SDK core first then suffix at end
+                            assert ua.startswith("azsdk-python-cosmos")
+                            assert ua.endswith("customString | F2")
+                        else:
+                            # Prefix case: custom part first
+                            assert ua.startswith("customString azsdk-python-cosmos")
+                            assert ua.endswith("| F2")
+
+        _run_case(False)  # prefix scenario
+        _run_case(True)   # suffix scenario
+
+
+if __name__ == '__main__':
+    unittest.main()