Implement "Automatic use_repo fixups for module extensions"
Proposal: https://docs.google.com/document/d/1dj8SN5L6nwhNOufNqjBhYkk5f-BJI_FPYWKxlB3GAmA/edit
src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD src/main/java/com/google/devtools/build/lib/bazel/bzlmod/DelegateTypeAdapterFactory.java src/main/java/com/google/devtools/build/lib/bazel/bzlmod/GsonTypeAdapterUtil.java src/main/java/com/google/devtools/build/lib/bazel/bzlmod/Module.java src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionContext.java src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionUsage.java src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleFileGlobals.java src/main/java/com/google/devtools/build/lib/bazel/bzlmod/SingleExtensionEvalFunction.java
src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD src/main/java/com/google/devtools/build/lib/bazel/bzlmod/DelegateTypeAdapterFactory.java src/main/java/com/google/devtools/build/lib/bazel/bzlmod/GsonTypeAdapterUtil.java src/main/java/com/google/devtools/build/lib/bazel/bzlmod/Module.java src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionContext.java src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionMetadata.java src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionUsage.java src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleFileGlobals.java src/main/java/com/google/devtools/build/lib/bazel/bzlmod/SingleExtensionEvalFunction.java
src/test/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD src/test/java/com/google/devtools/build/lib/bazel/bzlmod/BazelDepGraphFunctionTest.java src/test/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionResolutionTest.java src/test/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleFileFunctionTest.java src/test/java/com/google/devtools/build/lib/bazel/bzlmod/StarlarkBazelModuleTest.java
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD
index bc24b3319356fd..bd02cc51fcdec3 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD
@@ -142,6 +142,7 @@ java_library(
         "Discovery.java",
         "GsonTypeAdapterUtil.java",
         "ModuleExtensionContext.java",
+        "ModuleExtensionMetadata.java",
         "ModuleFileFunction.java",
         "ModuleFileGlobals.java",
         "Selection.java",
@@ -168,7 +169,6 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/events",
         "//src/main/java/com/google/devtools/build/lib/packages",
-        "//src/main/java/com/google/devtools/build/lib/packages/semantics",
         "//src/main/java/com/google/devtools/build/lib/rules:repository/repository_directory_value",
         "//src/main/java/com/google/devtools/build/lib/rules:repository/repository_function",
         "//src/main/java/com/google/devtools/build/lib/skyframe:bzl_load_value",
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/DelegateTypeAdapterFactory.java b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/DelegateTypeAdapterFactory.java
index d791789dc46bbb..b5beb7880a0ed8 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/DelegateTypeAdapterFactory.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/DelegateTypeAdapterFactory.java
@@ -18,6 +18,7 @@
 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.gson.Gson;
 import com.google.gson.TypeAdapter;
 import com.google.gson.TypeAdapterFactory;
@@ -29,8 +30,10 @@
 import java.lang.reflect.Type;
 import java.util.ArrayList;
 import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.function.Function;
 import javax.annotation.Nullable;
 import net.starlark.java.eval.Dict;
@@ -92,6 +95,14 @@ private DelegateTypeAdapterFactory(
           raw -> new ArrayList<>((List<?>) raw),
           delegate -> ImmutableList.copyOf((List<?>) delegate));
 
+  public static final TypeAdapterFactory IMMUTABLE_SET =
+      new DelegateTypeAdapterFactory<>(
+          ImmutableSet.class,
+          Set.class,
+          LinkedHashSet.class,
+          raw -> new LinkedHashSet<>((Set<?>) raw),
+          delegate -> ImmutableSet.copyOf((Set<?>) delegate));
+
   @SuppressWarnings("unchecked")
   @Override
   @Nullable
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 3cb6d22954126a..4a7a8b43908914 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
@@ -18,6 +18,7 @@
 import static com.google.devtools.build.lib.bazel.bzlmod.DelegateTypeAdapterFactory.IMMUTABLE_BIMAP;
 import static com.google.devtools.build.lib.bazel.bzlmod.DelegateTypeAdapterFactory.IMMUTABLE_LIST;
 import static com.google.devtools.build.lib.bazel.bzlmod.DelegateTypeAdapterFactory.IMMUTABLE_MAP;
+import static com.google.devtools.build.lib.bazel.bzlmod.DelegateTypeAdapterFactory.IMMUTABLE_SET;
 
 import com.google.common.base.Splitter;
 import com.google.devtools.build.lib.bazel.bzlmod.Version.ParseException;
@@ -95,6 +96,7 @@ public ModuleKey read(JsonReader jsonReader) throws IOException {
           .registerTypeAdapterFactory(IMMUTABLE_MAP)
           .registerTypeAdapterFactory(IMMUTABLE_LIST)
           .registerTypeAdapterFactory(IMMUTABLE_BIMAP)
+          .registerTypeAdapterFactory(IMMUTABLE_SET)
           .registerTypeAdapter(Version.class, VERSION_TYPE_ADAPTER)
           .registerTypeAdapter(ModuleKey.class, MODULE_KEY_TYPE_ADAPTER)
           .registerTypeAdapter(AttributeValues.class, new AttributeValuesAdapter())
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/Module.java b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/Module.java
index ca56fb4ac28bd5..86fae25df40b45 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/Module.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/Module.java
@@ -261,6 +261,8 @@ public Builder addExtensionUsage(ModuleExtensionUsage value) {
       return this;
     }
 
+    abstract ModuleKey getKey();
+
     abstract String getName();
 
     abstract Optional<String> getRepoName();
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionContext.java b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionContext.java
index 0cf2fd58a9a15d..049b93053ea03b 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionContext.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionContext.java
@@ -29,6 +29,8 @@
 import net.starlark.java.annot.StarlarkBuiltin;
 import net.starlark.java.annot.StarlarkMethod;
 import net.starlark.java.eval.EvalException;
+import net.starlark.java.eval.NoneType;
+import net.starlark.java.eval.Sequence;
 import net.starlark.java.eval.StarlarkList;
 import net.starlark.java.eval.StarlarkSemantics;
 
@@ -117,4 +119,68 @@ public StarlarkList<StarlarkBazelModule> getModules() {
   public boolean isDevDependency(TypeCheckedTag tag) {
     return tag.isDevDependency();
   }
+
+  @StarlarkMethod(
+      name = "extension_metadata",
+      doc =
+          "Constructs an opaque object that can be returned from the module extension's"
+              + " implementation function to provide metadata about the repositories generated by"
+              + " the extension to Bazel.",
+      parameters = {
+        @Param(
+            name = "root_module_direct_deps",
+            doc =
+                "The names of the repositories that the extension considers to be direct"
+                    + " dependencies of the root module. If the root module imports additional"
+                    + " repositories or does not import all of these repositories via <a"
+                    + " href=\"../globals/module.html#use_repo\"><code>use_repo</code></a>, Bazel"
+                    + " will print a warning and a fixup command when the extension is"
+                    + " evaluated.<p>If one of <code>root_module_direct_deps</code> and"
+                    + " <code>root_module_direct_dev_deps</code> is specified, the other has to be"
+                    + " as well. The lists specified by these two parameters must be"
+                    + " disjoint.<p>Exactly one of <code>root_module_direct_deps</code> and"
+                    + " <code>root_module_direct_dev_deps</code> can be set to the special value"
+                    + " <code>\"all\"</code>, which is treated as if a list with the names of"
+                    + " allrepositories generated by the extension was specified as the value.",
+            positional = false,
+            named = true,
+            defaultValue = "None",
+            allowedTypes = {
+              @ParamType(type = Sequence.class, generic1 = String.class),
+              @ParamType(type = String.class),
+              @ParamType(type = NoneType.class)
+            }),
+        @Param(
+            name = "root_module_direct_dev_deps",
+            doc =
+                "The names of the repositories that the extension considers to be direct dev"
+                    + " dependencies of the root module. If the root module imports additional"
+                    + " repositories or does not import all of these repositories via <a"
+                    + " href=\"../globals/module.html#use_repo\"><code>use_repo</code></a> on an"
+                    + " extension proxy created with <code><a"
+                    + " href=\"../globals/module.html#use_extension>use_extension</a>(...,"
+                    + " dev_dependency = True)</code>, Bazel will print a warning and a fixup"
+                    + " command when the extension is evaluated.<p>If one of"
+                    + " <code>root_module_direct_deps</code> and"
+                    + " <code>root_module_direct_dev_deps</code> is specified, the other has to be"
+                    + " as well. The lists specified by these two parameters must be"
+                    + " disjoint.<p>Exactly one of <code>root_module_direct_deps</code> and"
+                    + " <code>root_module_direct_dev_deps</code> can be set to the special value"
+                    + " <code>\"all\"</code>, which is treated as if a list with the names of"
+                    + " allrepositories generated by the extension was specified as the value.",
+            positional = false,
+            named = true,
+            defaultValue = "None",
+            allowedTypes = {
+              @ParamType(type = Sequence.class, generic1 = String.class),
+              @ParamType(type = String.class),
+              @ParamType(type = NoneType.class)
+            }),
+      })
+  public ModuleExtensionMetadata extensionMetadata(
+      Object rootModuleDirectDepsUnchecked, Object rootModuleDirectDevDepsUnchecked)
+      throws EvalException {
+    return ModuleExtensionMetadata.create(
+        rootModuleDirectDepsUnchecked, rootModuleDirectDevDepsUnchecked);
+  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionMetadata.java b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionMetadata.java
new file mode 100644
index 00000000000000..fde1c8fcd6887c
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionMetadata.java
@@ -0,0 +1,350 @@
+// Copyright 2021 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.bazel.bzlmod;
+
+import static com.google.common.collect.ImmutableList.toImmutableList;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
+
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableSortedSet;
+import com.google.common.collect.Sets;
+import com.google.devtools.build.docgen.annot.DocCategory;
+import com.google.devtools.build.lib.cmdline.RepositoryName;
+import com.google.devtools.build.lib.events.Event;
+import com.google.devtools.build.lib.events.EventHandler;
+import java.util.Collection;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import javax.annotation.Nullable;
+import net.starlark.java.annot.StarlarkBuiltin;
+import net.starlark.java.eval.EvalException;
+import net.starlark.java.eval.Sequence;
+import net.starlark.java.eval.Starlark;
+import net.starlark.java.eval.StarlarkList;
+import net.starlark.java.eval.StarlarkValue;
+import net.starlark.java.syntax.Location;
+
+/** The Starlark object passed to the implementation function of module extension metadata. */
+@StarlarkBuiltin(
+    name = "extension_metadata",
+    category = DocCategory.BUILTIN,
+    doc =
+        "Return values of this type from a module extension's implementation function to "
+            + "provide metadata about the repositories generated by the extension to Bazel.")
+public class ModuleExtensionMetadata implements StarlarkValue {
+  @Nullable private final ImmutableSet<String> explicitRootModuleDirectDeps;
+  @Nullable private final ImmutableSet<String> explicitRootModuleDirectDevDeps;
+  private final UseAllRepos useAllRepos;
+
+  private ModuleExtensionMetadata(
+      @Nullable Set<String> explicitRootModuleDirectDeps,
+      @Nullable Set<String> explicitRootModuleDirectDevDeps,
+      UseAllRepos useAllRepos) {
+    this.explicitRootModuleDirectDeps =
+        explicitRootModuleDirectDeps != null
+            ? ImmutableSet.copyOf(explicitRootModuleDirectDeps)
+            : null;
+    this.explicitRootModuleDirectDevDeps =
+        explicitRootModuleDirectDevDeps != null
+            ? ImmutableSet.copyOf(explicitRootModuleDirectDevDeps)
+            : null;
+    this.useAllRepos = useAllRepos;
+  }
+
+  static ModuleExtensionMetadata create(
+      Object rootModuleDirectDepsUnchecked, Object rootModuleDirectDevDepsUnchecked)
+      throws EvalException {
+    if (rootModuleDirectDepsUnchecked == Starlark.NONE
+        && rootModuleDirectDevDepsUnchecked == Starlark.NONE) {
+      return new ModuleExtensionMetadata(null, null, UseAllRepos.NO);
+    }
+
+    // When root_module_direct_deps = "all", accept both root_module_direct_dev_deps = None and
+    // root_module_direct_dev_deps = [], but not root_module_direct_dev_deps = ["some_repo"].
+    if (rootModuleDirectDepsUnchecked.equals("all")
+        && rootModuleDirectDevDepsUnchecked.equals(StarlarkList.immutableOf())) {
+      return new ModuleExtensionMetadata(null, null, UseAllRepos.REGULAR);
+    }
+
+    if (rootModuleDirectDevDepsUnchecked.equals("all")
+        && rootModuleDirectDepsUnchecked.equals(StarlarkList.immutableOf())) {
+      return new ModuleExtensionMetadata(null, null, UseAllRepos.DEV);
+    }
+
+    if (rootModuleDirectDepsUnchecked.equals("all")
+        || rootModuleDirectDevDepsUnchecked.equals("all")) {
+      throw Starlark.errorf(
+          "if one of root_module_direct_deps and root_module_direct_dev_deps is "
+              + "\"all\", the other must be an empty list");
+    }
+
+    if (rootModuleDirectDepsUnchecked instanceof String
+        || rootModuleDirectDevDepsUnchecked instanceof String) {
+      throw Starlark.errorf(
+          "root_module_direct_deps and root_module_direct_dev_deps must be "
+              + "None, \"all\", or a list of strings");
+    }
+    if ((rootModuleDirectDepsUnchecked == Starlark.NONE)
+        != (rootModuleDirectDevDepsUnchecked == Starlark.NONE)) {
+      throw Starlark.errorf(
+          "root_module_direct_deps and root_module_direct_dev_deps must both be "
+              + "specified or both be unspecified");
+    }
+
+    Sequence<String> rootModuleDirectDeps =
+        Sequence.cast(rootModuleDirectDepsUnchecked, String.class, "root_module_direct_deps");
+    Sequence<String> rootModuleDirectDevDeps =
+        Sequence.cast(
+            rootModuleDirectDevDepsUnchecked, String.class, "root_module_direct_dev_deps");
+
+    Set<String> explicitRootModuleDirectDeps = new LinkedHashSet<>();
+    for (String dep : rootModuleDirectDeps) {
+      try {
+        RepositoryName.validateUserProvidedRepoName(dep);
+      } catch (EvalException e) {
+        throw Starlark.errorf("in root_module_direct_deps: %s", e.getMessage());
+      }
+      if (!explicitRootModuleDirectDeps.add(dep)) {
+        throw Starlark.errorf("in root_module_direct_deps: duplicate entry '%s'", dep);
+      }
+    }
+
+    Set<String> explicitRootModuleDirectDevDeps = new LinkedHashSet<>();
+    for (String dep : rootModuleDirectDevDeps) {
+      try {
+        RepositoryName.validateUserProvidedRepoName(dep);
+      } catch (EvalException e) {
+        throw Starlark.errorf("in root_module_direct_dev_deps: %s", e.getMessage());
+      }
+      if (explicitRootModuleDirectDeps.contains(dep)) {
+        throw Starlark.errorf(
+            "in root_module_direct_dev_deps: entry '%s' is also in " + "root_module_direct_deps",
+            dep);
+      }
+      if (!explicitRootModuleDirectDevDeps.add(dep)) {
+        throw Starlark.errorf("in root_module_direct_dev_deps: duplicate entry '%s'", dep);
+      }
+    }
+
+    return new ModuleExtensionMetadata(
+        explicitRootModuleDirectDeps, explicitRootModuleDirectDevDeps, UseAllRepos.NO);
+  }
+
+  public void evaluate(
+      Collection<ModuleExtensionUsage> usages, Set<String> allRepos, EventHandler handler)
+      throws EvalException {
+    generateFixupMessage(usages, allRepos).ifPresent(handler::handle);
+  }
+
+  Optional<Event> generateFixupMessage(
+      Collection<ModuleExtensionUsage> usages, Set<String> allRepos) throws EvalException {
+    var rootUsages =
+        usages.stream()
+            .filter(usage -> usage.getUsingModule().equals(ModuleKey.ROOT))
+            .collect(toImmutableList());
+    if (rootUsages.isEmpty()) {
+      // The root module doesn't use the current extension. Do not suggest fixes as the user isn't
+      // expected to modify any other module's MODULE.bazel file.
+      return Optional.empty();
+    }
+
+    var rootModuleDirectDevDeps = getRootModuleDirectDevDeps(allRepos);
+    var rootModuleDirectDeps = getRootModuleDirectDeps(allRepos);
+    if (rootModuleDirectDevDeps.isEmpty() && rootModuleDirectDeps.isEmpty()) {
+      return Optional.empty();
+    }
+
+    Preconditions.checkState(
+        rootModuleDirectDevDeps.isPresent() && rootModuleDirectDeps.isPresent());
+    return generateFixupMessage(
+        rootUsages, allRepos, rootModuleDirectDeps.get(), rootModuleDirectDevDeps.get());
+  }
+
+  private static Optional<Event> generateFixupMessage(
+      List<ModuleExtensionUsage> rootUsages,
+      Set<String> allRepos,
+      Set<String> expectedImports,
+      Set<String> expectedDevImports) {
+    var actualDevImports =
+        rootUsages.stream()
+            .flatMap(usage -> usage.getDevImports().stream())
+            .collect(toImmutableSet());
+    var actualImports =
+        rootUsages.stream()
+            .flatMap(usage -> usage.getImports().keySet().stream())
+            .filter(repo -> !actualDevImports.contains(repo))
+            .collect(toImmutableSet());
+
+    // All label strings that map to the same Label are equivalent for buildozer as it implements
+    // the same normalization of label strings with no or empty repo name.
+    ModuleExtensionUsage firstUsage = rootUsages.get(0);
+    String extensionBzlFile = firstUsage.getExtensionBzlFile();
+    String extensionName = firstUsage.getExtensionName();
+    Location location = firstUsage.getLocation();
+
+    var importsToAdd = ImmutableSortedSet.copyOf(Sets.difference(expectedImports, actualImports));
+    var importsToRemove =
+        ImmutableSortedSet.copyOf(Sets.difference(actualImports, expectedImports));
+    var devImportsToAdd =
+        ImmutableSortedSet.copyOf(Sets.difference(expectedDevImports, actualDevImports));
+    var devImportsToRemove =
+        ImmutableSortedSet.copyOf(Sets.difference(actualDevImports, expectedDevImports));
+
+    if (importsToAdd.isEmpty()
+        && importsToRemove.isEmpty()
+        && devImportsToAdd.isEmpty()
+        && devImportsToRemove.isEmpty()) {
+      return Optional.empty();
+    }
+
+    var message =
+        String.format(
+            "The module extension %s defined in %s reported incorrect imports "
+                + "of repositories via use_repo():\n\n",
+            extensionName, extensionBzlFile);
+
+    var allActualImports = ImmutableSortedSet.copyOf(Sets.union(actualImports, actualDevImports));
+    var allExpectedImports =
+        ImmutableSortedSet.copyOf(Sets.union(expectedImports, expectedDevImports));
+
+    var invalidImports = ImmutableSortedSet.copyOf(Sets.difference(allActualImports, allRepos));
+    if (!invalidImports.isEmpty()) {
+      message +=
+          String.format(
+              "Imported, but not created by the extension (will cause the build to fail):\n"
+                  + "    %s\n\n",
+              String.join(", ", invalidImports));
+    }
+
+    var missingImports =
+        ImmutableSortedSet.copyOf(Sets.difference(allExpectedImports, allActualImports));
+    if (!missingImports.isEmpty()) {
+      message +=
+          String.format(
+              "Not imported, but reported as direct dependencies by the extension (may cause the"
+                  + " build to fail):\n"
+                  + "    %s\n\n",
+              String.join(", ", missingImports));
+    }
+
+    var indirectDepImports =
+        ImmutableSortedSet.copyOf(
+            Sets.difference(Sets.intersection(allActualImports, allRepos), allExpectedImports));
+    if (!indirectDepImports.isEmpty()) {
+      message +=
+          String.format(
+              "Imported, but reported as indirect dependencies by the extension:\n    %s\n\n",
+              String.join(", ", indirectDepImports));
+    }
+
+    var fixupCommands =
+        Stream.of(
+                makeUseRepoCommand(
+                    "use_repo_add", false, importsToAdd, extensionBzlFile, extensionName),
+                makeUseRepoCommand(
+                    "use_repo_remove", false, importsToRemove, extensionBzlFile, extensionName),
+                makeUseRepoCommand(
+                    "use_repo_add", true, devImportsToAdd, extensionBzlFile, extensionName),
+                makeUseRepoCommand(
+                    "use_repo_remove", true, devImportsToRemove, extensionBzlFile, extensionName))
+            .flatMap(Optional::stream);
+
+    return Optional.of(
+        Event.warn(
+            location,
+            message
+                + String.format(
+                    "%s ** You can use the following buildozer command(s) to fix these"
+                        + " issues:%s\n\n"
+                        + "%s",
+                    "\033[35m\033[1m",
+                    "\033[0m",
+                    fixupCommands.collect(Collectors.joining("\n")))));
+  }
+
+  private static Optional<String> makeUseRepoCommand(
+      String cmd,
+      boolean devDependency,
+      Collection<String> repos,
+      String extensionBzlFile,
+      String extensionName) {
+    if (repos.isEmpty()) {
+      return Optional.empty();
+    }
+    return Optional.of(
+        String.format(
+            "buildozer '%s%s %s %s %s' //MODULE.bazel:all",
+            cmd,
+            devDependency ? " dev" : "",
+            extensionBzlFile,
+            extensionName,
+            String.join(" ", repos)));
+  }
+
+  private Optional<ImmutableSet<String>> getRootModuleDirectDeps(Set<String> allRepos)
+      throws EvalException {
+    switch (useAllRepos) {
+      case NO:
+        if (explicitRootModuleDirectDeps != null) {
+          Set<String> invalidRepos = Sets.difference(explicitRootModuleDirectDeps, allRepos);
+          if (!invalidRepos.isEmpty()) {
+            throw Starlark.errorf(
+                "root_module_direct_deps contained the following repositories "
+                    + "not generated by the extension: %s",
+                String.join(", ", invalidRepos));
+          }
+        }
+        return Optional.ofNullable(explicitRootModuleDirectDeps);
+      case REGULAR:
+        return Optional.of(ImmutableSet.copyOf(allRepos));
+      case DEV:
+        return Optional.of(ImmutableSet.of());
+    }
+    throw new IllegalStateException("not reached");
+  }
+
+  private Optional<ImmutableSet<String>> getRootModuleDirectDevDeps(Set<String> allRepos)
+      throws EvalException {
+    switch (useAllRepos) {
+      case NO:
+        if (explicitRootModuleDirectDevDeps != null) {
+          Set<String> invalidRepos = Sets.difference(explicitRootModuleDirectDevDeps, allRepos);
+          if (!invalidRepos.isEmpty()) {
+            throw Starlark.errorf(
+                "root_module_direct_dev_deps contained the following "
+                    + "repositories not generated by the extension: %s",
+                String.join(", ", invalidRepos));
+          }
+        }
+        return Optional.ofNullable(explicitRootModuleDirectDevDeps);
+      case REGULAR:
+        return Optional.of(ImmutableSet.of());
+      case DEV:
+        return Optional.of(ImmutableSet.copyOf(allRepos));
+    }
+    throw new IllegalStateException("not reached");
+  }
+
+  private enum UseAllRepos {
+    NO,
+    REGULAR,
+    DEV,
+  }
+}
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionUsage.java b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionUsage.java
index 64be185428b300..31cff67ecace10 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionUsage.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionUsage.java
@@ -17,6 +17,7 @@
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableBiMap;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.ryanharter.auto.value.gson.GenerateTypeAdapter;
 import net.starlark.java.syntax.Location;
@@ -34,6 +35,9 @@ public abstract class ModuleExtensionUsage {
   /** The name of the extension. */
   public abstract String getExtensionName();
 
+  /** The module that contains this particular extension usage. */
+  public abstract ModuleKey getUsingModule();
+
   /**
    * The location where this proxy object was created (by the {@code use_extension} call). Note that
    * if there were multiple {@code use_extension} calls on same extension, then this only stores the
@@ -48,6 +52,12 @@ public abstract class ModuleExtensionUsage {
    */
   public abstract ImmutableBiMap<String, String> getImports();
 
+  /**
+   * The repo names as exported by the module extension that were imported using a proxy marked as a
+   * dev dependency.
+   */
+  public abstract ImmutableSet<String> getDevImports();
+
   /** All the tags specified by this module for this extension. */
   public abstract ImmutableList<Tag> getTags();
 
@@ -63,10 +73,14 @@ public abstract static class Builder {
 
     public abstract Builder setExtensionName(String value);
 
+    public abstract Builder setUsingModule(ModuleKey value);
+
     public abstract Builder setLocation(Location value);
 
     public abstract Builder setImports(ImmutableBiMap<String, String> value);
 
+    public abstract Builder setDevImports(ImmutableSet<String> value);
+
     public abstract Builder setTags(ImmutableList<Tag> value);
 
     abstract ImmutableList.Builder<Tag> tagsBuilder();
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleFileGlobals.java b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleFileGlobals.java
index 7e1909a82c5272..a141e2c4c7df49 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleFileGlobals.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleFileGlobals.java
@@ -24,6 +24,7 @@
 import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.docgen.annot.DocumentMethods;
 import com.google.devtools.build.lib.bazel.bzlmod.ModuleFileGlobals.ModuleExtensionUsageBuilder.ModuleExtensionProxy;
 import com.google.devtools.build.lib.bazel.bzlmod.Version.ParseException;
@@ -61,6 +62,7 @@ public class ModuleFileGlobals {
       Pattern.compile("(>|<|-|<=|>=)(\\d+\\.){2}\\d+");
 
   private boolean moduleCalled = false;
+  private boolean hadNonModuleCall = false;
   private final boolean ignoreDevDeps;
   private final Module.Builder module;
   private final Map<String, ModuleKey> deps = new LinkedHashMap<>();
@@ -206,6 +208,9 @@ public void module(
     if (moduleCalled) {
       throw Starlark.errorf("the module() directive can only be called once");
     }
+    if (hadNonModuleCall) {
+      throw Starlark.errorf("if module() is called, it must be called before any other functions");
+    }
     moduleCalled = true;
     if (!name.isEmpty()) {
       validateModuleName(name);
@@ -296,6 +301,7 @@ private static ImmutableList<String> checkAllCompatibilityVersions(
   public void bazelDep(
       String name, String version, String repoName, boolean devDependency, StarlarkThread thread)
       throws EvalException {
+    hadNonModuleCall = true;
     if (repoName.isEmpty()) {
       repoName = name;
     }
@@ -328,6 +334,7 @@ public void bazelDep(
               allowedTypes = {@ParamType(type = Sequence.class, generic1 = String.class)},
               doc = "The labels of the platforms to register."))
   public void registerExecutionPlatforms(Sequence<?> platformLabels) throws EvalException {
+    hadNonModuleCall = true;
     module.addExecutionPlatformsToRegister(
         checkAllAbsolutePatterns(platformLabels, "register_execution_platforms"));
   }
@@ -345,6 +352,7 @@ public void registerExecutionPlatforms(Sequence<?> platformLabels) throws EvalEx
               allowedTypes = {@ParamType(type = Sequence.class, generic1 = String.class)},
               doc = "The labels of the toolchains to register."))
   public void registerToolchains(Sequence<?> toolchainLabels) throws EvalException {
+    hadNonModuleCall = true;
     module.addToolchainsToRegister(
         checkAllAbsolutePatterns(toolchainLabels, "register_toolchains"));
   }
@@ -374,7 +382,14 @@ public void registerToolchains(Sequence<?> toolchainLabels) throws EvalException
       },
       useStarlarkThread = true)
   public ModuleExtensionProxy useExtension(
-      String extensionBzlFile, String extensionName, boolean devDependency, StarlarkThread thread) {
+      String rawExtensionBzlFile,
+      String extensionName,
+      boolean devDependency,
+      StarlarkThread thread) {
+    hadNonModuleCall = true;
+
+    String extensionBzlFile = normalizeLabelString(rawExtensionBzlFile);
+
     ModuleExtensionUsageBuilder newUsageBuilder =
         new ModuleExtensionUsageBuilder(
             extensionBzlFile, extensionName, thread.getCallerLocation());
@@ -397,11 +412,28 @@ public ModuleExtensionProxy useExtension(
     return newUsageBuilder.getProxy(devDependency);
   }
 
+  private String normalizeLabelString(String rawExtensionBzlFile) {
+    // Normalize the label by adding the current module's repo_name if the label doesn't specify a
+    // repository name. This is necessary as ModuleExtensionUsages are grouped by the string value
+    // of this label, but later mapped to their Label representation. If multiple strings map to the
+    // same Label, this would result in a crash.
+    // ownName can't change anymore as calling module() after this results in an error.
+    String ownName = module.getRepoName().orElse(module.getName());
+    if (module.getKey().equals(ModuleKey.ROOT) && rawExtensionBzlFile.startsWith("@//")) {
+      return "@" + ownName + rawExtensionBzlFile.substring(1);
+    } else if (rawExtensionBzlFile.startsWith("//")) {
+      return "@" + ownName + rawExtensionBzlFile;
+    } else {
+      return rawExtensionBzlFile;
+    }
+  }
+
   class ModuleExtensionUsageBuilder {
     private final String extensionBzlFile;
     private final String extensionName;
     private final Location location;
     private final HashBiMap<String, String> imports;
+    private final ImmutableSet.Builder<String> devImports;
     private final ImmutableList.Builder<Tag> tags;
 
     ModuleExtensionUsageBuilder(String extensionBzlFile, String extensionName, Location location) {
@@ -409,6 +441,7 @@ class ModuleExtensionUsageBuilder {
       this.extensionName = extensionName;
       this.location = location;
       this.imports = HashBiMap.create();
+      this.devImports = ImmutableSet.builder();
       this.tags = ImmutableList.builder();
     }
 
@@ -416,8 +449,10 @@ ModuleExtensionUsage buildUsage() {
       return ModuleExtensionUsage.builder()
           .setExtensionBzlFile(extensionBzlFile)
           .setExtensionName(extensionName)
+          .setUsingModule(module.getKey())
           .setLocation(location)
           .setImports(ImmutableBiMap.copyOf(imports))
+          .setDevImports(devImports.build())
           .setTags(tags.build())
           .build();
     }
@@ -451,6 +486,9 @@ void addImport(String localRepoName, String exportedName, Location location)
               exportedName, extensionName, repoNameUsages.get(collisionRepoName).getWhere());
         }
         imports.put(localRepoName, exportedName);
+        if (devDependency) {
+          devImports.add(exportedName);
+        }
       }
 
       @Nullable
@@ -514,6 +552,7 @@ public void useRepo(
       Dict<String, Object> kwargs,
       StarlarkThread thread)
       throws EvalException {
+    hadNonModuleCall = true;
     Location location = thread.getCallerLocation();
     for (String arg : Sequence.cast(args, String.class, "args")) {
       extensionProxy.addImport(arg, arg, location);
@@ -596,6 +635,7 @@ public void singleVersionOverride(
       Iterable<?> patchCmds,
       StarlarkInt patchStrip)
       throws EvalException {
+    hadNonModuleCall = true;
     Version parsedVersion;
     try {
       parsedVersion = Version.parse(version);
@@ -649,6 +689,7 @@ public void singleVersionOverride(
       })
   public void multipleVersionOverride(String moduleName, Iterable<?> versions, String registry)
       throws EvalException {
+    hadNonModuleCall = true;
     ImmutableList.Builder<Version> parsedVersionsBuilder = new ImmutableList.Builder<>();
     try {
       for (String version : Sequence.cast(versions, String.class, "versions").getImmutableList()) {
@@ -732,6 +773,7 @@ public void archiveOverride(
       Iterable<?> patchCmds,
       StarlarkInt patchStrip)
       throws EvalException {
+    hadNonModuleCall = true;
     ImmutableList<String> urlList =
         urls instanceof String
             ? ImmutableList.of((String) urls)
@@ -803,6 +845,7 @@ public void gitOverride(
       Iterable<?> patchCmds,
       StarlarkInt patchStrip)
       throws EvalException {
+    hadNonModuleCall = true;
     addOverride(
         moduleName,
         GitOverride.create(
@@ -832,6 +875,7 @@ public void gitOverride(
             positional = false),
       })
   public void localPathOverride(String moduleName, String path) throws EvalException {
+    hadNonModuleCall = true;
     addOverride(moduleName, LocalPathOverride.create(path));
   }
 
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 0418870181811b..38360edb1fd724 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
@@ -183,8 +183,26 @@ public SkyValue compute(SkyKey skyKey, Environment env)
           createContext(env, usagesValue, starlarkSemantics, extensionId, extension);
       threadContext.storeInThread(thread);
       try {
-        Starlark.fastcall(
-            thread, extension.getImplementation(), new Object[] {moduleContext}, new Object[0]);
+        Object returnValue =
+            Starlark.fastcall(
+                thread, extension.getImplementation(), new Object[] {moduleContext}, new Object[0]);
+        if (returnValue != Starlark.NONE && !(returnValue instanceof ModuleExtensionMetadata)) {
+          throw new SingleExtensionEvalFunctionException(
+              ExternalDepsException.withMessage(
+                  Code.BAD_MODULE,
+                  "expected module extension %s in %s to return None or extension_metadata, got %s",
+                  extensionId.getExtensionName(),
+                  extensionId.getBzlFileLabel(),
+                  Starlark.type(returnValue)),
+              Transience.PERSISTENT);
+        }
+        if (returnValue instanceof ModuleExtensionMetadata) {
+          ModuleExtensionMetadata metadata = (ModuleExtensionMetadata) returnValue;
+          metadata.evaluate(
+              usagesValue.getExtensionUsages().values(),
+              threadContext.getGeneratedRepos().keySet(),
+              env.getListener());
+        }
       } catch (NeedsSkyframeRestartException e) {
         // Clean up and restart by returning null.
         try {

diff --git a/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD b/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD
index 3762ba587a24e6..8c2c76bf5bc1d1 100644
--- a/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD
+++ b/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD
@@ -45,6 +45,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/bazel/repository/starlark",
         "//src/main/java/com/google/devtools/build/lib/clock",
         "//src/main/java/com/google/devtools/build/lib/cmdline",
+        "//src/main/java/com/google/devtools/build/lib/events",
         "//src/main/java/com/google/devtools/build/lib/packages",
         "//src/main/java/com/google/devtools/build/lib/packages/semantics",
         "//src/main/java/com/google/devtools/build/lib/pkgcache",
@@ -76,6 +77,7 @@ java_library(
         "//src/main/java/net/starlark/java/syntax",
         "//src/test/java/com/google/devtools/build/lib/analysis/util",
         "//src/test/java/com/google/devtools/build/lib/testutil",
+        "//src/test/java/com/google/devtools/build/lib/testutil:JunitUtils",
         "//third_party:auto_value",
         "//third_party:caffeine",
         "//third_party:gson",
diff --git a/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/BazelDepGraphFunctionTest.java b/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/BazelDepGraphFunctionTest.java
index bcacdab1c8bb5b..b1979e9a6515ef 100644
--- a/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/BazelDepGraphFunctionTest.java
+++ b/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/BazelDepGraphFunctionTest.java
@@ -25,6 +25,7 @@
 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.actions.FileValue;
 import com.google.devtools.build.lib.analysis.BlazeDirectories;
 import com.google.devtools.build.lib.analysis.ServerDirectories;
@@ -234,6 +235,8 @@ private static ModuleExtensionUsage createModuleExtensionUsage(
         .setExtensionBzlFile(bzlFile)
         .setExtensionName(name)
         .setImports(importsBuilder.buildOrThrow())
+        .setDevImports(ImmutableSet.of())
+        .setUsingModule(ModuleKey.ROOT)
         .setLocation(Location.BUILTIN)
         .build();
   }
diff --git a/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionResolutionTest.java b/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionResolutionTest.java
index c4ac3a67d72a88..06033115a3f683 100644
--- a/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionResolutionTest.java
+++ b/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionResolutionTest.java
@@ -16,6 +16,7 @@
 
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.devtools.build.lib.bazel.bzlmod.BzlmodTestUtil.createModuleKey;
+import static com.google.devtools.build.lib.testutil.MoreAsserts.assertEventCount;
 
 import com.github.benmanes.caffeine.cache.Caffeine;
 import com.google.common.base.Suppliers;
@@ -39,6 +40,7 @@
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.PackageIdentifier;
 import com.google.devtools.build.lib.cmdline.RepositoryName;
+import com.google.devtools.build.lib.events.EventKind;
 import com.google.devtools.build.lib.packages.PackageFactory;
 import com.google.devtools.build.lib.packages.WorkspaceFileValue;
 import com.google.devtools.build.lib.packages.semantics.BuildLanguageOptions;
@@ -332,6 +334,88 @@ public void simpleExtension() throws Exception {
     assertThat(result.get(skyKey).getModule().getGlobal("data")).isEqualTo("foo:fu bar:ba");
   }
 
+  @Test
+  public void simpleExtension_nonCanonicalLabel() throws Exception {
+    scratch.file(
+        workspaceRoot.getRelative("MODULE.bazel").getPathString(),
+        "module(name='my_module', version = '1.0')",
+        "bazel_dep(name='data_repo', version='1.0')",
+        "ext1 = use_extension('//:defs.bzl', 'ext')",
+        "ext1.tag(name='foo', data='fu')",
+        "use_repo(ext1, 'foo')",
+        "ext2 = use_extension('@my_module//:defs.bzl', 'ext')",
+        "ext2.tag(name='bar', data='ba')",
+        "use_repo(ext2, 'bar')",
+        "ext3 = use_extension('@//:defs.bzl', 'ext')",
+        "ext3.tag(name='quz', data='qu')",
+        "use_repo(ext3, 'quz')");
+    scratch.file(
+        workspaceRoot.getRelative("defs.bzl").getPathString(),
+        "load('@data_repo//:defs.bzl','data_repo')",
+        "tag = tag_class(attrs = {'name':attr.string(),'data':attr.string()})",
+        "def _ext_impl(ctx):",
+        "  for mod in ctx.modules:",
+        "    for tag in mod.tags.tag:",
+        "      data_repo(name=tag.name,data=tag.data)",
+        "ext = module_extension(implementation=_ext_impl, tag_classes={'tag':tag})");
+    scratch.file(workspaceRoot.getRelative("BUILD").getPathString());
+    scratch.file(
+        workspaceRoot.getRelative("data.bzl").getPathString(),
+        "load('@foo//:data.bzl', foo_data='data')",
+        "load('@bar//:data.bzl', bar_data='data')",
+        "load('@quz//:data.bzl', quz_data='data')",
+        "data = 'foo:'+foo_data+' bar:'+bar_data+' quz:'+quz_data");
+
+    SkyKey skyKey = BzlLoadValue.keyForBuild(Label.parseCanonical("//:data.bzl"));
+    EvaluationResult<BzlLoadValue> result =
+        evaluator.evaluate(ImmutableList.of(skyKey), evaluationContext);
+    if (result.hasError()) {
+      throw result.getError().getException();
+    }
+    assertThat(result.get(skyKey).getModule().getGlobal("data")).isEqualTo("foo:fu bar:ba quz:qu");
+  }
+
+  @Test
+  public void simpleExtension_nonCanonicalLabel_repoName() throws Exception {
+    scratch.file(
+        workspaceRoot.getRelative("MODULE.bazel").getPathString(),
+        "module(name='my_module', version = '1.0', repo_name='my_name')",
+        "bazel_dep(name='data_repo', version='1.0')",
+        "ext1 = use_extension('//:defs.bzl', 'ext')",
+        "ext1.tag(name='foo', data='fu')",
+        "use_repo(ext1, 'foo')",
+        "ext2 = use_extension('@my_name//:defs.bzl', 'ext')",
+        "ext2.tag(name='bar', data='ba')",
+        "use_repo(ext2, 'bar')",
+        "ext3 = use_extension('@//:defs.bzl', 'ext')",
+        "ext3.tag(name='quz', data='qu')",
+        "use_repo(ext3, 'quz')");
+    scratch.file(
+        workspaceRoot.getRelative("defs.bzl").getPathString(),
+        "load('@data_repo//:defs.bzl','data_repo')",
+        "tag = tag_class(attrs = {'name':attr.string(),'data':attr.string()})",
+        "def _ext_impl(ctx):",
+        "  for mod in ctx.modules:",
+        "    for tag in mod.tags.tag:",
+        "      data_repo(name=tag.name,data=tag.data)",
+        "ext = module_extension(implementation=_ext_impl, tag_classes={'tag':tag})");
+    scratch.file(workspaceRoot.getRelative("BUILD").getPathString());
+    scratch.file(
+        workspaceRoot.getRelative("data.bzl").getPathString(),
+        "load('@foo//:data.bzl', foo_data='data')",
+        "load('@bar//:data.bzl', bar_data='data')",
+        "load('@quz//:data.bzl', quz_data='data')",
+        "data = 'foo:'+foo_data+' bar:'+bar_data+' quz:'+quz_data");
+
+    SkyKey skyKey = BzlLoadValue.keyForBuild(Label.parseCanonical("//:data.bzl"));
+    EvaluationResult<BzlLoadValue> result =
+        evaluator.evaluate(ImmutableList.of(skyKey), evaluationContext);
+    if (result.hasError()) {
+      throw result.getError().getException();
+    }
+    assertThat(result.get(skyKey).getModule().getGlobal("data")).isEqualTo("foo:fu bar:ba quz:qu");
+  }
+
   @Test
   public void multipleModules() throws Exception {
     scratch.file(
@@ -1374,4 +1458,437 @@ public void testReportRepoAndBzlCycles_extRepoLoadSelfCycle() throws Exception {
             + "|   @_main~my_ext~candy1//:data.bzl\n"
             + "`-- @_main~my_ext~candy1");
   }
+
+  @Test
+  public void extensionMetadata_exactlyOneArgIsNone() throws Exception {
+    var result =
+        evaluateSimpleModuleExtension(
+            "return ctx.extension_metadata(root_module_direct_deps=['foo'])");
+
+    assertThat(result.hasError()).isTrue();
+    assertContainsEvent(
+        "root_module_direct_deps and root_module_direct_dev_deps must both be specified or both be"
+            + " unspecified");
+  }
+
+  @Test
+  public void extensionMetadata_exactlyOneArgIsNoneDev() throws Exception {
+    var result =
+        evaluateSimpleModuleExtension(
+            "return ctx.extension_metadata(root_module_direct_dev_deps=['foo'])");
+
+    assertThat(result.hasError()).isTrue();
+    assertContainsEvent(
+        "root_module_direct_deps and root_module_direct_dev_deps must both be specified or both be"
+            + " unspecified");
+  }
+
+  @Test
+  public void extensionMetadata_allUsedTwice() throws Exception {
+    var result =
+        evaluateSimpleModuleExtension(
+            "return"
+                + " ctx.extension_metadata(root_module_direct_deps='all',root_module_direct_dev_deps='all')");
+
+    assertThat(result.hasError()).isTrue();
+    assertContainsEvent(
+        "if one of root_module_direct_deps and root_module_direct_dev_deps is \"all\", the other"
+            + " must be an empty list");
+  }
+
+  @Test
+  public void extensionMetadata_allAndNone() throws Exception {
+    var result =
+        evaluateSimpleModuleExtension(
+            "return ctx.extension_metadata(root_module_direct_deps='all')");
+
+    assertThat(result.hasError()).isTrue();
+    assertContainsEvent(
+        "if one of root_module_direct_deps and root_module_direct_dev_deps is \"all\", the other"
+            + " must be an empty list");
+  }
+
+  @Test
+  public void extensionMetadata_unsupportedString() throws Exception {
+    var result =
+        evaluateSimpleModuleExtension(
+            "return ctx.extension_metadata(root_module_direct_deps='not_all')");
+
+    assertThat(result.hasError()).isTrue();
+    assertContainsEvent(
+        "root_module_direct_deps and root_module_direct_dev_deps must be None, \"all\", or a list"
+            + " of strings");
+  }
+
+  @Test
+  public void extensionMetadata_unsupportedStringDev() throws Exception {
+    var result =
+        evaluateSimpleModuleExtension(
+            "return ctx.extension_metadata(root_module_direct_dev_deps='not_all')");
+
+    assertThat(result.hasError()).isTrue();
+    assertContainsEvent(
+        "root_module_direct_deps and root_module_direct_dev_deps must be None, \"all\", or a list"
+            + " of strings");
+  }
+
+  @Test
+  public void extensionMetadata_invalidRepoName() throws Exception {
+    var result =
+        evaluateSimpleModuleExtension(
+            "return"
+                + " ctx.extension_metadata(root_module_direct_deps=['~invalid'],root_module_direct_dev_deps=[])");
+
+    assertThat(result.hasError()).isTrue();
+    assertContainsEvent(
+        "in root_module_direct_deps: invalid user-provided repo name '~invalid': valid names may"
+            + " contain only A-Z, a-z, 0-9, '-', '_', '.', and must start with a letter");
+  }
+
+  @Test
+  public void extensionMetadata_invalidDevRepoName() throws Exception {
+    var result =
+        evaluateSimpleModuleExtension(
+            "return"
+                + " ctx.extension_metadata(root_module_direct_dev_deps=['~invalid'],root_module_direct_deps=[])");
+
+    assertThat(result.hasError()).isTrue();
+    assertContainsEvent(
+        "in root_module_direct_dev_deps: invalid user-provided repo name '~invalid': valid names"
+            + " may contain only A-Z, a-z, 0-9, '-', '_', '.', and must start with a letter");
+  }
+
+  @Test
+  public void extensionMetadata_duplicateRepo() throws Exception {
+    var result =
+        evaluateSimpleModuleExtension(
+            "return"
+                + " ctx.extension_metadata(root_module_direct_deps=['dep','dep'],root_module_direct_dev_deps=[])");
+
+    assertThat(result.hasError()).isTrue();
+    assertContainsEvent("in root_module_direct_deps: duplicate entry 'dep'");
+  }
+
+  @Test
+  public void extensionMetadata_duplicateDevRepo() throws Exception {
+    var result =
+        evaluateSimpleModuleExtension(
+            "return"
+                + " ctx.extension_metadata(root_module_direct_deps=[],root_module_direct_dev_deps=['dep','dep'])");
+
+    assertThat(result.hasError()).isTrue();
+    assertContainsEvent("in root_module_direct_dev_deps: duplicate entry 'dep'");
+  }
+
+  @Test
+  public void extensionMetadata_duplicateRepoAcrossTypes() throws Exception {
+    var result =
+        evaluateSimpleModuleExtension(
+            "return"
+                + " ctx.extension_metadata(root_module_direct_deps=['dep'],root_module_direct_dev_deps=['dep'])");
+
+    assertThat(result.hasError()).isTrue();
+    assertContainsEvent(
+        "in root_module_direct_dev_deps: entry 'dep' is also in root_module_direct_deps");
+  }
+
+  @Test
+  public void extensionMetadata() throws Exception {
+    scratch.file(
+        workspaceRoot.getRelative("MODULE.bazel").getPathString(),
+        "bazel_dep(name='ext', version='1.0')",
+        "bazel_dep(name='data_repo',version='1.0')",
+        "ext = use_extension('@ext//:defs.bzl', 'ext')",
+        "use_repo(ext, 'direct_dep', 'indirect_dep', 'invalid_dep')",
+        "ext_dev = use_extension('@ext//:defs.bzl', 'ext', dev_dependency = True)",
+        "use_repo(ext_dev, 'direct_dev_dep', 'indirect_dev_dep', 'invalid_dev_dep')");
+    scratch.file(workspaceRoot.getRelative("BUILD").getPathString());
+    scratch.file(
+        workspaceRoot.getRelative("data.bzl").getPathString(),
+        "load('@direct_dep//:data.bzl', direct_dep_data='data')",
+        "data = direct_dep_data");
+
+    registry.addModule(
+        createModuleKey("ext", "1.0"),
+        "module(name='ext',version='1.0')",
+        "bazel_dep(name='data_repo',version='1.0')",
+        "ext = use_extension('//:defs.bzl', 'ext')",
+        "use_repo(ext, 'indirect_dep')",
+        "ext_dev = use_extension('//:defs.bzl', 'ext', dev_dependency = True)",
+        "use_repo(ext_dev, 'indirect_dev_dep')");
+    scratch.file(modulesRoot.getRelative("ext~1.0/WORKSPACE").getPathString());
+    scratch.file(modulesRoot.getRelative("ext~1.0/BUILD").getPathString());
+    scratch.file(
+        modulesRoot.getRelative("ext~1.0/defs.bzl").getPathString(),
+        "load('@data_repo//:defs.bzl','data_repo')",
+        "def _ext_impl(ctx):",
+        "  data_repo(name='direct_dep')",
+        "  data_repo(name='direct_dev_dep')",
+        "  data_repo(name='missing_direct_dep')",
+        "  data_repo(name='missing_direct_dev_dep')",
+        "  data_repo(name='indirect_dep')",
+        "  data_repo(name='indirect_dev_dep')",
+        "  return ctx.extension_metadata(",
+        "    root_module_direct_deps=['direct_dep', 'missing_direct_dep'],",
+        "    root_module_direct_dev_deps=['direct_dev_dep', 'missing_direct_dev_dep'],",
+        "  )",
+        "ext=module_extension(implementation=_ext_impl)");
+
+    SkyKey skyKey = BzlLoadValue.keyForBuild(Label.parseCanonical("//:data.bzl"));
+    // Evaluation fails due to the import of a repository not generated by the extension, but we
+    // only want to assert that the warning is emitted.
+    reporter.removeHandler(failFastHandler);
+    EvaluationResult<BzlLoadValue> result =
+        evaluator.evaluate(ImmutableList.of(skyKey), evaluationContext);
+    assertThat(result.hasError()).isTrue();
+
+    assertEventCount(1, eventCollector);
+    assertContainsEvent(
+        "WARNING <root>/MODULE.bazel:3:20: The module extension ext defined in @ext//:defs.bzl"
+            + " reported incorrect imports of repositories via use_repo():\n"
+            + "\n"
+            + "Imported, but not created by the extension (will cause the build to fail):\n"
+            + "    invalid_dep, invalid_dev_dep\n"
+            + "\n"
+            + "Not imported, but reported as direct dependencies by the extension (may cause the"
+            + " build to fail):\n"
+            + "    missing_direct_dep, missing_direct_dev_dep\n"
+            + "\n"
+            + "Imported, but reported as indirect dependencies by the extension:\n"
+            + "    indirect_dep, indirect_dev_dep\n"
+            + "\n"
+            + "\033[35m\033[1m ** You can use the following buildozer command(s) to fix these"
+            + " issues:\033[0m\n"
+            + "\n"
+            + "buildozer 'use_repo_add @ext//:defs.bzl ext missing_direct_dep' //MODULE.bazel:all\n"
+            + "buildozer 'use_repo_remove @ext//:defs.bzl ext indirect_dep invalid_dep'"
+            + " //MODULE.bazel:all\n"
+            + "buildozer 'use_repo_add dev @ext//:defs.bzl ext missing_direct_dev_dep'"
+            + " //MODULE.bazel:all\n"
+            + "buildozer 'use_repo_remove dev @ext//:defs.bzl ext indirect_dev_dep invalid_dev_dep'"
+            + " //MODULE.bazel:all",
+        ImmutableSet.of(EventKind.WARNING));
+  }
+
+  @Test
+  public void extensionMetadata_all() throws Exception {
+    scratch.file(
+        workspaceRoot.getRelative("MODULE.bazel").getPathString(),
+        "bazel_dep(name='ext', version='1.0')",
+        "bazel_dep(name='data_repo',version='1.0')",
+        "ext = use_extension('@ext//:defs.bzl', 'ext')",
+        "use_repo(ext, 'direct_dep', 'indirect_dep', 'invalid_dep')",
+        "ext_dev = use_extension('@ext//:defs.bzl', 'ext', dev_dependency = True)",
+        "use_repo(ext_dev, 'direct_dev_dep', 'indirect_dev_dep', 'invalid_dev_dep')");
+    scratch.file(workspaceRoot.getRelative("BUILD").getPathString());
+    scratch.file(
+        workspaceRoot.getRelative("data.bzl").getPathString(),
+        "load('@direct_dep//:data.bzl', direct_dep_data='data')",
+        "data = direct_dep_data");
+
+    registry.addModule(
+        createModuleKey("ext", "1.0"),
+        "module(name='ext',version='1.0')",
+        "bazel_dep(name='data_repo',version='1.0')",
+        "ext = use_extension('//:defs.bzl', 'ext')",
+        "use_repo(ext, 'indirect_dep')",
+        "ext_dev = use_extension('//:defs.bzl', 'ext', dev_dependency = True)",
+        "use_repo(ext_dev, 'indirect_dev_dep')");
+    scratch.file(modulesRoot.getRelative("ext~1.0/WORKSPACE").getPathString());
+    scratch.file(modulesRoot.getRelative("ext~1.0/BUILD").getPathString());
+    scratch.file(
+        modulesRoot.getRelative("ext~1.0/defs.bzl").getPathString(),
+        "load('@data_repo//:defs.bzl','data_repo')",
+        "def _ext_impl(ctx):",
+        "  data_repo(name='direct_dep')",
+        "  data_repo(name='direct_dev_dep')",
+        "  data_repo(name='missing_direct_dep')",
+        "  data_repo(name='missing_direct_dev_dep')",
+        "  data_repo(name='indirect_dep')",
+        "  data_repo(name='indirect_dev_dep')",
+        "  return ctx.extension_metadata(",
+        "    root_module_direct_deps='all',",
+        "    root_module_direct_dev_deps=[],",
+        "  )",
+        "ext=module_extension(implementation=_ext_impl)");
+
+    SkyKey skyKey = BzlLoadValue.keyForBuild(Label.parseCanonical("//:data.bzl"));
+    reporter.removeHandler(failFastHandler);
+    EvaluationResult<BzlLoadValue> result =
+        evaluator.evaluate(ImmutableList.of(skyKey), evaluationContext);
+    assertThat(result.hasError()).isTrue();
+    assertThat(result.getError().getException())
+        .hasMessageThat()
+        .isEqualTo(
+            "module extension \"ext\" from \"@ext~1.0//:defs.bzl\" does not generate repository "
+                + "\"invalid_dep\", yet it is imported as \"invalid_dep\" in the usage at "
+                + "<root>/MODULE.bazel:3:20");
+
+    assertEventCount(1, eventCollector);
+    assertContainsEvent(
+        "WARNING <root>/MODULE.bazel:3:20: The module extension ext defined in @ext//:defs.bzl"
+            + " reported incorrect imports of repositories via use_repo():\n"
+            + "\n"
+            + "Imported, but not created by the extension (will cause the build to fail):\n"
+            + "    invalid_dep, invalid_dev_dep\n"
+            + "\n"
+            + "Not imported, but reported as direct dependencies by the extension (may cause the"
+            + " build to fail):\n"
+            + "    missing_direct_dep, missing_direct_dev_dep\n"
+            + "\n"
+            + "\033[35m\033[1m ** You can use the following buildozer command(s) to fix these"
+            + " issues:\033[0m\n"
+            + "\n"
+            + "buildozer 'use_repo_add @ext//:defs.bzl ext direct_dev_dep indirect_dev_dep"
+            + " missing_direct_dep missing_direct_dev_dep' //MODULE.bazel:all\n"
+            + "buildozer 'use_repo_remove @ext//:defs.bzl ext invalid_dep' //MODULE.bazel:all\n"
+            + "buildozer 'use_repo_remove dev @ext//:defs.bzl ext direct_dev_dep indirect_dev_dep"
+            + " invalid_dev_dep' //MODULE.bazel:all",
+        ImmutableSet.of(EventKind.WARNING));
+  }
+
+  @Test
+  public void extensionMetadata_allDev() throws Exception {
+    scratch.file(
+        workspaceRoot.getRelative("MODULE.bazel").getPathString(),
+        "bazel_dep(name='ext', version='1.0')",
+        "bazel_dep(name='data_repo',version='1.0')",
+        "ext = use_extension('@ext//:defs.bzl', 'ext')",
+        "use_repo(ext, 'direct_dep', 'indirect_dep', 'invalid_dep')",
+        "ext_dev = use_extension('@ext//:defs.bzl', 'ext', dev_dependency = True)",
+        "use_repo(ext_dev, 'direct_dev_dep', 'indirect_dev_dep', 'invalid_dev_dep')");
+    scratch.file(workspaceRoot.getRelative("BUILD").getPathString());
+    scratch.file(
+        workspaceRoot.getRelative("data.bzl").getPathString(),
+        "load('@direct_dep//:data.bzl', direct_dep_data='data')",
+        "data = direct_dep_data");
+
+    registry.addModule(
+        createModuleKey("ext", "1.0"),
+        "module(name='ext',version='1.0')",
+        "bazel_dep(name='data_repo',version='1.0')",
+        "ext = use_extension('//:defs.bzl', 'ext')",
+        "use_repo(ext, 'indirect_dep')",
+        "ext_dev = use_extension('//:defs.bzl', 'ext', dev_dependency = True)",
+        "use_repo(ext_dev, 'indirect_dev_dep')");
+    scratch.file(modulesRoot.getRelative("ext~1.0/WORKSPACE").getPathString());
+    scratch.file(modulesRoot.getRelative("ext~1.0/BUILD").getPathString());
+    scratch.file(
+        modulesRoot.getRelative("ext~1.0/defs.bzl").getPathString(),
+        "load('@data_repo//:defs.bzl','data_repo')",
+        "def _ext_impl(ctx):",
+        "  data_repo(name='direct_dep')",
+        "  data_repo(name='direct_dev_dep')",
+        "  data_repo(name='missing_direct_dep')",
+        "  data_repo(name='missing_direct_dev_dep')",
+        "  data_repo(name='indirect_dep')",
+        "  data_repo(name='indirect_dev_dep')",
+        "  return ctx.extension_metadata(",
+        "    root_module_direct_deps=[],",
+        "    root_module_direct_dev_deps='all',",
+        "  )",
+        "ext=module_extension(implementation=_ext_impl)");
+
+    SkyKey skyKey = BzlLoadValue.keyForBuild(Label.parseCanonical("//:data.bzl"));
+    // Evaluation fails due to the import of a repository not generated by the extension, but we
+    // only want to assert that the warning is emitted.
+    reporter.removeHandler(failFastHandler);
+    EvaluationResult<BzlLoadValue> result =
+        evaluator.evaluate(ImmutableList.of(skyKey), evaluationContext);
+    assertThat(result.hasError()).isTrue();
+    assertThat(result.getError().getException())
+        .hasMessageThat()
+        .isEqualTo(
+            "module extension \"ext\" from \"@ext~1.0//:defs.bzl\" does not generate repository "
+                + "\"invalid_dep\", yet it is imported as \"invalid_dep\" in the usage at "
+                + "<root>/MODULE.bazel:3:20");
+
+    assertEventCount(1, eventCollector);
+    assertContainsEvent(
+        "WARNING <root>/MODULE.bazel:3:20: The module extension ext defined in @ext//:defs.bzl"
+            + " reported incorrect imports of repositories via use_repo():\n"
+            + "\n"
+            + "Imported, but not created by the extension (will cause the build to fail):\n"
+            + "    invalid_dep, invalid_dev_dep\n"
+            + "\n"
+            + "Not imported, but reported as direct dependencies by the extension (may cause the"
+            + " build to fail):\n"
+            + "    missing_direct_dep, missing_direct_dev_dep\n"
+            + "\n"
+            + "\033[35m\033[1m ** You can use the following buildozer command(s) to fix these"
+            + " issues:\033[0m\n"
+            + "\n"
+            + "buildozer 'use_repo_remove @ext//:defs.bzl ext direct_dep indirect_dep invalid_dep'"
+            + " //MODULE.bazel:all\n"
+            + "buildozer 'use_repo_add dev @ext//:defs.bzl ext direct_dep indirect_dep"
+            + " missing_direct_dep missing_direct_dev_dep' //MODULE.bazel:all\n"
+            + "buildozer 'use_repo_remove dev @ext//:defs.bzl ext invalid_dev_dep'"
+            + " //MODULE.bazel:all",
+        ImmutableSet.of(EventKind.WARNING));
+  }
+
+  @Test
+  public void extensionMetadata_noRootUsage() throws Exception {
+    scratch.file(
+        workspaceRoot.getRelative("MODULE.bazel").getPathString(),
+        "bazel_dep(name='ext', version='1.0')",
+        "bazel_dep(name='data_repo',version='1.0')");
+    scratch.file(workspaceRoot.getRelative("BUILD").getPathString());
+
+    registry.addModule(
+        createModuleKey("ext", "1.0"),
+        "module(name='ext',version='1.0')",
+        "bazel_dep(name='data_repo',version='1.0')",
+        "ext = use_extension('//:defs.bzl', 'ext')",
+        "use_repo(ext, 'indirect_dep')",
+        "ext_dev = use_extension('//:defs.bzl', 'ext', dev_dependency = True)",
+        "use_repo(ext_dev, 'indirect_dev_dep')");
+    scratch.file(modulesRoot.getRelative("ext~1.0/WORKSPACE").getPathString());
+    scratch.file(modulesRoot.getRelative("ext~1.0/BUILD").getPathString());
+    scratch.file(
+        modulesRoot.getRelative("ext~1.0/defs.bzl").getPathString(),
+        "load('@data_repo//:defs.bzl','data_repo')",
+        "def _ext_impl(ctx):",
+        "  data_repo(name='direct_dep')",
+        "  data_repo(name='direct_dev_dep')",
+        "  data_repo(name='missing_direct_dep')",
+        "  data_repo(name='missing_direct_dev_dep')",
+        "  data_repo(name='indirect_dep', data='indirect_dep_data')",
+        "  data_repo(name='indirect_dev_dep')",
+        "  return ctx.extension_metadata(",
+        "    root_module_direct_deps='all',",
+        "    root_module_direct_dev_deps=[],",
+        "  )",
+        "ext=module_extension(implementation=_ext_impl)");
+    scratch.file(
+        modulesRoot.getRelative("ext~1.0/data.bzl").getPathString(),
+        "load('@indirect_dep//:data.bzl', indirect_dep_data='data')",
+        "data = indirect_dep_data");
+
+    SkyKey skyKey = BzlLoadValue.keyForBuild(Label.parseCanonical("@ext~1.0//:data.bzl"));
+    EvaluationResult<BzlLoadValue> result =
+        evaluator.evaluate(ImmutableList.of(skyKey), evaluationContext);
+    assertThat(result.get(skyKey).getModule().getGlobal("data")).isEqualTo("indirect_dep_data");
+
+    assertEventCount(0, eventCollector);
+  }
+
+  private EvaluationResult<SingleExtensionEvalValue> evaluateSimpleModuleExtension(
+      String returnStatement) throws Exception {
+    scratch.file(
+        workspaceRoot.getRelative("MODULE.bazel").getPathString(),
+        "ext = use_extension('//:defs.bzl', 'ext')");
+    scratch.file(
+        workspaceRoot.getRelative("defs.bzl").getPathString(),
+        "def _ext_impl(ctx):",
+        "  " + returnStatement,
+        "ext = module_extension(implementation=_ext_impl)");
+    scratch.file(workspaceRoot.getRelative("BUILD").getPathString());
+
+    ModuleExtensionId extensionId =
+        ModuleExtensionId.create(Label.parseCanonical("//:defs.bzl"), "ext");
+    reporter.removeHandler(failFastHandler);
+    return evaluator.evaluate(
+        ImmutableList.of(SingleExtensionEvalValue.key(extensionId)), evaluationContext);
+  }
 }
diff --git a/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleFileFunctionTest.java b/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleFileFunctionTest.java
index ecd8353f8b4605..d55e990a487364 100644
--- a/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleFileFunctionTest.java
+++ b/src/test/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleFileFunctionTest.java
@@ -459,7 +459,8 @@ public void testModuleExtensions_good() throws Exception {
                 "maven.dep(coord='guava')");
     ModuleFileFunction.REGISTRIES.set(differencer, ImmutableList.of(registry.getUrl()));
 
-    SkyKey skyKey = ModuleFileValue.key(createModuleKey("mymod", "1.0"), null);
+    ModuleKey myMod = createModuleKey("mymod", "1.0");
+    SkyKey skyKey = ModuleFileValue.key(myMod, null);
     EvaluationResult<ModuleFileValue> result =
         evaluator.evaluate(ImmutableList.of(skyKey), evaluationContext);
     if (result.hasError()) {
@@ -473,10 +474,12 @@ public void testModuleExtensions_good() throws Exception {
                 .setRegistry(registry)
                 .addExtensionUsage(
                     ModuleExtensionUsage.builder()
-                        .setExtensionBzlFile("//:defs.bzl")
+                        .setExtensionBzlFile("@mymod//:defs.bzl")
                         .setExtensionName("myext1")
+                        .setUsingModule(myMod)
                         .setLocation(Location.fromFileLineColumn("mymod@1.0/MODULE.bazel", 2, 23))
                         .setImports(ImmutableBiMap.of("repo1", "repo1"))
+                        .setDevImports(ImmutableSet.of())
                         .addTag(
                             Tag.builder()
                                 .setTagName("tag")
@@ -492,10 +495,12 @@ public void testModuleExtensions_good() throws Exception {
                         .build())
                 .addExtensionUsage(
                     ModuleExtensionUsage.builder()
-                        .setExtensionBzlFile("//:defs.bzl")
+                        .setExtensionBzlFile("@mymod//:defs.bzl")
                         .setExtensionName("myext2")
+                        .setUsingModule(myMod)
                         .setLocation(Location.fromFileLineColumn("mymod@1.0/MODULE.bazel", 5, 23))
                         .setImports(ImmutableBiMap.of("other_repo1", "repo1", "repo2", "repo2"))
+                        .setDevImports(ImmutableSet.of())
                         .addTag(
                             Tag.builder()
                                 .setTagName("tag1")
@@ -525,9 +530,11 @@ public void testModuleExtensions_good() throws Exception {
                     ModuleExtensionUsage.builder()
                         .setExtensionBzlFile("@rules_jvm_external//:defs.bzl")
                         .setExtensionName("maven")
+                        .setUsingModule(myMod)
                         .setLocation(Location.fromFileLineColumn("mymod@1.0/MODULE.bazel", 10, 22))
                         .setImports(
                             ImmutableBiMap.of("mvn", "maven", "junit", "junit", "guava", "guava"))
+                        .setDevImports(ImmutableSet.of())
                         .addTag(
                             Tag.builder()
                                 .setTagName("dep")
@@ -587,13 +594,15 @@ public void testModuleExtensions_duplicateProxy_asRoot() throws Exception {
                 .setKey(ModuleKey.ROOT)
                 .addExtensionUsage(
                     ModuleExtensionUsage.builder()
-                        .setExtensionBzlFile("//:defs.bzl")
+                        .setExtensionBzlFile("@//:defs.bzl")
                         .setExtensionName("myext")
+                        .setUsingModule(ModuleKey.ROOT)
                         .setLocation(Location.fromFileLineColumn("<root>/MODULE.bazel", 1, 23))
                         .setImports(
                             ImmutableBiMap.of(
                                 "alpha", "alpha", "beta", "beta", "gamma", "gamma", "delta",
                                 "delta"))
+                        .setDevImports(ImmutableSet.of("alpha", "gamma"))
                         .addTag(
                             Tag.builder()
                                 .setTagName("tag")
@@ -668,7 +677,8 @@ public void testModuleExtensions_duplicateProxy_asDep() throws Exception {
                 "use_repo(myext4, 'delta')");
     ModuleFileFunction.REGISTRIES.set(differencer, ImmutableList.of(registry.getUrl()));
 
-    SkyKey skyKey = ModuleFileValue.key(createModuleKey("mymod", "1.0"), null);
+    ModuleKey myMod = createModuleKey("mymod", "1.0");
+    SkyKey skyKey = ModuleFileValue.key(myMod, null);
     EvaluationResult<ModuleFileValue> result =
         evaluator.evaluate(ImmutableList.of(skyKey), evaluationContext);
     if (result.hasError()) {
@@ -681,10 +691,12 @@ public void testModuleExtensions_duplicateProxy_asDep() throws Exception {
                 .setRegistry(registry)
                 .addExtensionUsage(
                     ModuleExtensionUsage.builder()
-                        .setExtensionBzlFile("//:defs.bzl")
+                        .setExtensionBzlFile("@mymod//:defs.bzl")
                         .setExtensionName("myext")
+                        .setUsingModule(myMod)
                         .setLocation(Location.fromFileLineColumn("mymod@1.0/MODULE.bazel", 5, 23))
                         .setImports(ImmutableBiMap.of("beta", "beta", "delta", "delta"))
+                        .setDevImports(ImmutableSet.of())
                         .addTag(
                             Tag.builder()
                                 .setTagName("tag")
@@ -967,4 +979,34 @@ public void moduleRepoName_conflict() throws Exception {
 
     assertContainsEvent("The repo name 'bbb' is already being used as the module's own repo name");
   }
+
+  @Test
+  public void module_calledTwice() throws Exception {
+    scratch.file(
+        rootDirectory.getRelative("MODULE.bazel").getPathString(),
+        "module(name='aaa',version='0.1',repo_name='bbb')",
+        "module(name='aaa',version='0.1',repo_name='bbb')");
+    FakeRegistry registry = registryFactory.newFakeRegistry("/foo");
+    ModuleFileFunction.REGISTRIES.set(differencer, ImmutableList.of(registry.getUrl()));
+
+    reporter.removeHandler(failFastHandler); // expect failures
+    evaluator.evaluate(ImmutableList.of(ModuleFileValue.KEY_FOR_ROOT_MODULE), evaluationContext);
+
+    assertContainsEvent("the module() directive can only be called once");
+  }
+
+  @Test
+  public void module_calledLate() throws Exception {
+    scratch.file(
+        rootDirectory.getRelative("MODULE.bazel").getPathString(),
+        "use_extension('//:extensions.bzl', 'my_ext')",
+        "module(name='aaa',version='0.1',repo_name='bbb')");
+    FakeRegistry registry = registryFactory.newFakeRegistry("/foo");
+    ModuleFileFunction.REGISTRIES.set(differencer, ImmutableList.of(registry.getUrl()));
+
+    reporter.removeHandler(failFastHandler); // expect failures
+    evaluator.evaluate(ImmutableList.of(ModuleFileValue.KEY_FOR_ROOT_MODULE), evaluationContext);
+
+    assertContainsEvent("if module() is called, it must be called before any other functions");
+  }
 }
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 16fea23fe95156..3f655551f24383 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 com.google.common.collect.ImmutableBiMap;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.packages.BuildType;
 import com.google.devtools.build.lib.packages.Type;
@@ -43,8 +44,10 @@ private static ModuleExtensionUsage.Builder getBaseUsageBuilder() {
     return ModuleExtensionUsage.builder()
         .setExtensionBzlFile("//:rje.bzl")
         .setExtensionName("maven")
+        .setUsingModule(ModuleKey.ROOT)
         .setLocation(Location.BUILTIN)
-        .setImports(ImmutableBiMap.of());
+        .setImports(ImmutableBiMap.of())
+        .setDevImports(ImmutableSet.of());
   }
 
   /** A builder for ModuleExtension that sets all the mandatory but irrelevant fields. */
