diff --git a/services/core/java/com/android/server/pm/PackageInstallerSession.java b/services/core/java/com/android/server/pm/PackageInstallerSession.java
index e6a62c5dce8c6..7efe0810946a1 100644
--- a/services/core/java/com/android/server/pm/PackageInstallerSession.java
+++ b/services/core/java/com/android/server/pm/PackageInstallerSession.java
@@ -28,6 +28,7 @@ import static android.content.pm.PackageManager.INSTALL_FAILED_INVALID_APK;
 import static android.content.pm.PackageManager.INSTALL_FAILED_MEDIA_UNAVAILABLE;
 import static android.content.pm.PackageManager.INSTALL_FAILED_MISSING_SPLIT;
 import static android.content.pm.PackageManager.INSTALL_PARSE_FAILED_CERTIFICATE_ENCODING;
+import static android.content.pm.PackageManager.INSTALL_STAGED;
 import static android.content.pm.PackageManager.INSTALL_SUCCEEDED;
 import static android.content.pm.PackageParser.APEX_FILE_EXTENSION;
 import static android.content.pm.PackageParser.APK_FILE_EXTENSION;
@@ -2232,6 +2233,10 @@ public class PackageInstallerSession extends IPackageInstallerSession.Stub {
             user = new UserHandle(userId);
         }
 
+        if (params.isStaged) {
+            params.installFlags |= INSTALL_STAGED;
+        }
+
         synchronized (mLock) {
             return mPm.new InstallParams(stageDir, localObserver, params, mInstallSource, user,
                     mSigningDetails, mInstallerUid);
diff --git a/services/core/java/com/android/server/pm/PackageManagerService.java b/services/core/java/com/android/server/pm/PackageManagerService.java
index cc1ef86bfb1a5..6da19df6ff67c 100644
--- a/services/core/java/com/android/server/pm/PackageManagerService.java
+++ b/services/core/java/com/android/server/pm/PackageManagerService.java
@@ -66,6 +66,7 @@ import static android.content.pm.PackageManager.INSTALL_PARSE_FAILED_INCONSISTEN
 import static android.content.pm.PackageManager.INSTALL_PARSE_FAILED_NO_CERTIFICATES;
 import static android.content.pm.PackageManager.INSTALL_REASON_DEVICE_RESTORE;
 import static android.content.pm.PackageManager.INSTALL_REASON_DEVICE_SETUP;
+import static android.content.pm.PackageManager.INSTALL_STAGED;
 import static android.content.pm.PackageManager.INSTALL_SUCCEEDED;
 import static android.content.pm.PackageManager.INTENT_FILTER_DOMAIN_VERIFICATION_STATUS_ALWAYS;
 import static android.content.pm.PackageManager.INTENT_FILTER_DOMAIN_VERIFICATION_STATUS_ALWAYS_ASK;
@@ -15977,7 +15978,7 @@ public class PackageManagerService extends IPackageManager.Stub
                 return false;
             }
 
-            final File targetDir = codeFile.getParentFile();
+            final File targetDir = resolveTargetDir();
             final File beforeCodeFile = codeFile;
             final File afterCodeFile = getNextCodePath(targetDir, parsedPackage.getPackageName());
 
@@ -16020,6 +16021,17 @@ public class PackageManagerService extends IPackageManager.Stub
             return true;
         }
 
+        // TODO(b/168126411): Once staged install flow starts using the same folder as non-staged
+        //  flow, we won't need this method anymore.
+        private File resolveTargetDir() {
+            boolean isStagedInstall = (installFlags & INSTALL_STAGED) != 0;
+            if (isStagedInstall) {
+                return Environment.getDataAppDirectory(null);
+            } else {
+                return codeFile.getParentFile();
+            }
+        }
+
         int doPostInstall(int status, int uid) {
             if (status != PackageManager.INSTALL_SUCCEEDED) {
                 cleanUp();