diff --git a/services/core/java/com/android/server/media/MediaSessionService.java b/services/core/java/com/android/server/media/MediaSessionService.java
index e5867e7939484..476c9f879ca21 100644
--- a/services/core/java/com/android/server/media/MediaSessionService.java
+++ b/services/core/java/com/android/server/media/MediaSessionService.java
@@ -165,6 +165,10 @@ public class MediaSessionService extends SystemService implements Monitor {
         mAudioPlayerStateMonitor = AudioPlayerStateMonitor.getInstance(mContext);
         mAudioPlayerStateMonitor.registerListener(
                 (config, isRemoved) -> {
+                    if (DEBUG) {
+                        Log.d(TAG, "Audio playback is changed, config=" + config
+                                + ", removed=" + isRemoved);
+                    }
                     if (config.getPlayerType()
                             == AudioPlaybackConfiguration.PLAYER_TYPE_JAM_SOUNDPOOL) {
                         return;
@@ -1993,7 +1997,7 @@ public class MediaSessionService extends SystemService implements Monitor {
                     FullUserRecord user = getFullUserRecordLocked(record.getUserId());
                     if (record != null && user != null) {
                         record.setSessionPolicies(policies);
-                        user.mPriorityStack.updateMediaButtonSessionIfNeeded();
+                        user.mPriorityStack.updateMediaButtonSessionBySessionPolicyChange(record);
                     }
                 }
             } finally {
diff --git a/services/core/java/com/android/server/media/MediaSessionStack.java b/services/core/java/com/android/server/media/MediaSessionStack.java
index bd98b9cb64b21..402355a0161ac 100644
--- a/services/core/java/com/android/server/media/MediaSessionStack.java
+++ b/services/core/java/com/android/server/media/MediaSessionStack.java
@@ -16,6 +16,8 @@
 
 package com.android.server.media;
 
+import static com.android.server.media.SessionPolicyProvider.SESSION_POLICY_IGNORE_BUTTON_SESSION;
+
 import android.media.Session2Token;
 import android.media.session.MediaSession;
 import android.os.Debug;
@@ -102,6 +104,7 @@ class MediaSessionStack {
             // When the media button session is removed, nullify the media button session and do not
             // search for the alternative media session within the app. It's because the alternative
             // media session might be a dummy which isn't able to handle the media key events.
+            // TODO(b/154456172): Make this decision unaltered by non-media app's playback.
             updateMediaButtonSession(null);
         }
         clearCache(record.getUserId());
@@ -158,7 +161,7 @@ class MediaSessionStack {
                     findMediaButtonSession(mMediaButtonSession.getUid());
             if (newMediaButtonSession != mMediaButtonSession
                     && (newMediaButtonSession.getSessionPolicies()
-                    & SessionPolicyProvider.SESSION_POLICY_IGNORE_BUTTON_SESSION) == 0) {
+                            & SESSION_POLICY_IGNORE_BUTTON_SESSION) == 0) {
                 // Check if the policy states that this session should not be updated as a media
                 // button session.
                 updateMediaButtonSession(newMediaButtonSession);
@@ -189,19 +192,43 @@ class MediaSessionStack {
         }
         IntArray audioPlaybackUids = mAudioPlayerStateMonitor.getSortedAudioPlaybackClientUids();
         for (int i = 0; i < audioPlaybackUids.size(); i++) {
-            MediaSessionRecordImpl mediaButtonSession =
-                    findMediaButtonSession(audioPlaybackUids.get(i));
-            if (mediaButtonSession == null) continue;
-            boolean ignoreButtonSession = (mediaButtonSession.getSessionPolicies()
-                    & SessionPolicyProvider.SESSION_POLICY_IGNORE_BUTTON_SESSION) != 0;
-            if (mediaButtonSession == mMediaButtonSession && ignoreButtonSession) {
+            int audioPlaybackUid = audioPlaybackUids.get(i);
+            MediaSessionRecordImpl mediaButtonSession = findMediaButtonSession(audioPlaybackUid);
+            if (mediaButtonSession == null) {
+                if (DEBUG) {
+                    Log.d(TAG, "updateMediaButtonSessionIfNeeded, skipping uid="
+                            + audioPlaybackUid);
+                }
+                // Ignore if the lastly played app isn't a media app (i.e. has no media session)
+                continue;
+            }
+            boolean ignoreButtonSession =
+                    (mediaButtonSession.getSessionPolicies()
+                            & SESSION_POLICY_IGNORE_BUTTON_SESSION) != 0;
+            if (DEBUG) {
+                Log.d(TAG, "updateMediaButtonSessionIfNeeded, checking uid=" + audioPlaybackUid
+                        + ", mediaButtonSession=" + mediaButtonSession
+                        + ", ignoreButtonSession=" + ignoreButtonSession);
+            }
+            if (!ignoreButtonSession) {
+                mAudioPlayerStateMonitor.cleanUpAudioPlaybackUids(mediaButtonSession.getUid());
+                if (mediaButtonSession != mMediaButtonSession) {
+                    updateMediaButtonSession(mediaButtonSession);
+                }
+                return;
+            }
+        }
+    }
+
+    // TODO: Remove this and make updateMediaButtonSessionIfNeeded() to also cover this case.
+    public void updateMediaButtonSessionBySessionPolicyChange(MediaSessionRecord record) {
+        if ((record.getSessionPolicies() & SESSION_POLICY_IGNORE_BUTTON_SESSION) != 0) {
+            if (record == mMediaButtonSession) {
+                // TODO(b/154456172): Make this decision unaltered by non-media app's playback.
                 updateMediaButtonSession(null);
-                return;
-            }
-            if (mediaButtonSession != mMediaButtonSession && !ignoreButtonSession) {
-                updateMediaButtonSession(mediaButtonSession);
-                return;
             }
+        } else {
+            updateMediaButtonSessionIfNeeded();
         }
     }
 
@@ -280,7 +307,7 @@ class MediaSessionStack {
         return mMediaButtonSession;
     }
 
-    private void updateMediaButtonSession(MediaSessionRecordImpl newMediaButtonSession) {
+    public void updateMediaButtonSession(MediaSessionRecordImpl newMediaButtonSession) {
         MediaSessionRecordImpl oldMediaButtonSession = mMediaButtonSession;
         mMediaButtonSession = newMediaButtonSession;
         mOnMediaButtonSessionChangedListener.onMediaButtonSessionChanged(