From c5b0f28b7e52ed7d286648bec8ca8400660acb00 Mon Sep 17 00:00:00 2001
From: Hector Hernandez <39923391+hectorhdzg@users.noreply.github.com>
Date: Fri, 29 May 2026 15:14:55 -0700
Subject: [PATCH] fix: correctly merge statsbeat feature flags from JSON env
 var

The statsbeat feature bitmap stored in AZURE_MONITOR_STATSBEAT_FEATURES is
written as JSON (via JSON.stringify), but the merge logic read it back with
Number(), which always returns NaN for JSON strings. This silently discarded
all previously-set feature flags on every export cycle, causing statsbeat
telemetry to lose feature-usage data for all customers.

The fix parses both formats:
- Plain number strings (backward compat with the shim)
- JSON objects (written by setSdkStatsFeatures and patchOpenTelemetryInstrumentationEnable)

Also fixes test isolation in main.test.ts where the env var leaked between
tests, and adds regression tests for the JSON-format merge path.
---
 src/utils/sdkStats.ts               | 16 +++++++++++---
 test/internal/unit/main.test.ts     |  1 +
 test/internal/unit/sdkStats.test.ts | 33 +++++++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/src/utils/sdkStats.ts b/src/utils/sdkStats.ts
index 0598d95..4dff617 100644
--- a/src/utils/sdkStats.ts
+++ b/src/utils/sdkStats.ts
@@ -132,9 +132,19 @@ class SdkStatsConfiguration {
 
     // Merge old SDK Stats options with new SDK Stats options overriding any common properties
     try {
-      const currentFeaturesBitMap = Number(process.env[AZURE_MONITOR_STATSBEAT_FEATURES]);
-      if (!isNaN(currentFeaturesBitMap)) {
-        featureBitMap |= currentFeaturesBitMap;
+      const envValue = process.env[AZURE_MONITOR_STATSBEAT_FEATURES];
+      if (envValue) {
+        const asNumber = Number(envValue);
+        if (!isNaN(asNumber)) {
+          // Plain number format (e.g. set by the shim)
+          featureBitMap |= asNumber;
+        } else {
+          // JSON format (e.g. set by a previous call to this function)
+          const parsed = JSON.parse(envValue);
+          if (parsed && typeof parsed.feature === "number") {
+            featureBitMap |= parsed.feature;
+          }
+        }
       }
       process.env[AZURE_MONITOR_STATSBEAT_FEATURES] = JSON.stringify({
         instrumentation: instrumentationBitMap,
diff --git a/test/internal/unit/main.test.ts b/test/internal/unit/main.test.ts
index 21b6fa0..f183933 100644
--- a/test/internal/unit/main.test.ts
+++ b/test/internal/unit/main.test.ts
@@ -62,6 +62,7 @@ describe("Main functions", () => {
 
   beforeEach(() => {
     originalEnv = process.env;
+    delete process.env[AZURE_MONITOR_STATSBEAT_FEATURES];
     // Preserve whatever the global OTel API object looks like before each test
     savedOTelGlobal = (globalThis as Record<symbol, unknown>)[GLOBAL_OPENTELEMETRY_API_KEY];
   });
diff --git a/test/internal/unit/sdkStats.test.ts b/test/internal/unit/sdkStats.test.ts
index e2e06fb..07f0a3b 100644
--- a/test/internal/unit/sdkStats.test.ts
+++ b/test/internal/unit/sdkStats.test.ts
@@ -79,4 +79,37 @@ describe("SdkStatsConfiguration — a365 and otlp feature flags", () => {
     expect(features & SdkStatsFeature.A365).toBeTruthy();
     expect(features & SdkStatsFeature.OTLP).toBeTruthy();
   });
+
+  it("should preserve feature bits from JSON-formatted env var", () => {
+    // Seed the env var with JSON format (as written by a previous setSdkStatsFeatures call)
+    process.env[AZURE_MONITOR_STATSBEAT_FEATURES] = JSON.stringify({
+      instrumentation: 0,
+      feature: SdkStatsFeature.AAD_HANDLING | SdkStatsFeature.DISK_RETRY,
+    });
+
+    const sb = getInstance();
+    sb.setSdkStatsFeatures({}, { a365: true });
+
+    const output = JSON.parse(String(process.env[AZURE_MONITOR_STATSBEAT_FEATURES]));
+    const features = Number(output.feature);
+    expect(features & SdkStatsFeature.AAD_HANDLING).toBeTruthy();
+    expect(features & SdkStatsFeature.DISK_RETRY).toBeTruthy();
+    expect(features & SdkStatsFeature.A365).toBeTruthy();
+  });
+
+  it("should not lose feature flags across consecutive setSdkStatsFeatures calls", () => {
+    const sb = getInstance();
+
+    // First call sets OTLP
+    sb.setSdkStatsFeatures({}, { otlp: true });
+    const firstOutput = JSON.parse(String(process.env[AZURE_MONITOR_STATSBEAT_FEATURES]));
+    expect(Number(firstOutput.feature) & SdkStatsFeature.OTLP).toBeTruthy();
+
+    // Second call sets A365 — OTLP should be preserved via env var merge
+    sb.setSdkStatsFeatures({}, { a365: true });
+    const secondOutput = JSON.parse(String(process.env[AZURE_MONITOR_STATSBEAT_FEATURES]));
+    const features = Number(secondOutput.feature);
+    expect(features & SdkStatsFeature.OTLP).toBeTruthy();
+    expect(features & SdkStatsFeature.A365).toBeTruthy();
+  });
 });