diff --git a/javalib/src/android/system/virtualmachine/VirtualMachine.java b/javalib/src/android/system/virtualmachine/VirtualMachine.java
index 955b350..9d8a4f7 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachine.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachine.java
@@ -100,21 +100,21 @@
     private final Object mLock = new Object();
 
     /** The package which owns this VM. */
-    private final @NonNull String mPackageName;
+    @NonNull private final String mPackageName;
 
     /** Name of this VM within the package. The name should be unique in the package. */
-    private final @NonNull String mName;
+    @NonNull private final String mName;
 
     /**
      * Path to the config file for this VM. The config file is where the configuration is persisted.
      */
-    private final @NonNull File mConfigFilePath;
+    @NonNull private final File mConfigFilePath;
 
     /** Path to the instance image file for this VM. */
-    private final @NonNull File mInstanceFilePath;
+    @NonNull private final File mInstanceFilePath;
 
     /** Path to the idsig file for this VM. */
-    private final @NonNull File mIdsigFilePath;
+    @NonNull private final File mIdsigFilePath;
 
     private static class ExtraApkSpec {
         public final File apk;
@@ -130,30 +130,32 @@
      * List of extra apks. Apks are specified by the vm config, and corresponding idsigs are to be
      * generated.
      */
-    private final @NonNull List<ExtraApkSpec> mExtraApks;
+    @NonNull private final List<ExtraApkSpec> mExtraApks;
 
     /** Size of the instance image. 10 MB. */
     private static final long INSTANCE_FILE_SIZE = 10 * 1024 * 1024;
 
     /** The configuration that is currently associated with this VM. */
-    private @NonNull VirtualMachineConfig mConfig;
+    @NonNull private VirtualMachineConfig mConfig;
 
     /** Handle to the "running" VM. */
-    private @Nullable IVirtualMachine mVirtualMachine;
+    @Nullable private IVirtualMachine mVirtualMachine;
 
     /** The registered callback */
     @GuardedBy("mLock")
-    private @Nullable VirtualMachineCallback mCallback;
+    @Nullable
+    private VirtualMachineCallback mCallback;
 
     /** The executor on which the callback will be executed */
     @GuardedBy("mLock")
-    private @Nullable Executor mCallbackExecutor;
+    @Nullable
+    private Executor mCallbackExecutor;
 
-    private @Nullable ParcelFileDescriptor mConsoleReader;
-    private @Nullable ParcelFileDescriptor mConsoleWriter;
+    @Nullable private ParcelFileDescriptor mConsoleReader;
+    @Nullable private ParcelFileDescriptor mConsoleWriter;
 
-    private @Nullable ParcelFileDescriptor mLogReader;
-    private @Nullable ParcelFileDescriptor mLogWriter;
+    @Nullable private ParcelFileDescriptor mLogReader;
+    @Nullable private ParcelFileDescriptor mLogWriter;
 
     private final ExecutorService mExecutorService = Executors.newCachedThreadPool();
 
@@ -181,7 +183,8 @@
      * it is persisted until it is deleted by calling {@link #delete()}. The created virtual machine
      * is in {@link Status#STOPPED} state. To run the VM, call {@link #run()}.
      */
-    /* package */ static @NonNull VirtualMachine create(
+    @NonNull
+    static VirtualMachine create(
             @NonNull Context context, @NonNull String name, @NonNull VirtualMachineConfig config)
             throws VirtualMachineException {
         if (config == null) {
@@ -232,7 +235,8 @@
     }
 
     /** Loads a virtual machine that is already created before. */
-    /* package */ static @Nullable VirtualMachine load(
+    @Nullable
+    static VirtualMachine load(
             @NonNull Context context, @NonNull String name) throws VirtualMachineException {
         File configFilePath = getConfigFilePath(context, name);
         VirtualMachineConfig config;
@@ -260,8 +264,11 @@
     /**
      * Returns the name of this virtual machine. The name is unique in the package and can't be
      * changed.
+     *
+     * @hide
      */
-    public @NonNull String getName() {
+    @NonNull
+    public String getName() {
         return mName;
     }
 
@@ -271,13 +278,21 @@
      * isolated from each other; one cannot share its secret to another virtual machine even if they
      * share the same config. It is also possible that a virtual machine can switch its config,
      * which can be done by calling {@link #setConfig(VirtualMachineConfig)}.
+     *
+     * @hide
      */
-    public @NonNull VirtualMachineConfig getConfig() {
+    @NonNull
+    public VirtualMachineConfig getConfig() {
         return mConfig;
     }
 
-    /** Returns the current status of this virtual machine. */
-    public @NonNull Status getStatus() throws VirtualMachineException {
+    /**
+     * Returns the current status of this virtual machine.
+     *
+     * @hide
+     */
+    @NonNull
+    public Status getStatus() throws VirtualMachineException {
         try {
             if (mVirtualMachine != null) {
                 switch (mVirtualMachine.getState()) {
@@ -304,6 +319,8 @@
     /**
      * Registers the callback object to get events from the virtual machine. If a callback was
      * already registered, it is replaced with the new one.
+     *
+     * @hide
      */
     public void setCallback(
             @NonNull @CallbackExecutor Executor executor,
@@ -314,7 +331,11 @@
         }
     }
 
-    /** Clears the currently registered callback. */
+    /**
+     * Clears the currently registered callback.
+     *
+     * @hide
+     */
     public void clearCallback() {
         synchronized (mLock) {
             mCallback = null;
@@ -345,6 +366,8 @@
      * Runs this virtual machine. The returning of this method however doesn't mean that the VM has
      * actually started running or the OS has booted there. Such events can be notified by
      * registering a callback object (not implemented currently).
+     *
+     * @hide
      */
     public void run() throws VirtualMachineException {
         if (getStatus() != Status.STOPPED) {
@@ -405,13 +428,10 @@
             // The VM should only be observed to die once
             AtomicBoolean onDiedCalled = new AtomicBoolean(false);
 
-            IBinder.DeathRecipient deathRecipient = new IBinder.DeathRecipient() {
-                @Override
-                public void binderDied() {
-                    if (onDiedCalled.compareAndSet(false, true)) {
-                        executeCallback((cb) -> cb.onDied(VirtualMachine.this,
-                                VirtualMachineCallback.DEATH_REASON_VIRTUALIZATIONSERVICE_DIED));
-                    }
+            IBinder.DeathRecipient deathRecipient = () -> {
+                if (onDiedCalled.compareAndSet(false, true)) {
+                    executeCallback((cb) -> cb.onDied(VirtualMachine.this,
+                            VirtualMachineCallback.DEATH_REASON_VIRTUALIZATIONSERVICE_DIED));
                 }
             };
 
@@ -454,23 +474,31 @@
             );
             service.asBinder().linkToDeath(deathRecipient, 0);
             mVirtualMachine.start();
-        } catch (IOException e) {
-            throw new VirtualMachineException(e);
-        } catch (RemoteException e) {
+        } catch (IOException | RemoteException e) {
             throw new VirtualMachineException(e);
         }
     }
 
-    /** Returns the stream object representing the console output from the virtual machine. */
-    public @NonNull InputStream getConsoleOutputStream() throws VirtualMachineException {
+    /**
+     * Returns the stream object representing the console output from the virtual machine.
+     *
+     * @hide
+     */
+    @NonNull
+    public InputStream getConsoleOutputStream() throws VirtualMachineException {
         if (mConsoleReader == null) {
             throw new VirtualMachineException("Console output not available");
         }
         return new FileInputStream(mConsoleReader.getFileDescriptor());
     }
 
-    /** Returns the stream object representing the log output from the virtual machine. */
-    public @NonNull InputStream getLogOutputStream() throws VirtualMachineException {
+    /**
+     * Returns the stream object representing the log output from the virtual machine.
+     *
+     * @hide
+     */
+    @NonNull
+    public InputStream getLogOutputStream() throws VirtualMachineException {
         if (mLogReader == null) {
             throw new VirtualMachineException("Log output not available");
         }
@@ -482,6 +510,8 @@
      * computer; the machine halts immediately. Software running on the virtual machine is not
      * notified with the event. A stopped virtual machine can be re-started by calling {@link
      * #run()}.
+     *
+     * @hide
      */
     public void stop() throws VirtualMachineException {
         if (mVirtualMachine == null) return;
@@ -498,6 +528,8 @@
      * associated with it including the per-VM secret. This is an irreversable action. A virtual
      * machine once deleted can never be restored. A new virtual machine created with the same name
      * and the same config is different from an already deleted virtual machine.
+     *
+     * @hide
      */
     public void delete() throws VirtualMachineException {
         if (getStatus() != Status.STOPPED) {
@@ -513,8 +545,13 @@
         vmRootDir.delete();
     }
 
-    /** Returns the CID of this virtual machine, if it is running. */
-    public @NonNull Optional<Integer> getCid() throws VirtualMachineException {
+    /**
+     * Returns the CID of this virtual machine, if it is running.
+     *
+     * @hide
+     */
+    @NonNull
+    public Optional<Integer> getCid() throws VirtualMachineException {
         if (getStatus() != Status.RUNNING) {
             return Optional.empty();
         }
@@ -535,8 +572,11 @@
      * when an incompatible config is attempted.
      *
      * @return the old config
+     *
+     * @hide
      */
-    public @NonNull VirtualMachineConfig setConfig(@NonNull VirtualMachineConfig newConfig)
+    @NonNull
+    public VirtualMachineConfig setConfig(@NonNull VirtualMachineConfig newConfig)
             throws VirtualMachineException {
         final VirtualMachineConfig oldConfig = getConfig();
         if (!oldConfig.isCompatibleWith(newConfig)) {
@@ -566,6 +606,8 @@
      * expected to set up vsock servers in their payload. After the host app receives the {@link
      * VirtualMachineCallback#onPayloadReady(VirtualMachine)}, it can use this method to
      * establish an RPC session to the guest VMs.
+     *
+     * @hide
      */
     public Future<IBinder> connectToVsockServer(int port) throws VirtualMachineException {
         if (getStatus() != Status.RUNNING) {
@@ -586,7 +628,7 @@
 
     private static List<String> parseExtraApkListFromPayloadConfig(JsonReader reader)
             throws VirtualMachineException {
-        /**
+        /*
          * JSON schema from packages/modules/Virtualization/microdroid/payload/config/src/lib.rs:
          *
          * <p>{ "extra_apks": [ { "path": "/system/app/foo.apk", }, ... ], ... }
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
index e0f74ec..b7c7a88 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineConfig.java
@@ -21,6 +21,7 @@
 import android.annotation.NonNull;
 import android.content.Context;
 import android.content.pm.PackageManager;
+import android.content.pm.PackageManager.PackageInfoFlags;
 import android.content.pm.Signature; // This actually is certificate!
 import android.os.ParcelFileDescriptor;
 import android.os.PersistableBundle;
@@ -59,10 +60,14 @@
     private static final String KEY_CPU_AFFINITY = "cpuAffinity";
 
     // Paths to the APK file of this application.
-    private final @NonNull String mApkPath;
-    private final @NonNull Signature[] mCerts;
+    @NonNull private final String mApkPath;
+    @NonNull private final Signature[] mCerts;
 
-    /** A debug level defines the set of debug features that the VM can be configured to. */
+    /**
+     * A debug level defines the set of debug features that the VM can be configured to.
+     *
+     * @hide
+     */
     public enum DebugLevel {
         /**
          * Not debuggable at all. No log is exported from the VM. Debugger can't be attached to the
@@ -110,7 +115,7 @@
     /**
      * Path within the APK to the payload config file that defines software aspects of this config.
      */
-    private final @NonNull String mPayloadConfigPath;
+    @NonNull private final String mPayloadConfigPath;
 
     private VirtualMachineConfig(
             @NonNull String apkPath,
@@ -132,7 +137,8 @@
     }
 
     /** Loads a config from a stream, for example a file. */
-    /* package */ static @NonNull VirtualMachineConfig from(@NonNull InputStream input)
+    @NonNull
+    static VirtualMachineConfig from(@NonNull InputStream input)
             throws IOException, VirtualMachineException {
         PersistableBundle b = PersistableBundle.readFromStream(input);
         final int version = b.getInt(KEY_VERSION);
@@ -186,8 +192,13 @@
         b.writeToStream(output);
     }
 
-    /** Returns the path to the payload config within the owning application. */
-    public @NonNull String getPayloadConfigPath() {
+    /**
+     * Returns the path to the payload config within the owning application.
+     *
+     * @hide
+     */
+    @NonNull
+    public String getPayloadConfigPath() {
         return mPayloadConfigPath;
     }
 
@@ -197,6 +208,8 @@
      * number of CPUs and the size of the RAM, and change of the payload as long as the payload is
      * signed by the same signer. All other changes (e.g. using a payload from a different signer,
      * change of the debug mode, etc.) are considered as incompatible.
+     *
+     * @hide
      */
     public boolean isCompatibleWith(@NonNull VirtualMachineConfig other) {
         if (!Arrays.equals(this.mCerts, other.mCerts)) {
@@ -243,17 +256,25 @@
         return parcel;
     }
 
-    /** A builder used to create a {@link VirtualMachineConfig}. */
+    /**
+     * A builder used to create a {@link VirtualMachineConfig}.
+     *
+     * @hide
+     */
     public static class Builder {
-        private Context mContext;
-        private String mPayloadConfigPath;
+        private final Context mContext;
+        private final String mPayloadConfigPath;
         private DebugLevel mDebugLevel;
         private boolean mProtectedVm;
         private int mMemoryMib;
         private int mNumCpus;
         private String mCpuAffinity;
 
-        /** Creates a builder for the given context (APK), and the payload config file in APK. */
+        /**
+         * Creates a builder for the given context (APK), and the payload config file in APK.
+         *
+         * @hide
+         */
         public Builder(@NonNull Context context, @NonNull String payloadConfigPath) {
             mContext = Objects.requireNonNull(context);
             mPayloadConfigPath = Objects.requireNonNull(payloadConfigPath);
@@ -263,13 +284,21 @@
             mCpuAffinity = null;
         }
 
-        /** Sets the debug level */
+        /**
+         * Sets the debug level
+         *
+         * @hide
+         */
         public Builder debugLevel(DebugLevel debugLevel) {
             mDebugLevel = debugLevel;
             return this;
         }
 
-        /** Sets whether to protect the VM memory from the host. Defaults to false. */
+        /**
+         *  Sets whether to protect the VM memory from the host. Defaults to false.
+         *
+         * @hide
+         */
         public Builder protectedVm(boolean protectedVm) {
             mProtectedVm = protectedVm;
             return this;
@@ -278,6 +307,8 @@
         /**
          * Sets the amount of RAM to give the VM. If this is zero or negative then the default will
          * be used.
+         *
+         * @hide
          */
         public Builder memoryMib(int memoryMib) {
             mMemoryMib = memoryMib;
@@ -286,6 +317,8 @@
 
         /**
          * Sets the number of vCPUs in the VM. Defaults to 1.
+         *
+         * @hide
          */
         public Builder numCpus(int num) {
             mNumCpus = num;
@@ -297,24 +330,30 @@
          * or CPU ranges to run vCPUs on. e.g. "0,1-3,5" to choose host CPUs 0, 1, 2, 3, and 5.
          * Or this can be a colon-separated list of assignments of vCPU to host CPU assignments.
          * e.g. "0=0:1=1:2=2" to map vCPU 0 to host CPU 0, and so on.
+         *
+         * @hide
          */
         public Builder cpuAffinity(String affinity) {
             mCpuAffinity = affinity;
             return this;
         }
 
-        /** Builds an immutable {@link VirtualMachineConfig} */
-        public @NonNull VirtualMachineConfig build() {
+        /**
+         * Builds an immutable {@link VirtualMachineConfig}
+         *
+         * @hide
+         */
+        @NonNull
+        public VirtualMachineConfig build() {
             final String apkPath = mContext.getPackageCodePath();
             final String packageName = mContext.getPackageName();
             Signature[] certs;
             try {
-                certs =
-                        mContext.getPackageManager()
-                                .getPackageInfo(
-                                        packageName, PackageManager.GET_SIGNING_CERTIFICATES)
-                                .signingInfo
-                                .getSigningCertificateHistory();
+                certs = mContext.getPackageManager()
+                        .getPackageInfo(packageName,
+                                PackageInfoFlags.of(PackageManager.GET_SIGNING_CERTIFICATES))
+                        .signingInfo
+                        .getSigningCertificateHistory();
             } catch (PackageManager.NameNotFoundException e) {
                 // This cannot happen as `packageName` is from this app.
                 throw new RuntimeException(e);
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineManager.java b/javalib/src/android/system/virtualmachine/VirtualMachineManager.java
index eaa383e..1ffc6bb 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineManager.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineManager.java
@@ -31,23 +31,29 @@
  * @hide
  */
 public class VirtualMachineManager {
-    private final @NonNull Context mContext;
+    @NonNull private final Context mContext;
 
     private VirtualMachineManager(@NonNull Context context) {
         mContext = context;
     }
 
-    static Map<Context, WeakReference<VirtualMachineManager>> sInstances = new WeakHashMap<>();
+    private static final Map<Context, WeakReference<VirtualMachineManager>> sInstances =
+            new WeakHashMap<>();
 
-    /** Returns the per-context instance. */
-    public static @NonNull VirtualMachineManager getInstance(@NonNull Context context) {
+    /**
+     * Returns the per-context instance.
+     *
+     * @hide
+     */
+    @NonNull
+    public static VirtualMachineManager getInstance(@NonNull Context context) {
         Objects.requireNonNull(context);
         synchronized (sInstances) {
             VirtualMachineManager vmm =
                     sInstances.containsKey(context) ? sInstances.get(context).get() : null;
             if (vmm == null) {
                 vmm = new VirtualMachineManager(context);
-                sInstances.put(context, new WeakReference(vmm));
+                sInstances.put(context, new WeakReference<>(vmm));
             }
             return vmm;
         }
@@ -62,8 +68,11 @@
      * machine has to be deleted before its name can be reused. Every call to this methods creates a
      * new (and different) virtual machine even if the name and the config are the same as the
      * deleted one.
+     *
+     * @hide
      */
-    public @NonNull VirtualMachine create(
+    @NonNull
+    public VirtualMachine create(
             @NonNull String name, @NonNull VirtualMachineConfig config)
             throws VirtualMachineException {
         synchronized (sCreateLock) {
@@ -74,16 +83,22 @@
     /**
      * Returns an existing {@link VirtualMachine} with the given name. Returns null if there is no
      * such virtual machine.
+     *
+     * @hide
      */
-    public @Nullable VirtualMachine get(@NonNull String name) throws VirtualMachineException {
+    @Nullable
+    public VirtualMachine get(@NonNull String name) throws VirtualMachineException {
         return VirtualMachine.load(mContext, name);
     }
 
     /**
      * Returns an existing {@link VirtualMachine} if it exists, or create a new one. The config
      * parameter is used only when a new virtual machine is created.
+     *
+     * @hide
      */
-    public @NonNull VirtualMachine getOrCreate(
+    @NonNull
+    public VirtualMachine getOrCreate(
             @NonNull String name, @NonNull VirtualMachineConfig config)
             throws VirtualMachineException {
         VirtualMachine vm;
diff --git a/tests/helper/src/java/com/android/microdroid/test/MicrodroidDeviceTestBase.java b/tests/helper/src/java/com/android/microdroid/test/MicrodroidDeviceTestBase.java
index efd7c85..1a573bb 100644
--- a/tests/helper/src/java/com/android/microdroid/test/MicrodroidDeviceTestBase.java
+++ b/tests/helper/src/java/com/android/microdroid/test/MicrodroidDeviceTestBase.java
@@ -23,7 +23,6 @@
 import android.os.ParcelFileDescriptor;
 import android.os.SystemProperties;
 import android.sysprop.HypervisorProperties;
-import android.system.virtualizationservice.DeathReason;
 import android.system.virtualmachine.VirtualMachine;
 import android.system.virtualmachine.VirtualMachineCallback;
 import android.system.virtualmachine.VirtualMachineConfig;
@@ -223,7 +222,7 @@
 
         @Override
         @CallSuper
-        public void onDied(VirtualMachine vm, @DeathReason int reason) {
+        public void onDied(VirtualMachine vm, int reason) {
             mExecutorService.shutdown();
         }
 
@@ -318,7 +317,7 @@
         listener.runToFinish(logTag, vm);
         return new BootResult(
                 payloadStarted.getNow(false),
-                deathReason.getNow(DeathReason.INFRASTRUCTURE_ERROR),
+                deathReason.getNow(VirtualMachineCallback.DEATH_REASON_INFRASTRUCTURE_ERROR),
                 apiCallNanoTime,
                 endTime.getNow(apiCallNanoTime) - apiCallNanoTime,
                 listener.getVcpuStartedNanoTime(),
diff --git a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
index 911efbb..2c6d139 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -23,8 +23,8 @@
 import android.os.Build;
 import android.os.ParcelFileDescriptor;
 import android.os.SystemProperties;
-import android.system.virtualizationservice.DeathReason;
 import android.system.virtualmachine.VirtualMachine;
+import android.system.virtualmachine.VirtualMachineCallback;
 import android.system.virtualmachine.VirtualMachineConfig;
 import android.system.virtualmachine.VirtualMachineConfig.DebugLevel;
 import android.system.virtualmachine.VirtualMachineException;
@@ -172,13 +172,14 @@
         VmEventListener listener =
                 new VmEventListener() {
                     @Override
-                    public void onDied(VirtualMachine vm, @DeathReason int reason) {
+                    public void onDied(VirtualMachine vm,  int reason) {
                         exception.complete(reason);
                         super.onDied(vm, reason);
                     }
                 };
         listener.runToFinish(TAG, vm);
-        assertThat(exception.getNow(0)).isAnyOf(DeathReason.REBOOT, DeathReason.HANGUP);
+        assertThat(exception.getNow(0)).isAnyOf(VirtualMachineCallback.DEATH_REASON_REBOOT,
+                VirtualMachineCallback.DEATH_REASON_HANGUP);
     }
 
     @Test
@@ -434,7 +435,7 @@
         assertThat(result.payloadStarted).isFalse();
 
         // This failure should shut the VM down immediately and shouldn't trigger a hangup.
-        assertThat(result.deathReason).isNotEqualTo(DeathReason.HANGUP);
+        assertThat(result.deathReason).isNotEqualTo(VirtualMachineCallback.DEATH_REASON_HANGUP);
     }
 
     @Test
@@ -502,6 +503,7 @@
 
         BootResult bootResult = tryBootVm(TAG, "test_vm_invalid_config");
         assertThat(bootResult.payloadStarted).isFalse();
-        assertThat(bootResult.deathReason).isEqualTo(DeathReason.MICRODROID_INVALID_PAYLOAD_CONFIG);
+        assertThat(bootResult.deathReason).isEqualTo(
+                VirtualMachineCallback.DEATH_REASON_MICRODROID_INVALID_PAYLOAD_CONFIG);
     }
 }