Merge pull request storybookjs#35158 from storybookjs/jeppe-cursor/docgen-hot-update-e2e-1f53

JReinhold · web-flow · commit a5177c493e35 · 2026-06-15T13:58:36.000+02:00
Open Service: Fix docgen hot update controls refresh
diff --git a/code/.storybook/main.ts b/code/.storybook/main.ts
@@ -156,9 +156,9 @@ const config = defineMain({
   features: {
     developmentModeForBuild: true,
     experimentalTestSyntax: true,
-    // Disabled for now: the docgen service does not yet work in production builds. Keeping it off
-    // ensures this branch exercises the normal (non-experimental) docgen path without regressions.
-    experimentalDocgenServer: false,
+    // Disabled by default for production builds; the internal dev server enables it via
+    // STORYBOOK_EXPERIMENTAL_DOCGEN_SERVER so hot-update e2e covers the open-service path.
+    experimentalDocgenServer: process.env.STORYBOOK_EXPERIMENTAL_DOCGEN_SERVER === 'true',
     experimentalReactComponentMeta: true,
     changeDetection: true,
   },
diff --git a/code/core/src/shared/open-service/service-runtime.test.ts b/code/core/src/shared/open-service/service-runtime.test.ts
@@ -80,6 +80,34 @@ describe('service runtime', () => {
       unsubscribe();
     });
 
+    it('emits detached snapshots when nested state changes', async () => {
+      const service = registerService(mutableRecordLookupServiceDef);
+      const calls: Array<Record<string, string> | null> = [];
+
+      const unsubscribe = service.queries.getRecordFields.subscribe(
+        { entryId: 'entry-a' },
+        (value) => {
+          calls.push(value);
+        }
+      );
+
+      await vi.waitFor(() => expect(calls).toEqual([null]));
+      await service.commands.assignRecordField({
+        entryId: 'entry-a',
+        fieldKey: 'marker',
+        fieldValue: 'first',
+      });
+      await service.commands.assignRecordField({
+        entryId: 'entry-a',
+        fieldKey: 'marker',
+        fieldValue: 'second',
+      });
+
+      expect(calls).toEqual([null, { marker: 'first' }, { marker: 'second' }]);
+      expect(calls[1]).not.toBe(calls[2]);
+      unsubscribe();
+    });
+
     it('stops notifying after unsubscribe', async () => {
       const service = registerService(mutableRecordLookupServiceDef);
       const calls: Array<Record<string, string> | null> = [];
diff --git a/code/core/src/shared/open-service/service-runtime.ts b/code/core/src/shared/open-service/service-runtime.ts
@@ -1102,7 +1102,7 @@ function subscribeToQuery<TState>(
         selector(output);
         return detachSnapshot(selector(validated));
       }
-      return validateQueryOutput(refs, queryName, queryDef, output);
+      return detachSnapshot(validateQueryOutput(refs, queryName, queryDef, output));
     });
 
     let hasEmitted = false;
diff --git a/code/core/src/shared/open-service/services/docgen/server.test.ts b/code/core/src/shared/open-service/services/docgen/server.test.ts
@@ -2,7 +2,8 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 
 import { Tag } from '../../../../shared/constants/tags.ts';
 import type { DocsIndexEntry, IndexEntry, StoryIndex } from '../../../../types/modules/indexer.ts';
-import { buildStaticFiles, clearRegistry } from '../../server.ts';
+import { buildStaticFiles, clearRegistry, getService } from '../../server.ts';
+import type { ModuleGraphService } from '../module-graph/definition.ts';
 import { registerTestModuleGraphService } from '../module-graph/module-graph.test-helpers.ts';
 import { registerDocgenService } from './server.ts';
 import type { DocgenPayload, DocgenProvider } from './types.ts';
@@ -195,6 +196,39 @@ describe('docgen open service', () => {
     });
   });
 
+  describe('module graph hot refresh', () => {
+    it('refreshes already-extracted components without loading every bumped component', async () => {
+      const buttonEntry = makeStoryEntry('button--primary', 'Button');
+      const cardEntry = makeStoryEntry('card--primary', 'Card');
+      const provider = vi.fn<DocgenProvider>(async ({ entry }) =>
+        makeDocgenPayload({
+          id: entry.id.split('--')[0],
+          name: entry.title,
+          path: entry.importPath,
+        })
+      );
+      const service = registerDocgenService({
+        getIndex: makeGetIndex([buttonEntry, cardEntry]),
+        provider,
+      });
+
+      await service.queries.getDocgen.loaded({ id: 'button' });
+
+      const moduleGraph = getService<ModuleGraphService>('core/module-graph');
+      await moduleGraph.commands._applyGraphUpdate({
+        storiesByFile: {},
+        bumpedStoryFiles: ['./button.stories.tsx', './card.stories.tsx'],
+      });
+
+      await vi.waitFor(() =>
+        expect(provider.mock.calls.map(([input]) => input.entry.importPath)).toEqual([
+          './button.stories.tsx',
+          './button.stories.tsx',
+        ])
+      );
+    });
+  });
+
   describe('static build', () => {
     it('does not request docgen for component ids that only exist on unattached docs entries', async () => {
       const storyEntry = makeStoryEntry('button--primary', 'Button');
diff --git a/code/core/src/shared/open-service/services/docgen/server.ts b/code/core/src/shared/open-service/services/docgen/server.ts
@@ -8,7 +8,7 @@ import { getService, registerService } from '../../server.ts';
 import type { ModuleGraphService } from '../module-graph/definition.ts';
 import { toStoryIndexPath } from '../module-graph/types.ts';
 import { docgenServiceDef } from './definition.ts';
-import type { DocgenPayload, DocgenProvider } from './types.ts';
+import type { DocgenProvider } from './types.ts';
 
 export type RegisterDocgenServiceOptions = {
   workingDir?: string;
@@ -41,6 +41,7 @@ export type RegisterDocgenServiceOptions = {
  */
 export function registerDocgenService(options: RegisterDocgenServiceOptions) {
   const workingDir = options.workingDir ?? process.cwd();
+  const extractedComponentIds = new Set<string>();
 
   const runtime = registerService(docgenServiceDef, {
     queries: {
@@ -77,6 +78,7 @@ export function registerDocgenService(options: RegisterDocgenServiceOptions) {
           ctx.self.setState((state) => {
             state.components[input.id] = payload;
           });
+          extractedComponentIds.add(input.id);
           return payload;
         },
       },
@@ -137,7 +139,7 @@ export function registerDocgenService(options: RegisterDocgenServiceOptions) {
     // Only refresh components that already have extracted docgen in service state, so we never
     // eagerly extract docgen nobody has requested yet.
     const componentIdsToRefresh = Array.from(bumpedComponentIds).filter((id) => {
-      return runtime.queries.getDocgen({ id }) !== undefined;
+      return extractedComponentIds.has(id);
     });
 
     if (componentIdsToRefresh.length === 0) {
diff --git a/code/core/src/shared/open-service/use-service-query.test.tsx b/code/core/src/shared/open-service/use-service-query.test.tsx
@@ -134,27 +134,34 @@ describe('useServiceQuery', () => {
     expect(subscribeSpy.mock.calls.length).toBe(subscribeCallsAfterMount);
   });
 
-  it('maintains referential stability when result is deeply equal', async () => {
-    const service = registerService(mutableRecordLookupServiceDef);
-
-    await service.commands.assignRecordField({ entryId: 'a', fieldKey: 'k', fieldValue: 'v' });
-
+  it('uses emitted snapshots without deep-equal output filtering', async () => {
+    const subscribers: Array<(value: { k: string }) => void> = [];
+    const getRecordFields = Object.assign(
+      vi.fn(() => ({ k: 'v' })),
+      {
+        subscribe: vi.fn(
+          (_input: { entryId: string }, callback: (value: { k: string }) => void) => {
+            subscribers.push(callback);
+            return () => {};
+          }
+        ),
+      }
+    );
+    const service = { queries: { getRecordFields } };
     let renderCount = 0;
     const { result } = renderHook(() => {
       renderCount++;
-      return useServiceQuery(service, 'getRecordFields', { entryId: 'a' });
+      return useServiceQuery(service, 'getRecordFields', { entryId: 'a' }) as { k: string };
     });
 
     const firstRef = result.current;
     const countAfterMount = renderCount;
 
-    // Assign the same value again — deeply equal, so no re-render.
-    await service.commands.assignRecordField({ entryId: 'a', fieldKey: 'k', fieldValue: 'v' });
+    await waitFor(() => expect(subscribers).toHaveLength(1));
+    subscribers[0]({ k: 'v' });
 
-    // Wait a tick to let any spurious re-renders fire.
-    await new Promise<void>((resolve) => setTimeout(resolve, 20));
-
-    expect(result.current).toBe(firstRef);
-    expect(renderCount).toBe(countAfterMount);
+    await waitFor(() => expect(renderCount).toBe(countAfterMount + 1));
+    expect(result.current).toEqual(firstRef);
+    expect(result.current).not.toBe(firstRef);
   });
 });
diff --git a/code/core/src/shared/open-service/use-service-query.ts b/code/core/src/shared/open-service/use-service-query.ts
@@ -4,11 +4,10 @@
  * Backed by `useSyncExternalStore`, so it integrates correctly with React 18+ concurrent
  * features and works in both manager and preview contexts without any adapter.
  *
- * Re-renders only when the specific query result changes by value. Signal-level dedup
- * inside the service runtime ensures that a load which rewrites a deeply-equal payload does
- * not re-fire the subscription; `isEqual` in the subscription callback provides an additional
- * referential-stability layer at the React boundary so the component sees a stable object
- * reference across updates that return the same logical value.
+ * Re-renders only when the subscribed service query emits a new snapshot. Signal-level dedup inside
+ * the service runtime ensures that a load which rewrites a deeply-equal payload does not re-fire the
+ * subscription; this hook intentionally trusts those emissions rather than deep-comparing outputs a
+ * second time.
  *
  * Object inputs are compared with deep equality when deciding whether to re-subscribe, so inline
  * literals at the call site are safe.
@@ -79,9 +78,6 @@ export function useServiceQuery<TKey extends string, TInput, TOutput>(
   const subscribe = React.useCallback(
     (listener: () => void): (() => void) =>
       queryFn.subscribe(subscriptionKey.input, (value) => {
-        if (isEqual(value, snapshotRef.current)) {
-          return;
-        }
         snapshotRef.current = value;
         listener();
       }),
diff --git a/code/e2e-internal/docgen-hot-update.spec.ts b/code/e2e-internal/docgen-hot-update.spec.ts
@@ -0,0 +1,108 @@
+import { readFile, writeFile } from 'node:fs/promises';
+import { join } from 'node:path';
+
+import { expect, test } from '@playwright/test';
+import process from 'process';
+
+import { PREVIEW_STORY_TIMEOUT, waitForPreviewReady } from './helpers.ts';
+
+const storybookUrl = process.env.STORYBOOK_URL || 'http://localhost:6006';
+const runsAgainstDevServer = !['build', 'static'].includes(process.env.STORYBOOK_TYPE || 'dev');
+
+const codeDir = process.cwd();
+const buttonSourcePath = join(codeDir, 'core/src/components/components/Button/Button.tsx');
+const storyPath = '/story/button-component--base';
+const hotUpdatePropName = 'e2eDocgenHotUpdateProp';
+const hotUpdatePropSource = `
+  /**
+   * E2E-only docgen hot update marker.
+   */
+  ${hotUpdatePropName}?: 'before' | 'after';
+`;
+
+// Start the internal dev server with STORYBOOK_EXPERIMENTAL_DOCGEN_SERVER=true before running this
+// spec. CI sets that env var in the internal Storybook e2e job.
+let originalButtonSource: string | undefined;
+
+async function restoreFile(path: string, contents: string) {
+  if ((await readFile(path, 'utf8')) !== contents) {
+    await writeFile(path, contents, 'utf8');
+  }
+}
+
+async function addHotUpdateProp() {
+  const current = await readFile(buttonSourcePath, 'utf8');
+  if (current.includes(hotUpdatePropName)) {
+    return;
+  }
+
+  const marker = '  shortcut?: API_KeyCollection;\n';
+  expect(current, `Could not find ButtonProps insertion marker in ${buttonSourcePath}`).toContain(
+    marker
+  );
+
+  await writeFile(buttonSourcePath, current.replace(marker, `${marker}${hotUpdatePropSource}`));
+}
+
+test.describe('docgen open service hot updates', () => {
+  test.describe.configure({ mode: 'serial' });
+  test.setTimeout(90_000);
+
+  test.beforeAll(async () => {
+    test.skip(!runsAgainstDevServer, 'Docgen hot updates require the dev server file watcher.');
+
+    originalButtonSource = await readFile(buttonSourcePath, 'utf8');
+  });
+
+  test.afterAll(async () => {
+    if (originalButtonSource) {
+      await restoreFile(buttonSourcePath, originalButtonSource);
+    }
+  });
+
+  test('updates manager Controls when a component prop type changes without navigation', async ({
+    page,
+  }) => {
+    await restoreFile(buttonSourcePath, originalButtonSource!);
+
+    await page.goto(`${storybookUrl}/?path=${storyPath}`);
+    await expect
+      .poll(
+        () =>
+          page.evaluate(() =>
+            Boolean(
+              (
+                globalThis as {
+                  FEATURES?: { experimentalDocgenServer?: boolean };
+                }
+              ).FEATURES?.experimentalDocgenServer
+            )
+          ),
+        { timeout: PREVIEW_STORY_TIMEOUT }
+      )
+      .toBe(true);
+    await waitForPreviewReady(page);
+    await expect(
+      page.frameLocator('#storybook-preview-iframe').getByRole('button', { name: 'Button' })
+    ).toBeVisible({ timeout: PREVIEW_STORY_TIMEOUT });
+
+    await page.getByRole('tab', { name: 'Controls' }).click();
+    const controlsPanel = page.getByRole('tabpanel', { name: 'Controls' });
+    await expect(
+      controlsPanel.getByRole('cell', { name: 'ariaLabel', exact: true }).first()
+    ).toBeVisible({
+      timeout: PREVIEW_STORY_TIMEOUT,
+    });
+    await expect(controlsPanel.getByText(hotUpdatePropName)).toHaveCount(0);
+
+    try {
+      await addHotUpdateProp();
+
+      await expect(
+        controlsPanel.getByRole('cell', { name: hotUpdatePropName, exact: true }).first()
+      ).toBeVisible({ timeout: PREVIEW_STORY_TIMEOUT });
+    } finally {
+      await restoreFile(buttonSourcePath, originalButtonSource!);
+    }
+  });
+});
diff --git a/scripts/ci/common-jobs.ts b/scripts/ci/common-jobs.ts
@@ -176,7 +176,7 @@ export const internalStorybookE2e = defineJob(
           name: 'Run internal Storybook',
           working_directory: 'code',
           background: true,
-          command: 'yarn storybook:ui',
+          command: 'STORYBOOK_EXPERIMENTAL_DOCGEN_SERVER=true yarn storybook:ui',
         },
       },
       server.wait(['6006']),

Original file line number	Diff line number	Diff line change
`@@ -1102,7 +1102,7 @@ function subscribeToQuery<TState>(`
`1102`	`1102`	`selector(output);`
`1103`	`1103`	`return detachSnapshot(selector(validated));`
`1104`	`1104`	`}`
`1105`		`- return validateQueryOutput(refs, queryName, queryDef, output);`
	`1105`	`+ return detachSnapshot(validateQueryOutput(refs, queryName, queryDef, output));`
`1106`	`1106`	`});`
`1107`	`1107`
`1108`	`1108`	`let hasEmitted = false;`