diff --git a/packages/SystemUI/src/com/android/systemui/media/MediaData.kt b/packages/SystemUI/src/com/android/systemui/media/MediaData.kt
index 330a5c0..a94f6a8 100644
--- a/packages/SystemUI/src/com/android/systemui/media/MediaData.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/MediaData.kt
@@ -35,7 +35,8 @@
     val packageName: String?,
     val token: MediaSession.Token?,
     val clickIntent: PendingIntent?,
-    val device: MediaDeviceData?
+    val device: MediaDeviceData?,
+    val notificationKey: String = "INVALID"
 )
 
 /** State of a media action. */
diff --git a/packages/SystemUI/src/com/android/systemui/media/MediaDataManager.kt b/packages/SystemUI/src/com/android/systemui/media/MediaDataManager.kt
index cf7fbfa..d949857 100644
--- a/packages/SystemUI/src/com/android/systemui/media/MediaDataManager.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/MediaDataManager.kt
@@ -35,6 +35,8 @@
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.statusbar.notification.MediaNotificationProcessor
+import com.android.systemui.statusbar.notification.NotificationEntryManager
+import com.android.systemui.statusbar.notification.NotificationEntryManager.UNDEFINED_DISMISS_REASON
 import com.android.systemui.statusbar.notification.row.HybridGroupManager
 import com.android.systemui.util.Assert
 import com.android.systemui.util.Utils
@@ -77,6 +79,8 @@
 class MediaDataManager @Inject constructor(
     private val context: Context,
     private val mediaControllerFactory: MediaControllerFactory,
+    private val mediaTimeoutListener: MediaTimeoutListener,
+    private val notificationEntryManager: NotificationEntryManager,
     @Background private val backgroundExecutor: Executor,
     @Main private val foregroundExecutor: Executor
 ) {
@@ -84,6 +88,12 @@
     private val listeners: MutableSet<Listener> = mutableSetOf()
     private val mediaEntries: LinkedHashMap<String, MediaData> = LinkedHashMap()
 
+    init {
+        mediaTimeoutListener.timeoutCallback = { token: String, timedOut: Boolean ->
+            setTimedOut(token, timedOut) }
+        addListener(mediaTimeoutListener)
+    }
+
     fun onNotificationAdded(key: String, sbn: StatusBarNotification) {
         if (Utils.useQsMediaPlayer(context) && isMediaNotification(sbn)) {
             Assert.isMainThread()
@@ -112,6 +122,16 @@
      */
     fun removeListener(listener: Listener) = listeners.remove(listener)
 
+    private fun setTimedOut(token: String, timedOut: Boolean) {
+        if (!timedOut) {
+            return
+        }
+        mediaEntries[token]?.let {
+            notificationEntryManager.removeNotification(it.notificationKey, null /* ranking */,
+                    UNDEFINED_DISMISS_REASON)
+        }
+    }
+
     private fun loadMediaDataInBg(key: String, sbn: StatusBarNotification) {
         val token = sbn.notification.extras.getParcelable(Notification.EXTRA_MEDIA_SESSION)
                 as MediaSession.Token?
@@ -223,7 +243,7 @@
         foregroundExecutor.execute {
             onMediaDataLoaded(key, MediaData(true, bgColor, app, smallIconDrawable, artist, song,
                     artWorkIcon, actionIcons, actionsToShowCollapsed, sbn.packageName, token,
-                    notif.contentIntent, null))
+                    notif.contentIntent, null, key))
         }
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/media/MediaTimeoutListener.kt b/packages/SystemUI/src/com/android/systemui/media/MediaTimeoutListener.kt
new file mode 100644
index 0000000..92a1ab1
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/media/MediaTimeoutListener.kt
@@ -0,0 +1,117 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.media
+
+import android.media.session.MediaController
+import android.media.session.PlaybackState
+import android.os.SystemProperties
+import android.util.Log
+import com.android.systemui.dagger.qualifiers.Main
+import com.android.systemui.statusbar.NotificationMediaManager.isPlayingState
+import com.android.systemui.util.concurrency.DelayableExecutor
+import java.util.concurrent.TimeUnit
+import javax.inject.Inject
+import javax.inject.Singleton
+
+private const val DEBUG = true
+private const val TAG = "MediaTimeout"
+private val PAUSED_MEDIA_TIMEOUT = SystemProperties
+        .getLong("debug.sysui.media_timeout", TimeUnit.MINUTES.toMillis(10))
+
+/**
+ * Controller responsible for keeping track of playback states and expiring inactive streams.
+ */
+@Singleton
+class MediaTimeoutListener @Inject constructor(
+    private val mediaControllerFactory: MediaControllerFactory,
+    @Main private val mainExecutor: DelayableExecutor
+) : MediaDataManager.Listener {
+
+    private val mediaListeners: MutableMap<String, PlaybackStateListener> = mutableMapOf()
+
+    lateinit var timeoutCallback: (String, Boolean) -> Unit
+
+    override fun onMediaDataLoaded(key: String, data: MediaData) {
+        if (mediaListeners.containsKey(key)) {
+            return
+        }
+        mediaListeners[key] = PlaybackStateListener(key, data)
+    }
+
+    override fun onMediaDataRemoved(key: String) {
+        mediaListeners.remove(key)?.destroy()
+    }
+
+    fun isTimedOut(key: String): Boolean {
+        return mediaListeners[key]?.timedOut ?: false
+    }
+
+    private inner class PlaybackStateListener(
+        private val key: String,
+        data: MediaData
+    ) : MediaController.Callback() {
+
+        var timedOut = false
+
+        private val mediaController = mediaControllerFactory.create(data.token)
+        private var cancellation: Runnable? = null
+
+        init {
+            mediaController.registerCallback(this)
+        }
+
+        fun destroy() {
+            mediaController.unregisterCallback(this)
+        }
+
+        override fun onPlaybackStateChanged(state: PlaybackState?) {
+            if (DEBUG) {
+                Log.v(TAG, "onPlaybackStateChanged: $state")
+            }
+            expireMediaTimeout(key, "playback state ativity - $state, $key")
+
+            if (state == null || !isPlayingState(state.state)) {
+                if (DEBUG) {
+                    Log.v(TAG, "schedule timeout for $key")
+                }
+                expireMediaTimeout(key, "PLAYBACK STATE CHANGED - $state")
+                cancellation = mainExecutor.executeDelayed({
+                    cancellation = null
+                    if (DEBUG) {
+                        Log.v(TAG, "Execute timeout for $key")
+                    }
+                    timedOut = true
+                    timeoutCallback(key, timedOut)
+                }, PAUSED_MEDIA_TIMEOUT)
+            } else {
+                timedOut = false
+                timeoutCallback(key, timedOut)
+            }
+        }
+
+        private fun expireMediaTimeout(mediaNotificationKey: String, reason: String) {
+            cancellation?.apply {
+                if (DEBUG) {
+                    Log.v(TAG,
+                            "media timeout cancelled for  $mediaNotificationKey, reason: $reason")
+                }
+                run()
+            }
+            cancellation = null
+        }
+    }
+}
\ No newline at end of file
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationMediaManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationMediaManager.java
index db5329a..217148d 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationMediaManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationMediaManager.java
@@ -16,7 +16,6 @@
 package com.android.systemui.statusbar;
 
 import static com.android.systemui.statusbar.StatusBarState.KEYGUARD;
-import static com.android.systemui.statusbar.notification.NotificationEntryManager.UNDEFINED_DISMISS_REASON;
 import static com.android.systemui.statusbar.phone.StatusBar.DEBUG_MEDIA_FAKE_ARTWORK;
 import static com.android.systemui.statusbar.phone.StatusBar.ENABLE_LOCKSCREEN_WALLPAPER;
 import static com.android.systemui.statusbar.phone.StatusBar.SHOW_LOCKSCREEN_MEDIA_ARTWORK;
@@ -36,7 +35,6 @@
 import android.media.session.MediaSessionManager;
 import android.media.session.PlaybackState;
 import android.os.AsyncTask;
-import android.os.SystemProperties;
 import android.os.Trace;
 import android.os.UserHandle;
 import android.provider.DeviceConfig;
@@ -80,7 +78,6 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
-import java.util.concurrent.TimeUnit;
 
 import dagger.Lazy;
 
@@ -91,8 +88,6 @@
 public class NotificationMediaManager implements Dumpable {
     private static final String TAG = "NotificationMediaManager";
     public static final boolean DEBUG_MEDIA = false;
-    private static final long PAUSED_MEDIA_TIMEOUT = SystemProperties
-            .getLong("debug.sysui.media_timeout", TimeUnit.MINUTES.toMillis(10));
 
     private final StatusBarStateController mStatusBarStateController
             = Dependency.get(StatusBarStateController.class);
@@ -134,7 +129,6 @@
     private MediaController mMediaController;
     private String mMediaNotificationKey;
     private MediaMetadata mMediaMetadata;
-    private Runnable mMediaTimeoutCancellation;
 
     private BackDropView mBackdrop;
     private ImageView mBackdropFront;
@@ -164,47 +158,11 @@
             if (DEBUG_MEDIA) {
                 Log.v(TAG, "DEBUG_MEDIA: onPlaybackStateChanged: " + state);
             }
-            if (mMediaTimeoutCancellation != null) {
-                if (DEBUG_MEDIA) {
-                    Log.v(TAG, "DEBUG_MEDIA: media timeout cancelled");
-                }
-                mMediaTimeoutCancellation.run();
-                mMediaTimeoutCancellation = null;
-            }
             if (state != null) {
                 if (!isPlaybackActive(state.getState())) {
                     clearCurrentMediaNotification();
                 }
                 findAndUpdateMediaNotifications();
-                scheduleMediaTimeout(state);
-            }
-        }
-
-        private void scheduleMediaTimeout(PlaybackState state) {
-            final NotificationEntry entry;
-            synchronized (mEntryManager) {
-                entry = mEntryManager.getActiveNotificationUnfiltered(mMediaNotificationKey);
-            }
-            if (entry != null) {
-                if (!isPlayingState(state.getState())) {
-                    if (DEBUG_MEDIA) {
-                        Log.v(TAG, "DEBUG_MEDIA: schedule timeout for "
-                                + mMediaNotificationKey);
-                    }
-                    mMediaTimeoutCancellation = mMainExecutor.executeDelayed(() -> {
-                        synchronized (mEntryManager) {
-                            if (DEBUG_MEDIA) {
-                                Log.v(TAG, "DEBUG_MEDIA: execute timeout for "
-                                        + mMediaNotificationKey);
-                            }
-                            if (mMediaNotificationKey == null) {
-                                return;
-                            }
-                            mEntryManager.removeNotification(mMediaNotificationKey, null,
-                                    UNDEFINED_DISMISS_REASON);
-                        }
-                    }, PAUSED_MEDIA_TIMEOUT);
-                }
             }
         }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/media/MediaDataCombineLatestTest.java b/packages/SystemUI/tests/src/com/android/systemui/media/MediaDataCombineLatestTest.java
index aa889a6..48e3b0a 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/media/MediaDataCombineLatestTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/media/MediaDataCombineLatestTest.java
@@ -79,7 +79,7 @@
         mManager.addListener(mListener);
 
         mMediaData = new MediaData(true, BG_COLOR, APP, null, ARTIST, TITLE, null,
-                new ArrayList<>(), new ArrayList<>(), PACKAGE, null, null, null);
+                new ArrayList<>(), new ArrayList<>(), PACKAGE, null, null, null, KEY);
         mDeviceData = new MediaDeviceData(true, null, DEVICE_NAME);
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/media/MediaTimeoutListenerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/media/MediaTimeoutListenerTest.kt
new file mode 100644
index 0000000..c21343c
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/media/MediaTimeoutListenerTest.kt
@@ -0,0 +1,130 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.media
+
+import android.media.session.MediaController
+import android.media.session.PlaybackState
+import android.testing.AndroidTestingRunner
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.util.concurrency.DelayableExecutor
+import com.android.systemui.util.mockito.capture
+import com.google.common.truth.Truth.assertThat
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.ArgumentCaptor
+import org.mockito.ArgumentMatchers.any
+import org.mockito.ArgumentMatchers.anyLong
+import org.mockito.Captor
+import org.mockito.Mock
+import org.mockito.Mockito
+import org.mockito.Mockito.`when`
+import org.mockito.Mockito.clearInvocations
+import org.mockito.Mockito.never
+import org.mockito.Mockito.verify
+import org.mockito.junit.MockitoJUnit
+
+private const val KEY = "KEY"
+
+private fun <T> eq(value: T): T = Mockito.eq(value) ?: value
+private fun <T> anyObject(): T {
+    return Mockito.anyObject<T>()
+}
+
+@SmallTest
+@RunWith(AndroidTestingRunner::class)
+class MediaTimeoutListenerTest : SysuiTestCase() {
+
+    @Mock private lateinit var mediaControllerFactory: MediaControllerFactory
+    @Mock private lateinit var mediaController: MediaController
+    @Mock private lateinit var executor: DelayableExecutor
+    @Mock private lateinit var mediaData: MediaData
+    @Mock private lateinit var timeoutCallback: (String, Boolean) -> Unit
+    @Mock private lateinit var cancellationRunnable: Runnable
+    @Captor private lateinit var timeoutCaptor: ArgumentCaptor<Runnable>
+    @Captor private lateinit var mediaCallbackCaptor: ArgumentCaptor<MediaController.Callback>
+    @JvmField @Rule val mockito = MockitoJUnit.rule()
+    private lateinit var mediaTimeoutListener: MediaTimeoutListener
+
+    @Before
+    fun setup() {
+        `when`(mediaControllerFactory.create(any())).thenReturn(mediaController)
+        `when`(executor.executeDelayed(any(), anyLong())).thenReturn(cancellationRunnable)
+        mediaTimeoutListener = MediaTimeoutListener(mediaControllerFactory, executor)
+        mediaTimeoutListener.timeoutCallback = timeoutCallback
+    }
+
+    @Test
+    fun testOnMediaDataLoaded_registersPlaybackListener() {
+        mediaTimeoutListener.onMediaDataLoaded(KEY, mediaData)
+        verify(mediaController).registerCallback(capture(mediaCallbackCaptor))
+
+        // Ignores is same key
+        clearInvocations(mediaController)
+        mediaTimeoutListener.onMediaDataLoaded(KEY, mediaData)
+        verify(mediaController, never()).registerCallback(anyObject())
+    }
+
+    @Test
+    fun testOnMediaDataRemoved_unregistersPlaybackListener() {
+        mediaTimeoutListener.onMediaDataLoaded(KEY, mediaData)
+        mediaTimeoutListener.onMediaDataRemoved(KEY)
+        verify(mediaController).unregisterCallback(anyObject())
+
+        // Ignores duplicate requests
+        clearInvocations(mediaController)
+        mediaTimeoutListener.onMediaDataRemoved(KEY)
+        verify(mediaController, never()).unregisterCallback(anyObject())
+    }
+
+    @Test
+    fun testOnPlaybackStateChanged_schedulesTimeout_whenPaused() {
+        // Assuming we're registered
+        testOnMediaDataLoaded_registersPlaybackListener()
+
+        mediaCallbackCaptor.value.onPlaybackStateChanged(PlaybackState.Builder()
+                .setState(PlaybackState.STATE_PAUSED, 0L, 0f).build())
+        verify(executor).executeDelayed(capture(timeoutCaptor), anyLong())
+    }
+
+    @Test
+    fun testOnPlaybackStateChanged_cancelsTimeout_whenResumed() {
+        // Assuming we're have a pending timeout
+        testOnPlaybackStateChanged_schedulesTimeout_whenPaused()
+
+        mediaCallbackCaptor.value.onPlaybackStateChanged(PlaybackState.Builder()
+                .setState(PlaybackState.STATE_PLAYING, 0L, 0f).build())
+        verify(cancellationRunnable).run()
+    }
+
+    @Test
+    fun testTimeoutCallback_invokedIfTimeout() {
+        // Assuming we're have a pending timeout
+        testOnPlaybackStateChanged_schedulesTimeout_whenPaused()
+
+        timeoutCaptor.value.run()
+        verify(timeoutCallback).invoke(eq(KEY), eq(true))
+    }
+
+    @Test
+    fun testIsTimedOut() {
+        mediaTimeoutListener.onMediaDataLoaded(KEY, mediaData)
+        assertThat(mediaTimeoutListener.isTimedOut(KEY)).isFalse()
+    }
+}
\ No newline at end of file