Addon-Vitest: Prevent manager "Server timed out" during large test runs

code/addons/vitest/src/constants.ts

@@ -18,6 +18,21 @@ export const DOCUMENTATION_FATAL_ERROR_LINK = `${DOCUMENTATION_LINK}#what-happen
 
 export const COVERAGE_DIRECTORY = 'coverage';
 
+/**
+ * How often (in ms) batched test-case results are flushed to the status stores and universal store
+ * during a run. Flushing serializes status events onto the dev-server WebSocket, so a low value on
+ * large storybooks keeps the channel saturated. It can be raised via
+ * {@link TEST_STATUS_FLUSH_INTERVAL_ENV_VAR} for very large projects.
+ */
+export const DEFAULT_TEST_STATUS_FLUSH_INTERVAL = 500;
+
+/**
+ * Environment variable to override {@link DEFAULT_TEST_STATUS_FLUSH_INTERVAL}. It is read inside the
+ * Vitest child process, which inherits the dev server's environment, so setting it before starting
+ * Storybook is enough.
+ */
+export const TEST_STATUS_FLUSH_INTERVAL_ENV_VAR = 'STORYBOOK_TEST_STATUS_FLUSH_INTERVAL';
+
 export const storeOptions = {
   id: ADDON_ID,
   initialState: {

code/addons/vitest/src/node/test-manager.test.ts

@@ -522,6 +522,29 @@ describe('TestManager', () => {
     expect(mockStore.getState().currentRun.totalTestCount).toBe(1);
   });
 
+  it('should batch test case result flushes according to the configured interval', () => {
+    vi.useFakeTimers();
+    try {
+      const testManager = new TestManager({ ...options, flushTestCaseResultsInterval: 2000 });
+      const passedResult = { state: 'passed', errors: [] } as unknown as TestResult;
+
+      // Leading edge: the first result flushes immediately
+      testManager.onTestCaseResult({ storyId: 'story--one', testResult: passedResult });
+      expect(mockComponentTestStatusStore.set).toHaveBeenCalledTimes(1);
+
+      // Subsequent results within the interval are batched, not flushed yet
+      testManager.onTestCaseResult({ storyId: 'story--two', testResult: passedResult });
+      vi.advanceTimersByTime(1999);
+      expect(mockComponentTestStatusStore.set).toHaveBeenCalledTimes(1);
+
+      // Trailing edge: the batch is flushed once the interval elapses
+      vi.advanceTimersByTime(1);
+      expect(mockComponentTestStatusStore.set).toHaveBeenCalledTimes(2);
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+
   it('should restart Vitest before a test run if coverage is enabled', async () => {
     const testManager = await TestManager.start(options);
     expect(createVitest).toHaveBeenCalledTimes(1);

code/addons/vitest/src/node/test-manager.ts

@@ -10,10 +10,15 @@ import type {
 
 import type { BuilderOptions } from '@storybook/builder-vite';
 
-import { throttle } from 'es-toolkit/function';
+import { type ThrottledFunction, throttle } from 'es-toolkit/function';
 import type { Report } from 'storybook/preview-api';
 
-import { STATUS_TYPE_ID_A11Y, STATUS_TYPE_ID_COMPONENT_TEST, storeOptions } from '../constants.ts';
+import {
+  DEFAULT_TEST_STATUS_FLUSH_INTERVAL,
+  STATUS_TYPE_ID_A11Y,
+  STATUS_TYPE_ID_COMPONENT_TEST,
+  storeOptions,
+} from '../constants.ts';
 import type {
   CurrentRun,
   RunConfig,
@@ -35,6 +40,12 @@ export type TestManagerOptions = {
   testProviderStore: TestProviderStoreById;
   onError?: (message: string, error: Error) => void;
   onReady?: () => void;
+  /**
+   * How often (in ms) batched test-case results are flushed onto the channel during a run. Defaults
+   * to {@link DEFAULT_TEST_STATUS_FLUSH_INTERVAL}. Larger values reduce WebSocket pressure on very
+   * large storybooks at the cost of less frequent sidebar/status updates.
+   */
+  flushTestCaseResultsInterval?: number;
 };
 
 const testStateToStatusValueMap: Record<TestState | 'warning', StatusValue> = {
@@ -77,6 +88,11 @@ export class TestManager {
     this.storybookOptions = options.storybookOptions;
     this.configLoader = options.configLoader;
 
+    this.throttledFlushTestCaseResults = throttle(
+      this.flushTestCaseResults,
+      options.flushTestCaseResultsInterval ?? DEFAULT_TEST_STATUS_FLUSH_INTERVAL
+    );
+
     this.vitestManager = new VitestManager(this);
 
     this.store.subscribe('TRIGGER_RUN', this.handleTriggerRunEvent.bind(this));
@@ -203,7 +219,19 @@ export class TestManager {
   }
 
   /**
-   * Throttled function to process batched test case results.
+   * Throttled wrapper around {@link flushTestCaseResults}, created in the constructor so the interval
+   * can be configured.
+   *
+   * Throttling is necessary because each flush serializes status events onto the dev-server
+   * WebSocket, and without it the channel gets overwhelmed with events. On large storybooks this
+   * backlog can delay heartbeat processing in the manager and cause it to drop the connection with a
+   * "Server timed out" error, so the interval is configurable via
+   * {@link TestManagerOptions.flushTestCaseResultsInterval}.
+   */
+  public throttledFlushTestCaseResults: ThrottledFunction<() => void>;
+
+  /**
+   * Processes batched test case results.
    *
    * This function:
    *
@@ -212,11 +240,8 @@ export class TestManager {
    * 3. Adjusts the totalTestCount if more tests were run than initially anticipated
    * 4. Creates status objects for component tests and updates the component test status store
    * 5. Creates status objects for a11y tests (if any) and updates the a11y status store
-   *
-   * The throttling (500ms) is necessary as the channel would otherwise get overwhelmed with events,
-   * eventually causing the manager and dev server to lose connection.
    */
-  throttledFlushTestCaseResults = throttle(() => {
+  private flushTestCaseResults = () => {
     const testCaseResultsToFlush = this.batchedTestCaseResults;
     this.batchedTestCaseResults = [];
 
@@ -318,7 +343,7 @@ export class TestManager {
         },
       };
     });
-  }, 500);
+  };
 
   onTestRunEnd(endResult: { totalTestCount: number; unhandledErrors: VitestError[] }) {
     this.throttledFlushTestCaseResults.flush();

code/addons/vitest/src/node/vitest.ts

@@ -13,6 +13,7 @@ import {
   ADDON_ID,
   STATUS_TYPE_ID_A11Y,
   STATUS_TYPE_ID_COMPONENT_TEST,
+  TEST_STATUS_FLUSH_INTERVAL_ENV_VAR,
   storeOptions,
 } from '../constants.ts';
 import type { ErrorLike, FatalErrorEvent, StoreEvent, StoreState } from '../types.ts';
@@ -39,6 +40,10 @@ const channel: Channel = new Channel({
 
 const store = UniversalStore.create<StoreState, StoreEvent>(storeOptions);
 
+const parsedFlushInterval = Number(process.env[TEST_STATUS_FLUSH_INTERVAL_ENV_VAR]);
+const flushTestCaseResultsInterval =
+  Number.isFinite(parsedFlushInterval) && parsedFlushInterval > 0 ? parsedFlushInterval : undefined;
+
 new TestManager({
   store,
   componentTestStatusStore: getStatusStore(STATUS_TYPE_ID_COMPONENT_TEST),
@@ -51,6 +56,7 @@ new TestManager({
     configDir: process.env.STORYBOOK_CONFIG_DIR || '',
   } as any,
   configLoader: process.env.STORYBOOK_CONFIG_LOADER as BuilderOptions['configLoader'],
+  flushTestCaseResultsInterval,
 });
 
 const exit = (code = 0) => {

code/core/src/channels/websocket/index.test.ts

@@ -0,0 +1,137 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { parse, stringify } from 'telejson';
+
+import { HEARTBEAT_INTERVAL, HEARTBEAT_MAX_LATENCY, WebsocketTransport } from './index.ts';
+
+const TIMEOUT = HEARTBEAT_INTERVAL + HEARTBEAT_MAX_LATENCY;
+
+const socketRef = { current: undefined as unknown as MockWebSocket };
+
+/**
+ * Minimal WebSocket stub that, unlike the shared one in `channels/index.test.ts`, keeps outgoing
+ * `send` calls separate from incoming messages. This lets us assert heartbeat behavior without a
+ * pong echoing straight back into the message handler.
+ */
+class MockWebSocket {
+  static OPEN = 1;
+
+  onopen?: () => void;
+
+  onmessage?: (event: { data: string }) => void;
+
+  onerror?: (event: unknown) => void;
+
+  onclose?: (event: { code: number; reason: string }) => void;
+
+  readyState = MockWebSocket.OPEN;
+
+  sent: string[] = [];
+
+  closed?: { code: number; reason: string };
+
+  constructor(public url: string) {
+    socketRef.current = this;
+  }
+
+  send(data: string) {
+    this.sent.push(data);
+  }
+
+  close(code: number, reason: string) {
+    this.closed = { code, reason };
+    this.onclose?.({ code, reason });
+  }
+
+  receive(event: unknown) {
+    this.onmessage?.({ data: stringify(event) });
+  }
+}
+
+vi.stubGlobal('WebSocket', MockWebSocket);
+
+const createConnectedTransport = () => {
+  const handler = vi.fn();
+  const onError = vi.fn();
+  const transport = new WebsocketTransport({
+    url: 'ws://localhost:6006',
+    page: 'manager',
+    onError,
+  });
+  transport.setHandler(handler);
+  // onopen marks the transport ready and starts the heartbeat watchdog
+  socketRef.current.onopen?.();
+  return { transport, handler, socket: socketRef.current };
+};
+
+const sentEvents = (socket: MockWebSocket) => socket.sent.map((data) => parse(data));
+
+describe('WebsocketTransport heartbeat', () => {
+  beforeEach(() => {
+    vi.useFakeTimers();
+  });
+
+  afterEach(() => {
+    vi.useRealTimers();
+    vi.clearAllMocks();
+  });
+
+  it('closes with code 3008 when no ping is received within the timeout window', () => {
+    const { socket } = createConnectedTransport();
+
+    vi.advanceTimersByTime(TIMEOUT - 1);
+    expect(socket.closed).toBeUndefined();
+
+    vi.advanceTimersByTime(1);
+    expect(socket.closed).toEqual({ code: 3008, reason: 'timeout' });
+  });
+
+  it('resets the heartbeat and replies with pong when a ping is received', () => {
+    const { socket } = createConnectedTransport();
+
+    vi.advanceTimersByTime(TIMEOUT - 1);
+    socket.receive({ type: 'ping' });
+
+    expect(sentEvents(socket)).toContainEqual({ type: 'pong' });
+
+    // The heartbeat was reset, so another near-full window can elapse without closing
+    vi.advanceTimersByTime(TIMEOUT - 1);
+    expect(socket.closed).toBeUndefined();
+
+    vi.advanceTimersByTime(1);
+    expect(socket.closed).toEqual({ code: 3008, reason: 'timeout' });
+  });
+
+  it('does not forward ping events to the channel handler', () => {
+    const { handler, socket } = createConnectedTransport();
+
+    socket.receive({ type: 'ping' });
+
+    expect(handler).not.toHaveBeenCalled();
+  });
+
+  it('forwards non-ping events to the channel handler', () => {
+    const { handler, socket } = createConnectedTransport();
+
+    socket.receive({ type: 'test', args: [], from: 'preview' });
+
+    expect(handler).toHaveBeenCalledWith(expect.objectContaining({ type: 'test' }));
+  });
+
+  it('resets the heartbeat for pings even when the channel handler throws', () => {
+    const { transport, socket } = createConnectedTransport();
+
+    // Simulate an overloaded/failing channel handler. Heartbeat processing must not depend on it,
+    // otherwise a busy channel could starve pings and close a healthy connection with code 3008.
+    transport.setHandler(() => {
+      throw new Error('handler boom');
+    });
+
+    vi.advanceTimersByTime(TIMEOUT - 1);
+    expect(() => socket.receive({ type: 'ping' })).not.toThrow();
+    expect(sentEvents(socket)).toContainEqual({ type: 'pong' });
+
+    vi.advanceTimersByTime(TIMEOUT - 1);
+    expect(socket.closed).toBeUndefined();
+  });
+});

code/core/src/channels/websocket/index.ts

@@ -50,11 +50,19 @@ export class WebsocketTransport implements ChannelTransport {
     this.socket.onmessage = ({ data }) => {
       const event = typeof data === 'string' && isJSON(data) ? parse(data) : data;
       invariant(this.handler, 'WebsocketTransport handler should be set');
-      this.handler(event);
+
+      // Reset the heartbeat before running the channel handler. The handler can be slow when the
+      // channel is busy (e.g. addon-vitest flushing status updates during a large test run), and a
+      // backlog of queued messages could otherwise delay ping processing past the heartbeat
+      // timeout, closing a healthy connection with code 3008. Pings are internal to the transport
+      // and have no channel listeners, so there is no reason to forward them to the handler.
       if (event.type === 'ping') {
         this.heartbeat();
         this.send({ type: 'pong' });
+        return;
       }
+
+      this.handler(event);
     };
     this.socket.onerror = (e) => {
       if (onError) {

docs/writing-tests/integrations/vitest-addon/index.mdx

@@ -392,6 +392,16 @@ vitest run --shard=3/3
 
 See [Vitest's sharding guide](https://vitest.dev/guide/improving-performance.html#sharding) for more details on how to integrate this into your CI pipeline.
 
+### Why does the Storybook UI show "Server timed out" during a large test run?
+
+On very large projects (thousands of stories), a running test suite streams a high volume of status updates from the dev server to the Storybook manager over a WebSocket. If that channel stays saturated, the manager's heartbeat can be delayed enough to disconnect and show a "Server timed out" notification, even though the dev server is healthy and the tests keep running in the background. Reloading the page reconnects the manager.
+
+To reduce the pressure on the channel, you can increase how often test-case results are batched before being flushed to the UI by setting the `STORYBOOK_TEST_STATUS_FLUSH_INTERVAL` environment variable (in milliseconds) before starting Storybook. The default is `500`; larger values (for example `2000`) reduce WebSocket traffic at the cost of less frequent sidebar and status updates during a run:
+
+```bash
+STORYBOOK_TEST_STATUS_FLUSH_INTERVAL=2000 storybook dev
+```
+
 ## API
 
 ### Exports