Simplify FeatureFlags ref counting by removing WeakMap

packages/react/src/FeatureFlags/FeatureFlags.tsx

@@ -1,17 +1,34 @@
 import type React from 'react'
-import {useContext, useMemo, useEffect, useRef} from 'react'
+import {useContext, useMemo, useEffect} from 'react'
 import {FeatureFlagContext} from './FeatureFlagContext'
 import {FeatureFlagScope, type FeatureFlags} from './FeatureFlagScope'
 
 export type FeatureFlagsProps = React.PropsWithChildren<{
   flags: FeatureFlags
 }>
 
-// WeakMap-based ref counting for data-dialog-scroll-optimized attribute
-// Keys are component instances, values track if the instance has contributed to the count
-const dialogScrollOptimizedInstances = new WeakMap<object, boolean>()
+/**
+ * Ref count for data-dialog-scroll-optimized attribute management.
+ *
+ * NOTE: This is temporary infrastructure while we feature flag the CSS :has()
+ * performance optimization (primer_react_css_has_selector_perf). Once the flag
+ * is removed and the optimization is the default behavior, this ref counting
+ * can be removed - the attribute can simply always be present.
+ *
+ * @internal - Not part of the public API
+ */
 let dialogScrollOptimizedCount = 0
 
+/**
+ * Reset the ref count for testing purposes only.
+ *
+ * @internal - Not part of the public API. Only exported for test isolation.
+ */
+export function __resetDialogScrollOptimizedCount(): void {
+  dialogScrollOptimizedCount = 0
+  document.body.removeAttribute('data-dialog-scroll-optimized')
+}
+
 export function FeatureFlags({children, flags}: FeatureFlagsProps) {
   const parentFeatureFlags = useContext(FeatureFlagContext)
   const value = useMemo(() => {
@@ -20,24 +37,17 @@ export function FeatureFlags({children, flags}: FeatureFlagsProps) {
   }, [parentFeatureFlags, flags])
 
   const isOptimizationEnabled = value.enabled('primer_react_css_has_selector_perf')
-  const instanceRef = useRef({})
 
   // Set body attribute for CSS :has() optimization when flag is enabled
   useEffect(() => {
     if (isOptimizationEnabled) {
-      const instance = instanceRef.current
-      if (!dialogScrollOptimizedInstances.get(instance)) {
-        dialogScrollOptimizedInstances.set(instance, true)
-        dialogScrollOptimizedCount++
-        document.body.setAttribute('data-dialog-scroll-optimized', '')
-      }
+      dialogScrollOptimizedCount++
+      document.body.setAttribute('data-dialog-scroll-optimized', '')
+
       return () => {
-        if (dialogScrollOptimizedInstances.get(instance)) {
-          dialogScrollOptimizedInstances.delete(instance)
-          dialogScrollOptimizedCount--
-          if (dialogScrollOptimizedCount === 0) {
-            document.body.removeAttribute('data-dialog-scroll-optimized')
-          }
+        dialogScrollOptimizedCount--
+        if (dialogScrollOptimizedCount === 0) {
+          document.body.removeAttribute('data-dialog-scroll-optimized')
         }
       }
     }

packages/react/src/FeatureFlags/__tests__/FeatureFlags.test.tsx

@@ -1,11 +1,12 @@
 import {describe, expect, it, beforeEach} from 'vitest'
 import {render} from '@testing-library/react'
 import {FeatureFlags, useFeatureFlag} from '../../FeatureFlags'
+import {__resetDialogScrollOptimizedCount} from '../FeatureFlags'
 
 describe('FeatureFlags', () => {
   beforeEach(() => {
-    // Clean up body attributes between tests
-    document.body.removeAttribute('data-dialog-scroll-optimized')
+    // Reset module state between tests for isolation
+    __resetDialogScrollOptimizedCount()
   })
 
   it('should allow a component to check if a feature flag is enabled', () => {