Implemented prepared context.

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

@@ -59,6 +59,7 @@ describe('StoryRender', () => {
       applyLoaders: jest.fn(),
       unboundStoryFn: jest.fn(),
       playFunction: jest.fn(),
+      prepareContext: jest.fn(),
     };
 
     const render = new StoryRender(
@@ -85,6 +86,7 @@ describe('StoryRender', () => {
       applyLoaders: jest.fn(),
       unboundStoryFn: jest.fn(),
       playFunction: jest.fn(),
+      prepareContext: jest.fn(),
     };
 
     const render = new StoryRender(

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

@@ -154,8 +154,17 @@ export class StoryRender<TRenderer extends Renderer> implements Render<TRenderer
     if (!this.story) throw new Error('cannot render when not prepared');
     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,
+      unboundStoryFn,
+      playFunction,
+      prepareContext,
+    } = this.story;
 
     if (forceRemount && !initial) {
       // NOTE: we don't check the cancel actually worked here, so the previous
@@ -181,7 +190,7 @@ export class StoryRender<TRenderer extends Renderer> implements Render<TRenderer
         return;
       }
 
-      const renderStoryContext: StoryContext<TRenderer> = {
+      const renderStoryContext: StoryContext<TRenderer> = prepareContext({
         ...loadedContext!,
         // By this stage, it is possible that new args/globals have been received for this story
         // and we need to ensure we render it with the new values
@@ -189,7 +198,7 @@ export class StoryRender<TRenderer extends Renderer> implements Render<TRenderer
         abortSignal,
         // We should consider parameterizing the story types with TRenderer['canvasElement'] in the future
         canvasElement: canvasElement as any,
-      };
+      });
       const renderContext: RenderContext<TRenderer> = {
         componentId,
         title,

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

@@ -468,8 +468,8 @@ describe('prepareStory', () => {
         { render: renderMock }
       );
 
-      const context = { args: story.initialArgs, ...story };
-      story.undecoratedStoryFn(context as any);
+      const context = story.prepareContext({ args: story.initialArgs, ...story } as any);
+      story.undecoratedStoryFn(context);
       expect(renderMock).toHaveBeenCalledWith(
         { one: 'mapped', two: 2, three: 3 },
         expect.objectContaining({ args: { one: 'mapped', two: 2, three: 3 } })
@@ -529,6 +529,50 @@ describe('prepareStory', () => {
 
       hooks.clean();
     });
+
+    it('prepared context is applied to decorators', () => {
+      const renderMock = jest.fn();
+      let ctx1;
+      let ctx2;
+      let ctx3;
+
+      const globalDecorator = jest.fn((fn, ctx) => {
+        ctx1 = ctx;
+        return fn();
+      });
+      const componentDecorator = jest.fn((fn, ctx) => {
+        ctx2 = ctx;
+        return fn();
+      });
+      const storyDecorator = jest.fn((fn, ctx) => {
+        ctx3 = ctx;
+        return fn();
+      });
+      const story = prepareStory(
+        {
+          id,
+          name,
+          argTypes: {
+            one: { name: 'one', type: { name: 'string' }, mapping: { 1: 'mapped-1' } },
+          },
+          args: { one: 1 },
+          decorators: [storyDecorator],
+          moduleExport,
+        },
+        { id, title, decorators: [componentDecorator] },
+        { render: renderMock, decorators: [globalDecorator] }
+      );
+
+      const hooks = new HooksContext();
+      const context = story.prepareContext({ args: story.initialArgs, hooks, ...story } as any);
+      story.unboundStoryFn(context);
+
+      expect(ctx1).toMatchObject({ args: { one: 'mapped-1' } });
+      expect(ctx2).toMatchObject({ args: { one: 'mapped-1' } });
+      expect(ctx3).toMatchObject({ args: { one: 'mapped-1' } });
+
+      hooks.clean();
+    });
   });
 
   describe('with `FEATURES.argTypeTargetsV7`', () => {
@@ -549,11 +593,12 @@ describe('prepareStory', () => {
         { render: renderMock }
       );
 
-      firstStory.unboundStoryFn({
+      const context = firstStory.prepareContext({
         args: firstStory.initialArgs,
         hooks: new HooksContext(),
         ...firstStory,
       } as any);
+      firstStory.unboundStoryFn(context);
       expect(renderMock).toHaveBeenCalledWith(
         { a: 1 },
         expect.objectContaining({ args: { a: 1 }, allArgs: { a: 1, b: 2 } })
@@ -574,11 +619,12 @@ describe('prepareStory', () => {
         { render: renderMock }
       );
 
-      firstStory.unboundStoryFn({
+      const context = firstStory.prepareContext({
         args: firstStory.initialArgs,
         hooks: new HooksContext(),
         ...firstStory,
       } as any);
+      firstStory.unboundStoryFn(context);
       expect(renderMock).toHaveBeenCalledWith(
         { a: 1 },
         expect.objectContaining({ args: { a: 1 }, allArgs: { a: 1, b: 2 } })
@@ -599,11 +645,12 @@ describe('prepareStory', () => {
         { render: renderMock }
       );
 
-      firstStory.unboundStoryFn({
+      const context = firstStory.prepareContext({
         args: firstStory.initialArgs,
         hooks: new HooksContext(),
         ...firstStory,
       } as any);
+      firstStory.unboundStoryFn(context);
       expect(renderMock).toHaveBeenCalledWith(
         { a: 1 },
         expect.objectContaining({ argsByTarget: { [NO_TARGET_NAME]: { a: 1 }, foo: { b: 2 } } })
@@ -624,11 +671,12 @@ describe('prepareStory', () => {
         { render: renderMock }
       );
 
-      firstStory.unboundStoryFn({
+      const context = firstStory.prepareContext({
         args: firstStory.initialArgs,
         hooks: new HooksContext(),
         ...firstStory,
       } as any);
+      firstStory.unboundStoryFn(context);
       expect(renderMock).toHaveBeenCalledWith(
         {},
         expect.objectContaining({ argsByTarget: { foo: { b: 2 } } })
@@ -647,11 +695,12 @@ describe('prepareStory', () => {
         { render: renderMock }
       );
 
-      firstStory.unboundStoryFn({
+      const context = firstStory.prepareContext({
         args: firstStory.initialArgs,
         hooks: new HooksContext(),
         ...firstStory,
       } as any);
+      firstStory.unboundStoryFn(context);
       expect(renderMock).toHaveBeenCalledWith({}, expect.objectContaining({ argsByTarget: {} }));
     });
   });
@@ -739,6 +788,7 @@ describe('prepareMeta', () => {
       unboundStoryFn,
       undecoratedStoryFn,
       playFunction,
+      prepareContext,
       ...expectedPreparedMeta
     } = preparedStory;
 

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

@@ -71,23 +71,10 @@ export function prepareStory<TRenderer extends Renderer>(
   };
 
   const undecoratedStoryFn: LegacyStoryFn<TRenderer> = (context: StoryContext<TRenderer>) => {
-    const mappedArgs = Object.entries(context.args).reduce((acc, [key, val]) => {
-      const mapping = context.argTypes[key]?.mapping;
-      acc[key] = mapping && val in mapping ? mapping[val] : val;
-      return acc;
-    }, {} as Args);
-
-    const includedArgs = Object.entries(mappedArgs).reduce((acc, [key, val]) => {
-      const argType = context.argTypes[key] || {};
-      if (includeConditionalArg(argType, mappedArgs, context.globals)) acc[key] = val;
-      return acc;
-    }, {} as Args);
-
-    const includedContext = { ...context, args: includedArgs };
     const { passArgsFirst: renderTimePassArgsFirst = true } = context.parameters;
     return renderTimePassArgsFirst
-      ? (render as ArgsStoryFn<TRenderer>)(includedContext.args, includedContext)
-      : (render as LegacyStoryFn<TRenderer>)(includedContext);
+      ? (render as ArgsStoryFn<TRenderer>)(context.args, context)
+      : (render as LegacyStoryFn<TRenderer>)(context);
   };
 
   // Currently it is only possible to set these globally
@@ -109,8 +96,15 @@ export function prepareStory<TRenderer extends Renderer>(
   if (!render) throw new Error(`No render function available for storyId '${id}'`);
 
   const decoratedStoryFn = applyHooks<TRenderer>(applyDecorators)(undecoratedStoryFn, decorators);
-  const unboundStoryFn = (context: StoryContext<TRenderer>) => {
+  const unboundStoryFn = (context: StoryContext<TRenderer>) => decoratedStoryFn(context);
+
+  // prepareContext is invoked at StoryRender.render()
+  // the context is prepared before invoking the render function, instead of here directly
+  // to ensure args don't loose there special properties set by the renderer
+  // eg. reactive proxies set by frameworks like SolidJS or Vue
+  const prepareContext = (context: StoryContext<TRenderer>) => {
     let finalContext: StoryContext<TRenderer> = context;
+
     if (global.FEATURES?.argTypeTargetsV7) {
       const argsByTarget = groupArgsByTarget(context);
       finalContext = {
@@ -121,7 +115,19 @@ export function prepareStory<TRenderer extends Renderer>(
       };
     }
 
-    return decoratedStoryFn(finalContext);
+    const mappedArgs = Object.entries(finalContext.args).reduce((acc, [key, val]) => {
+      const mapping = finalContext.argTypes[key]?.mapping;
+      acc[key] = mapping && val in mapping ? mapping[val] : val;
+      return acc;
+    }, {} as Args);
+
+    const includedArgs = Object.entries(mappedArgs).reduce((acc, [key, val]) => {
+      const argType = finalContext.argTypes[key] || {};
+      if (includeConditionalArg(argType, mappedArgs, finalContext.globals)) acc[key] = val;
+      return acc;
+    }, {} as Args);
+
+    return { ...finalContext, args: includedArgs };
   };
 
   const play = storyAnnotations?.play || componentAnnotations.play;
@@ -150,6 +156,7 @@ export function prepareStory<TRenderer extends Renderer>(
     unboundStoryFn,
     applyLoaders,
     playFunction,
+    prepareContext,
   };
 }
 

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

@@ -84,6 +84,7 @@ export type PreparedStory<TRenderer extends Renderer = Renderer> =
       context: StoryContextForLoaders<TRenderer>
     ) => Promise<StoryContextForLoaders<TRenderer> & { loaded: StoryContext<TRenderer>['loaded'] }>;
     playFunction?: (context: StoryContext<TRenderer>) => Promise<void> | void;
+    prepareContext: (context: StoryContext<TRenderer>) => StoryContext<TRenderer>;
   };
 
 export type PreparedMeta<TRenderer extends Renderer = Renderer> = Omit<