From 030f9d00edadd06b700cbaa438d4f1bb4fd55bdb Mon Sep 17 00:00:00 2001
From: Ian VanSchooten <ian.vanschooten@gmail.com>
Date: Tue, 13 Dec 2022 19:57:06 -0500
Subject: [PATCH] Vite: Use mdx2 babel pre-processing

---
 code/lib/builder-vite/package.json            |  3 +-
 .../builder-vite/src/plugins/mdx-plugin.ts    | 39 ++-----------------
 code/yarn.lock                                | 12 +++++-
 3 files changed, 14 insertions(+), 40 deletions(-)

diff --git a/code/lib/builder-vite/package.json b/code/lib/builder-vite/package.json
index dde8d2342084..dab66a485738 100644
--- a/code/lib/builder-vite/package.json
+++ b/code/lib/builder-vite/package.json
@@ -45,12 +45,11 @@
     "@storybook/client-logger": "7.0.0-beta.6",
     "@storybook/core-common": "7.0.0-beta.6",
     "@storybook/csf-plugin": "7.0.0-beta.6",
-    "@storybook/mdx2-csf": "next",
+    "@storybook/mdx2-csf": "0.0.4-canary.30.c978d7f.0",
     "@storybook/node-logger": "7.0.0-beta.6",
     "@storybook/preview": "7.0.0-beta.6",
     "@storybook/preview-api": "7.0.0-beta.6",
     "@storybook/types": "7.0.0-beta.6",
-    "@vitejs/plugin-react": "^2.0.0",
     "browser-assert": "^1.2.1",
     "es-module-lexer": "^0.9.3",
     "express": "^4.17.3",
diff --git a/code/lib/builder-vite/src/plugins/mdx-plugin.ts b/code/lib/builder-vite/src/plugins/mdx-plugin.ts
index ac16e629be22..5a00dab147ee 100644
--- a/code/lib/builder-vite/src/plugins/mdx-plugin.ts
+++ b/code/lib/builder-vite/src/plugins/mdx-plugin.ts
@@ -1,25 +1,9 @@
 import type { Options } from '@storybook/types';
 import type { Plugin } from 'vite';
 import { createFilter } from 'vite';
-import reactVite from '@vitejs/plugin-react';
 
 const isStorybookMdx = (id: string) => id.endsWith('stories.mdx') || id.endsWith('story.mdx');
 
-function injectRenderer(code: string) {
-  return `
-           import React from 'react';
-           ${code}
-           `;
-}
-
-// HACK: find a better way to do this, ideally avoiding @vitejs/plugin-react entirely.
-// We're just using it to run the mdx with jsx through babel
-// @ts-expect-error We're forcing the plugin shape here
-const viteBabel: Plugin | undefined = reactVite({ fastRefresh: false }).find(
-  // @ts-expect-error we know these have names, and what the shape will be
-  (p) => p.name === 'vite:react-babel'
-);
-
 /**
  * Storybook uses two different loaders when dealing with MDX:
  *
@@ -35,7 +19,7 @@ export function mdxPlugin(options: Options): Plugin {
   return {
     name: 'storybook:mdx-plugin',
     enforce: 'pre',
-    async transform(src, id, transformOptions) {
+    async transform(src, id) {
       if (!filter(id)) return undefined;
 
       const { compile } = await import('@storybook/mdx2-csf');
@@ -46,33 +30,16 @@ export function mdxPlugin(options: Options): Plugin {
         },
       });
 
-      const mdxCode = String(
+      const code = String(
         await compile(src, {
           skipCsf: !isStorybookMdx(id),
           ...mdxLoaderOptions,
         })
       );
 
-      const modifiedCode = injectRenderer(mdxCode);
-
-      // Hooks in recent rollup versions can be functions or objects, and though react hasn't changed, the typescript defs have
-      const rTransform = viteBabel?.transform;
-      const transform = rTransform && 'handler' in rTransform ? rTransform.handler : rTransform;
-
-      // It's safe to disable this, because we know it'll be there, since we added it ourselves.
-      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-      const result = await transform!.call(this, modifiedCode, `${id}.jsx`, transformOptions);
-
-      if (!result) return modifiedCode;
-
-      if (typeof result === 'string') return result;
-
-      const { code, map: resultMap } = result;
-
       return {
         code,
-        map:
-          !resultMap || typeof resultMap === 'string' ? resultMap : { ...resultMap, sources: [id] },
+        map: null, // TODO: update mdx2-csf to return the map
       };
     },
   };
diff --git a/code/yarn.lock b/code/yarn.lock
index 5fdab157e255..0c4015cdd55a 100644
--- a/code/yarn.lock
+++ b/code/yarn.lock
@@ -6021,14 +6021,13 @@ __metadata:
     "@storybook/client-logger": 7.0.0-beta.6
     "@storybook/core-common": 7.0.0-beta.6
     "@storybook/csf-plugin": 7.0.0-beta.6
-    "@storybook/mdx2-csf": next
+    "@storybook/mdx2-csf": 0.0.4-canary.30.c978d7f.0
     "@storybook/node-logger": 7.0.0-beta.6
     "@storybook/preview": 7.0.0-beta.6
     "@storybook/preview-api": 7.0.0-beta.6
     "@storybook/types": 7.0.0-beta.6
     "@types/express": ^4.17.13
     "@types/node": ^16.0.0
-    "@vitejs/plugin-react": ^2.0.0
     browser-assert: ^1.2.1
     es-module-lexer: ^0.9.3
     express: ^4.17.3
@@ -6779,6 +6778,15 @@ __metadata:
   languageName: unknown
   linkType: soft
 
+"@storybook/mdx2-csf@npm:0.0.4-canary.30.c978d7f.0":
+  version: 0.0.4-canary.30.c978d7f.0
+  resolution: "@storybook/mdx2-csf@npm:0.0.4-canary.30.c978d7f.0"
+  dependencies:
+    loader-utils: ^2.0.4
+  checksum: a9e8c36673f08ec294566535f61c20440ede8ccf834006506efd183ed0dc566dbb6863776088b628b8cd3f96662ed79c8c427de4617851eae5d60b3a0ae1b23f
+  languageName: node
+  linkType: hard
+
 "@storybook/mdx2-csf@npm:next":
   version: 0.1.0-next.7
   resolution: "@storybook/mdx2-csf@npm:0.1.0-next.7"