chore(vite): Use vite plugin for context wrapping (#2058)

Author: Tobbe

packages/babel-config/src/api.ts

@@ -168,7 +168,10 @@ export const getApiSideBabelConfigPath = () => {
   }
 }
 
-export const getApiSideBabelOverrides = ({ projectIsEsm = false } = {}) => {
+export const getApiSideBabelOverrides = ({
+  projectIsEsm = false,
+  forJest = false,
+} = {}) => {
   const overrides = [
     // Extract graphql options from the graphql function
     // NOTE: this must come before the context wrapping
@@ -177,8 +180,9 @@ export const getApiSideBabelOverrides = ({ projectIsEsm = false } = {}) => {
       test: /.+api(?:[\\|/])src(?:[\\|/])functions(?:[\\|/])graphql\.(?:js|ts)$/,
       plugins: [pluginCedarGraphqlOptionsExtract, pluginCedarGqlormInject],
     },
-    // Apply context wrapping to all functions
-    {
+    // Apply context wrapping to all functions (Jest only; Vite uses
+    // cedarContextWrappingPlugin instead)
+    forJest && {
       // match */api/src/functions/*.js|ts
       test: /.+api(?:[\\|/])src(?:[\\|/])functions(?:[\\|/]).+.(?:js|ts)$/,
       plugins: [
@@ -200,11 +204,14 @@ export const getApiSideBabelOverrides = ({ projectIsEsm = false } = {}) => {
   return overrides as TransformOptions[]
 }
 
-export const getApiSideDefaultBabelConfig = ({ projectIsEsm = false } = {}) => {
+export const getApiSideDefaultBabelConfig = ({
+  projectIsEsm = false,
+  forJest = false,
+} = {}) => {
   return {
     presets: getApiSideBabelPresets(),
     plugins: getApiSideBabelPlugins({ projectIsEsm }),
-    overrides: getApiSideBabelOverrides({ projectIsEsm }),
+    overrides: getApiSideBabelOverrides({ projectIsEsm, forJest }),
     extends: getApiSideBabelConfigPath(),
     babelrc: false,
     ignore: ['node_modules'],

packages/internal/src/build/api.ts

@@ -14,6 +14,8 @@ import { getConfig, getPaths, projectSideIsEsm } from '@cedarjs/project-config'
 
 import { findApiFiles } from '../files.js'
 
+import { applyContextWrapping } from './esbuild-plugin-cedar-context-wrapping.js'
+
 let BUILD_CTX: BuildContext | null = null
 
 export const buildApi = async () => {
@@ -55,8 +57,20 @@ const runCedarBabelTransformsPlugin = {
         }),
       )
       if (transformedCode?.code) {
+        // Apply the context-wrapping safeguard to API function handlers. This
+        // is the standalone-esbuild equivalent of the Vite
+        // cedarContextWrappingPlugin and the (Jest-only) babel plugin it
+        // replaced.
+        const functionsDir = normalizePath(
+          path.join(getPaths().api.src, 'functions'),
+        )
+        const code = normalizePath(args.path).startsWith(functionsDir + '/')
+          ? (applyContextWrapping(transformedCode.code, {
+              projectIsEsm: projectSideIsEsm('api'),
+            }) ?? transformedCode.code)
+          : transformedCode.code
         return {
-          contents: transformedCode.code,
+          contents: code,
           loader: 'js',
         }
       }
@@ -99,8 +113,16 @@ function createCedarViteApiPlugin(): Plugin {
       )
 
       if (transformedCode?.code) {
+        const functionsDir = normalizePath(
+          path.join(cedarPaths.api.src, 'functions'),
+        )
+        const code = normalizePath(id).startsWith(functionsDir + '/')
+          ? (applyContextWrapping(transformedCode.code, {
+              projectIsEsm: isEsm,
+            }) ?? transformedCode.code)
+          : transformedCode.code
         return {
-          code: transformedCode.code,
+          code,
           map: transformedCode.map ?? null,
         }
       }

packages/internal/src/build/esbuild-plugin-cedar-context-wrapping.ts

@@ -0,0 +1,53 @@
+// Context-wrapping safeguard for API function handlers. This mirrors the Vite
+// cedarContextWrappingPlugin (used by buildCedarApp) and the Jest-only babel
+// plugin it replaced. The legacy esbuild build and the standalone-Vite build
+// (buildApiWithVite) don't go through the Vite plugin pipeline, so they apply
+// it directly here. Keep this in sync with
+// packages/vite/src/plugins/vite-plugin-cedar-context-wrapping.ts.
+
+export function applyContextWrapping(
+  code: string,
+  { projectIsEsm = false }: { projectIsEsm?: boolean } = {},
+): string | null {
+  const handlerRe =
+    /^export\s+(?:const|let|var)\s+handler(?:[^=]|=>)*?=(?![>=])/m
+
+  const handlerMatch = handlerRe.exec(code)
+  if (!handlerMatch) {
+    return null
+  }
+
+  const afterEquals = code
+    .slice(handlerMatch.index + handlerMatch[0].length)
+    .trimStart()
+  const isAsync = /^async(?:\s*[\(\*]|\s+function)/.test(afterEquals)
+
+  const storePath = projectIsEsm
+    ? '@cedarjs/context/dist/store.js'
+    : '@cedarjs/context/dist/store'
+
+  const importStatement = `import { getAsyncStoreInstance as __rw_getAsyncStoreInstance } from '${storePath}'\n`
+
+  const handlerStart = handlerMatch.index
+  const before = code.slice(0, handlerStart)
+  const after = code.slice(handlerStart)
+
+  const renamed = after.replace(handlerRe, 'const __rw_handler =')
+
+  const wrappedHandler =
+    `\nexport const handler = ${isAsync ? 'async ' : ''}(__rw_event, __rw__context) => {\n` +
+    `  // The store will be undefined if no context isolation has been performed yet\n` +
+    `  const __rw_contextStore = __rw_getAsyncStoreInstance().getStore()\n` +
+    `  if (__rw_contextStore === undefined) {\n` +
+    `    return __rw_getAsyncStoreInstance().run(\n` +
+    `      new Map(),\n` +
+    `      __rw_handler,\n` +
+    `      __rw_event,\n` +
+    `      __rw__context\n` +
+    `    )\n` +
+    `  }\n` +
+    `  return __rw_handler(__rw_event, __rw__context)\n` +
+    `}\n`
+
+  return before + importStatement + renamed + wrappedHandler
+}

packages/testing/src/config/jest/api/apiBabelConfig.ts

@@ -7,7 +7,9 @@ import {
 // Since configFile and babelrc is already passed a level up, cleaning up these keys here.
 // babelrc can not reside inside "extend"ed
 // Ref: packages/testing/config/jest/api/index.js
-const { babelrc: _b, ...defaultBabelConfig } = getApiSideDefaultBabelConfig()
+const { babelrc: _b, ...defaultBabelConfig } = getApiSideDefaultBabelConfig({
+  forJest: true,
+})
 
 type ConfigType = Omit<
   ReturnType<typeof getApiSideDefaultBabelConfig>,

packages/testing/src/config/jest/api/jest-preset.ts

@@ -7,7 +7,7 @@ import { getPaths } from '@cedarjs/project-config'
 
 const rwjsPaths = getPaths()
 const NODE_MODULES_PATH = path.join(rwjsPaths.base, 'node_modules')
-const { babelrc } = getApiSideDefaultBabelConfig()
+const { babelrc } = getApiSideDefaultBabelConfig({ forJest: true })
 
 const config: Config = {
   // To make sure other config option which depends on rootDir use

packages/vite/src/buildApp.ts

@@ -14,6 +14,7 @@ import { findApiFiles } from '@cedarjs/internal/dist/files.js'
 import { getConfig, getPaths, projectSideIsEsm } from '@cedarjs/project-config'
 
 import { getWorkspacePackageAliases } from './lib/workspacePackageAliases.js'
+import { cedarContextWrappingPlugin } from './plugins/vite-plugin-cedar-context-wrapping.js'
 
 function resolveWithExtensions(id: string): string {
   if (fs.existsSync(id)) {
@@ -325,44 +326,48 @@ export async function buildCedarApp({
         return null
       },
     })
+  }
 
-    if (babelPlugins) {
-      plugins.push({
-        name: 'cedar-vite-api-babel-transform',
-        enforce: 'pre',
-        async transform(code: string, id: string) {
-          if (!/\.(js|ts|tsx|jsx)$/.test(id)) {
-            return null
-          }
+  if (workspace.includes('api')) {
+    plugins.push(
+      cedarContextWrappingPlugin({ projectIsEsm: projectSideIsEsm('api') }),
+    )
+  }
 
-          if (id.includes('node_modules')) {
-            return null
-          }
+  if (babelPlugins) {
+    plugins.push({
+      name: 'cedar-vite-api-babel-transform',
+      enforce: 'pre',
+      async transform(code: string, id: string) {
+        if (!/\.(js|ts|tsx|jsx)$/.test(id)) {
+          return null
+        }
 
-          if (
-            !normalizePath(id).startsWith(normalizePath(cedarPaths.api.base))
-          ) {
-            return null
-          }
+        if (id.includes('node_modules')) {
+          return null
+        }
 
-          const transformedCode = await transformWithBabel(
-            code,
-            id,
-            babelPlugins,
-            true,
-          )
+        if (!normalizePath(id).startsWith(normalizePath(cedarPaths.api.base))) {
+          return null
+        }
 
-          if (transformedCode?.code) {
-            return {
-              code: transformedCode.code,
-              map: transformedCode.map ?? null,
-            }
+        const transformedCode = await transformWithBabel(
+          code,
+          id,
+          babelPlugins,
+          true,
+        )
+
+        if (transformedCode?.code) {
+          return {
+            code: transformedCode.code,
+            map: transformedCode.map ?? null,
           }
+        }
 
-          return null
-        },
-      })
-    }
+        return null
+      },
+    })
   }
 
   const builder = await createBuilder({

packages/vite/src/index.ts

@@ -25,6 +25,8 @@ import { cedarSwapApolloProvider } from './plugins/vite-plugin-swap-apollo-provi
 export { cedarAutoImportsPlugin } from './plugins/vite-plugin-cedar-auto-import.js'
 export { cedarCjsCompatPlugin } from './plugins/vite-plugin-cedar-cjs-compat.js'
 export { cedarCellTransform } from './plugins/vite-plugin-cedar-cell.js'
+export { cedarContextWrappingPlugin } from './plugins/vite-plugin-cedar-context-wrapping.js'
+export { applyContextWrapping } from './plugins/vite-plugin-cedar-context-wrapping.js'
 export { cedarEntryInjectionPlugin } from './plugins/vite-plugin-cedar-entry-injection.js'
 export { cedarHtmlEnvPlugin } from './plugins/vite-plugin-cedar-html-env.js'
 export { cedarImportDirPlugin } from './plugins/vite-plugin-cedar-import-dir.js'