Cosmos: Add user agent tracking for azure-cosmos-encryption SDK

sdk/cosmos/azure-cosmos-encryption/CHANGELOG.md

@@ -3,6 +3,8 @@
 ### 2.28.0-beta.1 (Unreleased)
 
 #### Features Added
+* Added user agent tracking for the encryption SDK. The user agent string now includes `azure-cosmos-encryption/{version}` to enable telemetry tracking of encryption SDK adoption and version distribution. - See [PR 48505](https://github.com/Azure/azure-sdk-for-java/pull/48505)
+* GA'd `deleteAllItemsByPartitionKey` and `queryChangeFeed` APIs in `CosmosEncryptionAsyncContainer` and `CosmosEncryptionContainer`. - See [PR 48505](https://github.com/Azure/azure-sdk-for-java/pull/48505)
 
 #### Breaking Changes
 

sdk/cosmos/azure-cosmos-encryption/src/main/java/com/azure/cosmos/encryption/CosmosEncryptionAsyncClient.java

@@ -10,6 +10,7 @@
 import com.azure.cosmos.CosmosAsyncClientEncryptionKey;
 import com.azure.cosmos.CosmosAsyncContainer;
 import com.azure.cosmos.CosmosAsyncDatabase;
+import com.azure.cosmos.CosmosBridgeInternal;
 import com.azure.cosmos.CosmosException;
 import com.azure.cosmos.CosmosItemSerializer;
 import com.azure.cosmos.encryption.implementation.Constants;
@@ -67,6 +68,7 @@ public final class CosmosEncryptionAsyncClient implements Closeable {
         this.containerPropertiesCacheByContainerId = new AsyncCache<>();
         this.keyEncryptionKeyResolverName = keyEncryptionKeyResolverName;
         this.encryptionKeyStoreProviderImpl = new EncryptionKeyStoreProviderImpl(keyEncryptionKeyResolver, keyEncryptionKeyResolverName);
+        this.appendEncryptionUserAgentSuffix();
     }
 
     /**
@@ -191,6 +193,16 @@ public CosmosAsyncClient getCosmosAsyncClient() {
         return cosmosAsyncClient;
     }
 
+    private void appendEncryptionUserAgentSuffix() {
+        try {
+            CosmosBridgeInternal
+                .getAsyncDocumentClient(this.cosmosAsyncClient)
+                .appendUserAgentSuffix(Constants.USER_AGENT_SUFFIX);
+        } catch (Exception e) {
+            LOGGER.warn("Failed to append encryption SDK user agent suffix", e);
+        }
+    }
+
     /**
      * Gets a database with Encryption capabilities
      *

sdk/cosmos/azure-cosmos-encryption/src/main/java/com/azure/cosmos/encryption/CosmosEncryptionAsyncContainer.java

@@ -324,8 +324,7 @@ public <T> Mono<CosmosItemResponse<Object>> deleteItem(T item, CosmosItemRequest
      * @param requestOptions the request options.
      * @return an {@link Mono} containing the Cosmos item resource response.
      */
-    // TODO Make this api public once it is GA in cosmos core library
-    Mono<CosmosItemResponse<Object>> deleteAllItemsByPartitionKey(PartitionKey partitionKey, CosmosItemRequestOptions requestOptions) {
+    public Mono<CosmosItemResponse<Object>> deleteAllItemsByPartitionKey(PartitionKey partitionKey, CosmosItemRequestOptions requestOptions) {
         final CosmosItemRequestOptions options = Optional.ofNullable(requestOptions)
             .orElse(new CosmosItemRequestOptions());
 
@@ -630,8 +629,7 @@ public <T> CosmosPagedFlux<T> queryItemsOnEncryptedProperties(SqlQuerySpecWithEn
      * @return a {@link CosmosPagedFlux} containing one or several feed response pages of the obtained
      * items or an error.
      */
-    // TODO Make this api public once it is GA in cosmos core library
-    <T> CosmosPagedFlux<T> queryChangeFeed(CosmosChangeFeedRequestOptions options, Class<T> classType) {
+    public <T> CosmosPagedFlux<T> queryChangeFeed(CosmosChangeFeedRequestOptions options, Class<T> classType) {
         checkNotNull(options, "Argument 'options' must not be null.");
         checkNotNull(classType, "Argument 'classType' must not be null.");
 

sdk/cosmos/azure-cosmos-encryption/src/main/java/com/azure/cosmos/encryption/CosmosEncryptionContainer.java

@@ -123,8 +123,7 @@ public <T> CosmosItemResponse<Object> deleteItem(T item, CosmosItemRequestOption
      * @param options the options.
      * @return the Cosmos item response
      */
-    // TODO Make this api public once it is GA in cosmos core library
-    CosmosItemResponse<Object> deleteAllItemsByPartitionKey(PartitionKey partitionKey, CosmosItemRequestOptions options) {
+    public CosmosItemResponse<Object> deleteAllItemsByPartitionKey(PartitionKey partitionKey, CosmosItemRequestOptions options) {
         return this.blockDeleteItemResponse(this.cosmosEncryptionAsyncContainer.deleteAllItemsByPartitionKey(partitionKey, options));
     }
 
@@ -277,8 +276,7 @@ public <T> CosmosPagedIterable<T> queryItemsOnEncryptedProperties(SqlQuerySpecWi
      * @param classType the class type.
      * @return a {@link CosmosPagedFlux} containing one feed response page
      */
-    // TODO Make this api public once it is GA in cosmos core library
-    <T> CosmosPagedIterable<T> queryChangeFeed(
+    public <T> CosmosPagedIterable<T> queryChangeFeed(
         CosmosChangeFeedRequestOptions options,
         Class<T> classType) {
         checkNotNull(options, "Argument 'options' must not be null.");

sdk/cosmos/azure-cosmos-encryption/src/main/java/com/azure/cosmos/encryption/implementation/Constants.java

@@ -3,9 +3,19 @@
 
 package com.azure.cosmos.encryption.implementation;
 
+import com.azure.core.util.CoreUtils;
+
+import java.util.Map;
+
 public class Constants {
     public static final int CACHED_ENCRYPTION_SETTING_DEFAULT_DEFAULT_TTL_IN_MINUTES = 60;
 
+    public static final String PROPERTIES_FILE_NAME = "azure-cosmos-encryption.properties";
+    private static final Map<String, String> PROPERTIES = CoreUtils.getProperties(PROPERTIES_FILE_NAME);
+    public static final String CURRENT_NAME = PROPERTIES.getOrDefault("name", "azure-cosmos-encryption");
+    public static final String CURRENT_VERSION = PROPERTIES.getOrDefault("version", "unknown");
+    public static final String USER_AGENT_SUFFIX = CURRENT_NAME + "/" + CURRENT_VERSION;
+
     public static final String INTENDED_COLLECTION_RID_HEADER = "x-ms-cosmos-intended-collection-rid";
 
     public static final String IS_CLIENT_ENCRYPTED_HEADER = "x-ms-cosmos-is-client-encrypted";

sdk/cosmos/azure-cosmos-encryption/src/main/resources/azure-cosmos-encryption.properties

@@ -0,0 +1,2 @@
+name=${project.artifactId}
+version=${project.version}

sdk/cosmos/azure-cosmos-encryption/src/test/java/com/azure/cosmos/encryption/CosmosEncryptionAsyncClientUnitTest.java

@@ -8,6 +8,9 @@
 import com.azure.cosmos.CosmosAsyncClientEncryptionKey;
 import com.azure.cosmos.CosmosAsyncContainer;
 import com.azure.cosmos.CosmosAsyncDatabase;
+import com.azure.cosmos.CosmosBridgeInternal;
+import com.azure.cosmos.encryption.implementation.Constants;
+import com.azure.cosmos.implementation.AsyncDocumentClient;
 import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
 import com.azure.cosmos.implementation.RequestOptions;
 import com.azure.cosmos.models.CosmosClientEncryptionKeyProperties;
@@ -16,6 +19,8 @@
 import org.testng.annotations.Test;
 import reactor.core.publisher.Mono;
 
+import static org.assertj.core.api.Assertions.assertThat;
+
 public class CosmosEncryptionAsyncClientUnitTest {
 
     private final static ImplementationBridgeHelpers.CosmosAsyncClientEncryptionKeyHelper.CosmosAsyncClientEncryptionKeyAccessor cosmosAsyncClientEncryptionKeyAccessor = ImplementationBridgeHelpers.CosmosAsyncClientEncryptionKeyHelper.getCosmosAsyncClientEncryptionKeyAccessor();
@@ -54,4 +59,56 @@ public void clientEncryptionPropertiesAsync() {
         spyEncryptionAsyncClient.getClientEncryptionPropertiesAsync("testKey", "testDB", mockCosmosAsyncContainer, true, null, true).block();
         Mockito.verify(spyEncryptionAsyncClient, Mockito.times(2)).fetchClientEncryptionKeyPropertiesAsync(Mockito.any(CosmosAsyncContainer.class), Mockito.anyString(), Mockito.any(RequestOptions.class));
     }
+
+    @Test(groups = {"unit"}, timeOut = TestSuiteBase.TIMEOUT)
+    public void encryptionClientAppendsUserAgentSuffix() {
+        // Setup: mock CosmosAsyncClient with a real AsyncDocumentClient to verify UA suffix
+        CosmosAsyncClient mockAsyncClient = Mockito.mock(CosmosAsyncClient.class);
+        AsyncDocumentClient mockDocClient = Mockito.mock(AsyncDocumentClient.class);
+        KeyEncryptionKeyResolver mockKeyResolver = Mockito.mock(KeyEncryptionKeyResolver.class);
+
+        org.mockito.MockedStatic<CosmosBridgeInternal> bridgeMock = Mockito.mockStatic(CosmosBridgeInternal.class);
+        try {
+            bridgeMock.when(() -> CosmosBridgeInternal.getAsyncDocumentClient(mockAsyncClient))
+                .thenReturn(mockDocClient);
+
+            new CosmosEncryptionAsyncClient(mockAsyncClient, mockKeyResolver, "TEST_KEY_RESOLVER");
+
+            // Verify appendUserAgentSuffix was called with the encryption SDK suffix
+            Mockito.verify(mockDocClient, Mockito.times(1))
+                .appendUserAgentSuffix(Constants.USER_AGENT_SUFFIX);
+        } finally {
+            bridgeMock.close();
+        }
+    }
+
+    @Test(groups = {"unit"}, timeOut = TestSuiteBase.TIMEOUT)
+    public void encryptionUserAgentSuffixContainsVersionInfo() {
+        // Verify the suffix constants are properly loaded from properties
+        assertThat(Constants.CURRENT_NAME).isNotEmpty();
+        assertThat(Constants.CURRENT_VERSION).isNotEmpty();
+        assertThat(Constants.USER_AGENT_SUFFIX).isEqualTo(Constants.CURRENT_NAME + "/" + Constants.CURRENT_VERSION);
+        assertThat(Constants.USER_AGENT_SUFFIX).startsWith("azure-cosmos-encryption/");
+    }
+
+    @Test(groups = {"unit"}, timeOut = TestSuiteBase.TIMEOUT)
+    public void encryptionClientHandlesAppendFailureGracefully() {
+        // If getAsyncDocumentClient throws, encryption client should still be created
+        CosmosAsyncClient mockAsyncClient = Mockito.mock(CosmosAsyncClient.class);
+        KeyEncryptionKeyResolver mockKeyResolver = Mockito.mock(KeyEncryptionKeyResolver.class);
+
+        org.mockito.MockedStatic<CosmosBridgeInternal> bridgeMock = Mockito.mockStatic(CosmosBridgeInternal.class);
+        try {
+            bridgeMock.when(() -> CosmosBridgeInternal.getAsyncDocumentClient(mockAsyncClient))
+                .thenThrow(new RuntimeException("simulated failure"));
+
+            // Should not throw — the failure is caught and logged
+            CosmosEncryptionAsyncClient client =
+                new CosmosEncryptionAsyncClient(mockAsyncClient, mockKeyResolver, "TEST_KEY_RESOLVER");
+            assertThat(client).isNotNull();
+            assertThat(client.getCosmosAsyncClient()).isSameAs(mockAsyncClient);
+        } finally {
+            bridgeMock.close();
+        }
+    }
 }

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/UserAgentContainerTest.java

@@ -132,6 +132,46 @@ public void UserAgentIntegration() {
 
     }
 
+    @Test(groups = {"unit"})
+    public void appendUserAgentSuffix() {
+        String expectedStringFixedPart = getUserAgentFixedPart();
+
+        // Append to empty suffix
+        UserAgentContainer userAgentContainer = new UserAgentContainer();
+        String suffix = "azure-cosmos-encryption/2.28.0";
+        userAgentContainer.setSuffix(suffix);
+        String expectedString = expectedStringFixedPart + SPACE + suffix;
+        assertThat(userAgentContainer.getUserAgent()).isEqualTo(expectedString);
+
+        // Append to existing suffix (simulating appendUserAgentSuffix behavior)
+        userAgentContainer = new UserAgentContainer();
+        String customerSuffix = "my-app";
+        userAgentContainer.setSuffix(customerSuffix);
+        String combinedSuffix = customerSuffix + " " + suffix;
+        userAgentContainer.setSuffix(combinedSuffix);
+        expectedString = expectedStringFixedPart + SPACE + combinedSuffix;
+        assertThat(userAgentContainer.getUserAgent()).isEqualTo(expectedString);
+        assertThat(userAgentContainer.getUserAgent()).contains("my-app");
+        assertThat(userAgentContainer.getUserAgent()).contains("azure-cosmos-encryption/2.28.0");
+
+        // Feature flags are preserved after setSuffix + setFeatureEnabledFlagsAsSuffix
+        userAgentContainer = new UserAgentContainer();
+        userAgentContainer.setSuffix(customerSuffix);
+        Set<UserAgentFeatureFlags> flags = new HashSet<>(Arrays.asList(
+            UserAgentFeatureFlags.PerPartitionAutomaticFailover,
+            UserAgentFeatureFlags.PerPartitionCircuitBreaker));
+        userAgentContainer.setFeatureEnabledFlagsAsSuffix(flags);
+        assertThat(userAgentContainer.getUserAgent()).contains("|F3");
+        // After setSuffix, feature flags are cleared
+        userAgentContainer.setSuffix(combinedSuffix);
+        assertThat(userAgentContainer.getUserAgent()).doesNotContain("|F3");
+        // Re-applying feature flags restores them
+        userAgentContainer.setFeatureEnabledFlagsAsSuffix(flags);
+        assertThat(userAgentContainer.getUserAgent()).contains("|F3");
+        assertThat(userAgentContainer.getUserAgent()).contains("my-app");
+        assertThat(userAgentContainer.getUserAgent()).contains("azure-cosmos-encryption/2.28.0");
+    }
+
     private String getUserAgentFixedPart() {
         String osName = System.getProperty("os.name");
         if (osName == null) {

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -13,6 +13,7 @@
 * Fixed availability strategy for Gateway V2 (thin client) by ensuring `RegionalRoutingContext` identity is based only on the immutable gateway endpoint. - See [PR 48432](https://github.com/Azure/azure-sdk-for-java/pull/48432)
 
 #### Other Changes
+* Added `appendUserAgentSuffix` method to `AsyncDocumentClient` to allow downstream libraries to append to the user agent after client construction. - See [PR 48505](https://github.com/Azure/azure-sdk-for-java/pull/48505)
 * Added aggressive HTTP timeout policies for document operations routed to Gateway V2. - [PR 47879](https://github.com/Azure/azure-sdk-for-java/pull/47879)
 * Added a default connect timeout of 5s for Gateway V2 (thin client) data-plane endpoints. - See [PR 48174](https://github.com/Azure/azure-sdk-for-java/pull/48174)
 

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java

@@ -404,6 +404,15 @@ public Builder withOperationPolicies(List<CosmosOperationPolicy> operationPolici
 
     String getUserAgent();
 
+    /**
+     * Appends an additional suffix to the user agent string.
+     *
+     * @param suffix the suffix to append.
+     */
+    default void appendUserAgentSuffix(String suffix) {
+        // no-op default for binary compatibility
+    }
+
     /**
      * Gets the boolean which indicates whether to only return the headers and status code in Cosmos DB response
      * in case of Create, Update and Delete operations on CosmosItem.

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java

@@ -1078,6 +1078,52 @@ public String getUserAgent() {
         return this.userAgentContainer.getUserAgent();
     }
 
+    @Override
+    public void appendUserAgentSuffix(String suffix) {
+        if (StringUtils.isEmpty(suffix)) {
+            return;
+        }
+
+        String trimmedSuffix = suffix.trim();
+        if (trimmedSuffix.isEmpty()) {
+            return;
+        }
+
+        // Check for duplicate using token matching to prevent unbounded growth when
+        // multiple encryption clients wrap the same CosmosAsyncClient
+        String currentSuffix = this.userAgentContainer.getSuffix();
+        if (StringUtils.isNotEmpty(currentSuffix)) {
+            for (String token : currentSuffix.split("\\s+")) {
+                if (trimmedSuffix.equals(token)) {
+                    return;
+                }
+            }
+        }
+
+        // Preserve feature flags ("|F...") which are appended to userAgent directly
+        // by setFeatureEnabledFlagsAsSuffix and would be lost when setSuffix overwrites userAgent
+        String currentUserAgent = this.userAgentContainer.getUserAgent();
+        String featureFlagsSuffix = null;
+        int featureFlagsIndex = currentUserAgent.indexOf("|F");
+        if (featureFlagsIndex >= 0) {
+            featureFlagsSuffix = currentUserAgent.substring(featureFlagsIndex);
+        }
+
+        String newSuffix;
+        if (StringUtils.isNotEmpty(currentSuffix)) {
+            newSuffix = currentSuffix + " " + trimmedSuffix;
+        } else {
+            newSuffix = trimmedSuffix;
+        }
+
+        this.userAgentContainer.setSuffix(newSuffix);
+
+        // Re-apply feature flags since setSuffix overwrites the userAgent string
+        if (StringUtils.isNotEmpty(featureFlagsSuffix)) {
+            this.addUserAgentSuffix(this.userAgentContainer, EnumSet.allOf(UserAgentFeatureFlags.class));
+        }
+    }
+
     @Override
     public CosmosDiagnostics getMostRecentlyCreatedDiagnostics() {
         return mostRecentlyCreatedDiagnostics.get();