From 11e0386e98f4e49afc7b0cbc8360850d28ed83e7 Mon Sep 17 00:00:00 2001
From: atanasster <atanasster@gmail.com>
Date: Sun, 21 Jun 2020 21:04:12 -0400
Subject: [PATCH] fix: storybook beta breaking SET_CURRENT_STORY

---
 .../src/components/ManagerContainer.tsx        | 18 +++++++++++++++---
 misc/storybook-custom-docs/src/index.tsx       |  6 +++---
 misc/storybook-custom-docs/src/types.ts        |  2 ++
 ui/blocks/src/context/block/BlockContext.tsx   |  1 -
 4 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/misc/storybook-custom-docs/src/components/ManagerContainer.tsx b/misc/storybook-custom-docs/src/components/ManagerContainer.tsx
index c82bf30a9..007fb7269 100644
--- a/misc/storybook-custom-docs/src/components/ManagerContainer.tsx
+++ b/misc/storybook-custom-docs/src/components/ManagerContainer.tsx
@@ -1,5 +1,6 @@
 import * as React from 'react';
-import { API } from '@storybook/api';
+import { API, Consumer, Combo } from '@storybook/api';
+import { UPDATE_STORY_CONTEXT } from '../types';
 
 interface ManagerContainerProps {
   active?: boolean;
@@ -11,7 +12,7 @@ interface ManagerContainerProps {
 export const ManagerContainer: React.FC<ManagerContainerProps> = props => {
   const { active, api, route } = props;
   const channel = React.useMemo(() => api.getChannel(), [api]);
-
+  const storyRef = React.useRef('');
   React.useEffect(() => {
     const ATTACH_DOCS_PAGE = `attach_docs_page_${route}`;
 
@@ -29,5 +30,16 @@ export const ManagerContainer: React.FC<ManagerContainerProps> = props => {
     };
     updateDOM();
   }, [active, api, route, channel]);
-  return null;
+
+  return (
+    <Consumer>
+      {({ state }: Combo) => {
+        if (state.storyId !== storyRef.current) {
+          storyRef.current = state.storyId;
+          channel.emit(UPDATE_STORY_CONTEXT, state);
+        }
+        return null;
+      }}
+    </Consumer>
+  );
 };
diff --git a/misc/storybook-custom-docs/src/index.tsx b/misc/storybook-custom-docs/src/index.tsx
index fdf322866..f7456c013 100644
--- a/misc/storybook-custom-docs/src/index.tsx
+++ b/misc/storybook-custom-docs/src/index.tsx
@@ -1,7 +1,7 @@
 import * as React from 'react';
 import addons from '@storybook/addons';
-import { SET_CURRENT_STORY } from '@storybook/core-events';
 import { ConfigApi } from '@storybook/client-api';
+import { UPDATE_STORY_CONTEXT } from './types';
 export * from './types';
 
 /**
@@ -49,10 +49,10 @@ export const useStoryId = () => {
     const onStoryChange = ({ storyId: id }: { storyId: string }) => {
       setStoryId(id);
     };
-    channel.on(SET_CURRENT_STORY, onStoryChange);
+    channel.on(UPDATE_STORY_CONTEXT, onStoryChange);
 
     return () => {
-      channel.off(SET_CURRENT_STORY, onStoryChange);
+      channel.off(UPDATE_STORY_CONTEXT, onStoryChange);
     };
   }, [channel]);
   return storyId;
diff --git a/misc/storybook-custom-docs/src/types.ts b/misc/storybook-custom-docs/src/types.ts
index c613877bc..ad1a34863 100644
--- a/misc/storybook-custom-docs/src/types.ts
+++ b/misc/storybook-custom-docs/src/types.ts
@@ -47,3 +47,5 @@ export interface CustomPageDef {
   // Return an object that can be rendered from ReactDOM.render
   render: CustomPageRenderFn;
 }
+
+export const UPDATE_STORY_CONTEXT = `update_story_context`;
diff --git a/ui/blocks/src/context/block/BlockContext.tsx b/ui/blocks/src/context/block/BlockContext.tsx
index df320f8a4..9a5f54079 100644
--- a/ui/blocks/src/context/block/BlockContext.tsx
+++ b/ui/blocks/src/context/block/BlockContext.tsx
@@ -57,7 +57,6 @@ export const BlockContextProvider: React.FC<BlockContextInputProps> = ({
   options,
 }) => {
   let storyId = propsStoryId;
-  console.log(propsStoryId, propsDocId);
   const docId = storyId || propsDocId;
   if (!storyId && docId) {
     const doc = store.getStoryDoc(docId);