Merge pull request storybookjs#34496 from NYCU-Chung/fix/docs-blocks-custom-mdx

Addon Docs: Fix Primary and Controls blocks not rendering in custom MDX pages

Author: Sidnioulz

code/addons/docs/src/blocks/blocks/external/ExternalDocsContext.ts

@@ -16,6 +16,11 @@ export class ExternalDocsContext<TRenderer extends Renderer> extends DocsContext
       renderStoryToElement as DocsContextProps<TRenderer>['renderStoryToElement'],
       []
     );
+
+    // External docs are always an author-supplied page (`<ExternalDocs>` passes its children as the
+    // docs page), so `<Primary />` / `<Controls />` should respect the author's story selection
+    // rather than filtering to `autodocs`-tagged stories.
+    this.filterByAutodocs = false;
   }
 
   referenceMeta = (metaExports: ModuleExports, attach: boolean) => {

code/addons/docs/src/blocks/blocks/usePrimaryStory.test.tsx

@@ -19,16 +19,21 @@ const stories: Record<string, Partial<PreparedStory>> = {
   story4: { name: 'Story Four', tags: [] },
 };
 
-const createMockContext = (storyList: PreparedStory[]) => ({
+const createMockContext = (
+  storyList: PreparedStory[],
+  overrides: Partial<DocsContextProps> = {}
+) => ({
   componentStories: vi.fn(() => storyList),
+  filterByAutodocs: true,
+  ...overrides,
 });
 
 const Wrapper: FC<PropsWithChildren<{ context: Partial<DocsContextProps> }>> = ({
   children,
   context,
 }) => <DocsContext.Provider value={context as DocsContextProps}>{children}</DocsContext.Provider>;
 
-describe('usePrimaryStory', () => {
+describe('usePrimaryStory — autodocs page (filterByAutodocs: true)', () => {
   it('ignores !autodocs stories', () => {
     const mockContext = createMockContext([
       stories.story1,
@@ -49,7 +54,7 @@ describe('usePrimaryStory', () => {
     expect(result.current?.name).toBe('Story Two');
   });
 
-  it('returns undefined if no story has "autodocs" tag', () => {
+  it('returns undefined when no story has the autodocs tag', () => {
     const mockContext = createMockContext([stories.story1, stories.story4] as PreparedStory[]);
     const { result } = renderHook(() => usePrimaryStory(), {
       wrapper: ({ children }) => <Wrapper context={mockContext}>{children}</Wrapper>,
@@ -65,3 +70,16 @@ describe('usePrimaryStory', () => {
     expect(result.current).toBeUndefined();
   });
 });
+
+describe('usePrimaryStory — MDX / custom page (filterByAutodocs: false)', () => {
+  it('returns the first story regardless of autodocs tag', () => {
+    const mockContext = createMockContext(
+      [stories.story1, stories.story2, stories.story3] as PreparedStory[],
+      { filterByAutodocs: false }
+    );
+    const { result } = renderHook(() => usePrimaryStory(), {
+      wrapper: ({ children }) => <Wrapper context={mockContext}>{children}</Wrapper>,
+    });
+    expect(result.current?.name).toBe('Story One');
+  });
+});

code/addons/docs/src/blocks/blocks/usePrimaryStory.ts

@@ -6,11 +6,15 @@ import type { PreparedStory } from 'storybook/internal/types';
 import { DocsContext } from './DocsContext';
 
 /**
- * A hook to get the primary story for the current component's doc page. It defines the primary
- * story as the first story that includes the 'autodocs' tag
+ * Returns the primary story for the current docs page. Autodocs pages pick the first story tagged
+ * `autodocs`; MDX pages pick the first story regardless of tag (driven by
+ * `DocsContext.filterByAutodocs`, set by the docs render based on the entry type).
  */
 export const usePrimaryStory = (): PreparedStory | undefined => {
   const context = useContext(DocsContext);
   const stories = context.componentStories();
+  if (context.filterByAutodocs === false) {
+    return stories[0];
+  }
   return stories.find((story) => story.tags.includes(Tag.AUTODOCS));
 };

code/core/src/core-server/utils/StoryIndexGenerator.ts

@@ -33,6 +33,7 @@ import { resolveImport, supportedExtensions } from '../../common/index.ts';
 import { userOrAutoTitleFromSpecifier } from '../../preview-api/modules/store/autoTitle.ts';
 import { sortStoriesV7 } from '../../preview-api/modules/store/sortStories.ts';
 import { Tag } from '../../shared/constants/tags.ts';
+import { isMdxEntry } from '../../shared/utils/story-index-filters.ts';
 import { IndexingError, MultipleIndexingError } from './IndexingError.ts';
 import { autoName } from './autoName.ts';
 import { analyzeMdx } from './analyze-mdx.ts';
@@ -65,11 +66,6 @@ export type StoryIndexGeneratorOptions = {
   build?: StorybookConfigRaw['build'];
 };
 
-/** Was this docs entry generated by a .mdx file? (see discussion below) */
-export function isMdxEntry({ tags }: DocsIndexEntry) {
-  return tags?.includes(Tag.UNATTACHED_MDX) || tags?.includes(Tag.ATTACHED_MDX);
-}
-
 const makeAbsolute = (otherImport: Path, normalizedPath: Path, workingDir: Path) =>
   otherImport.startsWith('.')
     ? slash(resolve(workingDir, normalizeStoryPath(join(dirname(normalizedPath), otherImport))))

code/core/src/core-server/utils/summarizeIndex.ts

@@ -2,7 +2,7 @@ import { isExampleStoryId } from 'storybook/internal/telemetry';
 import type { IndexEntry, StoryIndex } from 'storybook/internal/types';
 
 import { Tag } from '../../shared/constants/tags.ts';
-import { isMdxEntry } from './StoryIndexGenerator.ts';
+import { isMdxEntry } from '../../shared/utils/story-index-filters.ts';
 
 const PAGE_REGEX = /(page|screen)/i;
 

code/core/src/preview-api/modules/preview-web/PreviewWithSelection.tsx

@@ -22,13 +22,14 @@ import {
   MdxFileWithNoCsfReferencesError,
   NoStoryMatchError,
 } from 'storybook/internal/preview-errors';
-import type { DocsIndexEntry, StoryIndex } from 'storybook/internal/types';
+import type { StoryIndex } from 'storybook/internal/types';
 import type { Args, Globals, Renderer, StoryId, ViewMode } from 'storybook/internal/types';
 import type { ModuleImportFn, ProjectAnnotations } from 'storybook/internal/types';
 
 import invariant from 'tiny-invariant';
 
 import { Tag } from '../../../shared/constants/tags.ts';
+import { isMdxEntry } from '../../../shared/utils/story-index-filters.ts';
 import type { StorySpecifier } from '../store/StoryIndexStore.ts';
 import type { MaybePromise } from './Preview.tsx';
 import { Preview } from './Preview.tsx';
@@ -46,11 +47,6 @@ function focusInInput(event: Event) {
   return /input|textarea/i.test(target.tagName) || target.getAttribute('contenteditable') !== null;
 }
 
-/** Was this docs entry generated by a .mdx file? (see discussion below) */
-export function isMdxEntry({ tags }: DocsIndexEntry) {
-  return tags?.includes(Tag.UNATTACHED_MDX) || tags?.includes(Tag.ATTACHED_MDX);
-}
-
 type PossibleRender<TRenderer extends Renderer> =
   | StoryRender<TRenderer>
   | CsfDocsRender<TRenderer>

code/core/src/preview-api/modules/preview-web/docs-context/DocsContext.ts

@@ -32,6 +32,10 @@ export class DocsContext<TRenderer extends Renderer> implements DocsContextProps
 
   private primaryStory?: PreparedStory<TRenderer>;
 
+  // Set by the docs render (autodocs vs MDX) so `usePrimaryStory` knows whether to filter the
+  // CSF file's stories to `autodocs`-tagged ones. See `DocsContextProps.filterByAutodocs`.
+  public filterByAutodocs?: boolean;
+
   constructor(
     public channel: Channel,
     protected store: StoryStore<TRenderer>,

code/core/src/preview-api/modules/preview-web/render/CsfDocsRender.test.ts

@@ -77,3 +77,29 @@ it('attached immediately', async () => {
 
   expect(context.storyById()).toEqual(story);
 });
+
+it('sets filterByAutodocs to true for autodocs pages', async () => {
+  const { story, csfFile, moduleExports } = csfFileParts();
+
+  const store = {
+    loadEntry: () => ({
+      entryExports: moduleExports,
+      csfFiles: [],
+    }),
+    processCSFFileWithCache: () => csfFile,
+    componentStoriesFromCSFFile: () => [story],
+    storyFromCSFFile: () => story,
+  } as unknown as StoryStore<Renderer>;
+
+  const render = new CsfDocsRender(
+    new Channel({}),
+    store,
+    entry,
+    {} as RenderContextCallbacks<Renderer>
+  );
+  await render.prepare();
+
+  const context = render.docsContext(vi.fn());
+
+  expect(context.filterByAutodocs).toBe(true);
+});

code/core/src/preview-api/modules/preview-web/render/CsfDocsRender.ts

@@ -5,6 +5,7 @@ import type { CSFFile, PreparedStory } from 'storybook/internal/types';
 import type { IndexEntry } from 'storybook/internal/types';
 import type { RenderContextCallbacks } from 'storybook/internal/types';
 
+import { isMdxEntry } from '../../../../shared/utils/story-index-filters.ts';
 import type { StoryStore } from '../../../store.ts';
 import { DocsContext } from '../docs-context/DocsContext.ts';
 import type { DocsContextProps } from '../docs-context/DocsContextProps.ts';
@@ -110,6 +111,12 @@ export class CsfDocsRender<TRenderer extends Renderer> implements Render<TRender
     //  - When you create two CSF files that both reference the same title, they are combined into
     //    a single CSF docs entry with a `storiesImport` defined.
     this.csfFiles.forEach((csfFile) => docsContext.attachCSFFile(csfFile));
+
+    // Autodocs pages filter the CSF file's stories down to `autodocs`-tagged ones when picking
+    // the `<Primary />` story; a custom `docs.page` template on an autodocs entry does not change
+    // that (the entry is still not an MDX entry).
+    docsContext.filterByAutodocs = !isMdxEntry(this.entry);
+
     return docsContext;
   }
 

code/core/src/preview-api/modules/preview-web/render/MdxDocsRender.test.ts

@@ -113,6 +113,20 @@ describe('attaching', () => {
 
     expect(context.storyById()).toEqual(story);
   });
+
+  it('sets filterByAutodocs to false for MDX pages', async () => {
+    const render = new MdxDocsRender(
+      new Channel({}),
+      store,
+      attachedEntry,
+      {} as RenderContextCallbacks<Renderer>
+    );
+    await render.prepare();
+
+    const context = render.docsContext(vi.fn());
+
+    expect(context.filterByAutodocs).toBe(false);
+  });
 });
 
 describe('docs parameters', () => {