From 00e3c13898a613cd0f95bb3fd858800195e6019c Mon Sep 17 00:00:00 2001
From: Tom Coleman <tom@thesnail.org>
Date: Fri, 10 Sep 2021 18:45:14 +1000
Subject: [PATCH 1/2] Docs changes for on demand store

---
 addons/docs/package.json                      |  10 +-
 addons/docs/src/blocks/ArgsTable.tsx          | 102 ++++++++++--------
 addons/docs/src/blocks/Canvas.tsx             |   8 +-
 addons/docs/src/blocks/Description.tsx        |   7 +-
 addons/docs/src/blocks/DocsContainer.tsx      |  16 ++-
 addons/docs/src/blocks/DocsContext.ts         |  20 +---
 addons/docs/src/blocks/DocsPage.test.ts       |   7 +-
 addons/docs/src/blocks/Meta.tsx               |  10 +-
 addons/docs/src/blocks/Primary.tsx            |   9 +-
 addons/docs/src/blocks/Source.tsx             |  44 +++-----
 addons/docs/src/blocks/Stories.tsx            |   6 +-
 addons/docs/src/blocks/Story.tsx              |  85 +++++++++------
 addons/docs/src/blocks/Subtitle.tsx           |   4 +-
 addons/docs/src/blocks/Title.tsx              |   6 +-
 addons/docs/src/blocks/enhanceSource.ts       |   9 +-
 addons/docs/src/blocks/mdx.tsx                |   2 +-
 addons/docs/src/blocks/useStory.ts            |  18 ++++
 addons/docs/src/blocks/utils.ts               |  15 +--
 .../docs/src/frameworks/angular/compodoc.ts   |   4 +-
 .../frameworks/angular/prepareForInline.ts    |   9 +-
 .../src/frameworks/angular/sourceDecorator.ts |  14 +--
 .../frameworks/common/enhanceArgTypes.test.ts |  41 ++++---
 .../src/frameworks/common/enhanceArgTypes.ts  |  21 ++--
 .../frameworks/common/normalizeArgTypes.ts    |  18 ----
 .../src/frameworks/html/prepareForInline.tsx  |   4 +-
 .../src/frameworks/html/sourceDecorator.ts    |  16 ++-
 addons/docs/src/frameworks/react/config.ts    |   6 +-
 .../src/frameworks/react/extractArgTypes.ts   |   4 +-
 .../src/frameworks/react/jsxDecorator.tsx     |  21 ++--
 .../react/react-argtypes.stories.tsx          |   4 +-
 .../frameworks/react/react-properties.test.ts |  10 +-
 .../frameworks/svelte/extractArgTypes.test.ts |   9 +-
 .../src/frameworks/svelte/extractArgTypes.ts  |  10 +-
 .../src/frameworks/svelte/prepareForInline.ts |   4 +-
 .../src/frameworks/svelte/sourceDecorator.ts  |   8 +-
 .../src/frameworks/vue/extractArgTypes.ts     |   4 +-
 .../src/frameworks/vue/prepareForInline.ts    |   8 +-
 .../src/frameworks/vue/sourceDecorator.ts     |   8 +-
 .../src/frameworks/vue3/extractArgTypes.ts    |   4 +-
 .../src/frameworks/vue3/prepareForInline.ts   |   9 +-
 .../web-components/prepareForInline.ts        |   6 +-
 .../web-components/sourceDecorator.ts         |  19 +++-
 addons/docs/src/lib/convert/flow/convert.ts   |   2 +-
 .../docs/src/lib/convert/proptypes/convert.ts |   2 +-
 .../src/lib/convert/typescript/convert.ts     |   2 +-
 addons/docs/src/lib/docgen/types.ts           |   4 +-
 46 files changed, 356 insertions(+), 293 deletions(-)
 create mode 100644 addons/docs/src/blocks/useStory.ts
 delete mode 100644 addons/docs/src/frameworks/common/normalizeArgTypes.ts

diff --git a/addons/docs/package.json b/addons/docs/package.json
index 057532365d3c..dfb57e93baad 100644
--- a/addons/docs/package.json
+++ b/addons/docs/package.json
@@ -71,11 +71,13 @@
     "@storybook/components": "6.4.0-alpha.34",
     "@storybook/core": "6.4.0-alpha.34",
     "@storybook/core-events": "6.4.0-alpha.34",
-    "@storybook/csf": "0.0.1",
+    "@storybook/csf": "0.0.2--canary.b1d5348.0",
     "@storybook/csf-tools": "6.4.0-alpha.34",
     "@storybook/node-logger": "6.4.0-alpha.34",
     "@storybook/postinstall": "6.4.0-alpha.34",
+    "@storybook/preview-web": "6.4.0-alpha.34",
     "@storybook/source-loader": "6.4.0-alpha.34",
+    "@storybook/store": "6.4.0-alpha.34",
     "@storybook/theming": "6.4.0-alpha.34",
     "acorn": "^7.4.1",
     "acorn-jsx": "^5.3.1",
@@ -106,6 +108,7 @@
     "@emotion/core": "^10.1.1",
     "@emotion/styled": "^10.0.27",
     "@storybook/angular": "6.4.0-alpha.34",
+    "@storybook/html": "6.4.0-alpha.34",
     "@storybook/react": "6.4.0-alpha.34",
     "@storybook/vue": "6.4.0-alpha.34",
     "@storybook/web-components": "6.4.0-alpha.34",
@@ -139,6 +142,8 @@
   },
   "peerDependencies": {
     "@storybook/angular": "6.4.0-alpha.34",
+    "@storybook/html": "6.4.0-alpha.34",
+    "@storybook/react": "6.4.0-alpha.34",
     "@storybook/vue": "6.4.0-alpha.34",
     "@storybook/vue3": "6.4.0-alpha.34",
     "@storybook/web-components": "6.4.0-alpha.34",
@@ -155,6 +160,9 @@
     "@storybook/angular": {
       "optional": true
     },
+    "@storybook/react": {
+      "optional": true
+    },
     "@storybook/vue": {
       "optional": true
     },
diff --git a/addons/docs/src/blocks/ArgsTable.tsx b/addons/docs/src/blocks/ArgsTable.tsx
index fe8702db0423..7db03289d8ef 100644
--- a/addons/docs/src/blocks/ArgsTable.tsx
+++ b/addons/docs/src/blocks/ArgsTable.tsx
@@ -1,24 +1,23 @@
-/* eslint-disable no-underscore-dangle */
 import React, { FC, useContext, useEffect, useState, useCallback } from 'react';
 import mapValues from 'lodash/mapValues';
 import {
   ArgsTable as PureArgsTable,
   ArgsTableProps as PureArgsTableProps,
   ArgsTableError,
-  ArgTypes,
   SortType,
   TabbedArgsTable,
 } from '@storybook/components';
-import { Args } from '@storybook/addons';
-import { StoryStore, filterArgTypes } from '@storybook/client-api';
-import type { PropDescriptor } from '@storybook/client-api';
+import { addons } from '@storybook/addons';
+import { filterArgTypes, PropDescriptor } from '@storybook/store';
 import Events from '@storybook/core-events';
+import { StrictArgTypes, Args } from '@storybook/csf';
 
 import { DocsContext, DocsContextProps } from './DocsContext';
 import { Component, CURRENT_SELECTION, PRIMARY_STORY } from './types';
-import { getComponentName, getDocsStories } from './utils';
+import { getComponentName } from './utils';
 import { ArgTypesExtractor } from '../lib/docgen/types';
 import { lookupStoryId } from './Story';
+import { useStory } from './useStory';
 
 interface BaseProps {
   include?: PropDescriptor;
@@ -45,29 +44,33 @@ type ArgsTableProps = BaseProps | OfProps | ComponentsProps | StoryProps;
 
 const useArgs = (
   storyId: string,
-  storyStore: StoryStore
+  context: DocsContextProps<any>
 ): [Args, (args: Args) => void, (argNames?: string[]) => void] => {
-  const story = storyStore.fromId(storyId);
+  const channel = addons.getChannel();
+
+  const story = context.storyById(storyId);
   if (!story) {
     throw new Error(`Unknown story: ${storyId}`);
   }
 
-  const { args: initialArgs } = story;
-  const [args, setArgs] = useState(initialArgs);
+  const storyContext = context.getStoryContext(story);
+
+  const [args, setArgs] = useState(storyContext.args);
   useEffect(() => {
     const cb = (changed: { storyId: string; args: Args }) => {
       if (changed.storyId === storyId) {
         setArgs(changed.args);
       }
     };
-    storyStore._channel.on(Events.STORY_ARGS_UPDATED, cb);
-    return () => storyStore._channel.off(Events.STORY_ARGS_UPDATED, cb);
+    channel.on(Events.STORY_ARGS_UPDATED, cb);
+    return () => channel.off(Events.STORY_ARGS_UPDATED, cb);
   }, [storyId]);
-  const updateArgs = useCallback((newArgs) => storyStore.updateStoryArgs(storyId, newArgs), [
-    storyId,
-  ]);
+  const updateArgs = useCallback(
+    (updatedArgs) => channel.emit(Events.UPDATE_STORY_ARGS, { storyId, updatedArgs }),
+    [storyId]
+  );
   const resetArgs = useCallback(
-    (argNames?: string[]) => storyStore.resetStoryArgs(storyId, argNames),
+    (argNames?: string[]) => channel.emit(Events.RESET_STORY_ARGS, { storyId, argNames }),
     [storyId]
   );
   return [args, updateArgs, resetArgs];
@@ -75,12 +78,12 @@ const useArgs = (
 
 export const extractComponentArgTypes = (
   component: Component,
-  { parameters }: DocsContextProps,
+  { id, storyById }: DocsContextProps<any>,
   include?: PropDescriptor,
   exclude?: PropDescriptor
-): ArgTypes => {
-  const params = parameters || {};
-  const { extractArgTypes }: { extractArgTypes: ArgTypesExtractor } = params.docs || {};
+): StrictArgTypes => {
+  const { parameters } = storyById(id);
+  const { extractArgTypes }: { extractArgTypes: ArgTypesExtractor } = parameters.docs || {};
   if (!extractArgTypes) {
     throw new Error(ArgsTableError.ARGS_UNSUPPORTED);
   }
@@ -94,11 +97,13 @@ const isShortcut = (value?: string) => {
   return value && [CURRENT_SELECTION, PRIMARY_STORY].includes(value);
 };
 
-export const getComponent = (props: ArgsTableProps = {}, context: DocsContextProps): Component => {
+export const getComponent = (
+  props: ArgsTableProps = {},
+  { id, storyById }: DocsContextProps<any>
+): Component => {
   const { of } = props as OfProps;
   const { story } = props as StoryProps;
-  const { parameters = {} } = context;
-  const { component } = parameters;
+  const { component } = storyById(id);
   if (isShortcut(of) || isShortcut(story)) {
     return component || null;
   }
@@ -111,7 +116,7 @@ export const getComponent = (props: ArgsTableProps = {}, context: DocsContextPro
 const addComponentTabs = (
   tabs: Record<string, PureArgsTableProps>,
   components: Record<string, Component>,
-  context: DocsContextProps,
+  context: DocsContextProps<any>,
   include?: PropDescriptor,
   exclude?: PropDescriptor,
   sort?: SortType
@@ -127,47 +132,50 @@ export const StoryTable: FC<
   StoryProps & { component: Component; subcomponents: Record<string, Component> }
 > = (props) => {
   const context = useContext(DocsContext);
+  const { id: currentId, componentStories } = context;
   const {
-    id: currentId,
-    parameters: { argTypes },
-    storyStore,
-  } = context;
-  const { story, component, subcomponents, showComponent, include, exclude, sort } = props;
-  let storyArgTypes;
+    story: storyName,
+    component,
+    subcomponents,
+    showComponent,
+    include,
+    exclude,
+    sort,
+  } = props;
   try {
     let storyId;
-    switch (story) {
+    switch (storyName) {
       case CURRENT_SELECTION: {
         storyId = currentId;
-        storyArgTypes = argTypes;
         break;
       }
       case PRIMARY_STORY: {
-        const primaryStory = getDocsStories(context)[0];
+        const primaryStory = componentStories()[0];
         storyId = primaryStory.id;
-        storyArgTypes = primaryStory.parameters.argTypes;
         break;
       }
       default: {
-        storyId = lookupStoryId(story, context);
-        const data = storyStore.fromId(storyId);
-        storyArgTypes = data.parameters.argTypes;
+        storyId = lookupStoryId(storyName, context);
       }
     }
-    storyArgTypes = filterArgTypes(storyArgTypes, include, exclude);
+    const story = useStory(storyId, context);
+    if (!story) {
+      return <div>Loading...</div>;
+    }
+
+    const argTypes = filterArgTypes(story.argTypes, include, exclude);
 
     const mainLabel = getComponentName(component) || 'Story';
 
     // eslint-disable-next-line prefer-const
-    let [args, updateArgs, resetArgs] = useArgs(storyId, storyStore);
-    let tabs = { [mainLabel]: { rows: storyArgTypes, args, updateArgs, resetArgs } } as Record<
+    let [args, updateArgs, resetArgs] = useArgs(storyId, context);
+    let tabs = { [mainLabel]: { rows: argTypes, args, updateArgs, resetArgs } } as Record<
       string,
       PureArgsTableProps
     >;
 
     // Use the dynamically generated component tabs if there are no controls
-    const storyHasArgsWithControls =
-      storyArgTypes && Object.values(storyArgTypes).find((v) => !!v?.control);
+    const storyHasArgsWithControls = argTypes && Object.values(argTypes).find((v) => !!v?.control);
 
     if (!storyHasArgsWithControls) {
       updateArgs = null;
@@ -203,15 +211,19 @@ export const ComponentsTable: FC<ComponentsProps> = (props) => {
 
 export const ArgsTable: FC<ArgsTableProps> = (props) => {
   const context = useContext(DocsContext);
-  const { parameters: { subcomponents, controls } = {} } = context;
+  const { id, storyById } = context;
+  const {
+    parameters: { controls },
+    subcomponents,
+  } = storyById(id);
 
   const { include, exclude, components, sort: sortProp } = props as ComponentsProps;
-  const { story } = props as StoryProps;
+  const { story: storyName } = props as StoryProps;
 
   const sort = sortProp || controls?.sort;
 
   const main = getComponent(props, context);
-  if (story) {
+  if (storyName) {
     return <StoryTable {...(props as StoryProps)} component={main} {...{ subcomponents, sort }} />;
   }
 
diff --git a/addons/docs/src/blocks/Canvas.tsx b/addons/docs/src/blocks/Canvas.tsx
index 41deb826d8da..1f42270c907a 100644
--- a/addons/docs/src/blocks/Canvas.tsx
+++ b/addons/docs/src/blocks/Canvas.tsx
@@ -1,6 +1,6 @@
 import React, { FC, ReactElement, ReactNode, ReactNodeArray, useContext } from 'react';
 import { MDXProvider } from '@mdx-js/react';
-import { toId, storyNameFromExport } from '@storybook/csf';
+import { toId, storyNameFromExport, AnyFramework } from '@storybook/csf';
 import {
   resetComponents,
   Preview as PurePreview,
@@ -19,10 +19,10 @@ type CanvasProps = PurePreviewProps & {
 
 const getPreviewProps = (
   { withSource, mdxSource, children, ...props }: CanvasProps & { children?: ReactNode },
-  docsContext: DocsContextProps,
+  docsContext: DocsContextProps<AnyFramework>,
   sourceContext: SourceContextProps
 ): PurePreviewProps => {
-  const { mdxComponentMeta, mdxStoryNameToKey } = docsContext;
+  const { mdxComponentAnnotations, mdxStoryNameToKey } = docsContext;
   let sourceState = withSource;
   if (sourceState === SourceState.NONE) {
     return props;
@@ -41,7 +41,7 @@ const getPreviewProps = (
     (s) =>
       s.props.id ||
       toId(
-        mdxComponentMeta.id || mdxComponentMeta.title,
+        mdxComponentAnnotations.id || mdxComponentAnnotations.title,
         storyNameFromExport(mdxStoryNameToKey[s.props.name])
       )
   );
diff --git a/addons/docs/src/blocks/Description.tsx b/addons/docs/src/blocks/Description.tsx
index f8bf8e711db0..2f92c105e160 100644
--- a/addons/docs/src/blocks/Description.tsx
+++ b/addons/docs/src/blocks/Description.tsx
@@ -31,12 +31,13 @@ const noDescription = (component?: Component): string | null => null;
 
 export const getDescriptionProps = (
   { of, type, markdown, children }: DescriptionProps,
-  { parameters }: DocsContextProps
+  { id, storyById }: DocsContextProps<any>
 ): PureDescriptionProps => {
+  const { component, parameters } = storyById(id);
   if (children || markdown) {
     return { markdown: children || markdown };
   }
-  const { component, notes, info, docs } = parameters;
+  const { notes, info, docs } = parameters;
   const { extractComponentDescription = noDescription, description } = docs || {};
   const target = of === CURRENT_SELECTION ? component : of;
 
@@ -63,7 +64,7 @@ ${extractComponentDescription(target) || ''}
     case DescriptionType.DOCGEN:
     case DescriptionType.AUTO:
     default:
-      return { markdown: extractComponentDescription(target, parameters) };
+      return { markdown: extractComponentDescription(target, { component, ...parameters }) };
   }
 };
 
diff --git a/addons/docs/src/blocks/DocsContainer.tsx b/addons/docs/src/blocks/DocsContainer.tsx
index 214090580a40..a3cf42db9ac0 100644
--- a/addons/docs/src/blocks/DocsContainer.tsx
+++ b/addons/docs/src/blocks/DocsContainer.tsx
@@ -5,6 +5,7 @@ import dedent from 'ts-dedent';
 import { MDXProvider } from '@mdx-js/react';
 import { ThemeProvider, ensure as ensureTheme } from '@storybook/theming';
 import { DocsWrapper, DocsContent, components as htmlComponents } from '@storybook/components';
+import { AnyFramework } from '@storybook/csf';
 import { DocsContextProps, DocsContext } from './DocsContext';
 import { anchorBlockIdFromId } from './Anchor';
 import { storyBlockIdFromId } from './Story';
@@ -14,8 +15,8 @@ import { scrollToElement } from './utils';
 
 const { document, window: globalWindow } = global;
 
-export interface DocsContainerProps {
-  context: DocsContextProps;
+export interface DocsContainerProps<TFramework extends AnyFramework> {
+  context: DocsContextProps<TFramework>;
 }
 
 const defaultComponents = {
@@ -34,9 +35,14 @@ const warnOptionsTheme = deprecate(
 `
 );
 
-export const DocsContainer: FunctionComponent<DocsContainerProps> = ({ context, children }) => {
-  const { id: storyId = null, parameters = {} } = context || {};
-  const { options = {}, docs = {} } = parameters;
+export const DocsContainer: FunctionComponent<DocsContainerProps<any>> = ({
+  context,
+  children,
+}) => {
+  const { id: storyId, storyById } = context;
+  const {
+    parameters: { options = {}, docs = {} },
+  } = storyById(storyId);
   let themeVars = docs.theme;
   if (!themeVars && options.theme) {
     warnOptionsTheme();
diff --git a/addons/docs/src/blocks/DocsContext.ts b/addons/docs/src/blocks/DocsContext.ts
index 9f66a2588552..3fada361a323 100644
--- a/addons/docs/src/blocks/DocsContext.ts
+++ b/addons/docs/src/blocks/DocsContext.ts
@@ -1,22 +1,10 @@
 import { Context, createContext } from 'react';
 import { window as globalWindow } from 'global';
 
-export interface DocsContextProps {
-  id?: string;
-  kind?: string;
-  name?: string;
+import { DocsContextProps } from '@storybook/preview-web';
+import { AnyFramework } from '@storybook/csf';
 
-  /**
-   * mdxStoryNameToKey is an MDX-compiler-generated mapping of an MDX story's
-   * display name to its story key for ID generation. It's used internally by the `<Story>`
-   * and `Preview` doc blocks.
-   */
-  mdxStoryNameToKey?: Record<string, string>;
-  mdxComponentMeta?: any;
-  parameters?: any;
-  storyStore?: any;
-  forceRender?: () => void;
-}
+export type { DocsContextProps };
 
 // We add DocsContext to window. The reason is that in case DocsContext.ts is
 // imported multiple times (maybe once directly, and another time from a minified bundle)
@@ -29,4 +17,4 @@ if (globalWindow.__DOCS_CONTEXT__ === undefined) {
   globalWindow.__DOCS_CONTEXT__.displayName = 'DocsContext';
 }
 
-export const DocsContext: Context<DocsContextProps> = globalWindow.__DOCS_CONTEXT__;
+export const DocsContext: Context<DocsContextProps<AnyFramework>> = globalWindow.__DOCS_CONTEXT__;
diff --git a/addons/docs/src/blocks/DocsPage.test.ts b/addons/docs/src/blocks/DocsPage.test.ts
index 34a120d6e90c..0e94fcb1d7d2 100644
--- a/addons/docs/src/blocks/DocsPage.test.ts
+++ b/addons/docs/src/blocks/DocsPage.test.ts
@@ -2,9 +2,8 @@ import { extractTitle } from './Title';
 
 describe('defaultTitleSlot', () => {
   it('splits on last /', () => {
-    const parameters = {};
-    expect(extractTitle({ kind: 'a/b/c', parameters })).toBe('c');
-    expect(extractTitle({ kind: 'a|b', parameters })).toBe('a|b');
-    expect(extractTitle({ kind: 'a/b/c.d', parameters })).toBe('c.d');
+    expect(extractTitle({ title: 'a/b/c' } as any)).toBe('c');
+    expect(extractTitle({ title: 'a|b' } as any)).toBe('a|b');
+    expect(extractTitle({ title: 'a/b/c.d' } as any)).toBe('c.d');
   });
 });
diff --git a/addons/docs/src/blocks/Meta.tsx b/addons/docs/src/blocks/Meta.tsx
index 5fd193e5df27..d2aca741e5c4 100644
--- a/addons/docs/src/blocks/Meta.tsx
+++ b/addons/docs/src/blocks/Meta.tsx
@@ -1,17 +1,15 @@
 import React, { FC, useContext } from 'react';
 import global from 'global';
-import { Args, BaseAnnotations, BaseMeta } from '@storybook/addons';
+import { BaseAnnotations } from '@storybook/csf';
 import { Anchor } from './Anchor';
 import { DocsContext, DocsContextProps } from './DocsContext';
-import { getDocsStories } from './utils';
-import { Component } from './types';
 
 const { document } = global;
 
-type MetaProps = BaseMeta<Component> & BaseAnnotations<Args, any>;
+type MetaProps = BaseAnnotations<any>;
 
-function getFirstStoryId(docsContext: DocsContextProps): string {
-  const stories = getDocsStories(docsContext);
+function getFirstStoryId(docsContext: DocsContextProps<any>): string {
+  const stories = docsContext.componentStories();
 
   return stories.length > 0 ? stories[0].id : null;
 }
diff --git a/addons/docs/src/blocks/Primary.tsx b/addons/docs/src/blocks/Primary.tsx
index 242e17794876..01252656a329 100644
--- a/addons/docs/src/blocks/Primary.tsx
+++ b/addons/docs/src/blocks/Primary.tsx
@@ -1,18 +1,19 @@
 import React, { useContext, FC } from 'react';
+import { Story } from '@storybook/store';
+
 import { DocsContext } from './DocsContext';
 import { DocsStory } from './DocsStory';
-import { getDocsStories } from './utils';
 
 interface PrimaryProps {
   name?: string;
 }
 
 export const Primary: FC<PrimaryProps> = ({ name }) => {
-  const context = useContext(DocsContext);
-  const componentStories = getDocsStories(context);
+  const { componentStories: getComponentStories } = useContext(DocsContext);
+  const componentStories = getComponentStories();
   let story;
   if (componentStories) {
-    story = name ? componentStories.find((s) => s.name === name) : componentStories[0];
+    story = name ? componentStories.find((s: Story<any>) => s.name === name) : componentStories[0];
   }
   return story ? <DocsStory {...story} expanded={false} withToolbar /> : null;
 };
diff --git a/addons/docs/src/blocks/Source.tsx b/addons/docs/src/blocks/Source.tsx
index 4777d11072eb..0b23310cd8fd 100644
--- a/addons/docs/src/blocks/Source.tsx
+++ b/addons/docs/src/blocks/Source.tsx
@@ -5,8 +5,7 @@ import {
   SourceProps as PureSourceProps,
 } from '@storybook/components';
 import { StoryId } from '@storybook/api';
-import { logger } from '@storybook/client-logger';
-import { StoryContext } from '@storybook/addons';
+import { Story } from '@storybook/store';
 
 import { DocsContext, DocsContextProps } from './DocsContext';
 import { SourceContext, SourceContextProps } from './SourceContainer';
@@ -43,25 +42,15 @@ type NoneProps = CommonProps;
 
 type SourceProps = SingleSourceProps | MultiSourceProps | CodeProps | NoneProps;
 
-const getStoryContext = (storyId: StoryId, docsContext: DocsContextProps): StoryContext | null => {
-  const { storyStore } = docsContext;
-  const storyContext = storyStore?.fromId(storyId);
-
-  if (!storyContext) {
-    // Fallback if we can't get the story data for this story
-    logger.warn(`Unable to find information for story ID '${storyId}'`);
-    return null;
-  }
-
-  return storyContext;
+const getStory = (storyId: StoryId, docsContext: DocsContextProps<any>): Story<any> | null => {
+  return docsContext.storyById(storyId);
 };
 
-const getSourceState = (storyIds: string[], docsContext: DocsContextProps) => {
+const getSourceState = (storyIds: string[], docsContext: DocsContextProps<any>) => {
   const states = storyIds
     .map((storyId) => {
-      const storyContext = getStoryContext(storyId, docsContext);
-      if (!storyContext) return null;
-      return storyContext.parameters.docs?.source?.state;
+      const story = getStory(storyId, docsContext);
+      return story?.parameters.docs?.source?.state;
     })
     .filter(Boolean);
 
@@ -77,12 +66,12 @@ const getStorySource = (storyId: StoryId, sourceContext: SourceContextProps): st
   return sources?.[storyId] || '';
 };
 
-const getSnippet = (snippet: string, storyContext?: StoryContext): string => {
-  if (!storyContext) {
+const getSnippet = (snippet: string, story?: Story<any>): string => {
+  if (!story) {
     return snippet;
   }
 
-  const { parameters } = storyContext;
+  const { parameters } = story;
   // eslint-disable-next-line no-underscore-dangle
   const isArgsStory = parameters.__isArgsStory;
   const type = parameters.docs?.source?.type || SourceType.AUTO;
@@ -95,16 +84,16 @@ const getSnippet = (snippet: string, storyContext?: StoryContext): string => {
 
   // if user has explicitly set this as dynamic, use snippet
   if (type === SourceType.DYNAMIC) {
-    return parameters.docs?.transformSource?.(snippet, storyContext) || snippet;
+    return parameters.docs?.transformSource?.(snippet, story) || snippet;
   }
 
   // if this is an args story and there's a snippet
   if (type === SourceType.AUTO && snippet && isArgsStory) {
-    return parameters.docs?.transformSource?.(snippet, storyContext) || snippet;
+    return parameters.docs?.transformSource?.(snippet, story) || snippet;
   }
 
   // otherwise, use the source code logic
-  const enhanced = enhanceSource(storyContext) || parameters;
+  const enhanced = enhanceSource(story) || parameters;
   return enhanced?.docs?.source?.code || '';
 };
 
@@ -112,10 +101,11 @@ type SourceStateProps = { state: SourceState };
 
 export const getSourceProps = (
   props: SourceProps,
-  docsContext: DocsContextProps,
+  docsContext: DocsContextProps<any>,
   sourceContext: SourceContextProps
 ): PureSourceProps & SourceStateProps => {
-  const { id: currentId, parameters = {} } = docsContext;
+  const { id: currentId, storyById } = docsContext;
+  const { parameters } = storyById(currentId);
 
   const codeProps = props as CodeProps;
   const singleProps = props as SingleSourceProps;
@@ -131,8 +121,8 @@ export const getSourceProps = (
     source = targetIds
       .map((storyId) => {
         const storySource = getStorySource(storyId, sourceContext);
-        const storyContext = getStoryContext(storyId, docsContext);
-        return getSnippet(storySource, storyContext);
+        const story = getStory(storyId, docsContext);
+        return getSnippet(storySource, story);
       })
       .join('\n\n');
   }
diff --git a/addons/docs/src/blocks/Stories.tsx b/addons/docs/src/blocks/Stories.tsx
index afb7de3e8ad8..9d460a5cd830 100644
--- a/addons/docs/src/blocks/Stories.tsx
+++ b/addons/docs/src/blocks/Stories.tsx
@@ -2,7 +2,6 @@ import React, { useContext, FunctionComponent } from 'react';
 import { DocsContext } from './DocsContext';
 import { DocsStory } from './DocsStory';
 import { Heading } from './Heading';
-import { getDocsStories } from './utils';
 import { DocsStoryProps } from './types';
 
 interface StoriesProps {
@@ -11,10 +10,9 @@ interface StoriesProps {
 }
 
 export const Stories: FunctionComponent<StoriesProps> = ({ title, includePrimary = false }) => {
-  const context = useContext(DocsContext);
-  const componentStories = getDocsStories(context);
+  const { componentStories } = useContext(DocsContext);
 
-  let stories: DocsStoryProps[] = componentStories;
+  let stories: DocsStoryProps[] = componentStories();
   if (!includePrimary) stories = stories.slice(1);
 
   if (!stories || stories.length === 0) {
diff --git a/addons/docs/src/blocks/Story.tsx b/addons/docs/src/blocks/Story.tsx
index 6e239f072ca1..7e2abf1e5f50 100644
--- a/addons/docs/src/blocks/Story.tsx
+++ b/addons/docs/src/blocks/Story.tsx
@@ -1,17 +1,28 @@
-import React, { FunctionComponent, ReactNode, ElementType, ComponentProps } from 'react';
+import React, {
+  FunctionComponent,
+  ReactNode,
+  ElementType,
+  ComponentProps,
+  useContext,
+} from 'react';
 import { MDXProvider } from '@mdx-js/react';
 import { resetComponents, Story as PureStory } from '@storybook/components';
-import { toId, storyNameFromExport } from '@storybook/csf';
-import { Args, BaseAnnotations } from '@storybook/addons';
-import { CURRENT_SELECTION } from './types';
+import { StoryId, toId, storyNameFromExport, StoryAnnotations, AnyFramework } from '@storybook/csf';
+import { Story as StoryType } from '@storybook/store';
 
+import { CURRENT_SELECTION } from './types';
 import { DocsContext, DocsContextProps } from './DocsContext';
+import { useStory } from './useStory';
 
 export const storyBlockIdFromId = (storyId: string) => `story--${storyId}`;
 
 type PureStoryProps = ComponentProps<typeof PureStory>;
 
-type CommonProps = BaseAnnotations<Args, any> & {
+type Annotations = Pick<
+  StoryAnnotations<AnyFramework>,
+  'decorators' | 'parameters' | 'args' | 'argTypes' | 'loaders'
+>;
+type CommonProps = Annotations & {
   height?: string;
   inline?: boolean;
 };
@@ -34,22 +45,26 @@ export type StoryProps = (StoryDefProps | StoryRefProps | StoryImportProps) & Co
 
 export const lookupStoryId = (
   storyName: string,
-  { mdxStoryNameToKey, mdxComponentMeta }: DocsContextProps
+  { mdxStoryNameToKey, mdxComponentAnnotations }: DocsContextProps<any>
 ) =>
   toId(
-    mdxComponentMeta.id || mdxComponentMeta.title,
+    mdxComponentAnnotations.id || mdxComponentAnnotations.title,
     storyNameFromExport(mdxStoryNameToKey[storyName])
   );
 
-export const getStoryProps = (props: StoryProps, context: DocsContextProps): PureStoryProps => {
+export const getStoryId = (props: StoryProps, context: DocsContextProps<AnyFramework>): StoryId => {
   const { id } = props as StoryRefProps;
   const { name } = props as StoryDefProps;
   const inputId = id === CURRENT_SELECTION ? context.id : id;
-  const previewId = inputId || lookupStoryId(name, context);
-  const data = context.storyStore.fromId(previewId) || {};
+  return inputId || lookupStoryId(name, context);
+};
 
-  const { height, inline } = props;
-  const { storyFn = undefined, name: storyName = undefined, parameters = {} } = data;
+export const getStoryProps = (
+  { height, inline }: StoryProps,
+  story: StoryType<any>,
+  context: DocsContextProps<any>
+): PureStoryProps => {
+  const { name: storyName, parameters } = story;
   const { docs = {} } = parameters;
 
   if (docs.disable) {
@@ -65,33 +80,41 @@ export const getStoryProps = (props: StoryProps, context: DocsContextProps): Pur
     );
   }
 
+  const boundStoryFn = () =>
+    story.unboundStoryFn({
+      ...context.getStoryContext(story),
+      loaded: {},
+    });
   return {
     parameters,
     inline: storyIsInline,
-    id: previewId,
-    storyFn: prepareForInline && storyFn ? () => prepareForInline(storyFn, data) : storyFn,
+    id: story.id,
+    storyFn: prepareForInline ? () => prepareForInline(boundStoryFn, story) : boundStoryFn,
     height: height || (storyIsInline ? undefined : iframeHeight),
     title: storyName,
   };
 };
 
-const Story: FunctionComponent<StoryProps> = (props) => (
-  <DocsContext.Consumer>
-    {(context) => {
-      const storyProps = getStoryProps(props, context);
-      if (!storyProps) {
-        return null;
-      }
-      return (
-        <div id={storyBlockIdFromId(storyProps.id)}>
-          <MDXProvider components={resetComponents}>
-            <PureStory {...storyProps} />
-          </MDXProvider>
-        </div>
-      );
-    }}
-  </DocsContext.Consumer>
-);
+const Story: FunctionComponent<StoryProps> = (props) => {
+  const context = useContext(DocsContext);
+  const story = useStory(getStoryId(props, context), context);
+
+  if (!story) {
+    return <div>Loading...</div>;
+  }
+
+  const storyProps = getStoryProps(props, story, context);
+  if (!storyProps) {
+    return null;
+  }
+  return (
+    <div id={storyBlockIdFromId(storyProps.id)}>
+      <MDXProvider components={resetComponents}>
+        <PureStory {...storyProps} />
+      </MDXProvider>
+    </div>
+  );
+};
 
 Story.defaultProps = {
   children: null,
diff --git a/addons/docs/src/blocks/Subtitle.tsx b/addons/docs/src/blocks/Subtitle.tsx
index 3fb44a8af5a6..64e160176dd5 100644
--- a/addons/docs/src/blocks/Subtitle.tsx
+++ b/addons/docs/src/blocks/Subtitle.tsx
@@ -7,8 +7,8 @@ interface SubtitleProps {
 }
 
 export const Subtitle: FunctionComponent<SubtitleProps> = ({ children }) => {
-  const context = useContext(DocsContext);
-  const { parameters } = context;
+  const { id, storyById } = useContext(DocsContext);
+  const { parameters } = storyById(id);
   let text: JSX.Element | string = children;
   if (!text) {
     text = parameters?.componentSubtitle;
diff --git a/addons/docs/src/blocks/Title.tsx b/addons/docs/src/blocks/Title.tsx
index 052f0dbe904b..a815a1894861 100644
--- a/addons/docs/src/blocks/Title.tsx
+++ b/addons/docs/src/blocks/Title.tsx
@@ -8,9 +8,9 @@ interface TitleProps {
 
 const STORY_KIND_PATH_SEPARATOR = /\s*\/\s*/;
 
-export const extractTitle = ({ kind }: DocsContextProps) => {
-  const groups = kind.trim().split(STORY_KIND_PATH_SEPARATOR);
-  return (groups && groups[groups.length - 1]) || kind;
+export const extractTitle = ({ title }: DocsContextProps<any>) => {
+  const groups = title.trim().split(STORY_KIND_PATH_SEPARATOR);
+  return (groups && groups[groups.length - 1]) || title;
 };
 
 export const Title: FunctionComponent<TitleProps> = ({ children }) => {
diff --git a/addons/docs/src/blocks/enhanceSource.ts b/addons/docs/src/blocks/enhanceSource.ts
index 42a87ca906a9..ea969c1d46ce 100644
--- a/addons/docs/src/blocks/enhanceSource.ts
+++ b/addons/docs/src/blocks/enhanceSource.ts
@@ -1,5 +1,6 @@
 import { combineParameters } from '@storybook/client-api';
-import { StoryContext, Parameters } from '@storybook/addons';
+import { Parameters } from '@storybook/addons';
+import { Story } from '@storybook/store';
 
 // ============================================================
 // START @storybook/source-loader/extract-source
@@ -76,8 +77,8 @@ const extract = (targetId: string, { source, locationsMap }: StorySource) => {
   return extractSource(location, lines);
 };
 
-export const enhanceSource = (context: StoryContext): Parameters => {
-  const { id, parameters } = context;
+export const enhanceSource = (story: Story<any>): Parameters => {
+  const { id, parameters } = story;
   const { storySource, docs = {} } = parameters;
   const { transformSource } = docs;
 
@@ -87,7 +88,7 @@ export const enhanceSource = (context: StoryContext): Parameters => {
   }
 
   const input = extract(id, storySource);
-  const code = transformSource ? transformSource(input, context) : input;
+  const code = transformSource ? transformSource(input, story) : input;
 
   return { docs: combineParameters(docs, { source: { code } }) };
 };
diff --git a/addons/docs/src/blocks/mdx.tsx b/addons/docs/src/blocks/mdx.tsx
index 3dbc6746ee99..59566c729309 100644
--- a/addons/docs/src/blocks/mdx.tsx
+++ b/addons/docs/src/blocks/mdx.tsx
@@ -17,7 +17,7 @@ export const assertIsFn = (val: any) => {
 };
 
 // Hacky utility for adding mdxStoryToId to the default context
-export const AddContext: FC<DocsContextProps> = (props) => {
+export const AddContext: FC<DocsContextProps<any>> = (props) => {
   const { children, ...rest } = props;
   const parentContext = React.useContext(DocsContext);
   return (
diff --git a/addons/docs/src/blocks/useStory.ts b/addons/docs/src/blocks/useStory.ts
new file mode 100644
index 000000000000..6c785247fe9c
--- /dev/null
+++ b/addons/docs/src/blocks/useStory.ts
@@ -0,0 +1,18 @@
+import { useState, useEffect } from 'react';
+import { StoryId, AnyFramework } from '@storybook/csf';
+import { Story } from '@storybook/store';
+
+import { DocsContextProps } from './DocsContext';
+
+export function useStory<TFramework extends AnyFramework>(
+  storyId: StoryId,
+  context: DocsContextProps<TFramework>
+): Story<TFramework> | void {
+  const [story, setStory] = useState(null);
+
+  useEffect(() => {
+    context.loadStory(storyId).then((s) => setStory(s));
+  });
+
+  return story;
+}
diff --git a/addons/docs/src/blocks/utils.ts b/addons/docs/src/blocks/utils.ts
index 036925af27e5..753959137741 100644
--- a/addons/docs/src/blocks/utils.ts
+++ b/addons/docs/src/blocks/utils.ts
@@ -1,18 +1,5 @@
 /* eslint-disable no-underscore-dangle */
-import { DocsContextProps } from './DocsContext';
-import { StoryData, Component } from './types';
-
-export const getDocsStories = (context: DocsContextProps): StoryData[] => {
-  const { storyStore, kind } = context;
-
-  if (!storyStore) {
-    return [];
-  }
-
-  return storyStore
-    .getStoriesForKind(kind)
-    .filter((s: any) => !(s.parameters && s.parameters.docs && s.parameters.docs.disable));
-};
+import { Component } from './types';
 
 const titleCase = (str: string): string =>
   str
diff --git a/addons/docs/src/frameworks/angular/compodoc.ts b/addons/docs/src/frameworks/angular/compodoc.ts
index 8a7d68c87af9..c6dd02d450c7 100644
--- a/addons/docs/src/frameworks/angular/compodoc.ts
+++ b/addons/docs/src/frameworks/angular/compodoc.ts
@@ -43,7 +43,7 @@ export const checkValidCompodocJson = (compodocJson: CompodocJson) => {
 const hasDecorator = (item: Property, decoratorName: string) =>
   item.decorators && item.decorators.find((x: any) => x.name === decoratorName);
 
-const mapPropertyToSection = (key: string, item: Property) => {
+const mapPropertyToSection = (item: Property) => {
   if (hasDecorator(item, 'ViewChild')) {
     return 'view child';
   }
@@ -73,7 +73,7 @@ const mapItemToSection = (key: string, item: Method | Property): string => {
       if (isMethod(item)) {
         throw new Error("Cannot be of type Method if key === 'propertiesClass'");
       }
-      return mapPropertyToSection(key, item);
+      return mapPropertyToSection(item);
     default:
       throw new Error(`Unknown key: ${key}`);
   }
diff --git a/addons/docs/src/frameworks/angular/prepareForInline.ts b/addons/docs/src/frameworks/angular/prepareForInline.ts
index fc9e4f0ab26c..76be79a3f4ac 100644
--- a/addons/docs/src/frameworks/angular/prepareForInline.ts
+++ b/addons/docs/src/frameworks/angular/prepareForInline.ts
@@ -2,16 +2,19 @@ import React from 'react';
 import pLimit from 'p-limit';
 import { nanoid } from 'nanoid';
 
-import { IStory, StoryContext } from '@storybook/angular';
+import { AngularFramework, StoryContext } from '@storybook/angular';
 import { rendererFactory } from '@storybook/angular/renderer';
-import { StoryFn } from '@storybook/addons';
+import { PartialStoryFn } from '@storybook/csf';
 
 const limit = pLimit(1);
 
 /**
  * Uses the angular renderer to generate a story. Uses p-limit to run synchronously
  */
-export const prepareForInline = (storyFn: StoryFn<IStory>, { id, parameters }: StoryContext) => {
+export const prepareForInline = (
+  storyFn: PartialStoryFn<AngularFramework>,
+  { id, parameters }: StoryContext
+) => {
   return React.createElement('div', {
     ref: async (node?: HTMLDivElement): Promise<void> => {
       if (!node) {
diff --git a/addons/docs/src/frameworks/angular/sourceDecorator.ts b/addons/docs/src/frameworks/angular/sourceDecorator.ts
index d5e3b1bd96d2..496a316f55c5 100644
--- a/addons/docs/src/frameworks/angular/sourceDecorator.ts
+++ b/addons/docs/src/frameworks/angular/sourceDecorator.ts
@@ -1,5 +1,6 @@
-import { addons, StoryContext, StoryFn } from '@storybook/addons';
-import { IStory } from '@storybook/angular';
+import { addons } from '@storybook/addons';
+import { PartialStoryFn } from '@storybook/csf';
+import { StoryContext, AngularFramework } from '@storybook/angular';
 import { computesTemplateSourceFromComponent } from '@storybook/angular/renderer';
 import prettierHtml from 'prettier/parser-html';
 import prettier from 'prettier/standalone';
@@ -30,7 +31,10 @@ const prettyUp = (source: string) => {
  * @param storyFn Fn
  * @param context  StoryContext
  */
-export const sourceDecorator = (storyFn: StoryFn<IStory>, context: StoryContext) => {
+export const sourceDecorator = (
+  storyFn: PartialStoryFn<AngularFramework>,
+  context: StoryContext
+) => {
   const story = storyFn();
   if (skipSourceRender(context)) {
     return story;
@@ -38,9 +42,7 @@ export const sourceDecorator = (storyFn: StoryFn<IStory>, context: StoryContext)
   const channel = addons.getChannel();
   const { props, template, userDefinedTemplate } = story;
 
-  const {
-    parameters: { component, argTypes },
-  } = context;
+  const { component, argTypes } = context;
 
   if (component && !userDefinedTemplate) {
     const source = computesTemplateSourceFromComponent(component, props, argTypes);
diff --git a/addons/docs/src/frameworks/common/enhanceArgTypes.test.ts b/addons/docs/src/frameworks/common/enhanceArgTypes.test.ts
index 252dd5cfe01d..43e0da967100 100644
--- a/addons/docs/src/frameworks/common/enhanceArgTypes.test.ts
+++ b/addons/docs/src/frameworks/common/enhanceArgTypes.test.ts
@@ -1,4 +1,5 @@
-import { ArgType, ArgTypes } from '@storybook/api';
+import { ArgTypes } from '@storybook/api';
+import { StrictInputType } from '@storybook/csf';
 import { enhanceArgTypes } from './enhanceArgTypes';
 
 expect.addSnapshotSerializer({
@@ -12,30 +13,28 @@ const enhance = ({
   extractedArgTypes,
   isArgsStory = true,
 }: {
-  argType?: ArgType;
+  argType?: StrictInputType;
   arg?: any;
   extractedArgTypes?: ArgTypes;
   isArgsStory?: boolean;
 }) => {
   const context = {
-    id: 'foo--bar',
+    componentId: 'foo',
+    title: 'foo',
     kind: 'foo',
+    id: 'foo--bar',
     name: 'bar',
+    story: 'bar',
+    component: 'dummy',
     parameters: {
-      component: 'dummy',
       __isArgsStory: isArgsStory,
       docs: {
         extractArgTypes: extractedArgTypes && (() => extractedArgTypes),
       },
-      argTypes: argType && {
-        input: argType,
-      },
-      args: {
-        input: arg,
-      },
     },
-    args: {},
-    argTypes: {},
+    argTypes: argType && { input: argType },
+    initialArgs: { input: arg },
+    args: { input: arg },
     globals: {},
   };
   return enhanceArgTypes(context);
@@ -46,7 +45,7 @@ describe('enhanceArgTypes', () => {
     it('should no-op', () => {
       expect(
         enhance({
-          argType: { foo: 'unmodified', type: { name: 'number' } },
+          argType: { name: 'input', foo: 'unmodified', type: { name: 'number' } },
           isArgsStory: false,
         }).input
       ).toMatchInlineSnapshot(`
@@ -66,7 +65,7 @@ describe('enhanceArgTypes', () => {
         it('number', () => {
           expect(
             enhance({
-              argType: { type: { name: 'number' } },
+              argType: { name: 'input', type: { name: 'number' } },
             }).input
           ).toMatchInlineSnapshot(`
             {
@@ -99,7 +98,7 @@ describe('enhanceArgTypes', () => {
         it('range', () => {
           expect(
             enhance({
-              argType: { control: { type: 'range', min: 0, max: 100 } },
+              argType: { name: 'input', control: { type: 'range', min: 0, max: 100 } },
             }).input
           ).toMatchInlineSnapshot(`
             {
@@ -115,7 +114,7 @@ describe('enhanceArgTypes', () => {
         it('options', () => {
           expect(
             enhance({
-              argType: { control: { type: 'radio', options: [1, 2] } },
+              argType: { name: 'input', control: { type: 'radio', options: [1, 2] } },
             }).input
           ).toMatchInlineSnapshot(`
             {
@@ -137,7 +136,7 @@ describe('enhanceArgTypes', () => {
       it('user-specified argTypes take precedence over extracted argTypes', () => {
         expect(
           enhance({
-            argType: { type: { name: 'number' } },
+            argType: { name: 'input', type: { name: 'number' } },
             extractedArgTypes: { input: { type: { name: 'string' } } },
           }).input
         ).toMatchInlineSnapshot(`
@@ -153,7 +152,7 @@ describe('enhanceArgTypes', () => {
       it('user-specified argTypes take precedence over inferred argTypes', () => {
         expect(
           enhance({
-            argType: { type: { name: 'number' } },
+            argType: { name: 'input', type: { name: 'number' } },
             arg: 'hello',
           }).input
         ).toMatchInlineSnapshot(`
@@ -184,7 +183,7 @@ describe('enhanceArgTypes', () => {
       it('user-specified controls take precedence over inferred controls', () => {
         expect(
           enhance({
-            argType: { defaultValue: 5, control: { type: 'range', step: 50 } },
+            argType: { name: 'input', defaultValue: 5, control: { type: 'range', step: 50 } },
             arg: 3,
             extractedArgTypes: { input: { name: 'input' } },
           }).input
@@ -223,7 +222,7 @@ describe('enhanceArgTypes', () => {
       it('includes extracted argTypes when user-specified argTypes match', () => {
         expect(
           enhance({
-            argType: { type: { name: 'number' } },
+            argType: { name: 'input', type: { name: 'number' } },
             extractedArgTypes: { input: { name: 'input' }, foo: { type: { name: 'number' } } },
           })
         ).toMatchInlineSnapshot(`
@@ -246,7 +245,7 @@ describe('enhanceArgTypes', () => {
       it('excludes extracted argTypes when user-specified argTypes do not match', () => {
         expect(
           enhance({
-            argType: { type: { name: 'number' } },
+            argType: { name: 'input', type: { name: 'number' } },
             extractedArgTypes: { foo: { type: { name: 'number' } } },
           })
         ).toMatchInlineSnapshot(`
diff --git a/addons/docs/src/frameworks/common/enhanceArgTypes.ts b/addons/docs/src/frameworks/common/enhanceArgTypes.ts
index f1ee242400e0..540b80e56b42 100644
--- a/addons/docs/src/frameworks/common/enhanceArgTypes.ts
+++ b/addons/docs/src/frameworks/common/enhanceArgTypes.ts
@@ -1,17 +1,20 @@
-import mapValues from 'lodash/mapValues';
-import { ArgTypesEnhancer, combineParameters } from '@storybook/client-api';
-import { normalizeArgTypes } from './normalizeArgTypes';
+import { AnyFramework, StoryContextForEnhancers } from '@storybook/csf';
+import { combineParameters } from '@storybook/store';
 
-export const enhanceArgTypes: ArgTypesEnhancer = (context) => {
-  const { component, argTypes: userArgTypes = {}, docs = {} } = context.parameters;
+export const enhanceArgTypes = <TFramework extends AnyFramework>(
+  context: StoryContextForEnhancers<TFramework>
+) => {
+  const {
+    component,
+    argTypes: userArgTypes,
+    parameters: { docs = {} },
+  } = context;
   const { extractArgTypes } = docs;
 
-  const normalizedArgTypes = normalizeArgTypes(userArgTypes);
-  const namedArgTypes = mapValues(normalizedArgTypes, (val, key) => ({ name: key, ...val }));
   const extractedArgTypes = extractArgTypes && component ? extractArgTypes(component) : {};
   const withExtractedTypes = extractedArgTypes
-    ? combineParameters(extractedArgTypes, namedArgTypes)
-    : namedArgTypes;
+    ? combineParameters(extractedArgTypes, userArgTypes)
+    : userArgTypes;
 
   return withExtractedTypes;
 };
diff --git a/addons/docs/src/frameworks/common/normalizeArgTypes.ts b/addons/docs/src/frameworks/common/normalizeArgTypes.ts
deleted file mode 100644
index 8af93c8e4142..000000000000
--- a/addons/docs/src/frameworks/common/normalizeArgTypes.ts
+++ /dev/null
@@ -1,18 +0,0 @@
-import mapValues from 'lodash/mapValues';
-import { ArgTypes } from '@storybook/api';
-import { SBType } from '@storybook/client-api';
-
-const normalizeType = (type: SBType | string) => (typeof type === 'string' ? { name: type } : type);
-
-const normalizeControl = (control?: any) =>
-  typeof control === 'string' ? { type: control } : control;
-
-export const normalizeArgTypes = (argTypes: ArgTypes) =>
-  mapValues(argTypes, (argType) => {
-    if (!argType) return argType;
-    const normalized = { ...argType };
-    const { type, control } = argType;
-    if (type) normalized.type = normalizeType(type);
-    if (control) normalized.control = normalizeControl(control);
-    return normalized;
-  });
diff --git a/addons/docs/src/frameworks/html/prepareForInline.tsx b/addons/docs/src/frameworks/html/prepareForInline.tsx
index 578f42c0bb34..421b9625f68c 100644
--- a/addons/docs/src/frameworks/html/prepareForInline.tsx
+++ b/addons/docs/src/frameworks/html/prepareForInline.tsx
@@ -1,7 +1,7 @@
 import React from 'react';
-import { StoryFn } from '@storybook/addons';
+import { PartialStoryFn } from '@storybook/csf';
 
-export function prepareForInline(storyFn: StoryFn<string>) {
+export function prepareForInline(storyFn: PartialStoryFn<any>) {
   const html = storyFn();
   if (typeof html === 'string') {
     // eslint-disable-next-line react/no-danger
diff --git a/addons/docs/src/frameworks/html/sourceDecorator.ts b/addons/docs/src/frameworks/html/sourceDecorator.ts
index fa45e820ed4c..c37143f4ee2a 100644
--- a/addons/docs/src/frameworks/html/sourceDecorator.ts
+++ b/addons/docs/src/frameworks/html/sourceDecorator.ts
@@ -1,9 +1,12 @@
 /* global window */
-import { addons, StoryContext, StoryFn } from '@storybook/addons';
+import { addons } from '@storybook/addons';
+import { ArgsStoryFn, PartialStoryFn, StoryContext } from '@storybook/csf';
 import dedent from 'ts-dedent';
+import { HtmlFramework } from '@storybook/html';
+
 import { SNIPPET_RENDERED, SourceType } from '../../shared';
 
-function skipSourceRender(context: StoryContext) {
+function skipSourceRender(context: StoryContext<HtmlFramework>) {
   const sourceParams = context?.parameters.docs?.source;
   const isArgsStory = context?.parameters.__isArgsStory;
 
@@ -23,15 +26,18 @@ function defaultTransformSource(source: string) {
   return dedent(source);
 }
 
-function applyTransformSource(source: string, context: StoryContext): string {
+function applyTransformSource(source: string, context: StoryContext<HtmlFramework>): string {
   const docs = context.parameters.docs ?? {};
   const transformSource = docs.transformSource ?? defaultTransformSource;
   return transformSource(source, context);
 }
 
-export function sourceDecorator(storyFn: StoryFn, context: StoryContext) {
+export function sourceDecorator(
+  storyFn: PartialStoryFn<HtmlFramework>,
+  context: StoryContext<HtmlFramework>
+) {
   const story = context?.parameters.docs?.source?.excludeDecorators
-    ? context.originalStoryFn(context.args)
+    ? (context.originalStoryFn as ArgsStoryFn<HtmlFramework>)(context.args, context)
     : storyFn();
 
   if (typeof story === 'string' && !skipSourceRender(context)) {
diff --git a/addons/docs/src/frameworks/react/config.ts b/addons/docs/src/frameworks/react/config.ts
index 61a01e8a3aff..28ef4afc0e47 100644
--- a/addons/docs/src/frameworks/react/config.ts
+++ b/addons/docs/src/frameworks/react/config.ts
@@ -1,4 +1,6 @@
-import { StoryFn } from '@storybook/addons';
+import { PartialStoryFn } from '@storybook/csf';
+import { ReactFramework } from '@storybook/react';
+
 import { extractArgTypes } from './extractArgTypes';
 import { extractComponentDescription } from '../../lib/docgen';
 import { jsxDecorator } from './jsxDecorator';
@@ -7,7 +9,7 @@ export const parameters = {
   docs: {
     inlineStories: true,
     // NOTE: that the result is a react element. Hooks support is provided by the outer code.
-    prepareForInline: (storyFn: StoryFn) => storyFn(),
+    prepareForInline: (storyFn: PartialStoryFn<ReactFramework>) => storyFn(),
     extractArgTypes,
     extractComponentDescription,
   },
diff --git a/addons/docs/src/frameworks/react/extractArgTypes.ts b/addons/docs/src/frameworks/react/extractArgTypes.ts
index 76768f268435..e86dcaa99b88 100644
--- a/addons/docs/src/frameworks/react/extractArgTypes.ts
+++ b/addons/docs/src/frameworks/react/extractArgTypes.ts
@@ -1,4 +1,4 @@
-import { ArgTypes } from '@storybook/api';
+import { StrictArgTypes } from '@storybook/csf';
 import { PropDef, ArgTypesExtractor } from '../../lib/docgen';
 import { extractProps } from './extractProps';
 
@@ -6,7 +6,7 @@ export const extractArgTypes: ArgTypesExtractor = (component) => {
   if (component) {
     const { rows } = extractProps(component);
     if (rows) {
-      return rows.reduce((acc: ArgTypes, row: PropDef) => {
+      return rows.reduce((acc: StrictArgTypes, row: PropDef) => {
         const {
           name,
           description,
diff --git a/addons/docs/src/frameworks/react/jsxDecorator.tsx b/addons/docs/src/frameworks/react/jsxDecorator.tsx
index dedb1505fe63..e3a3cc888460 100644
--- a/addons/docs/src/frameworks/react/jsxDecorator.tsx
+++ b/addons/docs/src/frameworks/react/jsxDecorator.tsx
@@ -3,8 +3,10 @@ import reactElementToJSXString, { Options } from 'react-element-to-jsx-string';
 import dedent from 'ts-dedent';
 import deprecate from 'util-deprecate';
 
-import { addons, StoryContext } from '@storybook/addons';
+import { addons } from '@storybook/addons';
+import { StoryContext, ArgsStoryFn, PartialStoryFn } from '@storybook/csf';
 import { logger } from '@storybook/client-logger';
+import { ReactFramework } from '@storybook/react';
 
 import { SourceType, SNIPPET_RENDERED } from '../../shared';
 import { getDocgenSection } from '../../lib/docgen';
@@ -22,7 +24,7 @@ type JSXOptions = Options & {
   /** Deprecated: A function ran after the story is rendered */
   onBeforeRender?(dom: string): string;
   /** A function ran after a story is rendered (prefer this over `onBeforeRender`) */
-  transformSource?(dom: string, context?: StoryContext): string;
+  transformSource?(dom: string, context?: StoryContext<ReactFramework>): string;
 };
 
 /** Run the user supplied onBeforeRender function if it exists */
@@ -44,7 +46,11 @@ const applyBeforeRender = (domString: string, options: JSXOptions) => {
 };
 
 /** Run the user supplied transformSource function if it exists */
-const applyTransformSource = (domString: string, options: JSXOptions, context?: StoryContext) => {
+const applyTransformSource = (
+  domString: string,
+  options: JSXOptions,
+  context?: StoryContext<ReactFramework>
+) => {
   if (typeof options.transformSource !== 'function') {
     return domString;
   }
@@ -138,7 +144,7 @@ const defaultOpts = {
   showDefaultProps: false,
 };
 
-export const skipJsxRender = (context: StoryContext) => {
+export const skipJsxRender = (context: StoryContext<ReactFramework>) => {
   const sourceParams = context?.parameters.docs?.source;
   const isArgsStory = context?.parameters.__isArgsStory;
 
@@ -165,7 +171,10 @@ const mdxToJsx = (node: any) => {
   return createElement(originalType, rest, ...jsxChildren);
 };
 
-export const jsxDecorator = (storyFn: any, context: StoryContext) => {
+export const jsxDecorator = (
+  storyFn: PartialStoryFn<ReactFramework>,
+  context: StoryContext<ReactFramework>
+) => {
   const story = storyFn();
 
   // We only need to render JSX if the source block is actually going to
@@ -183,7 +192,7 @@ export const jsxDecorator = (storyFn: any, context: StoryContext) => {
 
   // Exclude decorators from source code snippet by default
   const storyJsx = context?.parameters.docs?.source?.excludeDecorators
-    ? context.originalStoryFn(context.args)
+    ? (context.originalStoryFn as ArgsStoryFn<ReactFramework>)(context.args, context)
     : story;
 
   const sourceJsx = mdxToJsx(storyJsx);
diff --git a/addons/docs/src/frameworks/react/react-argtypes.stories.tsx b/addons/docs/src/frameworks/react/react-argtypes.stories.tsx
index 4006962f1fe0..8c5764331136 100644
--- a/addons/docs/src/frameworks/react/react-argtypes.stories.tsx
+++ b/addons/docs/src/frameworks/react/react-argtypes.stories.tsx
@@ -10,8 +10,8 @@ import { Component } from '../../blocks';
 
 const argsTableProps = (component: Component) => {
   const argTypes = extractArgTypes(component);
-  const parameters = { __isArgsStory: true, argTypes };
-  const rows = inferControls(({ parameters } as unknown) as StoryContext);
+  const parameters = { __isArgsStory: true };
+  const rows = inferControls(({ argTypes, parameters } as unknown) as StoryContext<any>);
   return { rows };
 };
 
diff --git a/addons/docs/src/frameworks/react/react-properties.test.ts b/addons/docs/src/frameworks/react/react-properties.test.ts
index 1417a82b5ab5..30bdc5aa69c1 100644
--- a/addons/docs/src/frameworks/react/react-properties.test.ts
+++ b/addons/docs/src/frameworks/react/react-properties.test.ts
@@ -3,8 +3,9 @@ import path from 'path';
 import fs from 'fs';
 
 import { transformFileSync, transformSync } from '@babel/core';
-import { inferControls } from '@storybook/client-api';
+import { inferControls } from '@storybook/store';
 import { StoryContext } from '@storybook/react';
+import { AnyFramework } from '@storybook/csf';
 import requireFromString from 'require-from-string';
 
 import { extractProps } from './extractProps';
@@ -69,8 +70,11 @@ describe('react component properties', () => {
 
           // snapshot the output of `extractArgTypes`
           const argTypes = extractArgTypes(component);
-          const parameters = { __isArgsStory: true, argTypes };
-          const rows = inferControls(({ parameters } as unknown) as StoryContext);
+          const parameters = { __isArgsStory: true };
+          const rows = inferControls(({
+            argTypes,
+            parameters,
+          } as unknown) as StoryContext<AnyFramework>);
           expect(rows).toMatchSpecificSnapshot(path.join(testDir, 'argTypes.snapshot'));
         });
       }
diff --git a/addons/docs/src/frameworks/svelte/extractArgTypes.test.ts b/addons/docs/src/frameworks/svelte/extractArgTypes.test.ts
index 0f96ae85552c..a67b08db35af 100644
--- a/addons/docs/src/frameworks/svelte/extractArgTypes.test.ts
+++ b/addons/docs/src/frameworks/svelte/extractArgTypes.test.ts
@@ -73,7 +73,8 @@ describe('Extracting Arguments', () => {
             "category": "events",
           },
           "type": Object {
-            "name": "void",
+            "name": "other",
+            "value": "void",
           },
         },
         "event_click": Object {
@@ -83,7 +84,8 @@ describe('Extracting Arguments', () => {
             "category": "events",
           },
           "type": Object {
-            "name": "void",
+            "name": "other",
+            "value": "void",
           },
         },
         "rounded": Object {
@@ -115,7 +117,8 @@ describe('Extracting Arguments', () => {
             "category": "slots",
           },
           "type": Object {
-            "name": "void",
+            "name": "other",
+            "value": "void",
           },
         },
         "text": Object {
diff --git a/addons/docs/src/frameworks/svelte/extractArgTypes.ts b/addons/docs/src/frameworks/svelte/extractArgTypes.ts
index 7105ad532366..ca1ee3e31522 100644
--- a/addons/docs/src/frameworks/svelte/extractArgTypes.ts
+++ b/addons/docs/src/frameworks/svelte/extractArgTypes.ts
@@ -1,4 +1,4 @@
-import { ArgTypes } from '@storybook/api';
+import { SBScalarType, StrictArgTypes } from '@storybook/csf';
 import { logger } from '@storybook/client-logger';
 import type {
   SvelteComponentDoc,
@@ -31,7 +31,7 @@ export const extractArgTypes: ArgTypesExtractor = (component: ComponentWithDocge
 };
 
 export const createArgTypes = (docgen: SvelteComponentDoc) => {
-  const results: ArgTypes = {};
+  const results: StrictArgTypes = {};
   docgen.data.forEach((item) => {
     results[item.name] = {
       control: parseTypeToControl(item.type),
@@ -39,7 +39,7 @@ export const createArgTypes = (docgen: SvelteComponentDoc) => {
       description: item.description,
       type: {
         required: hasKeyword('required', item.keywords),
-        name: item.type?.text,
+        name: item.type?.text as SBScalarType['name'],
       },
       table: {
         type: {
@@ -57,7 +57,7 @@ export const createArgTypes = (docgen: SvelteComponentDoc) => {
     results[`event_${item.name}`] = {
       name: item.name,
       description: item.description,
-      type: { name: 'void' },
+      type: { name: 'other', value: 'void' },
       table: {
         category: 'events',
       },
@@ -70,7 +70,7 @@ export const createArgTypes = (docgen: SvelteComponentDoc) => {
       description: [item.description, item.params?.map((p) => `\`${p.name}\``).join(' ')]
         .filter((p) => p)
         .join('\n\n'),
-      type: { name: 'void' },
+      type: { name: 'other', value: 'void' },
       table: {
         category: 'slots',
       },
diff --git a/addons/docs/src/frameworks/svelte/prepareForInline.ts b/addons/docs/src/frameworks/svelte/prepareForInline.ts
index 4d5b66403633..e8f4ce7c5ab2 100644
--- a/addons/docs/src/frameworks/svelte/prepareForInline.ts
+++ b/addons/docs/src/frameworks/svelte/prepareForInline.ts
@@ -1,11 +1,11 @@
-import { StoryFn } from '@storybook/addons';
+import { AnyFramework, StoryFn } from '@storybook/csf';
 
 import React from 'react';
 
 // @ts-ignore
 import HOC from './HOC.svelte';
 
-export const prepareForInline = (storyFn: StoryFn) => {
+export const prepareForInline = (storyFn: StoryFn<AnyFramework>) => {
   const el = React.useRef(null);
   React.useEffect(() => {
     const root = new HOC({
diff --git a/addons/docs/src/frameworks/svelte/sourceDecorator.ts b/addons/docs/src/frameworks/svelte/sourceDecorator.ts
index 40e96c0cbd1e..6eb2d41c2c0a 100644
--- a/addons/docs/src/frameworks/svelte/sourceDecorator.ts
+++ b/addons/docs/src/frameworks/svelte/sourceDecorator.ts
@@ -1,5 +1,5 @@
-import { addons, StoryContext } from '@storybook/addons';
-import { ArgTypes, Args } from '@storybook/api';
+import { addons } from '@storybook/addons';
+import { ArgTypes, Args, StoryContext, AnyFramework } from '@storybook/csf';
 
 import { SourceType, SNIPPET_RENDERED } from '../../shared';
 
@@ -8,7 +8,7 @@ import { SourceType, SNIPPET_RENDERED } from '../../shared';
  *
  * @param context StoryContext
  */
-const skipSourceRender = (context: StoryContext) => {
+const skipSourceRender = (context: StoryContext<AnyFramework>) => {
   const sourceParams = context?.parameters.docs?.source;
   const isArgsStory = context?.parameters.__isArgsStory;
 
@@ -144,7 +144,7 @@ function getWrapperProperties(component: any) {
  * @param storyFn Fn
  * @param context  StoryContext
  */
-export const sourceDecorator = (storyFn: any, context: StoryContext) => {
+export const sourceDecorator = (storyFn: any, context: StoryContext<AnyFramework>) => {
   const story = storyFn();
 
   if (skipSourceRender(context)) {
diff --git a/addons/docs/src/frameworks/vue/extractArgTypes.ts b/addons/docs/src/frameworks/vue/extractArgTypes.ts
index fbd1f85fc205..30457c43246f 100644
--- a/addons/docs/src/frameworks/vue/extractArgTypes.ts
+++ b/addons/docs/src/frameworks/vue/extractArgTypes.ts
@@ -1,4 +1,4 @@
-import { ArgTypes } from '@storybook/api';
+import { StrictArgTypes } from '@storybook/csf';
 import { ArgTypesExtractor, hasDocgen, extractComponentProps } from '../../lib/docgen';
 import { convert } from '../../lib/convert';
 
@@ -8,7 +8,7 @@ export const extractArgTypes: ArgTypesExtractor = (component) => {
   if (!hasDocgen(component)) {
     return null;
   }
-  const results: ArgTypes = {};
+  const results: StrictArgTypes = {};
   SECTIONS.forEach((section) => {
     const props = extractComponentProps(component, section);
     props.forEach(({ propDef, docgenInfo, jsDocTags }) => {
diff --git a/addons/docs/src/frameworks/vue/prepareForInline.ts b/addons/docs/src/frameworks/vue/prepareForInline.ts
index 98931991e96b..6627db5ec96f 100644
--- a/addons/docs/src/frameworks/vue/prepareForInline.ts
+++ b/addons/docs/src/frameworks/vue/prepareForInline.ts
@@ -1,6 +1,7 @@
 import React from 'react';
 import Vue from 'vue';
-import { StoryFn, StoryContext } from '@storybook/addons';
+import { StoryContext, PartialStoryFn } from '@storybook/csf';
+import { VueFramework } from '@storybook/vue';
 
 // Inspired by https://github.com/egoist/vue-to-react,
 // modified to store args as props in the root store
@@ -9,7 +10,10 @@ import { StoryFn, StoryContext } from '@storybook/addons';
 const COMPONENT = 'STORYBOOK_COMPONENT';
 const VALUES = 'STORYBOOK_VALUES';
 
-export const prepareForInline = (storyFn: StoryFn, { args }: StoryContext) => {
+export const prepareForInline = (
+  storyFn: PartialStoryFn<VueFramework>,
+  { args }: StoryContext<VueFramework>
+) => {
   const component = storyFn();
   const el = React.useRef(null);
 
diff --git a/addons/docs/src/frameworks/vue/sourceDecorator.ts b/addons/docs/src/frameworks/vue/sourceDecorator.ts
index 7965dd6beff1..0fbab7a1e4ea 100644
--- a/addons/docs/src/frameworks/vue/sourceDecorator.ts
+++ b/addons/docs/src/frameworks/vue/sourceDecorator.ts
@@ -1,14 +1,16 @@
 /* eslint no-underscore-dangle: ["error", { "allow": ["_vnode"] }] */
 
-import { addons, StoryContext } from '@storybook/addons';
+import { StoryContext } from '@storybook/csf';
+import { addons } from '@storybook/addons';
 import { logger } from '@storybook/client-logger';
 import prettier from 'prettier/standalone';
 import prettierHtml from 'prettier/parser-html';
 import type Vue from 'vue';
+import { VueFramework } from '@storybook/vue';
 
 import { SourceType, SNIPPET_RENDERED } from '../../shared';
 
-export const skipSourceRender = (context: StoryContext) => {
+export const skipSourceRender = (context: StoryContext<VueFramework>) => {
   const sourceParams = context?.parameters.docs?.source;
   const isArgsStory = context?.parameters.__isArgsStory;
 
@@ -22,7 +24,7 @@ export const skipSourceRender = (context: StoryContext) => {
   return !isArgsStory || sourceParams?.code || sourceParams?.type === SourceType.CODE;
 };
 
-export const sourceDecorator = (storyFn: any, context: StoryContext) => {
+export const sourceDecorator = (storyFn: any, context: StoryContext<VueFramework>) => {
   const story = storyFn();
 
   // See ../react/jsxDecorator.tsx
diff --git a/addons/docs/src/frameworks/vue3/extractArgTypes.ts b/addons/docs/src/frameworks/vue3/extractArgTypes.ts
index 1555b6d6acc0..039f0c7be2ae 100644
--- a/addons/docs/src/frameworks/vue3/extractArgTypes.ts
+++ b/addons/docs/src/frameworks/vue3/extractArgTypes.ts
@@ -1,4 +1,4 @@
-import { ArgTypes } from '@storybook/api';
+import { StrictArgTypes } from '@storybook/csf';
 import { ArgTypesExtractor, hasDocgen, extractComponentProps } from '../../lib/docgen';
 import { convert } from '../../lib/convert';
 
@@ -8,7 +8,7 @@ export const extractArgTypes: ArgTypesExtractor = (component) => {
   if (!hasDocgen(component)) {
     return null;
   }
-  const results: ArgTypes = {};
+  const results: StrictArgTypes = {};
   SECTIONS.forEach((section) => {
     const props = extractComponentProps(component, section);
     props.forEach(({ propDef, docgenInfo, jsDocTags }) => {
diff --git a/addons/docs/src/frameworks/vue3/prepareForInline.ts b/addons/docs/src/frameworks/vue3/prepareForInline.ts
index 096dcf1349e6..2f414504289e 100644
--- a/addons/docs/src/frameworks/vue3/prepareForInline.ts
+++ b/addons/docs/src/frameworks/vue3/prepareForInline.ts
@@ -1,12 +1,15 @@
 import React from 'react';
 import * as Vue from 'vue';
-import { StoryFn, StoryContext } from '@storybook/addons';
-import { app } from '@storybook/vue3';
+import { StoryContext, PartialStoryFn } from '@storybook/csf';
+import { app, VueFramework } from '@storybook/vue3';
 
 // This is cast as `any` to workaround type errors caused by Vue 2 types
 const { render, h } = Vue as any;
 
-export const prepareForInline = (storyFn: StoryFn, { args }: StoryContext) => {
+export const prepareForInline = (
+  storyFn: PartialStoryFn<VueFramework>,
+  { args }: StoryContext<VueFramework>
+) => {
   const component = storyFn();
 
   const vnode = h(component, args);
diff --git a/addons/docs/src/frameworks/web-components/prepareForInline.ts b/addons/docs/src/frameworks/web-components/prepareForInline.ts
index 715bdb8afef6..dcdb86c3432b 100644
--- a/addons/docs/src/frameworks/web-components/prepareForInline.ts
+++ b/addons/docs/src/frameworks/web-components/prepareForInline.ts
@@ -1,8 +1,10 @@
-import type { StoryFn } from '@storybook/addons';
+import type { PartialStoryFn } from '@storybook/csf';
+import { WebComponentsFramework } from '@storybook/web-components';
 import React from 'react';
+
 import { render } from 'lit-html';
 
-export const prepareForInline = (storyFn: StoryFn) => {
+export const prepareForInline = (storyFn: PartialStoryFn<WebComponentsFramework>) => {
   class Story extends React.Component {
     wrapperRef = React.createRef<HTMLElement>();
 
diff --git a/addons/docs/src/frameworks/web-components/sourceDecorator.ts b/addons/docs/src/frameworks/web-components/sourceDecorator.ts
index d1ab37f8ec2e..a6df7744944f 100644
--- a/addons/docs/src/frameworks/web-components/sourceDecorator.ts
+++ b/addons/docs/src/frameworks/web-components/sourceDecorator.ts
@@ -1,9 +1,12 @@
 /* global window */
 import { render } from 'lit-html';
-import { addons, StoryContext, StoryFn } from '@storybook/addons';
+import { ArgsStoryFn, PartialStoryFn, StoryContext } from '@storybook/csf';
+import { addons } from '@storybook/addons';
+import { WebComponentsFramework } from '@storybook/web-components';
+
 import { SNIPPET_RENDERED, SourceType } from '../../shared';
 
-function skipSourceRender(context: StoryContext) {
+function skipSourceRender(context: StoryContext<WebComponentsFramework>) {
   const sourceParams = context?.parameters.docs?.source;
   const isArgsStory = context?.parameters.__isArgsStory;
 
@@ -17,15 +20,21 @@ function skipSourceRender(context: StoryContext) {
   return !isArgsStory || sourceParams?.code || sourceParams?.type === SourceType.CODE;
 }
 
-function applyTransformSource(source: string, context: StoryContext): string {
+function applyTransformSource(
+  source: string,
+  context: StoryContext<WebComponentsFramework>
+): string {
   const { transformSource } = context.parameters.docs ?? {};
   if (typeof transformSource !== 'function') return source;
   return transformSource(source, context);
 }
 
-export function sourceDecorator(storyFn: StoryFn, context: StoryContext) {
+export function sourceDecorator(
+  storyFn: PartialStoryFn<WebComponentsFramework>,
+  context: StoryContext<WebComponentsFramework>
+) {
   const story = context?.parameters.docs?.source?.excludeDecorators
-    ? context.originalStoryFn(context.args)
+    ? (context.originalStoryFn as ArgsStoryFn<WebComponentsFramework>)(context.args, context)
     : storyFn();
 
   if (!skipSourceRender(context)) {
diff --git a/addons/docs/src/lib/convert/flow/convert.ts b/addons/docs/src/lib/convert/flow/convert.ts
index b9ef15867ea7..e61a8c6acd4d 100644
--- a/addons/docs/src/lib/convert/flow/convert.ts
+++ b/addons/docs/src/lib/convert/flow/convert.ts
@@ -1,5 +1,5 @@
 /* eslint-disable no-case-declarations */
-import { SBType } from '@storybook/client-api';
+import { SBType } from '@storybook/csf';
 import { FlowType, FlowSigType, FlowLiteralType } from './types';
 
 const isLiteral = (type: FlowType) => type.name === 'literal';
diff --git a/addons/docs/src/lib/convert/proptypes/convert.ts b/addons/docs/src/lib/convert/proptypes/convert.ts
index bc4da9573401..7a7bbd9b3f51 100644
--- a/addons/docs/src/lib/convert/proptypes/convert.ts
+++ b/addons/docs/src/lib/convert/proptypes/convert.ts
@@ -1,6 +1,6 @@
 /* eslint-disable no-case-declarations */
 import mapValues from 'lodash/mapValues';
-import { SBType } from '@storybook/client-api';
+import { SBType } from '@storybook/csf';
 import { PTType } from './types';
 import { trimQuotes } from '../utils';
 
diff --git a/addons/docs/src/lib/convert/typescript/convert.ts b/addons/docs/src/lib/convert/typescript/convert.ts
index 3affe3e70487..bae2e7235e9c 100644
--- a/addons/docs/src/lib/convert/typescript/convert.ts
+++ b/addons/docs/src/lib/convert/typescript/convert.ts
@@ -1,5 +1,5 @@
 /* eslint-disable no-case-declarations */
-import { SBType } from '@storybook/client-api';
+import { SBType } from '@storybook/csf';
 import { TSType, TSSigType } from './types';
 
 const convertSig = (type: TSSigType) => {
diff --git a/addons/docs/src/lib/docgen/types.ts b/addons/docs/src/lib/docgen/types.ts
index fcae772fa5b9..0fb5fa82373c 100644
--- a/addons/docs/src/lib/docgen/types.ts
+++ b/addons/docs/src/lib/docgen/types.ts
@@ -1,10 +1,10 @@
-import { ArgTypes } from '@storybook/api';
+import { StrictArgTypes } from '@storybook/csf';
 import { PropDef } from './PropDef';
 import { Component } from '../../blocks/types';
 
 export type PropsExtractor = (component: Component) => { rows?: PropDef[] } | null;
 
-export type ArgTypesExtractor = (component: Component) => ArgTypes | null;
+export type ArgTypesExtractor = (component: Component) => StrictArgTypes | null;
 
 export interface DocgenType {
   name: string;

From 9f799b2b2f9d1470e7dda123031ac5f4c6c0d0ce Mon Sep 17 00:00:00 2001
From: Tom Coleman <tom@thesnail.org>
Date: Sat, 11 Sep 2021 19:50:48 +1000
Subject: [PATCH 2/2] Update most <any> to nothing which defaults to
 <AnyFramework>

---
 addons/docs/package.json                 |  3 +++
 addons/docs/src/blocks/ArgsTable.tsx     | 10 +++++-----
 addons/docs/src/blocks/DocsContainer.tsx |  7 ++-----
 addons/docs/src/blocks/Meta.tsx          |  6 +++---
 addons/docs/src/blocks/Source.tsx        |  4 ++--
 addons/docs/src/blocks/Story.tsx         | 12 ++++++------
 addons/docs/src/blocks/Title.tsx         |  2 +-
 addons/docs/src/blocks/mdx.tsx           |  2 +-
 lib/preview-web/src/types.ts             |  2 +-
 lib/store/src/types.ts                   | 20 ++++++++++++--------
 10 files changed, 36 insertions(+), 32 deletions(-)

diff --git a/addons/docs/package.json b/addons/docs/package.json
index dfb57e93baad..2ec46af49c7c 100644
--- a/addons/docs/package.json
+++ b/addons/docs/package.json
@@ -160,6 +160,9 @@
     "@storybook/angular": {
       "optional": true
     },
+    "@storybook/html": {
+      "optional": true
+    },
     "@storybook/react": {
       "optional": true
     },
diff --git a/addons/docs/src/blocks/ArgsTable.tsx b/addons/docs/src/blocks/ArgsTable.tsx
index 7db03289d8ef..7dac7f5a9728 100644
--- a/addons/docs/src/blocks/ArgsTable.tsx
+++ b/addons/docs/src/blocks/ArgsTable.tsx
@@ -10,7 +10,7 @@ import {
 import { addons } from '@storybook/addons';
 import { filterArgTypes, PropDescriptor } from '@storybook/store';
 import Events from '@storybook/core-events';
-import { StrictArgTypes, Args } from '@storybook/csf';
+import { StrictArgTypes, Args, AnyFramework } from '@storybook/csf';
 
 import { DocsContext, DocsContextProps } from './DocsContext';
 import { Component, CURRENT_SELECTION, PRIMARY_STORY } from './types';
@@ -44,7 +44,7 @@ type ArgsTableProps = BaseProps | OfProps | ComponentsProps | StoryProps;
 
 const useArgs = (
   storyId: string,
-  context: DocsContextProps<any>
+  context: DocsContextProps
 ): [Args, (args: Args) => void, (argNames?: string[]) => void] => {
   const channel = addons.getChannel();
 
@@ -78,7 +78,7 @@ const useArgs = (
 
 export const extractComponentArgTypes = (
   component: Component,
-  { id, storyById }: DocsContextProps<any>,
+  { id, storyById }: DocsContextProps,
   include?: PropDescriptor,
   exclude?: PropDescriptor
 ): StrictArgTypes => {
@@ -99,7 +99,7 @@ const isShortcut = (value?: string) => {
 
 export const getComponent = (
   props: ArgsTableProps = {},
-  { id, storyById }: DocsContextProps<any>
+  { id, storyById }: DocsContextProps
 ): Component => {
   const { of } = props as OfProps;
   const { story } = props as StoryProps;
@@ -116,7 +116,7 @@ export const getComponent = (
 const addComponentTabs = (
   tabs: Record<string, PureArgsTableProps>,
   components: Record<string, Component>,
-  context: DocsContextProps<any>,
+  context: DocsContextProps,
   include?: PropDescriptor,
   exclude?: PropDescriptor,
   sort?: SortType
diff --git a/addons/docs/src/blocks/DocsContainer.tsx b/addons/docs/src/blocks/DocsContainer.tsx
index a3cf42db9ac0..66d904b1f5df 100644
--- a/addons/docs/src/blocks/DocsContainer.tsx
+++ b/addons/docs/src/blocks/DocsContainer.tsx
@@ -15,7 +15,7 @@ import { scrollToElement } from './utils';
 
 const { document, window: globalWindow } = global;
 
-export interface DocsContainerProps<TFramework extends AnyFramework> {
+export interface DocsContainerProps<TFramework extends AnyFramework = AnyFramework> {
   context: DocsContextProps<TFramework>;
 }
 
@@ -35,10 +35,7 @@ const warnOptionsTheme = deprecate(
 `
 );
 
-export const DocsContainer: FunctionComponent<DocsContainerProps<any>> = ({
-  context,
-  children,
-}) => {
+export const DocsContainer: FunctionComponent<DocsContainerProps> = ({ context, children }) => {
   const { id: storyId, storyById } = context;
   const {
     parameters: { options = {}, docs = {} },
diff --git a/addons/docs/src/blocks/Meta.tsx b/addons/docs/src/blocks/Meta.tsx
index d2aca741e5c4..206496266e87 100644
--- a/addons/docs/src/blocks/Meta.tsx
+++ b/addons/docs/src/blocks/Meta.tsx
@@ -1,14 +1,14 @@
 import React, { FC, useContext } from 'react';
 import global from 'global';
-import { BaseAnnotations } from '@storybook/csf';
+import { AnyFramework, BaseAnnotations } from '@storybook/csf';
 import { Anchor } from './Anchor';
 import { DocsContext, DocsContextProps } from './DocsContext';
 
 const { document } = global;
 
-type MetaProps = BaseAnnotations<any>;
+type MetaProps = BaseAnnotations;
 
-function getFirstStoryId(docsContext: DocsContextProps<any>): string {
+function getFirstStoryId(docsContext: DocsContextProps): string {
   const stories = docsContext.componentStories();
 
   return stories.length > 0 ? stories[0].id : null;
diff --git a/addons/docs/src/blocks/Source.tsx b/addons/docs/src/blocks/Source.tsx
index 0b23310cd8fd..1c521af34268 100644
--- a/addons/docs/src/blocks/Source.tsx
+++ b/addons/docs/src/blocks/Source.tsx
@@ -42,11 +42,11 @@ type NoneProps = CommonProps;
 
 type SourceProps = SingleSourceProps | MultiSourceProps | CodeProps | NoneProps;
 
-const getStory = (storyId: StoryId, docsContext: DocsContextProps<any>): Story<any> | null => {
+const getStory = (storyId: StoryId, docsContext: DocsContextProps): Story | null => {
   return docsContext.storyById(storyId);
 };
 
-const getSourceState = (storyIds: string[], docsContext: DocsContextProps<any>) => {
+const getSourceState = (storyIds: string[], docsContext: DocsContextProps) => {
   const states = storyIds
     .map((storyId) => {
       const story = getStory(storyId, docsContext);
diff --git a/addons/docs/src/blocks/Story.tsx b/addons/docs/src/blocks/Story.tsx
index 7e2abf1e5f50..7d88e533ef48 100644
--- a/addons/docs/src/blocks/Story.tsx
+++ b/addons/docs/src/blocks/Story.tsx
@@ -19,7 +19,7 @@ export const storyBlockIdFromId = (storyId: string) => `story--${storyId}`;
 type PureStoryProps = ComponentProps<typeof PureStory>;
 
 type Annotations = Pick<
-  StoryAnnotations<AnyFramework>,
+  StoryAnnotations,
   'decorators' | 'parameters' | 'args' | 'argTypes' | 'loaders'
 >;
 type CommonProps = Annotations & {
@@ -45,24 +45,24 @@ export type StoryProps = (StoryDefProps | StoryRefProps | StoryImportProps) & Co
 
 export const lookupStoryId = (
   storyName: string,
-  { mdxStoryNameToKey, mdxComponentAnnotations }: DocsContextProps<any>
+  { mdxStoryNameToKey, mdxComponentAnnotations }: DocsContextProps
 ) =>
   toId(
     mdxComponentAnnotations.id || mdxComponentAnnotations.title,
     storyNameFromExport(mdxStoryNameToKey[storyName])
   );
 
-export const getStoryId = (props: StoryProps, context: DocsContextProps<AnyFramework>): StoryId => {
+export const getStoryId = (props: StoryProps, context: DocsContextProps): StoryId => {
   const { id } = props as StoryRefProps;
   const { name } = props as StoryDefProps;
   const inputId = id === CURRENT_SELECTION ? context.id : id;
   return inputId || lookupStoryId(name, context);
 };
 
-export const getStoryProps = (
+export const getStoryProps = <TFramework extends AnyFramework>(
   { height, inline }: StoryProps,
-  story: StoryType<any>,
-  context: DocsContextProps<any>
+  story: StoryType<TFramework>,
+  context: DocsContextProps<TFramework>
 ): PureStoryProps => {
   const { name: storyName, parameters } = story;
   const { docs = {} } = parameters;
diff --git a/addons/docs/src/blocks/Title.tsx b/addons/docs/src/blocks/Title.tsx
index a815a1894861..100b8b2926df 100644
--- a/addons/docs/src/blocks/Title.tsx
+++ b/addons/docs/src/blocks/Title.tsx
@@ -8,7 +8,7 @@ interface TitleProps {
 
 const STORY_KIND_PATH_SEPARATOR = /\s*\/\s*/;
 
-export const extractTitle = ({ title }: DocsContextProps<any>) => {
+export const extractTitle = ({ title }: DocsContextProps) => {
   const groups = title.trim().split(STORY_KIND_PATH_SEPARATOR);
   return (groups && groups[groups.length - 1]) || title;
 };
diff --git a/addons/docs/src/blocks/mdx.tsx b/addons/docs/src/blocks/mdx.tsx
index 59566c729309..3dbc6746ee99 100644
--- a/addons/docs/src/blocks/mdx.tsx
+++ b/addons/docs/src/blocks/mdx.tsx
@@ -17,7 +17,7 @@ export const assertIsFn = (val: any) => {
 };
 
 // Hacky utility for adding mdxStoryToId to the default context
-export const AddContext: FC<DocsContextProps<any>> = (props) => {
+export const AddContext: FC<DocsContextProps> = (props) => {
   const { children, ...rest } = props;
   const parentContext = React.useContext(DocsContext);
   return (
diff --git a/lib/preview-web/src/types.ts b/lib/preview-web/src/types.ts
index db389e2518a5..95992f3eaf2f 100644
--- a/lib/preview-web/src/types.ts
+++ b/lib/preview-web/src/types.ts
@@ -8,7 +8,7 @@ export type WebProjectAnnotations<
   renderToDOM?: (context: RenderContext<TFramework>, element: Element) => Promise<void> | void;
 };
 
-export interface DocsContextProps<TFramework extends AnyFramework> {
+export interface DocsContextProps<TFramework extends AnyFramework = AnyFramework> {
   id: string;
   title: string;
   name: string;
diff --git a/lib/store/src/types.ts b/lib/store/src/types.ts
index fb8d4d4f9535..7484a518a9d6 100644
--- a/lib/store/src/types.ts
+++ b/lib/store/src/types.ts
@@ -26,21 +26,21 @@ export type ModuleExports = Record<string, any>;
 export type ModuleImportFn = (path: Path) => Promise<ModuleExports> | ModuleExports;
 
 export type NormalizedProjectAnnotations<
-  TFramework extends AnyFramework
+  TFramework extends AnyFramework = AnyFramework
 > = ProjectAnnotations<TFramework> & {
   argTypes?: StrictArgTypes;
   globalTypes?: StrictGlobalTypes;
 };
 
 export type NormalizedComponentAnnotations<
-  TFramework extends AnyFramework
+  TFramework extends AnyFramework = AnyFramework
 > = ComponentAnnotations<TFramework> & {
   // Useful to guarantee that id exists
   id: ComponentId;
   argTypes?: StrictArgTypes;
 };
 
-export type NormalizedStoryAnnotations<TFramework extends AnyFramework> = Omit<
+export type NormalizedStoryAnnotations<TFramework extends AnyFramework = AnyFramework> = Omit<
   StoryAnnotations<TFramework>,
   'storyName' | 'story'
 > & {
@@ -50,12 +50,14 @@ export type NormalizedStoryAnnotations<TFramework extends AnyFramework> = Omit<
   userStoryFn?: StoryFn<TFramework>;
 };
 
-export type CSFFile<TFramework extends AnyFramework> = {
+export type CSFFile<TFramework extends AnyFramework = AnyFramework> = {
   meta: NormalizedComponentAnnotations<TFramework>;
   stories: Record<StoryId, NormalizedStoryAnnotations<TFramework>>;
 };
 
-export type Story<TFramework extends AnyFramework> = StoryContextForEnhancers<TFramework> & {
+export type Story<
+  TFramework extends AnyFramework = AnyFramework
+> = StoryContextForEnhancers<TFramework> & {
   originalStoryFn: StoryFn<TFramework>;
   undecoratedStoryFn: LegacyStoryFn<TFramework>;
   unboundStoryFn: LegacyStoryFn<TFramework>;
@@ -63,11 +65,13 @@ export type Story<TFramework extends AnyFramework> = StoryContextForEnhancers<TF
   runPlayFunction: () => Promise<void>;
 };
 
-export type BoundStory<TFramework extends AnyFramework> = Story<TFramework> & {
+export type BoundStory<TFramework extends AnyFramework = AnyFramework> = Story<TFramework> & {
   storyFn: PartialStoryFn<TFramework>;
 };
 
-export declare type RenderContext<TFramework extends AnyFramework> = StoryIdentifier & {
+export declare type RenderContext<
+  TFramework extends AnyFramework = AnyFramework
+> = StoryIdentifier & {
   showMain: () => void;
   showError: (error: { title: string; description: string }) => void;
   showException: (err: Error) => void;
@@ -102,7 +106,7 @@ export interface Selection {
   viewMode: ViewMode;
 }
 
-export type DecoratorApplicator<TFramework extends AnyFramework> = (
+export type DecoratorApplicator<TFramework extends AnyFramework = AnyFramework> = (
   storyFn: LegacyStoryFn<TFramework>,
   decorators: DecoratorFunction<TFramework>[]
 ) => LegacyStoryFn<TFramework>;