diff --git a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleFileFunction.java b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleFileFunction.java
index 1bace018024857..523193892d5977 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleFileFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleFileFunction.java
@@ -245,7 +245,7 @@ public SkyValue compute(SkyKey skyKey, Environment env)
     } else {
       moduleThreadContext =
           execNonRegistryModuleFile(
-              moduleKey, starlarkSemantics, env, SymbolGenerator.create(skyKey));
+              moduleKey, state, starlarkSemantics, env, SymbolGenerator.create(skyKey));
       if (moduleThreadContext == null) {
         return null;
       }
@@ -319,7 +319,7 @@ private SkyValue computeForRootModule(
       }
     }
     var moduleThreadContext =
-        execNonRegistryModuleFile(ModuleKey.ROOT, starlarkSemantics, env, symbolGenerator);
+        execNonRegistryModuleFile(ModuleKey.ROOT, state, starlarkSemantics, env, symbolGenerator);
     if (moduleThreadContext == null) {
       return null;
     }
@@ -334,11 +334,11 @@ private SkyValue computeForRootModule(
   @Nullable
   private ModuleThreadContext execNonRegistryModuleFile(
       ModuleKey moduleKey,
+      State state,
       StarlarkSemantics starlarkSemantics,
       Environment env,
       SymbolGenerator<?> symbolGenerator)
       throws ModuleFileFunctionException, InterruptedException {
-    var state = env.getState(State::new);
     Preconditions.checkNotNull(state.compiledModuleFile);
     if (state.horizon == null) {
       state.horizon = state.compiledModuleFile.includeStatements();