diff --git a/packages/tailwindcss-language-service/src/util/jit.ts b/packages/tailwindcss-language-service/src/util/jit.ts
index 311319d0faf6ef7152d73644c767a32ef8c4f225..482dcf8284386389be6c230a14fc8f350250e2d4 100644
--- a/packages/tailwindcss-language-service/src/util/jit.ts
+++ b/packages/tailwindcss-language-service/src/util/jit.ts
@@ -13,24 +13,20 @@   let rules: [bigint, Rule][] = state.modules.jit.generateRules
     .module(new Set(classNames), state.jitContext)
     .sort(([a], [z]) => bigSign(a - z))
 
-  let actualRules: Rule[] = []
-
+  let root = state.modules.postcss.module.root({ nodes: rules.map(([, rule]) => rule) })
+  // @ts-ignore
 import { State } from './state'
 
 import { State } from './state'
-export function bigSign(bigIntValue) {
+import { remToPx } from './remToPx'
-import { State } from './state'
   // @ts-ignore
-    } else if (rule.walkRules) {
-      rule.walkRules((subRule) => {
 import type { Container, Root, Rule } from 'postcss'
-      })
-    }
+    actualRules.push(subRule)
-import type { Container, Root, Rule } from 'postcss'
+
 import dlv from 'dlv'
 
   return {
-    root: state.modules.postcss.module.root({ nodes: rules.map(([, rule]) => rule) }),
+    root,
     rules: actualRules,
   }
 }

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

diff --git a/src/server.ts b/src/server.ts
index 0a928f527f0cadbe8419bef46fd8e648818a989e..6f6e317c70a8a7e62409a8cb12c222afd3bbc61c 100644
--- a/src/server.ts
+++ b/src/server.ts
@@ -445,6 +445,11 @@             module: __non_webpack_require__(
               resolveFrom(configDir, 'tailwindcss/lib/jit/lib/setupContext')
             ).default,
           },
+          expandApplyAtRules: {
+            module: __non_webpack_require__(
+              resolveFrom(configDir, 'tailwindcss/lib/jit/lib/expandApplyAtRules')
+            ).default,
+          },
         }
       } catch (_) {
         try {
@@ -457,6 +462,11 @@             },
             setupContext: {
               module: __non_webpack_require__(
                 resolveFrom(configDir, 'tailwindcss/jit/lib/setupContext')
+              ),
+            },
+            expandApplyAtRules: {
+              module: __non_webpack_require__(
+                resolveFrom(configDir, 'tailwindcss/jit/lib/expandApplyAtRules')
               ),
             },
           }
@@ -491,6 +501,10 @@
     if (applyComplexClasses && !applyComplexClasses.default.__patched) {
       let _applyComplexClasses = applyComplexClasses.default
       applyComplexClasses.default = (config, ...args) => {
+        if (state.jit) {
+          return state.modules.jit.expandApplyAtRules.module(state.jitContext)
+        }
+
         let configClone = klona(config)
         configClone.separator = typeof state.separator === 'undefined' ? ':' : state.separator
 
@@ -629,6 +643,15 @@       hook.unhook()
       throw error
     }
 
+    if (state.jit) {
+      state.jitContext = state.modules.jit.setupContext.module(state.configPath)(
+        { opts: {}, messages: [] },
+        state.modules.postcss.module.root()
+      )
+      state.jitContext.tailwindConfig.separator =
+        typeof userSeperator === 'undefined' ? ':' : userSeperator
+    }
+
     let postcssResult: Result
     try {
       postcssResult = await postcss
@@ -689,13 +712,6 @@     state.config = resolveConfig.module(config)
     state.separator = typeof userSeperator === 'string' ? userSeperator : ':'
     state.plugins = await getPlugins(config)
     state.classNames = (await extractClassNames(postcssResult.root)) as ClassNames
-
-    if (state.jit) {
-      state.jitContext = state.modules.jit.setupContext.module(state.configPath)(
-        { opts: {}, messages: [] },
-        state.modules.postcss.module.root()
-      )
-    }
 
     state.variants = getVariants(state)