From 243efbb7e3f551d6f1b62c7314d23d27d0db839f Mon Sep 17 00:00:00 2001
From: Tom Coleman <tom@thesnail.org>
Date: Fri, 23 Apr 2021 13:57:36 +1000
Subject: [PATCH] Bind the context *inside* the  function

---
 lib/addons/src/types.ts           | 14 +++++--
 lib/client-api/src/decorators.ts  | 64 +++++++++++++++++++------------
 lib/client-api/src/story_store.ts |  4 +-
 3 files changed, 52 insertions(+), 30 deletions(-)

diff --git a/lib/addons/src/types.ts b/lib/addons/src/types.ts
index 445e8b5731c..7d9e2e6e737 100644
--- a/lib/addons/src/types.ts
+++ b/lib/addons/src/types.ts
@@ -53,6 +53,7 @@ export interface StoryIdentifier {
   name: StoryName;
 }
 
+export type StoryContextUpdate = Partial<StoryContext>;
 export type StoryContext = StoryIdentifier & {
   [key: string]: any;
   parameters: Parameters;
@@ -93,8 +94,13 @@ export interface OptionsParameter extends Object {
 
 export type StoryGetter = (context: StoryContext) => any;
 
+// This is the type of story function passed to a decorator -- does not rely on being passed any context
+export type PartialStoryFn<ReturnType = unknown> = (p?: StoryContextUpdate) => ReturnType;
+// This is a passArgsFirst: false user story function
 export type LegacyStoryFn<ReturnType = unknown> = (p?: StoryContext) => ReturnType;
+// This is a passArgsFirst: true user story function
 export type ArgsStoryFn<ReturnType = unknown> = (a?: Args, p?: StoryContext) => ReturnType;
+// This is either type of user story function
 export type StoryFn<ReturnType = unknown> = LegacyStoryFn<ReturnType> | ArgsStoryFn<ReturnType>;
 
 export type StoryWrapper = (
@@ -136,16 +142,16 @@ export interface StoryApi<StoryFnReturnType = unknown> {
 }
 
 export type DecoratorFunction<StoryFnReturnType = unknown> = (
-  fn: StoryFn<StoryFnReturnType>,
+  fn: PartialStoryFn<StoryFnReturnType>,
   c: StoryContext
-) => ReturnType<StoryFn<StoryFnReturnType>>;
+) => ReturnType<LegacyStoryFn<StoryFnReturnType>>;
 
 export type LoaderFunction = (c: StoryContext) => Promise<Record<string, any>>;
 
 export type DecorateStoryFunction<StoryFnReturnType = unknown> = (
-  storyFn: StoryFn<StoryFnReturnType>,
+  storyFn: LegacyStoryFn<StoryFnReturnType>,
   decorators: DecoratorFunction<StoryFnReturnType>[]
-) => StoryFn<StoryFnReturnType>;
+) => LegacyStoryFn<StoryFnReturnType>;
 
 export interface ClientStoryApi<StoryFnReturnType = unknown> {
   storiesOf(kind: StoryKind, module: NodeModule): StoryApi<StoryFnReturnType>;
diff --git a/lib/client-api/src/decorators.ts b/lib/client-api/src/decorators.ts
index 8e51b7d580d..eef4b9bb89a 100644
--- a/lib/client-api/src/decorators.ts
+++ b/lib/client-api/src/decorators.ts
@@ -1,20 +1,6 @@
-import { StoryContext, StoryFn } from '@storybook/addons';
+import { StoryContext, StoryContextUpdate, PartialStoryFn, LegacyStoryFn } from '@storybook/addons';
 import { DecoratorFunction } from './types';
 
-interface StoryContextUpdate {
-  [key: string]: any;
-}
-
-const defaultContext: StoryContext = {
-  id: 'unspecified',
-  name: 'unspecified',
-  kind: 'unspecified',
-  parameters: {},
-  args: {},
-  argTypes: {},
-  globals: {},
-};
-
 /**
  * When you call the story function inside a decorator, e.g.:
  *
@@ -25,15 +11,43 @@ const defaultContext: StoryContext = {
  * This will override the `foo` property on the `innerContext`, which gets
  * merged in with the default context
  */
-export const decorateStory = (storyFn: StoryFn, decorator: DecoratorFunction) => {
-  return (context: StoryContext = defaultContext) =>
-    decorator(
-      // You cannot override the parameters key, it is fixed
-      ({ parameters, ...innerContext }: StoryContextUpdate = {}) =>
-        storyFn({ ...context, ...innerContext }),
-      context
-    );
+const bindWithContext = (
+  storyFn: LegacyStoryFn,
+  getStoryContext: () => StoryContext
+): PartialStoryFn =>
+  // (NOTE: You cannot override the parameters key, it is fixed)
+  ({ parameters, ...innerContext }: StoryContextUpdate = {}) =>
+    storyFn({ ...getStoryContext(), ...innerContext });
+
+export const decorateStory = (
+  storyFn: LegacyStoryFn,
+  decorator: DecoratorFunction,
+  getStoryContext: () => StoryContext
+): LegacyStoryFn => {
+  // Bind the partially decorated storyFn so that when it is called it always knows about the story context,
+  // no matter what it is passed directly. This is because we cannot guarantee a decorator will
+  // pass the context down to the next decorator in the chain.
+  const boundStoryFunction = bindWithContext(storyFn, getStoryContext);
+
+  return (context: StoryContext) => decorator(boundStoryFunction, context);
 };
 
-export const defaultDecorateStory = (storyFn: StoryFn, decorators: DecoratorFunction[]) =>
-  decorators.reduce(decorateStory, storyFn);
+export const defaultDecorateStory = (
+  storyFn: LegacyStoryFn,
+  decorators: DecoratorFunction[]
+): LegacyStoryFn => {
+  // We use a trick to avoid recreating the bound story function inside `decorateStory`.
+  // Instead we pass it a context "getter", which is defined once (at "decoration time")
+  // The getter reads a variable which is scoped to this call of `decorateStory`
+  // (ie to this story), so there is no possibility of overlap.
+  // This will break if you call the same story twice interleaved.
+  let contextStore: StoryContext;
+  const decoratedWithContextStore = decorators.reduce(
+    (story, decorator) => decorateStory(story, decorator, () => contextStore),
+    storyFn
+  );
+  return (context) => {
+    contextStore = context;
+    return decoratedWithContextStore(context);
+  };
+};
diff --git a/lib/client-api/src/story_store.ts b/lib/client-api/src/story_store.ts
index 1c85fcd2ae5..12a6f175dbe 100644
--- a/lib/client-api/src/story_store.ts
+++ b/lib/client-api/src/story_store.ts
@@ -394,7 +394,9 @@ export default class StoryStore {
           return acc;
         }, {} as Args),
       };
-      return passArgsFirst ? (original as ArgsStoryFn)(mapped.args, mapped) : original(mapped);
+      return passArgsFirst
+        ? (original as ArgsStoryFn)(mapped.args, mapped)
+        : (original as LegacyStoryFn)(mapped);
     };
 
     // lazily decorate the story when it's loaded