Combine MediaSession Callback and TransportControlsCallback

This combines them into a single Callback class and adds default
handling to media buttons to check the available actions and call
one of the other methods if appropriate.

Change-Id: If9897d8cf6d8d8046aa85a646c22382f1db1461b
diff --git a/api/current.txt b/api/current.txt
index 5710415..680d315 100644
--- a/api/current.txt
+++ b/api/current.txt
@@ -16717,18 +16717,14 @@
 
   public final class MediaSession {
     ctor public MediaSession(android.content.Context, java.lang.String);
-    method public void addCallback(android.media.session.MediaSession.Callback);
-    method public void addCallback(android.media.session.MediaSession.Callback, android.os.Handler);
-    method public void addTransportControlsCallback(android.media.session.MediaSession.TransportControlsCallback);
-    method public void addTransportControlsCallback(android.media.session.MediaSession.TransportControlsCallback, android.os.Handler);
     method public android.media.session.MediaController getController();
     method public android.media.session.MediaSession.Token getSessionToken();
     method public boolean isActive();
     method public void release();
-    method public void removeCallback(android.media.session.MediaSession.Callback);
-    method public void removeTransportControlsCallback(android.media.session.MediaSession.TransportControlsCallback);
     method public void sendSessionEvent(java.lang.String, android.os.Bundle);
     method public void setActive(boolean);
+    method public void setCallback(android.media.session.MediaSession.Callback);
+    method public void setCallback(android.media.session.MediaSession.Callback, android.os.Handler);
     method public void setExtras(android.os.Bundle);
     method public void setFlags(int);
     method public void setLaunchActivity(android.app.PendingIntent);
@@ -16749,7 +16745,20 @@
   public static abstract class MediaSession.Callback {
     ctor public MediaSession.Callback();
     method public void onCommand(java.lang.String, android.os.Bundle, android.os.ResultReceiver);
-    method public void onMediaButtonEvent(android.content.Intent);
+    method public void onCustomAction(java.lang.String, android.os.Bundle);
+    method public void onFastForward();
+    method public boolean onMediaButtonEvent(android.content.Intent);
+    method public void onPause();
+    method public void onPlay();
+    method public void onPlayFromSearch(java.lang.String, android.os.Bundle);
+    method public void onPlayUri(android.net.Uri, android.os.Bundle);
+    method public void onRewind();
+    method public void onSeekTo(long);
+    method public void onSetRating(android.media.Rating);
+    method public void onSkipToNext();
+    method public void onSkipToPrevious();
+    method public void onSkipToTrack(long);
+    method public void onStop();
   }
 
   public static final class MediaSession.Token implements android.os.Parcelable {
@@ -16775,23 +16784,6 @@
     method public android.media.session.MediaSession.Track.Builder setExtras(android.os.Bundle);
   }
 
-  public static abstract class MediaSession.TransportControlsCallback {
-    ctor public MediaSession.TransportControlsCallback();
-    method public void onCustomAction(java.lang.String, android.os.Bundle);
-    method public void onFastForward();
-    method public void onPause();
-    method public void onPlay();
-    method public void onPlayFromSearch(java.lang.String, android.os.Bundle);
-    method public void onPlayUri(android.net.Uri, android.os.Bundle);
-    method public void onRewind();
-    method public void onSeekTo(long);
-    method public void onSetRating(android.media.Rating);
-    method public void onSkipToNext();
-    method public void onSkipToPrevious();
-    method public void onSkipToTrack(long);
-    method public void onStop();
-  }
-
   public final class MediaSessionManager {
     method public void addActiveSessionsListener(android.media.session.MediaSessionManager.SessionListener, android.content.ComponentName);
     method public java.util.List<android.media.session.MediaController> getActiveSessions(android.content.ComponentName);
diff --git a/media/java/android/media/RemoteControlClient.java b/media/java/android/media/RemoteControlClient.java
index 96c66c5..2a0fd83 100644
--- a/media/java/android/media/RemoteControlClient.java
+++ b/media/java/android/media/RemoteControlClient.java
@@ -970,8 +970,7 @@
     public final static int RCSE_ID_UNREGISTERED = -1;
 
     // USE_SESSIONS
-    private MediaSession.TransportControlsCallback mTransportListener
-            = new MediaSession.TransportControlsCallback() {
+    private MediaSession.Callback mTransportListener = new MediaSession.Callback() {
 
         @Override
         public void onSeekTo(long pos) {
diff --git a/media/java/android/media/session/MediaController.java b/media/java/android/media/session/MediaController.java
index e3c198e..f6e189a 100644
--- a/media/java/android/media/session/MediaController.java
+++ b/media/java/android/media/session/MediaController.java
@@ -179,7 +179,8 @@
     }
 
     /**
-     * Get the current play queue for this session.
+     * Get the current play queue for this session if one is set. If you only
+     * care about the current item {@link #getMetadata()} should be used.
      *
      * @return The current play queue or null.
      */
diff --git a/media/java/android/media/session/MediaSession.java b/media/java/android/media/session/MediaSession.java
index cf8e3dd..cf73c2a 100644
--- a/media/java/android/media/session/MediaSession.java
+++ b/media/java/android/media/session/MediaSession.java
@@ -21,12 +21,10 @@
 import android.annotation.Nullable;
 import android.app.Activity;
 import android.app.PendingIntent;
-import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
 import android.content.pm.ParceledListSlice;
 import android.media.AudioAttributes;
-import android.media.AudioManager;
 import android.media.MediaMetadata;
 import android.media.Rating;
 import android.media.VolumeProvider;
@@ -43,11 +41,11 @@
 import android.os.UserHandle;
 import android.text.TextUtils;
 import android.util.Log;
+import android.view.KeyEvent;
 
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.lang.ref.WeakReference;
-import java.util.ArrayList;
 import java.util.List;
 
 /**
@@ -64,10 +62,8 @@
  * create a {@link MediaController} to interact with the session.
  * <p>
  * To receive commands, media keys, and other events a {@link Callback} must be
- * set with {@link #addCallback(Callback)} and {@link #setActive(boolean)
- * setActive(true)} must be called. To receive transport control commands a
- * {@link TransportControlsCallback} must be set with
- * {@link #addTransportControlsCallback}.
+ * set with {@link #setCallback(Callback)} and {@link #setActive(boolean)
+ * setActive(true)} must be called.
  * <p>
  * When an app is finished performing playback it must call {@link #release()}
  * to clean up the session and notify any controllers.
@@ -85,8 +81,7 @@
 
     /**
      * Set this flag on the session to indicate that it handles transport
-     * control commands through a {@link TransportControlsCallback}.
-     * The callback can be retrieved by calling {@link #addTransportControlsCallback}.
+     * control commands through its {@link Callback}.
      */
     public static final int FLAG_HANDLES_TRANSPORT_CONTROLS = 1 << 1;
 
@@ -124,12 +119,9 @@
     private final ISession mBinder;
     private final CallbackStub mCbStub;
 
-    private final ArrayList<CallbackMessageHandler> mCallbacks
-            = new ArrayList<CallbackMessageHandler>();
-    private final ArrayList<TransportMessageHandler> mTransportCallbacks
-            = new ArrayList<TransportMessageHandler>();
-
+    private CallbackMessageHandler mCallback;
     private VolumeProvider mVolumeProvider;
+    private PlaybackState mPlaybackState;
 
     private boolean mActive = false;
 
@@ -177,30 +169,35 @@
     }
 
     /**
-     * Add a callback to receive updates on for the MediaSession. This includes
-     * media button and volume events. The caller's thread will be used to post
-     * events.
+     * Set the callback to receive updates for the MediaSession. This includes
+     * media button events and transport controls. The caller's thread will be
+     * used to post updates.
+     * <p>
+     * Set the callback to null to stop receiving updates.
      *
      * @param callback The callback object
      */
-    public void addCallback(@NonNull Callback callback) {
-        addCallback(callback, null);
+    public void setCallback(@Nullable Callback callback) {
+        setCallback(callback, null);
     }
 
     /**
-     * Add a callback to receive updates for the MediaSession. This includes
-     * media button and volume events.
+     * Set the callback to receive updates for the MediaSession. This includes
+     * media button events and transport controls.
+     * <p>
+     * Set the callback to null to stop receiving updates.
      *
      * @param callback The callback to receive updates on.
      * @param handler The handler that events should be posted on.
      */
-    public void addCallback(@NonNull Callback callback, @Nullable Handler handler) {
+    public void setCallback(@Nullable Callback callback, @Nullable Handler handler) {
         if (callback == null) {
-            throw new IllegalArgumentException("Callback cannot be null");
+            mCallback = null;
+            return;
         }
         synchronized (mLock) {
-            if (getHandlerForCallbackLocked(callback) != null) {
-                Log.w(TAG, "Callback is already added, ignoring");
+            if (mCallback != null && mCallback.mCallback == callback) {
+                Log.w(TAG, "Tried to set same callback, ignoring");
                 return;
             }
             if (handler == null) {
@@ -208,18 +205,7 @@
             }
             CallbackMessageHandler msgHandler = new CallbackMessageHandler(handler.getLooper(),
                     callback);
-            mCallbacks.add(msgHandler);
-        }
-    }
-
-    /**
-     * Remove a callback. It will no longer receive updates.
-     *
-     * @param callback The callback to remove.
-     */
-    public void removeCallback(@NonNull Callback callback) {
-        synchronized (mLock) {
-            removeCallbackLocked(callback);
+            mCallback = msgHandler;
         }
     }
 
@@ -421,63 +407,12 @@
     }
 
     /**
-     * Add a callback to receive transport controls on, such as play, rewind, or
-     * fast forward.
-     *
-     * @param callback The callback object
-     */
-    public void addTransportControlsCallback(@NonNull TransportControlsCallback callback) {
-        addTransportControlsCallback(callback, null);
-    }
-
-    /**
-     * Add a callback to receive transport controls on, such as play, rewind, or
-     * fast forward. The updates will be posted to the specified handler. If no
-     * handler is provided they will be posted to the caller's thread.
-     *
-     * @param callback The callback to receive updates on
-     * @param handler The handler to post the updates on
-     */
-    public void addTransportControlsCallback(@NonNull TransportControlsCallback callback,
-            @Nullable Handler handler) {
-        if (callback == null) {
-            throw new IllegalArgumentException("Callback cannot be null");
-        }
-        synchronized (mLock) {
-            if (getTransportControlsHandlerForCallbackLocked(callback) != null) {
-                Log.w(TAG, "Callback is already added, ignoring");
-                return;
-            }
-            if (handler == null) {
-                handler = new Handler();
-            }
-            TransportMessageHandler msgHandler = new TransportMessageHandler(handler.getLooper(),
-                    callback);
-            mTransportCallbacks.add(msgHandler);
-        }
-    }
-
-    /**
-     * Stop receiving transport controls on the specified callback. If an update
-     * has already been posted you may still receive it after this call returns.
-     *
-     * @param callback The callback to stop receiving updates on
-     */
-    public void removeTransportControlsCallback(@NonNull TransportControlsCallback callback) {
-        if (callback == null) {
-            throw new IllegalArgumentException("Callback cannot be null");
-        }
-        synchronized (mLock) {
-            removeTransportControlsCallbackLocked(callback);
-        }
-    }
-
-    /**
      * Update the current playback state.
      *
      * @param state The current state of playback
      */
     public void setPlaybackState(@Nullable PlaybackState state) {
+        mPlaybackState = state;
         try {
             mBinder.setPlaybackState(state);
         } catch (RemoteException e) {
@@ -566,138 +501,78 @@
     }
 
     private void dispatchPlay() {
-        postToTransportCallbacks(TransportMessageHandler.MSG_PLAY);
+        postToCallback(CallbackMessageHandler.MSG_PLAY);
     }
 
     private void dispatchPlayUri(Uri uri, Bundle extras) {
-        postToTransportCallbacks(TransportMessageHandler.MSG_PLAY_URI, uri, extras);
+        postToCallback(CallbackMessageHandler.MSG_PLAY_URI, uri, extras);
     }
 
     private void dispatchPlayFromSearch(String query, Bundle extras) {
-        postToTransportCallbacks(TransportMessageHandler.MSG_PLAY_SEARCH, query, extras);
+        postToCallback(CallbackMessageHandler.MSG_PLAY_SEARCH, query, extras);
     }
 
     private void dispatchSkipToTrack(long id) {
-        postToTransportCallbacks(TransportMessageHandler.MSG_SKIP_TO_TRACK, id);
+        postToCallback(CallbackMessageHandler.MSG_SKIP_TO_TRACK, id);
     }
 
     private void dispatchPause() {
-        postToTransportCallbacks(TransportMessageHandler.MSG_PAUSE);
+        postToCallback(CallbackMessageHandler.MSG_PAUSE);
     }
 
     private void dispatchStop() {
-        postToTransportCallbacks(TransportMessageHandler.MSG_STOP);
+        postToCallback(CallbackMessageHandler.MSG_STOP);
     }
 
     private void dispatchNext() {
-        postToTransportCallbacks(TransportMessageHandler.MSG_NEXT);
+        postToCallback(CallbackMessageHandler.MSG_NEXT);
     }
 
     private void dispatchPrevious() {
-        postToTransportCallbacks(TransportMessageHandler.MSG_PREVIOUS);
+        postToCallback(CallbackMessageHandler.MSG_PREVIOUS);
     }
 
     private void dispatchFastForward() {
-        postToTransportCallbacks(TransportMessageHandler.MSG_FAST_FORWARD);
+        postToCallback(CallbackMessageHandler.MSG_FAST_FORWARD);
     }
 
     private void dispatchRewind() {
-        postToTransportCallbacks(TransportMessageHandler.MSG_REWIND);
+        postToCallback(CallbackMessageHandler.MSG_REWIND);
     }
 
     private void dispatchSeekTo(long pos) {
-        postToTransportCallbacks(TransportMessageHandler.MSG_SEEK_TO, pos);
+        postToCallback(CallbackMessageHandler.MSG_SEEK_TO, pos);
     }
 
     private void dispatchRate(Rating rating) {
-        postToTransportCallbacks(TransportMessageHandler.MSG_RATE, rating);
+        postToCallback(CallbackMessageHandler.MSG_RATE, rating);
     }
 
     private void dispatchCustomAction(String action, Bundle args) {
-        postToTransportCallbacks(TransportMessageHandler.MSG_CUSTOM_ACTION, action, args);
+        postToCallback(CallbackMessageHandler.MSG_CUSTOM_ACTION, action, args);
     }
 
-    private TransportMessageHandler getTransportControlsHandlerForCallbackLocked(
-            TransportControlsCallback callback) {
-        for (int i = mTransportCallbacks.size() - 1; i >= 0; i--) {
-            TransportMessageHandler handler = mTransportCallbacks.get(i);
-            if (callback == handler.mCallback) {
-                return handler;
-            }
-        }
-        return null;
+    private void dispatchMediaButton(Intent mediaButtonIntent) {
+        postToCallback(CallbackMessageHandler.MSG_MEDIA_BUTTON, mediaButtonIntent);
     }
 
-    private boolean removeTransportControlsCallbackLocked(TransportControlsCallback callback) {
-        for (int i = mTransportCallbacks.size() - 1; i >= 0; i--) {
-            if (callback == mTransportCallbacks.get(i).mCallback) {
-                mTransportCallbacks.remove(i);
-                return true;
-            }
-        }
-        return false;
-    }
-
-    private void postToTransportCallbacks(int what, Object obj) {
-        synchronized (mLock) {
-            for (int i = mTransportCallbacks.size() - 1; i >= 0; i--) {
-                mTransportCallbacks.get(i).post(what, obj);
-            }
-        }
-    }
-
-    private void postToTransportCallbacks(int what, Object obj, Bundle args) {
-        synchronized (mLock) {
-            for (int i = mTransportCallbacks.size() - 1; i >= 0; i--) {
-                mTransportCallbacks.get(i).post(what, obj, args);
-            }
-        }
-    }
-
-    private void postToTransportCallbacks(int what) {
-        postToTransportCallbacks(what, null);
-    }
-
-    private CallbackMessageHandler getHandlerForCallbackLocked(Callback cb) {
-        if (cb == null) {
-            throw new IllegalArgumentException("Callback cannot be null");
-        }
-        for (int i = mCallbacks.size() - 1; i >= 0; i--) {
-            CallbackMessageHandler handler = mCallbacks.get(i);
-            if (cb == handler.mCallback) {
-                return handler;
-            }
-        }
-        return null;
-    }
-
-    private boolean removeCallbackLocked(Callback cb) {
-        if (cb == null) {
-            throw new IllegalArgumentException("Callback cannot be null");
-        }
-        for (int i = mCallbacks.size() - 1; i >= 0; i--) {
-            CallbackMessageHandler handler = mCallbacks.get(i);
-            if (cb == handler.mCallback) {
-                mCallbacks.remove(i);
-                return true;
-            }
-        }
-        return false;
+    private void postToCallback(int what) {
+        postToCallback(what, null);
     }
 
     private void postCommand(String command, Bundle args, ResultReceiver resultCb) {
         Command cmd = new Command(command, args, resultCb);
-        synchronized (mLock) {
-            for (int i = mCallbacks.size() - 1; i >= 0; i--) {
-                mCallbacks.get(i).post(CallbackMessageHandler.MSG_COMMAND, cmd);
-            }
-        }
+        postToCallback(CallbackMessageHandler.MSG_COMMAND, cmd);
     }
 
-    private void postMediaButton(Intent mediaButtonIntent) {
+    private void postToCallback(int what, Object obj) {
+        postToCallback(what, obj, null);
+    }
+
+    private void postToCallback(int what, Object obj, Bundle extras) {
         synchronized (mLock) {
-            for (int i = mCallbacks.size() - 1; i >= 0; i--) {
-                mCallbacks.get(i).post(CallbackMessageHandler.MSG_MEDIA_BUTTON, mediaButtonIntent);
+            if (mCallback != null) {
+                mCallback.post(what, obj, extras);
             }
         }
     }
@@ -791,30 +666,16 @@
     }
 
     /**
-     * Receives generic commands or updates from controllers and the system.
-     * Callbacks may be registered using {@link #addCallback}.
+     * Receives media buttons, transport controls, and commands from controllers
+     * and the system. A callback may be set using {@link #setCallback}.
      */
     public abstract static class Callback {
+        private MediaSession mSession;
 
         public Callback() {
         }
 
         /**
-         * Called when a media button is pressed and this session has the
-         * highest priority or a controller sends a media button event to the
-         * session. TODO determine if using Intents identical to the ones
-         * RemoteControlClient receives is useful
-         * <p>
-         * The intent will be of type {@link Intent#ACTION_MEDIA_BUTTON} with a
-         * KeyEvent in {@link Intent#EXTRA_KEY_EVENT}
-         *
-         * @param mediaButtonIntent an intent containing the KeyEvent as an
-         *            extra
-         */
-        public void onMediaButtonEvent(@NonNull Intent mediaButtonIntent) {
-        }
-
-        /**
          * Called when a controller has sent a command to this session.
          * The owner of the session may handle custom commands but is not
          * required to.
@@ -826,13 +687,81 @@
         public void onCommand(@NonNull String command, @Nullable Bundle args,
                 @Nullable ResultReceiver cb) {
         }
-    }
 
-    /**
-     * Receives transport control commands. Callbacks may be registered using
-     * {@link #addTransportControlsCallback}.
-     */
-    public static abstract class TransportControlsCallback {
+        /**
+         * Called when a media button is pressed and this session has the
+         * highest priority or a controller sends a media button event to the
+         * session. The default behavior will call the relevant method if the
+         * action for it was set.
+         * <p>
+         * The intent will be of type {@link Intent#ACTION_MEDIA_BUTTON} with a
+         * KeyEvent in {@link Intent#EXTRA_KEY_EVENT}
+         *
+         * @param mediaButtonIntent an intent containing the KeyEvent as an
+         *            extra
+         */
+        public boolean onMediaButtonEvent(@NonNull Intent mediaButtonIntent) {
+            if (mSession != null
+                    && Intent.ACTION_MEDIA_BUTTON.equals(mediaButtonIntent.getAction())) {
+                KeyEvent ke = mediaButtonIntent.getParcelableExtra(Intent.EXTRA_KEY_EVENT);
+                if (ke != null && ke.getAction() == KeyEvent.ACTION_DOWN) {
+                    PlaybackState state = mSession.mPlaybackState;
+                    long validActions = state == null ? 0 : state.getActions();
+                    switch (ke.getKeyCode()) {
+                        case KeyEvent.KEYCODE_MEDIA_PLAY:
+                            if ((validActions & PlaybackState.ACTION_PLAY) != 0) {
+                                onPlay();
+                            }
+                            break;
+                        case KeyEvent.KEYCODE_MEDIA_PAUSE:
+                            if ((validActions & PlaybackState.ACTION_PAUSE) != 0) {
+                                onPause();
+                            }
+                            break;
+                        case KeyEvent.KEYCODE_MEDIA_NEXT:
+                            if ((validActions & PlaybackState.ACTION_SKIP_TO_NEXT) != 0) {
+                                onSkipToNext();
+                            }
+                            break;
+                        case KeyEvent.KEYCODE_MEDIA_PREVIOUS:
+                            if ((validActions & PlaybackState.ACTION_SKIP_TO_PREVIOUS) != 0) {
+                                onSkipToPrevious();
+                            }
+                            break;
+                        case KeyEvent.KEYCODE_MEDIA_STOP:
+                            if ((validActions & PlaybackState.ACTION_STOP) != 0) {
+                                onStop();
+                            }
+                            break;
+                        case KeyEvent.KEYCODE_MEDIA_FAST_FORWARD:
+                            if ((validActions & PlaybackState.ACTION_FAST_FORWARD) != 0) {
+                                onFastForward();
+                            }
+                            break;
+                        case KeyEvent.KEYCODE_MEDIA_REWIND:
+                            if ((validActions & PlaybackState.ACTION_REWIND) != 0) {
+                                onRewind();
+                            }
+                            break;
+                        case KeyEvent.KEYCODE_MEDIA_PLAY_PAUSE:
+                        case KeyEvent.KEYCODE_HEADSETHOOK:
+                            boolean isPlaying = state == null ? false
+                                    : state.getState() == PlaybackState.STATE_PLAYING;
+                            boolean canPlay = (validActions & (PlaybackState.ACTION_PLAY_PAUSE
+                                    | PlaybackState.ACTION_PLAY)) != 0;
+                            boolean canPause = (validActions & (PlaybackState.ACTION_PLAY_PAUSE
+                                    | PlaybackState.ACTION_PAUSE)) != 0;
+                            if (isPlaying && canPause) {
+                                onPause();
+                            } else if (!isPlaying && canPlay) {
+                                onPlay();
+                            }
+                            break;
+                    }
+                }
+            }
+            return false;
+        }
 
         /**
          * Override to handle requests to begin playback.
@@ -920,6 +849,10 @@
          */
         public void onCustomAction(@NonNull String action, @Nullable Bundle extras) {
         }
+
+        private void setSession(MediaSession session) {
+            mSession = session;
+        }
     }
 
     /**
@@ -946,7 +879,7 @@
             MediaSession session = mMediaSession.get();
             try {
                 if (session != null) {
-                    session.postMediaButton(mediaButtonIntent);
+                    session.dispatchMediaButton(mediaButtonIntent);
                 }
             } finally {
                 if (cb != null) {
@@ -1232,44 +1165,6 @@
         }
     }
 
-    private class CallbackMessageHandler extends Handler {
-        private static final int MSG_MEDIA_BUTTON = 1;
-        private static final int MSG_COMMAND = 2;
-
-        private MediaSession.Callback mCallback;
-
-        public CallbackMessageHandler(Looper looper, MediaSession.Callback callback) {
-            super(looper, null, true);
-            mCallback = callback;
-        }
-
-        @Override
-        public void handleMessage(Message msg) {
-            synchronized (mLock) {
-                if (mCallback == null) {
-                    return;
-                }
-                switch (msg.what) {
-                    case MSG_MEDIA_BUTTON:
-                        mCallback.onMediaButtonEvent((Intent) msg.obj);
-                        break;
-                    case MSG_COMMAND:
-                        Command cmd = (Command) msg.obj;
-                        mCallback.onCommand(cmd.command, cmd.extras, cmd.stub);
-                        break;
-                }
-            }
-        }
-
-        public void post(int what, Object obj) {
-            obtainMessage(what, obj).sendToTarget();
-        }
-
-        public void post(int what, Object obj, int arg1) {
-            obtainMessage(what, arg1, 0, obj).sendToTarget();
-        }
-    }
-
     private static final class Command {
         public final String command;
         public final Bundle extras;
@@ -1282,7 +1177,8 @@
         }
     }
 
-    private class TransportMessageHandler extends Handler {
+    private class CallbackMessageHandler extends Handler {
+
         private static final int MSG_PLAY = 1;
         private static final int MSG_PLAY_URI = 2;
         private static final int MSG_PLAY_SEARCH = 3;
@@ -1296,12 +1192,14 @@
         private static final int MSG_SEEK_TO = 11;
         private static final int MSG_RATE = 12;
         private static final int MSG_CUSTOM_ACTION = 13;
+        private static final int MSG_MEDIA_BUTTON = 14;
+        private static final int MSG_COMMAND = 15;
 
-        private TransportControlsCallback mCallback;
+        private MediaSession.Callback mCallback;
 
-        public TransportMessageHandler(Looper looper, TransportControlsCallback cb) {
-            super(looper);
-            mCallback = cb;
+        public CallbackMessageHandler(Looper looper, MediaSession.Callback callback) {
+            super(looper, null, true);
+            mCallback = callback;
         }
 
         public void post(int what, Object obj, Bundle bundle) {
@@ -1318,6 +1216,10 @@
             post(what, null);
         }
 
+        public void post(int what, Object obj, int arg1) {
+            obtainMessage(what, arg1, 0, obj).sendToTarget();
+        }
+
         @Override
         public void handleMessage(Message msg) {
             switch (msg.what) {
@@ -1359,6 +1261,13 @@
                 case MSG_CUSTOM_ACTION:
                     mCallback.onCustomAction((String) msg.obj, msg.getData());
                     break;
+                case MSG_MEDIA_BUTTON:
+                    mCallback.onMediaButtonEvent((Intent) msg.obj);
+                    break;
+                case MSG_COMMAND:
+                    Command cmd = (Command) msg.obj;
+                    mCallback.onCommand(cmd.command, cmd.extras, cmd.stub);
+                    break;
             }
         }
     }
diff --git a/media/java/android/media/session/MediaSessionLegacyHelper.java b/media/java/android/media/session/MediaSessionLegacyHelper.java
index f075ded..a182982 100644
--- a/media/java/android/media/session/MediaSessionLegacyHelper.java
+++ b/media/java/android/media/session/MediaSessionLegacyHelper.java
@@ -29,6 +29,9 @@
 import android.media.MediaMetadata;
 import android.media.MediaMetadataEditor;
 import android.media.MediaMetadataRetriever;
+import android.media.Rating;
+import android.media.RemoteControlClient;
+import android.media.RemoteControlClient.MetadataEditor;
 import android.os.Bundle;
 import android.os.Handler;
 import android.os.Looper;
@@ -229,8 +232,7 @@
         }
     }
 
-    public void addRccListener(PendingIntent pi,
-            MediaSession.TransportControlsCallback listener) {
+    public void addRccListener(PendingIntent pi, MediaSession.Callback listener) {
         if (pi == null) {
             Log.w(TAG, "Pending intent was null, can't add rcc listener.");
             return;
@@ -247,10 +249,7 @@
                 // This is already the registered listener, ignore
                 return;
             }
-            // Otherwise it changed so we need to switch to the new one
-            holder.mSession.removeTransportControlsCallback(holder.mRccListener);
         }
-        holder.mSession.addTransportControlsCallback(listener, mHandler);
         holder.mRccListener = listener;
         holder.mFlags |= MediaSession.FLAG_HANDLES_TRANSPORT_CONTROLS;
         holder.mSession.setFlags(holder.mFlags);
@@ -266,7 +265,6 @@
         }
         SessionHolder holder = getHolder(pi, false);
         if (holder != null && holder.mRccListener != null) {
-            holder.mSession.removeTransportControlsCallback(holder.mRccListener);
             holder.mRccListener = null;
             holder.mFlags &= ~MediaSession.FLAG_HANDLES_TRANSPORT_CONTROLS;
             holder.mSession.setFlags(holder.mFlags);
@@ -288,8 +286,7 @@
             return;
         }
         if (holder.mMediaButtonListener != null) {
-            // Already have this listener registered, but update it anyway as
-            // the extras may have changed.
+            // Already have this listener registered
             if (DEBUG) {
                 Log.d(TAG, "addMediaButtonListener already added " + pi);
             }
@@ -300,11 +297,8 @@
         // set this flag
         holder.mFlags |= MediaSession.FLAG_HANDLES_MEDIA_BUTTONS;
         holder.mSession.setFlags(holder.mFlags);
-        holder.mSession.addTransportControlsCallback(holder.mMediaButtonListener, mHandler);
-
-        holder.mMediaButtonReceiver = new MediaButtonReceiver(pi, context);
-        holder.mSession.addCallback(holder.mMediaButtonReceiver, mHandler);
         holder.mSession.setMediaButtonReceiver(pi);
+        holder.update();
         if (DEBUG) {
             Log.d(TAG, "addMediaButtonListener added " + pi);
         }
@@ -316,13 +310,10 @@
         }
         SessionHolder holder = getHolder(pi, false);
         if (holder != null && holder.mMediaButtonListener != null) {
-            holder.mSession.removeTransportControlsCallback(holder.mMediaButtonListener);
             holder.mFlags &= ~MediaSession.FLAG_HANDLES_MEDIA_BUTTONS;
             holder.mSession.setFlags(holder.mFlags);
             holder.mMediaButtonListener = null;
 
-            holder.mSession.removeCallback(holder.mMediaButtonReceiver);
-            holder.mMediaButtonReceiver = null;
             holder.update();
             if (DEBUG) {
                 Log.d(TAG, "removeMediaButtonListener removed " + pi);
@@ -387,22 +378,7 @@
         }
     }
 
-    private static final class MediaButtonReceiver extends MediaSession.Callback {
-        private final PendingIntent mPendingIntent;
-        private final Context mContext;
-
-        public MediaButtonReceiver(PendingIntent pi, Context context) {
-            mPendingIntent = pi;
-            mContext = context;
-        }
-
-        @Override
-        public void onMediaButtonEvent(Intent mediaButtonIntent) {
-            MediaSessionLegacyHelper.sendKeyEvent(mPendingIntent, mContext, mediaButtonIntent);
-        }
-    }
-
-    private static final class MediaButtonListener extends MediaSession.TransportControlsCallback {
+    private static final class MediaButtonListener extends MediaSession.Callback {
         private final PendingIntent mPendingIntent;
         private final Context mContext;
 
@@ -412,6 +388,12 @@
         }
 
         @Override
+        public boolean onMediaButtonEvent(Intent mediaButtonIntent) {
+            MediaSessionLegacyHelper.sendKeyEvent(mPendingIntent, mContext, mediaButtonIntent);
+            return true;
+        }
+
+        @Override
         public void onPlay() {
             sendKeyEvent(KeyEvent.KEYCODE_MEDIA_PLAY);
         }
@@ -468,10 +450,11 @@
         public final MediaSession mSession;
         public final PendingIntent mPi;
         public MediaButtonListener mMediaButtonListener;
-        public MediaButtonReceiver mMediaButtonReceiver;
-        public MediaSession.TransportControlsCallback mRccListener;
+        public MediaSession.Callback mRccListener;
         public int mFlags;
 
+        public SessionCallback mCb;
+
         public SessionHolder(MediaSession session, PendingIntent pi) {
             mSession = session;
             mPi = pi;
@@ -479,8 +462,87 @@
 
         public void update() {
             if (mMediaButtonListener == null && mRccListener == null) {
+                mSession.setCallback(null);
                 mSession.release();
+                mCb = null;
                 mSessions.remove(mPi);
+            } else if (mCb == null) {
+                mCb = new SessionCallback();
+                mSession.setCallback(mCb);
+            }
+        }
+
+        private class SessionCallback extends MediaSession.Callback {
+
+            @Override
+            public boolean onMediaButtonEvent(Intent mediaButtonIntent) {
+                if (mMediaButtonListener != null) {
+                    mMediaButtonListener.onMediaButtonEvent(mediaButtonIntent);
+                }
+                return true;
+            }
+
+            @Override
+            public void onPlay() {
+                if (mMediaButtonListener != null) {
+                    mMediaButtonListener.onPlay();
+                }
+            }
+
+            @Override
+            public void onPause() {
+                if (mMediaButtonListener != null) {
+                    mMediaButtonListener.onPause();
+                }
+            }
+
+            @Override
+            public void onSkipToNext() {
+                if (mMediaButtonListener != null) {
+                    mMediaButtonListener.onSkipToNext();
+                }
+            }
+
+            @Override
+            public void onSkipToPrevious() {
+                if (mMediaButtonListener != null) {
+                    mMediaButtonListener.onSkipToPrevious();
+                }
+            }
+
+            @Override
+            public void onFastForward() {
+                if (mMediaButtonListener != null) {
+                    mMediaButtonListener.onFastForward();
+                }
+            }
+
+            @Override
+            public void onRewind() {
+                if (mMediaButtonListener != null) {
+                    mMediaButtonListener.onRewind();
+                }
+            }
+
+            @Override
+            public void onStop() {
+                if (mMediaButtonListener != null) {
+                    mMediaButtonListener.onStop();
+                }
+            }
+
+            @Override
+            public void onSeekTo(long pos) {
+                if (mRccListener != null) {
+                    mRccListener.onSeekTo(pos);
+                }
+            }
+
+            @Override
+            public void onSetRating(Rating rating) {
+                if (mRccListener != null) {
+                    mRccListener.onSetRating(rating);
+                }
             }
         }
     }
diff --git a/media/java/android/media/session/PlaybackState.java b/media/java/android/media/session/PlaybackState.java
index 65bd677..2ad8eae 100644
--- a/media/java/android/media/session/PlaybackState.java
+++ b/media/java/android/media/session/PlaybackState.java
@@ -845,20 +845,23 @@
         }
 
         /**
-         * Add a custom action to the playback state. Actions can be used to expose additional
-         * functionality to {@link MediaController MediaControllers} beyond what is offered by the
-         * standard transport controls.
+         * Add a custom action to the playback state. Actions can be used to
+         * expose additional functionality to {@link MediaController
+         * MediaControllers} beyond what is offered by the standard transport
+         * controls.
          * <p>
-         * e.g. start a radio station based on the current item or skip ahead by 30 seconds.
+         * e.g. start a radio station based on the current item or skip ahead by
+         * 30 seconds.
          *
-         * @param action An identifier for this action. It will be sent back to the
-         *               {@link MediaSession} through
-         *               {@link
-         *               MediaSession.TransportControlsCallback#onCustomAction(String, Bundle)}.
-         * @param name The display name for the action. If text is shown with the action or used
-         *             for accessibility, this is what should be used.
-         * @param icon The resource action of the icon that should be displayed for the action. The
-         *             resource should be in the package of the {@link MediaSession}.
+         * @param action An identifier for this action. It can be sent back to
+         *            the {@link MediaSession} through
+         *            {@link MediaController.TransportControls#sendCustomAction(String, Bundle)}.
+         * @param name The display name for the action. If text is shown with
+         *            the action or used for accessibility, this is what should
+         *            be used.
+         * @param icon The resource action of the icon that should be displayed
+         *            for the action. The resource should be in the package of
+         *            the {@link MediaSession}.
          * @return this
          */
         public Builder addCustomAction(String action, String name, int icon) {
diff --git a/tests/OneMedia/src/com/android/onemedia/PlayerSession.java b/tests/OneMedia/src/com/android/onemedia/PlayerSession.java
index feecfde..890d68d 100644
--- a/tests/OneMedia/src/com/android/onemedia/PlayerSession.java
+++ b/tests/OneMedia/src/com/android/onemedia/PlayerSession.java
@@ -86,10 +86,10 @@
         mRouter.setRoutingCallback(new RoutingCallback(), null);
 
         mSession = new MediaSession(mContext, "OneMedia");
-        mSession.addCallback(mCallback);
-        mSession.addTransportControlsCallback(new TransportCallback());
+        mSession.setCallback(mCallback);
         mSession.setPlaybackState(mPlaybackState);
-        mSession.setFlags(MediaSession.FLAG_HANDLES_TRANSPORT_CONTROLS);
+        mSession.setFlags(MediaSession.FLAG_HANDLES_TRANSPORT_CONTROLS
+                | MediaSession.FLAG_HANDLES_MEDIA_BUTTONS);
         mSession.setMediaRouter(mRouter);
         mSession.setActive(true);
     }
@@ -230,26 +230,6 @@
 
     private class SessionCb extends MediaSession.Callback {
         @Override
-        public void onMediaButtonEvent(Intent mediaRequestIntent) {
-            if (Intent.ACTION_MEDIA_BUTTON.equals(mediaRequestIntent.getAction())) {
-                KeyEvent event = (KeyEvent) mediaRequestIntent
-                        .getParcelableExtra(Intent.EXTRA_KEY_EVENT);
-                switch (event.getKeyCode()) {
-                    case KeyEvent.KEYCODE_MEDIA_PLAY:
-                        Log.d(TAG, "play button received");
-                        mRenderer.onPlay();
-                        break;
-                    case KeyEvent.KEYCODE_MEDIA_PAUSE:
-                        Log.d(TAG, "pause button received");
-                        mRenderer.onPause();
-                        break;
-                }
-            }
-        }
-    }
-
-    private class TransportCallback extends MediaSession.TransportControlsCallback {
-        @Override
         public void onPlay() {
             mRenderer.onPlay();
         }
@@ -315,7 +295,7 @@
                         updateState(PlaybackState.STATE_NONE);
                         break;
                 }
-            } 
+            }
         }
     }
 }