diff --git a/code/addons/docs/src/docgen.ts b/code/addons/docs/src/docgen.ts deleted file mode 100644 index 4a21de036141..000000000000 --- a/code/addons/docs/src/docgen.ts +++ /dev/null @@ -1,23 +0,0 @@ -import type { DocgenProviderPreset } from 'storybook/internal/types'; - -/** - * Addon-docs docgen provider. - * - * A small enrichment layer: appends a `(docs enabled)` marker to the downstream description so - * consumers can tell that addon-docs participated. Does NOT produce docgen on its own — when no - * downstream provider supplied a payload it returns undefined so the chain falls through. - */ -export const experimental_docgenProvider: DocgenProviderPreset = async (nextDocgen) => { - return async (input) => { - const downstream = await nextDocgen(input); - if (!downstream) { - return undefined; - } - return { - ...downstream, - description: downstream.description - ? `${downstream.description} (docs enabled)` - : 'docs enabled', - }; - }; -}; diff --git a/code/addons/docs/src/preset.ts b/code/addons/docs/src/preset.ts index 5789a41393d6..d8f6d6a5521a 100644 --- a/code/addons/docs/src/preset.ts +++ b/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'; diff --git a/code/core/build-config.ts b/code/core/build-config.ts index 4f31aad36c01..5bbf398ad9de 100644 --- a/code/core/build-config.ts +++ b/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'], diff --git a/code/core/package.json b/code/core/package.json index 313849066b29..e8a767434562 100644 --- a/code/core/package.json +++ b/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", diff --git a/code/core/src/builder-manager/utils/template.ts b/code/core/src/builder-manager/utils/template.ts index 9fa806bd34a7..f4a23dcfdd60 100644 --- a/code/core/src/builder-manager/utils/template.ts +++ b/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 diff --git a/code/core/src/controls/components/ControlsPanel.stories.tsx b/code/core/src/controls/components/ControlsPanel.stories.tsx index 50f64c4b0c6b..79a9755a52d3 100644 --- a/code/core/src/controls/components/ControlsPanel.stories.tsx +++ b/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 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) => ( {storyFn()} @@ -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) => ( + {storyFn()} + ), + ], + 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) => ( + + {storyFn()} + + ), + ], + 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) => ( + + {storyFn()} + + ), + ], + 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) => ( + + {storyFn()} + + ), + ], + 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) => ( + + {storyFn()} + + ), + ], + 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). diff --git a/code/core/src/controls/components/ControlsPanel.tsx b/code/core/src/controls/components/ControlsPanel.tsx index 1779ad499800..b37124132fc4 100644 --- a/code/core/src/controls/components/ControlsPanel.tsx +++ b/code/core/src/controls/components/ControlsPanel.tsx @@ -1,6 +1,7 @@ import React, { useEffect, useMemo, useState } from 'react'; -import type { ArgTypes } from 'storybook/internal/types'; +import { STORY_FINISHED, STORY_PREPARED } from 'storybook/internal/core-events'; +import type { ArgTypes, StoryId } from 'storybook/internal/types'; import { global } from '@storybook/global'; @@ -9,6 +10,7 @@ import { mergeServiceArgTypes } from '../../docs-tools/argTypes/docgenServiceArg import { useArgTypes, useArgs, + useChannel, useGlobals, useParameter, useServiceQuery, @@ -33,6 +35,11 @@ const clean = (obj: { [key: string]: any }) => {} as typeof obj ); +// True when at least one row carries an interactive control, i.e. the args table renders rows instead +// of its "No controls" empty state. Used to decide between showing controls and the loading skeleton +// while docgen for the story is still pending, so the empty state never flashes before docgen lands. +const hasAnyControl = (rows: ArgTypes) => Object.values(rows).some((arg) => arg?.control); + const AddonWrapper = styled.div<{ showSaveFromUI: boolean }>(({ showSaveFromUI, theme }) => ({ // `minHeight` (not `height`) so the wrapper fills a short panel but grows with tall // content; otherwise a fixed height clips the bottom padding into the middle of the @@ -92,7 +99,7 @@ function ControlsPanelTable({ const api = useStorybookApi(); const storyData = api.getCurrentStoryData(); - const hasControls = Object.values(rows).some((arg) => arg?.control); + const hasControls = hasAnyControl(rows); const withPresetColors = applyPresetColors(rows, presetColors); const hasUpdatedArgs = useMemo( @@ -150,22 +157,25 @@ function LegacyControlsPanel(props: ControlsPanelProps) { return ; } -function ServiceControlsPanel({ +function LoadedServiceControlsPanel({ + customArgTypes, docgenService, + id, + initialArgs, + isStoryPrepared, + storyData, ...props -}: ControlsPanelProps & { docgenService: DocgenService }) { - const api = useStorybookApi(); - const storyData = api.getCurrentStoryData(); - const [, , , initialArgs] = useArgs(); - // Custom argTypes (project + meta + story annotations) for the selected story arrive over the - // channel via STORY_PREPARED. With experimentalDocgenServer, prepareStory skips second-pass - // enhancers so these stay annotation-only; mergeServiceArgTypes layers them on server docgen. - const customArgTypes = useArgTypes(); - const id = storyData.id.split('--')[0]; +}: ControlsPanelProps & { + customArgTypes: ArgTypes; + docgenService: DocgenService; + id: string; + initialArgs: ReturnType[3]; + isStoryPrepared: boolean; + storyData: ReturnType['getCurrentStoryData']>; +}) { const { data: docgenPayload, isInitialLoading } = useServiceQuery(docgenService.queries.docgen, { id, }); - const isStoryPrepared = storyData.type === 'story' ? storyData.prepared : true; // The manager Controls panel only ever shows the main component's rows; subcomponent tabs are a // docs-blocks-only feature, so this intentionally ignores `payload.subcomponents` to match the @@ -184,8 +194,100 @@ function ServiceControlsPanel({ [docgenPayload, initialArgs, storyData.id, storyData.parameters, customArgTypes] ); + // Keep the skeleton up only while there is genuinely nothing to show: the story isn't prepared, or + // docgen is still doing its first load and there are no annotation controls to fall back on. Once + // docgen resolves (even to nothing) the table or its "No controls" empty state is the real answer — + // never a flash. While docgen loads over already-available annotation controls, show those controls + // rather than skeletoning over them. + return ( + + ); +} + +// Tracks whether it is safe to query docgen for the selected story, gating the CPU-bound worker +// extraction out of the first-render window so it never starves the dev server's bundling or the +// preview's render/paint. The gate resets on every story change. +// +// In a static build there is no bundling to contend with and docgen is precomputed JSON, so the gate +// is open from the start. In dev it opens once the story reaches a safe point in its lifecycle: +// STORY_FINISHED by default, or — with the STORYBOOK_DOCGEN_STORY_PREPARED env var — the earlier +// STORY_PREPARED (the story module has been delivered to the browser, before it mounts and paints), +// trading a slice of the first-render budget for an earlier Controls population. The server reads the +// env var and injects it into the manager as the DOCGEN_STORY_PREPARED global so it can be evaluated +// per project. +function useStoryDocgenGateReady(storyId: StoryId): boolean { + const isDevelopment = global.CONFIG_TYPE === 'DEVELOPMENT'; + const requestAtStoryPrepared = global.DOCGEN_STORY_PREPARED === true; + // Key readiness by the story it was opened for (derived during render, not via an effect). On a + // story switch this is immediately false for the new id, so we never render one frame with a stale + // `ready=true` that would mount the query before the new story emits its gate event. + const [readyStoryId, setReadyStoryId] = useState(null); + const ready = !isDevelopment || readyStoryId === storyId; + + // STORY_PREPARED carries `{ id }`; STORY_FINISHED carries `{ storyId }`. + const gateEvent = requestAtStoryPrepared ? STORY_PREPARED : STORY_FINISHED; + useChannel( + { + [gateEvent]: (payload: { id?: StoryId; storyId?: StoryId }) => { + const eventStoryId = requestAtStoryPrepared ? payload?.id : payload?.storyId; + if (eventStoryId === storyId) { + setReadyStoryId(storyId); + } + }, + }, + [storyId, requestAtStoryPrepared, gateEvent] + ); + + return ready; +} + +function ServiceControlsPanel({ + docgenService, + ...props +}: ControlsPanelProps & { docgenService: DocgenService }) { + const api = useStorybookApi(); + const storyData = api.getCurrentStoryData(); + const [, , , initialArgs] = useArgs(); + // Custom argTypes (project + meta + story annotations) for the selected story arrive over the + // channel via STORY_PREPARED. With experimentalDocgenServer, prepareStory skips second-pass + // enhancers so these stay annotation-only; mergeServiceArgTypes layers them on server docgen. + const customArgTypes = useArgTypes(); + const id = storyData.id.split('--')[0]; + const isStory = storyData.type === 'story'; + // In a static build docgen is precomputed and there is no gate, so the loading decision must not + // hang on `prepared` — treat non-development as always prepared. In dev this is unchanged. + const isDevelopment = global.CONFIG_TYPE === 'DEVELOPMENT'; + const isStoryPrepared = !isDevelopment || (isStory ? storyData.prepared : true); + // Docs entries don't emit the story lifecycle events the gate listens for, so it only applies to + // actual stories; everything else falls through and queries docgen right away. + const gateReady = useStoryDocgenGateReady(storyData.id); + if (isStory && !gateReady) { + // Docgen hasn't been queried yet (the query is gated until the story reaches a safe lifecycle + // point). Show the story's own annotation argTypes when it has controls; otherwise keep the + // loading skeleton rather than flashing the "No controls" empty state before docgen resolves. + return ( + + ); + } + return ( - + ); } diff --git a/code/core/src/controls/typings.d.ts b/code/core/src/controls/typings.d.ts index 4a32763646ea..370581ae71e9 100644 --- a/code/core/src/controls/typings.d.ts +++ b/code/core/src/controls/typings.d.ts @@ -1,5 +1,6 @@ declare var DOCS_OPTIONS: any; declare var CONFIG_TYPE: 'DEVELOPMENT' | 'PRODUCTION'; +declare var DOCGEN_STORY_PREPARED: boolean; declare var PREVIEW_URL: any; declare var __STORYBOOK_ADDONS_MANAGER: any; diff --git a/code/core/src/core-server/presets/common-preset.ts b/code/core/src/core-server/presets/common-preset.ts index cc06fdbfa9b6..c5d7ac4d9ec6 100644 --- a/code/core/src/core-server/presets/common-preset.ts +++ b/code/core/src/core-server/presets/common-preset.ts @@ -21,7 +21,7 @@ import { logger } from 'storybook/internal/node-logger'; import { telemetry } from 'storybook/internal/telemetry'; import type { CoreConfig, - DocgenProvider, + DocgenProviderDescriptor, Indexer, Options, PresetProperty, @@ -31,6 +31,7 @@ import type { } from 'storybook/internal/types'; import { registerDocgenService } from '../../shared/open-service/services/docgen/server.ts'; +import { createDocgenWorkerClient } from '../../shared/open-service/services/docgen/worker/docgen-worker-client.ts'; import { registerModuleGraphService } from '../../shared/open-service/services/module-graph/server.ts'; import { registerStoryDocsService } from '../../shared/open-service/services/story-docs/server.ts'; @@ -359,19 +360,30 @@ export const services = async (_value: void, options: Options): Promise => // produce docgen files that wouldn't be served anywhere). Mirrors the !options.ignorePreview // gate around index.json and writeManifests in build-static.ts. if (features?.experimentalDocgenServer && !options.ignorePreview) { - const [docgenProvider, storyDocsProvider] = await Promise.all([ - options.presets.apply('experimental_docgenProvider', async () => undefined), + const [docgenDescriptors, storyDocsProvider] = await Promise.all([ + options.presets.apply('experimental_docgenProvider', []), options.presets.apply( 'experimental_storyDocsProvider', async () => undefined ), ]); - registerDocgenService({ - getIndex: () => storyIndexGenerator.getIndex(), - docgenProvider, - workingDir: process.cwd(), - }); + // Docgen extraction runs in a long-lived worker so its CPU-bound TypeScript work never starves + // the dev-server event loop. The worker composes the descriptor chain; here we forward one + // component to it. When the compiled worker script is unavailable (e.g. running from source + // without a build) the client is undefined and we skip docgen registration rather than + // extracting on the main thread. + const docgenWorker = + docgenDescriptors.length > 0 ? createDocgenWorkerClient(docgenDescriptors) : undefined; + + if (docgenWorker) { + registerDocgenService({ + getIndex: () => storyIndexGenerator.getIndex(), + docgenProvider: (input) => docgenWorker.extract(input.entry), + workingDir: process.cwd(), + }); + } + registerStoryDocsService({ getIndex: () => storyIndexGenerator.getIndex(), storyDocsProvider, diff --git a/code/core/src/shared/open-service/services/docgen/types.ts b/code/core/src/shared/open-service/services/docgen/types.ts index 26159d89825d..36067272055a 100644 --- a/code/core/src/shared/open-service/services/docgen/types.ts +++ b/code/core/src/shared/open-service/services/docgen/types.ts @@ -1,5 +1,4 @@ import type { StrictArgTypes } from '../../../../types/modules/csf.ts'; -import type { Options } from '../../../../types/modules/core-common.ts'; import type { IndexEntry } from '../../../../types/modules/indexer.ts'; /** @@ -59,11 +58,11 @@ export interface DocgenSubcomponent { } /** - * Middleware-style provider function registered through the `experimental_docgenProvider` preset. + * A docgen provider: given a component's index entry, returns a complete {@link DocgenPayload} or + * `undefined` when no docgen is available for the file. * - * Each registrant returns a wrapper around the previous accumulated provider; it may call that - * inner provider to merge with downstream output, and either returns a complete - * {@link DocgenPayload} or `undefined` when no docgen is available for the given file. + * Providers are composed middleware-style inside the docgen worker — each wraps the previous one in + * the chain and may delegate to it to merge with downstream output. * * **Merge convention.** When combining your output with downstream's, use spread * (`{ ...downstream, ...yourOverrides }`) and `downstream?.field ?? yours` rather than rebuilding @@ -75,16 +74,33 @@ export interface DocgenSubcomponent { export type DocgenProvider = (input: DocgenProviderInput) => Promise; /** - * Preset signature for `experimental_docgenProvider`. + * Middleware that wraps the next provider in the docgen chain. * - * Like `PresetPropertyFn<'experimental_docgenProvider'>` but with `nextDocgen` typed as - * non-nullable. Core's `services` preset always seeds the middleware chain with an identity - * provider, so the optional typing inherited from `StorybookConfigRaw` is impossible-state - * defense at the provider-author level — use this type to drop the `?.` noise. If the seed is - * ever missing at runtime, that's a preset-wiring bug and the provider will throw on the first - * `nextDocgen(...)` call rather than silently degrading. + * The worker seeds the chain with an identity provider (returns `undefined`) and folds each + * descriptor's middleware over it in registration order, so a later registrant wraps (and can + * delegate to) the earlier ones. */ -export type DocgenProviderPreset = ( - nextDocgen: DocgenProvider, - options: Options -) => DocgenProvider | Promise; +export type DocgenMiddleware = (nextDocgen: DocgenProvider) => DocgenProvider; + +/** + * Serializable descriptor a renderer or addon contributes via the `experimental_docgenProvider` + * preset. + * + * Docgen extraction runs off the main thread in a long-lived worker owned by core. Because a + * closure cannot cross a worker boundary, integrations describe their provider as data: a + * `moduleSpecifier` pointing at a module that satisfies {@link DocgenWorkerModule}. Core collects + * these descriptors (preserving preset order) and the worker imports and composes them. + */ +export interface DocgenProviderDescriptor { + /** Absolute path to a module that exports {@link DocgenWorkerModule.createDocgenProvider}. */ + moduleSpecifier: string; +} + +/** + * Contract a worker-target docgen module must satisfy. The worker imports the descriptor's + * `moduleSpecifier` and calls `createDocgenProvider()` once to build the middleware it folds into + * the provider chain. Integrations implement only this factory — they never touch threading. + */ +export interface DocgenWorkerModule { + createDocgenProvider: () => DocgenMiddleware | Promise; +} diff --git a/code/core/src/shared/open-service/services/docgen/worker/docgen-worker-client.test.ts b/code/core/src/shared/open-service/services/docgen/worker/docgen-worker-client.test.ts new file mode 100644 index 000000000000..b3ce685e9999 --- /dev/null +++ b/code/core/src/shared/open-service/services/docgen/worker/docgen-worker-client.test.ts @@ -0,0 +1,197 @@ +// Unit tests for the docgen worker client's dispatch / readiness / failure plumbing. A real worker is +// replaced with an EventEmitter stand-in so the tests stay hermetic. +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; + +import type { DocgenProviderDescriptor } from '../types.ts'; +import type { DocgenWorkerRequest, DocgenWorkerResponse } from './protocol.ts'; + +/** Only the surface the tests touch; the runtime instance is a real Node EventEmitter subclass. */ +interface FakeWorker { + posted: DocgenWorkerRequest[]; + postMessage: ReturnType; + terminate: ReturnType; + unref: ReturnType; + emit: (event: string, ...args: unknown[]) => boolean; +} + +const fakeWorkers: FakeWorker[] = []; + +vi.mock('node:worker_threads', async () => { + const { EventEmitter: NodeEventEmitter } = await import('node:events'); + + class FakeWorkerImpl extends NodeEventEmitter { + posted: DocgenWorkerRequest[] = []; + constructor() { + super(); + fakeWorkers.push(this as unknown as FakeWorker); + } + postMessage = vi.fn((msg: DocgenWorkerRequest) => { + this.posted.push(msg); + }); + terminate = vi.fn(async () => 0); + unref = vi.fn(); + } + + return { Worker: FakeWorkerImpl }; +}); + +vi.mock('node:fs', () => ({ existsSync: vi.fn(() => true) })); + +vi.mock('../../../../utils/module.ts', () => ({ + // Include a drive letter so `fileURLToPath()` accepts it on Windows too (a driveless file URL + // throws ERR_INVALID_FILE_URL_PATH there, which would make the client resolve to undefined). + importMetaResolve: vi.fn(() => 'file:///C:/fake/storybook/docgen-worker.js'), +})); + +vi.mock('storybook/internal/node-logger', { spy: true }); + +const loadModule = async () => import('./docgen-worker-client.ts'); + +const DESCRIPTORS: DocgenProviderDescriptor[] = [ + { moduleSpecifier: '/fake/react/docgen-worker.js' }, +]; + +/** Resolve the worker's `init` ack so awaiting `extract` calls can proceed. */ +function ackInit(worker: FakeWorker, error?: { name: string; message: string }) { + worker.emit( + 'message', + (error ? { type: 'init', error } : { type: 'init' }) satisfies DocgenWorkerResponse + ); +} + +beforeEach(() => { + fakeWorkers.length = 0; +}); + +afterEach(() => { + vi.clearAllMocks(); + fakeWorkers.length = 0; +}); + +describe('createDocgenWorkerClient', () => { + it('returns undefined when the compiled worker script is missing', async () => { + const fs = await import('node:fs'); + vi.mocked(fs.existsSync).mockReturnValueOnce(false); + + const { createDocgenWorkerClient } = await loadModule(); + expect(createDocgenWorkerClient(DESCRIPTORS)).toBeUndefined(); + expect(fakeWorkers).toHaveLength(0); + }); + + it('does not spawn a worker until the first extract', async () => { + const { createDocgenWorkerClient } = await loadModule(); + const client = createDocgenWorkerClient(DESCRIPTORS); + + expect(client).toBeDefined(); + expect(fakeWorkers).toHaveLength(0); + }); + + it('spawns the worker lazily on the first extract and posts init with descriptors', async () => { + const { createDocgenWorkerClient } = await loadModule(); + const client = createDocgenWorkerClient(DESCRIPTORS)!; + + const promise = client.extract({ id: 'button--primary' } as any); + const worker = fakeWorkers[0]; + + expect(fakeWorkers).toHaveLength(1); + expect(worker.posted[0]).toEqual({ type: 'init', descriptors: DESCRIPTORS }); + expect(worker.unref).toHaveBeenCalled(); + + // Drive the extract to completion so dispose isn't racing a not-yet-queued request. + ackInit(worker); + await Promise.resolve(); + const extractMsg = worker.posted.find((m) => m.type === 'extract') as { id: number }; + worker.emit('message', { + type: 'extract', + id: extractMsg.id, + payload: { id: 'button', name: 'Button', path: './button.stories.tsx', jsDocTags: {} }, + } satisfies DocgenWorkerResponse); + await expect(promise).resolves.toMatchObject({ id: 'button' }); + }); +}); + +describe('DocgenWorkerClient.extract', () => { + it('waits for init, dispatches the entry, and resolves the payload', async () => { + const { createDocgenWorkerClient } = await loadModule(); + const client = createDocgenWorkerClient(DESCRIPTORS)!; + + const entry = { id: 'button--primary', importPath: './button.stories.tsx' } as any; + const promise = client.extract(entry); + const worker = fakeWorkers[0]; + + // Nothing is dispatched until init is acked. + expect(worker.posted.filter((m) => m.type === 'extract')).toHaveLength(0); + ackInit(worker); + await Promise.resolve(); + + const extractMsg = worker.posted.find((m) => m.type === 'extract'); + expect(extractMsg).toMatchObject({ type: 'extract', entry }); + + worker.emit('message', { + type: 'extract', + id: (extractMsg as { id: number }).id, + payload: { id: 'button', name: 'Button', path: './button.stories.tsx', jsDocTags: {} }, + } satisfies DocgenWorkerResponse); + + await expect(promise).resolves.toMatchObject({ id: 'button', name: 'Button' }); + }); + + it('rejects extract calls when init fails', async () => { + const { createDocgenWorkerClient } = await loadModule(); + const client = createDocgenWorkerClient(DESCRIPTORS)!; + + const promise = client.extract({ id: 'x' } as any); + const worker = fakeWorkers[0]; + ackInit(worker, { name: 'Error', message: 'boom' }); + + await expect(promise).rejects.toThrow('boom'); + }); + + it('rejects an extract awaiting init when the worker exits before init', async () => { + const { createDocgenWorkerClient } = await loadModule(); + const client = createDocgenWorkerClient(DESCRIPTORS)!; + + const promise = client.extract({ id: 'x' } as any); + const worker = fakeWorkers[0]; + // Worker dies during boot, before its `init` ack — `ready` must reject so the awaiting extract + // fails fast instead of hanging forever. + worker.emit('exit', 1); + + await expect(promise).rejects.toThrow(/exited unexpectedly/); + }); + + it('rejects with the error name/message from a failed extraction', async () => { + const { createDocgenWorkerClient } = await loadModule(); + const client = createDocgenWorkerClient(DESCRIPTORS)!; + + const promise = client.extract({ id: 'x' } as any); + const worker = fakeWorkers[0]; + ackInit(worker); + await Promise.resolve(); + + const extractMsg = worker.posted.find((m) => m.type === 'extract') as { id: number }; + worker.emit('message', { + type: 'extract', + id: extractMsg.id, + error: { name: 'DocgenError', message: 'extraction exploded' }, + } satisfies DocgenWorkerResponse); + + await expect(promise).rejects.toThrowError( + expect.objectContaining({ name: 'DocgenError', message: 'extraction exploded' }) + ); + }); + + it('rejects in-flight extractions when the worker exits unexpectedly', async () => { + const { createDocgenWorkerClient } = await loadModule(); + const client = createDocgenWorkerClient(DESCRIPTORS)!; + + const promise = client.extract({ id: 'x' } as any); + const worker = fakeWorkers[0]; + ackInit(worker); + await Promise.resolve(); + + worker.emit('exit', 1); + + await expect(promise).rejects.toThrow(/exited unexpectedly/); + }); +}); diff --git a/code/core/src/shared/open-service/services/docgen/worker/docgen-worker-client.ts b/code/core/src/shared/open-service/services/docgen/worker/docgen-worker-client.ts new file mode 100644 index 000000000000..9b95e4cd9f66 --- /dev/null +++ b/code/core/src/shared/open-service/services/docgen/worker/docgen-worker-client.ts @@ -0,0 +1,206 @@ +/** + * Main-thread client for the long-lived docgen worker. + * + * Owns a single worker (docgen extraction serializes on one warm TypeScript program, so a pool + * would only duplicate multi-second program builds and memory). Spawned once per process when the + * compiled worker script is present; when it is missing — e.g. running from source without a build — + * {@link createDocgenWorkerClient} returns `undefined` and the caller skips docgen registration + * rather than silently extracting on the main thread. + */ +import { existsSync } from 'node:fs'; +import { fileURLToPath } from 'node:url'; +import { Worker } from 'node:worker_threads'; + +import { logger } from 'storybook/internal/node-logger'; + +import type { IndexEntry } from '../../../../../types/modules/indexer.ts'; +import { importMetaResolve } from '../../../../utils/module.ts'; +import type { ErrorLike } from '../../module-graph/types.ts'; +import type { DocgenPayload, DocgenProviderDescriptor } from '../types.ts'; +import type { DocgenWorkerRequest, DocgenWorkerResponse } from './protocol.ts'; + +/** + * Package-relative specifier for the compiled worker. Resolved via the package export map (not a + * hard-coded dist path) so strict package managers like pnpm resolve it correctly. + */ +const WORKER_SPECIFIER = 'storybook/internal/docgen-worker'; + +const DEFAULT_TASK_TIMEOUT_MS = 120_000; + +interface Pending { + resolve: (payload: DocgenPayload | undefined) => void; + reject: (error: unknown) => void; + timer?: NodeJS.Timeout; +} + +export interface DocgenWorkerClient { + /** Extracts docgen for one component entry off the main thread. */ + extract(entry: IndexEntry): Promise; +} + +/** Rebuild an Error from a worker {@link ErrorLike} so the original name/message/stack survive. */ +function errorLikeToError(errorLike: ErrorLike): Error { + const error = new Error(errorLike.message); + if (errorLike.name) { + error.name = errorLike.name; + } + if (errorLike.stack) { + error.stack = errorLike.stack; + } + return error; +} + +class DocgenWorker implements DocgenWorkerClient { + private readonly worker: Worker; + private readonly pending = new Map(); + private readonly ready: Promise; + // Captured from the `ready` promise executor (which runs synchronously) so `fail()` can settle + // `ready` if the worker dies before sending `init`. No-op default keeps TS definite-assignment happy. + private rejectReady: (error: unknown) => void = () => undefined; + private nextId = 0; + private dead = false; + + constructor( + scriptPath: string, + descriptors: DocgenProviderDescriptor[], + private readonly taskTimeoutMs = DEFAULT_TASK_TIMEOUT_MS + ) { + this.worker = new Worker(scriptPath); + // Never let an idle worker keep the process alive. + this.worker.unref(); + this.worker.on('message', (msg: DocgenWorkerResponse) => this.handleMessage(msg)); + this.worker.on('error', (error) => this.fail(error)); + this.worker.on('exit', (code) => { + if (!this.dead) { + this.fail(new Error(`docgen worker exited unexpectedly with code ${code}`)); + } + }); + + this.ready = new Promise((resolve, reject) => { + this.rejectReady = reject; + const onMessage = (msg: DocgenWorkerResponse) => { + if (msg.type !== 'init') { + return; + } + this.worker.off('message', onMessage); + if (msg.error) { + reject(errorLikeToError(msg.error)); + } else { + resolve(); + } + }; + this.worker.on('message', onMessage); + }); + // Surface late init rejections instead of leaving an unhandled rejection. + this.ready.catch(() => undefined); + + this.post({ type: 'init', descriptors }); + } + + async extract(entry: IndexEntry): Promise { + if (this.dead) { + throw new Error('docgen worker is no longer running'); + } + await this.ready; + return new Promise((resolve, reject) => { + const id = this.nextId++; + const pending: Pending = { resolve, reject }; + if (this.taskTimeoutMs > 0) { + pending.timer = setTimeout(() => { + this.pending.delete(id); + reject(new Error(`docgen worker extract ${id} timed out after ${this.taskTimeoutMs}ms`)); + }, this.taskTimeoutMs); + pending.timer.unref?.(); + } + this.pending.set(id, pending); + this.post({ type: 'extract', id, entry }); + }); + } + + private post(msg: DocgenWorkerRequest): void { + this.worker.postMessage(msg); + } + + private handleMessage(msg: DocgenWorkerResponse): void { + if (msg.type !== 'extract') { + return; + } + const pending = this.pending.get(msg.id); + if (!pending) { + return; + } + if (pending.timer) { + clearTimeout(pending.timer); + } + this.pending.delete(msg.id); + if (msg.error) { + pending.reject(errorLikeToError(msg.error)); + } else { + pending.resolve(msg.payload); + } + } + + /** Reject everything in flight and tear the worker down; used on fatal worker failure. */ + private fail(error: Error): void { + if (this.dead) { + return; + } + logger.debug(`docgen worker failure: ${error.message}`); + this.dead = true; + // If the worker dies before `init`, `ready` would otherwise never settle and an in-flight + // `extract` awaiting it would hang. Rejecting is a no-op once `ready` has already settled. + this.rejectReady(error); + this.rejectAllPending(error); + this.worker.terminate().catch(() => 0); + } + + private rejectAllPending(error: Error): void { + for (const [, pending] of this.pending) { + if (pending.timer) { + clearTimeout(pending.timer); + } + pending.reject(error); + } + this.pending.clear(); + } +} + +function resolveWorkerScriptPath(): string | undefined { + try { + const scriptPath = fileURLToPath(importMetaResolve(WORKER_SPECIFIER)); + return existsSync(scriptPath) ? scriptPath : undefined; + } catch { + return undefined; + } +} + +/** + * Creates the docgen worker client for `descriptors`. Returns `undefined` when the compiled worker + * script is unavailable (no fallback — the caller skips docgen registration). The worker lives for the + * process lifetime and is torn down with it, so there is nothing to dispose explicitly. + * + * The worker thread is spawned lazily on the first {@link DocgenWorkerClient.extract}, not here: + * spawning the thread and importing the provider module graph is CPU work, and extraction is gated + * until a story renders, so there is nothing to keep warm at dev-server boot — spawning eagerly would + * only contend with Vite's cold start. + */ +export function createDocgenWorkerClient( + descriptors: DocgenProviderDescriptor[] +): DocgenWorkerClient | undefined { + const scriptPath = resolveWorkerScriptPath(); + if (!scriptPath) { + logger.debug( + 'docgen worker disabled: compiled worker script not found (running from source without a build?)' + ); + return undefined; + } + + let worker: DocgenWorker | undefined; + + return { + async extract(entry) { + worker ??= new DocgenWorker(scriptPath, descriptors); + return worker.extract(entry); + }, + }; +} diff --git a/code/core/src/shared/open-service/services/docgen/worker/docgen-worker.ts b/code/core/src/shared/open-service/services/docgen/worker/docgen-worker.ts new file mode 100644 index 000000000000..72b60776ba0b --- /dev/null +++ b/code/core/src/shared/open-service/services/docgen/worker/docgen-worker.ts @@ -0,0 +1,98 @@ +/** + * Worker-thread entry point for docgen extraction. + * + * Spawned once by {@link ./docgen-worker-client.ts} and kept alive for the dev server's lifetime so + * the TypeScript program(s) each provider builds stay warm across requests. On `init` it imports + * every {@link DocgenProviderDescriptor}'s module and folds them into a single provider chain; on + * `extract` it runs that chain for one component and posts the resulting payload back. Running here + * keeps the synchronous, CPU-bound TypeScript work off the main event loop so it never starves the + * Vite dev server during first render. + */ +import { parentPort } from 'node:worker_threads'; +import { pathToFileURL } from 'node:url'; + +import type { IndexEntry } from '../../../../../types/modules/indexer.ts'; +import { errorToErrorLike } from '../../module-graph/types.ts'; +import type { + DocgenMiddleware, + DocgenProvider, + DocgenProviderDescriptor, + DocgenWorkerModule, +} from '../types.ts'; +import type { DocgenWorkerRequest, DocgenWorkerResponse } from './protocol.ts'; + +if (!parentPort) { + throw new Error('docgen worker must be run as a worker thread'); +} + +// Capture into a const so TypeScript keeps the non-null narrowing inside the async handlers and the +// message listener below (it won't narrow the mutable `parentPort` binding across those closures). +const port = parentPort; + +/** Identity provider that seeds the chain; the bottom of the stack has no docgen to contribute. */ +const seedProvider: DocgenProvider = async () => undefined; + +/** Resolved once per worker on `init`; every `extract` awaits it. */ +let providerPromise: Promise | undefined; + +async function composeProvider(descriptors: DocgenProviderDescriptor[]): Promise { + let provider = seedProvider; + for (const descriptor of descriptors) { + const moduleUrl = pathToFileURL(descriptor.moduleSpecifier).href; + const mod = (await import(moduleUrl)) as Partial; + if (typeof mod.createDocgenProvider !== 'function') { + throw new Error( + `docgen worker module "${descriptor.moduleSpecifier}" does not export createDocgenProvider` + ); + } + const middleware: DocgenMiddleware = await mod.createDocgenProvider(); + provider = middleware(provider); + } + return provider; +} + +async function handleInit(descriptors: DocgenProviderDescriptor[]): Promise { + try { + providerPromise = composeProvider(descriptors); + await providerPromise; + port.postMessage({ type: 'init' } satisfies DocgenWorkerResponse); + } catch (error) { + providerPromise = undefined; + port.postMessage({ + type: 'init', + error: errorToErrorLike(error), + } satisfies DocgenWorkerResponse); + } +} + +async function handleExtract(id: number, entry: IndexEntry): Promise { + try { + if (!providerPromise) { + throw new Error('docgen worker received an extract request before init'); + } + const provider = await providerPromise; + const payload = await provider({ entry }); + port.postMessage({ type: 'extract', id, payload } satisfies DocgenWorkerResponse); + } catch (error) { + port.postMessage({ + type: 'extract', + id, + error: errorToErrorLike(error), + } satisfies DocgenWorkerResponse); + } +} + +port.on('message', (msg: DocgenWorkerRequest) => { + switch (msg.type) { + case 'init': + void handleInit(msg.descriptors); + return; + case 'extract': + void handleExtract(msg.id, msg.entry); + return; + default: { + const _exhaustive: never = msg; + throw new Error(`docgen worker received unknown message: ${JSON.stringify(_exhaustive)}`); + } + } +}); diff --git a/code/core/src/shared/open-service/services/docgen/worker/protocol.ts b/code/core/src/shared/open-service/services/docgen/worker/protocol.ts new file mode 100644 index 000000000000..91d7a077c56f --- /dev/null +++ b/code/core/src/shared/open-service/services/docgen/worker/protocol.ts @@ -0,0 +1,60 @@ +/** + * Message protocol between the docgen worker client (main thread) and the worker + * ({@link ./docgen-worker.ts}). All payloads are structured-clone-safe so they can cross the + * `postMessage` boundary. + * + * Responses are success/failure unions discriminated by the presence of an {@link ErrorLike} + * `error` (mirroring `ModuleGraphStatus`), rather than an `ok` flag with nullable fields — a + * response either carries its result or an error, never an ambiguous mix. + */ +import type { IndexEntry } from '../../../../../types/modules/indexer.ts'; +import type { ErrorLike } from '../../module-graph/types.ts'; +import type { DocgenPayload, DocgenProviderDescriptor } from '../types.ts'; + +/** Sent once, right after spawn, to compose the provider chain inside the worker. */ +export interface DocgenWorkerInitRequest { + type: 'init'; + descriptors: DocgenProviderDescriptor[]; +} + +/** Sent per component to extract its docgen payload. `id` correlates the response. */ +export interface DocgenWorkerExtractRequest { + type: 'extract'; + id: number; + entry: IndexEntry; +} + +export type DocgenWorkerRequest = DocgenWorkerInitRequest | DocgenWorkerExtractRequest; + +/** Chain composition succeeded; the worker is ready for extract requests. */ +export interface DocgenWorkerInitSuccess { + type: 'init'; + error?: undefined; +} + +/** Chain composition failed; the worker is unusable. */ +export interface DocgenWorkerInitFailure { + type: 'init'; + error: ErrorLike; +} + +export type DocgenWorkerInitResponse = DocgenWorkerInitSuccess | DocgenWorkerInitFailure; + +/** Extraction succeeded; `payload` is the result (absent when no provider produced docgen). */ +export interface DocgenWorkerExtractSuccess { + type: 'extract'; + id: number; + payload?: DocgenPayload; + error?: undefined; +} + +/** Extraction threw; `error` describes the failure. */ +export interface DocgenWorkerExtractFailure { + type: 'extract'; + id: number; + error: ErrorLike; +} + +export type DocgenWorkerExtractResponse = DocgenWorkerExtractSuccess | DocgenWorkerExtractFailure; + +export type DocgenWorkerResponse = DocgenWorkerInitResponse | DocgenWorkerExtractResponse; diff --git a/code/core/src/types/modules/core-common.ts b/code/core/src/types/modules/core-common.ts index e37ee57bfe64..6b1f7b5a2193 100644 --- a/code/core/src/types/modules/core-common.ts +++ b/code/core/src/types/modules/core-common.ts @@ -10,7 +10,7 @@ import type { Server as NetServer } from 'net'; import type { Options as TelejsonOptions } from 'telejson'; import type { PackageJson as PackageJsonFromTypeFest } from 'type-fest'; -import type { DocgenProvider } from '../../shared/open-service/services/docgen/types.ts'; +import type { DocgenProviderDescriptor } from '../../shared/open-service/services/docgen/types.ts'; import type { StoryDocsProvider } from '../../shared/open-service/services/story-docs/types.ts'; import type { SupportedBuilder } from './builders.ts'; import type { SupportedFramework } from './frameworks.ts'; @@ -20,11 +20,13 @@ import type { SupportedRenderer } from './renderers.ts'; export type { DocgenError, DocgenJsDocTags, + DocgenMiddleware, DocgenPayload, DocgenProvider, + DocgenProviderDescriptor, DocgenProviderInput, - DocgenProviderPreset, DocgenSubcomponent, + DocgenWorkerModule, } from '../../shared/open-service/services/docgen/types.ts'; export type { StoryDoc, @@ -137,9 +139,9 @@ export interface Presets { apply(extension: 'services', config?: StorybookConfigRaw['services'], args?: any): Promise; apply( extension: 'experimental_docgenProvider', - config: DocgenProvider, + config: DocgenProviderDescriptor[], args?: any - ): Promise; + ): Promise; apply( extension: 'experimental_storyDocsProvider', config: StoryDocsProvider, @@ -469,7 +471,7 @@ export interface StorybookConfigRaw { core?: CoreConfig; experimental_manifests?: Manifests; experimental_enrichCsf?: CsfEnricher; - experimental_docgenProvider?: DocgenProvider; + experimental_docgenProvider?: DocgenProviderDescriptor[]; experimental_storyDocsProvider?: StoryDocsProvider; staticDirs?: (DirectoryMapping | string)[]; logLevel?: string; @@ -803,9 +805,9 @@ export interface StorybookConfig { services?: PresetValue; /** - * Middleware-style provider for the experimental docgen service. Each registrant receives the - * previously accumulated provider as its config argument and returns a wrapping provider that - * may delegate to it via the input forwarding pattern. + * Provider descriptors for the experimental docgen service. Each registrant appends a + * structured-clone-safe {@link DocgenProviderDescriptor} (a module specifier) to the accumulated + * array; core's docgen worker imports and composes them middleware-style off the main thread. */ experimental_docgenProvider?: PresetValue; diff --git a/code/renderers/react/build-config.ts b/code/renderers/react/build-config.ts index b14582ea919a..47b92afe5d61 100644 --- a/code/renderers/react/build-config.ts +++ b/code/renderers/react/build-config.ts @@ -42,6 +42,14 @@ const config: BuildEntries = { entryPoint: './src/preset.ts', dts: false, }, + { + // Worker-target docgen module imported by core's docgen worker. Exposed as an internal + // export so the preset resolves it via import.meta.resolve (package map) instead of a + // hard-coded dist path. + exportEntries: ['./internal/docgen-worker'], + entryPoint: './src/docgen/docgen-worker.ts', + dts: false, + }, ], }, }; diff --git a/code/renderers/react/package.json b/code/renderers/react/package.json index 407b25436843..42d2d06a679e 100644 --- a/code/renderers/react/package.json +++ b/code/renderers/react/package.json @@ -35,6 +35,7 @@ "code": "./src/playwright.ts", "default": "./dist/playwright.js" }, + "./internal/docgen-worker": "./dist/docgen/docgen-worker.js", "./package.json": "./package.json", "./preset": "./dist/preset.js", "./preview": { diff --git a/code/renderers/react/src/componentManifest/componentMetaManagerSingleton.ts b/code/renderers/react/src/componentManifest/componentMetaManagerSingleton.ts index b8c570cdcdce..fcd5c5b826ab 100644 --- a/code/renderers/react/src/componentManifest/componentMetaManagerSingleton.ts +++ b/code/renderers/react/src/componentManifest/componentMetaManagerSingleton.ts @@ -5,13 +5,16 @@ import { ComponentMetaManager } from './componentMeta/ComponentMetaManager.ts'; let componentMetaManagerPromise: Promise | undefined; /** - * Process-wide {@link ComponentMetaManager} shared by the experimental manifest generator and the - * docgen provider so static builds do not construct duplicate TypeScript language services. + * Process-wide {@link ComponentMetaManager} for the experimental manifest generator: building a full + * TypeScript program over every tsconfig project in the workspace is expensive, so the manifest + * generator keeps a single set of programs (and one file-snapshot cache) hot for the lifetime of the + * process rather than rebuilding per request. * - * Both features build full TypeScript programs over every tsconfig project in the workspace. When - * each owned its own manager, enabling both at once kept two complete sets of programs resident - * simultaneously and exhausted the heap. Sharing one manager keeps a single set of programs (and - * one file-snapshot cache) hot for the lifetime of the process. + * `experimentalDocgenServer` deliberately does NOT share this manager: its extraction runs in a + * worker thread ({@link ../docgen/docgen-worker.ts}) with its own manager so the synchronous program + * build stays off the main event loop. Worker threads have separate V8 heaps but share process RSS, + * so enabling both the manifest generator and the docgen server keeps two full program sets resident + * at once — acceptable today, but the reason this singleton is scoped to the manifest path only. */ export function getSharedComponentMetaManager(): Promise { if (!componentMetaManagerPromise) { diff --git a/code/renderers/react/src/docgen/docgen-worker.ts b/code/renderers/react/src/docgen/docgen-worker.ts new file mode 100644 index 000000000000..2e793f2af380 --- /dev/null +++ b/code/renderers/react/src/docgen/docgen-worker.ts @@ -0,0 +1,65 @@ +/** + * Worker-target docgen module for the React renderer. + * + * Core's docgen worker imports this module and calls {@link createDocgenProvider} once to build the + * middleware it folds into the provider chain. Everything here runs inside that worker thread: the + * heavy `react-component-meta` extraction (a full TypeScript program build via + * {@link ComponentMetaManager}) happens off the main event loop. The React renderer implements no + * threading itself — it only describes this module to core through the `experimental_docgenProvider` + * preset. + * + * This provider needs no configuration: `react-component-meta` extraction reads nothing from the + * user's `typescript` options. `ComponentMetaManager` builds its own program, and the only + * docgen-relevant TS option (`reactDocgenTypescriptOptions`, including `propFilter`) applies solely + * to the `react-docgen-typescript` engine used by the legacy manifest path — not here. + */ +import { STORY_FILE_TEST_REGEXP, getStoryImportPathFromEntry } from 'storybook/internal/common'; +import type { DocgenMiddleware, DocgenProvider } from 'storybook/internal/types'; + +import { ComponentMetaManager } from '../componentManifest/componentMeta/ComponentMetaManager.ts'; +import { buildDocgenPayload } from './buildDocgen.ts'; + +/** + * Builds the React docgen middleware. Owns one {@link ComponentMetaManager} for the worker's + * lifetime so the TypeScript language service / program stays warm across components. The manager is + * created lazily on the first eligible request and memoized; when TypeScript is unavailable the + * middleware passes through to the rest of the chain. + */ +export const createDocgenProvider = (): DocgenMiddleware => { + let managerPromise: Promise | undefined; + + const getManager = (): Promise => { + if (!managerPromise) { + managerPromise = (async () => { + try { + const ts = await import('typescript'); + return new ComponentMetaManager(ts); + } catch { + return undefined; + } + })(); + } + return managerPromise; + }; + + return (nextDocgen: DocgenProvider): DocgenProvider => + async (input) => { + const storyImportPath = getStoryImportPathFromEntry(input.entry); + if (!storyImportPath || !STORY_FILE_TEST_REGEXP.test(storyImportPath)) { + return nextDocgen(input); + } + + const componentMetaManager = await getManager(); + if (!componentMetaManager) { + return nextDocgen(input); + } + + const ours = await buildDocgenPayload(input, { componentMetaManager }); + if (!ours) { + return nextDocgen(input); + } + + const downstream = await nextDocgen(input); + return { ...downstream, ...ours }; + }; +}; diff --git a/code/renderers/react/src/docgen/preset.ts b/code/renderers/react/src/docgen/preset.ts index 9819efd89f36..fa9e3348a8fc 100644 --- a/code/renderers/react/src/docgen/preset.ts +++ b/code/renderers/react/src/docgen/preset.ts @@ -1,58 +1,24 @@ -import { STORY_FILE_TEST_REGEXP, getStoryImportPathFromEntry } from 'storybook/internal/common'; -import type { DocgenProviderPreset } from 'storybook/internal/types'; +import { fileURLToPath } from 'node:url'; -import { getSharedComponentMetaManager } from '../componentManifest/componentMetaManagerSingleton.ts'; -import type { TypescriptOptions } from '../componentManifest/getComponentImports.ts'; -import { buildDocgenPayload } from './buildDocgen.ts'; +import type { DocgenProviderDescriptor } from 'storybook/internal/types'; /** - * React renderer docgen provider — phase 3: real RCM-backed extraction. + * React renderer docgen provider. * - * Receives the authoritative index `entry` from the docgen service. Bails to `nextDocgen` when - * the entry does not resolve to a CSF story file and when TypeScript isn't available. Otherwise - * delegates to {@link buildDocgenPayload} which runs RCM against the file and returns a complete - * {@link DocgenPayload}, or falls through to `nextDocgen` when nothing extractable is found. + * Contributes a {@link DocgenProviderDescriptor} pointing at {@link ./docgen-worker.ts}, which + * core's docgen worker imports and runs off the main thread. The renderer owns no threading code — + * it only resolves the worker module path; nothing else needs to cross the worker boundary, since + * `react-component-meta` extraction reads nothing from the user's `typescript` options (see + * {@link ./docgen-worker.ts}). * - * When extraction succeeds, the payload is merged with downstream via the documented - * `{ ...downstream, ...ours }` spread idiom so any fields a future provider sets and we don't - * know about survive intact. + * The descriptor is appended to the accumulated array so addon providers can stack on top, mirroring + * the previous middleware chain — just composed inside the worker instead of in this process. */ -export const experimental_docgenProvider: DocgenProviderPreset = async (nextDocgen, options) => { - // Resolve the typescript options preset once at chain-build time so the provider closure does - // not re-read it on every call. - const typescriptOptionsPromise = (options.presets?.apply>( - 'typescript', - {} - ) ?? Promise.resolve({})) as Promise>; - - // Pre-boot the component meta manager in the background, off the critical path. - // `import('typescript')` loads a multi-MB module; deferring it to the first docgen request adds - // several seconds of latency to that request. Kicking it off here (fire-and-forget) lets it warm - // up in parallel with the rest of server startup, so `await getSharedComponentMetaManager()` - // below usually resolves instantly. The singleton memoizes and swallows its own errors, so this - // is safe to leave unawaited. - void getSharedComponentMetaManager(); - - return async (input) => { - const storyImportPath = getStoryImportPathFromEntry(input.entry); - if (!storyImportPath || !STORY_FILE_TEST_REGEXP.test(storyImportPath)) { - return nextDocgen(input); - } - - const componentMetaManager = await getSharedComponentMetaManager(); - if (!componentMetaManager) { - return nextDocgen(input); - } - - const ours = await buildDocgenPayload(input, { - componentMetaManager, - typescriptOptions: await typescriptOptionsPromise, - }); - if (!ours) { - return nextDocgen(input); - } - - const downstream = await nextDocgen(input); - return { ...downstream, ...ours }; - }; -}; +export const experimental_docgenProvider = async ( + existing: DocgenProviderDescriptor[] = [] +): Promise => [ + ...existing, + { + moduleSpecifier: fileURLToPath(import.meta.resolve('@storybook/react/internal/docgen-worker')), + }, +];