diff --git a/src/com/android/server/telecom/HeadsetMediaButton.java b/src/com/android/server/telecom/HeadsetMediaButton.java
index 7458f54..8e9caff 100644
--- a/src/com/android/server/telecom/HeadsetMediaButton.java
+++ b/src/com/android/server/telecom/HeadsetMediaButton.java
@@ -23,11 +23,16 @@
 import android.os.Handler;
 import android.os.Looper;
 import android.os.Message;
+import android.telecom.CallAudioState;
+import android.telecom.CallEndpoint;
 import android.telecom.Log;
+import android.util.ArraySet;
 import android.view.KeyEvent;
 
 import com.android.internal.annotations.VisibleForTesting;
 
+import java.util.Set;
+
 /**
  * Static class to handle listening to the headset media buttons.
  */
@@ -149,8 +154,10 @@
     private final Context mContext;
     private final CallsManager mCallsManager;
     private final TelecomSystem.SyncRoot mLock;
+    private final Set<Call> mCalls = new ArraySet<>();
     private MediaSessionAdapter mSession;
     private KeyEvent mLastHookEvent;
+    private @CallEndpoint.EndpointType int mCurrentEndpointType;
 
     /**
      * Constructor used for testing purposes to initialize a {@link HeadsetMediaButton} with a
@@ -212,7 +219,7 @@
             return mCallsManager.onMediaButton(LONG_PRESS);
         } else if (event.getAction() == KeyEvent.ACTION_UP) {
             // We should not judge SHORT_PRESS by ACTION_UP event repeatCount, because it always
-            // return 0.
+            // returns 0.
             // Actually ACTION_DOWN event repeatCount only increases when LONG_PRESS performed.
             if (mLastHookEvent != null && mLastHookEvent.getRepeatCount() == 0) {
                 return mCallsManager.onMediaButton(SHORT_PRESS);
@@ -226,52 +233,72 @@
         return true;
     }
 
+    @Override
+    public void onCallEndpointChanged(CallEndpoint callEndpoint) {
+        mCurrentEndpointType = callEndpoint.getEndpointType();
+        Log.i(this, "onCallEndpointChanged: endPoint=%s", callEndpoint);
+        maybeChangeSessionState();
+    }
+
     /** ${inheritDoc} */
     @Override
     public void onCallAdded(Call call) {
-        if (call.isExternalCall()) {
-            return;
-        }
-        handleCallAddition();
+        handleCallAddition(call);
     }
 
     /**
      * Triggers session activation due to call addition.
      */
-    private void handleCallAddition() {
-        mMediaSessionHandler.obtainMessage(MSG_MEDIA_SESSION_SET_ACTIVE, 1, 0).sendToTarget();
-    }
-
-    /** ${inheritDoc} */
-    @Override
-    public void onCallRemoved(Call call) {
-        if (call.isExternalCall()) {
-            return;
-        }
-        handleCallRemoval();
+    private void handleCallAddition(Call call) {
+        mCalls.add(call);
+        maybeChangeSessionState();
     }
 
     /**
-     * Triggers session deactivation due to call removal.
+     * Based on whether there are tracked calls and the audio is routed to a wired headset,
+     * potentially activate or deactive the media session.
      */
-    private void handleCallRemoval() {
-        if (!mCallsManager.hasAnyCalls()) {
+    private void maybeChangeSessionState() {
+        boolean hasNonExternalCalls = !mCalls.isEmpty()
+                && mCalls.stream().anyMatch(c -> !c.isExternalCall());
+        if (hasNonExternalCalls && mCurrentEndpointType == CallEndpoint.TYPE_WIRED_HEADSET) {
+            Log.i(this, "maybeChangeSessionState: hasCalls=%b, currentEndpointType=%s, ACTIVATE",
+                    hasNonExternalCalls, CallEndpoint.endpointTypeToString(mCurrentEndpointType));
+            mMediaSessionHandler.obtainMessage(MSG_MEDIA_SESSION_SET_ACTIVE, 1, 0).sendToTarget();
+        } else {
+            Log.i(this, "maybeChangeSessionState: hasCalls=%b, currentEndpointType=%s, DEACTIVATE",
+                    hasNonExternalCalls, CallEndpoint.endpointTypeToString(mCurrentEndpointType));
             mMediaSessionHandler.obtainMessage(MSG_MEDIA_SESSION_SET_ACTIVE, 0, 0).sendToTarget();
         }
     }
 
     /** ${inheritDoc} */
     @Override
-    public void onExternalCallChanged(Call call, boolean isExternalCall) {
-        // Note: We don't use the onCallAdded/onCallRemoved methods here since they do checks to see
-        // if the call is external or not and would skip the session activation/deactivation.
-        if (isExternalCall) {
-            handleCallRemoval();
-        } else {
-            handleCallAddition();
+    public void onCallRemoved(Call call) {
+        handleCallRemoval(call);
+    }
+
+    /**
+     * Triggers session deactivation due to call removal.
+     */
+    private void handleCallRemoval(Call call) {
+        // If we were tracking the call, potentially change session state.
+        if (mCalls.remove(call)) {
+            if (mCalls.isEmpty()) {
+                // When there are no calls, don't cache that we previously had a wired headset
+                // connected; we'll be updated on the next call.
+                mCurrentEndpointType = CallEndpoint.TYPE_UNKNOWN;
+            }
+            maybeChangeSessionState();
         }
     }
 
+    /** ${inheritDoc} */
+    @Override
+    public void onExternalCallChanged(Call call, boolean isExternalCall) {
+        maybeChangeSessionState();
+    }
+
     @VisibleForTesting
     /**
      * @return the handler this class instance uses for operation; used for unit testing.
diff --git a/tests/src/com/android/server/telecom/tests/HeadsetMediaButtonTest.java b/tests/src/com/android/server/telecom/tests/HeadsetMediaButtonTest.java
index 0bfa987..ce23724 100644
--- a/tests/src/com/android/server/telecom/tests/HeadsetMediaButtonTest.java
+++ b/tests/src/com/android/server/telecom/tests/HeadsetMediaButtonTest.java
@@ -18,6 +18,7 @@
 
 import android.content.Intent;
 import android.media.session.MediaSession;
+import android.telecom.CallEndpoint;
 import android.test.suitebuilder.annotation.MediumTest;
 import android.test.suitebuilder.annotation.SmallTest;
 import android.view.KeyEvent;
@@ -80,7 +81,7 @@
     }
 
     /**
-     * Nominal case; just add a call and remove it.
+     * Nominal case; just add a call and remove it; this happens when the audio state is earpiece.
      */
     @SmallTest
     @Test
@@ -90,14 +91,95 @@
         when(mMockCallsManager.hasAnyCalls()).thenReturn(true);
         mHeadsetMediaButton.onCallAdded(regularCall);
         waitForHandlerAction(mHeadsetMediaButton.getHandler(), TEST_TIMEOUT_MILLIS);
+        verify(mMediaSessionAdapter, never()).setActive(eq(true));
+
+        // Report that the endpoint is earpiece and other routes that don't matter
+        mHeadsetMediaButton.onCallEndpointChanged(
+                new CallEndpoint("Earpiece", CallEndpoint.TYPE_EARPIECE));
+        mHeadsetMediaButton.onCallEndpointChanged(
+                new CallEndpoint("Speaker", CallEndpoint.TYPE_SPEAKER));
+        mHeadsetMediaButton.onCallEndpointChanged(
+                new CallEndpoint("BT", CallEndpoint.TYPE_BLUETOOTH));
+        waitForHandlerAction(mHeadsetMediaButton.getHandler(), TEST_TIMEOUT_MILLIS);
+        verify(mMediaSessionAdapter, never()).setActive(eq(true));
+
+        // ... and thus we see how the original code isn't amenable to tests.
+        when(mMediaSessionAdapter.isActive()).thenReturn(false);
+
+        // Still should not have done anything; we never hit wired headset
+        mHeadsetMediaButton.onCallRemoved(regularCall);
+        waitForHandlerAction(mHeadsetMediaButton.getHandler(), TEST_TIMEOUT_MILLIS);
+        verify(mMediaSessionAdapter, never()).setActive(eq(false));
+    }
+
+    /**
+     * Call is added and then routed to headset after call start
+     */
+    @SmallTest
+    @Test
+    public void testAddCallThenRouteToHeadset() {
+        Call regularCall = getRegularCall();
+
+        mHeadsetMediaButton.onCallAdded(regularCall);
+        waitForHandlerAction(mHeadsetMediaButton.getHandler(), TEST_TIMEOUT_MILLIS);
+        verify(mMediaSessionAdapter, never()).setActive(eq(true));
+
+        mHeadsetMediaButton.onCallEndpointChanged(
+                new CallEndpoint("Wired Headset", CallEndpoint.TYPE_WIRED_HEADSET));
+        waitForHandlerAction(mHeadsetMediaButton.getHandler(), TEST_TIMEOUT_MILLIS);
         verify(mMediaSessionAdapter).setActive(eq(true));
+
         // ... and thus we see how the original code isn't amenable to tests.
         when(mMediaSessionAdapter.isActive()).thenReturn(true);
 
-        when(mMockCallsManager.hasAnyCalls()).thenReturn(false);
+        // Remove the one call; we should release the session.
         mHeadsetMediaButton.onCallRemoved(regularCall);
         waitForHandlerAction(mHeadsetMediaButton.getHandler(), TEST_TIMEOUT_MILLIS);
         verify(mMediaSessionAdapter).setActive(eq(false));
+        when(mMediaSessionAdapter.isActive()).thenReturn(false);
+
+        // Add a new call; make sure we go active once more.
+        mHeadsetMediaButton.onCallAdded(regularCall);
+        mHeadsetMediaButton.onCallEndpointChanged(
+                new CallEndpoint("Wired Headset", CallEndpoint.TYPE_WIRED_HEADSET));
+        waitForHandlerAction(mHeadsetMediaButton.getHandler(), TEST_TIMEOUT_MILLIS);
+        verify(mMediaSessionAdapter, times(2)).setActive(eq(true));
+    }
+
+    /**
+     * Call is added and then routed to headset after call start
+     */
+    @SmallTest
+    @Test
+    public void testAddCallThenRouteToHeadsetAndBack() {
+        Call regularCall = getRegularCall();
+
+        when(mMockCallsManager.hasAnyCalls()).thenReturn(true);
+        mHeadsetMediaButton.onCallAdded(regularCall);
+        waitForHandlerAction(mHeadsetMediaButton.getHandler(), TEST_TIMEOUT_MILLIS);
+        verify(mMediaSessionAdapter, never()).setActive(eq(true));
+
+        mHeadsetMediaButton.onCallEndpointChanged(
+                new CallEndpoint("Wired Headset", CallEndpoint.TYPE_WIRED_HEADSET));
+        waitForHandlerAction(mHeadsetMediaButton.getHandler(), TEST_TIMEOUT_MILLIS);
+        verify(mMediaSessionAdapter).setActive(eq(true));
+
+        // ... and thus we see how the original code isn't amenable to tests.
+        when(mMediaSessionAdapter.isActive()).thenReturn(true);
+
+        mHeadsetMediaButton.onCallEndpointChanged(
+                new CallEndpoint("Earpiece", CallEndpoint.TYPE_EARPIECE));
+        waitForHandlerAction(mHeadsetMediaButton.getHandler(), TEST_TIMEOUT_MILLIS);
+        verify(mMediaSessionAdapter).setActive(eq(false));
+        when(mMediaSessionAdapter.isActive()).thenReturn(false);
+
+        // Remove the one call; we should not release again.
+        mHeadsetMediaButton.onCallRemoved(regularCall);
+        waitForHandlerAction(mHeadsetMediaButton.getHandler(), TEST_TIMEOUT_MILLIS);
+        // Remember, mockito counts total invocations; we should have went active once and then
+        // inactive again when we hit earpiece.
+        verify(mMediaSessionAdapter, times(1)).setActive(eq(true));
+        verify(mMediaSessionAdapter, times(1)).setActive(eq(false));
     }
 
     /**
@@ -111,6 +193,8 @@
         // Start with a regular old call.
         when(mMockCallsManager.hasAnyCalls()).thenReturn(true);
         mHeadsetMediaButton.onCallAdded(regularCall);
+        mHeadsetMediaButton.onCallEndpointChanged(
+                new CallEndpoint("Wired Headset", CallEndpoint.TYPE_WIRED_HEADSET));
         waitForHandlerAction(mHeadsetMediaButton.getHandler(), TEST_TIMEOUT_MILLIS);
         verify(mMediaSessionAdapter).setActive(eq(true));
         when(mMediaSessionAdapter.isActive()).thenReturn(true);
@@ -122,6 +206,7 @@
         // Expect to set session inactive.
         waitForHandlerAction(mHeadsetMediaButton.getHandler(), TEST_TIMEOUT_MILLIS);
         verify(mMediaSessionAdapter).setActive(eq(false));
+        when(mMediaSessionAdapter.isActive()).thenReturn(false);
 
         // For good measure lets make it non-external again.
         when(regularCall.isExternalCall()).thenReturn(false);
@@ -129,7 +214,7 @@
         mHeadsetMediaButton.onExternalCallChanged(regularCall, false);
         // Expect to set session active.
         waitForHandlerAction(mHeadsetMediaButton.getHandler(), TEST_TIMEOUT_MILLIS);
-        verify(mMediaSessionAdapter).setActive(eq(true));
+        verify(mMediaSessionAdapter, times(2)).setActive(eq(true));
     }
 
     @MediumTest
@@ -139,6 +224,8 @@
         when(externalCall.isExternalCall()).thenReturn(true);
 
         mHeadsetMediaButton.onCallAdded(externalCall);
+        mHeadsetMediaButton.onCallEndpointChanged(
+                new CallEndpoint("Wired Headset", CallEndpoint.TYPE_WIRED_HEADSET));
         waitForHandlerAction(mHeadsetMediaButton.getHandler(), TEST_TIMEOUT_MILLIS);
         verify(mMediaSessionAdapter, never()).setActive(eq(true));