MediaStyle notification is not responding to Remot

We are in the process of migrating our ongoing playback notification to MediaStyle notifications introduced in Lollipop. RemoteControlClient seems to be deprecated, and the MediaStyle notification is not handling the media button events (such as pause/play through headphones remote).

Did anyone get this work? None of the events in MediaSessionCallback are called.

Here is how the media session is initialized

    mSession = new MediaSessionCompat(this, TAG);
    mSession.setCallback(new MediaSessionCallback());
    mSession.setFlags(MediaSessionCompat.FLAG_HANDLES_MEDIA_BUTTONS | MediaSessionCompat.FLAG_HANDLES_TRANSPORT_CONTROLS);
    mSession.setPlaybackToLocal(AudioManager.STREAM_MUSIC);
    mSession.setActive(true);

Here is how metadata is is set

    MediaMetadataCompat.Builder metadataBuilder = new MediaMetadataCompat.Builder();
    metadataBuilder
            .putLong(MediaMetadata.METADATA_KEY_DURATION, clip.getDuration())
            .putString(MediaMetadata.METADATA_KEY_MEDIA_ID, clip.getClipId())
            .putString(MediaMetadata.METADATA_KEY_TITLE, clip.getTitle())
            .putString(MediaMetadata.METADATA_KEY_ARTIST, clip.getSourceName())
            .putString(MediaMetadata.METADATA_KEY_ALBUM_ART_URI, clip.getImageUrl())
            .putLong(MediaMetadata.METADATA_KEY_DURATION, clip.getDuration());
    mSession.setMetadata(metadataBuilder.build());

Finally, the notification code:

        MediaSession mediaSession = (MediaSession) session.getMediaSession();
        Notification.Builder builder =
                new Notification.Builder(c)
                        .setDefaults(0)
                        .setSmallIcon(R.drawable.ic_notif)
                        .setVisibility(NotificationCompat.VISIBILITY_PUBLIC)
                        .setContentTitle(clip.getTitle())
                        .setContentText(clip.getSourceName())
                        .setProgress((int)duration, (int)progress, false)
                        .setWhen(0)
                        .setContentIntent(pendingIntent);

        if (playing) {
            builder.addAction(R.drawable.ic_media_pause, c.getString(R.string.media_pause),
                    getPendingIntentForKeyCode(app.getApplicationContext(), KeyEvent.KEYCODE_MEDIA_PAUSE));
        } else {
            builder.addAction(R.drawable.ic_media_play, c.getString(R.string.media_play),
                    getPendingIntentForKeyCode(app.getApplicationContext(), KeyEvent.KEYCODE_MEDIA_PLAY));
        }
        builder.addAction(R.drawable.ic_media_next, c.getString(R.string.media_next),
                    getPendingIntentForKeyCode(app.getApplicationContext(), KeyEvent.KEYCODE_MEDIA_NEXT));

        builder.setStyle(new Notification.MediaStyle()
                .setMediaSession(mediaSession.getSessionToken())
                .setShowActionsInCompactView(new int[] {1, 2})
                )
        );

        notification = builder.build();

Set the playback state in your MediaSession with the actions that you support:

PlaybackState state = new PlaybackState.Builder()
        .setActions(
                PlaybackState.ACTION_PLAY | PlaybackState.ACTION_PLAY_PAUSE |
                PlaybackState.ACTION_PLAY_FROM_MEDIA_ID | PlaybackState.ACTION_PAUSE |
                PlaybackState.ACTION_SKIP_TO_NEXT | PlaybackState.ACTION_SKIP_TO_PREVIOUS)
        .setState(PlaybackState.STATE_PLAYING, position, speed, SystemClock.elapsedRealtime())
        .build();
mSession.setPlaybackState(state);