diff --git a/core/Makefile b/core/Makefile
index 7b4f58d..447dcb6 100644
--- a/core/Makefile
+++ b/core/Makefile
@@ -84,21 +84,6 @@
 endif
 
 
-###########################################################
-# Get the module names suitable for ALL_MODULES.* variables that are installed
-# for a given partition
-#
-# $(1): Partition
-###########################################################
-define register-names-for-partition
-$(sort $(foreach m,$(product_MODULES),\
-	$(if $(filter $(PRODUCT_OUT)/$(strip $(1))/%, $(ALL_MODULES.$(m).INSTALLED)), \
-		$(m)
-	) \
-))
-endef
-
-
 # Release & Aconfig Flags
 # -----------------------------------------------------------------
 include $(BUILD_SYSTEM)/packaging/flags.mk
diff --git a/core/packaging/flags.mk b/core/packaging/flags.mk
index a371a00..19068f4 100644
--- a/core/packaging/flags.mk
+++ b/core/packaging/flags.mk
@@ -17,9 +17,8 @@
 # the combined flags files.
 #
 
-# TODO: Should we do all of the images in $(IMAGES_TO_BUILD)?
-_FLAG_PARTITIONS := product system vendor
-
+# TODO: Should we do all of the images?
+_FLAG_PARTITIONS := product system system_ext vendor
 
 # -----------------------------------------------------------------
 # Aconfig Flags
@@ -62,28 +61,38 @@
 $(call copy-one-file, $(1), $(2))
 endef
 
+define out-dir-for-partition
+$(TARGET_COPY_OUT_$(call to-upper,$(1)))
+endef
+
+# Get the module names suitable for ALL_MODULES.* variables that are installed
+# for a given container
+# $(1): container
+define register-names-for-container
+$(sort $(foreach m,$(product_MODULES),\
+	$(if $(filter $(PRODUCT_OUT)/$(call out-dir-for-partition,$(strip $(1)))/%, $(ALL_MODULES.$(m).INSTALLED)), \
+		$(m)
+	) \
+))
+endef
+
 $(foreach partition, $(_FLAG_PARTITIONS), \
-	$(eval aconfig_flag_summaries_protobuf.$(partition) := $(PRODUCT_OUT)/$(partition)/etc/aconfig_flags.pb) \
+	$(eval aconfig_flag_summaries_protobuf.$(partition) := $(PRODUCT_OUT)/$(call out-dir-for-partition,$(partition))/etc/aconfig_flags.pb) \
 	$(eval $(call generate-partition-aconfig-flag-file, \
 			$(TARGET_OUT_FLAGS)/$(partition)/aconfig_flags.pb, \
 			$(aconfig_flag_summaries_protobuf.$(partition)), \
 			$(partition), \
 			$(sort \
-				$(foreach m, $(call register-names-for-partition, $(partition)), \
+				$(foreach m, $(call register-names-for-container, $(partition)), \
 					$(ALL_MODULES.$(m).ACONFIG_FILES) \
 				) \
-				$(if $(filter system, $(partition)), \
-					$(foreach m, $(call register-names-for-partition, system_ext), \
-						$(ALL_MODULES.$(m).ACONFIG_FILES) \
-					) \
-				) \
 			) \
 	)) \
 )
 
 # Collect the on-device flags into a single file, similar to all_aconfig_declarations.
 required_aconfig_flags_files := \
-		$(sort $(foreach partition, $(filter $(IMAGES_TO_BUILD), $(_FLAG_PARTITIONS)), \
+		$(sort $(foreach partition, $(_FLAG_PARTITIONS), \
 			$(aconfig_flag_summaries_protobuf.$(partition)) \
 		))
 
@@ -158,10 +167,10 @@
 
 ifeq ($(RELEASE_CREATE_ACONFIG_STORAGE_FILE),true)
 $(foreach partition, $(_FLAG_PARTITIONS), \
-	$(eval aconfig_storage_package_map.$(partition) := $(PRODUCT_OUT)/$(partition)/etc/aconfig/package.map) \
-	$(eval aconfig_storage_flag_map.$(partition) := $(PRODUCT_OUT)/$(partition)/etc/aconfig/flag.map) \
-	$(eval aconfig_storage_flag_val.$(partition) := $(PRODUCT_OUT)/$(partition)/etc/aconfig/flag.val) \
-	$(eval aconfig_storage_flag_info.$(partition) := $(PRODUCT_OUT)/$(partition)/etc/aconfig/flag.info) \
+	$(eval aconfig_storage_package_map.$(partition) := $(PRODUCT_OUT)/$(call out-dir-for-partition,$(partition))/etc/aconfig/package.map) \
+	$(eval aconfig_storage_flag_map.$(partition) := $(PRODUCT_OUT)/$(call out-dir-for-partition,$(partition))/etc/aconfig/flag.map) \
+	$(eval aconfig_storage_flag_val.$(partition) := $(PRODUCT_OUT)/$(call out-dir-for-partition,$(partition))/etc/aconfig/flag.val) \
+	$(eval aconfig_storage_flag_info.$(partition) := $(PRODUCT_OUT)/$(call out-dir-for-partition,$(partition))/etc/aconfig/flag.info) \
 	$(eval $(call generate-partition-aconfig-storage-file, \
 				$(TARGET_OUT_FLAGS)/$(partition)/package.map, \
 				$(TARGET_OUT_FLAGS)/$(partition)/flag.map, \
@@ -180,7 +189,7 @@
 # -----------------------------------------------------------------
 # Install the ones we need for the configured product
 required_flags_files := \
-		$(sort $(foreach partition, $(filter $(IMAGES_TO_BUILD), $(_FLAG_PARTITIONS)), \
+		$(sort $(foreach partition, $(_FLAG_PARTITIONS), \
 			$(build_flag_summaries.$(partition)) \
 			$(aconfig_flag_summaries_protobuf.$(partition)) \
 			$(aconfig_storage_package_map.$(partition)) \
@@ -198,6 +207,8 @@
 
 
 # Clean up
+out-dir-for-partition:=
+register-names-for-container:=
 required_flags_files:=
 required_aconfig_flags_files:=
 $(foreach partition, $(_FLAG_PARTITIONS), \
diff --git a/tools/aconfig/aconfig/src/commands.rs b/tools/aconfig/aconfig/src/commands.rs
index ab726aa..64fb2b6 100644
--- a/tools/aconfig/aconfig/src/commands.rs
+++ b/tools/aconfig/aconfig/src/commands.rs
@@ -80,18 +80,8 @@
             .read_to_string(&mut contents)
             .with_context(|| format!("failed to read {}", input.source))?;
 
-        let mut flag_declarations =
-            aconfig_protos::flag_declarations::try_from_text_proto(&contents)
-                .with_context(|| input.error_context())?;
-
-        // system_ext flags should be treated as system flags as we are combining /system_ext
-        // and /system as one container
-        // TODO: remove this logic when we start enforcing that system_ext cannot be set as
-        // container in aconfig declaration files.
-        if flag_declarations.container() == "system_ext" {
-            flag_declarations.set_container(String::from("system"));
-        }
-
+        let flag_declarations = aconfig_protos::flag_declarations::try_from_text_proto(&contents)
+            .with_context(|| input.error_context())?;
         ensure!(
             package == flag_declarations.package(),
             "failed to parse {}: expected package {}, got {}",
@@ -459,6 +449,7 @@
 pub fn should_include_flag(pf: &ProtoParsedFlag) -> bool {
     let should_filter_container = pf.container == Some("vendor".to_string())
         || pf.container == Some("system".to_string())
+        || pf.container == Some("system_ext".to_string())
         || pf.container == Some("product".to_string());
 
     let disabled_ro = pf.state == Some(ProtoFlagState::DISABLED.into())
diff --git a/tools/aconfig/aconfig_device_paths/partition_aconfig_flags_paths.txt b/tools/aconfig/aconfig_device_paths/partition_aconfig_flags_paths.txt
index e997e3d..140cd21 100644
--- a/tools/aconfig/aconfig_device_paths/partition_aconfig_flags_paths.txt
+++ b/tools/aconfig/aconfig_device_paths/partition_aconfig_flags_paths.txt
@@ -1,3 +1,4 @@
 "/system/etc/aconfig_flags.pb",
+"/system_ext/etc/aconfig_flags.pb",
 "/product/etc/aconfig_flags.pb",
 "/vendor/etc/aconfig_flags.pb",
diff --git a/tools/aconfig/aconfig_device_paths/src/lib.rs b/tools/aconfig/aconfig_device_paths/src/lib.rs
index 8871b4f..9ab9cea 100644
--- a/tools/aconfig/aconfig_device_paths/src/lib.rs
+++ b/tools/aconfig/aconfig_device_paths/src/lib.rs
@@ -62,12 +62,13 @@
 
     #[test]
     fn test_read_partition_paths() {
-        assert_eq!(read_partition_paths().len(), 3);
+        assert_eq!(read_partition_paths().len(), 4);
 
         assert_eq!(
             read_partition_paths(),
             vec![
                 PathBuf::from("/system/etc/aconfig_flags.pb"),
+                PathBuf::from("/system_ext/etc/aconfig_flags.pb"),
                 PathBuf::from("/product/etc/aconfig_flags.pb"),
                 PathBuf::from("/vendor/etc/aconfig_flags.pb")
             ]