React: Fix act environment leak causing act(async) without await warnings in multi-file Vitest runs

code/renderers/react/src/act-compat.test.ts

@@ -0,0 +1,181 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { getAct, getReactActEnvironment, setReactActEnvironment } from './act-compat.ts';
+
+// Top-level module mocks. `react` keeps its real exports, but the fallback test
+// can hide `act` (via the hoisted flag, so it's available to the hoisted mock
+// factory) to route getAct onto the deprecated react-dom/test-utils branch.
+const mockState = vi.hoisted(() => ({ reactExposesAct: true }));
+
+vi.mock('react', async (importOriginal) => {
+  const actual = await importOriginal<typeof import('react')>();
+  return {
+    ...actual,
+    get act() {
+      return mockState.reactExposesAct ? actual.act : undefined;
+    },
+  };
+});
+
+vi.mock('react-dom/test-utils', { spy: true });
+
+// Regression coverage for https://github.com/storybookjs/storybook/issues/34708:
+// `IS_REACT_ACT_ENVIRONMENT` must be restored once the act work settles (so it
+// can't leak `true` across story boundaries) and stay set while concurrent acts
+// overlap. A stuck flag makes React log "act(async () => ...) without await".
+describe('act-compat', () => {
+  let errors: string[] = [];
+
+  beforeEach(() => {
+    errors = [];
+    setReactActEnvironment(false);
+    vi.spyOn(console, 'error').mockImplementation((...args: unknown[]) => {
+      errors.push(args.map((arg) => String(arg)).join(' '));
+    });
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+    mockState.reactExposesAct = true;
+  });
+
+  describe('restores the environment once the act settles', () => {
+    it('after an un-awaited async act, without warning', async () => {
+      const act = await getAct();
+
+      // Trigger an async act but never await the returned thenable, mimicking a
+      // pipeline caller that discards the result.
+      void (act(async () => {
+        await Promise.resolve();
+      }) as unknown as Promise<unknown>);
+
+      // The flag is set synchronously while act is in progress.
+      expect(getReactActEnvironment()).toBe(true);
+
+      // Let the act work and React's deferred await-tracking check settle without
+      // the caller ever awaiting the thenable.
+      await new Promise((resolve) => setTimeout(resolve, 0));
+
+      const actWarnings = errors.filter((error) =>
+        error.includes('You called act(async () => ...) without await')
+      );
+
+      expect(getReactActEnvironment()).toBe(false);
+      expect(actWarnings).toEqual([]);
+    });
+
+    it('after an awaited async act resolves', async () => {
+      const act = await getAct();
+
+      await act(async () => {
+        await Promise.resolve();
+      });
+
+      expect(getReactActEnvironment()).toBe(false);
+    });
+
+    it('after an awaited async act rejects', async () => {
+      const act = await getAct();
+
+      await expect(
+        act(async () => {
+          throw new Error('boom');
+        })
+      ).rejects.toThrow('boom');
+
+      expect(getReactActEnvironment()).toBe(false);
+    });
+
+    it('after a synchronous act callback', async () => {
+      const act = await getAct();
+
+      act(() => {
+        expect(getReactActEnvironment()).toBe(true);
+      });
+
+      expect(getReactActEnvironment()).toBe(false);
+    });
+
+    it('after a synchronous act callback throws', async () => {
+      const act = await getAct();
+
+      expect(() =>
+        act(() => {
+          throw new Error('boom');
+        })
+      ).toThrow('boom');
+
+      expect(getReactActEnvironment()).toBe(false);
+    });
+  });
+
+  describe('ref-counts across overlapping acts', () => {
+    it('stays set until every interleaved act has settled', async () => {
+      const act = await getAct();
+
+      let resolveFirst: () => void = () => {};
+      let resolveSecond: () => void = () => {};
+
+      const first = act(() => new Promise<void>((resolve) => (resolveFirst = resolve)));
+      const second = act(() => new Promise<void>((resolve) => (resolveSecond = resolve)));
+
+      expect(getReactActEnvironment()).toBe(true);
+
+      resolveFirst();
+      await first;
+
+      // The second act is still in flight, so the flag must remain set.
+      expect(getReactActEnvironment()).toBe(true);
+
+      resolveSecond();
+      await second;
+
+      expect(getReactActEnvironment()).toBe(false);
+    });
+
+    it('restores the pre-existing value when a nested act settles', async () => {
+      const act = await getAct();
+
+      // An outer act already enabled the environment; the nested act must leave
+      // it enabled rather than hardcode it back to `false`.
+      setReactActEnvironment(true);
+
+      await act(async () => {
+        expect(getReactActEnvironment()).toBe(true);
+        await Promise.resolve();
+      });
+
+      expect(getReactActEnvironment()).toBe(true);
+    });
+  });
+
+  describe('getAct resolution', () => {
+    it('returns a no-op that leaves the environment untouched when act is disabled', async () => {
+      const act = await getAct({ disableAct: true });
+
+      let called = false;
+      act(() => {
+        called = true;
+        expect(getReactActEnvironment()).toBe(false);
+      });
+
+      expect(called).toBe(true);
+      expect(getReactActEnvironment()).toBe(false);
+    });
+
+    it('falls back to react-dom/test-utils when React.act is unavailable', async () => {
+      // Hide `act` on the mocked react so getAct takes the fallback branch, then
+      // re-evaluate act-compat (it snapshots `{ ...React }` at module load).
+      mockState.reactExposesAct = false;
+      vi.resetModules();
+
+      const { getAct: getActFresh } = await import('./act-compat.ts');
+      const testUtils = await import('react-dom/test-utils');
+      const act = await getActFresh();
+
+      act(() => {});
+
+      expect(vi.mocked(testUtils).act).toHaveBeenCalledOnce();
+    });
+  });
+});

code/renderers/react/src/act-compat.ts

@@ -1,4 +1,4 @@
-// Copied from
+// Adapted from
 // https://github.com/testing-library/react-testing-library/blob/3dcd8a9649e25054c0e650d95fca2317b7008576/src/act-compat.js
 import * as React from 'react';
 
@@ -19,10 +19,32 @@ export function getReactActEnvironment() {
   return globalThis.IS_REACT_ACT_ENVIRONMENT;
 }
 
+// Storybook interleaves act calls across stories, so a per-call save/restore
+// (as react-testing-library does) is wrong: the first call to settle would
+// clear the flag while others are still in flight. Ref-count instead, and only
+// restore the captured baseline once the last concurrent act settles.
+// See https://github.com/storybookjs/storybook/issues/34708.
+const actEnvironment = {
+  depth: 0,
+  baseline: false,
+  enter() {
+    if (this.depth === 0) {
+      this.baseline = getReactActEnvironment();
+    }
+    this.depth = this.depth + 1;
+    setReactActEnvironment(true);
+  },
+  exit() {
+    this.depth = this.depth - 1;
+    if (this.depth === 0) {
+      setReactActEnvironment(this.baseline);
+    }
+  },
+};
+
 function withGlobalActEnvironment(actImplementation: (callback: () => void) => Promise<any>) {
   return (callback: () => any) => {
-    const previousActEnvironment = getReactActEnvironment();
-    setReactActEnvironment(true);
+    actEnvironment.enter();
     try {
       // The return value of `act` is always a thenable.
       let callbackNeedsToBeAwaited = false;
@@ -35,28 +57,32 @@ function withGlobalActEnvironment(actImplementation: (callback: () => void) => P
       });
       if (callbackNeedsToBeAwaited) {
         const thenable = actResult;
-        return {
-          then: (resolve: (param: any) => void, reject: (param: any) => void) => {
-            thenable.then(
-              (returnValue: any) => {
-                setReactActEnvironment(previousActEnvironment);
-                resolve(returnValue);
-              },
-              (error: any) => {
-                setReactActEnvironment(previousActEnvironment);
-                reject(error);
-              }
-            );
-          },
-        };
+        // Attach to React's act thenable eagerly (the executor runs
+        // synchronously): some pipeline callers (e.g. the testing-library
+        // `eventWrapper`) discard the result without awaiting, yet the
+        // environment reset must still be tied to the act work settling. We wrap
+        // it in a real Promise because React's act thenable is non-conformant —
+        // its `.then` returns `undefined` rather than a chainable promise.
+        return new Promise((resolve, reject) => {
+          thenable.then(
+            (returnValue: any) => {
+              actEnvironment.exit();
+              resolve(returnValue);
+            },
+            (error: any) => {
+              actEnvironment.exit();
+              reject(error);
+            }
+          );
+        });
       } else {
-        setReactActEnvironment(previousActEnvironment);
+        actEnvironment.exit();
         return actResult;
       }
     } catch (error) {
-      // Can't be a `finally {}` block since we don't know if we have to immediately restore IS_REACT_ACT_ENVIRONMENT
-      // or if we have to await the callback first.
-      setReactActEnvironment(previousActEnvironment);
+      // Not a `finally`: the async branch defers exit() until the act thenable
+      // settles, so finally would call exit() twice. This only covers sync throws.
+      actEnvironment.exit();
       throw error;
     }
   };