diff --git a/packages/tailwindcss-language-service/src/util/state.ts b/packages/tailwindcss-language-service/src/util/state.ts
index c1e58e4..8ff1a3c 100644
--- a/packages/tailwindcss-language-service/src/util/state.ts
+++ b/packages/tailwindcss-language-service/src/util/state.ts
@@ -76,7 +76,7 @@ export interface State {
     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 96df13a..92902a7 100644
--- a/src/server.ts
+++ b/src/server.ts
@@ -438,6 +438,11 @@ async function createProjectService(
               resolveFrom(configDir, 'tailwindcss/lib/jit/lib/generateRules')
             ).generateRules,
           },
+          setupContext: {
+            module: __non_webpack_require__(
+              resolveFrom(configDir, 'tailwindcss/lib/jit/lib/setupContext')
+            ).default,
+          },
         }
       } catch (_) {
         try {
@@ -447,6 +452,11 @@ async function createProjectService(
                 resolveFrom(configDir, 'tailwindcss/jit/lib/generateRules')
               ).generateRules,
             },
+            setupContext: {
+              module: __non_webpack_require__(
+                resolveFrom(configDir, 'tailwindcss/jit/lib/setupContext')
+              ),
+            },
           }
         } catch (_) {}
       }
@@ -676,19 +686,7 @@ async function createProjectService(
     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')
-        )
-      }
-
-      state.jitContext = setupContext(state.configPath)(
+      state.jitContext = state.modules.jit.setupContext.module(state.configPath)(
         { opts: {}, messages: [] },
         state.modules.postcss.module.root()
       )