fix: improve iframe theme initialization

Signed-off-by: axel7083 <[email protected]>

website/src/pages/storybook/index.tsx

@@ -27,13 +27,6 @@ import React, { useEffect, useRef } from 'react';
 import items from './sidebar.cjs'; // generated by storybook plugin
 import styles from './styles.module.css';
 
-// Get the query containing the id of the doc we want to display
-function useQuery(): URLSearchParams {
- const { search } = useLocation();
-
- return React.useMemo(() => new URLSearchParams(search), [search]);
-}
-
 // the iframe is by default not taking all height
 // adding an observer to ensure it always takes all the space it need
 function observe(iframe: HTMLIFrameElement): void {
@@ -53,27 +46,57 @@ function observe(iframe: HTMLIFrameElement): void {
 }
 
 function StorybookRoot(): JSX.Element {
- const query = useQuery();
  // eslint-disable-next-line no-null/no-null
  const iframeRef = useRef<HTMLIFrameElement>(null);
  const { isDarkTheme } = useColorMode();
- const id = query.get('id');
+
+ const { search } = useLocation();
+ const [id, setId] = React.useState<string | undefined>(undefined);
 
  useEffect(() => {
+ const queryId = new URLSearchParams(search).get('id');
+ if (queryId) {
+ setId(queryId);
+ }
+ }, [search]);
+
+ const storybookItems: PropSidebarItem[] = items.map(item => ({
+ type: 'link',
+ href: `/storybook?id=${item.id}`,
+ label: item.label,
+ }));
+
+ const notifyIframe = (): void => {
  // we send the iframe the dark mode change https://storybook.js.org/addons/storybook-dark-mode
  iframeRef?.current?.contentWindow?.postMessage(
  JSON.stringify({
  key: 'storybook-channel',
  event: { type: 'DARK_MODE', args: [isDarkTheme] },
  }),
  );
- }, [isDarkTheme]);
+ };
 
- const storybookItems: PropSidebarItem[] = items.map(item => ({
- type: 'link',
- href: `/storybook?id=${item.id}`,
- label: item.label,
- }));
+ const onLoad = (e: React.SyntheticEvent<HTMLIFrameElement>): void => {
+ // observe resize
+ observe(e.currentTarget);
+
+ // https://github.com/storybookjs/storybook/blob/1943ee6b88d89c963f15ef4087aeabe64d05c9a1/code/lib/core-events/src/index.ts#L65
+ window.addEventListener('message', message => {
+ if (message.source !== iframeRef?.current?.contentWindow) {
+ return;
+ }
+
+ const data = JSON.parse(message.data);
+ if (!('key' in data) || data['key'] !== 'storybook-channel') return;
+ if (!('event' in data) || typeof data['event'] !== 'object' || data['event']['type'] !== 'docsRendered') return;
+
+ notifyIframe();
+ });
+ };
+
+ useEffect(() => {
+ notifyIframe();
+ }, [isDarkTheme]);
 
  return (
  <div className={clsx(styles.storybookRoot)}>
@@ -82,7 +105,7 @@ function StorybookRoot(): JSX.Element {
  </aside>
  <iframe
  ref={iframeRef}
- onLoad={e => observe(e.currentTarget)}
+ onLoad={onLoad}
  src={`/storybook-iframe?id=${id}`}
  style={{ width: '100%', height: '100%' }}
  />