From 1bc0f87fcd05c655b6f1ed0ffcabad47e8026baa Mon Sep 17 00:00:00 2001
From: Salma Samy <salmasamy@google.com>
Date: Fri, 4 Aug 2023 03:17:32 +0300
Subject: [PATCH] [6.4.0] Lockfile updates & fixes (#19153)

* Add 'environ' attribute to module extensions and update lockfile

- Add 'environ' to module extensions
- Store the variables and their values in the lockfile
- Compare the variables and values with the lockfile and re-evaluate or error if they differ

PiperOrigin-RevId: 548964193
Change-Id: Ic2d52fe3332e93095c414d8bca4c9b4312bca8c2

# Conflicts:
#	src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtension.java
#	src/main/java/com/google/devtools/build/lib/bazel/bzlmod/SingleExtensionEvalFunction.java
#	src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkRepositoryModule.java
#	src/main/java/com/google/devtools/build/lib/skyframe/ActionEnvironmentFunction.java
#	src/main/java/com/google/devtools/build/lib/skyframe/ClientEnvironmentFunction.java
#	src/main/java/com/google/devtools/build/lib/starlarkbuildapi/repository/RepositoryModuleApi.java
#	src/main/java/com/google/devtools/build/skydoc/fakebuildapi/repository/FakeRepositoryModule.java
#	src/test/java/com/google/devtools/build/lib/bazel/bzlmod/StarlarkBazelModuleTest.java
#	src/test/py/bazel/bzlmod/bazel_lockfile_test.py

* Remove extra changes

* Re-run extension if its files change
fixes https://github.com/bazelbuild/bazel/issues/19068

PiperOrigin-RevId: 552823534
Change-Id: I87256b2cf954b932e24c70e22386020599f21a6f

* Do not fail on new fields added to lockfile data
Fixes https://github.com/bazelbuild/bazel/issues/19105

PiperOrigin-RevId: 553068023
Change-Id: I877bc8ece0641c01119a9295e09175a2d0a3a0c1

---------

Co-authored-by: Ian (Hee) Cha <heec@google.com>
---
 .../lib/bazel/bzlmod/BazelLockFileValue.java  |  31 ++-
 .../lib/bazel/bzlmod/GsonTypeAdapterUtil.java |  14 +
 .../bazel/bzlmod/LockFileModuleExtension.java |  30 ++-
 .../lib/bazel/bzlmod/ModuleExtension.java     |   8 +-
 .../bzlmod/SingleExtensionEvalFunction.java   | 139 ++++++++--
 .../starlark/StarlarkRepositoryModule.java    |  11 +
 .../rules/repository/RepositoryFunction.java  |  52 ++--
 .../skyframe/ActionEnvironmentFunction.java   |  19 +-
 .../skyframe/ClientEnvironmentFunction.java   |   3 +-
 .../repository/FakeRepositoryModule.java      |   3 +-
 .../bazel/bzlmod/StarlarkBazelModuleTest.java |   2 +
 .../py/bazel/bzlmod/bazel_lockfile_test.py    | 249 +++++++++++++-----
 12 files changed, 428 insertions(+), 133 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BazelLockFileValue.java b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BazelLockFileValue.java
index c0d0c06aed41e4..26c674fb6bc14f 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BazelLockFileValue.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BazelLockFileValue.java
@@ -109,25 +109,34 @@ public ImmutableList<String> getModuleAndFlagsDiff(
     return moduleDiff.build();
   }
 
+  /** Returns the differences between an extension and its locked data */
   public ImmutableList<String> getModuleExtensionDiff(
-      LockFileModuleExtension lockedExtension,
-      ImmutableMap<ModuleKey, ModuleExtensionUsage> lockedExtensionUsages,
       ModuleExtensionId extensionId,
       byte[] transitiveDigest,
-      ImmutableMap<ModuleKey, ModuleExtensionUsage> extensionUsages) {
+      boolean filesChanged,
+      ImmutableMap<String, String> envVariables,
+      ImmutableMap<ModuleKey, ModuleExtensionUsage> extensionUsages,
+      ImmutableMap<ModuleKey, ModuleExtensionUsage> lockedExtensionUsages) {
+    LockFileModuleExtension lockedExtension = getModuleExtensions().get(extensionId);
+
     ImmutableList.Builder<String> extDiff = new ImmutableList.Builder<>();
-    if (lockedExtension == null) {
-      extDiff.add("The module extension '" + extensionId + "' does not exist in the lockfile");
-    } else {
-      if (!Arrays.equals(transitiveDigest, lockedExtension.getBzlTransitiveDigest())) {
+    if (!Arrays.equals(transitiveDigest, lockedExtension.getBzlTransitiveDigest())) {
         extDiff.add(
             "The implementation of the extension '"
                 + extensionId
                 + "' or one of its transitive .bzl files has changed");
-      }
-      if (!extensionUsages.equals(lockedExtensionUsages)) {
-        extDiff.add("The usages of the extension named '" + extensionId + "' has changed");
-      }
+    }
+    if (filesChanged) {
+      extDiff.add("One or more files the extension '" + extensionId + "' is using have changed");
+    }
+    if (!extensionUsages.equals(lockedExtensionUsages)) {
+      extDiff.add("The usages of the extension '" + extensionId + "' has changed");
+    }
+    if (!envVariables.equals(lockedExtension.getEnvVariables())) {
+      extDiff.add(
+          "The environment variables the extension '"
+              + extensionId
+              + "' depends on (or their values) have changed");
     }
     return extDiff.build();
   }
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/GsonTypeAdapterUtil.java b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/GsonTypeAdapterUtil.java
index f0190f4c9fb66f..334484716f8bc2 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/GsonTypeAdapterUtil.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/GsonTypeAdapterUtil.java
@@ -91,6 +91,19 @@ public ModuleKey read(JsonReader jsonReader) throws IOException {
         }
       };
 
+  public static final TypeAdapter<Label> LABEL_TYPE_ADAPTER =
+      new TypeAdapter<Label>() {
+        @Override
+        public void write(JsonWriter jsonWriter, Label label) throws IOException {
+          jsonWriter.value(label.getUnambiguousCanonicalForm());
+        }
+
+        @Override
+        public Label read(JsonReader jsonReader) throws IOException {
+          return Label.parseCanonicalUnchecked(jsonReader.nextString());
+        }
+      };
+
   public static final TypeAdapter<ModuleExtensionId> MODULE_EXTENSION_ID_TYPE_ADAPTER =
       new TypeAdapter<>() {
         @Override
@@ -308,6 +321,7 @@ public static Gson createLockFileGson(Path moduleFilePath) {
         .registerTypeAdapterFactory(IMMUTABLE_SET)
         .registerTypeAdapterFactory(OPTIONAL)
         .registerTypeAdapterFactory(new LocationTypeAdapterFactory(moduleFilePath))
+        .registerTypeAdapter(Label.class, LABEL_TYPE_ADAPTER)
         .registerTypeAdapter(Version.class, VERSION_TYPE_ADAPTER)
         .registerTypeAdapter(ModuleKey.class, MODULE_KEY_TYPE_ADAPTER)
         .registerTypeAdapter(ModuleExtensionId.class, MODULE_EXTENSION_ID_TYPE_ADAPTER)
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/LockFileModuleExtension.java b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/LockFileModuleExtension.java
index 5431d6cf0fabe6..1b5eb71d60c35b 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/LockFileModuleExtension.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/LockFileModuleExtension.java
@@ -16,6 +16,7 @@
 
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableMap;
+import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.events.ExtendedEventHandler.Postable;
 import com.ryanharter.auto.value.gson.GenerateTypeAdapter;
 
@@ -28,13 +29,36 @@
 @GenerateTypeAdapter
 public abstract class LockFileModuleExtension implements Postable {
 
+  public static Builder builder() {
+    return new AutoValue_LockFileModuleExtension.Builder()
+        // TODO(salmasamy) can be removed when updating lockfile version
+        .setEnvVariables(ImmutableMap.of())
+        .setAccumulatedFileDigests(ImmutableMap.of());
+  }
+
   @SuppressWarnings("mutable")
   public abstract byte[] getBzlTransitiveDigest();
 
+  public abstract ImmutableMap<Label, String> getAccumulatedFileDigests();
+
+  public abstract ImmutableMap<String, String> getEnvVariables();
+
   public abstract ImmutableMap<String, RepoSpec> getGeneratedRepoSpecs();
 
-  public static LockFileModuleExtension create(
-      byte[] transitiveDigest, ImmutableMap<String, RepoSpec> generatedRepoSpecs) {
-    return new AutoValue_LockFileModuleExtension(transitiveDigest, generatedRepoSpecs);
+  public abstract Builder toBuilder();
+
+  /** Builder type for {@link LockFileModuleExtension}. */
+  @AutoValue.Builder
+  public abstract static class Builder {
+
+    public abstract Builder setBzlTransitiveDigest(byte[] digest);
+
+    public abstract Builder setAccumulatedFileDigests(ImmutableMap<Label, String> value);
+
+    public abstract Builder setEnvVariables(ImmutableMap<String, String> value);
+
+    public abstract Builder setGeneratedRepoSpecs(ImmutableMap<String, RepoSpec> value);
+
+    public abstract LockFileModuleExtension build();
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtension.java b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtension.java
index a55de9583b2545..99219efde8f595 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtension.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtension.java
@@ -16,6 +16,7 @@
 
 import com.google.auto.value.AutoValue;
 import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.events.EventHandler;
@@ -23,6 +24,7 @@
 import javax.annotation.Nullable;
 import net.starlark.java.annot.StarlarkBuiltin;
 import net.starlark.java.eval.StarlarkCallable;
+import net.starlark.java.eval.StarlarkValue;
 import net.starlark.java.syntax.Location;
 
 /**
@@ -30,7 +32,7 @@
  * or register toolchains and execution platforms.
  */
 @AutoValue
-public abstract class ModuleExtension {
+public abstract class ModuleExtension implements StarlarkValue {
   public abstract String getName();
 
   public abstract StarlarkCallable getImplementation();
@@ -43,6 +45,8 @@ public abstract class ModuleExtension {
 
   public abstract Location getLocation();
 
+  public abstract ImmutableList<String> getEnvVariables();
+
   public static Builder builder() {
     return new AutoValue_ModuleExtension.Builder();
   }
@@ -63,6 +67,8 @@ public abstract static class Builder {
 
     public abstract Builder setTagClasses(ImmutableMap<String, TagClass> value);
 
+    public abstract Builder setEnvVariables(ImmutableList<String> value);
+
     public abstract ModuleExtension build();
   }
 
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/SingleExtensionEvalFunction.java b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/SingleExtensionEvalFunction.java
index 8bb26e5a45f26a..4d72d021d9541e 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/SingleExtensionEvalFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/SingleExtensionEvalFunction.java
@@ -18,10 +18,12 @@
 import static com.google.common.collect.ImmutableBiMap.toImmutableBiMap;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 
+import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableTable;
+import com.google.devtools.build.lib.actions.FileValue;
 import com.google.devtools.build.lib.analysis.BlazeDirectories;
 import com.google.devtools.build.lib.bazel.repository.RepositoryOptions.LockfileMode;
 import com.google.devtools.build.lib.bazel.repository.downloader.DownloadManager;
@@ -32,6 +34,7 @@
 import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.rules.repository.NeedsSkyframeRestartException;
+import com.google.devtools.build.lib.rules.repository.RepositoryFunction;
 import com.google.devtools.build.lib.runtime.ProcessWrapper;
 import com.google.devtools.build.lib.runtime.RepositoryRemoteExecutor;
 import com.google.devtools.build.lib.server.FailureDetails.ExternalDeps;
@@ -41,20 +44,24 @@
 import com.google.devtools.build.lib.skyframe.BzlLoadValue;
 import com.google.devtools.build.lib.skyframe.PrecomputedValue;
 import com.google.devtools.build.lib.vfs.Path;
+import com.google.devtools.build.lib.vfs.RootedPath;
 import com.google.devtools.build.skyframe.SkyFunction;
 import com.google.devtools.build.skyframe.SkyFunctionException;
 import com.google.devtools.build.skyframe.SkyFunctionException.Transience;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
+import com.google.devtools.build.skyframe.SkyframeLookupResult;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.function.Function;
 import java.util.function.Supplier;
 import javax.annotation.Nullable;
 import net.starlark.java.eval.EvalException;
+import net.starlark.java.eval.Module;
 import net.starlark.java.eval.Mutability;
 import net.starlark.java.eval.Starlark;
 import net.starlark.java.eval.StarlarkList;
@@ -144,9 +151,16 @@ public SkyValue compute(SkyKey skyKey, Environment env)
           Transience.PERSISTENT);
     }
 
-    // Check the lockfile first for that module extension
+    ModuleExtension extension = ((ModuleExtension.InStarlark) exported).get();
+    ImmutableMap<String, String> extensionEnvVars =
+        RepositoryFunction.getEnvVarValues(env, extension.getEnvVariables());
+    if (extensionEnvVars == null) {
+      return null;
+    }
     byte[] bzlTransitiveDigest =
         BazelModuleContext.of(bzlLoadValue.getModule()).bzlTransitiveDigest();
+
+    // Check the lockfile first for that module extension
     LockfileMode lockfileMode = BazelLockFileFunction.LOCKFILE_MODE.get(env);
     if (!lockfileMode.equals(LockfileMode.OFF)) {
       BazelLockFileValue lockfile = (BazelLockFileValue) env.getValue(BazelLockFileValue.KEY);
@@ -155,20 +169,27 @@ public SkyValue compute(SkyKey skyKey, Environment env)
       }
       SingleExtensionEvalValue singleExtensionEvalValue =
           tryGettingValueFromLockFile(
-              extensionId, usagesValue, bzlTransitiveDigest, lockfileMode, lockfile);
+              env,
+              extensionId,
+              extensionEnvVars,
+              usagesValue,
+              bzlTransitiveDigest,
+              lockfileMode,
+              lockfile);
       if (singleExtensionEvalValue != null) {
         return singleExtensionEvalValue;
       }
     }
 
     // Run that extension!
-    ModuleExtension extension = ((ModuleExtension.InStarlark) exported).get();
-    ImmutableMap<String, RepoSpec> generatedRepoSpecs =
+    RunModuleExtensionResult moduleExtensionResult =
         runModuleExtension(
             extensionId, extension, usagesValue, bzlLoadValue.getModule(), starlarkSemantics, env);
-    if (generatedRepoSpecs == null) {
+    if (moduleExtensionResult == null) {
       return null;
     }
+    ImmutableMap<String, RepoSpec> generatedRepoSpecs =
+        moduleExtensionResult.getGeneratedRepoSpecs();
     // Check that all imported repos have been actually generated
     validateAllImportsAreGenerated(generatedRepoSpecs, usagesValue, extensionId);
 
@@ -177,20 +198,39 @@ public SkyValue compute(SkyKey skyKey, Environment env)
           .post(
               ModuleExtensionResolutionEvent.create(
                   extensionId,
-                  LockFileModuleExtension.create(bzlTransitiveDigest, generatedRepoSpecs)));
+                  LockFileModuleExtension.builder()
+                      .setBzlTransitiveDigest(bzlTransitiveDigest)
+                      .setAccumulatedFileDigests(moduleExtensionResult.getAccumulatedFileDigests())
+                      .setEnvVariables(extensionEnvVars)
+                      .setGeneratedRepoSpecs(generatedRepoSpecs)
+                      .build()));
     }
     return createSingleExtentionValue(generatedRepoSpecs, usagesValue);
   }
 
   @Nullable
   private SingleExtensionEvalValue tryGettingValueFromLockFile(
+      Environment env,
       ModuleExtensionId extensionId,
+      ImmutableMap<String, String> envVariables,
       SingleExtensionUsagesValue usagesValue,
       byte[] bzlTransitiveDigest,
       LockfileMode lockfileMode,
       BazelLockFileValue lockfile)
-      throws SingleExtensionEvalFunctionException {
+      throws SingleExtensionEvalFunctionException, InterruptedException {
     LockFileModuleExtension lockedExtension = lockfile.getModuleExtensions().get(extensionId);
+    if (lockedExtension == null) {
+      if (lockfileMode.equals(LockfileMode.ERROR)) {
+        throw new SingleExtensionEvalFunctionException(
+            ExternalDepsException.withMessage(
+                Code.BAD_MODULE,
+                "The module extension '%s' does not exist in the lockfile",
+                extensionId),
+            Transience.PERSISTENT);
+      }
+      return null;
+    }
+
     ImmutableMap<ModuleKey, ModuleExtensionUsage> lockedExtensionUsages;
     try {
       // TODO(salmasamy) might be nicer to precompute this table when we construct
@@ -202,19 +242,26 @@ private SingleExtensionEvalValue tryGettingValueFromLockFile(
       throw new SingleExtensionEvalFunctionException(e, Transience.PERSISTENT);
     }
 
-    // If we have the extension, check if the implementation and usage haven't changed
-    if (lockedExtension != null
+    Boolean filesChanged = didFilesChange(env, lockedExtension.getAccumulatedFileDigests());
+    if (filesChanged == null) { // still calculating file changes
+      return null;
+    }
+
+    // Check extension data in lockfile still valid
+    if (!filesChanged
         && Arrays.equals(bzlTransitiveDigest, lockedExtension.getBzlTransitiveDigest())
-        && usagesValue.getExtensionUsages().equals(lockedExtensionUsages)) {
+        && usagesValue.getExtensionUsages().equals(lockedExtensionUsages)
+        && envVariables.equals(lockedExtension.getEnvVariables())) {
       return createSingleExtentionValue(lockedExtension.getGeneratedRepoSpecs(), usagesValue);
     } else if (lockfileMode.equals(LockfileMode.ERROR)) {
       ImmutableList<String> extDiff =
           lockfile.getModuleExtensionDiff(
-              lockedExtension,
-              lockedExtensionUsages,
               extensionId,
               bzlTransitiveDigest,
-              usagesValue.getExtensionUsages());
+              filesChanged,
+              envVariables,
+              usagesValue.getExtensionUsages(),
+              lockedExtensionUsages);
       throw new SingleExtensionEvalFunctionException(
           ExternalDepsException.withMessage(
               Code.BAD_MODULE,
@@ -225,6 +272,43 @@ private SingleExtensionEvalValue tryGettingValueFromLockFile(
     return null;
   }
 
+  @Nullable
+  private Boolean didFilesChange(
+      Environment env, ImmutableMap<Label, String> accumulatedFileDigests)
+      throws InterruptedException {
+    // Turn labels into FileValue keys & get those values
+    Map<Label, FileValue.Key> fileKeys = new HashMap<>();
+    for (Label label : accumulatedFileDigests.keySet()) {
+      try {
+        RootedPath rootedPath = RepositoryFunction.getRootedPathFromLabel(label, env);
+        fileKeys.put(label, FileValue.key(rootedPath));
+      } catch (NeedsSkyframeRestartException e) {
+        return null;
+      } catch (EvalException e) {
+        // Consider those exception to be a cause for invalidation
+        return true;
+      }
+    }
+    SkyframeLookupResult result = env.getValuesAndExceptions(fileKeys.values());
+    if (env.valuesMissing()) {
+      return null;
+    }
+
+    // Compare the collected file values with the hashes stored in the lockfile
+    for (Entry<Label, String> entry : accumulatedFileDigests.entrySet()) {
+      FileValue fileValue = (FileValue) result.get(fileKeys.get(entry.getKey()));
+      try {
+        if (!entry.getValue().equals(RepositoryFunction.fileValueToMarkerValue(fileValue))) {
+          return true;
+        }
+      } catch (IOException e) {
+        // Consider those exception to be a cause for invalidation
+        return true;
+      }
+    }
+    return false;
+  }
+
   private SingleExtensionEvalValue createSingleExtentionValue(
       ImmutableMap<String, RepoSpec> generatedRepoSpecs, SingleExtensionUsagesValue usagesValue) {
     return SingleExtensionEvalValue.create(
@@ -301,11 +385,11 @@ private BzlLoadValue loadBzlFile(
   }
 
   @Nullable
-  private ImmutableMap<String, RepoSpec> runModuleExtension(
+  private RunModuleExtensionResult runModuleExtension(
       ModuleExtensionId extensionId,
       ModuleExtension extension,
       SingleExtensionUsagesValue usagesValue,
-      net.starlark.java.eval.Module module,
+      Module module,
       StarlarkSemantics starlarkSemantics,
       Environment env)
       throws SingleExtensionEvalFunctionException, InterruptedException {
@@ -316,12 +400,12 @@ private ImmutableMap<String, RepoSpec> runModuleExtension(
             BazelModuleContext.of(module).repoMapping(),
             directories,
             env.getListener());
+    ModuleExtensionContext moduleContext;
     try (Mutability mu =
         Mutability.create("module extension", usagesValue.getExtensionUniqueName())) {
       StarlarkThread thread = new StarlarkThread(mu, starlarkSemantics);
       thread.setPrintHandler(Event.makeDebugPrintHandler(env.getListener()));
-      ModuleExtensionContext moduleContext =
-          createContext(env, usagesValue, starlarkSemantics, extensionId, extension);
+      moduleContext = createContext(env, usagesValue, starlarkSemantics, extensionId, extension);
       threadContext.storeInThread(thread);
       try {
         Object returnValue =
@@ -371,7 +455,8 @@ private ImmutableMap<String, RepoSpec> runModuleExtension(
             Transience.TRANSIENT);
       }
     }
-    return threadContext.getGeneratedRepoSpecs();
+    return RunModuleExtensionResult.create(
+        moduleContext.getAccumulatedFileDigests(), threadContext.getGeneratedRepoSpecs());
   }
 
   private ModuleExtensionContext createContext(
@@ -423,4 +508,22 @@ static final class SingleExtensionEvalFunctionException extends SkyFunctionExcep
       super(cause, transience);
     }
   }
+
+  /* Holds the result data from running a module extension */
+  @AutoValue
+  abstract static class RunModuleExtensionResult {
+
+    abstract ImmutableMap<Label, String> getAccumulatedFileDigests();
+
+    abstract ImmutableMap<String, RepoSpec> getGeneratedRepoSpecs();
+
+    static RunModuleExtensionResult create(
+        ImmutableMap<Label, String> accumulatedFileDigests,
+        ImmutableMap<String, RepoSpec> generatedRepoSpecs) {
+      return new AutoValue_SingleExtensionEvalFunction_RunModuleExtensionResult(
+          accumulatedFileDigests, generatedRepoSpecs);
+    }
+  }
 }
+
+
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkRepositoryModule.java b/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkRepositoryModule.java
index 3dd3e9d8f6d67e..540b556c070f2e 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkRepositoryModule.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/repository/starlark/StarlarkRepositoryModule.java
@@ -296,6 +296,15 @@ public void failWithIncompatibleUseCcConfigureFromRulesCc(StarlarkThread thread)
                 "A description of the module extension that can be extracted by documentation"
                     + " generating tools.",
             named = true,
+            positional = false),
+        @Param(
+            name = "environ",
+            defaultValue = "[]",
+            doc =
+                "Provides a list of environment variable that this module extension depends on. If "
+                    + "an environment variable in that list changes, the extension will be "
+                    + "re-evaluated.",
+            named = true,
             positional = false)
       },
       useStarlarkThread = true)
@@ -303,6 +312,7 @@ public Object moduleExtension(
       StarlarkCallable implementation,
       Dict<?, ?> tagClasses, // Dict<String, TagClass>
       String doc,
+      Sequence<?> environ, // <String>
       StarlarkThread thread)
       throws EvalException {
     ModuleExtension.InStarlark inStarlark = new ModuleExtension.InStarlark();
@@ -314,6 +324,7 @@ public Object moduleExtension(
         .setDoc(doc)
         .setDefinitionEnvironmentLabel(
             BazelModuleContext.of(Module.ofInnermostEnclosingStarlarkFunction(thread)).label())
+        .setEnvVariables(ImmutableList.copyOf(Sequence.cast(environ, String.class, "environ")))
         .setLocation(thread.getCallerLocation());
     return inStarlark;
   }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/repository/RepositoryFunction.java b/src/main/java/com/google/devtools/build/lib/rules/repository/RepositoryFunction.java
index 55084b91d90d3a..26d04ae28c2288 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/repository/RepositoryFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/repository/RepositoryFunction.java
@@ -20,6 +20,7 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.io.BaseEncoding;
 import com.google.devtools.build.lib.actions.FileValue;
 import com.google.devtools.build.lib.analysis.BlazeDirectories;
@@ -50,6 +51,7 @@
 import com.google.devtools.build.skyframe.SkyFunctionException.Transience;
 import com.google.devtools.build.skyframe.SkyKey;
 import java.io.IOException;
+import java.util.Collection;
 import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.Objects;
@@ -170,9 +172,9 @@ public abstract RepositoryDirectoryValue.Builder fetch(
       throws InterruptedException, RepositoryFunctionException;
 
   @SuppressWarnings("unchecked")
-  private static Iterable<String> getEnviron(Rule rule) {
+  private static Collection<String> getEnviron(Rule rule) {
     if (rule.isAttrDefined("$environ", Type.STRING_LIST)) {
-      return (Iterable<String>) rule.getAttr("$environ");
+      return (Collection<String>) rule.getAttr("$environ");
     }
     return ImmutableList.of();
   }
@@ -183,7 +185,7 @@ private static Iterable<String> getEnviron(Rule rule) {
    * is needed.
    */
   public boolean verifyMarkerData(Rule rule, Map<String, String> markerData, Environment env)
-      throws InterruptedException, RepositoryFunctionException {
+      throws InterruptedException {
     return verifyEnvironMarkerData(markerData, env, getEnviron(rule))
         && verifyMarkerDataForFiles(rule, markerData, env)
         && verifySemanticsMarkerData(markerData, env);
@@ -288,33 +290,37 @@ public static RootedPath getRootedPathFromLabel(Label label, Environment env)
   protected Map<String, String> declareEnvironmentDependencies(
       Map<String, String> markerData, Environment env, Iterable<String> keys)
       throws InterruptedException {
-    Map<String, String> environ = ActionEnvironmentFunction.getEnvironmentView(env, keys);
+    ImmutableMap<String, String> envDep = getEnvVarValues(env, keys);
+    if (envDep == null) {
+      return null;
+    }
+    // Add the dependencies to the marker file
+    keys.forEach(key -> markerData.put("ENV:" + key, envDep.get(key)));
+    return envDep;
+  }
 
-    // Returns true if there is a null value and we need to wait for some dependencies.
+  @Nullable
+  public static ImmutableMap<String, String> getEnvVarValues(Environment env, Iterable<String> keys)
+      throws InterruptedException {
+    ImmutableMap<String, String> environ = ActionEnvironmentFunction.getEnvironmentView(env, keys);
     if (environ == null) {
       return null;
     }
-
     Map<String, String> repoEnvOverride = PrecomputedValue.REPO_ENV.get(env);
     if (repoEnvOverride == null) {
       return null;
     }
 
     // Only depend on --repo_env values that are specified in the "environ" attribute.
-    Map<String, String> repoEnv = new LinkedHashMap<String, String>(environ);
+    ImmutableMap.Builder<String, String> repoEnv = ImmutableMap.builder();
+    repoEnv.putAll(environ);
     for (String key : keys) {
       String value = repoEnvOverride.get(key);
       if (value != null) {
         repoEnv.put(key, value);
       }
     }
-
-    // Add the dependencies to the marker file
-    for (Map.Entry<String, String> value : repoEnv.entrySet()) {
-      markerData.put("ENV:" + value.getKey(), value.getValue());
-    }
-
-    return repoEnv;
+    return repoEnv.buildKeepingLast();
   }
 
   /**
@@ -323,9 +329,9 @@ protected Map<String, String> declareEnvironmentDependencies(
    * Environment)} function to verify the values for environment variables.
    */
   protected boolean verifyEnvironMarkerData(
-      Map<String, String> markerData, Environment env, Iterable<String> keys)
+      Map<String, String> markerData, Environment env, Collection<String> keys)
       throws InterruptedException {
-    Map<String, String> environ = ActionEnvironmentFunction.getEnvironmentView(env, keys);
+    ImmutableMap<String, String> environ = ActionEnvironmentFunction.getEnvironmentView(env, keys);
     if (env.valuesMissing()) {
       return false; // Returns false so caller knows to return immediately
     }
@@ -346,17 +352,18 @@ protected boolean verifyEnvironMarkerData(
 
     // Verify that all environment variable in the marker file are also in keys
     for (String key : markerData.keySet()) {
-      if (key.startsWith("ENV:") && !repoEnv.containsKey(key.substring(4))) {
+      if (key.startsWith("ENV:") && !keys.contains(key.substring(4))) {
         return false;
       }
     }
+
     // Now verify the values of the marker data
-    for (Map.Entry<String, String> value : repoEnv.entrySet()) {
-      if (!markerData.containsKey("ENV:" + value.getKey())) {
+    for (String key : keys) {
+      if (!markerData.containsKey("ENV:" + key)) {
         return false;
       }
-      String markerValue = markerData.get("ENV:" + value.getKey());
-      if (!Objects.equals(markerValue, value.getValue())) {
+      String markerValue = markerData.get("ENV:" + key);
+      if (!Objects.equals(markerValue, repoEnv.get(key))) {
         return false;
       }
     }
@@ -439,8 +446,7 @@ protected static String getPathAttr(Rule rule) throws RepositoryFunctionExceptio
   }
 
   @VisibleForTesting
-  protected static PathFragment getTargetPath(String userDefinedPath, Path workspace)
-      throws RepositoryFunctionException {
+  protected static PathFragment getTargetPath(String userDefinedPath, Path workspace) {
     PathFragment pathFragment = PathFragment.create(userDefinedPath);
     return workspace.getRelative(pathFragment).asFragment();
   }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ActionEnvironmentFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/ActionEnvironmentFunction.java
index 25d12db34afbcd..4edf0fb1b31078 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ActionEnvironmentFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ActionEnvironmentFunction.java
@@ -15,6 +15,7 @@
 package com.google.devtools.build.lib.skyframe;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Interner;
 import com.google.devtools.build.lib.bugreport.BugReport;
 import com.google.devtools.build.lib.concurrent.BlazeInterners;
@@ -25,8 +26,6 @@
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
 import com.google.devtools.build.skyframe.SkyframeLookupResult;
-import java.util.Collections;
-import java.util.LinkedHashMap;
 import java.util.Map;
 import javax.annotation.Nullable;
 
@@ -47,7 +46,7 @@ public SkyValue compute(SkyKey skyKey, Environment env) throws InterruptedExcept
     return env.getValue(ClientEnvironmentFunction.key(key));
   }
 
-  /** @return the SkyKey to invoke this function for the environment variable {@code variable}. */
+  /** Returns the SkyKey to invoke this function for the environment variable {@code variable}. */
   public static Key key(String variable) {
     return Key.create(variable);
   }
@@ -78,8 +77,8 @@ public SkyFunctionName functionName() {
    * if and only if some dependencies from Skyframe still need to be resolved.
    */
   @Nullable
-  public static Map<String, String> getEnvironmentView(Environment env, Iterable<String> keys)
-      throws InterruptedException {
+  public static ImmutableMap<String, String> getEnvironmentView(
+      Environment env, Iterable<String> keys) throws InterruptedException {
     ImmutableList.Builder<SkyKey> skyframeKeysBuilder = ImmutableList.builder();
     for (String key : keys) {
       skyframeKeysBuilder.add(key(key));
@@ -89,8 +88,8 @@ public static Map<String, String> getEnvironmentView(Environment env, Iterable<S
     if (env.valuesMissing()) {
       return null;
     }
-    // To return the initial order and support null values, we use a LinkedHashMap.
-    LinkedHashMap<String, String> result = new LinkedHashMap<>();
+
+    ImmutableMap.Builder<String, String> result = ImmutableMap.builder();
     for (SkyKey key : skyframeKeys) {
       ClientEnvironmentValue value = (ClientEnvironmentValue) values.get(key);
       if (value == null) {
@@ -99,8 +98,10 @@ public static Map<String, String> getEnvironmentView(Environment env, Iterable<S
                 "ClientEnvironmentValue " + key + " was missing, this should never happen"));
         return null;
       }
-      result.put(key.argument().toString(), value.getValue());
+      if (value.getValue() != null) {
+        result.put(key.argument().toString(), value.getValue());
+      }
     }
-    return Collections.unmodifiableMap(result);
+    return result.buildOrThrow();
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ClientEnvironmentFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/ClientEnvironmentFunction.java
index 60bc5ad34481b0..510cfce6886f8a 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ClientEnvironmentFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ClientEnvironmentFunction.java
@@ -32,6 +32,7 @@ public static Key key(String keyString) {
     return ClientEnvironmentFunction.Key.create(keyString);
   }
 
+  /** The Skyframe key for the client environment function. */
   @AutoCodec.VisibleForSerialization
   @AutoCodec
   static class Key extends AbstractSkyKey<String> {
@@ -61,7 +62,7 @@ public ClientEnvironmentFunction(AtomicReference<Map<String, String>> clientEnv)
 
   @Nullable
   @Override
-  public SkyValue compute(SkyKey key, Environment env) throws InterruptedException {
+  public SkyValue compute(SkyKey key, Environment env) {
     return new ClientEnvironmentValue(clientEnv.get().get((String) key.argument()));
   }
 }
diff --git a/src/main/java/com/google/devtools/build/skydoc/fakebuildapi/repository/FakeRepositoryModule.java b/src/main/java/com/google/devtools/build/skydoc/fakebuildapi/repository/FakeRepositoryModule.java
index 8f9edc31cb8be9..2cbf475276304a 100644
--- a/src/main/java/com/google/devtools/build/skydoc/fakebuildapi/repository/FakeRepositoryModule.java
+++ b/src/main/java/com/google/devtools/build/skydoc/fakebuildapi/repository/FakeRepositoryModule.java
@@ -72,7 +72,6 @@ public StarlarkCallable repositoryRule(
       String doc,
       StarlarkThread thread)
       throws EvalException {
-    List<AttributeInfo> attrInfos;
     ImmutableMap.Builder<String, FakeDescriptor> attrsMapBuilder = ImmutableMap.builder();
     if (attrs != null && attrs != Starlark.NONE) {
       attrsMapBuilder.putAll(Dict.cast(attrs, String.class, FakeDescriptor.class, "attrs"));
@@ -80,7 +79,7 @@ public StarlarkCallable repositoryRule(
 
     attrsMapBuilder.put("name", IMPLICIT_NAME_ATTRIBUTE_DESCRIPTOR);
     attrsMapBuilder.put("repo_mapping", IMPLICIT_REPO_MAPPING_ATTRIBUTE_DESCRIPTOR);
-    attrInfos =
+    List<AttributeInfo> attrInfos =
         attrsMapBuilder.build().entrySet().stream()
             .filter(entry -> !entry.getKey().startsWith("_"))
             .map(entry -> entry.getValue().asAttributeInfo(entry.getKey()))
diff --git a/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/StarlarkBazelModuleTest.java b/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/StarlarkBazelModuleTest.java
index f20e7c092d8c62..6e0dd26752ed64 100644
--- a/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/StarlarkBazelModuleTest.java
+++ b/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/StarlarkBazelModuleTest.java
@@ -24,6 +24,7 @@
 import static org.junit.Assert.assertThrows;
 
 import com.google.common.collect.ImmutableBiMap;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.cmdline.Label;
@@ -61,6 +62,7 @@ private static ModuleExtension.Builder getBaseExtensionBuilder() {
         .setName("maven")
         .setDoc("")
         .setLocation(Location.BUILTIN)
+        .setEnvVariables(ImmutableList.of())
         .setDefinitionEnvironmentLabel(Label.parseAbsoluteUnchecked("//:rje.bzl"))
         .setImplementation(() -> "maven");
   }
diff --git a/src/test/py/bazel/bzlmod/bazel_lockfile_test.py b/src/test/py/bazel/bzlmod/bazel_lockfile_test.py
index f145c633d301b5..ce5700804d38f9 100644
--- a/src/test/py/bazel/bzlmod/bazel_lockfile_test.py
+++ b/src/test/py/bazel/bzlmod/bazel_lockfile_test.py
@@ -282,23 +282,27 @@ def testModuleExtension(self):
             '        for dep in mod.tags.dep:',
             '            print("Name:", dep.name, ", Versions:", dep.versions)',
             '',
-            (
-                '_dep = tag_class(attrs={"name": attr.string(), "versions":'
-                ' attr.string_list()})'
-            ),
+            '_dep = tag_class(attrs={"name": attr.string(), "versions":',
+            ' attr.string_list()})',
             'lockfile_ext = module_extension(',
             '    implementation=_module_ext_impl,',
             '    tag_classes={"dep": _dep},',
+            '    environ=["GREEN_TREES", "NOT_SET"],',
             ')',
         ],
     )
 
-    _, _, stderr = self.RunBazel(['build', '@hello//:all'])
+    # Only set one env var, to make sure null variables don't crash
+    _, _, stderr = self.RunBazel(
+        ['build', '@hello//:all'], env_add={'GREEN_TREES': 'In the city'}
+    )
     self.assertIn('Hello from the other side!', ''.join(stderr))
     self.assertIn('Name: bmbm , Versions: ["v1", "v2"]', ''.join(stderr))
 
     self.RunBazel(['shutdown'])
-    _, _, stderr = self.RunBazel(['build', '@hello//:all'])
+    _, _, stderr = self.RunBazel(
+        ['build', '@hello//:all'], env_add={'GREEN_TREES': 'In the city'}
+    )
     self.assertNotIn('Hello from the other side!', ''.join(stderr))
 
   def testIsolatedModuleExtension(self):
@@ -403,13 +407,10 @@ def testUpdateModuleExtension(self):
             '    ctx.file("WORKSPACE")',
             '    ctx.file("BUILD", "filegroup(name=\\"lala\\")")',
             'repo_rule = repository_rule(implementation = _repo_rule_impl)',
-            'def _module_ext_impl(ctx):',
+            'def _mod_ext_impl(ctx):',
             '    print("Hello from the other side!")',
             '    repo_rule(name= "hello")',
-            (
-                'lockfile_ext = module_extension(implementation ='
-                ' _module_ext_impl)'
-            ),
+            'lockfile_ext = module_extension(implementation = _mod_ext_impl)',
         ],
     )
     _, _, stderr = self.RunBazel(['build', '@hello//:all'])
@@ -434,13 +435,10 @@ def testUpdateModuleExtension(self):
             '    ctx.file("WORKSPACE")',
             '    ctx.file("BUILD", "filegroup(name=\\"lala\\")")',
             'repo_rule = repository_rule(implementation = _repo_rule_impl)',
-            'def _module_ext_impl(ctx):',
+            'def _mod_ext_impl(ctx):',
             '    print("Hello from the other town!")',
             '    repo_rule(name= "hello")',
-            (
-                'lockfile_ext = module_extension(implementation ='
-                ' _module_ext_impl)'
-            ),
+            'lockfile_ext = module_extension(implementation = _mod_ext_impl)',
         ],
     )
     _, _, stderr = self.RunBazel(['build', '@hello//:all'])
@@ -468,13 +466,10 @@ def testUpdateModuleExtensionErrorMode(self):
             '    ctx.file("WORKSPACE")',
             '    ctx.file("BUILD", "filegroup(name=\\"lala\\")")',
             'repo_rule = repository_rule(implementation = _repo_rule_impl)',
-            'def _module_ext_impl(ctx):',
+            'def _mod_ext_impl(ctx):',
             '    print("Hello from the other side!")',
             '    repo_rule(name= "hello")',
-            (
-                'lockfile_ext = module_extension(implementation ='
-                ' _module_ext_impl)'
-            ),
+            'lockfile_ext = module_extension(implementation = _mod_ext_impl)',
         ],
     )
     _, _, stderr = self.RunBazel(['build', '@hello//:all'])
@@ -489,13 +484,10 @@ def testUpdateModuleExtensionErrorMode(self):
             '    ctx.file("WORKSPACE")',
             '    ctx.file("BUILD", "filegroup(name=\\"lalo\\")")',
             'repo_rule = repository_rule(implementation = _repo_rule_impl)',
-            'def _module_ext_impl(ctx):',
+            'def _mod_ext_impl(ctx):',
             '    print("Hello from the other town!")',
             '    repo_rule(name= "hello")',
-            (
-                'lockfile_ext = module_extension(implementation ='
-                ' _module_ext_impl)'
-            ),
+            'lockfile_ext = module_extension(implementation = _mod_ext_impl)',
         ],
     )
 
@@ -554,42 +546,42 @@ def testRemoveModuleExtensionsNotUsed(self):
 
   def testNoAbsoluteRootModuleFilePath(self):
     self.ScratchFile(
-        'MODULE.bazel',
-        [
-            'ext = use_extension("extension.bzl", "ext")',
-            'ext.dep(generate = True)',
-            'use_repo(ext, ext_hello = "hello")',
-            'other_ext = use_extension("extension.bzl", "other_ext")',
-            'other_ext.dep(generate = False)',
-            'use_repo(other_ext, other_ext_hello = "hello")',
-        ],
+      'MODULE.bazel',
+      [
+        'ext = use_extension("extension.bzl", "ext")',
+        'ext.dep(generate = True)',
+        'use_repo(ext, ext_hello = "hello")',
+        'other_ext = use_extension("extension.bzl", "other_ext")',
+        'other_ext.dep(generate = False)',
+        'use_repo(other_ext, other_ext_hello = "hello")',
+      ],
     )
     self.ScratchFile('BUILD.bazel')
     self.ScratchFile(
-        'extension.bzl',
-        [
-            'def _repo_rule_impl(ctx):',
-            '    ctx.file("WORKSPACE")',
-            '    ctx.file("BUILD", "filegroup(name=\'lala\')")',
-            '',
-            'repo_rule = repository_rule(implementation=_repo_rule_impl)',
-            '',
-            'def _module_ext_impl(ctx):',
-            '    for mod in ctx.modules:',
-            '        for dep in mod.tags.dep:',
-            '            if dep.generate:',
-            '                repo_rule(name="hello")',
-            '',
-            '_dep = tag_class(attrs={"generate": attr.bool()})',
-            'ext = module_extension(',
-            '    implementation=_module_ext_impl,',
-            '    tag_classes={"dep": _dep},',
-            ')',
-            'other_ext = module_extension(',
-            '    implementation=_module_ext_impl,',
-            '    tag_classes={"dep": _dep},',
-            ')',
-        ],
+      'extension.bzl',
+      [
+        'def _repo_rule_impl(ctx):',
+        '    ctx.file("WORKSPACE")',
+        '    ctx.file("BUILD", "filegroup(name=\'lala\')")',
+        '',
+        'repo_rule = repository_rule(implementation=_repo_rule_impl)',
+        '',
+        'def _module_ext_impl(ctx):',
+        '    for mod in ctx.modules:',
+        '        for dep in mod.tags.dep:',
+        '            if dep.generate:',
+        '                repo_rule(name="hello")',
+        '',
+        '_dep = tag_class(attrs={"generate": attr.bool()})',
+        'ext = module_extension(',
+        '    implementation=_module_ext_impl,',
+        '    tag_classes={"dep": _dep},',
+        ')',
+        'other_ext = module_extension(',
+        '    implementation=_module_ext_impl,',
+        '    tag_classes={"dep": _dep},',
+        ')',
+      ],
     )
 
     # Paths to module files in error message always use forward slashes as
@@ -602,20 +594,147 @@ def testNoAbsoluteRootModuleFilePath(self):
 
     self.RunBazel(['shutdown'])
     exit_code, _, stderr = self.RunBazel(
-        ['build', '--nobuild', '@other_ext_hello//:all'], allow_failure=True
+      ['build', '--nobuild', '@other_ext_hello//:all'], allow_failure=True
     )
     self.AssertNotExitCode(exit_code, 0, stderr)
+    self.assertIn(
+      (
+          'ERROR: module extension "other_ext" from "//:extension.bzl" does '
+          'not generate repository "hello", yet it is imported as '
+          '"other_ext_hello" in the usage at '
+          + module_file_path
+          + ':4:26'
+      ),
+      stderr,
+    )
+
+  def testModuleExtensionEnvVariable(self):
+    self.ScratchFile(
+        'MODULE.bazel',
+        [
+            'lockfile_ext = use_extension("extension.bzl", "lockfile_ext")',
+            'use_repo(lockfile_ext, "hello")',
+        ],
+    )
+    self.ScratchFile('BUILD.bazel')
+    self.ScratchFile(
+        'extension.bzl',
+        [
+            'def _repo_rule_impl(ctx):',
+            '    ctx.file("WORKSPACE")',
+            '    ctx.file("BUILD", "filegroup(name=\\"lala\\")")',
+            'repo_rule = repository_rule(implementation = _repo_rule_impl)',
+            'def _module_ext_impl(ctx):',
+            '    print("Hello from the other side!")',
+            '    repo_rule(name= "hello")',
+            'lockfile_ext = module_extension(',
+            '    implementation = _module_ext_impl,',
+            '    environ = ["SET_ME"]',
+            ')',
+        ],
+    )
+    _, _, stderr = self.RunBazel(
+        ['build', '@hello//:all'], env_add={'SET_ME': 'High in sky'}
+    )
+    self.assertIn('Hello from the other side!', ''.join(stderr))
+    # Run with same value, no evaluated
+    _, _, stderr = self.RunBazel(
+        ['build', '@hello//:all'], env_add={'SET_ME': 'High in sky'}
+    )
+    self.assertNotIn('Hello from the other side!', ''.join(stderr))
+    # Run with different value, will be re-evaluated
+    _, _, stderr = self.RunBazel(
+        ['build', '@hello//:all'], env_add={'SET_ME': 'Down to earth'}
+    )
+    self.assertIn('Hello from the other side!', ''.join(stderr))
+
+  def testChangeEnvVariableInErrorMode(self):
+    # If environ is set up in module extension, it should be re-evaluated if its
+    # value changed
+    self.ScratchFile(
+        'MODULE.bazel',
+        [
+            'lockfile_ext = use_extension("extension.bzl", "lockfile_ext")',
+            'use_repo(lockfile_ext, "hello")',
+        ],
+    )
+    self.ScratchFile('BUILD.bazel')
+    self.ScratchFile(
+        'extension.bzl',
+        [
+            'def _repo_rule_impl(ctx):',
+            '    ctx.file("WORKSPACE")',
+            '    ctx.file("BUILD", "filegroup(name=\\"lala\\")")',
+            'repo_rule = repository_rule(implementation = _repo_rule_impl)',
+            'def _module_ext_impl(ctx):',
+            '    print("Hello from the other side!")',
+            '    repo_rule(name= "hello")',
+            'lockfile_ext = module_extension(',
+            '    implementation = _module_ext_impl,',
+            '    environ = ["SET_ME"]',
+            ')',
+        ],
+    )
+    self.RunBazel(['build', '@hello//:all'], env_add={'SET_ME': 'High in sky'})
+    exit_code, _, stderr = self.RunBazel(
+        ['build', '--lockfile_mode=error', '@hello//:all'],
+        env_add={'SET_ME': 'Down to earth'},
+        allow_failure=True,
+    )
+    self.AssertExitCode(exit_code, 48, stderr)
     self.assertIn(
         (
-            'ERROR: module extension "other_ext" from "//:extension.bzl" does '
-            'not generate repository "hello", yet it is imported as '
-            '"other_ext_hello" in the usage at '
-            + module_file_path
-            + ':4:26'
+            'ERROR: Lock file is no longer up-to-date because: The environment'
+            ' variables the extension'
+            " 'ModuleExtensionId{bzlFileLabel=//:extension.bzl,"
+            " extensionName=lockfile_ext, isolationKey=Optional.empty}' depends"
+            ' on (or their values) have changed'
         ),
         stderr,
     )
 
+  def testModuleExtensionWithFile(self):
+    self.ScratchFile(
+        'MODULE.bazel',
+        [
+            'lockfile_ext = use_extension("extension.bzl", "lockfile_ext")',
+            'use_repo(lockfile_ext, "hello")',
+        ],
+    )
+    self.ScratchFile('BUILD.bazel')
+    self.ScratchFile(
+        'extension.bzl',
+        [
+            'def _repo_rule_impl(ctx):',
+            '    ctx.file("WORKSPACE")',
+            '    ctx.file("BUILD", "filegroup(name=\'lala\')")',
+            '',
+            'repo_rule = repository_rule(implementation=_repo_rule_impl)',
+            '',
+            'def _module_ext_impl(ctx):',
+            '    print(ctx.read(Label("//:hello.txt")))',
+            '    repo_rule(name="hello")',
+            '',
+            'lockfile_ext = module_extension(',
+            '    implementation=_module_ext_impl',
+            ')',
+        ],
+    )
+
+    self.ScratchFile('hello.txt', ['I will not stay the same.'])
+    _, _, stderr = self.RunBazel(['build', '@hello//:all'])
+    self.assertIn('I will not stay the same.', ''.join(stderr))
+
+    # Shutdown bazel to empty cache and run with no changes
+    self.RunBazel(['shutdown'])
+    _, _, stderr = self.RunBazel(['build', '@hello//:all'])
+    self.assertNotIn('I will not stay the same.', ''.join(stderr))
+
+    # Update file and rerun
+    self.ScratchFile('hello.txt', ['I have changed now!'])
+    _, _, stderr = self.RunBazel(['build', '@hello//:all'])
+    self.assertIn('I have changed now!', ''.join(stderr))
+
 
 if __name__ == '__main__':
   unittest.main()