diff --git a/packages/SystemUI/src/com/android/systemui/fragments/ExtensionFragmentListener.java b/packages/SystemUI/src/com/android/systemui/fragments/ExtensionFragmentListener.java
index 880951036661b..4ff10e975b1a7 100644
--- a/packages/SystemUI/src/com/android/systemui/fragments/ExtensionFragmentListener.java
+++ b/packages/SystemUI/src/com/android/systemui/fragments/ExtensionFragmentListener.java
@@ -34,12 +34,14 @@ public class ExtensionFragmentListener<T extends FragmentBase> implements Consum
     private final FragmentHostManager mFragmentHostManager;
     private final String mTag;
     private final Extension<T> mExtension;
+    private final int mId;
     private String mOldClass;
 
     private ExtensionFragmentListener(View view, String tag, int id, Extension<T> extension) {
         mTag = tag;
         mFragmentHostManager = FragmentHostManager.get(view);
         mExtension = extension;
+        mId = id;
         mFragmentHostManager.getFragmentManager().beginTransaction()
                 .replace(id, (Fragment) mExtension.get(), mTag)
                 .commit();
@@ -49,7 +51,7 @@ public class ExtensionFragmentListener<T extends FragmentBase> implements Consum
     public void accept(T extension) {
         try {
             Fragment.class.cast(extension);
-            mFragmentHostManager.getExtensionManager().setCurrentExtension(mTag,
+            mFragmentHostManager.getExtensionManager().setCurrentExtension(mId, mTag,
                     mOldClass, extension.getClass().getName(), mExtension.getContext());
             mOldClass = extension.getClass().getName();
         } catch (ClassCastException e) {
diff --git a/packages/SystemUI/src/com/android/systemui/fragments/FragmentHostManager.java b/packages/SystemUI/src/com/android/systemui/fragments/FragmentHostManager.java
index 871f113ac95e9..f8f364da089b9 100644
--- a/packages/SystemUI/src/com/android/systemui/fragments/FragmentHostManager.java
+++ b/packages/SystemUI/src/com/android/systemui/fragments/FragmentHostManager.java
@@ -266,16 +266,14 @@ public class FragmentHostManager {
     class ExtensionFragmentManager {
         private final ArrayMap<String, Context> mExtensionLookup = new ArrayMap<>();
 
-        public void setCurrentExtension(@NonNull  String tag, @Nullable String oldClass,
+        public void setCurrentExtension(int id, @NonNull  String tag, @Nullable String oldClass,
                 @NonNull String currentClass, @Nullable Context context) {
-            Fragment fragment = getFragmentManager().findFragmentByTag(tag);
             if (oldClass != null) {
                 mExtensionLookup.remove(oldClass);
             }
             mExtensionLookup.put(currentClass, context);
             getFragmentManager().beginTransaction()
-                    .replace(((View) fragment.getView().getParent()).getId(),
-                            instantiate(context, currentClass, null), tag)
+                    .replace(id, instantiate(context, currentClass, null), tag)
                     .commit();
             reloadFragments();
         }