diff --git a/Android.bp b/Android.bp
index b52a4e1..4e2b156 100644
--- a/Android.bp
+++ b/Android.bp
@@ -303,6 +303,7 @@
             "sax/java",
             "telecomm/java",
             "telephony/java",
+            "wifi/java",
         ],
     },
 
diff --git a/wifi/Android.bp b/wifi/Android.bp
index d14a726..70c9bef 100644
--- a/wifi/Android.bp
+++ b/wifi/Android.bp
@@ -113,8 +113,6 @@
         ":framework-annotations",
         ":framework-wifi-updatable-sources",
     ],
-    // This is needed as IOnWifiActivityEnergyInfoListener.aidl in framework-wifi depends on
-    // WifiActivityEnergyInfo.aidl in framework-minus-apex
     aidl: {
         include_dirs: ["frameworks/base/core/java"],
     },
@@ -126,6 +124,11 @@
 java_library {
     name: "framework-wifi-stubs",
     srcs: [":framework-wifi-stubs-srcs"],
+    aidl: {
+        export_include_dirs: [
+            "java",
+        ],
+    },
     sdk_version: "core_current",
     libs: ["android_system_stubs_current"],
     installable: false,