Fix partitionLevelCircuitBreakerCfg missing from CosmosDiagnostics clientCfgs

The partitionLevelCircuitBreakerCfg field disappeared from the clientCfgs
section of CosmosDiagnostics when a customer explicitly enabled Per-Partition
Circuit Breaker (PPCB) client-side.

Root cause: the diagnostics write was coupled to the Per-Partition Automatic
Failover (PPAF) initialization path, so the field was only populated when the
service mandated PPAF, not when PPCB was configured client-side.

Fix: move the diagnostics write into initializePerPartitionCircuitBreaker(),
which is invoked unconditionally at client init, so the field appears whenever
the circuit breaker is configured client-side.

Adds a CI-runnable regression test asserting all clientCfgs keys are present,
including partitionLevelCircuitBreakerCfg.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Author: jeet1995

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

@@ -5,6 +5,7 @@
 import com.azure.core.credential.AzureKeyCredential;
 import com.azure.core.http.ProxyOptions;
 import com.azure.cosmos.BridgeInternal;
+import com.azure.cosmos.ConnectionMode;
 import com.azure.cosmos.ConsistencyLevel;
 import com.azure.cosmos.CosmosContainerProactiveInitConfig;
 import com.azure.cosmos.CosmosDiagnostics;
@@ -35,6 +36,11 @@
 import com.azure.cosmos.models.ModelBridgeInternal;
 import com.azure.cosmos.models.PartitionKey;
 import com.azure.cosmos.models.PartitionKeyDefinition;
+import com.fasterxml.jackson.core.JsonFactory;
+import com.fasterxml.jackson.core.JsonGenerator;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.SerializerProvider;
+import com.fasterxml.jackson.databind.node.ObjectNode;
 import io.netty.buffer.ByteBufInputStream;
 import io.netty.buffer.Unpooled;
 import io.netty.handler.codec.http.HttpResponseStatus;
@@ -47,6 +53,8 @@
 import reactor.test.StepVerifier;
 
 import java.net.URI;
+import java.io.StringWriter;
+import java.lang.reflect.Method;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.util.ArrayList;
@@ -324,6 +332,120 @@ public void readMany() {
         }
     }
 
+    // Regression test for the "partitionLevelCircuitBreakerCfg" diagnostics field silently disappearing from the
+    // CosmosDiagnostics "clientCfgs" section. Prior to the fix, the field was only written on the PPAF
+    // (service-mandated) path, so a client that explicitly enabled Per-Partition Circuit Breaker never surfaced it.
+    // This test constructs a real RxDocumentClientImpl (exercising the actual constructor wiring), drives the
+    // private initializePerPartitionCircuitBreaker() init path, serializes the resulting DiagnosticsClientConfig,
+    // and asserts that every expected "clientCfgs" key is present (guarding against future serialization
+    // truncation as well as the specific regression). It also asserts the effective PPCB config string.
+    @Test(groups = {"unit"})
+    public void diagnosticsClientConfigContainsAllClientCfgKeysIncludingPartitionLevelCircuitBreaker() throws Exception {
+        System.setProperty(
+            "COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG",
+            "{\"isPartitionLevelCircuitBreakerEnabled\": true, "
+                + "\"circuitBreakerType\": \"CONSECUTIVE_EXCEPTION_COUNT_BASED\","
+                + "\"consecutiveExceptionCountToleratedForReads\": 10,"
+                + "\"consecutiveExceptionCountToleratedForWrites\": 5,}");
+
+        Mockito.when(this.connectionPolicyMock.getIdleHttpConnectionTimeout()).thenReturn(Duration.ZERO);
+        Mockito.when(this.connectionPolicyMock.getMaxConnectionPoolSize()).thenReturn(1);
+        Mockito.when(this.connectionPolicyMock.getProxy()).thenReturn(null);
+        Mockito.when(this.connectionPolicyMock.getHttpNetworkRequestTimeout()).thenReturn(Duration.ZERO);
+        Mockito.when(this.connectionPolicyMock.getHttp2ConnectionConfig()).thenReturn(new Http2ConnectionConfig());
+        // The serializer eagerly calls getConnectionMode().toString() for the very first "connectionMode" key; if this
+        // returns null (Mockito default), serialization would NPE and silently drop every subsequent key.
+        Mockito.when(this.connectionPolicyMock.getConnectionMode()).thenReturn(ConnectionMode.DIRECT);
+
+        MockedStatic<HttpClient> httpClientMock = Mockito.mockStatic(HttpClient.class);
+        httpClientMock
+            .when(() -> HttpClient.createFixed(Mockito.any(HttpClientConfig.class)))
+            .thenReturn(dummyHttpClient());
+
+        RxDocumentClientImpl rxDocumentClient = null;
+
+        try {
+            rxDocumentClient = new RxDocumentClientImpl(
+                this.serviceEndpointMock,
+                this.masterKeyOrResourceTokenMock,
+                this.permissionFeedMock,
+                this.connectionPolicyMock,
+                this.consistencyLevelMock,
+                null,
+                this.configsMock,
+                this.cosmosAuthorizationTokenResolverMock,
+                this.azureKeyCredentialMock,
+                false,
+                false,
+                false,
+                this.metadataCachesSnapshotMock,
+                this.apiTypeMock,
+                this.cosmosClientTelemetryConfigMock,
+                this.clientCorrelationIdMock,
+                this.endToEndOperationLatencyPolicyConfig,
+                this.sessionRetryOptionsMock,
+                this.containerProactiveInitConfigMock,
+                this.defaultItemSerializer,
+                false
+            );
+
+            // Drive the exact wiring that regressed: explicit (client-side) Per-Partition Circuit Breaker
+            // initialization. The constructor does not invoke init() (which would require network), so invoke the
+            // private no-arg initializer reflectively.
+            Method initPpcb = RxDocumentClientImpl.class.getDeclaredMethod("initializePerPartitionCircuitBreaker");
+            initPpcb.setAccessible(true);
+            initPpcb.invoke(rxDocumentClient);
+
+            ObjectMapper objectMapper = new ObjectMapper();
+            StringWriter jsonWriter = new StringWriter();
+            JsonGenerator jsonGenerator = new JsonFactory().createGenerator(jsonWriter);
+            SerializerProvider serializerProvider = objectMapper.getSerializerProvider();
+            DiagnosticsClientContext.DiagnosticsClientConfigSerializer.INSTANCE
+                .serialize(rxDocumentClient.getConfig(), jsonGenerator, serializerProvider);
+            jsonGenerator.flush();
+            ObjectNode clientCfgs = (ObjectNode) objectMapper.readTree(jsonWriter.toString());
+
+            String serializedJson = clientCfgs.toString();
+
+            // Every key the serializer unconditionally writes, plus the (previously regressed)
+            // partitionLevelCircuitBreakerCfg which is present whenever PPCB is enabled.
+            String[] expectedKeys = new String[] {
+                "id",
+                "machineId",
+                "connectionMode",
+                "numberOfClients",
+                "isPpafEnabled",
+                "isFalseProgSessionTokenMergeEnabled",
+                "excrgns",
+                "clientEndpoints",
+                "connCfg",
+                "consistencyCfg",
+                "proactiveInitCfg",
+                "e2ePolicyCfg",
+                "sessionRetryCfg",
+                "partitionLevelCircuitBreakerCfg"
+            };
+
+            for (String expectedKey : expectedKeys) {
+                assertThat(clientCfgs.has(expectedKey))
+                    .withFailMessage("Expected clientCfgs key '%s' to be present. Serialized clientCfgs: %s",
+                        expectedKey, serializedJson)
+                    .isTrue();
+            }
+
+            assertThat(clientCfgs.get("partitionLevelCircuitBreakerCfg").asText())
+                .withFailMessage("Unexpected partitionLevelCircuitBreakerCfg value. Serialized clientCfgs: %s",
+                    serializedJson)
+                .isEqualTo("(cb: true, type: CONSECUTIVE_EXCEPTION_COUNT_BASED, rexcntt: 10, wexcntt: 5)");
+        } finally {
+            System.clearProperty("COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG");
+            if (rxDocumentClient != null) {
+                rxDocumentClient.close();
+            }
+            httpClientMock.close();
+        }
+    }
+
     private static HttpClient dummyHttpClient() {
         return new HttpClient() {
             @Override

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

@@ -8968,7 +8968,6 @@ private synchronized void initializePerPartitionFailover(DatabaseAccount databas
         checkNotNull(this.globalPartitionEndpointManagerForPerPartitionAutomaticFailover, "Argument 'globalPartitionEndpointManagerForPerPartitionAutomaticFailover' cannot be null.");
         checkNotNull(this.globalPartitionEndpointManagerForPerPartitionCircuitBreaker, "Argument 'globalPartitionEndpointManagerForPerPartitionCircuitBreaker' cannot be null.");
 
-        this.diagnosticsClientConfig.withPartitionLevelCircuitBreakerConfig(this.globalPartitionEndpointManagerForPerPartitionCircuitBreaker.getCircuitBreakerConfig());
         this.diagnosticsClientConfig.withIsPerPartitionAutomaticFailoverEnabled(this.globalPartitionEndpointManagerForPerPartitionAutomaticFailover.isPerPartitionAutomaticFailoverEnabled());
     }
 
@@ -8997,6 +8996,12 @@ private void initializePerPartitionCircuitBreaker() {
 
         this.globalPartitionEndpointManagerForPerPartitionCircuitBreaker.resetCircuitBreakerConfig(partitionLevelCircuitBreakerConfig);
         this.globalPartitionEndpointManagerForPerPartitionCircuitBreaker.init();
+
+        // Populate the circuit breaker config in the diagnostics client config here (rather than in
+        // initializePerPartitionFailover) so the "partitionLevelCircuitBreakerCfg" field appears in
+        // CosmosDiagnostics whenever the circuit breaker is configured client-side, not only when
+        // Per-Partition Automatic Failover is mandated by the service.
+        this.diagnosticsClientConfig.withPartitionLevelCircuitBreakerConfig(this.globalPartitionEndpointManagerForPerPartitionCircuitBreaker.getCircuitBreakerConfig());
     }
 
     private void enableAvailabilityStrategyForReads() {