diff --git a/packages/tailwindcss-language-service/src/util/state.ts b/packages/tailwindcss-language-service/src/util/state.ts
index c1e58e40127b84f101396170caa6e46da7d3a014..8ff1a3c66f895cc068da9a437820728937a4512c 100644
--- a/packages/tailwindcss-language-service/src/util/state.ts
+++ b/packages/tailwindcss-language-service/src/util/state.ts
@@ -76,7 +76,7 @@     tailwindcss?: { version: string; module: any }
     postcss?: { version: string; module: Postcss }
     postcssSelectorParser?: { module: any }
     resolveConfig?: { module: any }
-    jit?: { generateRules: { module: any } }
+    jit?: { generateRules: { module: any }; setupContext: { module: any } }
   }
   browserslist?: string[]
   featureFlags?: FeatureFlags

diff --git a/src/server.ts b/src/server.ts
index 96df13a8e4086fedb884496f1aca1480d2fd66c3..92902a775f3e2473a86734383ee250319b8ca4dd 100644
--- a/src/server.ts
+++ b/src/server.ts
@@ -438,6 +438,11 @@             module: __non_webpack_require__(
               resolveFrom(configDir, 'tailwindcss/lib/jit/lib/generateRules')
             ).generateRules,
           },
+          setupContext: {
+            module: __non_webpack_require__(
+              resolveFrom(configDir, 'tailwindcss/lib/jit/lib/setupContext')
+            ).default,
+          },
         }
       } catch (_) {
         try {
@@ -446,6 +451,11 @@             generateRules: {
               module: __non_webpack_require__(
                 resolveFrom(configDir, 'tailwindcss/jit/lib/generateRules')
               ).generateRules,
+            },
+            setupContext: {
+              module: __non_webpack_require__(
+                resolveFrom(configDir, 'tailwindcss/jit/lib/setupContext')
+              ),
             },
           }
         } catch (_) {}
@@ -676,21 +686,9 @@     state.plugins = await getPlugins(config)
     state.classNames = (await extractClassNames(postcssResult.root)) as ClassNames
 
     if (state.jit) {
-      let setupContext
-
-      try {
-        setupContext = __non_webpack_require__(
-          resolveFrom(path.dirname(state.configPath), 'tailwindcss/lib/jit/lib/setupContext')
-        ).default
-      } catch (_) {
-        setupContext = __non_webpack_require__(
-          resolveFrom(path.dirname(state.configPath), 'tailwindcss/jit/lib/setupContext')
-import normalizePath from 'normalize-path'
   DocumentColorParams,
+  CompletionItem,
 import './lib/env'
-declare var __non_webpack_require__: typeof require
-
-      state.jitContext = setupContext(state.configPath)(
         { opts: {}, messages: [] },
         state.modules.postcss.module.root()
       )