Implement beforeEach cleanup callbacks

code/addons/links/package.json

@@ -67,7 +67,7 @@
     "prep": "node --loader ../../../scripts/node_modules/esbuild-register/loader.js -r ../../../scripts/node_modules/esbuild-register/register.js ../../../scripts/prepare/addon-bundle.ts"
   },
   "dependencies": {
-    "@storybook/csf": "0.1.4--canary.82.941efd8.0",
+    "@storybook/csf": "0.1.4--canary.82.5e73744.0",
     "@storybook/global": "^5.0.0",
     "ts-dedent": "^2.0.0"
   },

code/lib/codemod/package.json

@@ -57,7 +57,7 @@
     "@babel/core": "^7.23.2",
     "@babel/preset-env": "^7.23.2",
     "@babel/types": "^7.23.0",
-    "@storybook/csf": "0.1.4--canary.82.941efd8.0",
+    "@storybook/csf": "0.1.4--canary.82.5e73744.0",
     "@storybook/csf-tools": "workspace:*",
     "@storybook/node-logger": "workspace:*",
     "@storybook/types": "workspace:*",

code/lib/core-server/package.json

@@ -62,7 +62,7 @@
     "@storybook/channels": "workspace:*",
     "@storybook/core-common": "workspace:*",
     "@storybook/core-events": "workspace:*",
-    "@storybook/csf": "0.1.4--canary.82.941efd8.0",
+    "@storybook/csf": "0.1.4--canary.82.5e73744.0",
     "@storybook/csf-tools": "workspace:*",
     "@storybook/docs-mdx": "3.0.0",
     "@storybook/global": "^5.0.0",

code/lib/csf-tools/package.json

@@ -46,7 +46,7 @@
     "@babel/parser": "^7.23.0",
     "@babel/traverse": "^7.23.2",
     "@babel/types": "^7.23.0",
-    "@storybook/csf": "0.1.4--canary.82.941efd8.0",
+    "@storybook/csf": "0.1.4--canary.82.5e73744.0",
     "@storybook/types": "workspace:*",
     "fs-extra": "^11.1.0",
     "recast": "^0.23.5",

code/lib/manager-api/package.json

@@ -47,7 +47,7 @@
     "@storybook/channels": "workspace:*",
     "@storybook/client-logger": "workspace:*",
     "@storybook/core-events": "workspace:*",
-    "@storybook/csf": "0.1.4--canary.82.941efd8.0",
+    "@storybook/csf": "0.1.4--canary.82.5e73744.0",
     "@storybook/global": "^5.0.0",
     "@storybook/icons": "^1.2.5",
     "@storybook/router": "workspace:*",

code/lib/preview-api/package.json

@@ -47,7 +47,7 @@
     "@storybook/channels": "workspace:*",
     "@storybook/client-logger": "workspace:*",
     "@storybook/core-events": "workspace:*",
-    "@storybook/csf": "0.1.4--canary.82.941efd8.0",
+    "@storybook/csf": "0.1.4--canary.82.5e73744.0",
     "@storybook/global": "^5.0.0",
     "@storybook/types": "workspace:*",
     "@types/qs": "^6.9.5",

code/lib/preview-api/src/modules/preview-web/render/StoryRender.ts

@@ -27,6 +27,7 @@ const { AbortController } = globalThis;
 export type RenderPhase =
   | 'preparing'
   | 'loading'
+  | 'beforeEach'
   | 'rendering'
   | 'playing'
   | 'played'
@@ -101,7 +102,7 @@ export class StoryRender<TRenderer extends Renderer> implements Render<TRenderer
     });
 
     if ((this.abortController as AbortController).signal.aborted) {
-      this.store.cleanupStory(this.story as PreparedStory<TRenderer>);
+      await this.store.cleanupStory(this.story as PreparedStory<TRenderer>);
       throw PREPARE_ABORTED;
     }
   }
@@ -120,7 +121,7 @@ export class StoryRender<TRenderer extends Renderer> implements Render<TRenderer
   }
 
   isPending() {
-    return ['rendering', 'playing'].includes(this.phase as RenderPhase);
+    return ['loading', 'beforeEach', 'rendering', 'playing'].includes(this.phase as RenderPhase);
   }
 
   async renderToElement(canvasElement: TRenderer['canvasElement']) {
@@ -149,10 +150,20 @@ export class StoryRender<TRenderer extends Renderer> implements Render<TRenderer
   } = {}) {
     const { canvasElement } = this;
     if (!this.story) throw new Error('cannot render when not prepared');
+    const story = this.story;
     if (!canvasElement) throw new Error('cannot render when canvasElement is unset');
 
-    const { id, componentId, title, name, tags, applyLoaders, unboundStoryFn, playFunction } =
-      this.story;
+    const {
+      id,
+      componentId,
+      title,
+      name,
+      tags,
+      applyLoaders,
+      applyBeforeEach,
+      unboundStoryFn,
+      playFunction,
+    } = story;
 
     if (forceRemount && !initial) {
       // NOTE: we don't check the cancel actually worked here, so the previous
@@ -176,9 +187,7 @@ export class StoryRender<TRenderer extends Renderer> implements Render<TRenderer
           canvasElement,
         } as unknown as StoryContextForLoaders<TRenderer>);
       });
-      if (abortSignal.aborted) {
-        return;
-      }
+      if (abortSignal.aborted) return;
 
       const renderStoryContext: StoryContext<TRenderer> = {
         ...loadedContext!,
@@ -189,6 +198,14 @@ export class StoryRender<TRenderer extends Renderer> implements Render<TRenderer
         // We should consider parameterizing the story types with TRenderer['canvasElement'] in the future
         canvasElement: canvasElement as any,
       };
+
+      await this.runPhase(abortSignal, 'beforeEach', async () => {
+        const cleanupCallbacks = await applyBeforeEach(renderStoryContext);
+        this.store.addCleanupCallbacks(story, cleanupCallbacks);
+      });
+
+      if (abortSignal.aborted) return;
+
       const renderContext: RenderContext<TRenderer> = {
         componentId,
         title,
@@ -290,8 +307,10 @@ export class StoryRender<TRenderer extends Renderer> implements Render<TRenderer
     this.torndown = true;
     this.cancelRender();
 
-    // If the story has loaded, we need to cleanup
-    if (this.story) this.store.cleanupStory(this.story);
+    // If the story has loaded, we need to clean up
+    if (this.story) {
+      await this.store.cleanupStory(this.story);
+    }
 
     // Check if we're done rendering/playing. If not, we may have to reload the page.
     // Wait several ticks that may be needed to handle the abort, then try again.

code/lib/preview-api/src/modules/store/StoryStore.test.ts

@@ -409,7 +409,7 @@ describe('StoryStore', () => {
 
       const { hooks } = store.getStoryContext(story) as { hooks: HooksContext<Renderer> };
       hooks.clean = vi.fn();
-      store.cleanupStory(story);
+      await store.cleanupStory(story);
       expect(hooks.clean).toHaveBeenCalled();
     });
   });

code/lib/preview-api/src/modules/store/StoryStore.ts

@@ -40,6 +40,7 @@ import {
   normalizeProjectAnnotations,
   prepareContext,
 } from './csf';
+import type { CleanupCallback } from '@storybook/csf';
 
 // TODO -- what are reasonable values for these?
 const CSF_CACHE_SIZE = 1000;
@@ -56,6 +57,8 @@ export class StoryStore<TRenderer extends Renderer> {
 
   hooks: Record<StoryId, HooksContext<TRenderer>>;
 
+  cleanupCallbacks: Record<StoryId, CleanupCallback[] | undefined>;
+
   cachedCSFFiles?: Record<Path, CSFFile<TRenderer>>;
 
   processCSFFileWithCache: typeof processCSFFile;
@@ -79,6 +82,7 @@ export class StoryStore<TRenderer extends Renderer> {
     this.args = new ArgsStore();
     this.globals = new GlobalsStore({ globals, globalTypes });
     this.hooks = {};
+    this.cleanupCallbacks = {};
 
     // We use a cache for these two functions for two reasons:
     //  1. For performance
@@ -234,8 +238,17 @@ export class StoryStore<TRenderer extends Renderer> {
     });
   }
 
-  cleanupStory(story: PreparedStory<TRenderer>): void {
+  addCleanupCallbacks(story: PreparedStory<TRenderer>, callbacks: CleanupCallback[]) {
+    this.cleanupCallbacks[story.id] = callbacks;
+  }
+
+  async cleanupStory(story: PreparedStory<TRenderer>): Promise<void> {
     this.hooks[story.id].clean();
+
+    const callbacks = this.cleanupCallbacks[story.id];
+    if (callbacks) for (const callback of [...callbacks].reverse()) await callback();
+
+    delete this.cleanupCallbacks[story.id];
   }
 
   extract(

code/lib/preview-api/src/modules/store/csf/portable-stories.ts

@@ -1,6 +1,6 @@
 /* eslint-disable no-underscore-dangle */
 /* eslint-disable @typescript-eslint/naming-convention */
-import { isExportStory } from '@storybook/csf';
+import { type CleanupCallback, isExportStory } from '@storybook/csf';
 import dedent from 'ts-dedent';
 import type {
   Renderer,
@@ -47,6 +47,8 @@ export function setProjectAnnotations<TRenderer extends Renderer = Renderer>(
   globalProjectAnnotations = composeConfigs(annotations.map(extractAnnotation));
 }
 
+const cleanupCallbacks: CleanupCallback[] = [];
+
 export function composeStory<TRenderer extends Renderer = Renderer, TArgs extends Args = Args>(
   storyAnnotations: LegacyStoryAnnotationsOrFn<TRenderer>,
   componentAnnotations: ComponentAnnotations<TRenderer, TArgs>,
@@ -126,8 +128,14 @@ export function composeStory<TRenderer extends Renderer = Renderer, TArgs extend
       id: story.id,
       storyName,
       load: async () => {
+        // First run any registerd cleanup function
+        for (const callback of [...cleanupCallbacks].reverse()) await callback();
+        cleanupCallbacks.length = 0;
+
         const loadedContext = await story.applyLoaders(context);
         context.loaded = loadedContext.loaded;
+
+        cleanupCallbacks.push(...(await story.applyBeforeEach(context)));
       },
       args: story.initialArgs as Partial<TArgs>,
       parameters: story.parameters as Parameters,

code/lib/preview-api/src/modules/store/csf/prepareStory.ts

@@ -20,7 +20,7 @@ import type {
   StoryContextForLoaders,
   StrictArgTypes,
 } from '@storybook/types';
-import { includeConditionalArg } from '@storybook/csf';
+import { type CleanupCallback, includeConditionalArg } from '@storybook/csf';
 
 import { applyHooks } from '../../addons';
 import { combineParameters } from '../parameters';
@@ -66,19 +66,20 @@ export function prepareStory<TRenderer extends Renderer>(
       updatedContext = { ...updatedContext, loaded: { ...updatedContext.loaded, ...loaded } };
     }
 
-    // TODO: What to do with those?
-    const cleanups = new Array<() => unknown>();
+    return updatedContext;
+  };
+
+  const applyBeforeEach = async (context: StoryContext<TRenderer>): Promise<CleanupCallback[]> => {
+    const cleanupCallbacks = new Array<() => unknown>();
     for (const beforeEach of [
       ...normalizeArrays(projectAnnotations.beforeEach),
       ...normalizeArrays(componentAnnotations.beforeEach),
       ...normalizeArrays(storyAnnotations.beforeEach),
     ]) {
-      const cleanup = await beforeEach(updatedContext);
-      if (cleanup) {
-        cleanups.push(cleanup);
-      }
+      const cleanup = await beforeEach(context);
+      if (cleanup) cleanupCallbacks.push(cleanup);
     }
-    return updatedContext;
+    return cleanupCallbacks;
   };
 
   const undecoratedStoryFn = (context: StoryContext<TRenderer>) =>
@@ -130,6 +131,7 @@ export function prepareStory<TRenderer extends Renderer>(
     undecoratedStoryFn,
     unboundStoryFn,
     applyLoaders,
+    applyBeforeEach,
     playFunction,
   };
 }

code/lib/source-loader/package.json

@@ -45,7 +45,7 @@
     "prep": "node --loader ../../../scripts/node_modules/esbuild-register/loader.js -r ../../../scripts/node_modules/esbuild-register/register.js ../../../scripts/prepare/bundle.ts"
   },
   "dependencies": {
-    "@storybook/csf": "0.1.4--canary.82.941efd8.0",
+    "@storybook/csf": "0.1.4--canary.82.5e73744.0",
     "@storybook/types": "workspace:*",
     "estraverse": "^5.2.0",
     "lodash": "^4.17.21",

code/lib/types/package.json

@@ -49,7 +49,7 @@
     "file-system-cache": "2.3.0"
   },
   "devDependencies": {
-    "@storybook/csf": "0.1.4--canary.82.941efd8.0",
+    "@storybook/csf": "0.1.4--canary.82.5e73744.0",
     "@types/fs-extra": "^11.0.1",
     "@types/node": "^18.0.0",
     "typescript": "^5.3.2"

code/lib/types/src/modules/story.ts

@@ -3,6 +3,7 @@ import type {
   ProjectAnnotations as CsfProjectAnnotations,
   DecoratorFunction,
   LoaderFunction,
+  CleanupCallback,
 } from '@storybook/csf';
 
 import type {
@@ -103,6 +104,7 @@ export type PreparedStory<TRenderer extends Renderer = Renderer> =
     applyLoaders: (
       context: StoryContextForLoaders<TRenderer>
     ) => Promise<StoryContextForLoaders<TRenderer> & { loaded: StoryContext<TRenderer>['loaded'] }>;
+    applyBeforeEach: (context: StoryContext<TRenderer>) => Promise<CleanupCallback[]>;
     playFunction?: (context: StoryContext<TRenderer>) => Promise<void> | void;
   };
 

code/package.json

@@ -127,7 +127,7 @@
     "@storybook/core-events": "workspace:*",
     "@storybook/core-server": "workspace:*",
     "@storybook/core-webpack": "workspace:*",
-    "@storybook/csf": "0.1.4--canary.82.941efd8.0",
+    "@storybook/csf": "0.1.4--canary.82.5e73744.0",
     "@storybook/csf-plugin": "workspace:*",
     "@storybook/csf-tools": "workspace:*",
     "@storybook/docs-tools": "workspace:*",

code/renderers/server/package.json

@@ -46,7 +46,7 @@
     "prep": "node --loader ../../../scripts/node_modules/esbuild-register/loader.js -r ../../../scripts/node_modules/esbuild-register/register.js ../../../scripts/prepare/bundle.ts"
   },
   "dependencies": {
-    "@storybook/csf": "0.1.4--canary.82.941efd8.0",
+    "@storybook/csf": "0.1.4--canary.82.5e73744.0",
     "@storybook/csf-tools": "workspace:*",
     "@storybook/global": "^5.0.0",
     "@storybook/preview-api": "workspace:*",

code/ui/blocks/package.json

@@ -48,7 +48,7 @@
     "@storybook/client-logger": "workspace:*",
     "@storybook/components": "workspace:*",
     "@storybook/core-events": "workspace:*",
-    "@storybook/csf": "0.1.4--canary.82.941efd8.0",
+    "@storybook/csf": "0.1.4--canary.82.5e73744.0",
     "@storybook/docs-tools": "workspace:*",
     "@storybook/global": "^5.0.0",
     "@storybook/icons": "^1.2.5",

code/ui/components/package.json

@@ -61,7 +61,7 @@
   "dependencies": {
     "@radix-ui/react-slot": "^1.0.2",
     "@storybook/client-logger": "workspace:*",
-    "@storybook/csf": "0.1.4--canary.82.941efd8.0",
+    "@storybook/csf": "0.1.4--canary.82.5e73744.0",
     "@storybook/global": "^5.0.0",
     "@storybook/icons": "^1.2.5",
     "@storybook/theming": "workspace:*",