Docgen: Run experimentalDocgenServer extraction in a worker thread

code/addons/docs/src/preset.ts

@@ -243,4 +243,3 @@ const optimizeViteDeps = [
 
 export { webpackX as webpack, docsX as docs, optimizeViteDeps };
 export { manifests as experimental_manifests } from './manifest';
-export { experimental_docgenProvider } from './docgen';

code/core/build-config.ts

@@ -45,6 +45,14 @@ const config: BuildEntries = {
         entryPoint: './src/oxc-parser/worker.ts',
         dts: false,
       },
+      {
+        // Long-lived worker that runs docgen extraction off the main thread. Exposed as an internal
+        // export so docgen-worker-client.ts resolves it via the package map (import.meta.resolve)
+        // rather than a hard-coded dist path, keeping strict package managers (pnpm) happy.
+        exportEntries: ['./internal/docgen-worker'],
+        entryPoint: './src/shared/open-service/services/docgen/worker/docgen-worker.ts',
+        dts: false,
+      },
       {
         entryPoint: './src/core-server/presets/common-override-preset.ts',
         exportEntries: ['./internal/core-server/presets/common-override-preset'],

code/core/package.json

@@ -119,6 +119,7 @@
       "code": "./src/csf-tools/index.ts",
       "default": "./dist/csf-tools/index.js"
     },
+    "./internal/docgen-worker": "./dist/shared/open-service/services/docgen/worker/docgen-worker.js",
     "./internal/docs-tools": {
       "types": "./dist/docs-tools/index.d.ts",
       "code": "./src/docs-tools/index.ts",

code/core/src/builder-manager/utils/template.ts

@@ -65,6 +65,9 @@ export const renderHTML = async (
       LOGLEVEL: JSON.stringify(await logLevel, null, 2),
       DOCS_OPTIONS: JSON.stringify(await docsOptions, null, 2),
       CONFIG_TYPE: JSON.stringify(await configType, null, 2),
+      // Opt-in via STORYBOOK_DOCGEN_STORY_PREPARED: request server docgen for the Controls panel at
+      // STORY_PREPARED instead of the default STORY_RENDERED. See useStoryDocgenGateReady.
+      DOCGEN_STORY_PREPARED: JSON.stringify(process.env.STORYBOOK_DOCGEN_STORY_PREPARED === 'true'),
       // These two need to be double stringified because the UI expects a string
       VERSIONCHECK: JSON.stringify(JSON.stringify(versionCheck), null, 2),
       PREVIEW_URL: JSON.stringify(previewUrl, null, 2), // global preview URL

code/core/src/controls/components/ControlsPanel.stories.tsx

@@ -1,11 +1,13 @@
 import React from 'react';
 
+import { STORY_FINISHED, STORY_PREPARED } from 'storybook/internal/core-events';
+
 import type { Meta, StoryObj } from '@storybook/react-vite';
 
 import { global } from '@storybook/global';
 
 import { ManagerContext } from 'storybook/manager-api';
-import { expect, fn, within } from 'storybook/test';
+import { expect, fn, waitFor, within } from 'storybook/test';
 
 import { buildQueryState } from '../../shared/open-service/query-state.ts';
 import type { QueryState } from '../../shared/open-service/types.ts';
@@ -68,6 +70,65 @@ const serviceManagerContext: any = {
   },
 };
 
+// Same story, but not yet prepared. Used by the build-mode story to prove docgen is fetched
+// immediately, without waiting for prepare or render.
+const serviceUnpreparedStoryData = {
+  ...serviceStoryData,
+  prepared: false,
+};
+
+const serviceUnpreparedManagerContext: any = {
+  ...serviceManagerContext,
+  api: {
+    ...serviceManagerContext.api,
+    getCurrentStoryData: fn(() => serviceUnpreparedStoryData).mockName('api::getCurrentStoryData'),
+  },
+};
+
+// A real (tiny) channel backing `api.on`/`off`/`emit`, so a story's `play` can emit story lifecycle
+// events (STORY_FINISHED / STORY_PREPARED) and drive the dev-only docgen gate in `ControlsPanel`.
+function createApiChannel() {
+  const listeners = new Map<string, Set<(...args: any[]) => void>>();
+  return {
+    on: fn((type: string, cb: (...args: any[]) => void) => {
+      const set = listeners.get(type) ?? new Set();
+      set.add(cb);
+      listeners.set(type, set);
+    }).mockName('api::on'),
+    off: fn((type: string, cb: (...args: any[]) => void) => {
+      listeners.get(type)?.delete(cb);
+    }).mockName('api::off'),
+    emit: fn((type: string, ...args: any[]) => {
+      listeners.get(type)?.forEach((cb) => cb(...args));
+    }).mockName('api::emit'),
+  };
+}
+
+const serviceRenderChannel = createApiChannel();
+const serviceRenderManagerContext: any = {
+  ...serviceManagerContext,
+  api: {
+    ...serviceManagerContext.api,
+    ...serviceRenderChannel,
+  },
+};
+
+// A prepared story whose annotation argTypes carry no controls — its controls come entirely from
+// docgen. Used to prove the gate shows the loading skeleton (not the "No controls" empty state) while
+// docgen is still pending.
+const serviceNoControlsStoryData = {
+  ...serviceStoryData,
+  argTypes: {},
+};
+
+const serviceNoControlsManagerContext: any = {
+  ...serviceManagerContext,
+  api: {
+    ...serviceManagerContext.api,
+    getCurrentStoryData: fn(() => serviceNoControlsStoryData).mockName('api::getCurrentStoryData'),
+  },
+};
+
 const serviceGetDocgen = Object.assign(
   fn((_input?: { id: string }) => ({
     id: 'example-button',
@@ -131,6 +192,14 @@ export const RefStoryControlsLoad: Story = {
 
 export const ServiceDocgenControlsLoad: Story = {
   args: { docgenService },
+  // Build mode (no dev render gate), so docgen loads and merges into the table immediately.
+  beforeEach: () => {
+    const original = global.CONFIG_TYPE;
+    global.CONFIG_TYPE = 'PRODUCTION';
+    return () => {
+      global.CONFIG_TYPE = original;
+    };
+  },
   decorators: [
     (storyFn) => (
       <ManagerContext.Provider value={serviceManagerContext}>{storyFn()}</ManagerContext.Provider>
@@ -142,6 +211,180 @@ export const ServiceDocgenControlsLoad: Story = {
   },
 };
 
+/**
+ * In dev, the panel must not query docgen until the story has finished (STORY_FINISHED) — even when
+ * the story is already prepared. Querying earlier would make the docgen worker's CPU-bound extraction
+ * compete with the dev server's bundling and the preview's render during the slowest part of first
+ * render. Until then the panel shows the story's own annotation argTypes (here, the `variant` radio),
+ * or a loading skeleton when the story has none — never the "No controls" empty state.
+ */
+export const ServiceDocgenWaitsForStoryFinishedInDev: Story = {
+  args: { docgenService },
+  beforeEach: () => {
+    serviceGetDocgen.mockClear();
+    serviceGetDocgen.get.mockClear();
+    serviceGetDocgen.subscribe.mockClear();
+    const original = global.CONFIG_TYPE;
+    global.CONFIG_TYPE = 'DEVELOPMENT';
+    return () => {
+      global.CONFIG_TYPE = original;
+    };
+  },
+  decorators: [
+    (storyFn) => (
+      <ManagerContext.Provider value={serviceManagerContext}>{storyFn()}</ManagerContext.Provider>
+    ),
+  ],
+  play: async ({ canvas }) => {
+    // No STORY_FINISHED is emitted, so the gate stays closed and the panel shows the story's own
+    // annotation controls. Wait for that row to render (a stable signal that the gate hook mounted),
+    // then assert synchronously that docgen was never queried — a `waitFor` on the negative passes
+    // before mount settles and would miss a next-tick subscription regression.
+    await expect(await canvas.findByText('variant')).toBeInTheDocument();
+    expect(serviceGetDocgen.subscribe).not.toHaveBeenCalled();
+  },
+};
+
+/**
+ * Regression test: while the dev gate is closed and the story has no annotation controls (its
+ * controls come entirely from docgen), the panel must show the loading skeleton — not the "No
+ * controls" empty state. Previously it flashed that empty state for the whole gate window before
+ * docgen resolved and filled the table.
+ */
+export const ServiceDocgenSkeletonWhileGatedWithoutControls: Story = {
+  args: { docgenService },
+  beforeEach: () => {
+    serviceGetDocgen.mockClear();
+    serviceGetDocgen.get.mockClear();
+    serviceGetDocgen.subscribe.mockClear();
+    const original = global.CONFIG_TYPE;
+    global.CONFIG_TYPE = 'DEVELOPMENT';
+    return () => {
+      global.CONFIG_TYPE = original;
+    };
+  },
+  decorators: [
+    (storyFn) => (
+      <ManagerContext.Provider value={serviceNoControlsManagerContext}>
+        {storyFn()}
+      </ManagerContext.Provider>
+    ),
+  ],
+  play: async ({ canvas }) => {
+    // The empty state debounces for 100ms before rendering, so wait past it, then assert it never
+    // appears while the gate is closed and docgen hasn't been queried.
+    await new Promise((resolve) => setTimeout(resolve, 250));
+    await expect(canvas.queryByText('This story has no controls')).not.toBeInTheDocument();
+    await expect(serviceGetDocgen.subscribe).not.toHaveBeenCalled();
+  },
+};
+
+/**
+ * Once the story finishes in dev (STORY_FINISHED), the gate opens and the panel subscribes to docgen,
+ * merging the extracted argTypes into the table.
+ */
+export const ServiceDocgenLoadsAfterStoryFinishedInDev: Story = {
+  args: { docgenService },
+  beforeEach: () => {
+    serviceGetDocgen.mockClear();
+    serviceGetDocgen.get.mockClear();
+    serviceGetDocgen.subscribe.mockClear();
+    const original = global.CONFIG_TYPE;
+    global.CONFIG_TYPE = 'DEVELOPMENT';
+    return () => {
+      global.CONFIG_TYPE = original;
+    };
+  },
+  decorators: [
+    (storyFn) => (
+      <ManagerContext.Provider value={serviceRenderManagerContext}>
+        {storyFn()}
+      </ManagerContext.Provider>
+    ),
+  ],
+  play: async ({ canvas }) => {
+    // Before the story finishes: the gate is closed, so only annotation controls render. Anchor on
+    // that row, then assert synchronously that docgen wasn't queried (deferred, no subscription).
+    await expect(await canvas.findByText('variant')).toBeInTheDocument();
+    expect(serviceGetDocgen.subscribe).not.toHaveBeenCalled();
+    // Simulate the preview reporting the story as finished.
+    serviceRenderChannel.emit(STORY_FINISHED, { storyId: serviceStoryData.id });
+    await expect(await canvas.findByRole('radio', { name: 'primary' })).toBeInTheDocument();
+    await expect(serviceGetDocgen).toHaveBeenCalledWith({ id: 'example-button' });
+  },
+};
+
+/**
+ * With STORYBOOK_DOCGEN_STORY_PREPARED="true" the dev gate opens earlier, at STORY_PREPARED instead of
+ * STORY_FINISHED, so the panel subscribes to docgen as soon as the story module is delivered.
+ */
+export const ServiceDocgenLoadsAfterStoryPreparedInDev: Story = {
+  args: { docgenService },
+  beforeEach: () => {
+    serviceGetDocgen.mockClear();
+    serviceGetDocgen.get.mockClear();
+    serviceGetDocgen.subscribe.mockClear();
+    const originalConfig = global.CONFIG_TYPE;
+    const originalFlag = global.DOCGEN_STORY_PREPARED;
+    global.CONFIG_TYPE = 'DEVELOPMENT';
+    global.DOCGEN_STORY_PREPARED = true;
+    return () => {
+      global.CONFIG_TYPE = originalConfig;
+      global.DOCGEN_STORY_PREPARED = originalFlag;
+    };
+  },
+  decorators: [
+    (storyFn) => (
+      <ManagerContext.Provider value={serviceRenderManagerContext}>
+        {storyFn()}
+      </ManagerContext.Provider>
+    ),
+  ],
+  play: async ({ canvas }) => {
+    // Before any lifecycle event: the gate is closed, only annotation controls render. Anchor on that
+    // row, then assert synchronously that docgen wasn't queried (deferred, no subscription).
+    await expect(await canvas.findByText('variant')).toBeInTheDocument();
+    expect(serviceGetDocgen.subscribe).not.toHaveBeenCalled();
+    // STORY_FINISHED must not open the gate while waiting for STORY_PREPARED. Give a potential
+    // erroneous subscription time to land, then assert it never happened.
+    serviceRenderChannel.emit(STORY_FINISHED, { storyId: serviceStoryData.id });
+    await new Promise((resolve) => setTimeout(resolve, 100));
+    expect(serviceGetDocgen.subscribe).not.toHaveBeenCalled();
+    // STORY_PREPARED opens it.
+    serviceRenderChannel.emit(STORY_PREPARED, { id: serviceStoryData.id });
+    await expect(await canvas.findByRole('radio', { name: 'primary' })).toBeInTheDocument();
+    await expect(serviceGetDocgen).toHaveBeenCalledWith({ id: 'example-button' });
+  },
+};
+
+/**
+ * In a static build docgen is precomputed JSON with no bundling to contend with, so the panel
+ * subscribes to docgen immediately without waiting for the story to prepare or render.
+ */
+export const ServiceDocgenLoadsImmediatelyInBuild: Story = {
+  args: { docgenService },
+  beforeEach: () => {
+    serviceGetDocgen.mockClear();
+    serviceGetDocgen.get.mockClear();
+    serviceGetDocgen.subscribe.mockClear();
+    const original = global.CONFIG_TYPE;
+    global.CONFIG_TYPE = 'PRODUCTION';
+    return () => {
+      global.CONFIG_TYPE = original;
+    };
+  },
+  decorators: [
+    (storyFn) => (
+      <ManagerContext.Provider value={serviceUnpreparedManagerContext}>
+        {storyFn()}
+      </ManagerContext.Provider>
+    ),
+  ],
+  play: async () => {
+    await waitFor(() => expect(serviceGetDocgen).toHaveBeenCalledWith({ id: 'example-button' }));
+  },
+};
+
 // A panel with more controls than fit, plus one edited arg so the "save from controls" bar
 // renders. The bar is absolutely positioned at the bottom of the panel, so the scrollable
 // content must reserve room for it — otherwise the last control sits under the bar (#34531).