nix: resolve buildbot extension from buildbot python set

flake.nix

@@ -48,8 +48,28 @@
           nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ final.setuptools ];
         });
       };
+      mkBuildbotExtPythonSet =
+        pkgs: python:
+        (pkgs.callPackage pyproject-nix.build.packages {
+          inherit python;
+        }).overrideScope
+          (
+            pkgs.lib.composeManyExtensions [
+              pyproject-build-systems.overlays.default
+              buildbotExtOverlay
+              pyprojectOverrides
+            ]
+          );
+      mkBuildbotAutoscaleExtPythonPackagesExtension = pkgs: pyFinal: _pyPrev: {
+        buildbot-autoscale-ext =
+          pyFinal.toPythonModule (mkBuildbotExtPythonSet pkgs pyFinal.python)."buildbot-autoscale-ext";
+      };
     in
     {
+      lib = {
+        inherit mkBuildbotAutoscaleExtPythonPackagesExtension;
+      };
+
       formatter = forAllSystems (
         pkgs: (treefmt-nix.lib.evalModule pkgs ./treefmt.nix).config.build.wrapper
       );
@@ -75,16 +95,7 @@
                 ]
               );
           buildbotExtPythonSet =
-            (pkgs.callPackage pyproject-nix.build.packages {
+            mkBuildbotExtPythonSet pkgs pkgs.python3;
-              python = pkgs.python3;
-            }).overrideScope
-              (
-                pkgs.lib.composeManyExtensions [
-                  pyproject-build-systems.overlays.default
-                  buildbotExtOverlay
-                  pyprojectOverrides
-                ]
-              );
           venv = agentPythonSet.mkVirtualEnv "nix-builder-autoscaler-env" agentWorkspace.deps.default;
           buildbotExtVenv = buildbotExtPythonSet.mkVirtualEnv "buildbot-autoscale-ext-env" buildbotExtWorkspace.deps.default;
         in
@@ -135,16 +146,7 @@
                 ]
               );
           buildbotExtPythonSet =
-            (pkgs.callPackage pyproject-nix.build.packages {
+            mkBuildbotExtPythonSet pkgs pkgs.python3;
-              python = pkgs.python3;
-            }).overrideScope
-              (
-                pkgs.lib.composeManyExtensions [
-                  pyproject-build-systems.overlays.default
-                  buildbotExtOverlay
-                  pyprojectOverrides
-                ]
-              );
           testVenv = agentPythonSet.mkVirtualEnv "nix-builder-autoscaler-test-env" {
             nix-builder-autoscaler = [ "dev" ];
           };

nix/modules/nixos/services/buildbot-nix-autoscaler.nix

@@ -1,21 +1,10 @@
 {
   config,
   lib,
-  pkgs,
   ...
 }:
 let
   cfg = config.services.buildbot-nix.nix-build-autoscaler;
-  defaultExtensionPackage =
-    if builtins.hasAttr "buildbot-autoscale-ext" pkgs then
-      pkgs."buildbot-autoscale-ext"
-    else if
-      builtins.hasAttr "python3Packages" pkgs
-      && builtins.hasAttr "buildbot-autoscale-ext" pkgs.python3Packages
-    then
-      pkgs.python3Packages."buildbot-autoscale-ext"
-    else
-      null;
 in
 {
   options.services.buildbot-nix.nix-build-autoscaler = {
@@ -23,8 +12,11 @@ in
 
     extensionPackage = lib.mkOption {
       type = lib.types.nullOr lib.types.package;
-      default = defaultExtensionPackage;
+      default = null;
-      description = "Package providing buildbot_autoscale_ext.";
+      description = ''
+        Optional explicit package override for buildbot_autoscale_ext.
+        Leave unset to resolve buildbot-autoscale-ext from Buildbot's pythonPackages set.
+      '';
     };
 
     daemonSocket = lib.mkOption {
@@ -135,20 +127,28 @@ in
 
   config = lib.mkIf cfg.enable {
     assertions = [
-      {
-        assertion = cfg.extensionPackage != null;
-        message = ''
-          services.buildbot-nix.nix-build-autoscaler.extensionPackage is not set and
-          pkgs.buildbot-autoscale-ext was not found. Configure extensionPackage explicitly.
-        '';
-      }
       {
         assertion = cfg.builderClusterHost != null;
         message = "services.buildbot-nix.nix-build-autoscaler.builderClusterHost must be set.";
       }
     ];
 
-    services.buildbot-master.pythonPackages = _ps: [ cfg.extensionPackage ];
+    services.buildbot-master.pythonPackages =
+      ps:
+      [
+        (
+          if cfg.extensionPackage != null then
+            ps.toPythonModule cfg.extensionPackage
+          else if builtins.hasAttr "buildbot-autoscale-ext" ps then
+            ps."buildbot-autoscale-ext"
+          else
+            throw ''
+              services.buildbot-nix.nix-build-autoscaler requires buildbot-autoscale-ext in
+              services.buildbot-master.pythonPackages set. Add a pythonPackagesExtensions overlay
+              providing buildbot-autoscale-ext, or set extensionPackage explicitly.
+            ''
+        )
+      ];
 
     services.buildbot-master.extraImports = ''
       from buildbot_autoscale_ext.configurator import AutoscaleConfigurator