Move prepareContext into store.storyContext()

This way any consumer of the context will get the properly mapped args. Also store the unmapped args on the context

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

@@ -380,7 +380,7 @@ export class PreviewWithSelection<TFramework extends Renderer> extends Preview<T
 
     if (isStoryRender(render)) {
       if (!render.story) throw new Error('Render has not been prepared!');
-      const { parameters, initialArgs, argTypes, args } = this.storyStore.getStoryContext(
+      const { parameters, initialArgs, argTypes, unmappedArgs } = this.storyStore.getStoryContext(
         render.story
       );
 
@@ -390,15 +390,15 @@ export class PreviewWithSelection<TFramework extends Renderer> extends Preview<T
           parameters,
           initialArgs,
           argTypes,
-          args,
+          args: unmappedArgs,
         });
       }
 
       // For v6 mode / compatibility
       // If the implementation changed, or args were persisted, the args may have changed,
       // and the STORY_PREPARED event above may not be respected.
       if (implementationChanged || persistedArgs) {
-        this.channel.emit(STORY_ARGS_UPDATED, { storyId, args });
+        this.channel.emit(STORY_ARGS_UPDATED, { storyId, args: unmappedArgs });
       }
     } else if (global.FEATURES?.storyStoreV7) {
       if (!this.storyStore.projectAnnotations) throw new Error('Store not initialized');

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

@@ -103,48 +103,4 @@ describe('StoryRender', () => {
     await render.renderToElement({} as any);
     expect(story.playFunction).not.toHaveBeenCalled();
   });
-
-  it('passes the initialArgs to loaders and render function if forceInitialArgs is true', async () => {
-    const story = {
-      id: 'id',
-      title: 'title',
-      name: 'name',
-      tags: [],
-      initialArgs: { a: 'b' },
-      applyLoaders: jest.fn(),
-      unboundStoryFn: jest.fn(),
-      playFunction: jest.fn(),
-      prepareContext: jest.fn((ctx) => ctx),
-    };
-
-    const renderToScreen = jest.fn();
-
-    const render = new StoryRender(
-      new Channel(),
-      { getStoryContext: () => ({ args: { a: 'c ' } }) } as any,
-      renderToScreen as any,
-      {} as any,
-      entry.id,
-      'story',
-      { forceInitialArgs: true },
-      story as any
-    );
-
-    await render.renderToElement({} as any);
-
-    expect(story.applyLoaders).toHaveBeenCalledWith(
-      expect.objectContaining({
-        args: { a: 'b' },
-      })
-    );
-
-    expect(renderToScreen).toHaveBeenCalledWith(
-      expect.objectContaining({
-        storyContext: expect.objectContaining({
-          args: { a: 'b' },
-        }),
-      }),
-      expect.any(Object)
-    );
-  });
 });

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

@@ -136,7 +136,8 @@ export class StoryRender<TRenderer extends Renderer> implements Render<TRenderer
 
   private storyContext() {
     if (!this.story) throw new Error(`Cannot call storyContext before preparing`);
-    return this.store.getStoryContext(this.story);
+    const { forceInitialArgs } = this.renderOptions;
+    return this.store.getStoryContext(this.story, { forceInitialArgs });
   }
 
   async render({
@@ -150,18 +151,8 @@ 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,
-      prepareContext,
-      initialArgs,
-    } = this.story;
+    const { id, componentId, title, name, tags, applyLoaders, unboundStoryFn, playFunction } =
+      this.story;
 
     if (forceRemount && !initial) {
       // NOTE: we don't check the cancel actually worked here, so the previous
@@ -176,16 +167,10 @@ export class StoryRender<TRenderer extends Renderer> implements Render<TRenderer
     const abortSignal = (this.abortController as AbortController).signal;
 
     try {
-      const getCurrentContext = () =>
-        prepareContext({
-          ...this.storyContext(),
-          ...(this.renderOptions.forceInitialArgs && { args: initialArgs }),
-        } as StoryContext);
-
       let loadedContext: Awaited<ReturnType<typeof applyLoaders>>;
       await this.runPhase(abortSignal, 'loading', async () => {
         loadedContext = await applyLoaders({
-          ...getCurrentContext(),
+          ...this.storyContext(),
           viewMode: this.viewMode,
         } as StoryContextForLoaders<TRenderer>);
       });
@@ -197,7 +182,7 @@ export class StoryRender<TRenderer extends Renderer> implements Render<TRenderer
         ...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
-        ...getCurrentContext(),
+        ...this.storyContext(),
         abortSignal,
         // We should consider parameterizing the story types with TRenderer['canvasElement'] in the future
         canvasElement: canvasElement as any,

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

@@ -8,6 +8,7 @@ import type { HooksContext } from './hooks';
 
 // Spy on prepareStory/processCSFFile
 jest.mock('./csf/prepareStory', () => ({
+  ...jest.requireActual('./csf/prepareStory'),
   prepareStory: jest.fn(jest.requireActual('./csf/prepareStory').prepareStory),
 }));
 jest.mock('./csf/processCSFFile', () => ({
@@ -425,6 +426,20 @@ describe('StoryStore', () => {
       });
     });
 
+    it('can force initial args', async () => {
+      const store = new StoryStore();
+      store.setProjectAnnotations(projectAnnotations);
+      store.initialize({ storyIndex, importFn, cache: false });
+
+      const story = await store.loadStory({ storyId: 'component-one--a' });
+
+      store.args.update(story.id, { foo: 'bar' });
+
+      expect(store.getStoryContext(story, { forceInitialArgs: true })).toMatchObject({
+        args: { foo: 'a' },
+      });
+    });
+
     it('returns the same hooks each time', async () => {
       const store = new StoryStore();
       store.setProjectAnnotations(projectAnnotations);
@@ -735,7 +750,6 @@ describe('StoryStore', () => {
               "fileName": "./src/ComponentOne.stories.js",
             },
             "playFunction": undefined,
-            "prepareContext": [Function],
             "story": "A",
             "storyFn": [Function],
             "subcomponents": undefined,
@@ -781,7 +795,6 @@ describe('StoryStore', () => {
               "fileName": "./src/ComponentOne.stories.js",
             },
             "playFunction": undefined,
-            "prepareContext": [Function],
             "story": "B",
             "storyFn": [Function],
             "subcomponents": undefined,
@@ -827,7 +840,6 @@ describe('StoryStore', () => {
               "fileName": "./src/ComponentTwo.stories.js",
             },
             "playFunction": undefined,
-            "prepareContext": [Function],
             "story": "C",
             "storyFn": [Function],
             "subcomponents": undefined,

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

@@ -30,7 +30,13 @@ import { HooksContext } from '../addons';
 import { StoryIndexStore } from './StoryIndexStore';
 import { ArgsStore } from './ArgsStore';
 import { GlobalsStore } from './GlobalsStore';
-import { processCSFFile, prepareStory, prepareMeta, normalizeProjectAnnotations } from './csf';
+import {
+  processCSFFile,
+  prepareStory,
+  prepareMeta,
+  normalizeProjectAnnotations,
+  prepareContext,
+} from './csf';
 
 const CSF_CACHE_SIZE = 1000;
 const STORY_CACHE_SIZE = 10000;
@@ -276,16 +282,17 @@ export class StoryStore<TRenderer extends Renderer> {
   // A prepared story does not include args, globals or hooks. These are stored in the story store
   // and updated separtely to the (immutable) story.
   getStoryContext(
-    story: PreparedStory<TRenderer>
+    story: PreparedStory<TRenderer>,
+    { forceInitialArgs = false } = {}
   ): Omit<StoryContextForLoaders<TRenderer>, 'viewMode'> {
     if (!this.globals) throw new Error(`getStoryContext called before initialization`);
 
-    return {
+    return prepareContext({
       ...story,
-      args: this.args.get(story.id),
+      args: forceInitialArgs ? story.initialArgs : this.args.get(story.id),
       globals: this.globals.get(),
       hooks: this.hooks[story.id] as unknown,
-    };
+    });
   }
 
   cleanupStory(story: PreparedStory<TRenderer>): void {

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

@@ -148,7 +148,10 @@ export const UNTARGETED = 'UNTARGETED';
 export function groupArgsByTarget<TArgs extends Args = Args>({
   args,
   argTypes,
-}: StoryContext<Renderer, TArgs>) {
+}: {
+  args: TArgs;
+  argTypes: ArgTypes<TArgs>;
+}) {
   const groupedArgs: Record<string, Partial<TArgs>> = {};
   (Object.entries(args) as [keyof TArgs, any][]).forEach(([name, value]) => {
     const { target = UNTARGETED } = (argTypes[name] || {}) as { target?: string };

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

@@ -2,11 +2,17 @@
 
 import { global } from '@storybook/global';
 import { expect } from '@jest/globals';
-import type { Renderer, ArgsEnhancer, PlayFunctionContext, SBScalarType } from '@storybook/types';
+import type {
+  Renderer,
+  ArgsEnhancer,
+  PlayFunctionContext,
+  SBScalarType,
+  StoryContextForLoaders,
+} from '@storybook/types';
 import { addons, HooksContext } from '../../addons';
 
 import { UNTARGETED } from '../args';
-import { prepareStory, prepareMeta } from './prepareStory';
+import { prepareStory, prepareMeta, prepareContext } from './prepareStory';
 
 jest.mock('@storybook/global', () => ({
   global: {
@@ -403,8 +409,8 @@ describe('prepareStory', () => {
         { render: renderMock }
       );
 
-      const context = story.prepareContext({ args: story.initialArgs, ...story } as any);
-      story.undecoratedStoryFn(context);
+      const context = prepareContext({ args: story.initialArgs, ...story } as any);
+      story.undecoratedStoryFn(context as any);
       expect(renderMock).toHaveBeenCalledWith(
         { one: 'mapped', two: 2, three: 3 },
         expect.objectContaining({ args: { one: 'mapped', two: 2, three: 3 } })
@@ -499,12 +505,12 @@ describe('prepareStory', () => {
       );
 
       const hooks = new HooksContext();
-      const context = story.prepareContext({ args: story.initialArgs, hooks, ...story } as any);
-      story.unboundStoryFn(context);
+      const context = prepareContext({ args: story.initialArgs, hooks, ...story } as any);
+      story.unboundStoryFn(context as any);
 
-      expect(ctx1).toMatchObject({ args: { one: 'mapped-1' } });
-      expect(ctx2).toMatchObject({ args: { one: 'mapped-1' } });
-      expect(ctx3).toMatchObject({ args: { one: 'mapped-1' } });
+      expect(ctx1).toMatchObject({ unmappedArgs: { one: 1 }, args: { one: 'mapped-1' } });
+      expect(ctx2).toMatchObject({ unmappedArgs: { one: 1 }, args: { one: 'mapped-1' } });
+      expect(ctx3).toMatchObject({ unmappedArgs: { one: 1 }, args: { one: 'mapped-1' } });
 
       hooks.clean();
     });
@@ -528,12 +534,12 @@ describe('prepareStory', () => {
         { render: renderMock }
       );
 
-      const context = firstStory.prepareContext({
+      const context = prepareContext({
         args: firstStory.initialArgs,
         hooks: new HooksContext(),
         ...firstStory,
       } as any);
-      firstStory.unboundStoryFn(context);
+      firstStory.unboundStoryFn(context as any);
       expect(renderMock).toHaveBeenCalledWith(
         { a: 1 },
         expect.objectContaining({ args: { a: 1 }, allArgs: { a: 1, b: 2 } })
@@ -554,12 +560,12 @@ describe('prepareStory', () => {
         { render: renderMock }
       );
 
-      const context = firstStory.prepareContext({
+      const context = prepareContext({
         args: firstStory.initialArgs,
         hooks: new HooksContext(),
         ...firstStory,
       } as any);
-      firstStory.unboundStoryFn(context);
+      firstStory.unboundStoryFn(context as any);
       expect(renderMock).toHaveBeenCalledWith(
         { a: 1 },
         expect.objectContaining({ args: { a: 1 }, allArgs: { a: 1, b: 2 } })
@@ -580,12 +586,12 @@ describe('prepareStory', () => {
         { render: renderMock }
       );
 
-      const context = firstStory.prepareContext({
+      const context = prepareContext({
         args: firstStory.initialArgs,
         hooks: new HooksContext(),
         ...firstStory,
       } as any);
-      firstStory.unboundStoryFn(context);
+      firstStory.unboundStoryFn(context as any);
       expect(renderMock).toHaveBeenCalledWith(
         { a: 1 },
         expect.objectContaining({ argsByTarget: { [UNTARGETED]: { a: 1 }, foo: { b: 2 } } })
@@ -606,12 +612,12 @@ describe('prepareStory', () => {
         { render: renderMock }
       );
 
-      const context = firstStory.prepareContext({
+      const context = prepareContext({
         args: firstStory.initialArgs,
         hooks: new HooksContext(),
         ...firstStory,
       } as any);
-      firstStory.unboundStoryFn(context);
+      firstStory.unboundStoryFn(context as any);
       expect(renderMock).toHaveBeenCalledWith(
         {},
         expect.objectContaining({ argsByTarget: { foo: { b: 2 } } })
@@ -630,12 +636,12 @@ describe('prepareStory', () => {
         { render: renderMock }
       );
 
-      const context = firstStory.prepareContext({
+      const context = prepareContext({
         args: firstStory.initialArgs,
         hooks: new HooksContext(),
         ...firstStory,
       } as any);
-      firstStory.unboundStoryFn(context);
+      firstStory.unboundStoryFn(context as any);
       expect(renderMock).toHaveBeenCalledWith({}, expect.objectContaining({ argsByTarget: {} }));
     });
   });
@@ -723,7 +729,6 @@ describe('prepareMeta', () => {
       unboundStoryFn,
       undecoratedStoryFn,
       playFunction,
-      prepareContext,
       // eslint-disable-next-line @typescript-eslint/naming-convention
       parameters: { __isArgsStory, ...parameters },
       ...expectedPreparedMeta