Move prepareContext into store.storyContext()

This way any consumer of the context will get the properly mapped args.

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

@@ -136,7 +136,10 @@ 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);
+    return {
+      ...this.store.getStoryContext(this.story),
+      ...(this.renderOptions.forceInitialArgs && { args: this.story.initialArgs }),
+    };
   }
 
   async render({
@@ -150,18 +153,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 +169,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 +184,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.ts

@@ -29,7 +29,7 @@ import { HooksContext } from '../addons';
 import { StoryIndexStore } from './StoryIndexStore';
 import { ArgsStore } from './ArgsStore';
 import { GlobalsStore } from './GlobalsStore';
-import { processCSFFile, prepareStory, normalizeProjectAnnotations } from './csf';
+import { processCSFFile, prepareStory, normalizeProjectAnnotations, prepareContext } from './csf';
 
 const CSF_CACHE_SIZE = 1000;
 const STORY_CACHE_SIZE = 10000;
@@ -264,12 +264,12 @@ export class StoryStore<TRenderer extends Renderer> {
   ): Omit<StoryContextForLoaders<TRenderer>, 'viewMode'> {
     if (!this.globals) throw new Error(`getStoryContext called before initialization`);
 
-    return {
+    return prepareContext({
       ...story,
       args: 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.ts

@@ -20,6 +20,7 @@ import type {
   PreparedMeta,
   ModuleExport,
 } from '@storybook/types';
+import type { StoryContextUpdate } from '@storybook/csf';
 import { includeConditionalArg } from '@storybook/csf';
 
 import { applyHooks } from '../../addons';
@@ -87,38 +88,6 @@ export function prepareStory<TRenderer extends Renderer>(
   const decoratedStoryFn = applyHooks<TRenderer>(applyDecorators)(undecoratedStoryFn, decorators);
   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 = {
-        ...context,
-        allArgs: context.args,
-        argsByTarget,
-        args: argsByTarget[UNTARGETED] || {},
-      };
-    }
-
-    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;
 
   const playFunction =
@@ -145,7 +114,6 @@ export function prepareStory<TRenderer extends Renderer>(
     unboundStoryFn,
     applyLoaders,
     playFunction,
-    prepareContext,
   };
 }
 
@@ -247,3 +215,37 @@ function preparePartialAnnotations<TRenderer extends Renderer>(
 
   return withoutStoryIdentifiers;
 }
+
+// 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
+export function prepareContext<TRenderer extends Renderer>(
+  // We use this slightly weird type because `Omit<..., 'viewMode'>` breaks the type somehow
+  context: StoryContextForEnhancers<TRenderer> & Required<StoryContextUpdate>
+) {
+  let finalContext = context;
+
+  if (global.FEATURES?.argTypeTargetsV7) {
+    const argsByTarget = groupArgsByTarget(context);
+    finalContext = {
+      ...context,
+      allArgs: context.args,
+      argsByTarget,
+      args: argsByTarget[UNTARGETED] || {},
+    };
+  }
+
+  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 };
+}