diff --git a/android/VmLauncherApp/java/com/android/virtualization/vmlauncher/DisplayProvider.java b/android/VmLauncherApp/java/com/android/virtualization/vmlauncher/DisplayProvider.java
new file mode 100644
index 0000000..6eba709
--- /dev/null
+++ b/android/VmLauncherApp/java/com/android/virtualization/vmlauncher/DisplayProvider.java
@@ -0,0 +1,210 @@
+/*
+ * Copyright (C) 2024 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.virtualization.vmlauncher;
+
+import android.crosvm.ICrosvmAndroidDisplayService;
+import android.graphics.PixelFormat;
+import android.os.IBinder;
+import android.os.ParcelFileDescriptor;
+import android.os.RemoteException;
+import android.os.ServiceManager;
+import android.system.virtualizationservice_internal.IVirtualizationServiceInternal;
+import android.util.Log;
+import android.view.SurfaceControl;
+import android.view.SurfaceHolder;
+import android.view.SurfaceView;
+
+import libcore.io.IoBridge;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+
+/** Presents Android-side surface where VM can use as a display */
+class DisplayProvider {
+    private static final String TAG = MainActivity.TAG;
+    private final SurfaceView mMainView;
+    private final SurfaceView mCursorView;
+    private final IVirtualizationServiceInternal mVirtService;
+    private CursorHandler mCursorHandler;
+
+    DisplayProvider(SurfaceView mainView, SurfaceView cursorView) {
+        mMainView = mainView;
+        mCursorView = cursorView;
+
+        mMainView.setSurfaceLifecycle(SurfaceView.SURFACE_LIFECYCLE_FOLLOWS_ATTACHMENT);
+        mMainView.getHolder().addCallback(new Callback(Callback.SurfaceKind.MAIN));
+
+        mCursorView.setSurfaceLifecycle(SurfaceView.SURFACE_LIFECYCLE_FOLLOWS_ATTACHMENT);
+        mCursorView.getHolder().addCallback(new Callback(Callback.SurfaceKind.CURSOR));
+        mCursorView.getHolder().setFormat(PixelFormat.RGBA_8888);
+        // TODO: do we need this z-order?
+        mCursorView.setZOrderMediaOverlay(true);
+
+        IBinder b = ServiceManager.waitForService("android.system.virtualizationservice");
+        mVirtService = IVirtualizationServiceInternal.Stub.asInterface(b);
+        try {
+            // To ensure that the previous display service is removed.
+            mVirtService.clearDisplayService();
+        } catch (RemoteException e) {
+            throw new RuntimeException("Failed to clear prior display service", e);
+        }
+    }
+
+    void notifyDisplayIsGoingToInvisible() {
+        // When the display is going to be invisible (by putting in the background), save the frame
+        // of the main surface so that we can re-draw it next time the display becomes visible. This
+        // is to save the duration of time where nothing is drawn by VM.
+        try {
+            getDisplayService().saveFrameForSurface(false /* forCursor */);
+        } catch (RemoteException e) {
+            throw new RuntimeException("Failed to save frame for the main surface", e);
+        }
+    }
+
+    private synchronized ICrosvmAndroidDisplayService getDisplayService() {
+        try {
+            IBinder b = mVirtService.waitDisplayService();
+            return ICrosvmAndroidDisplayService.Stub.asInterface(b);
+        } catch (Exception e) {
+            throw new RuntimeException("Error while getting display service", e);
+        }
+    }
+
+    private class Callback implements SurfaceHolder.Callback {
+        enum SurfaceKind {
+            MAIN,
+            CURSOR
+        }
+
+        private final SurfaceKind mSurfaceKind;
+
+        Callback(SurfaceKind kind) {
+            mSurfaceKind = kind;
+        }
+
+        private boolean isForCursor() {
+            return mSurfaceKind == SurfaceKind.CURSOR;
+        }
+
+        @Override
+        public void surfaceCreated(SurfaceHolder holder) {
+            try {
+                getDisplayService().setSurface(holder.getSurface(), isForCursor());
+            } catch (Exception e) {
+                // TODO: don't consume this exception silently. For some unknown reason, setSurface
+                // call above throws IllegalArgumentException and that fails the surface
+                // configuration.
+                Log.e(TAG, "Failed to present surface " + mSurfaceKind + " to VM", e);
+            }
+
+            try {
+                switch (mSurfaceKind) {
+                    case MAIN:
+                        getDisplayService().drawSavedFrameForSurface(isForCursor());
+                        break;
+                    case CURSOR:
+                        ParcelFileDescriptor stream = createNewCursorStream();
+                        getDisplayService().setCursorStream(stream);
+                        break;
+                }
+            } catch (Exception e) {
+                // TODO: don't consume exceptions here too
+                Log.e(TAG, "Failed to configure surface " + mSurfaceKind, e);
+            }
+        }
+
+        @Override
+        public void surfaceChanged(SurfaceHolder holder, int format, int width, int height) {
+            // TODO: support resizeable display. We could actually change the display size that the
+            // VM sees, or keep the size and render it by fitting it in the new surface.
+        }
+
+        @Override
+        public void surfaceDestroyed(SurfaceHolder holder) {
+            try {
+                getDisplayService().removeSurface(isForCursor());
+            } catch (RemoteException e) {
+                throw new RuntimeException("Error while destroying surface for " + mSurfaceKind, e);
+            }
+        }
+    }
+
+    private ParcelFileDescriptor createNewCursorStream() {
+        if (mCursorHandler != null) {
+            mCursorHandler.interrupt();
+        }
+        ParcelFileDescriptor[] pfds;
+        try {
+            pfds = ParcelFileDescriptor.createSocketPair();
+        } catch (IOException e) {
+            throw new RuntimeException("Failed to create socketpair for cursor stream", e);
+        }
+        mCursorHandler = new CursorHandler(pfds[0]);
+        mCursorHandler.start();
+        return pfds[1];
+    }
+
+    /**
+     * Thread reading cursor coordinate from a stream, and updating the position of the cursor
+     * surface accordingly.
+     */
+    private class CursorHandler extends Thread {
+        private final ParcelFileDescriptor mStream;
+        private final SurfaceControl mCursor;
+        private final SurfaceControl.Transaction mTransaction;
+
+        CursorHandler(ParcelFileDescriptor stream) {
+            mStream = stream;
+            mCursor = DisplayProvider.this.mCursorView.getSurfaceControl();
+            mTransaction = new SurfaceControl.Transaction();
+
+            SurfaceControl main = DisplayProvider.this.mMainView.getSurfaceControl();
+            mTransaction.reparent(mCursor, main).apply();
+        }
+
+        @Override
+        public void run() {
+            try {
+                ByteBuffer byteBuffer = ByteBuffer.allocate(8 /* (x: u32, y: u32) */);
+                byteBuffer.order(ByteOrder.LITTLE_ENDIAN);
+                while (true) {
+                    if (Thread.interrupted()) {
+                        Log.d(TAG, "CursorHandler thread interrupted!");
+                        return;
+                    }
+                    byteBuffer.clear();
+                    int bytes =
+                            IoBridge.read(
+                                    mStream.getFileDescriptor(),
+                                    byteBuffer.array(),
+                                    0,
+                                    byteBuffer.array().length);
+                    if (bytes == -1) {
+                        Log.e(TAG, "cannot read from cursor stream, stop the handler");
+                        return;
+                    }
+                    float x = (float) (byteBuffer.getInt() & 0xFFFFFFFF);
+                    float y = (float) (byteBuffer.getInt() & 0xFFFFFFFF);
+                    mTransaction.setPosition(mCursor, x, y).apply();
+                }
+            } catch (IOException e) {
+                Log.e(TAG, "failed to run CursorHandler", e);
+            }
+        }
+    }
+}
diff --git a/android/VmLauncherApp/java/com/android/virtualization/vmlauncher/MainActivity.java b/android/VmLauncherApp/java/com/android/virtualization/vmlauncher/MainActivity.java
index 29726f5..82331b3 100644
--- a/android/VmLauncherApp/java/com/android/virtualization/vmlauncher/MainActivity.java
+++ b/android/VmLauncherApp/java/com/android/virtualization/vmlauncher/MainActivity.java
@@ -23,8 +23,6 @@
 import android.content.ClipData;
 import android.content.ClipboardManager;
 import android.content.Intent;
-import android.crosvm.ICrosvmAndroidDisplayService;
-import android.graphics.PixelFormat;
 import android.os.Bundle;
 import android.os.ParcelFileDescriptor;
 import android.os.RemoteException;
@@ -36,16 +34,12 @@
 import android.system.virtualmachine.VirtualMachineException;
 import android.system.virtualmachine.VirtualMachineManager;
 import android.util.Log;
-import android.view.SurfaceControl;
-import android.view.SurfaceHolder;
 import android.view.SurfaceView;
 import android.view.View;
 import android.view.WindowInsets;
 import android.view.WindowInsetsController;
 import android.view.WindowManager;
 
-import libcore.io.IoBridge;
-
 import java.io.BufferedOutputStream;
 import java.io.BufferedReader;
 import java.io.FileInputStream;
@@ -72,9 +66,9 @@
 
     private ExecutorService mExecutorService;
     private VirtualMachine mVirtualMachine;
-    private CursorHandler mCursorHandler;
     private ClipboardManager mClipboardManager;
     private InputForwarder mInputForwarder;
+    private DisplayProvider mDisplayProvider;
 
     @Override
     protected void onCreate(Bundle savedInstanceState) {
@@ -87,14 +81,6 @@
         }
         checkAndRequestRecordAudioPermission();
         mExecutorService = Executors.newCachedThreadPool();
-        try {
-            // To ensure that the previous display service is removed.
-            IVirtualizationServiceInternal.Stub.asInterface(
-                            ServiceManager.waitForService("android.system.virtualizationservice"))
-                    .clearDisplayService();
-        } catch (RemoteException e) {
-            Log.d(TAG, "failed to clearDisplayService");
-        }
         getWindow().setDecorFitsSystemWindows(false);
         setContentView(R.layout.activity_main);
         VirtualMachineCallback callback =
@@ -171,105 +157,7 @@
 
         SurfaceView surfaceView = findViewById(R.id.surface_view);
         SurfaceView cursorSurfaceView = findViewById(R.id.cursor_surface_view);
-        cursorSurfaceView.setZOrderMediaOverlay(true);
         View backgroundTouchView = findViewById(R.id.background_touch_view);
-        surfaceView
-                .getHolder()
-                .addCallback(
-                        // TODO(b/331708504): it should be handled in AVF framework.
-                        new SurfaceHolder.Callback() {
-                            @Override
-                            public void surfaceCreated(SurfaceHolder holder) {
-                                Log.d(
-                                        TAG,
-                                        "surface size: "
-                                                + holder.getSurfaceFrame().flattenToString());
-                                Log.d(
-                                        TAG,
-                                        "ICrosvmAndroidDisplayService.setSurface("
-                                                + holder.getSurface()
-                                                + ")");
-                                runWithDisplayService(
-                                        s ->
-                                                s.setSurface(
-                                                        holder.getSurface(),
-                                                        false /* forCursor */));
-                                // TODO  execute the above and the below togther with the same call
-                                // to runWithDisplayService. Currently this doesn't work because
-                                // setSurface somtimes trigger an exception and as a result
-                                // drawSavedFrameForSurface is skipped.
-                                runWithDisplayService(
-                                        s -> s.drawSavedFrameForSurface(false /* forCursor */));
-                            }
-
-                            @Override
-                            public void surfaceChanged(
-                                    SurfaceHolder holder, int format, int width, int height) {
-                                Log.d(
-                                        TAG,
-                                        "surface changed, width: " + width + ", height: " + height);
-                            }
-
-                            @Override
-                            public void surfaceDestroyed(SurfaceHolder holder) {
-                                Log.d(TAG, "ICrosvmAndroidDisplayService.removeSurface()");
-                                runWithDisplayService(
-                                        (service) -> service.removeSurface(false /* forCursor */));
-                            }
-                        });
-        cursorSurfaceView.getHolder().setFormat(PixelFormat.RGBA_8888);
-        cursorSurfaceView
-                .getHolder()
-                .addCallback(
-                        new SurfaceHolder.Callback() {
-                            @Override
-                            public void surfaceCreated(SurfaceHolder holder) {
-                                try {
-                                    ParcelFileDescriptor[] pfds =
-                                            ParcelFileDescriptor.createSocketPair();
-                                    if (mCursorHandler != null) {
-                                        mCursorHandler.interrupt();
-                                    }
-                                    mCursorHandler =
-                                            new CursorHandler(
-                                                    surfaceView.getSurfaceControl(),
-                                                    cursorSurfaceView.getSurfaceControl(),
-                                                    pfds[0]);
-                                    mCursorHandler.start();
-                                    runWithDisplayService(
-                                            (service) -> service.setCursorStream(pfds[1]));
-                                } catch (Exception e) {
-                                    Log.d(TAG, "failed to run cursor stream handler", e);
-                                }
-                                Log.d(
-                                        TAG,
-                                        "ICrosvmAndroidDisplayService.setSurface("
-                                                + holder.getSurface()
-                                                + ")");
-                                runWithDisplayService(
-                                        (service) ->
-                                                service.setSurface(
-                                                        holder.getSurface(), true /* forCursor */));
-                            }
-
-                            @Override
-                            public void surfaceChanged(
-                                    SurfaceHolder holder, int format, int width, int height) {
-                                Log.d(
-                                        TAG,
-                                        "cursor surface changed, width: "
-                                                + width
-                                                + ", height: "
-                                                + height);
-                            }
-
-                            @Override
-                            public void surfaceDestroyed(SurfaceHolder holder) {
-                                Log.d(TAG, "ICrosvmAndroidDisplayService.removeSurface()");
-                                runWithDisplayService(
-                                        (service) -> service.removeSurface(true /* forCursor */));
-                            }
-                        });
         getWindow().addFlags(WindowManager.LayoutParams.FLAG_KEEP_SCREEN_ON);
 
         // Fullscreen:
@@ -284,6 +172,8 @@
         mInputForwarder =
                 new InputForwarder(
                         this, mVirtualMachine, touchReceiver, mouseReceiver, keyReceiver);
+
+        mDisplayProvider = new DisplayProvider(surfaceView, cursorSurfaceView);
     }
 
     @Override
@@ -295,7 +185,7 @@
     @Override
     protected void onPause() {
         super.onPause();
-        runWithDisplayService(s -> s.saveFrameForSurface(false /* forCursor */));
+        mDisplayProvider.notifyDisplayIsGoingToInvisible();
     }
 
     @Override
@@ -483,73 +373,6 @@
         }
     }
 
-    @FunctionalInterface
-    public interface RemoteExceptionCheckedFunction<T> {
-        void apply(T t) throws RemoteException;
-    }
-
-    private void runWithDisplayService(
-            RemoteExceptionCheckedFunction<ICrosvmAndroidDisplayService> func) {
-        IVirtualizationServiceInternal vs =
-                IVirtualizationServiceInternal.Stub.asInterface(
-                        ServiceManager.waitForService("android.system.virtualizationservice"));
-        try {
-            Log.d(TAG, "wait for the display service");
-            ICrosvmAndroidDisplayService service =
-                    ICrosvmAndroidDisplayService.Stub.asInterface(vs.waitDisplayService());
-            assert service != null;
-            func.apply(service);
-            Log.d(TAG, "display service runs successfully");
-        } catch (Exception e) {
-            Log.d(TAG, "error on running display service", e);
-        }
-    }
-
-    static class CursorHandler extends Thread {
-        private final SurfaceControl mCursor;
-        private final ParcelFileDescriptor mStream;
-        private final SurfaceControl.Transaction mTransaction;
-
-        CursorHandler(SurfaceControl main, SurfaceControl cursor, ParcelFileDescriptor stream) {
-            mCursor = cursor;
-            mStream = stream;
-            mTransaction = new SurfaceControl.Transaction();
-
-            mTransaction.reparent(cursor, main).apply();
-        }
-
-        @Override
-        public void run() {
-            Log.d(TAG, "running CursorHandler");
-            try {
-                ByteBuffer byteBuffer = ByteBuffer.allocate(8 /* (x: u32, y: u32) */);
-                byteBuffer.order(ByteOrder.LITTLE_ENDIAN);
-                while (true) {
-                    if (Thread.interrupted()) {
-                        Log.d(TAG, "interrupted: exiting CursorHandler");
-                        return;
-                    }
-                    byteBuffer.clear();
-                    int bytes =
-                            IoBridge.read(
-                                    mStream.getFileDescriptor(),
-                                    byteBuffer.array(),
-                                    0,
-                                    byteBuffer.array().length);
-                    if (bytes == -1) {
-                        Log.e(TAG, "cannot read from cursor stream, stop the handler");
-                        return;
-                    }
-                    float x = (float) (byteBuffer.getInt() & 0xFFFFFFFF);
-                    float y = (float) (byteBuffer.getInt() & 0xFFFFFFFF);
-                    mTransaction.setPosition(mCursor, x, y).apply();
-                }
-            } catch (IOException e) {
-                Log.e(TAG, "failed to run CursorHandler", e);
-            }
-        }
-    }
-
     private void checkAndRequestRecordAudioPermission() {
         if (getApplicationContext().checkSelfPermission(permission.RECORD_AUDIO)
                 != PERMISSION_GRANTED) {