Merge branch 'on-demand-store' into fix/on-demand-store-addon-docs

addons/docs/src/blocks/ArgsTable.tsx

@@ -162,7 +162,6 @@ export const StoryTable: FC<
     const story = useStory(storyId, context);
     // eslint-disable-next-line prefer-const
     let [args, updateArgs, resetArgs] = useArgs(storyId, context);
-
     if (!story) {
       return <div>Loading...</div>;
     }

addons/docs/src/blocks/SourceContainer.tsx

@@ -18,24 +18,18 @@ export const SourceContainer: FC<{}> = ({ children }) => {
   const [sources, setSources] = useState<StorySources>({});
   const channel = addons.getChannel();
 
-  const sourcesRef = React.useRef<StorySources>();
-  const handleSnippetRendered = (id: StoryId, newItem: SourceItem) => {
-    if (newItem !== sources[id]) {
-      const newSources = { ...sourcesRef.current, [id]: newItem };
-      sourcesRef.current = newSources;
-    }
-  };
-
-  // Bind this early (instead of inside `useEffect`), because the `SNIPPET_RENDERED` event
-  // is triggered *during* the rendering process, not after. We have to use the ref
-  // to ensure we don't end up calling setState outside the effect though.
-  channel.on(SNIPPET_RENDERED, handleSnippetRendered);
-
   useEffect(() => {
-    const current = sourcesRef.current || {};
-    if (!deepEqual(sources, current)) {
-      setSources(current);
-    }
+    const handleSnippetRendered = (id: StoryId, newItem: SourceItem) => {
+      if (newItem !== sources[id]) {
+        const newSources = { ...sources, [id]: newItem };
+
+        if (!deepEqual(sources, newSources)) {
+          setSources(newSources);
+        }
+      }
+    };
+
+    channel.on(SNIPPET_RENDERED, handleSnippetRendered);
 
     return () => channel.off(SNIPPET_RENDERED, handleSnippetRendered);
   });

addons/docs/src/blocks/Story.tsx

@@ -6,6 +6,7 @@ import React, {
   useContext,
   useRef,
   useEffect,
+  useMemo,
 } from 'react';
 import { MDXProvider } from '@mdx-js/react';
 import { resetComponents, Story as PureStory } from '@storybook/components';
@@ -105,6 +106,15 @@ const Story: FunctionComponent<StoryProps> = (props) => {
   const ref = useRef();
   const story = useStory(getStoryId(props, context), context);
 
+  // Ensure we wait until this story is properly rendered in the docs context.
+  // The purpose of this is to ensure that that the `DOCS_RENDERED` event isn't emitted
+  // until all stories on the page have rendered.
+  const { id: storyId, registerRenderingStory } = context;
+  const storyRendered = useMemo(registerRenderingStory, [storyId]);
+  useEffect(() => {
+    if (story) storyRendered();
+  }, [story]);
+
   useEffect(() => {
     let cleanup: () => void;
     if (story && ref.current) {

addons/docs/src/frameworks/angular/sourceDecorator.ts

@@ -1,4 +1,4 @@
-import { addons } from '@storybook/addons';
+import { addons, useEffect } from '@storybook/addons';
 import { PartialStoryFn } from '@storybook/csf';
 import { StoryContext, AngularFramework } from '@storybook/angular';
 import { computesTemplateSourceFromComponent } from '@storybook/angular/renderer';
@@ -44,20 +44,21 @@ export const sourceDecorator = (
 
   const { component, argTypes } = context;
 
+  let toEmit: string;
+  useEffect(() => {
+    if (toEmit) channel.emit(SNIPPET_RENDERED, context.id, prettyUp(template));
+  });
+
   if (component && !userDefinedTemplate) {
     const source = computesTemplateSourceFromComponent(component, props, argTypes);
 
     // We might have a story with a Directive or Service defined as the component
     // In these cases there might exist a template, even if we aren't able to create source from component
     if (source || template) {
-      channel.emit(SNIPPET_RENDERED, context.id, prettyUp(source || template));
+      toEmit = prettyUp(source || template);
     }
-    return story;
-  }
-
-  if (template) {
-    channel.emit(SNIPPET_RENDERED, context.id, prettyUp(template));
-    return story;
+  } else if (template) {
+    toEmit = prettyUp(template);
   }
 
   return story;

addons/docs/src/frameworks/html/sourceDecorator.test.ts

@@ -1,15 +1,18 @@
-import { addons, StoryContext } from '@storybook/addons';
+import { addons, StoryContext, useEffect } from '@storybook/addons';
 import { sourceDecorator } from './sourceDecorator';
 import { SNIPPET_RENDERED } from '../../shared';
 
 jest.mock('@storybook/addons');
 const mockedAddons = addons as jest.Mocked<typeof addons>;
+const mockedUseEffect = useEffect as jest.Mocked<typeof useEffect>;
 
 expect.addSnapshotSerializer({
   print: (val: any) => val,
   test: (val) => typeof val === 'string',
 });
 
+const tick = () => new Promise((r) => setTimeout(r, 0));
+
 const makeContext = (name: string, parameters: any, args: any, extra?: object): StoryContext => ({
   id: `html-test--${name}`,
   kind: 'js-text',
@@ -25,15 +28,17 @@ describe('sourceDecorator', () => {
   let mockChannel: { on: jest.Mock; emit?: jest.Mock };
   beforeEach(() => {
     mockedAddons.getChannel.mockReset();
+    mockedUseEffect.mockImplementation((cb) => setTimeout(cb, 0));
 
     mockChannel = { on: jest.fn(), emit: jest.fn() };
     mockedAddons.getChannel.mockReturnValue(mockChannel as any);
   });
 
-  it('should render dynamically for args stories', () => {
+  it('should render dynamically for args stories', async () => {
     const storyFn = (args: any) => `<div>args story</div>`;
     const context = makeContext('args', { __isArgsStory: true }, {});
     sourceDecorator(storyFn, context);
+    await tick();
     expect(mockChannel.emit).toHaveBeenCalledWith(
       SNIPPET_RENDERED,
       'html-test--args',
@@ -41,7 +46,7 @@ describe('sourceDecorator', () => {
     );
   });
 
-  it('should dedent source by default', () => {
+  it('should dedent source by default', async () => {
     const storyFn = (args: any) => `
       <div>
         args story
@@ -49,6 +54,7 @@ describe('sourceDecorator', () => {
     `;
     const context = makeContext('args', { __isArgsStory: true }, {});
     sourceDecorator(storyFn, context);
+    await tick();
     expect(mockChannel.emit).toHaveBeenCalledWith(
       SNIPPET_RENDERED,
       'html-test--args',
@@ -56,14 +62,15 @@ describe('sourceDecorator', () => {
     );
   });
 
-  it('should skip dynamic rendering for no-args stories', () => {
+  it('should skip dynamic rendering for no-args stories', async () => {
     const storyFn = () => `<div>classic story</div>`;
     const context = makeContext('classic', {}, {});
     sourceDecorator(storyFn, context);
+    await tick();
     expect(mockChannel.emit).not.toHaveBeenCalled();
   });
 
-  it('should use the originalStoryFn if excludeDecorators is set', () => {
+  it('should use the originalStoryFn if excludeDecorators is set', async () => {
     const storyFn = (args: any) => `<div>args story</div>`;
     const decoratedStoryFn = (args: any) => `
       <div style="padding: 25px; border: 3px solid red;">${storyFn(args)}</div>
@@ -82,6 +89,7 @@ describe('sourceDecorator', () => {
       { originalStoryFn: storyFn }
     );
     sourceDecorator(decoratedStoryFn, context);
+    await tick();
     expect(mockChannel.emit).toHaveBeenCalledWith(
       SNIPPET_RENDERED,
       'html-test--args',
@@ -89,12 +97,13 @@ describe('sourceDecorator', () => {
     );
   });
 
-  it('allows the snippet output to be modified by transformSource', () => {
+  it('allows the snippet output to be modified by transformSource', async () => {
     const storyFn = (args: any) => `<div>args story</div>`;
     const transformSource = (dom: string) => `<p>${dom}</p>`;
     const docs = { transformSource };
     const context = makeContext('args', { __isArgsStory: true, docs }, {});
     sourceDecorator(storyFn, context);
+    await tick();
     expect(mockChannel.emit).toHaveBeenCalledWith(
       SNIPPET_RENDERED,
       'html-test--args',

addons/docs/src/frameworks/html/sourceDecorator.ts

@@ -1,5 +1,5 @@
 /* global window */
-import { addons } from '@storybook/addons';
+import { addons, useEffect } from '@storybook/addons';
 import { ArgsStoryFn, PartialStoryFn, StoryContext } from '@storybook/csf';
 import dedent from 'ts-dedent';
 import { HtmlFramework } from '@storybook/html';
@@ -40,11 +40,13 @@ export function sourceDecorator(
     ? (context.originalStoryFn as ArgsStoryFn<HtmlFramework>)(context.args, context)
     : storyFn();
 
+  let source: string;
   if (typeof story === 'string' && !skipSourceRender(context)) {
-    const source = applyTransformSource(story, context);
-
-    addons.getChannel().emit(SNIPPET_RENDERED, context.id, source);
+    source = applyTransformSource(story, context);
   }
+  useEffect(() => {
+    if (source) addons.getChannel().emit(SNIPPET_RENDERED, context.id, source);
+  });
 
   return story;
 }

addons/docs/src/frameworks/react/jsxDecorator.test.tsx

@@ -1,12 +1,13 @@
 /* eslint-disable jsx-a11y/no-static-element-interactions, jsx-a11y/click-events-have-key-events */
 import React from 'react';
 import PropTypes from 'prop-types';
-import { addons, StoryContext } from '@storybook/addons';
+import { addons, StoryContext, useEffect } from '@storybook/addons';
 import { renderJsx, jsxDecorator } from './jsxDecorator';
 import { SNIPPET_RENDERED } from '../../shared';
 
 jest.mock('@storybook/addons');
 const mockedAddons = addons as jest.Mocked<typeof addons>;
+const mockedUseEffect = useEffect as jest.Mocked<typeof useEffect>;
 
 expect.addSnapshotSerializer({
   print: (val: any) => val,
@@ -168,15 +169,17 @@ describe('jsxDecorator', () => {
   let mockChannel: { on: jest.Mock; emit?: jest.Mock };
   beforeEach(() => {
     mockedAddons.getChannel.mockReset();
+    mockedUseEffect.mockImplementation((cb) => setTimeout(cb, 0));
 
     mockChannel = { on: jest.fn(), emit: jest.fn() };
     mockedAddons.getChannel.mockReturnValue(mockChannel as any);
   });
 
-  it('should render dynamically for args stories', () => {
+  it('should render dynamically for args stories', async () => {
     const storyFn = (args: any) => <div>args story</div>;
     const context = makeContext('args', { __isArgsStory: true }, {});
     jsxDecorator(storyFn, context);
+    await new Promise((r) => setTimeout(r, 0));
     expect(mockChannel.emit).toHaveBeenCalledWith(
       SNIPPET_RENDERED,
       'jsx-test--args',
@@ -184,7 +187,7 @@ describe('jsxDecorator', () => {
     );
   });
 
-  it('should not render decorators when provided excludeDecorators parameter', () => {
+  it('should not render decorators when provided excludeDecorators parameter', async () => {
     const storyFn = (args: any) => <div>args story</div>;
     const decoratedStoryFn = (args: any) => (
       <div style={{ padding: 25, border: '3px solid red' }}>{storyFn(args)}</div>
@@ -203,6 +206,8 @@ describe('jsxDecorator', () => {
       { originalStoryFn: storyFn }
     );
     jsxDecorator(decoratedStoryFn, context);
+    await new Promise((r) => setTimeout(r, 0));
+
     expect(mockChannel.emit).toHaveBeenCalledWith(
       SNIPPET_RENDERED,
       'jsx-test--args',
@@ -210,20 +215,24 @@ describe('jsxDecorator', () => {
     );
   });
 
-  it('should skip dynamic rendering for no-args stories', () => {
+  it('should skip dynamic rendering for no-args stories', async () => {
     const storyFn = () => <div>classic story</div>;
     const context = makeContext('classic', {}, {});
     jsxDecorator(storyFn, context);
+    await new Promise((r) => setTimeout(r, 0));
+
     expect(mockChannel.emit).not.toHaveBeenCalled();
   });
 
   // This is deprecated, but still test it
-  it('allows the snippet output to be modified by onBeforeRender', () => {
+  it('allows the snippet output to be modified by onBeforeRender', async () => {
     const storyFn = (args: any) => <div>args story</div>;
     const onBeforeRender = (dom: string) => `<p>${dom}</p>`;
     const jsx = { onBeforeRender };
     const context = makeContext('args', { __isArgsStory: true, jsx }, {});
     jsxDecorator(storyFn, context);
+    await new Promise((r) => setTimeout(r, 0));
+
     expect(mockChannel.emit).toHaveBeenCalledWith(
       SNIPPET_RENDERED,
       'jsx-test--args',
@@ -231,12 +240,14 @@ describe('jsxDecorator', () => {
     );
   });
 
-  it('allows the snippet output to be modified by transformSource', () => {
+  it('allows the snippet output to be modified by transformSource', async () => {
     const storyFn = (args: any) => <div>args story</div>;
     const transformSource = (dom: string) => `<p>${dom}</p>`;
     const jsx = { transformSource };
     const context = makeContext('args', { __isArgsStory: true, jsx }, {});
     jsxDecorator(storyFn, context);
+    await new Promise((r) => setTimeout(r, 0));
+
     expect(mockChannel.emit).toHaveBeenCalledWith(
       SNIPPET_RENDERED,
       'jsx-test--args',
@@ -253,7 +264,7 @@ describe('jsxDecorator', () => {
     expect(transformSource).toHaveBeenCalledWith('<div>\n  args story\n</div>', context);
   });
 
-  it('renders MDX properly', () => {
+  it('renders MDX properly', async () => {
     // FIXME: generate this from actual MDX
     const mdxElement = {
       type: { displayName: 'MDXCreateElement' },
@@ -265,6 +276,7 @@ describe('jsxDecorator', () => {
     };
 
     jsxDecorator(() => mdxElement, makeContext('mdx-args', { __isArgsStory: true }, {}));
+    await new Promise((r) => setTimeout(r, 0));
 
     expect(mockChannel.emit).toHaveBeenCalledWith(
       SNIPPET_RENDERED,

addons/docs/src/frameworks/react/jsxDecorator.tsx

@@ -3,7 +3,7 @@ import reactElementToJSXString, { Options } from 'react-element-to-jsx-string';
 import dedent from 'ts-dedent';
 import deprecate from 'util-deprecate';
 
-import { addons } from '@storybook/addons';
+import { addons, useEffect } from '@storybook/addons';
 import { StoryContext, ArgsStoryFn, PartialStoryFn } from '@storybook/csf';
 import { logger } from '@storybook/client-logger';
 import { ReactFramework } from '@storybook/react';
@@ -175,16 +175,22 @@ export const jsxDecorator = (
   storyFn: PartialStoryFn<ReactFramework>,
   context: StoryContext<ReactFramework>
 ) => {
+  const channel = addons.getChannel();
+  const skip = skipJsxRender(context);
   const story = storyFn();
 
+  let jsx = '';
+
+  useEffect(() => {
+    if (!skip) channel.emit(SNIPPET_RENDERED, (context || {}).id, jsx);
+  });
+
   // We only need to render JSX if the source block is actually going to
   // consume it. Otherwise it's just slowing us down.
-  if (skipJsxRender(context)) {
+  if (skip) {
     return story;
   }
 
-  const channel = addons.getChannel();
-
   const options = {
     ...defaultOpts,
     ...(context?.parameters.jsx || {}),
@@ -197,13 +203,10 @@ export const jsxDecorator = (
 
   const sourceJsx = mdxToJsx(storyJsx);
 
-  let jsx = '';
   const rendered = renderJsx(sourceJsx, options);
   if (rendered) {
     jsx = applyTransformSource(rendered, options, context);
   }
 
-  channel.emit(SNIPPET_RENDERED, (context || {}).id, jsx);
-
   return story;
 };

addons/docs/src/frameworks/svelte/sourceDecorator.ts

@@ -1,4 +1,4 @@
-import { addons } from '@storybook/addons';
+import { addons, useEffect } from '@storybook/addons';
 import { ArgTypes, Args, StoryContext, AnyFramework } from '@storybook/csf';
 
 import { SourceType, SNIPPET_RENDERED } from '../../shared';
@@ -145,14 +145,21 @@ function getWrapperProperties(component: any) {
  * @param context  StoryContext
  */
 export const sourceDecorator = (storyFn: any, context: StoryContext<AnyFramework>) => {
+  const channel = addons.getChannel();
+  const skip = skipSourceRender(context);
   const story = storyFn();
 
-  if (skipSourceRender(context)) {
+  let source: string;
+  useEffect(() => {
+    if (!skip && source) {
+      channel.emit(SNIPPET_RENDERED, (context || {}).id, source);
+    }
+  });
+
+  if (skip) {
     return story;
   }
 
-  const channel = addons.getChannel();
-
   const { parameters = {}, args = {} } = context || {};
   let { Component: component = {} } = story;
 
@@ -161,11 +168,7 @@ export const sourceDecorator = (storyFn: any, context: StoryContext<AnyFramework
     component = parameters.component;
   }
 
-  const source = generateSvelteSource(component, args, context?.argTypes, slotProperty);
-
-  if (source) {
-    channel.emit(SNIPPET_RENDERED, (context || {}).id, source);
-  }
+  source = generateSvelteSource(component, args, context?.argTypes, slotProperty);
 
   return story;
 };

addons/docs/src/frameworks/web-components/sourceDecorator.test.ts

@@ -1,17 +1,20 @@
 import { html } from 'lit-html';
 import { styleMap } from 'lit-html/directives/style-map';
-import { addons, StoryContext } from '@storybook/addons';
+import { addons, StoryContext, useEffect } from '@storybook/addons';
 import { sourceDecorator } from './sourceDecorator';
 import { SNIPPET_RENDERED } from '../../shared';
 
 jest.mock('@storybook/addons');
 const mockedAddons = addons as jest.Mocked<typeof addons>;
+const mockedUseEffect = useEffect as jest.Mocked<typeof useEffect>;
 
 expect.addSnapshotSerializer({
   print: (val: any) => val,
   test: (val) => typeof val === 'string',
 });
 
+const tick = () => new Promise((r) => setTimeout(r, 0));
+
 const makeContext = (name: string, parameters: any, args: any, extra?: object): StoryContext => ({
   id: `lit-test--${name}`,
   kind: 'js-text',
@@ -27,15 +30,17 @@ describe('sourceDecorator', () => {
   let mockChannel: { on: jest.Mock; emit?: jest.Mock };
   beforeEach(() => {
     mockedAddons.getChannel.mockReset();
+    mockedUseEffect.mockImplementation((cb) => setTimeout(cb, 0));
 
     mockChannel = { on: jest.fn(), emit: jest.fn() };
     mockedAddons.getChannel.mockReturnValue(mockChannel as any);
   });
 
-  it('should render dynamically for args stories', () => {
+  it('should render dynamically for args stories', async () => {
     const storyFn = (args: any) => html`<div>args story</div>`;
     const context = makeContext('args', { __isArgsStory: true }, {});
     sourceDecorator(storyFn, context);
+    await tick();
     expect(mockChannel.emit).toHaveBeenCalledWith(
       SNIPPET_RENDERED,
       'lit-test--args',
@@ -43,14 +48,15 @@ describe('sourceDecorator', () => {
     );
   });
 
-  it('should skip dynamic rendering for no-args stories', () => {
+  it('should skip dynamic rendering for no-args stories', async () => {
     const storyFn = () => html`<div>classic story</div>`;
     const context = makeContext('classic', {}, {});
     sourceDecorator(storyFn, context);
+    await tick();
     expect(mockChannel.emit).not.toHaveBeenCalled();
   });
 
-  it('should use the originalStoryFn if excludeDecorators is set', () => {
+  it('should use the originalStoryFn if excludeDecorators is set', async () => {
     const storyFn = (args: any) => html`<div>args story</div>`;
     const decoratedStoryFn = (args: any) => html`
       <div style=${styleMap({ padding: `${25}px`, border: '3px solid red' })}>${storyFn(args)}</div>
@@ -69,6 +75,7 @@ describe('sourceDecorator', () => {
       { originalStoryFn: storyFn }
     );
     sourceDecorator(decoratedStoryFn, context);
+    await tick();
     expect(mockChannel.emit).toHaveBeenCalledWith(
       SNIPPET_RENDERED,
       'lit-test--args',
@@ -76,12 +83,13 @@ describe('sourceDecorator', () => {
     );
   });
 
-  it('allows the snippet output to be modified by transformSource', () => {
+  it('allows the snippet output to be modified by transformSource', async () => {
     const storyFn = (args: any) => html`<div>args story</div>`;
     const transformSource = (dom: string) => `<p>${dom}</p>`;
     const docs = { transformSource };
     const context = makeContext('args', { __isArgsStory: true, docs }, {});
     sourceDecorator(storyFn, context);
+    await tick();
     expect(mockChannel.emit).toHaveBeenCalledWith(
       SNIPPET_RENDERED,
       'lit-test--args',

addons/docs/src/frameworks/web-components/sourceDecorator.ts

@@ -1,7 +1,7 @@
 /* global window */
 import { render } from 'lit-html';
 import { ArgsStoryFn, PartialStoryFn, StoryContext } from '@storybook/csf';
-import { addons } from '@storybook/addons';
+import { addons, useEffect } from '@storybook/addons';
 import { WebComponentsFramework } from '@storybook/web-components';
 
 import { SNIPPET_RENDERED, SourceType } from '../../shared';
@@ -37,11 +37,14 @@ export function sourceDecorator(
     ? (context.originalStoryFn as ArgsStoryFn<WebComponentsFramework>)(context.args, context)
     : storyFn();
 
+  let source: string;
+  useEffect(() => {
+    if (source) addons.getChannel().emit(SNIPPET_RENDERED, context.id, source);
+  });
   if (!skipSourceRender(context)) {
     const container = window.document.createElement('div');
     render(story, container);
-    const source = applyTransformSource(container.innerHTML.replace(/<!---->/g, ''), context);
-    if (source) addons.getChannel().emit(SNIPPET_RENDERED, context.id, source);
+    source = applyTransformSource(container.innerHTML.replace(/<!---->/g, ''), context);
   }
 
   return story;

lib/builder-webpack4/src/preview/iframe-webpack.config.ts

@@ -133,7 +133,7 @@ export default async (options: Options & Record<string, any>): Promise<Configura
       );
       entries.push(`${configFilename}-generated-config-entry.js`);
     });
-    if (stories) {
+    if (stories.length > 0) {
       const storyTemplate = await readTemplate(
         path.join(__dirname, 'virtualModuleStory.template.js')
       );

lib/builder-webpack5/src/preview/iframe-webpack.config.ts

@@ -128,7 +128,7 @@ export default async (options: Options & Record<string, any>): Promise<Configura
       );
       entries.push(`${configFilename}-generated-config-entry.js`);
     });
-    if (stories) {
+    if (stories.length > 0) {
       const storyTemplate = await readTemplate(
         path.join(__dirname, 'virtualModuleStory.template.js')
       );

lib/preview-web/src/PreviewWeb.test.ts

@@ -484,6 +484,28 @@ describe('PreviewWeb', () => {
 
         expect(mockChannel.emit).toHaveBeenCalledWith(Events.DOCS_RENDERED, 'component-one--a');
       });
+
+      it('emits DOCS_RENDERED after all stories are rendered', async () => {
+        document.location.search = '?id=component-one--a&viewMode=docs';
+        const [reactDomGate, openReactDomGate] = createGate();
+
+        let rendered;
+        ReactDOM.render.mockImplementationOnce((docsElement, element, cb) => {
+          rendered = docsElement.props.context.registerRenderingStory();
+          openReactDomGate();
+          cb();
+        });
+
+        await new PreviewWeb({ importFn, fetchStoryIndex }).initialize({ getProjectAnnotations });
+
+        // Wait for `ReactDOM.render()` to be called. We should still be waiting for the story
+        await reactDomGate;
+        expect(mockChannel.emit).not.toHaveBeenCalledWith(Events.DOCS_RENDERED, 'component-one--a');
+
+        rendered();
+        await waitForRender();
+        expect(mockChannel.emit).toHaveBeenCalledWith(Events.DOCS_RENDERED, 'component-one--a');
+      });
     });
   });
 

lib/preview-web/src/PreviewWeb.tsx

@@ -320,6 +320,7 @@ export class PreviewWeb<TFramework extends AnyFramework> {
     const csfFile: CSFFile<TFramework> = await this.storyStore.loadCSFFileByStoryId(id, {
       sync: false,
     });
+    const renderingStoryPromises: Promise<void>[] = [];
     const docsContext = {
       id,
       title,
@@ -329,6 +330,17 @@ export class PreviewWeb<TFramework extends AnyFramework> {
       componentStories: () => this.storyStore.componentStoriesFromCSFFile({ csfFile }),
       loadStory: (storyId: StoryId) => this.storyStore.loadStory({ storyId }),
       renderStoryToElement: this.renderStoryToElement.bind(this),
+      // Keep track of the stories that are rendered by the <Story/> component and don't emit
+      // the DOCS_RENDERED event(below) until they have all marked themselves as rendered.
+      registerRenderingStory: () => {
+        let rendered: (v: void) => void;
+        renderingStoryPromises.push(
+          new Promise((resolve) => {
+            rendered = resolve;
+          })
+        );
+        return rendered;
+      },
       getStoryContext: (renderedStory: Story<TFramework>) =>
         ({
           ...this.storyStore.getStoryContext(renderedStory),
@@ -350,7 +362,10 @@ export class PreviewWeb<TFramework extends AnyFramework> {
         <Page />
       </DocsContainer>
     );
-    ReactDOM.render(docsElement, element, () => this.channel.emit(Events.DOCS_RENDERED, id));
+    ReactDOM.render(docsElement, element, async () => {
+      await Promise.all(renderingStoryPromises);
+      this.channel.emit(Events.DOCS_RENDERED, id);
+    });
   }
 
   renderStory({ story }: { story: Story<TFramework> }) {

lib/preview-web/src/types.ts

@@ -17,6 +17,7 @@ export interface DocsContextProps<TFramework extends AnyFramework = AnyFramework
   loadStory: (id: StoryId) => Promise<Story<TFramework>>;
   renderStoryToElement: PreviewWeb<TFramework>['renderStoryToElement'];
   getStoryContext: (story: Story<TFramework>) => StoryContextForLoaders<TFramework>;
+  registerRenderingStory: () => (v: void) => void;
 
   /**
    * mdxStoryNameToKey is an MDX-compiler-generated mapping of an MDX story's