diff --git a/framework/Android.bp b/framework/Android.bp
index 441ad1f..c05ff69 100644
--- a/framework/Android.bp
+++ b/framework/Android.bp
@@ -44,6 +44,7 @@
         ":framework-wifi-updatable-java-sources",
         ":framework-wifi-updatable-exported-aidl-sources",
         ":module-utils-os-aidls",
+        ":wifi_aconfig_flags_lib{.generated_srcjars}",
     ],
 }
 
@@ -103,7 +104,7 @@
 // `framework-wifi-pre-jarjar`
 java_defaults {
     name: "framework-wifi-defaults",
-    defaults: ["wifi-module-sdk-version-defaults", "wifi-framework-aconfig-libraries"],
+    defaults: ["wifi-module-sdk-version-defaults"],
     static_libs: [
         "framework-wifi-util-lib",
         "android.hardware.wifi-V1.0-java-constants",
@@ -114,6 +115,9 @@
         "androidx.annotation_annotation",
         "unsupportedappusage", // for android.compat.annotation.UnsupportedAppUsage
         "app-compat-annotations",
+        // Add aconfig-annotations-lib as a dependency for the optimization
+        "aconfig-annotations-lib",
+        "framework-configinfrastructure",
     ],
     aidl: {
         generate_get_transaction_name: true,