fix: define in environment config was not working (#18515)

Author: sapphi-red

packages/vite/src/node/config.ts

@@ -637,6 +637,7 @@ function resolveEnvironmentOptions(
   const consumer =
     options.consumer ?? (isClientEnvironment ? 'client' : 'server')
   return {
+    define: options.define,
     resolve,
     consumer,
     webCompatible: options.webCompatible ?? consumer === 'client',

packages/vite/src/node/plugins/clientInjections.ts

@@ -3,6 +3,7 @@ import type { Plugin } from '../plugin'
 import type { ResolvedConfig } from '../config'
 import { CLIENT_ENTRY, ENV_ENTRY } from '../constants'
 import { isObject, normalizePath, resolveHostname } from '../utils'
+import { usePerEnvironmentState } from '../environment'
 import { replaceDefine, serializeDefine } from './define'
 
 // ids in transform are normalized to unix style
@@ -16,6 +17,19 @@ const normalizedEnvEntry = normalizePath(ENV_ENTRY)
 export function clientInjectionsPlugin(config: ResolvedConfig): Plugin {
   let injectConfigValues: (code: string) => string
 
+  const getDefineReplacer = usePerEnvironmentState((environment) => {
+    const userDefine: Record<string, any> = {}
+    for (const key in environment.config.define) {
+      // import.meta.env.* is handled in `importAnalysis` plugin
+      if (!key.startsWith('import.meta.env.')) {
+        userDefine[key] = environment.config.define[key]
+      }
+    }
+    const serializedDefines = serializeDefine(userDefine)
+    const definesReplacement = () => serializedDefines
+    return (code: string) => code.replace(`__DEFINES__`, definesReplacement)
+  })
+
   return {
     name: 'vite:client-inject',
     async buildStart() {
@@ -51,18 +65,8 @@ export function clientInjectionsPlugin(config: ResolvedConfig): Plugin {
         hmrBase = path.posix.join(hmrBase, hmrConfig.path)
       }
 
-      const userDefine: Record<string, any> = {}
-      for (const key in config.define) {
-        // import.meta.env.* is handled in `importAnalysis` plugin
-        if (!key.startsWith('import.meta.env.')) {
-          userDefine[key] = config.define[key]
-        }
-      }
-      const serializedDefines = serializeDefine(userDefine)
-
       const modeReplacement = escapeReplacement(config.mode)
       const baseReplacement = escapeReplacement(devBase)
-      const definesReplacement = () => serializedDefines
       const serverHostReplacement = escapeReplacement(serverHost)
       const hmrProtocolReplacement = escapeReplacement(protocol)
       const hmrHostnameReplacement = escapeReplacement(host)
@@ -77,7 +81,6 @@ export function clientInjectionsPlugin(config: ResolvedConfig): Plugin {
         return code
           .replace(`__MODE__`, modeReplacement)
           .replace(/__BASE__/g, baseReplacement)
-          .replace(`__DEFINES__`, definesReplacement)
           .replace(`__SERVER_HOST__`, serverHostReplacement)
           .replace(`__HMR_PROTOCOL__`, hmrProtocolReplacement)
           .replace(`__HMR_HOSTNAME__`, hmrHostnameReplacement)
@@ -93,13 +96,14 @@ export function clientInjectionsPlugin(config: ResolvedConfig): Plugin {
       // TODO: Remove options?.ssr, Vitest currently hijacks this plugin
       const ssr = options?.ssr ?? this.environment.config.consumer === 'server'
       if (id === normalizedClientEntry || id === normalizedEnvEntry) {
-        return injectConfigValues(code)
+        const defineReplacer = getDefineReplacer(this)
+        return defineReplacer(injectConfigValues(code))
       } else if (!ssr && code.includes('process.env.NODE_ENV')) {
         // replace process.env.NODE_ENV instead of defining a global
         // for it to avoid shimming a `process` object during dev,
         // avoiding inconsistencies between dev and build
         const nodeEnv =
-          config.define?.['process.env.NODE_ENV'] ||
+          this.environment.config.define?.['process.env.NODE_ENV'] ||
           JSON.stringify(process.env.NODE_ENV || config.mode)
         return await replaceDefine(this.environment, code, id, {
           'process.env.NODE_ENV': nodeEnv,

packages/vite/src/node/plugins/define.ts

@@ -46,18 +46,18 @@ export function definePlugin(config: ResolvedConfig): Plugin {
     importMetaFallbackKeys['import.meta.env'] = `undefined`
   }
 
-  const userDefine: Record<string, string> = {}
-  const userDefineEnv: Record<string, any> = {}
-  for (const key in config.define) {
-    userDefine[key] = handleDefineValue(config.define[key])
-
-    // make sure `import.meta.env` object has user define properties
-    if (isBuild && key.startsWith('import.meta.env.')) {
-      userDefineEnv[key.slice(16)] = config.define[key]
+  function generatePattern(environment: Environment) {
+    const userDefine: Record<string, string> = {}
+    const userDefineEnv: Record<string, any> = {}
+    for (const key in environment.config.define) {
+      userDefine[key] = handleDefineValue(environment.config.define[key])
+
+      // make sure `import.meta.env` object has user define properties
+      if (isBuild && key.startsWith('import.meta.env.')) {
+        userDefineEnv[key.slice(16)] = environment.config.define[key]
+      }
     }
-  }
 
-  function generatePattern(environment: Environment) {
     const replaceProcessEnv = environment.config.webCompatible
 
     const define: Record<string, string> = {

playground/define/__tests__/define.spec.ts

@@ -3,6 +3,7 @@ import viteConfig from '../vite.config'
 import { page } from '~utils'
 
 const defines = viteConfig.define
+const envDefines = viteConfig.environments.client.define
 
 test('string', async () => {
   expect(await page.textContent('.exp')).toBe(
@@ -48,6 +49,9 @@ test('string', async () => {
   expect(await page.textContent('.define-in-dep')).toBe(
     defines.__STRINGIFIED_OBJ__,
   )
+  expect(await page.textContent('.define-in-environment')).toBe(
+    envDefines.__DEFINE_IN_ENVIRONMENT__,
+  )
 })
 
 test('ignores constants in string literals', async () => {

playground/define/index.html

@@ -21,6 +21,7 @@ <h1>Define</h1>
 <p>define variable in html: <code class="exp-define">__EXP__</code></p>
 <p>import json: <code class="import-json"></code></p>
 <p>define in dep: <code class="define-in-dep"></code></p>
+<p>define in environment: <code class="define-in-environment"></code></p>
 
 <h2>Define ignores string literals</h2>
 <section class="ignores-string-literals">
@@ -114,6 +115,8 @@ <h2>Define undefined constants on import.meta.env when it's a invalid json</h2>
   import { defined } from '@vitejs/test-commonjs-dep'
   text('.define-in-dep', JSON.stringify(defined))
 
+  text('.define-in-environment', JSON.stringify(__DEFINE_IN_ENVIRONMENT__))
+
   text('.ignores-string-literals .process-env-dot', 'process.env.')
   text(
     '.ignores-string-literals .global-process-env-dot',

playground/define/vite.config.js

@@ -31,4 +31,11 @@ export default defineConfig({
     __STRINGIFIED_OBJ__: JSON.stringify({ foo: true }),
     'import.meta.env.SOME_IDENTIFIER': '__VITE_SOME_IDENTIFIER__',
   },
+  environments: {
+    client: {
+      define: {
+        __DEFINE_IN_ENVIRONMENT__: '"defined only in client"',
+      },
+    },
+  },
 })