Merge cherrypicks of ['googleplex-android-review.googlesource.com/30803769', 'googleplex-android-review.googlesource.com/30983917', 'googleplex-android-review.googlesource.com/31234195', 'googleplex-android-review.googlesource.com/31234196'] into 25Q1-release.

Change-Id: I13bffdb5ac851386be64bb8985cf50cb3e544764
diff --git a/core/java/android/view/inputmethod/InputMethodManager.java b/core/java/android/view/inputmethod/InputMethodManager.java
index 6d89f3d..d5f471e 100644
--- a/core/java/android/view/inputmethod/InputMethodManager.java
+++ b/core/java/android/view/inputmethod/InputMethodManager.java
@@ -938,27 +938,6 @@
             synchronized (mH) {
                 if (mCurRootView == viewRootImpl) {
                     mCurRootViewWindowFocused = false;
-
-                    if (Flags.refactorInsetsController() && mCurRootView != null) {
-                        final int softInputMode = mCurRootView.mWindowAttributes.softInputMode;
-                        final int state =
-                                softInputMode & WindowManager.LayoutParams.SOFT_INPUT_MASK_STATE;
-                        if (state == WindowManager.LayoutParams.SOFT_INPUT_STATE_ALWAYS_HIDDEN) {
-                            // when losing focus (e.g., by going to another window), we reset the
-                            // requestedVisibleTypes of WindowInsetsController by hiding the IME
-                            final var statsToken = ImeTracker.forLogging().onStart(
-                                    ImeTracker.TYPE_HIDE, ImeTracker.ORIGIN_CLIENT,
-                                    SoftInputShowHideReason.REASON_HIDE_WINDOW_LOST_FOCUS,
-                                    false /* fromUser */);
-                            if (DEBUG) {
-                                Log.d(TAG, "onWindowLostFocus, hiding IME because "
-                                        + "of STATE_ALWAYS_HIDDEN");
-                            }
-                            mCurRootView.getInsetsController().hide(WindowInsets.Type.ime(),
-                                    false /* fromIme */, statsToken);
-                        }
-                    }
-
                     clearCurRootViewIfNeeded();
                 }
             }
@@ -1012,6 +991,26 @@
         @GuardedBy("mH")
         private void setCurrentRootViewLocked(ViewRootImpl rootView) {
             final boolean wasEmpty = mCurRootView == null;
+            if (Flags.refactorInsetsController() && !wasEmpty && mCurRootView != rootView) {
+                final int softInputMode = mCurRootView.mWindowAttributes.softInputMode;
+                final int state =
+                        softInputMode & WindowManager.LayoutParams.SOFT_INPUT_MASK_STATE;
+                if (state == WindowManager.LayoutParams.SOFT_INPUT_STATE_ALWAYS_HIDDEN) {
+                    // when losing input focus (e.g., by going to another window), we reset the
+                    // requestedVisibleTypes of WindowInsetsController by hiding the IME
+                    final var statsToken = ImeTracker.forLogging().onStart(
+                            ImeTracker.TYPE_HIDE, ImeTracker.ORIGIN_CLIENT,
+                            SoftInputShowHideReason.HIDE_WINDOW_LOST_FOCUS,
+                            false /* fromUser */);
+                    if (DEBUG) {
+                        Log.d(TAG, "setCurrentRootViewLocked, hiding IME because "
+                                + "of STATE_ALWAYS_HIDDEN");
+                    }
+                    mCurRootView.getInsetsController().hide(WindowInsets.Type.ime(),
+                            false /* fromIme */, statsToken);
+                }
+            }
+
             mImeDispatcher.switchRootView(mCurRootView, rootView);
             mCurRootView = rootView;
             if (wasEmpty && mCurRootView != null) {
diff --git a/core/java/com/android/internal/inputmethod/InputMethodDebug.java b/core/java/com/android/internal/inputmethod/InputMethodDebug.java
index 2a5593f..4d5e67a 100644
--- a/core/java/com/android/internal/inputmethod/InputMethodDebug.java
+++ b/core/java/com/android/internal/inputmethod/InputMethodDebug.java
@@ -299,6 +299,12 @@
                 return "SHOW_SOFT_INPUT_IMM_DEPRECATION";
             case SoftInputShowHideReason.CONTROL_WINDOW_INSETS_ANIMATION:
                 return "CONTROL_WINDOW_INSETS_ANIMATION";
+            case SoftInputShowHideReason.SHOW_INPUT_TARGET_CHANGED:
+                return "SHOW_INPUT_TARGET_CHANGED";
+            case SoftInputShowHideReason.HIDE_INPUT_TARGET_CHANGED:
+                return "HIDE_INPUT_TARGET_CHANGED";
+            case SoftInputShowHideReason.HIDE_WINDOW_LOST_FOCUS:
+                return "HIDE_WINDOW_LOST_FOCUS";
             default:
                 return "Unknown=" + reason;
         }
diff --git a/core/java/com/android/internal/inputmethod/SoftInputShowHideReason.java b/core/java/com/android/internal/inputmethod/SoftInputShowHideReason.java
index 592ea9e..cf0580c 100644
--- a/core/java/com/android/internal/inputmethod/SoftInputShowHideReason.java
+++ b/core/java/com/android/internal/inputmethod/SoftInputShowHideReason.java
@@ -91,7 +91,7 @@
         SoftInputShowHideReason.CONTROL_WINDOW_INSETS_ANIMATION,
         SoftInputShowHideReason.SHOW_INPUT_TARGET_CHANGED,
         SoftInputShowHideReason.HIDE_INPUT_TARGET_CHANGED,
-        SoftInputShowHideReason.REASON_HIDE_WINDOW_LOST_FOCUS,
+        SoftInputShowHideReason.HIDE_WINDOW_LOST_FOCUS,
 })
 public @interface SoftInputShowHideReason {
     /** Default, undefined reason. */
@@ -340,18 +340,6 @@
     int HIDE_WINDOW_LEGACY_DIRECT = ImeProtoEnums.REASON_HIDE_WINDOW_LEGACY_DIRECT;
 
     /**
-     * Show soft input because the input target changed
-     * {@link com.android.server.wm.ImeInsetsSourceProvider#onInputTargetChanged}.
-     */
-    int SHOW_INPUT_TARGET_CHANGED = ImeProtoEnums.REASON_SHOW_INPUT_TARGET_CHANGED;
-
-    /**
-     * Hide soft input because the input target changed by
-     * {@link com.android.server.wm.ImeInsetsSourceProvider#onInputTargetChanged}.
-     */
-    int HIDE_INPUT_TARGET_CHANGED = ImeProtoEnums.REASON_HIDE_INPUT_TARGET_CHANGED;
-
-    /**
      * Show / Hide soft input by
      * {@link android.inputmethodservice.InputMethodService#resetStateForNewConfiguration}.
      */
@@ -420,6 +408,18 @@
      */
     int CONTROL_WINDOW_INSETS_ANIMATION = ImeProtoEnums.REASON_CONTROL_WINDOW_INSETS_ANIMATION;
 
+    /**
+     * Show soft input because the input target changed
+     * {@link com.android.server.wm.ImeInsetsSourceProvider#onInputTargetChanged}.
+     */
+    int SHOW_INPUT_TARGET_CHANGED = ImeProtoEnums.REASON_SHOW_INPUT_TARGET_CHANGED;
+
+    /**
+     * Hide soft input because the input target changed by
+     * {@link com.android.server.wm.ImeInsetsSourceProvider#onInputTargetChanged}.
+     */
+    int HIDE_INPUT_TARGET_CHANGED = ImeProtoEnums.REASON_HIDE_INPUT_TARGET_CHANGED;
+
     /** Hide soft input when the window lost focus. */
-    int REASON_HIDE_WINDOW_LOST_FOCUS = ImeProtoEnums.REASON_HIDE_WINDOW_LOST_FOCUS;
+    int HIDE_WINDOW_LOST_FOCUS = ImeProtoEnums.REASON_HIDE_WINDOW_LOST_FOCUS;
 }
diff --git a/packages/NeuralNetworks/framework/Android.bp b/packages/NeuralNetworks/framework/Android.bp
index 6f45daa..af071ba 100644
--- a/packages/NeuralNetworks/framework/Android.bp
+++ b/packages/NeuralNetworks/framework/Android.bp
@@ -19,10 +19,21 @@
 filegroup {
     name: "framework-ondeviceintelligence-sources",
     srcs: [
-        "java/**/*.aidl",
-        "java/**/*.java",
+        "module/java/**/*.aidl",
+        "module/java/**/*.java",
     ],
-    path: "java",
+    visibility: [
+        "//frameworks/base:__subpackages__",
+        "//packages/modules/NeuralNetworks:__subpackages__",
+    ],
+}
+
+filegroup {
+    name: "framework-ondeviceintelligence-sources-platform",
+    srcs: [
+        "platform/java/**/*.aidl",
+        "platform/java/**/*.java",
+    ],
     visibility: [
         "//frameworks/base:__subpackages__",
         "//packages/modules/NeuralNetworks:__subpackages__",
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/DownloadCallback.java b/packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/DownloadCallback.java
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/DownloadCallback.java
rename to packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/DownloadCallback.java
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/Feature.aidl b/packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/Feature.aidl
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/Feature.aidl
rename to packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/Feature.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/Feature.java b/packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/Feature.java
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/Feature.java
rename to packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/Feature.java
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/FeatureDetails.aidl b/packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/FeatureDetails.aidl
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/FeatureDetails.aidl
rename to packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/FeatureDetails.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/FeatureDetails.java b/packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/FeatureDetails.java
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/FeatureDetails.java
rename to packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/FeatureDetails.java
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/ICancellationSignal.aidl b/packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/ICancellationSignal.aidl
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/ICancellationSignal.aidl
rename to packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/ICancellationSignal.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IDownloadCallback.aidl b/packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/IDownloadCallback.aidl
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IDownloadCallback.aidl
rename to packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/IDownloadCallback.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IFeatureCallback.aidl b/packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/IFeatureCallback.aidl
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IFeatureCallback.aidl
rename to packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/IFeatureCallback.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IFeatureDetailsCallback.aidl b/packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/IFeatureDetailsCallback.aidl
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IFeatureDetailsCallback.aidl
rename to packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/IFeatureDetailsCallback.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IListFeaturesCallback.aidl b/packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/IListFeaturesCallback.aidl
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IListFeaturesCallback.aidl
rename to packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/IListFeaturesCallback.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IOnDeviceIntelligenceManager.aidl b/packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/IOnDeviceIntelligenceManager.aidl
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IOnDeviceIntelligenceManager.aidl
rename to packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/IOnDeviceIntelligenceManager.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IProcessingSignal.aidl b/packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/IProcessingSignal.aidl
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IProcessingSignal.aidl
rename to packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/IProcessingSignal.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IRemoteCallback.aidl b/packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/IRemoteCallback.aidl
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IRemoteCallback.aidl
rename to packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/IRemoteCallback.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IResponseCallback.aidl b/packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/IResponseCallback.aidl
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IResponseCallback.aidl
rename to packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/IResponseCallback.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IStreamingResponseCallback.aidl b/packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/IStreamingResponseCallback.aidl
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IStreamingResponseCallback.aidl
rename to packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/IStreamingResponseCallback.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/ITokenInfoCallback.aidl b/packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/ITokenInfoCallback.aidl
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/ITokenInfoCallback.aidl
rename to packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/ITokenInfoCallback.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/InferenceInfo.aidl b/packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/InferenceInfo.aidl
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/InferenceInfo.aidl
rename to packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/InferenceInfo.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/InferenceInfo.java b/packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/InferenceInfo.java
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/InferenceInfo.java
rename to packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/InferenceInfo.java
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/OnDeviceIntelligenceException.java b/packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/OnDeviceIntelligenceException.java
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/OnDeviceIntelligenceException.java
rename to packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/OnDeviceIntelligenceException.java
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/OnDeviceIntelligenceFrameworkInitializer.java b/packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/OnDeviceIntelligenceFrameworkInitializer.java
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/OnDeviceIntelligenceFrameworkInitializer.java
rename to packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/OnDeviceIntelligenceFrameworkInitializer.java
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/OnDeviceIntelligenceManager.java b/packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/OnDeviceIntelligenceManager.java
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/OnDeviceIntelligenceManager.java
rename to packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/OnDeviceIntelligenceManager.java
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/ProcessingCallback.java b/packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/ProcessingCallback.java
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/ProcessingCallback.java
rename to packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/ProcessingCallback.java
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/ProcessingSignal.java b/packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/ProcessingSignal.java
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/ProcessingSignal.java
rename to packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/ProcessingSignal.java
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/StreamingProcessingCallback.java b/packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/StreamingProcessingCallback.java
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/StreamingProcessingCallback.java
rename to packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/StreamingProcessingCallback.java
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/TokenInfo.aidl b/packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/TokenInfo.aidl
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/TokenInfo.aidl
rename to packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/TokenInfo.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/TokenInfo.java b/packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/TokenInfo.java
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/TokenInfo.java
rename to packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/TokenInfo.java
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/utils/BinderUtils.java b/packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/utils/BinderUtils.java
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/utils/BinderUtils.java
rename to packages/NeuralNetworks/framework/module/java/android/app/ondeviceintelligence/utils/BinderUtils.java
diff --git a/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/IOnDeviceIntelligenceService.aidl b/packages/NeuralNetworks/framework/module/java/android/service/ondeviceintelligence/IOnDeviceIntelligenceService.aidl
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/IOnDeviceIntelligenceService.aidl
rename to packages/NeuralNetworks/framework/module/java/android/service/ondeviceintelligence/IOnDeviceIntelligenceService.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/IOnDeviceSandboxedInferenceService.aidl b/packages/NeuralNetworks/framework/module/java/android/service/ondeviceintelligence/IOnDeviceSandboxedInferenceService.aidl
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/IOnDeviceSandboxedInferenceService.aidl
rename to packages/NeuralNetworks/framework/module/java/android/service/ondeviceintelligence/IOnDeviceSandboxedInferenceService.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/IProcessingUpdateStatusCallback.aidl b/packages/NeuralNetworks/framework/module/java/android/service/ondeviceintelligence/IProcessingUpdateStatusCallback.aidl
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/IProcessingUpdateStatusCallback.aidl
rename to packages/NeuralNetworks/framework/module/java/android/service/ondeviceintelligence/IProcessingUpdateStatusCallback.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/IRemoteProcessingService.aidl b/packages/NeuralNetworks/framework/module/java/android/service/ondeviceintelligence/IRemoteProcessingService.aidl
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/IRemoteProcessingService.aidl
rename to packages/NeuralNetworks/framework/module/java/android/service/ondeviceintelligence/IRemoteProcessingService.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/IRemoteStorageService.aidl b/packages/NeuralNetworks/framework/module/java/android/service/ondeviceintelligence/IRemoteStorageService.aidl
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/IRemoteStorageService.aidl
rename to packages/NeuralNetworks/framework/module/java/android/service/ondeviceintelligence/IRemoteStorageService.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/OnDeviceIntelligenceService.java b/packages/NeuralNetworks/framework/module/java/android/service/ondeviceintelligence/OnDeviceIntelligenceService.java
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/OnDeviceIntelligenceService.java
rename to packages/NeuralNetworks/framework/module/java/android/service/ondeviceintelligence/OnDeviceIntelligenceService.java
diff --git a/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/OnDeviceSandboxedInferenceService.java b/packages/NeuralNetworks/framework/module/java/android/service/ondeviceintelligence/OnDeviceSandboxedInferenceService.java
similarity index 100%
rename from packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/OnDeviceSandboxedInferenceService.java
rename to packages/NeuralNetworks/framework/module/java/android/service/ondeviceintelligence/OnDeviceSandboxedInferenceService.java
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/DownloadCallback.java b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/DownloadCallback.java
similarity index 100%
copy from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/DownloadCallback.java
copy to packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/DownloadCallback.java
diff --git a/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/Feature.aidl b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/Feature.aidl
new file mode 100644
index 0000000..18494d7
--- /dev/null
+++ b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/Feature.aidl
@@ -0,0 +1,22 @@
+/*
+ * 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 android.app.ondeviceintelligence;
+
+/**
+  * @hide
+  */
+parcelable Feature;
diff --git a/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/Feature.java b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/Feature.java
new file mode 100644
index 0000000..bcc56073
--- /dev/null
+++ b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/Feature.java
@@ -0,0 +1,267 @@
+/*
+ * 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 android.app.ondeviceintelligence;
+
+import static android.app.ondeviceintelligence.flags.Flags.FLAG_ENABLE_ON_DEVICE_INTELLIGENCE;
+
+import android.annotation.FlaggedApi;
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.annotation.SystemApi;
+import android.os.Parcel;
+import android.os.Parcelable;
+import android.os.PersistableBundle;
+
+/**
+ * Represents a typical feature associated with on-device intelligence.
+ *
+ * @hide
+ */
+@SystemApi
+@FlaggedApi(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
+public final class Feature implements Parcelable {
+    private final int mId;
+    @Nullable
+    private final String mName;
+    @Nullable
+    private final String mModelName;
+    private final int mType;
+    private final int mVariant;
+    @NonNull
+    private final PersistableBundle mFeatureParams;
+
+    /* package-private */ Feature(
+            int id,
+            @Nullable String name,
+            @Nullable String modelName,
+            int type,
+            int variant,
+            @NonNull PersistableBundle featureParams) {
+        this.mId = id;
+        this.mName = name;
+        this.mModelName = modelName;
+        this.mType = type;
+        this.mVariant = variant;
+        this.mFeatureParams = featureParams;
+        com.android.internal.util.AnnotationValidations.validate(
+                NonNull.class, null, mFeatureParams);
+    }
+
+    /** Returns the unique and immutable identifier of this feature. */
+    public int getId() {
+        return mId;
+    }
+
+    /** Returns human-readable name of this feature. */
+    public @Nullable String getName() {
+        return mName;
+    }
+
+    /** Returns base model name of this feature. */
+    public @Nullable String getModelName() {
+        return mModelName;
+    }
+
+    /** Returns type identifier of this feature. */
+    public int getType() {
+        return mType;
+    }
+
+    /** Returns variant kind for this feature. */
+    public int getVariant() {
+        return mVariant;
+    }
+
+    public @NonNull PersistableBundle getFeatureParams() {
+        return mFeatureParams;
+    }
+
+    @Override
+    public String toString() {
+        return "Feature { " +
+                "id = " + mId + ", " +
+                "name = " + mName + ", " +
+                "modelName = " + mModelName + ", " +
+                "type = " + mType + ", " +
+                "variant = " + mVariant + ", " +
+                "featureParams = " + mFeatureParams +
+                " }";
+    }
+
+    @Override
+    public boolean equals(@Nullable Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        @SuppressWarnings("unchecked")
+        Feature that = (Feature) o;
+        //noinspection PointlessBooleanExpression
+        return true
+                && mId == that.mId
+                && java.util.Objects.equals(mName, that.mName)
+                && java.util.Objects.equals(mModelName, that.mModelName)
+                && mType == that.mType
+                && mVariant == that.mVariant
+                && java.util.Objects.equals(mFeatureParams, that.mFeatureParams);
+    }
+
+    @Override
+    public int hashCode() {
+        int _hash = 1;
+        _hash = 31 * _hash + mId;
+        _hash = 31 * _hash + java.util.Objects.hashCode(mName);
+        _hash = 31 * _hash + java.util.Objects.hashCode(mModelName);
+        _hash = 31 * _hash + mType;
+        _hash = 31 * _hash + mVariant;
+        _hash = 31 * _hash + java.util.Objects.hashCode(mFeatureParams);
+        return _hash;
+    }
+
+    @Override
+    public void writeToParcel(@NonNull Parcel dest, int flags) {
+        byte flg = 0;
+        if (mName != null) flg |= 0x2;
+        if (mModelName != null) flg |= 0x4;
+        dest.writeByte(flg);
+        dest.writeInt(mId);
+        if (mName != null) dest.writeString(mName);
+        if (mModelName != null) dest.writeString(mModelName);
+        dest.writeInt(mType);
+        dest.writeInt(mVariant);
+        dest.writeTypedObject(mFeatureParams, flags);
+    }
+
+    @Override
+    public int describeContents() {
+        return 0;
+    }
+
+    /** @hide */
+    @SuppressWarnings({"unchecked", "RedundantCast"})
+    /* package-private */ Feature(@NonNull Parcel in) {
+        byte flg = in.readByte();
+        int id = in.readInt();
+        String name = (flg & 0x2) == 0 ? null : in.readString();
+        String modelName = (flg & 0x4) == 0 ? null : in.readString();
+        int type = in.readInt();
+        int variant = in.readInt();
+        PersistableBundle featureParams = (PersistableBundle) in.readTypedObject(
+                PersistableBundle.CREATOR);
+
+        this.mId = id;
+        this.mName = name;
+        this.mModelName = modelName;
+        this.mType = type;
+        this.mVariant = variant;
+        this.mFeatureParams = featureParams;
+        com.android.internal.util.AnnotationValidations.validate(
+                NonNull.class, null, mFeatureParams);
+    }
+
+    public static final @NonNull Parcelable.Creator<Feature> CREATOR
+            = new Parcelable.Creator<Feature>() {
+        @Override
+        public Feature[] newArray(int size) {
+            return new Feature[size];
+        }
+
+        @Override
+        public Feature createFromParcel(@NonNull Parcel in) {
+            return new Feature(in);
+        }
+    };
+
+    /**
+     * A builder for {@link Feature}
+     */
+    @SuppressWarnings("WeakerAccess")
+    public static final class Builder {
+        private int mId;
+        private @Nullable String mName;
+        private @Nullable String mModelName;
+        private int mType;
+        private int mVariant;
+        private @NonNull PersistableBundle mFeatureParams;
+
+        private long mBuilderFieldsSet = 0L;
+
+        /**
+         * Provides a builder instance to create a feature for given id.
+         * @param id the unique identifier for the feature.
+         */
+        public Builder(int id) {
+            mId = id;
+            mFeatureParams = new PersistableBundle();
+        }
+
+        public @NonNull Builder setName(@NonNull String value) {
+            checkNotUsed();
+            mBuilderFieldsSet |= 0x2;
+            mName = value;
+            return this;
+        }
+
+        public @NonNull Builder setModelName(@NonNull String value) {
+            checkNotUsed();
+            mBuilderFieldsSet |= 0x4;
+            mModelName = value;
+            return this;
+        }
+
+        public @NonNull Builder setType(int value) {
+            checkNotUsed();
+            mBuilderFieldsSet |= 0x8;
+            mType = value;
+            return this;
+        }
+
+        public @NonNull Builder setVariant(int value) {
+            checkNotUsed();
+            mBuilderFieldsSet |= 0x10;
+            mVariant = value;
+            return this;
+        }
+
+        public @NonNull Builder setFeatureParams(@NonNull PersistableBundle value) {
+            checkNotUsed();
+            mBuilderFieldsSet |= 0x20;
+            mFeatureParams = value;
+            return this;
+        }
+
+        /** Builds the instance. This builder should not be touched after calling this! */
+        public @NonNull Feature build() {
+            checkNotUsed();
+            mBuilderFieldsSet |= 0x40; // Mark builder used
+
+            Feature o = new Feature(
+                    mId,
+                    mName,
+                    mModelName,
+                    mType,
+                    mVariant,
+                    mFeatureParams);
+            return o;
+        }
+
+        private void checkNotUsed() {
+            if ((mBuilderFieldsSet & 0x40) != 0) {
+                throw new IllegalStateException(
+                        "This Builder should not be reused. Use a new Builder instance instead");
+            }
+        }
+    }
+}
diff --git a/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/FeatureDetails.aidl b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/FeatureDetails.aidl
new file mode 100644
index 0000000..0589bf8
--- /dev/null
+++ b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/FeatureDetails.aidl
@@ -0,0 +1,22 @@
+/*
+ * 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 android.app.ondeviceintelligence;
+
+/**
+  * @hide
+  */
+parcelable FeatureDetails;
diff --git a/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/FeatureDetails.java b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/FeatureDetails.java
new file mode 100644
index 0000000..0ee0cc3
--- /dev/null
+++ b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/FeatureDetails.java
@@ -0,0 +1,178 @@
+/*
+ * 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 android.app.ondeviceintelligence;
+
+import static android.app.ondeviceintelligence.flags.Flags.FLAG_ENABLE_ON_DEVICE_INTELLIGENCE;
+
+import android.annotation.FlaggedApi;
+import android.annotation.IntDef;
+import android.annotation.NonNull;
+import android.annotation.SystemApi;
+import android.os.Parcelable;
+import android.os.PersistableBundle;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+import java.text.MessageFormat;
+
+/**
+ * Represents a status of a requested {@link Feature}.
+ *
+ * @hide
+ */
+@SystemApi
+@FlaggedApi(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
+public final class FeatureDetails implements Parcelable {
+    @Status
+    private final int mFeatureStatus;
+    @NonNull
+    private final PersistableBundle mFeatureDetailParams;
+
+    /** Invalid or unavailable {@code AiFeature}. */
+    public static final int FEATURE_STATUS_UNAVAILABLE = 0;
+
+    /** Feature can be downloaded on request. */
+    public static final int FEATURE_STATUS_DOWNLOADABLE = 1;
+
+    /** Feature is being downloaded. */
+    public static final int FEATURE_STATUS_DOWNLOADING = 2;
+
+    /** Feature is fully downloaded and ready to use. */
+    public static final int FEATURE_STATUS_AVAILABLE = 3;
+
+    /** Underlying service is unavailable and feature status cannot be fetched. */
+    public static final int FEATURE_STATUS_SERVICE_UNAVAILABLE = 4;
+
+    /**
+     * @hide
+     */
+    @IntDef(value = {
+            FEATURE_STATUS_UNAVAILABLE,
+            FEATURE_STATUS_DOWNLOADABLE,
+            FEATURE_STATUS_DOWNLOADING,
+            FEATURE_STATUS_AVAILABLE,
+            FEATURE_STATUS_SERVICE_UNAVAILABLE
+    })
+    @Target({ElementType.TYPE_USE, ElementType.METHOD, ElementType.PARAMETER, ElementType.FIELD})
+    @Retention(RetentionPolicy.SOURCE)
+    public @interface Status {
+    }
+
+    public FeatureDetails(
+            @Status int featureStatus,
+            @NonNull PersistableBundle featureDetailParams) {
+        this.mFeatureStatus = featureStatus;
+        com.android.internal.util.AnnotationValidations.validate(
+                Status.class, null, mFeatureStatus);
+        this.mFeatureDetailParams = featureDetailParams;
+        com.android.internal.util.AnnotationValidations.validate(
+                NonNull.class, null, mFeatureDetailParams);
+    }
+
+    public FeatureDetails(
+            @Status int featureStatus) {
+        this.mFeatureStatus = featureStatus;
+        com.android.internal.util.AnnotationValidations.validate(
+                Status.class, null, mFeatureStatus);
+        this.mFeatureDetailParams = new PersistableBundle();
+    }
+
+
+    /**
+     * Returns an integer value associated with the feature status.
+     */
+    public @Status int getFeatureStatus() {
+        return mFeatureStatus;
+    }
+
+
+    /**
+     * Returns a persistable bundle contain any additional status related params.
+     */
+    public @NonNull PersistableBundle getFeatureDetailParams() {
+        return mFeatureDetailParams;
+    }
+
+    @Override
+    public String toString() {
+        return MessageFormat.format("FeatureDetails '{' status = {0}, "
+                        + "persistableBundle = {1} '}'",
+                mFeatureStatus,
+                mFeatureDetailParams);
+    }
+
+    @Override
+    public boolean equals(@android.annotation.Nullable Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        @SuppressWarnings("unchecked")
+        FeatureDetails that = (FeatureDetails) o;
+        return mFeatureStatus == that.mFeatureStatus
+                && java.util.Objects.equals(mFeatureDetailParams, that.mFeatureDetailParams);
+    }
+
+    @Override
+    public int hashCode() {
+        int _hash = 1;
+        _hash = 31 * _hash + mFeatureStatus;
+        _hash = 31 * _hash + java.util.Objects.hashCode(mFeatureDetailParams);
+        return _hash;
+    }
+
+    @Override
+    public void writeToParcel(@NonNull android.os.Parcel dest, int flags) {
+        dest.writeInt(mFeatureStatus);
+        dest.writeTypedObject(mFeatureDetailParams, flags);
+    }
+
+    @Override
+    public int describeContents() {
+        return 0;
+    }
+
+    /** @hide */
+    @SuppressWarnings({"unchecked", "RedundantCast"})
+    FeatureDetails(@NonNull android.os.Parcel in) {
+        int status = in.readInt();
+        PersistableBundle persistableBundle = (PersistableBundle) in.readTypedObject(
+                PersistableBundle.CREATOR);
+
+        this.mFeatureStatus = status;
+        com.android.internal.util.AnnotationValidations.validate(
+                Status.class, null, mFeatureStatus);
+        this.mFeatureDetailParams = persistableBundle;
+        com.android.internal.util.AnnotationValidations.validate(
+                NonNull.class, null, mFeatureDetailParams);
+    }
+
+
+    public static final @NonNull Parcelable.Creator<FeatureDetails> CREATOR =
+            new Parcelable.Creator<>() {
+                @Override
+                public FeatureDetails[] newArray(int size) {
+                    return new FeatureDetails[size];
+                }
+
+                @Override
+                public FeatureDetails createFromParcel(@NonNull android.os.Parcel in) {
+                    return new FeatureDetails(in);
+                }
+            };
+
+}
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/ICancellationSignal.aidl b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/ICancellationSignal.aidl
similarity index 100%
copy from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/ICancellationSignal.aidl
copy to packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/ICancellationSignal.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IDownloadCallback.aidl b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/IDownloadCallback.aidl
similarity index 100%
copy from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IDownloadCallback.aidl
copy to packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/IDownloadCallback.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IFeatureCallback.aidl b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/IFeatureCallback.aidl
similarity index 100%
copy from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IFeatureCallback.aidl
copy to packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/IFeatureCallback.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IFeatureDetailsCallback.aidl b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/IFeatureDetailsCallback.aidl
similarity index 100%
copy from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IFeatureDetailsCallback.aidl
copy to packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/IFeatureDetailsCallback.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IListFeaturesCallback.aidl b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/IListFeaturesCallback.aidl
similarity index 100%
copy from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IListFeaturesCallback.aidl
copy to packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/IListFeaturesCallback.aidl
diff --git a/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/IOnDeviceIntelligenceManager.aidl b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/IOnDeviceIntelligenceManager.aidl
new file mode 100644
index 0000000..1977a39
--- /dev/null
+++ b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/IOnDeviceIntelligenceManager.aidl
@@ -0,0 +1,79 @@
+/*
+ * Copyright (C) 2023 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 android.app.ondeviceintelligence;
+
+ import com.android.internal.infra.AndroidFuture;
+ import android.os.ICancellationSignal;
+ import android.os.ParcelFileDescriptor;
+ import android.os.PersistableBundle;
+ import android.os.RemoteCallback;
+ import android.os.Bundle;
+ import android.app.ondeviceintelligence.Feature;
+ import android.app.ondeviceintelligence.FeatureDetails;
+ import android.app.ondeviceintelligence.InferenceInfo;
+ import java.util.List;
+ import android.app.ondeviceintelligence.IDownloadCallback;
+ import android.app.ondeviceintelligence.IListFeaturesCallback;
+ import android.app.ondeviceintelligence.IFeatureCallback;
+ import android.app.ondeviceintelligence.IFeatureDetailsCallback;
+ import android.app.ondeviceintelligence.IResponseCallback;
+ import android.app.ondeviceintelligence.IStreamingResponseCallback;
+ import android.app.ondeviceintelligence.IProcessingSignal;
+ import android.app.ondeviceintelligence.ITokenInfoCallback;
+
+
+ /**
+  * Interface for a OnDeviceIntelligenceManager for managing OnDeviceIntelligenceService and OnDeviceSandboxedInferenceService.
+  *
+  * @hide
+  */
+interface IOnDeviceIntelligenceManager {
+      @JavaPassthrough(annotation="@android.annotation.RequiresPermission(android.Manifest.permission.USE_ON_DEVICE_INTELLIGENCE)")
+      void getVersion(in RemoteCallback remoteCallback) = 1;
+
+      @JavaPassthrough(annotation="@android.annotation.RequiresPermission(android.Manifest.permission.USE_ON_DEVICE_INTELLIGENCE)")
+      void getFeature(in int featureId, in IFeatureCallback remoteCallback) = 2;
+
+      @JavaPassthrough(annotation="@android.annotation.RequiresPermission(android.Manifest.permission.USE_ON_DEVICE_INTELLIGENCE)")
+      void listFeatures(in IListFeaturesCallback listFeaturesCallback) = 3;
+
+      @JavaPassthrough(annotation="@android.annotation.RequiresPermission(android.Manifest.permission.USE_ON_DEVICE_INTELLIGENCE)")
+      void getFeatureDetails(in Feature feature, in IFeatureDetailsCallback featureDetailsCallback) = 4;
+
+      @JavaPassthrough(annotation="@android.annotation.RequiresPermission(android.Manifest.permission.USE_ON_DEVICE_INTELLIGENCE)")
+      void requestFeatureDownload(in Feature feature, in  AndroidFuture cancellationSignalFuture, in IDownloadCallback callback) = 5;
+
+      @JavaPassthrough(annotation="@android.annotation.RequiresPermission(android.Manifest.permission.USE_ON_DEVICE_INTELLIGENCE)")
+      void requestTokenInfo(in Feature feature, in Bundle requestBundle, in  AndroidFuture cancellationSignalFuture,
+                                                        in ITokenInfoCallback tokenInfocallback) = 6;
+
+      @JavaPassthrough(annotation="@android.annotation.RequiresPermission(android.Manifest.permission.USE_ON_DEVICE_INTELLIGENCE)")
+      void processRequest(in Feature feature, in Bundle requestBundle, int requestType,
+                                                in  AndroidFuture cancellationSignalFuture,
+                                                in AndroidFuture processingSignalFuture,
+                                                in IResponseCallback responseCallback) = 7;
+
+      @JavaPassthrough(annotation="@android.annotation.RequiresPermission(android.Manifest.permission.USE_ON_DEVICE_INTELLIGENCE)")
+      void processRequestStreaming(in Feature feature,
+                    in Bundle requestBundle, int requestType, in  AndroidFuture cancellationSignalFuture,
+                    in  AndroidFuture processingSignalFuture,
+                    in IStreamingResponseCallback streamingCallback) = 8;
+
+      String getRemoteServicePackageName() = 9;
+
+      List<InferenceInfo> getLatestInferenceInfo(long startTimeEpochMillis) = 10;
+ }
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IProcessingSignal.aidl b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/IProcessingSignal.aidl
similarity index 100%
copy from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IProcessingSignal.aidl
copy to packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/IProcessingSignal.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IRemoteCallback.aidl b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/IRemoteCallback.aidl
similarity index 100%
copy from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IRemoteCallback.aidl
copy to packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/IRemoteCallback.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IResponseCallback.aidl b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/IResponseCallback.aidl
similarity index 100%
copy from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IResponseCallback.aidl
copy to packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/IResponseCallback.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IStreamingResponseCallback.aidl b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/IStreamingResponseCallback.aidl
similarity index 100%
copy from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IStreamingResponseCallback.aidl
copy to packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/IStreamingResponseCallback.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/ITokenInfoCallback.aidl b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/ITokenInfoCallback.aidl
similarity index 100%
copy from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/ITokenInfoCallback.aidl
copy to packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/ITokenInfoCallback.aidl
diff --git a/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/InferenceInfo.aidl b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/InferenceInfo.aidl
new file mode 100644
index 0000000..6d70fc4
--- /dev/null
+++ b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/InferenceInfo.aidl
@@ -0,0 +1,22 @@
+/*
+ * 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 android.app.ondeviceintelligence;
+
+/**
+  * @hide
+  */
+parcelable InferenceInfo;
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/InferenceInfo.java b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/InferenceInfo.java
similarity index 100%
copy from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/InferenceInfo.java
copy to packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/InferenceInfo.java
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/OnDeviceIntelligenceException.java b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/OnDeviceIntelligenceException.java
similarity index 100%
copy from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/OnDeviceIntelligenceException.java
copy to packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/OnDeviceIntelligenceException.java
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/OnDeviceIntelligenceFrameworkInitializer.java b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/OnDeviceIntelligenceFrameworkInitializer.java
similarity index 100%
copy from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/OnDeviceIntelligenceFrameworkInitializer.java
copy to packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/OnDeviceIntelligenceFrameworkInitializer.java
diff --git a/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/OnDeviceIntelligenceManager.java b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/OnDeviceIntelligenceManager.java
new file mode 100644
index 0000000..78cf1d7
--- /dev/null
+++ b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/OnDeviceIntelligenceManager.java
@@ -0,0 +1,642 @@
+/*
+ * 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 android.app.ondeviceintelligence;
+
+import static android.app.ondeviceintelligence.flags.Flags.FLAG_ENABLE_ON_DEVICE_INTELLIGENCE;
+import static android.app.ondeviceintelligence.flags.Flags.FLAG_ENABLE_ON_DEVICE_INTELLIGENCE_MODULE;
+
+import android.Manifest;
+import android.annotation.CallbackExecutor;
+import android.annotation.CurrentTimeMillisLong;
+import android.annotation.FlaggedApi;
+import android.annotation.IntDef;
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.annotation.RequiresPermission;
+import android.annotation.SystemApi;
+import android.annotation.SystemService;
+import android.content.Context;
+import android.graphics.Bitmap;
+import android.os.Binder;
+import android.os.Bundle;
+import android.os.CancellationSignal;
+import android.os.IBinder;
+import android.os.ICancellationSignal;
+import android.os.OutcomeReceiver;
+import android.os.PersistableBundle;
+import android.os.RemoteCallback;
+import android.os.RemoteException;
+import android.system.OsConstants;
+import android.util.Log;
+
+import com.android.internal.infra.AndroidFuture;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+import java.util.List;
+import java.util.concurrent.Executor;
+import java.util.function.LongConsumer;
+
+/**
+ * Allows granted apps to manage on-device intelligence service configured on the device. Typical
+ * calling pattern will be to query and setup a required feature before proceeding to request
+ * processing.
+ *
+ * The contracts in this Manager class are designed to be open-ended in general, to allow
+ * interoperability. Therefore, it is recommended that implementations of this system-service
+ * expose this API to the clients via a separate sdk or library which has more defined contract.
+ *
+ * @hide
+ */
+@SystemApi
+@SystemService(Context.ON_DEVICE_INTELLIGENCE_SERVICE)
+@FlaggedApi(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
+public final class OnDeviceIntelligenceManager {
+    /**
+     * @hide
+     */
+    public static final String API_VERSION_BUNDLE_KEY = "ApiVersionBundleKey";
+
+    /**
+     * @hide
+     */
+    public static final String AUGMENT_REQUEST_CONTENT_BUNDLE_KEY =
+            "AugmentRequestContentBundleKey";
+
+    private static final String TAG = "OnDeviceIntelligence";
+    private final Context mContext;
+    private final IOnDeviceIntelligenceManager mService;
+
+    /**
+     * @hide
+     */
+    public OnDeviceIntelligenceManager(Context context, IOnDeviceIntelligenceManager service) {
+        mContext = context;
+        mService = service;
+    }
+
+    /**
+     * Asynchronously get the version of the underlying remote implementation.
+     *
+     * @param versionConsumer  consumer to populate the version of remote implementation.
+     * @param callbackExecutor executor to run the callback on.
+     */
+    @RequiresPermission(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE)
+    public void getVersion(
+            @NonNull @CallbackExecutor Executor callbackExecutor,
+            @NonNull LongConsumer versionConsumer) {
+        try {
+            RemoteCallback callback = new RemoteCallback(result -> {
+                if (result == null) {
+                    Binder.withCleanCallingIdentity(
+                            () -> callbackExecutor.execute(() -> versionConsumer.accept(0)));
+                }
+                long version = result.getLong(API_VERSION_BUNDLE_KEY);
+                Binder.withCleanCallingIdentity(
+                        () -> callbackExecutor.execute(() -> versionConsumer.accept(version)));
+            });
+            mService.getVersion(callback);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+
+    /**
+     * Get package name configured for providing the remote implementation for this system service.
+     */
+    @Nullable
+    @RequiresPermission(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE)
+    public String getRemoteServicePackageName() {
+        String result;
+        try {
+            result = mService.getRemoteServicePackageName();
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+        return result;
+    }
+
+    /**
+     * Asynchronously get feature for a given id.
+     *
+     * @param featureId        the identifier pointing to the feature.
+     * @param featureReceiver  callback to populate the feature object for given identifier.
+     * @param callbackExecutor executor to run the callback on.
+     */
+    @RequiresPermission(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE)
+    public void getFeature(
+            int featureId,
+            @NonNull @CallbackExecutor Executor callbackExecutor,
+            @NonNull OutcomeReceiver<Feature, OnDeviceIntelligenceException> featureReceiver) {
+        try {
+            IFeatureCallback callback =
+                    new IFeatureCallback.Stub() {
+                        @Override
+                        public void onSuccess(Feature result) {
+                            Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                                    () -> featureReceiver.onResult(result)));
+                        }
+
+                        @Override
+                        public void onFailure(int errorCode, String errorMessage,
+                                PersistableBundle errorParams) {
+                            Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                                    () -> featureReceiver.onError(
+                                            new OnDeviceIntelligenceException(
+                                                    errorCode, errorMessage, errorParams))));
+                        }
+                    };
+            mService.getFeature(featureId, callback);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+    /**
+     * Asynchronously get a list of features that are supported for the caller.
+     *
+     * @param featureListReceiver callback to populate the list of features.
+     * @param callbackExecutor    executor to run the callback on.
+     */
+    @RequiresPermission(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE)
+    public void listFeatures(
+            @NonNull @CallbackExecutor Executor callbackExecutor,
+            @NonNull OutcomeReceiver<List<Feature>, OnDeviceIntelligenceException> featureListReceiver) {
+        try {
+            IListFeaturesCallback callback =
+                    new IListFeaturesCallback.Stub() {
+                        @Override
+                        public void onSuccess(List<Feature> result) {
+                            Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                                    () -> featureListReceiver.onResult(result)));
+                        }
+
+                        @Override
+                        public void onFailure(int errorCode, String errorMessage,
+                                PersistableBundle errorParams) {
+                            Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                                    () -> featureListReceiver.onError(
+                                            new OnDeviceIntelligenceException(
+                                                    errorCode, errorMessage, errorParams))));
+                        }
+                    };
+            mService.listFeatures(callback);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+    /**
+     * This method should be used to fetch details about a feature which need some additional
+     * computation, that can be inefficient to return in all calls to {@link #getFeature}. Callers
+     * and implementation can utilize the {@link Feature#getFeatureParams()} to pass hint on what
+     * details are expected by the caller.
+     *
+     * @param feature                the feature to check status for.
+     * @param featureDetailsReceiver callback to populate the feature details to.
+     * @param callbackExecutor       executor to run the callback on.
+     */
+    @RequiresPermission(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE)
+    public void getFeatureDetails(@NonNull Feature feature,
+            @NonNull @CallbackExecutor Executor callbackExecutor,
+            @NonNull OutcomeReceiver<FeatureDetails, OnDeviceIntelligenceException> featureDetailsReceiver) {
+        try {
+            IFeatureDetailsCallback callback = new IFeatureDetailsCallback.Stub() {
+
+                @Override
+                public void onSuccess(FeatureDetails result) {
+                    Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                            () -> featureDetailsReceiver.onResult(result)));
+                }
+
+                @Override
+                public void onFailure(int errorCode, String errorMessage,
+                        PersistableBundle errorParams) {
+                    Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                            () -> featureDetailsReceiver.onError(
+                                    new OnDeviceIntelligenceException(errorCode,
+                                            errorMessage, errorParams))));
+                }
+            };
+            mService.getFeatureDetails(feature, callback);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+    /**
+     * This method handles downloading all model and config files required to process requests
+     * sent against a given feature. The caller can listen to updates on the download status via
+     * the callback.
+     *
+     * Note: If a feature was already requested for downloaded previously, the onDownloadFailed
+     * callback would be invoked with {@link DownloadCallback#DOWNLOAD_FAILURE_STATUS_DOWNLOADING}.
+     * In such cases, clients should query the feature status via {@link #getFeatureDetails} to
+     * check on the feature's download status.
+     *
+     * @param feature            feature to request download for.
+     * @param callback           callback to populate updates about download status.
+     * @param cancellationSignal signal to invoke cancellation on the operation in the remote
+     *                           implementation.
+     * @param callbackExecutor   executor to run the callback on.
+     */
+    @RequiresPermission(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE)
+    public void requestFeatureDownload(@NonNull Feature feature,
+            @Nullable CancellationSignal cancellationSignal,
+            @NonNull @CallbackExecutor Executor callbackExecutor,
+            @NonNull DownloadCallback callback) {
+        try {
+            IDownloadCallback downloadCallback = new IDownloadCallback.Stub() {
+
+                @Override
+                public void onDownloadStarted(long bytesToDownload) {
+                    Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                            () -> callback.onDownloadStarted(bytesToDownload)));
+                }
+
+                @Override
+                public void onDownloadProgress(long bytesDownloaded) {
+                    Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                            () -> callback.onDownloadProgress(bytesDownloaded)));
+                }
+
+                @Override
+                public void onDownloadFailed(int failureStatus, String errorMessage,
+                        PersistableBundle errorParams) {
+                    Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                            () -> callback.onDownloadFailed(failureStatus, errorMessage,
+                                    errorParams)));
+                }
+
+                @Override
+                public void onDownloadCompleted(PersistableBundle downloadParams) {
+                    Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                            () -> callback.onDownloadCompleted(downloadParams)));
+                }
+            };
+
+            mService.requestFeatureDownload(feature,
+                    configureRemoteCancellationFuture(cancellationSignal, callbackExecutor),
+                    downloadCallback);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+
+    /**
+     * The methods computes the token related information for a given request payload using the
+     * provided {@link Feature}.
+     *
+     * @param feature            feature associated with the request.
+     * @param request            request and associated params represented by the Bundle
+     *                           data.
+     * @param outcomeReceiver    callback to populate the token info or exception in case of
+     *                           failure.
+     * @param cancellationSignal signal to invoke cancellation on the operation in the remote
+     *                           implementation.
+     * @param callbackExecutor   executor to run the callback on.
+     */
+    @RequiresPermission(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE)
+    public void requestTokenInfo(@NonNull Feature feature, @NonNull @InferenceParams Bundle request,
+            @Nullable CancellationSignal cancellationSignal,
+            @NonNull @CallbackExecutor Executor callbackExecutor,
+            @NonNull OutcomeReceiver<TokenInfo,
+                    OnDeviceIntelligenceException> outcomeReceiver) {
+        try {
+            ITokenInfoCallback callback = new ITokenInfoCallback.Stub() {
+                @Override
+                public void onSuccess(TokenInfo tokenInfo) {
+                    Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                            () -> outcomeReceiver.onResult(tokenInfo)));
+                }
+
+                @Override
+                public void onFailure(int errorCode, String errorMessage,
+                        PersistableBundle errorParams) {
+                    Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                            () -> outcomeReceiver.onError(
+                                    new OnDeviceIntelligenceException(
+                                            errorCode, errorMessage, errorParams))));
+                }
+            };
+
+            mService.requestTokenInfo(feature, request,
+                    configureRemoteCancellationFuture(cancellationSignal, callbackExecutor),
+                    callback);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+
+    /**
+     * Asynchronously Process a request based on the associated params, to populate a
+     * response in
+     * {@link OutcomeReceiver#onResult} callback or failure callback status code if there
+     * was a
+     * failure.
+     *
+     * @param feature            feature associated with the request.
+     * @param request            request and associated params represented by the Bundle
+     *                           data.
+     * @param requestType        type of request being sent for processing the content.
+     * @param cancellationSignal signal to invoke cancellation.
+     * @param processingSignal   signal to send custom signals in the
+     *                           remote implementation.
+     * @param callbackExecutor   executor to run the callback on.
+     * @param processingCallback callback to populate the response content and
+     *                           associated params.
+     */
+    @RequiresPermission(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE)
+    public void processRequest(@NonNull Feature feature, @NonNull @InferenceParams Bundle request,
+            @RequestType int requestType,
+            @Nullable CancellationSignal cancellationSignal,
+            @Nullable ProcessingSignal processingSignal,
+            @NonNull @CallbackExecutor Executor callbackExecutor,
+            @NonNull ProcessingCallback processingCallback) {
+        try {
+            IResponseCallback callback = new IResponseCallback.Stub() {
+                @Override
+                public void onSuccess(@InferenceParams Bundle result) {
+                    Binder.withCleanCallingIdentity(() -> {
+                        callbackExecutor.execute(() -> processingCallback.onResult(result));
+                    });
+                }
+
+                @Override
+                public void onFailure(int errorCode, String errorMessage,
+                        PersistableBundle errorParams) {
+                    Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                            () -> processingCallback.onError(
+                                    new OnDeviceIntelligenceException(
+                                            errorCode, errorMessage, errorParams))));
+                }
+
+                @Override
+                public void onDataAugmentRequest(@NonNull @InferenceParams Bundle request,
+                        @NonNull RemoteCallback contentCallback) {
+                    Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                            () -> processingCallback.onDataAugmentRequest(request, result -> {
+                                Bundle bundle = new Bundle();
+                                bundle.putParcelable(AUGMENT_REQUEST_CONTENT_BUNDLE_KEY, result);
+                                callbackExecutor.execute(() -> contentCallback.sendResult(bundle));
+                            })));
+                }
+            };
+
+
+            mService.processRequest(feature, request, requestType,
+                    configureRemoteCancellationFuture(cancellationSignal, callbackExecutor),
+                    configureRemoteProcessingSignalFuture(processingSignal, callbackExecutor),
+                    callback);
+
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+    /**
+     * Variation of {@link #processRequest} that asynchronously processes a request in a
+     * streaming
+     * fashion, where new content is pushed to caller in chunks via the
+     * {@link StreamingProcessingCallback#onPartialResult}. After the streaming is complete,
+     * the service should call {@link StreamingProcessingCallback#onResult} and can optionally
+     * populate the complete the full response {@link Bundle} as part of the callback in cases
+     * when the final response contains an enhanced aggregation of the contents already
+     * streamed.
+     *
+     * @param feature                     feature associated with the request.
+     * @param request                     request and associated params represented by the Bundle
+     *                                    data.
+     * @param requestType                 type of request being sent for processing the content.
+     * @param cancellationSignal          signal to invoke cancellation.
+     * @param processingSignal            signal to send custom signals in the
+     *                                    remote implementation.
+     * @param streamingProcessingCallback streaming callback to populate the response content and
+     *                                    associated params.
+     * @param callbackExecutor            executor to run the callback on.
+     */
+    @RequiresPermission(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE)
+    public void processRequestStreaming(@NonNull Feature feature,
+            @NonNull @InferenceParams Bundle request,
+            @RequestType int requestType,
+            @Nullable CancellationSignal cancellationSignal,
+            @Nullable ProcessingSignal processingSignal,
+            @NonNull @CallbackExecutor Executor callbackExecutor,
+            @NonNull StreamingProcessingCallback streamingProcessingCallback) {
+        try {
+            IStreamingResponseCallback callback = new IStreamingResponseCallback.Stub() {
+                @Override
+                public void onNewContent(@InferenceParams Bundle result) {
+                    Binder.withCleanCallingIdentity(() -> {
+                        callbackExecutor.execute(
+                                () -> streamingProcessingCallback.onPartialResult(result));
+                    });
+                }
+
+                @Override
+                public void onSuccess(@InferenceParams Bundle result) {
+                    Binder.withCleanCallingIdentity(() -> {
+                        callbackExecutor.execute(
+                                () -> streamingProcessingCallback.onResult(result));
+                    });
+                }
+
+                @Override
+                public void onFailure(int errorCode, String errorMessage,
+                        PersistableBundle errorParams) {
+                    Binder.withCleanCallingIdentity(() -> {
+                        callbackExecutor.execute(
+                                () -> streamingProcessingCallback.onError(
+                                        new OnDeviceIntelligenceException(
+                                                errorCode, errorMessage, errorParams)));
+                    });
+                }
+
+
+                @Override
+                public void onDataAugmentRequest(@NonNull @InferenceParams Bundle content,
+                        @NonNull RemoteCallback contentCallback) {
+                    Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                            () -> streamingProcessingCallback.onDataAugmentRequest(content,
+                                    contentResponse -> {
+                                        Bundle bundle = new Bundle();
+                                        bundle.putParcelable(AUGMENT_REQUEST_CONTENT_BUNDLE_KEY,
+                                                contentResponse);
+                                        callbackExecutor.execute(
+                                                () -> contentCallback.sendResult(bundle));
+                                    })));
+                }
+            };
+
+            mService.processRequestStreaming(
+                    feature, request, requestType,
+                    configureRemoteCancellationFuture(cancellationSignal, callbackExecutor),
+                    configureRemoteProcessingSignalFuture(processingSignal, callbackExecutor),
+                    callback);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+    /**
+     * This is primarily intended to be used to attribute/blame on-device intelligence power usage,
+     * via the configured remote implementation, to its actual caller.
+     *
+     * @param startTimeEpochMillis epoch millis used to filter the InferenceInfo events.
+     * @return InferenceInfo events since the passed in startTimeEpochMillis.
+     */
+    @RequiresPermission(Manifest.permission.DUMP)
+    @FlaggedApi(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE_MODULE)
+    public @NonNull List<InferenceInfo> getLatestInferenceInfo(@CurrentTimeMillisLong long startTimeEpochMillis) {
+        try {
+            return mService.getLatestInferenceInfo(startTimeEpochMillis);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+
+    /** Request inference with provided Bundle and Params. */
+    public static final int REQUEST_TYPE_INFERENCE = 0;
+
+    /**
+     * Prepares the remote implementation environment for e.g.loading inference runtime etc
+     * .which
+     * are time consuming beforehand to remove overhead and allow quick processing of requests
+     * thereof.
+     */
+    public static final int REQUEST_TYPE_PREPARE = 1;
+
+    /** Request Embeddings of the passed-in Bundle. */
+    public static final int REQUEST_TYPE_EMBEDDINGS = 2;
+
+    /**
+     * @hide
+     */
+    @IntDef(value = {
+            REQUEST_TYPE_INFERENCE,
+            REQUEST_TYPE_PREPARE,
+            REQUEST_TYPE_EMBEDDINGS
+    })
+    @Target({ElementType.TYPE_USE, ElementType.METHOD, ElementType.PARAMETER,
+            ElementType.FIELD})
+    @Retention(RetentionPolicy.SOURCE)
+    public @interface RequestType {
+    }
+
+    /**
+     * {@link Bundle}s annotated with this type will be validated that they are in-effect read-only
+     * when passed via Binder IPC. Following restrictions apply :
+     * <ul>
+     * <li> {@link PersistableBundle}s are allowed.</li>
+     * <li> Any primitive types or their collections can be added as usual.</li>
+     * <li>IBinder objects should *not* be added.</li>
+     * <li>Parcelable data which has no active-objects, should be added as
+     * {@link Bundle#putByteArray}</li>
+     * <li>Parcelables have active-objects, only following types will be allowed</li>
+     * <ul>
+     *  <li>{@link android.os.ParcelFileDescriptor} opened in
+     *  {@link android.os.ParcelFileDescriptor#MODE_READ_ONLY}</li>
+     * </ul>
+     * </ul>
+     *
+     * In all other scenarios the system-server might throw a
+     * {@link android.os.BadParcelableException} if the Bundle validation fails.
+     *
+     * @hide
+     */
+    @Target({ElementType.PARAMETER, ElementType.FIELD})
+    public @interface StateParams {
+    }
+
+    /**
+     * This is an extension of {@link StateParams} but for purpose of inference few other types are
+     * also allowed as read-only, as listed below.
+     *
+     * <li>{@link Bitmap} set as immutable.</li>
+     * <li>{@link android.database.CursorWindow}</li>
+     * <li>{@link android.os.SharedMemory} set to {@link OsConstants#PROT_READ}</li>
+     * </ul>
+     * </ul>
+     *
+     * In all other scenarios the system-server might throw a
+     * {@link android.os.BadParcelableException} if the Bundle validation fails.
+     *
+     * @hide
+     */
+    @Target({ElementType.PARAMETER, ElementType.FIELD, ElementType.TYPE_USE})
+    public @interface InferenceParams {
+    }
+
+    /**
+     * This is an extension of {@link StateParams} with the exception that it allows writing
+     * {@link Bitmap} as part of the response.
+     *
+     * In all other scenarios the system-server might throw a
+     * {@link android.os.BadParcelableException} if the Bundle validation fails.
+     *
+     * @hide
+     */
+    @Target({ElementType.PARAMETER, ElementType.FIELD})
+    public @interface ResponseParams {
+    }
+
+    @Nullable
+    private static AndroidFuture<IBinder> configureRemoteCancellationFuture(
+            @Nullable CancellationSignal cancellationSignal,
+            @NonNull Executor callbackExecutor) {
+        if (cancellationSignal == null) {
+            return null;
+        }
+        AndroidFuture<IBinder> cancellationFuture = new AndroidFuture<>();
+        cancellationFuture.whenCompleteAsync(
+                (cancellationTransport, error) -> {
+                    if (error != null || cancellationTransport == null) {
+                        Log.e(TAG, "Unable to receive the remote cancellation signal.", error);
+                    } else {
+                        cancellationSignal.setRemote(
+                                ICancellationSignal.Stub.asInterface(cancellationTransport));
+                    }
+                }, callbackExecutor);
+        return cancellationFuture;
+    }
+
+    @Nullable
+    private static AndroidFuture<IBinder> configureRemoteProcessingSignalFuture(
+            ProcessingSignal processingSignal, Executor executor) {
+        if (processingSignal == null) {
+            return null;
+        }
+        AndroidFuture<IBinder> processingSignalFuture = new AndroidFuture<>();
+        processingSignalFuture.whenCompleteAsync(
+                (transport, error) -> {
+                    if (error != null || transport == null) {
+                        Log.e(TAG, "Unable to receive the remote processing signal.", error);
+                    } else {
+                        processingSignal.setRemote(IProcessingSignal.Stub.asInterface(transport));
+                    }
+                }, executor);
+        return processingSignalFuture;
+    }
+
+
+}
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/ProcessingCallback.java b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/ProcessingCallback.java
similarity index 100%
copy from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/ProcessingCallback.java
copy to packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/ProcessingCallback.java
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/ProcessingSignal.java b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/ProcessingSignal.java
similarity index 100%
copy from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/ProcessingSignal.java
copy to packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/ProcessingSignal.java
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/StreamingProcessingCallback.java b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/StreamingProcessingCallback.java
similarity index 100%
copy from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/StreamingProcessingCallback.java
copy to packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/StreamingProcessingCallback.java
diff --git a/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/TokenInfo.aidl b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/TokenInfo.aidl
new file mode 100644
index 0000000..2c19c1e
--- /dev/null
+++ b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/TokenInfo.aidl
@@ -0,0 +1,22 @@
+/*
+ * 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 android.app.ondeviceintelligence;
+
+/**
+  * @hide
+  */
+parcelable TokenInfo;
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/TokenInfo.java b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/TokenInfo.java
similarity index 100%
copy from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/TokenInfo.java
copy to packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/TokenInfo.java
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/utils/BinderUtils.java b/packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/utils/BinderUtils.java
similarity index 100%
copy from packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/utils/BinderUtils.java
copy to packages/NeuralNetworks/framework/platform/java/android/app/ondeviceintelligence/utils/BinderUtils.java
diff --git a/packages/NeuralNetworks/framework/platform/java/android/service/ondeviceintelligence/IOnDeviceIntelligenceService.aidl b/packages/NeuralNetworks/framework/platform/java/android/service/ondeviceintelligence/IOnDeviceIntelligenceService.aidl
new file mode 100644
index 0000000..45c4350
--- /dev/null
+++ b/packages/NeuralNetworks/framework/platform/java/android/service/ondeviceintelligence/IOnDeviceIntelligenceService.aidl
@@ -0,0 +1,51 @@
+/*
+ * Copyright (C) 2022 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 android.service.ondeviceintelligence;
+
+import android.os.PersistableBundle;
+import android.os.ParcelFileDescriptor;
+import android.os.ICancellationSignal;
+import android.os.RemoteCallback;
+import android.app.ondeviceintelligence.IDownloadCallback;
+import android.app.ondeviceintelligence.Feature;
+import android.app.ondeviceintelligence.IFeatureCallback;
+import android.app.ondeviceintelligence.IListFeaturesCallback;
+import android.app.ondeviceintelligence.IFeatureDetailsCallback;
+import com.android.internal.infra.AndroidFuture;
+import android.service.ondeviceintelligence.IRemoteProcessingService;
+
+
+/**
+ * Interface for a concrete implementation to provide on device intelligence services.
+ *
+ * @hide
+ */
+oneway interface IOnDeviceIntelligenceService {
+    void getVersion(in RemoteCallback remoteCallback);
+    void getFeature(int callerUid, int featureId, in IFeatureCallback featureCallback);
+    void listFeatures(int callerUid, in IListFeaturesCallback listFeaturesCallback);
+    void getFeatureDetails(int callerUid, in Feature feature, in IFeatureDetailsCallback featureDetailsCallback);
+    void getReadOnlyFileDescriptor(in String fileName, in AndroidFuture<ParcelFileDescriptor> future);
+    void getReadOnlyFeatureFileDescriptorMap(in Feature feature, in RemoteCallback remoteCallback);
+    void requestFeatureDownload(int callerUid, in Feature feature,
+                                in AndroidFuture cancellationSignal,
+                                in IDownloadCallback downloadCallback);
+    void registerRemoteServices(in IRemoteProcessingService remoteProcessingService);
+    void notifyInferenceServiceConnected();
+    void notifyInferenceServiceDisconnected();
+    void ready();
+}
\ No newline at end of file
diff --git a/packages/NeuralNetworks/framework/platform/java/android/service/ondeviceintelligence/IOnDeviceSandboxedInferenceService.aidl b/packages/NeuralNetworks/framework/platform/java/android/service/ondeviceintelligence/IOnDeviceSandboxedInferenceService.aidl
new file mode 100644
index 0000000..1af3b0f
--- /dev/null
+++ b/packages/NeuralNetworks/framework/platform/java/android/service/ondeviceintelligence/IOnDeviceSandboxedInferenceService.aidl
@@ -0,0 +1,53 @@
+/*
+ * Copyright (C) 2022 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 android.service.ondeviceintelligence;
+
+import android.app.ondeviceintelligence.IStreamingResponseCallback;
+import android.app.ondeviceintelligence.IResponseCallback;
+import android.app.ondeviceintelligence.ITokenInfoCallback;
+import android.app.ondeviceintelligence.IProcessingSignal;
+import android.app.ondeviceintelligence.Feature;
+import android.os.IRemoteCallback;
+import android.os.ICancellationSignal;
+import android.os.PersistableBundle;
+import android.os.Bundle;
+import com.android.internal.infra.AndroidFuture;
+import android.service.ondeviceintelligence.IRemoteStorageService;
+import android.service.ondeviceintelligence.IProcessingUpdateStatusCallback;
+
+/**
+ * Interface for a concrete implementation to provide on-device sandboxed inference.
+ *
+ * @hide
+ */
+oneway interface IOnDeviceSandboxedInferenceService {
+    void registerRemoteStorageService(in IRemoteStorageService storageService,
+                                        in IRemoteCallback remoteCallback) = 0;
+    void requestTokenInfo(int callerUid, in Feature feature, in Bundle request,
+                            in AndroidFuture cancellationSignal,
+                            in ITokenInfoCallback tokenInfoCallback) = 1;
+    void processRequest(int callerUid, in Feature feature, in Bundle request, in int requestType,
+                        in AndroidFuture cancellationSignal,
+                        in AndroidFuture processingSignal,
+                        in IResponseCallback callback) = 2;
+    void processRequestStreaming(int callerUid, in Feature feature, in Bundle request, in int requestType,
+                                in AndroidFuture cancellationSignal,
+                                in AndroidFuture processingSignal,
+                                in IStreamingResponseCallback callback) = 3;
+    void updateProcessingState(in Bundle processingState,
+                                     in IProcessingUpdateStatusCallback callback) = 4;
+}
\ No newline at end of file
diff --git a/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/IProcessingUpdateStatusCallback.aidl b/packages/NeuralNetworks/framework/platform/java/android/service/ondeviceintelligence/IProcessingUpdateStatusCallback.aidl
similarity index 100%
copy from packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/IProcessingUpdateStatusCallback.aidl
copy to packages/NeuralNetworks/framework/platform/java/android/service/ondeviceintelligence/IProcessingUpdateStatusCallback.aidl
diff --git a/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/IRemoteProcessingService.aidl b/packages/NeuralNetworks/framework/platform/java/android/service/ondeviceintelligence/IRemoteProcessingService.aidl
similarity index 100%
copy from packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/IRemoteProcessingService.aidl
copy to packages/NeuralNetworks/framework/platform/java/android/service/ondeviceintelligence/IRemoteProcessingService.aidl
diff --git a/packages/NeuralNetworks/framework/platform/java/android/service/ondeviceintelligence/IRemoteStorageService.aidl b/packages/NeuralNetworks/framework/platform/java/android/service/ondeviceintelligence/IRemoteStorageService.aidl
new file mode 100644
index 0000000..a6f49e1
--- /dev/null
+++ b/packages/NeuralNetworks/framework/platform/java/android/service/ondeviceintelligence/IRemoteStorageService.aidl
@@ -0,0 +1,34 @@
+/*
+ * Copyright (C) 2022 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 android.service.ondeviceintelligence;
+
+import android.app.ondeviceintelligence.Feature;
+import android.os.ParcelFileDescriptor;
+import android.os.RemoteCallback;
+
+import com.android.internal.infra.AndroidFuture;
+
+/**
+ * Interface for a concrete implementation to provide access to storage read access
+ * for the isolated process.
+ *
+ * @hide
+ */
+oneway interface IRemoteStorageService {
+    void getReadOnlyFileDescriptor(in String filePath, in AndroidFuture<ParcelFileDescriptor> future);
+    void getReadOnlyFeatureFileDescriptorMap(in Feature feature, in RemoteCallback remoteCallback);
+}
\ No newline at end of file
diff --git a/packages/NeuralNetworks/framework/platform/java/android/service/ondeviceintelligence/OnDeviceIntelligenceService.java b/packages/NeuralNetworks/framework/platform/java/android/service/ondeviceintelligence/OnDeviceIntelligenceService.java
new file mode 100644
index 0000000..618d2a0
--- /dev/null
+++ b/packages/NeuralNetworks/framework/platform/java/android/service/ondeviceintelligence/OnDeviceIntelligenceService.java
@@ -0,0 +1,552 @@
+/*
+ * Copyright (C) 2023 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 android.service.ondeviceintelligence;
+
+import static android.app.ondeviceintelligence.flags.Flags.FLAG_ENABLE_ON_DEVICE_INTELLIGENCE;
+
+import static com.android.internal.util.function.pooled.PooledLambda.obtainMessage;
+
+import android.annotation.CallSuper;
+import android.annotation.CallbackExecutor;
+import android.annotation.FlaggedApi;
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.annotation.SdkConstant;
+import android.annotation.SystemApi;
+import android.app.Service;
+import android.app.ondeviceintelligence.DownloadCallback;
+import android.app.ondeviceintelligence.Feature;
+import android.app.ondeviceintelligence.FeatureDetails;
+import android.app.ondeviceintelligence.IDownloadCallback;
+import android.app.ondeviceintelligence.IFeatureCallback;
+import android.app.ondeviceintelligence.IFeatureDetailsCallback;
+import android.app.ondeviceintelligence.IListFeaturesCallback;
+import android.app.ondeviceintelligence.OnDeviceIntelligenceException;
+import android.app.ondeviceintelligence.OnDeviceIntelligenceManager;
+import android.app.ondeviceintelligence.OnDeviceIntelligenceManager.StateParams;
+import android.content.Intent;
+import android.os.Binder;
+import android.os.Bundle;
+import android.os.CancellationSignal;
+import android.os.Handler;
+import android.os.IBinder;
+import android.os.ICancellationSignal;
+import android.os.Looper;
+import android.os.OutcomeReceiver;
+import android.os.ParcelFileDescriptor;
+import android.os.PersistableBundle;
+import android.os.RemoteCallback;
+import android.os.RemoteException;
+import android.util.Log;
+import android.util.Slog;
+
+import com.android.internal.infra.AndroidFuture;
+
+import java.io.File;
+import java.io.FileNotFoundException;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.Executor;
+import java.util.function.Consumer;
+import java.util.function.LongConsumer;
+
+/**
+ * Abstract base class for performing setup for on-device inference and providing file access to
+ * the isolated counter part {@link OnDeviceSandboxedInferenceService}.
+ *
+ * <p> A service that provides configuration and model files relevant to performing inference on
+ * device. The system's default OnDeviceIntelligenceService implementation is configured in
+ * {@code config_defaultOnDeviceIntelligenceService}. If this config has no value, a stub is
+ * returned.
+ *
+ * <p> Similar to {@link OnDeviceIntelligenceManager} class, the contracts in this service are
+ * defined to be open-ended in general, to allow interoperability. Therefore, it is recommended
+ * that implementations of this system-service expose this API to the clients via a library which
+ * has more defined contract.</p>
+ * <pre>
+ * {@literal
+ * <service android:name=".SampleOnDeviceIntelligenceService"
+ *          android:permission="android.permission.BIND_ON_DEVICE_INTELLIGENCE_SERVICE">
+ * </service>}
+ * </pre>
+ *
+ * @hide
+ */
+@SystemApi
+@FlaggedApi(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
+public abstract class OnDeviceIntelligenceService extends Service {
+    private static final String TAG = OnDeviceIntelligenceService.class.getSimpleName();
+
+    private volatile IRemoteProcessingService mRemoteProcessingService;
+    private Handler mHandler;
+
+    @CallSuper
+    @Override
+    public void onCreate() {
+        super.onCreate();
+        mHandler = new Handler(Looper.getMainLooper(), null /* callback */, true /* async */);
+    }
+
+    /**
+     * The {@link Intent} that must be declared as handled by the service. To be supported, the
+     * service must also require the
+     * {@link android.Manifest.permission#BIND_ON_DEVICE_INTELLIGENCE_SERVICE}
+     * permission so that other applications can not abuse it.
+     */
+    @SdkConstant(SdkConstant.SdkConstantType.SERVICE_ACTION)
+    public static final String SERVICE_INTERFACE =
+            "android.service.ondeviceintelligence.OnDeviceIntelligenceService";
+
+
+    /**
+     * @hide
+     */
+    @Nullable
+    @Override
+    public final IBinder onBind(@NonNull Intent intent) {
+        if (SERVICE_INTERFACE.equals(intent.getAction())) {
+            return new IOnDeviceIntelligenceService.Stub() {
+                /** {@inheritDoc} */
+                @Override
+                public void ready() {
+                    mHandler.executeOrSendMessage(
+                            obtainMessage(OnDeviceIntelligenceService::onReady,
+                                    OnDeviceIntelligenceService.this));
+                }
+
+                @Override
+                public void getVersion(RemoteCallback remoteCallback) {
+                    Objects.requireNonNull(remoteCallback);
+                    mHandler.executeOrSendMessage(
+                            obtainMessage(
+                                    OnDeviceIntelligenceService::onGetVersion,
+                                    OnDeviceIntelligenceService.this, l -> {
+                                        Bundle b = new Bundle();
+                                        b.putLong(
+                                                OnDeviceIntelligenceManager.API_VERSION_BUNDLE_KEY,
+                                                l);
+                                        remoteCallback.sendResult(b);
+                                    }));
+                }
+
+                @Override
+                public void listFeatures(int callerUid,
+                        IListFeaturesCallback listFeaturesCallback) {
+                    Objects.requireNonNull(listFeaturesCallback);
+                    mHandler.executeOrSendMessage(
+                            obtainMessage(
+                                    OnDeviceIntelligenceService::onListFeatures,
+                                    OnDeviceIntelligenceService.this, callerUid,
+                                    wrapListFeaturesCallback(listFeaturesCallback)));
+                }
+
+                @Override
+                public void getFeature(int callerUid, int id, IFeatureCallback featureCallback) {
+                    Objects.requireNonNull(featureCallback);
+                    mHandler.executeOrSendMessage(
+                            obtainMessage(
+                                    OnDeviceIntelligenceService::onGetFeature,
+                                    OnDeviceIntelligenceService.this, callerUid,
+                                    id, wrapFeatureCallback(featureCallback)));
+                }
+
+
+                @Override
+                public void getFeatureDetails(int callerUid, Feature feature,
+                        IFeatureDetailsCallback featureDetailsCallback) {
+                    Objects.requireNonNull(feature);
+                    Objects.requireNonNull(featureDetailsCallback);
+                    mHandler.executeOrSendMessage(
+                            obtainMessage(
+                                    OnDeviceIntelligenceService::onGetFeatureDetails,
+                                    OnDeviceIntelligenceService.this, callerUid,
+                                    feature, wrapFeatureDetailsCallback(featureDetailsCallback)));
+                }
+
+                @Override
+                public void requestFeatureDownload(int callerUid, Feature feature,
+                        AndroidFuture cancellationSignalFuture,
+                        IDownloadCallback downloadCallback) {
+                    Objects.requireNonNull(feature);
+                    Objects.requireNonNull(downloadCallback);
+                    ICancellationSignal transport = null;
+                    if (cancellationSignalFuture != null) {
+                        transport = CancellationSignal.createTransport();
+                        cancellationSignalFuture.complete(transport);
+                    }
+                    mHandler.executeOrSendMessage(
+                            obtainMessage(
+                                    OnDeviceIntelligenceService::onDownloadFeature,
+                                    OnDeviceIntelligenceService.this, callerUid,
+                                    feature,
+                                    CancellationSignal.fromTransport(transport),
+                                    wrapDownloadCallback(downloadCallback)));
+                }
+
+                @Override
+                public void getReadOnlyFileDescriptor(String fileName,
+                        AndroidFuture<ParcelFileDescriptor> future) {
+                    Objects.requireNonNull(fileName);
+                    Objects.requireNonNull(future);
+                    mHandler.executeOrSendMessage(
+                            obtainMessage(
+                                    OnDeviceIntelligenceService::onGetReadOnlyFileDescriptor,
+                                    OnDeviceIntelligenceService.this, fileName,
+                                    future));
+                }
+
+                @Override
+                public void getReadOnlyFeatureFileDescriptorMap(
+                        Feature feature, RemoteCallback remoteCallback) {
+                    Objects.requireNonNull(feature);
+                    Objects.requireNonNull(remoteCallback);
+                    mHandler.executeOrSendMessage(
+                            obtainMessage(
+                                    OnDeviceIntelligenceService::onGetReadOnlyFeatureFileDescriptorMap,
+                                    OnDeviceIntelligenceService.this, feature,
+                                    parcelFileDescriptorMap -> {
+                                        Bundle bundle = new Bundle();
+                                        parcelFileDescriptorMap.forEach(bundle::putParcelable);
+                                        remoteCallback.sendResult(bundle);
+                                        tryClosePfds(parcelFileDescriptorMap.values());
+                                    }));
+                }
+
+                @Override
+                public void registerRemoteServices(
+                        IRemoteProcessingService remoteProcessingService) {
+                    mRemoteProcessingService = remoteProcessingService;
+                }
+
+                @Override
+                public void notifyInferenceServiceConnected() {
+                    mHandler.executeOrSendMessage(
+                            obtainMessage(
+                                    OnDeviceIntelligenceService::onInferenceServiceConnected,
+                                    OnDeviceIntelligenceService.this));
+                }
+
+                @Override
+                public void notifyInferenceServiceDisconnected() {
+                    mHandler.executeOrSendMessage(
+                            obtainMessage(
+                                    OnDeviceIntelligenceService::onInferenceServiceDisconnected,
+                                    OnDeviceIntelligenceService.this));
+                }
+            };
+        }
+        Slog.w(TAG, "Incorrect service interface, returning null.");
+        return null;
+    }
+
+    /**
+     * Using this signal to assertively a signal each time service binds successfully, used only in
+     * tests to get a signal that service instance is ready. This is needed because we cannot rely
+     * on {@link #onCreate} or {@link #onBind} to be invoke on each binding.
+     */
+    public void onReady() {
+    }
+
+
+    /**
+     * Invoked when a new instance of the remote inference service is created.
+     * This method should be used as a signal to perform any initialization operations, for e.g. by
+     * invoking the {@link #updateProcessingState} method to initialize the remote processing
+     * service.
+     */
+    public abstract void onInferenceServiceConnected();
+
+
+    /**
+     * Invoked when an instance of the remote inference service is disconnected.
+     */
+    public abstract void onInferenceServiceDisconnected();
+
+
+    /**
+     * Invoked by the {@link OnDeviceIntelligenceService} inorder to send updates to the inference
+     * service if there is a state change to be performed. State change could be config updates,
+     * performing initialization or cleanup tasks in the remote inference service.
+     * The Bundle passed in here is expected to be read-only and will be rejected if it has any
+     * writable fields as detailed under {@link StateParams}.
+     *
+     * @param processingState  the updated state to be applied.
+     * @param callbackExecutor executor to the run status callback on.
+     * @param statusReceiver   receiver to get status of the update state operation.
+     */
+    public final void updateProcessingState(@NonNull @StateParams Bundle processingState,
+            @NonNull @CallbackExecutor Executor callbackExecutor,
+            @NonNull OutcomeReceiver<PersistableBundle, OnDeviceIntelligenceException> statusReceiver) {
+        Objects.requireNonNull(callbackExecutor);
+        if (mRemoteProcessingService == null) {
+            throw new IllegalStateException("Remote processing service is unavailable.");
+        }
+        try {
+            mRemoteProcessingService.updateProcessingState(processingState,
+                    new IProcessingUpdateStatusCallback.Stub() {
+                        @Override
+                        public void onSuccess(PersistableBundle result) {
+                            Binder.withCleanCallingIdentity(() -> {
+                                callbackExecutor.execute(
+                                        () -> statusReceiver.onResult(result));
+                            });
+                        }
+
+                        @Override
+                        public void onFailure(int errorCode, String errorMessage) {
+                            Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                                    () -> statusReceiver.onError(
+                                            new OnDeviceIntelligenceException(
+                                                    errorCode, errorMessage))));
+                        }
+                    });
+        } catch (RemoteException e) {
+            Slog.e(TAG, "Error in updateProcessingState: " + e);
+            throw new RuntimeException(e);
+        }
+    }
+
+    private OutcomeReceiver<Feature,
+            OnDeviceIntelligenceException> wrapFeatureCallback(
+            IFeatureCallback featureCallback) {
+        return new OutcomeReceiver<>() {
+            @Override
+            public void onResult(@NonNull Feature feature) {
+                try {
+                    featureCallback.onSuccess(feature);
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Error sending feature: " + e);
+                }
+            }
+
+            @Override
+            public void onError(
+                    @NonNull OnDeviceIntelligenceException exception) {
+                try {
+                    featureCallback.onFailure(exception.getErrorCode(), exception.getMessage(),
+                            exception.getErrorParams());
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Error sending download feature: " + e);
+                }
+            }
+        };
+    }
+
+    private OutcomeReceiver<List<Feature>,
+            OnDeviceIntelligenceException> wrapListFeaturesCallback(
+            IListFeaturesCallback listFeaturesCallback) {
+        return new OutcomeReceiver<>() {
+            @Override
+            public void onResult(@NonNull List<Feature> features) {
+                try {
+                    listFeaturesCallback.onSuccess(features);
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Error sending feature: " + e);
+                }
+            }
+
+            @Override
+            public void onError(
+                    @NonNull OnDeviceIntelligenceException exception) {
+                try {
+                    listFeaturesCallback.onFailure(exception.getErrorCode(), exception.getMessage(),
+                            exception.getErrorParams());
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Error sending download feature: " + e);
+                }
+            }
+        };
+    }
+
+    private OutcomeReceiver<FeatureDetails,
+            OnDeviceIntelligenceException> wrapFeatureDetailsCallback(
+            IFeatureDetailsCallback featureStatusCallback) {
+        return new OutcomeReceiver<>() {
+            @Override
+            public void onResult(FeatureDetails result) {
+                try {
+                    featureStatusCallback.onSuccess(result);
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Error sending feature status: " + e);
+                }
+            }
+
+            @Override
+            public void onError(
+                    @NonNull OnDeviceIntelligenceException exception) {
+                try {
+                    featureStatusCallback.onFailure(exception.getErrorCode(),
+                            exception.getMessage(), exception.getErrorParams());
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Error sending feature status: " + e);
+                }
+            }
+        };
+    }
+
+
+    private DownloadCallback wrapDownloadCallback(IDownloadCallback downloadCallback) {
+        return new DownloadCallback() {
+            @Override
+            public void onDownloadStarted(long bytesToDownload) {
+                try {
+                    downloadCallback.onDownloadStarted(bytesToDownload);
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Error sending download status: " + e);
+                }
+            }
+
+            @Override
+            public void onDownloadFailed(int failureStatus,
+                    String errorMessage, @NonNull PersistableBundle errorParams) {
+                try {
+                    downloadCallback.onDownloadFailed(failureStatus, errorMessage, errorParams);
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Error sending download status: " + e);
+                }
+            }
+
+            @Override
+            public void onDownloadProgress(long totalBytesDownloaded) {
+                try {
+                    downloadCallback.onDownloadProgress(totalBytesDownloaded);
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Error sending download status: " + e);
+                }
+            }
+
+            @Override
+            public void onDownloadCompleted(@NonNull PersistableBundle persistableBundle) {
+                try {
+                    downloadCallback.onDownloadCompleted(persistableBundle);
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Error sending download status: " + e);
+                }
+            }
+        };
+    }
+
+    private static void tryClosePfds(Collection<ParcelFileDescriptor> pfds) {
+        pfds.forEach(pfd -> {
+            try {
+                pfd.close();
+            } catch (Exception e) {
+                Log.w(TAG, "Error closing FD", e);
+            }
+        });
+    }
+
+    private void onGetReadOnlyFileDescriptor(@NonNull String fileName,
+            @NonNull AndroidFuture<ParcelFileDescriptor> future) {
+        Slog.v(TAG, "onGetReadOnlyFileDescriptor " + fileName);
+        Binder.withCleanCallingIdentity(() -> {
+            Slog.v(TAG,
+                    "onGetReadOnlyFileDescriptor: " + fileName + " under internal app storage.");
+            File f = new File(getBaseContext().getFilesDir(), fileName);
+            if (!f.exists()) {
+                f = new File(fileName);
+            }
+            ParcelFileDescriptor pfd = null;
+            try {
+                pfd = ParcelFileDescriptor.open(f, ParcelFileDescriptor.MODE_READ_ONLY);
+                Slog.d(TAG, "Successfully opened a file with ParcelFileDescriptor.");
+            } catch (FileNotFoundException e) {
+                Slog.e(TAG, "Cannot open file. No ParcelFileDescriptor returned.");
+                future.completeExceptionally(e);
+            } finally {
+                future.complete(pfd);
+                if (pfd != null) {
+                    pfd.close();
+                }
+            }
+        });
+    }
+
+    /**
+     * Provide implementation for a scenario when caller wants to get all feature related
+     * file-descriptors that might be required for processing a request for the corresponding the
+     * feature.
+     *
+     * @param feature                   the feature for which files need to be opened.
+     * @param fileDescriptorMapConsumer callback to be populated with a map of file-path and
+     *                                  corresponding ParcelDescriptor to be used in a remote
+     *                                  service.
+     */
+    public abstract void onGetReadOnlyFeatureFileDescriptorMap(
+            @NonNull Feature feature,
+            @NonNull Consumer<Map<String, ParcelFileDescriptor>> fileDescriptorMapConsumer);
+
+    /**
+     * Request download for feature that is requested and listen to download progress updates. If
+     * the download completes successfully, success callback should be populated.
+     *
+     * @param callerUid          UID of the caller that initiated this call chain.
+     * @param feature            the feature for which files need to be downlaoded.
+     *                           process.
+     * @param cancellationSignal signal to attach a listener to, and receive cancellation signals
+     *                           from thw client.
+     * @param downloadCallback   callback to populate download updates for clients to listen on..
+     */
+    public abstract void onDownloadFeature(
+            int callerUid, @NonNull Feature feature,
+            @Nullable CancellationSignal cancellationSignal,
+            @NonNull DownloadCallback downloadCallback);
+
+    /**
+     * Provide feature details for the passed in feature. Usually the client and remote
+     * implementation use the {@link Feature#getFeatureParams()} as a hint to communicate what
+     * details the client is looking for.
+     *
+     * @param callerUid              UID of the caller that initiated this call chain.
+     * @param feature                the feature for which status needs to be known.
+     * @param featureDetailsCallback callback to populate the resulting feature status.
+     */
+    public abstract void onGetFeatureDetails(int callerUid, @NonNull Feature feature,
+            @NonNull OutcomeReceiver<FeatureDetails,
+                    OnDeviceIntelligenceException> featureDetailsCallback);
+
+
+    /**
+     * Get feature using the provided identifier to the remote implementation.
+     *
+     * @param callerUid       UID of the caller that initiated this call chain.
+     * @param featureCallback callback to populate the features list.
+     */
+    public abstract void onGetFeature(int callerUid, int featureId,
+            @NonNull OutcomeReceiver<Feature,
+                    OnDeviceIntelligenceException> featureCallback);
+
+    /**
+     * List all features which are available in the remote implementation. The implementation might
+     * choose to provide only a certain list of features based on the caller.
+     *
+     * @param callerUid            UID of the caller that initiated this call chain.
+     * @param listFeaturesCallback callback to populate the features list.
+     */
+    public abstract void onListFeatures(int callerUid, @NonNull OutcomeReceiver<List<Feature>,
+            OnDeviceIntelligenceException> listFeaturesCallback);
+
+    /**
+     * Provides a long value representing the version of the remote implementation processing
+     * requests.
+     *
+     * @param versionConsumer consumer to populate the version.
+     */
+    public abstract void onGetVersion(@NonNull LongConsumer versionConsumer);
+}
diff --git a/packages/NeuralNetworks/framework/platform/java/android/service/ondeviceintelligence/OnDeviceSandboxedInferenceService.java b/packages/NeuralNetworks/framework/platform/java/android/service/ondeviceintelligence/OnDeviceSandboxedInferenceService.java
new file mode 100644
index 0000000..949fb8d
--- /dev/null
+++ b/packages/NeuralNetworks/framework/platform/java/android/service/ondeviceintelligence/OnDeviceSandboxedInferenceService.java
@@ -0,0 +1,617 @@
+/*
+ * Copyright (C) 2023 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 android.service.ondeviceintelligence;
+
+import static android.app.ondeviceintelligence.OnDeviceIntelligenceManager.AUGMENT_REQUEST_CONTENT_BUNDLE_KEY;
+import static android.app.ondeviceintelligence.flags.Flags.FLAG_ENABLE_ON_DEVICE_INTELLIGENCE;
+
+import static com.android.internal.util.function.pooled.PooledLambda.obtainMessage;
+
+import android.annotation.CallSuper;
+import android.annotation.CallbackExecutor;
+import android.annotation.FlaggedApi;
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.annotation.SdkConstant;
+import android.annotation.SuppressLint;
+import android.annotation.SystemApi;
+import android.app.Service;
+import android.app.ondeviceintelligence.Feature;
+import android.app.ondeviceintelligence.IProcessingSignal;
+import android.app.ondeviceintelligence.IResponseCallback;
+import android.app.ondeviceintelligence.IStreamingResponseCallback;
+import android.app.ondeviceintelligence.ITokenInfoCallback;
+import android.app.ondeviceintelligence.OnDeviceIntelligenceException;
+import android.app.ondeviceintelligence.OnDeviceIntelligenceManager;
+import android.app.ondeviceintelligence.OnDeviceIntelligenceManager.InferenceParams;
+import android.app.ondeviceintelligence.OnDeviceIntelligenceManager.StateParams;
+import android.app.ondeviceintelligence.ProcessingCallback;
+import android.app.ondeviceintelligence.ProcessingSignal;
+import android.app.ondeviceintelligence.StreamingProcessingCallback;
+import android.app.ondeviceintelligence.TokenInfo;
+import android.content.Context;
+import android.content.Intent;
+import android.os.Bundle;
+import android.os.CancellationSignal;
+import android.os.Handler;
+import android.os.HandlerExecutor;
+import android.os.IBinder;
+import android.os.ICancellationSignal;
+import android.os.IRemoteCallback;
+import android.os.Looper;
+import android.os.OutcomeReceiver;
+import android.os.ParcelFileDescriptor;
+import android.os.PersistableBundle;
+import android.os.RemoteCallback;
+import android.os.RemoteException;
+import android.util.Log;
+import android.util.Slog;
+
+import com.android.internal.infra.AndroidFuture;
+
+import java.io.FileInputStream;
+import java.io.FileNotFoundException;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Executor;
+import java.util.function.Consumer;
+
+/**
+ * Abstract base class for performing inference in a isolated process. This service exposes its
+ * methods via {@link OnDeviceIntelligenceManager}.
+ *
+ * <p> A service that provides methods to perform on-device inference both in streaming and
+ * non-streaming fashion. Also, provides a way to register a storage service that will be used to
+ * read-only access files from the {@link OnDeviceIntelligenceService} counterpart. </p>
+ *
+ * <p> Similar to {@link OnDeviceIntelligenceManager} class, the contracts in this service are
+ * defined to be open-ended in general, to allow interoperability. Therefore, it is recommended
+ * that implementations of this system-service expose this API to the clients via a library which
+ * has more defined contract.</p>
+ *
+ * <pre>
+ * {@literal
+ * <service android:name=".SampleSandboxedInferenceService"
+ *          android:permission="android.permission.BIND_ONDEVICE_SANDBOXED_INFERENCE_SERVICE"
+ *          android:isolatedProcess="true">
+ * </service>}
+ * </pre>
+ *
+ * @hide
+ */
+@SystemApi
+@FlaggedApi(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
+public abstract class OnDeviceSandboxedInferenceService extends Service {
+    private static final String TAG = OnDeviceSandboxedInferenceService.class.getSimpleName();
+
+    /**
+     * @hide
+     */
+    public static final String INFERENCE_INFO_BUNDLE_KEY = "inference_info";
+
+    /**
+     * The {@link Intent} that must be declared as handled by the service. To be supported, the
+     * service must also require the
+     * {@link android.Manifest.permission#BIND_ON_DEVICE_SANDBOXED_INFERENCE_SERVICE}
+     * permission so that other applications can not abuse it.
+     */
+    @SdkConstant(SdkConstant.SdkConstantType.SERVICE_ACTION)
+    public static final String SERVICE_INTERFACE =
+            "android.service.ondeviceintelligence.OnDeviceSandboxedInferenceService";
+
+    // TODO(339594686): make API
+    /**
+     * @hide
+     */
+    public static final String REGISTER_MODEL_UPDATE_CALLBACK_BUNDLE_KEY =
+            "register_model_update_callback";
+    /**
+     * @hide
+     */
+    public static final String MODEL_LOADED_BUNDLE_KEY = "model_loaded";
+    /**
+     * @hide
+     */
+    public static final String MODEL_UNLOADED_BUNDLE_KEY = "model_unloaded";
+    /**
+     * @hide
+     */
+    public static final String MODEL_LOADED_BROADCAST_INTENT =
+        "android.service.ondeviceintelligence.MODEL_LOADED";
+    /**
+     * @hide
+     */
+    public static final String MODEL_UNLOADED_BROADCAST_INTENT =
+        "android.service.ondeviceintelligence.MODEL_UNLOADED";
+
+    /**
+     * @hide
+     */
+    public static final String DEVICE_CONFIG_UPDATE_BUNDLE_KEY = "device_config_update";
+
+    private IRemoteStorageService mRemoteStorageService;
+    private Handler mHandler;
+
+    @CallSuper
+    @Override
+    public void onCreate() {
+        super.onCreate();
+        mHandler = new Handler(Looper.getMainLooper(), null /* callback */, true /* async */);
+    }
+
+    /**
+     * @hide
+     */
+    @Nullable
+    @Override
+    public final IBinder onBind(@NonNull Intent intent) {
+        if (SERVICE_INTERFACE.equals(intent.getAction())) {
+            return new IOnDeviceSandboxedInferenceService.Stub() {
+                @Override
+                public void registerRemoteStorageService(IRemoteStorageService storageService,
+                        IRemoteCallback remoteCallback) throws RemoteException {
+                    Objects.requireNonNull(storageService);
+                    mRemoteStorageService = storageService;
+                    remoteCallback.sendResult(
+                            Bundle.EMPTY); //to notify caller uid to system-server.
+                }
+
+                @Override
+                public void requestTokenInfo(int callerUid, Feature feature, Bundle request,
+                        AndroidFuture cancellationSignalFuture,
+                        ITokenInfoCallback tokenInfoCallback) {
+                    Objects.requireNonNull(feature);
+                    Objects.requireNonNull(tokenInfoCallback);
+                    ICancellationSignal transport = null;
+                    if (cancellationSignalFuture != null) {
+                        transport = CancellationSignal.createTransport();
+                        cancellationSignalFuture.complete(transport);
+                    }
+
+                    mHandler.executeOrSendMessage(
+                            obtainMessage(
+                                    OnDeviceSandboxedInferenceService::onTokenInfoRequest,
+                                    OnDeviceSandboxedInferenceService.this,
+                                    callerUid, feature,
+                                    request,
+                                    CancellationSignal.fromTransport(transport),
+                                    wrapTokenInfoCallback(tokenInfoCallback)));
+                }
+
+                @Override
+                public void processRequestStreaming(int callerUid, Feature feature, Bundle request,
+                        int requestType,
+                        AndroidFuture cancellationSignalFuture,
+                        AndroidFuture processingSignalFuture,
+                        IStreamingResponseCallback callback) {
+                    Objects.requireNonNull(feature);
+                    Objects.requireNonNull(callback);
+
+                    ICancellationSignal transport = null;
+                    if (cancellationSignalFuture != null) {
+                        transport = CancellationSignal.createTransport();
+                        cancellationSignalFuture.complete(transport);
+                    }
+                    IProcessingSignal processingSignalTransport = null;
+                    if (processingSignalFuture != null) {
+                        processingSignalTransport = ProcessingSignal.createTransport();
+                        processingSignalFuture.complete(processingSignalTransport);
+                    }
+
+
+                    mHandler.executeOrSendMessage(
+                            obtainMessage(
+                                    OnDeviceSandboxedInferenceService::onProcessRequestStreaming,
+                                    OnDeviceSandboxedInferenceService.this, callerUid,
+                                    feature,
+                                    request,
+                                    requestType,
+                                    CancellationSignal.fromTransport(transport),
+                                    ProcessingSignal.fromTransport(processingSignalTransport),
+                                    wrapStreamingResponseCallback(callback)));
+                }
+
+                @Override
+                public void processRequest(int callerUid, Feature feature, Bundle request,
+                        int requestType,
+                        AndroidFuture cancellationSignalFuture,
+                        AndroidFuture processingSignalFuture,
+                        IResponseCallback callback) {
+                    Objects.requireNonNull(feature);
+                    Objects.requireNonNull(callback);
+                    ICancellationSignal transport = null;
+                    if (cancellationSignalFuture != null) {
+                        transport = CancellationSignal.createTransport();
+                        cancellationSignalFuture.complete(transport);
+                    }
+                    IProcessingSignal processingSignalTransport = null;
+                    if (processingSignalFuture != null) {
+                        processingSignalTransport = ProcessingSignal.createTransport();
+                        processingSignalFuture.complete(processingSignalTransport);
+                    }
+                    mHandler.executeOrSendMessage(
+                            obtainMessage(
+                                    OnDeviceSandboxedInferenceService::onProcessRequest,
+                                    OnDeviceSandboxedInferenceService.this, callerUid, feature,
+                                    request, requestType,
+                                    CancellationSignal.fromTransport(transport),
+                                    ProcessingSignal.fromTransport(processingSignalTransport),
+                                    wrapResponseCallback(callback)));
+                }
+
+                @Override
+                public void updateProcessingState(Bundle processingState,
+                        IProcessingUpdateStatusCallback callback) {
+                    Objects.requireNonNull(processingState);
+                    Objects.requireNonNull(callback);
+                    mHandler.executeOrSendMessage(
+                            obtainMessage(
+                                    OnDeviceSandboxedInferenceService::onUpdateProcessingState,
+                                    OnDeviceSandboxedInferenceService.this, processingState,
+                                    wrapOutcomeReceiver(callback)));
+                }
+            };
+        }
+        Slog.w(TAG, "Incorrect service interface, returning null.");
+        return null;
+    }
+
+    /**
+     * Invoked when caller  wants to obtain token info related to the payload in the passed
+     * content, associated with the provided feature.
+     * The expectation from the implementation is that when processing is complete, it
+     * should provide the token info in the {@link OutcomeReceiver#onResult}.
+     *
+     * @param callerUid          UID of the caller that initiated this call chain.
+     * @param feature            feature which is associated with the request.
+     * @param request            request that requires processing.
+     * @param cancellationSignal Cancellation Signal to receive cancellation events from client and
+     *                           configure a listener to.
+     * @param callback           callback to populate failure or the token info for the provided
+     *                           request.
+     */
+    @NonNull
+    public abstract void onTokenInfoRequest(
+            int callerUid, @NonNull Feature feature,
+            @NonNull @InferenceParams Bundle request,
+            @Nullable CancellationSignal cancellationSignal,
+            @NonNull OutcomeReceiver<TokenInfo, OnDeviceIntelligenceException> callback);
+
+    /**
+     * Invoked when caller provides a request for a particular feature to be processed in a
+     * streaming manner. The expectation from the implementation is that when processing the
+     * request,
+     * it periodically populates the {@link StreamingProcessingCallback#onPartialResult} to
+     * continuously
+     * provide partial Bundle results for the caller to utilize. Optionally the implementation can
+     * provide the complete response in the {@link StreamingProcessingCallback#onResult} upon
+     * processing completion.
+     *
+     * @param callerUid          UID of the caller that initiated this call chain.
+     * @param feature            feature which is associated with the request.
+     * @param request            request that requires processing.
+     * @param requestType        identifier representing the type of request.
+     * @param cancellationSignal Cancellation Signal to receive cancellation events from client and
+     *                           configure a listener to.
+     * @param processingSignal   Signal to receive custom action instructions from client.
+     * @param callback           callback to populate the partial responses, failure and optionally
+     *                           full response for the provided request.
+     */
+    @NonNull
+    public abstract void onProcessRequestStreaming(
+            int callerUid, @NonNull Feature feature,
+            @NonNull @InferenceParams Bundle request,
+            @OnDeviceIntelligenceManager.RequestType int requestType,
+            @Nullable CancellationSignal cancellationSignal,
+            @Nullable ProcessingSignal processingSignal,
+            @NonNull StreamingProcessingCallback callback);
+
+    /**
+     * Invoked when caller provides a request for a particular feature to be processed in one shot
+     * completely.
+     * The expectation from the implementation is that when processing the request is complete, it
+     * should
+     * provide the complete response in the {@link OutcomeReceiver#onResult}.
+     *
+     * @param callerUid          UID of the caller that initiated this call chain.
+     * @param feature            feature which is associated with the request.
+     * @param request            request that requires processing.
+     * @param requestType        identifier representing the type of request.
+     * @param cancellationSignal Cancellation Signal to receive cancellation events from client and
+     *                           configure a listener to.
+     * @param processingSignal   Signal to receive custom action instructions from client.
+     * @param callback           callback to populate failure and full response for the provided
+     *                           request.
+     */
+    @NonNull
+    public abstract void onProcessRequest(
+            int callerUid, @NonNull Feature feature,
+            @NonNull @InferenceParams Bundle request,
+            @OnDeviceIntelligenceManager.RequestType int requestType,
+            @Nullable CancellationSignal cancellationSignal,
+            @Nullable ProcessingSignal processingSignal,
+            @NonNull ProcessingCallback callback);
+
+
+    /**
+     * Invoked when processing environment needs to be updated or refreshed with fresh
+     * configuration, files or state.
+     *
+     * @param processingState contains updated state and params that are to be applied to the
+     *                        processing environmment,
+     * @param callback        callback to populate the update status and if there are params
+     *                        associated with the status.
+     */
+    public abstract void onUpdateProcessingState(@NonNull @StateParams Bundle processingState,
+            @NonNull OutcomeReceiver<PersistableBundle,
+                    OnDeviceIntelligenceException> callback);
+
+
+    /**
+     * Overrides {@link Context#openFileInput} to read files with the given file names under the
+     * internal app storage of the {@link OnDeviceIntelligenceService}, i.e., only files stored in
+     * {@link Context#getFilesDir()} can be opened.
+     */
+    @Override
+    public final FileInputStream openFileInput(@NonNull String filename) throws
+            FileNotFoundException {
+        try {
+            AndroidFuture<ParcelFileDescriptor> future = new AndroidFuture<>();
+            mRemoteStorageService.getReadOnlyFileDescriptor(filename, future);
+            ParcelFileDescriptor pfd = future.get();
+            return new FileInputStream(pfd.getFileDescriptor());
+        } catch (RemoteException | ExecutionException | InterruptedException e) {
+            Log.w(TAG, "Cannot open file due to remote service failure");
+            throw new FileNotFoundException(e.getMessage());
+        }
+    }
+
+    /**
+     * Provides read-only access to the internal app storage via the
+     * {@link OnDeviceIntelligenceService}. This is an asynchronous alternative for
+     * {@link #openFileInput(String)}.
+     *
+     * @param fileName       File name relative to the {@link Context#getFilesDir()}.
+     * @param resultConsumer Consumer to populate the corresponding file descriptor in.
+     */
+    public final void getReadOnlyFileDescriptor(@NonNull String fileName,
+            @NonNull @CallbackExecutor Executor executor,
+            @NonNull Consumer<ParcelFileDescriptor> resultConsumer) throws FileNotFoundException {
+        AndroidFuture<ParcelFileDescriptor> future = new AndroidFuture<>();
+        try {
+            mRemoteStorageService.getReadOnlyFileDescriptor(fileName, future);
+        } catch (RemoteException e) {
+            Log.w(TAG, "Cannot open file due to remote service failure");
+            throw new FileNotFoundException(e.getMessage());
+        }
+        future.whenCompleteAsync((pfd, err) -> {
+            if (err != null) {
+                Log.e(TAG, "Failure when reading file: " + fileName + err);
+                executor.execute(() -> resultConsumer.accept(null));
+            } else {
+                executor.execute(
+                        () -> resultConsumer.accept(pfd));
+            }
+        }, executor);
+    }
+
+    /**
+     * Provides access to all file streams required for feature via the
+     * {@link OnDeviceIntelligenceService}.
+     *
+     * @param feature        Feature for which the associated files should be fetched.
+     * @param executor       Executor to run the consumer callback on.
+     * @param resultConsumer Consumer to receive a map of filePath to the corresponding file input
+     *                       stream.
+     */
+    public final void fetchFeatureFileDescriptorMap(@NonNull Feature feature,
+            @NonNull @CallbackExecutor Executor executor,
+            @NonNull Consumer<Map<String, ParcelFileDescriptor>> resultConsumer) {
+        try {
+            mRemoteStorageService.getReadOnlyFeatureFileDescriptorMap(feature,
+                    wrapAsRemoteCallback(resultConsumer, executor));
+        } catch (RemoteException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+
+    /**
+     * Returns the {@link Executor} to use for incoming IPC from request sender into your service
+     * implementation. For e.g. see
+     * {@link ProcessingCallback#onDataAugmentRequest(Bundle,
+     * Consumer)} where we use the executor to populate the consumer.
+     * <p>
+     * Override this method in your {@link OnDeviceSandboxedInferenceService} implementation to
+     * provide the executor you want to use for incoming IPC.
+     *
+     * @return the {@link Executor} to use for incoming IPC from {@link OnDeviceIntelligenceManager}
+     * to {@link OnDeviceSandboxedInferenceService}.
+     */
+    @SuppressLint("OnNameExpected")
+    @NonNull
+    public Executor getCallbackExecutor() {
+        return new HandlerExecutor(Handler.createAsync(getMainLooper()));
+    }
+
+
+    private RemoteCallback wrapAsRemoteCallback(
+            @NonNull Consumer<Map<String, ParcelFileDescriptor>> resultConsumer,
+            @NonNull Executor executor) {
+        return new RemoteCallback(result -> {
+            if (result == null) {
+                executor.execute(() -> resultConsumer.accept(new HashMap<>()));
+            } else {
+                Map<String, ParcelFileDescriptor> pfdMap = new HashMap<>();
+                result.keySet().forEach(key ->
+                        pfdMap.put(key, result.getParcelable(key,
+                                ParcelFileDescriptor.class)));
+                executor.execute(() -> resultConsumer.accept(pfdMap));
+            }
+        });
+    }
+
+    private ProcessingCallback wrapResponseCallback(
+            IResponseCallback callback) {
+        return new ProcessingCallback() {
+            @Override
+            public void onResult(@NonNull Bundle result) {
+                try {
+                    callback.onSuccess(result);
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Error sending result: " + e);
+                }
+            }
+
+            @Override
+            public void onError(
+                    OnDeviceIntelligenceException exception) {
+                try {
+                    callback.onFailure(exception.getErrorCode(), exception.getMessage(),
+                            exception.getErrorParams());
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Error sending result: " + e);
+                }
+            }
+
+            @Override
+            public void onDataAugmentRequest(@NonNull Bundle content,
+                    @NonNull Consumer<Bundle> contentCallback) {
+                try {
+                    callback.onDataAugmentRequest(content, wrapRemoteCallback(contentCallback));
+
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Error sending augment request: " + e);
+                }
+            }
+        };
+    }
+
+    private StreamingProcessingCallback wrapStreamingResponseCallback(
+            IStreamingResponseCallback callback) {
+        return new StreamingProcessingCallback() {
+            @Override
+            public void onPartialResult(@NonNull Bundle partialResult) {
+                try {
+                    callback.onNewContent(partialResult);
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Error sending result: " + e);
+                }
+            }
+
+            @Override
+            public void onResult(@NonNull Bundle result) {
+                try {
+                    callback.onSuccess(result);
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Error sending result: " + e);
+                }
+            }
+
+            @Override
+            public void onError(
+                    OnDeviceIntelligenceException exception) {
+                try {
+                    callback.onFailure(exception.getErrorCode(), exception.getMessage(),
+                            exception.getErrorParams());
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Error sending result: " + e);
+                }
+            }
+
+            @Override
+            public void onDataAugmentRequest(@NonNull Bundle content,
+                    @NonNull Consumer<Bundle> contentCallback) {
+                try {
+                    callback.onDataAugmentRequest(content, wrapRemoteCallback(contentCallback));
+
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Error sending augment request: " + e);
+                }
+            }
+        };
+    }
+
+    private RemoteCallback wrapRemoteCallback(
+            @NonNull Consumer<Bundle> contentCallback) {
+        return new RemoteCallback(
+                result -> {
+                    if (result != null) {
+                        getCallbackExecutor().execute(() -> contentCallback.accept(
+                                result.getParcelable(AUGMENT_REQUEST_CONTENT_BUNDLE_KEY,
+                                        Bundle.class)));
+                    } else {
+                        getCallbackExecutor().execute(
+                                () -> contentCallback.accept(null));
+                    }
+                });
+    }
+
+    private OutcomeReceiver<TokenInfo, OnDeviceIntelligenceException> wrapTokenInfoCallback(
+            ITokenInfoCallback tokenInfoCallback) {
+        return new OutcomeReceiver<>() {
+            @Override
+            public void onResult(TokenInfo tokenInfo) {
+                try {
+                    tokenInfoCallback.onSuccess(tokenInfo);
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Error sending result: " + e);
+                }
+            }
+
+            @Override
+            public void onError(
+                    OnDeviceIntelligenceException exception) {
+                try {
+                    tokenInfoCallback.onFailure(exception.getErrorCode(), exception.getMessage(),
+                            exception.getErrorParams());
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Error sending failure: " + e);
+                }
+            }
+        };
+    }
+
+    @NonNull
+    private static OutcomeReceiver<PersistableBundle, OnDeviceIntelligenceException> wrapOutcomeReceiver(
+            IProcessingUpdateStatusCallback callback) {
+        return new OutcomeReceiver<>() {
+            @Override
+            public void onResult(@NonNull PersistableBundle result) {
+                try {
+                    callback.onSuccess(result);
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Error sending result: " + e);
+
+                }
+            }
+
+            @Override
+            public void onError(
+                    @NonNull OnDeviceIntelligenceException error) {
+                try {
+                    callback.onFailure(error.getErrorCode(), error.getMessage());
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Error sending exception details: " + e);
+                }
+            }
+        };
+    }
+
+}
diff --git a/packages/NeuralNetworks/service/Android.bp b/packages/NeuralNetworks/service/Android.bp
index 05c603f..cfdc1af 100644
--- a/packages/NeuralNetworks/service/Android.bp
+++ b/packages/NeuralNetworks/service/Android.bp
@@ -19,11 +19,20 @@
 filegroup {
     name: "service-ondeviceintelligence-sources",
     srcs: [
-        "java/**/*.java",
+        "module/java/**/*.java",
     ],
-    path: "java",
     visibility: [
         "//frameworks/base:__subpackages__",
         "//packages/modules/NeuralNetworks:__subpackages__",
     ],
 }
+
+filegroup {
+    name: "service-ondeviceintelligence-sources-platform",
+    srcs: [
+        "platform/java/**/*.java",
+    ],
+    visibility: [
+        "//frameworks/base:__subpackages__",
+    ],
+}
diff --git a/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/BundleUtil.java b/packages/NeuralNetworks/service/module/java/com/android/server/ondeviceintelligence/BundleUtil.java
similarity index 100%
rename from packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/BundleUtil.java
rename to packages/NeuralNetworks/service/module/java/com/android/server/ondeviceintelligence/BundleUtil.java
diff --git a/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/InferenceInfoStore.java b/packages/NeuralNetworks/service/module/java/com/android/server/ondeviceintelligence/InferenceInfoStore.java
similarity index 100%
rename from packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/InferenceInfoStore.java
rename to packages/NeuralNetworks/service/module/java/com/android/server/ondeviceintelligence/InferenceInfoStore.java
diff --git a/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerLocal.java b/packages/NeuralNetworks/service/module/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerLocal.java
similarity index 100%
rename from packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerLocal.java
rename to packages/NeuralNetworks/service/module/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerLocal.java
diff --git a/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerService.java b/packages/NeuralNetworks/service/module/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerService.java
similarity index 100%
rename from packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerService.java
rename to packages/NeuralNetworks/service/module/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerService.java
diff --git a/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceShellCommand.java b/packages/NeuralNetworks/service/module/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceShellCommand.java
similarity index 100%
rename from packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceShellCommand.java
rename to packages/NeuralNetworks/service/module/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceShellCommand.java
diff --git a/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/RemoteOnDeviceIntelligenceService.java b/packages/NeuralNetworks/service/module/java/com/android/server/ondeviceintelligence/RemoteOnDeviceIntelligenceService.java
similarity index 100%
rename from packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/RemoteOnDeviceIntelligenceService.java
rename to packages/NeuralNetworks/service/module/java/com/android/server/ondeviceintelligence/RemoteOnDeviceIntelligenceService.java
diff --git a/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/RemoteOnDeviceSandboxedInferenceService.java b/packages/NeuralNetworks/service/module/java/com/android/server/ondeviceintelligence/RemoteOnDeviceSandboxedInferenceService.java
similarity index 100%
rename from packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/RemoteOnDeviceSandboxedInferenceService.java
rename to packages/NeuralNetworks/service/module/java/com/android/server/ondeviceintelligence/RemoteOnDeviceSandboxedInferenceService.java
diff --git a/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/callbacks/ListenableDownloadCallback.java b/packages/NeuralNetworks/service/module/java/com/android/server/ondeviceintelligence/callbacks/ListenableDownloadCallback.java
similarity index 100%
rename from packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/callbacks/ListenableDownloadCallback.java
rename to packages/NeuralNetworks/service/module/java/com/android/server/ondeviceintelligence/callbacks/ListenableDownloadCallback.java
diff --git a/packages/NeuralNetworks/service/platform/java/com/android/server/ondeviceintelligence/BundleUtil.java b/packages/NeuralNetworks/service/platform/java/com/android/server/ondeviceintelligence/BundleUtil.java
new file mode 100644
index 0000000..7dd8f2f
--- /dev/null
+++ b/packages/NeuralNetworks/service/platform/java/com/android/server/ondeviceintelligence/BundleUtil.java
@@ -0,0 +1,407 @@
+/*
+ * 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.server.ondeviceintelligence;
+
+import static android.system.OsConstants.F_GETFL;
+import static android.system.OsConstants.O_ACCMODE;
+import static android.system.OsConstants.O_RDONLY;
+import static android.system.OsConstants.PROT_READ;
+
+import android.app.ondeviceintelligence.IResponseCallback;
+import android.app.ondeviceintelligence.IStreamingResponseCallback;
+import android.app.ondeviceintelligence.ITokenInfoCallback;
+import android.app.ondeviceintelligence.OnDeviceIntelligenceManager.InferenceParams;
+import android.app.ondeviceintelligence.OnDeviceIntelligenceManager.ResponseParams;
+import android.app.ondeviceintelligence.OnDeviceIntelligenceManager.StateParams;
+import android.app.ondeviceintelligence.TokenInfo;
+import android.database.CursorWindow;
+import android.graphics.Bitmap;
+import android.os.BadParcelableException;
+import android.os.Bundle;
+import android.os.ParcelFileDescriptor;
+import android.os.Parcelable;
+import android.os.PersistableBundle;
+import android.os.RemoteCallback;
+import android.os.RemoteException;
+import android.os.SharedMemory;
+import android.system.ErrnoException;
+import android.system.Os;
+import android.util.Log;
+
+import com.android.internal.infra.AndroidFuture;
+
+import java.util.concurrent.Executor;
+import java.util.concurrent.TimeoutException;
+
+/**
+ * Util methods for ensuring the Bundle passed in various methods are read-only and restricted to
+ * some known types.
+ */
+public class BundleUtil {
+    private static final String TAG = "BundleUtil";
+
+    /**
+     * Validation of the inference request payload as described in {@link InferenceParams}
+     * description.
+     *
+     * @throws BadParcelableException when the bundle does not meet the read-only requirements.
+     */
+    public static void sanitizeInferenceParams(
+            @InferenceParams Bundle bundle) {
+        ensureValidBundle(bundle);
+
+        if (!bundle.hasFileDescriptors()) {
+            return; //safe to exit if there are no FDs and Binders
+        }
+
+        for (String key : bundle.keySet()) {
+            Object obj = bundle.get(key);
+            if (obj == null) {
+                /* Null value here could also mean deserializing a custom parcelable has failed,
+                 *  and since {@link Bundle} is marked as defusable in system-server - the
+                 * {@link ClassNotFoundException} exception is swallowed and `null` is returned
+                 * instead. We want to ensure cleanup of null entries in such case.
+                 */
+                bundle.putObject(key, null);
+                continue;
+            }
+            if (canMarshall(obj) || obj instanceof CursorWindow) {
+                continue;
+            }
+            if (obj instanceof Bundle) {
+              sanitizeInferenceParams((Bundle) obj);
+            } else if (obj instanceof ParcelFileDescriptor) {
+                validatePfdReadOnly((ParcelFileDescriptor) obj);
+            } else if (obj instanceof SharedMemory) {
+                ((SharedMemory) obj).setProtect(PROT_READ);
+            } else if (obj instanceof Bitmap) {
+                validateBitmap((Bitmap) obj);
+            } else if (obj instanceof Parcelable[]) {
+                validateParcelableArray((Parcelable[]) obj);
+            } else {
+                throw new BadParcelableException(
+                        "Unsupported Parcelable type encountered in the Bundle: "
+                                + obj.getClass().getSimpleName());
+            }
+        }
+    }
+
+    /**
+     * Validation of the inference request payload as described in {@link ResponseParams}
+     * description.
+     *
+     * @throws BadParcelableException when the bundle does not meet the read-only requirements.
+     */
+    public static void sanitizeResponseParams(
+            @ResponseParams Bundle bundle) {
+        ensureValidBundle(bundle);
+
+        if (!bundle.hasFileDescriptors()) {
+            return; //safe to exit if there are no FDs and Binders
+        }
+
+        for (String key : bundle.keySet()) {
+            Object obj = bundle.get(key);
+            if (obj == null) {
+                /* Null value here could also mean deserializing a custom parcelable has failed,
+                 *  and since {@link Bundle} is marked as defusable in system-server - the
+                 * {@link ClassNotFoundException} exception is swallowed and `null` is returned
+                 * instead. We want to ensure cleanup of null entries in such case.
+                 */
+                bundle.putObject(key, null);
+                continue;
+            }
+            if (canMarshall(obj)) {
+                continue;
+            }
+
+            if (obj instanceof Bundle) {
+                sanitizeResponseParams((Bundle) obj);
+            } else if (obj instanceof ParcelFileDescriptor) {
+                validatePfdReadOnly((ParcelFileDescriptor) obj);
+            } else if (obj instanceof Bitmap) {
+                validateBitmap((Bitmap) obj);
+            } else if (obj instanceof Parcelable[]) {
+                validateParcelableArray((Parcelable[]) obj);
+            } else {
+                throw new BadParcelableException(
+                        "Unsupported Parcelable type encountered in the Bundle: "
+                                + obj.getClass().getSimpleName());
+            }
+        }
+    }
+
+    /**
+     * Validation of the inference request payload as described in {@link StateParams}
+     * description.
+     *
+     * @throws BadParcelableException when the bundle does not meet the read-only requirements.
+     */
+    public static void sanitizeStateParams(
+            @StateParams Bundle bundle) {
+        ensureValidBundle(bundle);
+
+        if (!bundle.hasFileDescriptors()) {
+            return; //safe to exit if there are no FDs and Binders
+        }
+
+        for (String key : bundle.keySet()) {
+            Object obj = bundle.get(key);
+            if (obj == null) {
+                /* Null value here could also mean deserializing a custom parcelable has failed,
+                 *  and since {@link Bundle} is marked as defusable in system-server - the
+                 * {@link ClassNotFoundException} exception is swallowed and `null` is returned
+                 * instead. We want to ensure cleanup of null entries in such case.
+                 */
+                bundle.putObject(key, null);
+                continue;
+            }
+            if (canMarshall(obj)) {
+                continue;
+            }
+
+            if (obj instanceof ParcelFileDescriptor) {
+                validatePfdReadOnly((ParcelFileDescriptor) obj);
+            } else {
+                throw new BadParcelableException(
+                        "Unsupported Parcelable type encountered in the Bundle: "
+                                + obj.getClass().getSimpleName());
+            }
+        }
+    }
+
+
+    public static IStreamingResponseCallback wrapWithValidation(
+            IStreamingResponseCallback streamingResponseCallback,
+            Executor resourceClosingExecutor,
+            AndroidFuture future,
+            InferenceInfoStore inferenceInfoStore) {
+        return new IStreamingResponseCallback.Stub() {
+            @Override
+            public void onNewContent(Bundle processedResult) throws RemoteException {
+                try {
+                    sanitizeResponseParams(processedResult);
+                    streamingResponseCallback.onNewContent(processedResult);
+                } finally {
+                    resourceClosingExecutor.execute(() -> tryCloseResource(processedResult));
+                }
+            }
+
+            @Override
+            public void onSuccess(Bundle resultBundle)
+                    throws RemoteException {
+                try {
+                    sanitizeResponseParams(resultBundle);
+                    streamingResponseCallback.onSuccess(resultBundle);
+                } finally {
+                    inferenceInfoStore.addInferenceInfoFromBundle(resultBundle);
+                    resourceClosingExecutor.execute(() -> tryCloseResource(resultBundle));
+                    future.complete(null);
+                }
+            }
+
+            @Override
+            public void onFailure(int errorCode, String errorMessage,
+                    PersistableBundle errorParams) throws RemoteException {
+                streamingResponseCallback.onFailure(errorCode, errorMessage, errorParams);
+                inferenceInfoStore.addInferenceInfoFromBundle(errorParams);
+                future.completeExceptionally(new TimeoutException());
+            }
+
+            @Override
+            public void onDataAugmentRequest(Bundle processedContent,
+                    RemoteCallback remoteCallback)
+                    throws RemoteException {
+                try {
+                    sanitizeResponseParams(processedContent);
+                    streamingResponseCallback.onDataAugmentRequest(processedContent,
+                            new RemoteCallback(
+                                    augmentedData -> {
+                                        try {
+                                            sanitizeInferenceParams(augmentedData);
+                                            remoteCallback.sendResult(augmentedData);
+                                        } finally {
+                                            resourceClosingExecutor.execute(
+                                                    () -> tryCloseResource(augmentedData));
+                                        }
+                                    }));
+                } finally {
+                    resourceClosingExecutor.execute(() -> tryCloseResource(processedContent));
+                }
+            }
+        };
+    }
+
+    public static IResponseCallback wrapWithValidation(IResponseCallback responseCallback,
+            Executor resourceClosingExecutor,
+            AndroidFuture future,
+            InferenceInfoStore inferenceInfoStore) {
+        return new IResponseCallback.Stub() {
+            @Override
+            public void onSuccess(Bundle resultBundle)
+                    throws RemoteException {
+                try {
+                    sanitizeResponseParams(resultBundle);
+                    responseCallback.onSuccess(resultBundle);
+                } finally {
+                    inferenceInfoStore.addInferenceInfoFromBundle(resultBundle);
+                    resourceClosingExecutor.execute(() -> tryCloseResource(resultBundle));
+                    future.complete(null);
+                }
+            }
+
+            @Override
+            public void onFailure(int errorCode, String errorMessage,
+                    PersistableBundle errorParams) throws RemoteException {
+                responseCallback.onFailure(errorCode, errorMessage, errorParams);
+                inferenceInfoStore.addInferenceInfoFromBundle(errorParams);
+                future.completeExceptionally(new TimeoutException());
+            }
+
+            @Override
+            public void onDataAugmentRequest(Bundle processedContent,
+                    RemoteCallback remoteCallback)
+                    throws RemoteException {
+                try {
+                    sanitizeResponseParams(processedContent);
+                    responseCallback.onDataAugmentRequest(processedContent, new RemoteCallback(
+                            augmentedData -> {
+                                try {
+                                    sanitizeInferenceParams(augmentedData);
+                                    remoteCallback.sendResult(augmentedData);
+                                } finally {
+                                    resourceClosingExecutor.execute(
+                                            () -> tryCloseResource(augmentedData));
+                                }
+                            }));
+                } finally {
+                    resourceClosingExecutor.execute(() -> tryCloseResource(processedContent));
+                }
+            }
+        };
+    }
+
+
+    public static ITokenInfoCallback wrapWithValidation(ITokenInfoCallback responseCallback,
+            AndroidFuture future,
+            InferenceInfoStore inferenceInfoStore) {
+        return new ITokenInfoCallback.Stub() {
+            @Override
+            public void onSuccess(TokenInfo tokenInfo) throws RemoteException {
+                responseCallback.onSuccess(tokenInfo);
+                inferenceInfoStore.addInferenceInfoFromBundle(tokenInfo.getInfoParams());
+                future.complete(null);
+            }
+
+            @Override
+            public void onFailure(int errorCode, String errorMessage, PersistableBundle errorParams)
+                    throws RemoteException {
+                responseCallback.onFailure(errorCode, errorMessage, errorParams);
+                inferenceInfoStore.addInferenceInfoFromBundle(errorParams);
+                future.completeExceptionally(new TimeoutException());
+            }
+        };
+    }
+
+    private static boolean canMarshall(Object obj) {
+        return obj instanceof byte[] || obj instanceof PersistableBundle
+                || PersistableBundle.isValidType(obj);
+    }
+
+    private static void ensureValidBundle(Bundle bundle) {
+        if (bundle == null) {
+            throw new IllegalArgumentException("Request passed is expected to be non-null");
+        }
+
+        if (bundle.hasBinders() != Bundle.STATUS_BINDERS_NOT_PRESENT) {
+            throw new BadParcelableException("Bundle should not contain IBinder objects.");
+        }
+    }
+
+    private static void validateParcelableArray(Parcelable[] parcelables) {
+        if (parcelables.length > 0
+                && parcelables[0] instanceof ParcelFileDescriptor) {
+            // Safe to cast
+            validatePfdsReadOnly(parcelables);
+        } else if (parcelables.length > 0
+                && parcelables[0] instanceof Bitmap) {
+            validateBitmapsImmutable(parcelables);
+        } else {
+            throw new BadParcelableException(
+                    "Could not cast to any known parcelable array");
+        }
+    }
+
+    public static void validatePfdsReadOnly(Parcelable[] pfds) {
+        for (Parcelable pfd : pfds) {
+            validatePfdReadOnly((ParcelFileDescriptor) pfd);
+        }
+    }
+
+    public static void validatePfdReadOnly(ParcelFileDescriptor pfd) {
+        if (pfd == null) {
+            return;
+        }
+        try {
+            int readMode = Os.fcntlInt(pfd.getFileDescriptor(), F_GETFL, 0) & O_ACCMODE;
+            if (readMode != O_RDONLY) {
+                throw new BadParcelableException(
+                        "Bundle contains a parcel file descriptor which is not read-only.");
+            }
+        } catch (ErrnoException e) {
+            throw new BadParcelableException(
+                    "Invalid File descriptor passed in the Bundle.", e);
+        }
+    }
+
+    private static void validateBitmap(Bitmap obj) {
+        if (obj.isMutable()) {
+            throw new BadParcelableException(
+                    "Encountered a mutable Bitmap in the Bundle at key : " + obj);
+        }
+    }
+
+    private static void validateBitmapsImmutable(Parcelable[] bitmaps) {
+        for (Parcelable bitmap : bitmaps) {
+            validateBitmap((Bitmap) bitmap);
+        }
+    }
+
+    public static void tryCloseResource(Bundle bundle) {
+        if (bundle == null || bundle.isEmpty() || !bundle.hasFileDescriptors()) {
+            return;
+        }
+
+        for (String key : bundle.keySet()) {
+            Object obj = bundle.get(key);
+
+            try {
+                // TODO(b/329898589) : This can be cleaned up after the flag passing is fixed.
+                if (obj instanceof ParcelFileDescriptor) {
+                    ((ParcelFileDescriptor) obj).close();
+                } else if (obj instanceof CursorWindow) {
+                    ((CursorWindow) obj).close();
+                } else if (obj instanceof SharedMemory) {
+                    // TODO(b/331796886) : Shared memory should honour parcelable flags.
+                    ((SharedMemory) obj).close();
+                }
+            } catch (Exception e) {
+                Log.e(TAG, "Error closing resource with key: " + key, e);
+            }
+        }
+    }
+}
diff --git a/packages/NeuralNetworks/service/platform/java/com/android/server/ondeviceintelligence/InferenceInfoStore.java b/packages/NeuralNetworks/service/platform/java/com/android/server/ondeviceintelligence/InferenceInfoStore.java
new file mode 100644
index 0000000..bef3f80
--- /dev/null
+++ b/packages/NeuralNetworks/service/platform/java/com/android/server/ondeviceintelligence/InferenceInfoStore.java
@@ -0,0 +1,101 @@
+/*
+ * 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.server.ondeviceintelligence;
+
+import android.app.ondeviceintelligence.InferenceInfo;
+import android.os.Bundle;
+import android.os.PersistableBundle;
+import android.service.ondeviceintelligence.OnDeviceSandboxedInferenceService;
+import android.util.Base64;
+import android.util.Slog;
+
+import java.io.IOException;
+import java.util.Comparator;
+import java.util.List;
+import java.util.TreeSet;
+
+public class InferenceInfoStore {
+    private static final String TAG = "InferenceInfoStore";
+    private final TreeSet<InferenceInfo> inferenceInfos;
+    private final long maxAgeMs;
+
+    public InferenceInfoStore(long maxAgeMs) {
+        this.maxAgeMs = maxAgeMs;
+        this.inferenceInfos = new TreeSet<>(
+                Comparator.comparingLong(InferenceInfo::getStartTimeMillis));
+    }
+
+    public List<InferenceInfo> getLatestInferenceInfo(long startTimeEpochMillis) {
+        return inferenceInfos.stream().filter(
+                info -> info.getStartTimeMillis() > startTimeEpochMillis).toList();
+    }
+
+    public void addInferenceInfoFromBundle(PersistableBundle pb) {
+        if (!pb.containsKey(OnDeviceSandboxedInferenceService.INFERENCE_INFO_BUNDLE_KEY)) {
+            return;
+        }
+
+        try {
+            String infoBytesBase64String = pb.getString(
+                    OnDeviceSandboxedInferenceService.INFERENCE_INFO_BUNDLE_KEY);
+            if (infoBytesBase64String != null) {
+                byte[] infoBytes = Base64.decode(infoBytesBase64String, Base64.DEFAULT);
+                com.android.server.ondeviceintelligence.nano.InferenceInfo inferenceInfo =
+                        com.android.server.ondeviceintelligence.nano.InferenceInfo.parseFrom(
+                                infoBytes);
+                add(inferenceInfo);
+            }
+        } catch (IOException e) {
+            Slog.e(TAG, "Unable to parse InferenceInfo from the received bytes.");
+        }
+    }
+
+    public void addInferenceInfoFromBundle(Bundle b) {
+        if (!b.containsKey(OnDeviceSandboxedInferenceService.INFERENCE_INFO_BUNDLE_KEY)) {
+            return;
+        }
+
+        try {
+            byte[] infoBytes = b.getByteArray(
+                    OnDeviceSandboxedInferenceService.INFERENCE_INFO_BUNDLE_KEY);
+            if (infoBytes != null) {
+                com.android.server.ondeviceintelligence.nano.InferenceInfo inferenceInfo =
+                        com.android.server.ondeviceintelligence.nano.InferenceInfo.parseFrom(
+                                infoBytes);
+                add(inferenceInfo);
+            }
+        } catch (IOException e) {
+            Slog.e(TAG, "Unable to parse InferenceInfo from the received bytes.");
+        }
+    }
+
+    private synchronized void add(com.android.server.ondeviceintelligence.nano.InferenceInfo info) {
+        while (!inferenceInfos.isEmpty()
+                && System.currentTimeMillis() - inferenceInfos.first().getStartTimeMillis()
+                > maxAgeMs) {
+            inferenceInfos.pollFirst();
+        }
+        inferenceInfos.add(toInferenceInfo(info));
+    }
+
+    private static InferenceInfo toInferenceInfo(
+            com.android.server.ondeviceintelligence.nano.InferenceInfo info) {
+        return new InferenceInfo.Builder(info.uid).setStartTimeMillis(
+                info.startTimeMs).setEndTimeMillis(info.endTimeMs).setSuspendedTimeMillis(
+                info.suspendedTimeMs).build();
+    }
+}
\ No newline at end of file
diff --git a/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerLocal.java b/packages/NeuralNetworks/service/platform/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerLocal.java
similarity index 100%
copy from packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerLocal.java
copy to packages/NeuralNetworks/service/platform/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerLocal.java
diff --git a/packages/NeuralNetworks/service/platform/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerService.java b/packages/NeuralNetworks/service/platform/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerService.java
new file mode 100644
index 0000000..0a69af6
--- /dev/null
+++ b/packages/NeuralNetworks/service/platform/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerService.java
@@ -0,0 +1,1096 @@
+/*
+ * 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.server.ondeviceintelligence;
+
+import static android.service.ondeviceintelligence.OnDeviceSandboxedInferenceService.DEVICE_CONFIG_UPDATE_BUNDLE_KEY;
+import static android.service.ondeviceintelligence.OnDeviceSandboxedInferenceService.MODEL_LOADED_BUNDLE_KEY;
+import static android.service.ondeviceintelligence.OnDeviceSandboxedInferenceService.MODEL_UNLOADED_BUNDLE_KEY;
+import static android.service.ondeviceintelligence.OnDeviceSandboxedInferenceService.MODEL_LOADED_BROADCAST_INTENT;
+import static android.service.ondeviceintelligence.OnDeviceSandboxedInferenceService.MODEL_UNLOADED_BROADCAST_INTENT;
+import static android.service.ondeviceintelligence.OnDeviceSandboxedInferenceService.REGISTER_MODEL_UPDATE_CALLBACK_BUNDLE_KEY;
+
+import static com.android.server.ondeviceintelligence.BundleUtil.sanitizeInferenceParams;
+import static com.android.server.ondeviceintelligence.BundleUtil.validatePfdReadOnly;
+import static com.android.server.ondeviceintelligence.BundleUtil.sanitizeStateParams;
+import static com.android.server.ondeviceintelligence.BundleUtil.wrapWithValidation;
+
+
+import android.Manifest;
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.annotation.RequiresPermission;
+import android.app.AppGlobals;
+import android.app.ondeviceintelligence.DownloadCallback;
+import android.app.ondeviceintelligence.Feature;
+import android.app.ondeviceintelligence.FeatureDetails;
+import android.app.ondeviceintelligence.IDownloadCallback;
+import android.app.ondeviceintelligence.IFeatureCallback;
+import android.app.ondeviceintelligence.IFeatureDetailsCallback;
+import android.app.ondeviceintelligence.IListFeaturesCallback;
+import android.app.ondeviceintelligence.IOnDeviceIntelligenceManager;
+import android.app.ondeviceintelligence.IProcessingSignal;
+import android.app.ondeviceintelligence.IResponseCallback;
+import android.app.ondeviceintelligence.IStreamingResponseCallback;
+import android.app.ondeviceintelligence.ITokenInfoCallback;
+import android.app.ondeviceintelligence.InferenceInfo;
+import android.app.ondeviceintelligence.OnDeviceIntelligenceException;
+import android.content.ComponentName;
+import android.content.Context;
+import android.content.Intent;
+import android.content.pm.PackageManager;
+import android.content.pm.ServiceInfo;
+import android.content.res.Resources;
+import android.os.Binder;
+import android.os.Bundle;
+import android.os.Handler;
+import android.os.IBinder;
+import android.os.ICancellationSignal;
+import android.os.IRemoteCallback;
+import android.os.Looper;
+import android.os.Message;
+import android.os.ParcelFileDescriptor;
+import android.os.PersistableBundle;
+import android.os.RemoteCallback;
+import android.os.RemoteException;
+import android.os.ResultReceiver;
+import android.os.ShellCallback;
+import android.os.UserHandle;
+import android.provider.DeviceConfig;
+import android.provider.Settings;
+import android.service.ondeviceintelligence.IOnDeviceIntelligenceService;
+import android.service.ondeviceintelligence.IOnDeviceSandboxedInferenceService;
+import android.service.ondeviceintelligence.IProcessingUpdateStatusCallback;
+import android.service.ondeviceintelligence.IRemoteProcessingService;
+import android.service.ondeviceintelligence.IRemoteStorageService;
+import android.service.ondeviceintelligence.OnDeviceIntelligenceService;
+import android.service.ondeviceintelligence.OnDeviceSandboxedInferenceService;
+import android.text.TextUtils;
+import android.util.Log;
+import android.util.Slog;
+
+import com.android.internal.R;
+import com.android.internal.annotations.GuardedBy;
+import com.android.internal.infra.AndroidFuture;
+import com.android.internal.infra.ServiceConnector;
+import com.android.internal.os.BackgroundThread;
+import com.android.server.LocalManagerRegistry;
+import com.android.server.SystemService;
+import com.android.server.SystemService.TargetUser;
+import com.android.server.ondeviceintelligence.callbacks.ListenableDownloadCallback;
+
+import java.io.FileDescriptor;
+import java.io.IOException;
+import java.util.List;
+import java.util.Objects;
+import java.util.Set;
+import java.util.concurrent.Executor;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+/**
+ * This is the system service for handling calls on the
+ * {@link android.app.ondeviceintelligence.OnDeviceIntelligenceManager}. This
+ * service holds connection references to the underlying remote services i.e. the isolated service
+ * {@link OnDeviceSandboxedInferenceService} and a regular
+ * service counter part {@link OnDeviceIntelligenceService}.
+ *
+ * Note: Both the remote services run under the SYSTEM user, as we cannot have separate instance of
+ * the Inference service for each user, due to possible high memory footprint.
+ *
+ * @hide
+ */
+public class OnDeviceIntelligenceManagerService extends SystemService {
+
+    private static final String TAG = OnDeviceIntelligenceManagerService.class.getSimpleName();
+    private static final String KEY_SERVICE_ENABLED = "service_enabled";
+
+    /** Handler message to {@link #resetTemporaryServices()} */
+    private static final int MSG_RESET_TEMPORARY_SERVICE = 0;
+    /** Handler message to clean up temporary broadcast keys. */
+    private static final int MSG_RESET_BROADCAST_KEYS = 1;
+    /** Handler message to clean up temporary config namespace. */
+    private static final int MSG_RESET_CONFIG_NAMESPACE = 2;
+
+    /** Default value in absence of {@link DeviceConfig} override. */
+    private static final boolean DEFAULT_SERVICE_ENABLED = true;
+    private static final String NAMESPACE_ON_DEVICE_INTELLIGENCE = "ondeviceintelligence";
+
+    private static final String SYSTEM_PACKAGE = "android";
+    private static final long MAX_AGE_MS = TimeUnit.HOURS.toMillis(3);
+
+
+    private final Executor resourceClosingExecutor = Executors.newCachedThreadPool();
+    private final Executor callbackExecutor = Executors.newCachedThreadPool();
+    private final Executor broadcastExecutor = Executors.newCachedThreadPool();
+    private final Executor mConfigExecutor = Executors.newCachedThreadPool();
+
+
+    private final Context mContext;
+    protected final Object mLock = new Object();
+
+    private final InferenceInfoStore mInferenceInfoStore;
+    private RemoteOnDeviceSandboxedInferenceService mRemoteInferenceService;
+    private RemoteOnDeviceIntelligenceService mRemoteOnDeviceIntelligenceService;
+    volatile boolean mIsServiceEnabled;
+
+    @GuardedBy("mLock")
+    private int remoteInferenceServiceUid = -1;
+
+    @GuardedBy("mLock")
+    private String[] mTemporaryServiceNames;
+    @GuardedBy("mLock")
+    private String[] mTemporaryBroadcastKeys;
+    @GuardedBy("mLock")
+    private String mBroadcastPackageName = SYSTEM_PACKAGE;
+    @GuardedBy("mLock")
+    private String mTemporaryConfigNamespace;
+
+    private final DeviceConfig.OnPropertiesChangedListener mOnPropertiesChangedListener =
+            this::sendUpdatedConfig;
+
+
+    /**
+     * Handler used to reset the temporary service names.
+     */
+    private Handler mTemporaryHandler;
+    private final @NonNull Handler mMainHandler = new Handler(Looper.getMainLooper());
+
+
+    public OnDeviceIntelligenceManagerService(Context context) {
+        super(context);
+        mContext = context;
+        mTemporaryServiceNames = new String[0];
+        mInferenceInfoStore = new InferenceInfoStore(MAX_AGE_MS);
+    }
+
+    @Override
+    public void onStart() {
+        publishBinderService(
+                Context.ON_DEVICE_INTELLIGENCE_SERVICE, getOnDeviceIntelligenceManagerService(),
+                /* allowIsolated = */ true);
+        LocalManagerRegistry.addManager(OnDeviceIntelligenceManagerLocal.class,
+                    this::getRemoteInferenceServiceUid);
+    }
+
+    @Override
+    public void onBootPhase(int phase) {
+        if (phase == SystemService.PHASE_SYSTEM_SERVICES_READY) {
+            DeviceConfig.addOnPropertiesChangedListener(
+                    NAMESPACE_ON_DEVICE_INTELLIGENCE,
+                    BackgroundThread.getExecutor(),
+                    (properties) -> onDeviceConfigChange(properties.getKeyset()));
+
+            mIsServiceEnabled = isServiceEnabled();
+        }
+    }
+
+    private void onDeviceConfigChange(@NonNull Set<String> keys) {
+        if (keys.contains(KEY_SERVICE_ENABLED)) {
+            mIsServiceEnabled = isServiceEnabled();
+        }
+    }
+
+    private boolean isServiceEnabled() {
+        return DeviceConfig.getBoolean(
+                NAMESPACE_ON_DEVICE_INTELLIGENCE,
+                KEY_SERVICE_ENABLED, DEFAULT_SERVICE_ENABLED);
+    }
+
+    private IBinder getOnDeviceIntelligenceManagerService() {
+        return new IOnDeviceIntelligenceManager.Stub() {
+            @Override
+            public String getRemoteServicePackageName() {
+                return OnDeviceIntelligenceManagerService.this.getRemoteConfiguredPackageName();
+            }
+
+            @Override
+            public List<InferenceInfo> getLatestInferenceInfo(long startTimeEpochMillis) {
+                mContext.enforceCallingPermission(
+                        Manifest.permission.DUMP, TAG);
+                return OnDeviceIntelligenceManagerService.this.getLatestInferenceInfo(
+                        startTimeEpochMillis);
+            }
+
+            @Override
+            public void getVersion(RemoteCallback remoteCallback) {
+                Slog.i(TAG, "OnDeviceIntelligenceManagerInternal getVersion");
+                Objects.requireNonNull(remoteCallback);
+                mContext.enforceCallingPermission(
+                        Manifest.permission.USE_ON_DEVICE_INTELLIGENCE, TAG);
+                if (!mIsServiceEnabled) {
+                    Slog.w(TAG, "Service not available");
+                    remoteCallback.sendResult(null);
+                    return;
+                }
+                ensureRemoteIntelligenceServiceInitialized();
+                mRemoteOnDeviceIntelligenceService.postAsync(
+                        service -> {
+                            AndroidFuture future = new AndroidFuture();
+                            service.getVersion(new RemoteCallback(
+                                    result -> {
+                                        remoteCallback.sendResult(result);
+                                        future.complete(null);
+                                    }));
+                            return future.orTimeout(getIdleTimeoutMs(), TimeUnit.MILLISECONDS);
+                        });
+            }
+
+            @Override
+            public void getFeature(int id, IFeatureCallback featureCallback)
+                    throws RemoteException {
+                Slog.i(TAG, "OnDeviceIntelligenceManagerInternal getFeatures");
+                Objects.requireNonNull(featureCallback);
+                mContext.enforceCallingPermission(
+                        Manifest.permission.USE_ON_DEVICE_INTELLIGENCE, TAG);
+                if (!mIsServiceEnabled) {
+                    Slog.w(TAG, "Service not available");
+                    featureCallback.onFailure(
+                            OnDeviceIntelligenceException.ON_DEVICE_INTELLIGENCE_SERVICE_UNAVAILABLE,
+                            "OnDeviceIntelligenceManagerService is unavailable",
+                            PersistableBundle.EMPTY);
+                    return;
+                }
+                ensureRemoteIntelligenceServiceInitialized();
+                int callerUid = Binder.getCallingUid();
+                mRemoteOnDeviceIntelligenceService.postAsync(
+                        service -> {
+                            AndroidFuture future = new AndroidFuture();
+                            service.getFeature(callerUid, id, new IFeatureCallback.Stub() {
+                                @Override
+                                public void onSuccess(Feature result) throws RemoteException {
+                                    featureCallback.onSuccess(result);
+                                    future.complete(null);
+                                }
+
+                                @Override
+                                public void onFailure(int errorCode, String errorMessage,
+                                        PersistableBundle errorParams) throws RemoteException {
+                                    featureCallback.onFailure(errorCode, errorMessage, errorParams);
+                                    future.completeExceptionally(new TimeoutException());
+                                }
+                            });
+                            return future.orTimeout(getIdleTimeoutMs(), TimeUnit.MILLISECONDS);
+                        });
+            }
+
+            @Override
+            public void listFeatures(IListFeaturesCallback listFeaturesCallback)
+                    throws RemoteException {
+                Slog.i(TAG, "OnDeviceIntelligenceManagerInternal getFeatures");
+                Objects.requireNonNull(listFeaturesCallback);
+                mContext.enforceCallingPermission(
+                        Manifest.permission.USE_ON_DEVICE_INTELLIGENCE, TAG);
+                if (!mIsServiceEnabled) {
+                    Slog.w(TAG, "Service not available");
+                    listFeaturesCallback.onFailure(
+                            OnDeviceIntelligenceException.ON_DEVICE_INTELLIGENCE_SERVICE_UNAVAILABLE,
+                            "OnDeviceIntelligenceManagerService is unavailable",
+                            PersistableBundle.EMPTY);
+                    return;
+                }
+                ensureRemoteIntelligenceServiceInitialized();
+                int callerUid = Binder.getCallingUid();
+                mRemoteOnDeviceIntelligenceService.postAsync(
+                        service -> {
+                            AndroidFuture future = new AndroidFuture();
+                            service.listFeatures(callerUid,
+                                    new IListFeaturesCallback.Stub() {
+                                        @Override
+                                        public void onSuccess(List<Feature> result)
+                                                throws RemoteException {
+                                            listFeaturesCallback.onSuccess(result);
+                                            future.complete(null);
+                                        }
+
+                                        @Override
+                                        public void onFailure(int errorCode, String errorMessage,
+                                                PersistableBundle errorParams)
+                                                throws RemoteException {
+                                            listFeaturesCallback.onFailure(errorCode, errorMessage,
+                                                    errorParams);
+                                            future.completeExceptionally(new TimeoutException());
+                                        }
+                                    });
+                            return future.orTimeout(getIdleTimeoutMs(), TimeUnit.MILLISECONDS);
+                        });
+            }
+
+            @Override
+            public void getFeatureDetails(Feature feature,
+                    IFeatureDetailsCallback featureDetailsCallback)
+                    throws RemoteException {
+                Slog.i(TAG, "OnDeviceIntelligenceManagerInternal getFeatureStatus");
+                Objects.requireNonNull(feature);
+                Objects.requireNonNull(featureDetailsCallback);
+                mContext.enforceCallingPermission(
+                        Manifest.permission.USE_ON_DEVICE_INTELLIGENCE, TAG);
+                if (!mIsServiceEnabled) {
+                    Slog.w(TAG, "Service not available");
+                    featureDetailsCallback.onFailure(
+                            OnDeviceIntelligenceException.ON_DEVICE_INTELLIGENCE_SERVICE_UNAVAILABLE,
+                            "OnDeviceIntelligenceManagerService is unavailable",
+                            PersistableBundle.EMPTY);
+                    return;
+                }
+                ensureRemoteIntelligenceServiceInitialized();
+                int callerUid = Binder.getCallingUid();
+                mRemoteOnDeviceIntelligenceService.postAsync(
+                        service -> {
+                            AndroidFuture future = new AndroidFuture();
+                            service.getFeatureDetails(callerUid, feature,
+                                    new IFeatureDetailsCallback.Stub() {
+                                        @Override
+                                        public void onSuccess(FeatureDetails result)
+                                                throws RemoteException {
+                                            future.complete(null);
+                                            featureDetailsCallback.onSuccess(result);
+                                        }
+
+                                        @Override
+                                        public void onFailure(int errorCode, String errorMessage,
+                                                PersistableBundle errorParams)
+                                                throws RemoteException {
+                                            future.completeExceptionally(null);
+                                            featureDetailsCallback.onFailure(errorCode,
+                                                    errorMessage, errorParams);
+                                        }
+                                    });
+                            return future.orTimeout(getIdleTimeoutMs(), TimeUnit.MILLISECONDS);
+                        });
+            }
+
+            @Override
+            public void requestFeatureDownload(Feature feature,
+                    AndroidFuture cancellationSignalFuture,
+                    IDownloadCallback downloadCallback) throws RemoteException {
+                Slog.i(TAG, "OnDeviceIntelligenceManagerInternal requestFeatureDownload");
+                Objects.requireNonNull(feature);
+                Objects.requireNonNull(downloadCallback);
+                mContext.enforceCallingPermission(
+                        Manifest.permission.USE_ON_DEVICE_INTELLIGENCE, TAG);
+                if (!mIsServiceEnabled) {
+                    Slog.w(TAG, "Service not available");
+                    downloadCallback.onDownloadFailed(
+                            DownloadCallback.DOWNLOAD_FAILURE_STATUS_UNAVAILABLE,
+                            "OnDeviceIntelligenceManagerService is unavailable",
+                            PersistableBundle.EMPTY);
+                }
+                ensureRemoteIntelligenceServiceInitialized();
+                int callerUid = Binder.getCallingUid();
+                mRemoteOnDeviceIntelligenceService.postAsync(
+                        service -> {
+                            AndroidFuture future = new AndroidFuture();
+                            ListenableDownloadCallback listenableDownloadCallback =
+                                    new ListenableDownloadCallback(
+                                            downloadCallback,
+                                            mMainHandler, future, getIdleTimeoutMs());
+                            service.requestFeatureDownload(callerUid, feature,
+                                    wrapCancellationFuture(cancellationSignalFuture),
+                                    listenableDownloadCallback);
+                            return future; // this future has no timeout because, actual download
+                            // might take long, but we fail early if there is no progress callbacks.
+                        }
+                );
+            }
+
+
+            @Override
+            public void requestTokenInfo(Feature feature,
+                    Bundle request,
+                    AndroidFuture cancellationSignalFuture,
+                    ITokenInfoCallback tokenInfoCallback) throws RemoteException {
+                Slog.i(TAG, "OnDeviceIntelligenceManagerInternal requestTokenInfo");
+                AndroidFuture<Void> result = null;
+                try {
+                    Objects.requireNonNull(feature);
+                    sanitizeInferenceParams(request);
+                    Objects.requireNonNull(tokenInfoCallback);
+
+                    mContext.enforceCallingPermission(
+                            Manifest.permission.USE_ON_DEVICE_INTELLIGENCE, TAG);
+                    if (!mIsServiceEnabled) {
+                        Slog.w(TAG, "Service not available");
+                        tokenInfoCallback.onFailure(
+                                OnDeviceIntelligenceException.ON_DEVICE_INTELLIGENCE_SERVICE_UNAVAILABLE,
+                                "OnDeviceIntelligenceManagerService is unavailable",
+                                PersistableBundle.EMPTY);
+                    }
+                    ensureRemoteInferenceServiceInitialized();
+                    int callerUid = Binder.getCallingUid();
+                    result = mRemoteInferenceService.postAsync(
+                            service -> {
+                                AndroidFuture future = new AndroidFuture();
+                                service.requestTokenInfo(callerUid, feature,
+                                        request,
+                                        wrapCancellationFuture(cancellationSignalFuture),
+                                        wrapWithValidation(tokenInfoCallback, future,
+                                                mInferenceInfoStore));
+                                return future.orTimeout(getIdleTimeoutMs(), TimeUnit.MILLISECONDS);
+                            });
+                    result.whenCompleteAsync((c, e) -> BundleUtil.tryCloseResource(request),
+                            resourceClosingExecutor);
+                } finally {
+                    if (result == null) {
+                        resourceClosingExecutor.execute(() -> BundleUtil.tryCloseResource(request));
+                    }
+                }
+            }
+
+            @Override
+            public void processRequest(Feature feature,
+                    Bundle request,
+                    int requestType,
+                    AndroidFuture cancellationSignalFuture,
+                    AndroidFuture processingSignalFuture,
+                    IResponseCallback responseCallback)
+                    throws RemoteException {
+                AndroidFuture<Void> result = null;
+                try {
+                    Slog.i(TAG, "OnDeviceIntelligenceManagerInternal processRequest");
+                    Objects.requireNonNull(feature);
+                    sanitizeInferenceParams(request);
+                    Objects.requireNonNull(responseCallback);
+                    mContext.enforceCallingPermission(
+                            Manifest.permission.USE_ON_DEVICE_INTELLIGENCE, TAG);
+                    if (!mIsServiceEnabled) {
+                        Slog.w(TAG, "Service not available");
+                        responseCallback.onFailure(
+                                OnDeviceIntelligenceException.PROCESSING_ERROR_SERVICE_UNAVAILABLE,
+                                "OnDeviceIntelligenceManagerService is unavailable",
+                                PersistableBundle.EMPTY);
+                    }
+                    ensureRemoteInferenceServiceInitialized();
+                    int callerUid = Binder.getCallingUid();
+                    result = mRemoteInferenceService.postAsync(
+                            service -> {
+                                AndroidFuture future = new AndroidFuture();
+                                service.processRequest(callerUid, feature,
+                                        request,
+                                        requestType,
+                                        wrapCancellationFuture(cancellationSignalFuture),
+                                        wrapProcessingFuture(processingSignalFuture),
+                                        wrapWithValidation(responseCallback,
+                                                resourceClosingExecutor, future,
+                                                mInferenceInfoStore));
+                                return future.orTimeout(getIdleTimeoutMs(), TimeUnit.MILLISECONDS);
+                            });
+                    result.whenCompleteAsync((c, e) -> BundleUtil.tryCloseResource(request),
+                            resourceClosingExecutor);
+                } finally {
+                    if (result == null) {
+                        resourceClosingExecutor.execute(() -> BundleUtil.tryCloseResource(request));
+                    }
+                }
+            }
+
+            @Override
+            public void processRequestStreaming(Feature feature,
+                    Bundle request,
+                    int requestType,
+                    AndroidFuture cancellationSignalFuture,
+                    AndroidFuture processingSignalFuture,
+                    IStreamingResponseCallback streamingCallback) throws RemoteException {
+                AndroidFuture<Void> result = null;
+                try {
+                    Slog.i(TAG, "OnDeviceIntelligenceManagerInternal processRequestStreaming");
+                    Objects.requireNonNull(feature);
+                    sanitizeInferenceParams(request);
+                    Objects.requireNonNull(streamingCallback);
+                    mContext.enforceCallingPermission(
+                            Manifest.permission.USE_ON_DEVICE_INTELLIGENCE, TAG);
+                    if (!mIsServiceEnabled) {
+                        Slog.w(TAG, "Service not available");
+                        streamingCallback.onFailure(
+                                OnDeviceIntelligenceException.PROCESSING_ERROR_SERVICE_UNAVAILABLE,
+                                "OnDeviceIntelligenceManagerService is unavailable",
+                                PersistableBundle.EMPTY);
+                    }
+                    ensureRemoteInferenceServiceInitialized();
+                    int callerUid = Binder.getCallingUid();
+                    result = mRemoteInferenceService.postAsync(
+                            service -> {
+                                AndroidFuture future = new AndroidFuture();
+                                service.processRequestStreaming(callerUid,
+                                        feature,
+                                        request, requestType,
+                                        wrapCancellationFuture(cancellationSignalFuture),
+                                        wrapProcessingFuture(processingSignalFuture),
+                                        wrapWithValidation(streamingCallback,
+                                                resourceClosingExecutor, future,
+                                                mInferenceInfoStore));
+                                return future.orTimeout(getIdleTimeoutMs(), TimeUnit.MILLISECONDS);
+                            });
+                    result.whenCompleteAsync((c, e) -> BundleUtil.tryCloseResource(request),
+                            resourceClosingExecutor);
+                } finally {
+                    if (result == null) {
+                        resourceClosingExecutor.execute(() -> BundleUtil.tryCloseResource(request));
+                    }
+                }
+            }
+
+            @Override
+            public void onShellCommand(FileDescriptor in, FileDescriptor out, FileDescriptor err,
+                    String[] args, ShellCallback callback, ResultReceiver resultReceiver) {
+                new OnDeviceIntelligenceShellCommand(OnDeviceIntelligenceManagerService.this).exec(
+                        this, in, out, err, args, callback, resultReceiver);
+            }
+        };
+    }
+
+    private void ensureRemoteIntelligenceServiceInitialized() {
+        synchronized (mLock) {
+            if (mRemoteOnDeviceIntelligenceService == null) {
+                String serviceName = getServiceNames()[0];
+                Binder.withCleanCallingIdentity(() -> validateServiceElevated(serviceName, false));
+                mRemoteOnDeviceIntelligenceService = new RemoteOnDeviceIntelligenceService(mContext,
+                        ComponentName.unflattenFromString(serviceName),
+                        UserHandle.SYSTEM.getIdentifier());
+                mRemoteOnDeviceIntelligenceService.setServiceLifecycleCallbacks(
+                        new ServiceConnector.ServiceLifecycleCallbacks<>() {
+                            @Override
+                            public void onConnected(
+                                    @NonNull IOnDeviceIntelligenceService service) {
+                                try {
+                                    service.registerRemoteServices(
+                                            getRemoteProcessingService());
+                                    service.ready();
+                                } catch (RemoteException ex) {
+                                    Slog.w(TAG, "Failed to send connected event", ex);
+                                }
+                            }
+                        });
+            }
+        }
+    }
+
+    @NonNull
+    private IRemoteProcessingService.Stub getRemoteProcessingService() {
+        return new IRemoteProcessingService.Stub() {
+            @Override
+            public void updateProcessingState(
+                    Bundle processingState,
+                    IProcessingUpdateStatusCallback callback) {
+                callbackExecutor.execute(() -> {
+                    AndroidFuture<Void> result = null;
+                    try {
+                        sanitizeStateParams(processingState);
+                        ensureRemoteInferenceServiceInitialized();
+                        result = mRemoteInferenceService.post(
+                                service -> service.updateProcessingState(
+                                        processingState, callback));
+                        result.whenCompleteAsync(
+                                (c, e) -> BundleUtil.tryCloseResource(processingState),
+                                resourceClosingExecutor);
+                    } finally {
+                        if (result == null) {
+                            resourceClosingExecutor.execute(
+                                    () -> BundleUtil.tryCloseResource(processingState));
+                        }
+                    }
+                });
+            }
+        };
+    }
+
+    private void ensureRemoteInferenceServiceInitialized() {
+        synchronized (mLock) {
+            if (mRemoteInferenceService == null) {
+                String serviceName = getServiceNames()[1];
+                Binder.withCleanCallingIdentity(() -> validateServiceElevated(serviceName, true));
+                mRemoteInferenceService = new RemoteOnDeviceSandboxedInferenceService(mContext,
+                        ComponentName.unflattenFromString(serviceName),
+                        UserHandle.SYSTEM.getIdentifier());
+                mRemoteInferenceService.setServiceLifecycleCallbacks(
+                        new ServiceConnector.ServiceLifecycleCallbacks<>() {
+                            @Override
+                            public void onConnected(
+                                    @NonNull IOnDeviceSandboxedInferenceService service) {
+                                try {
+                                    ensureRemoteIntelligenceServiceInitialized();
+                                    service.registerRemoteStorageService(
+                                            getIRemoteStorageService(), new IRemoteCallback.Stub() {
+                                                @Override
+                                                public void sendResult(Bundle bundle) {
+                                                    final int uid = Binder.getCallingUid();
+                                                    setRemoteInferenceServiceUid(uid);
+                                                }
+                                            });
+                                    mRemoteOnDeviceIntelligenceService.run(
+                                            IOnDeviceIntelligenceService::notifyInferenceServiceConnected);
+                                    broadcastExecutor.execute(
+                                            () -> registerModelLoadingBroadcasts(service));
+                                    mConfigExecutor.execute(
+                                            () -> registerDeviceConfigChangeListener());
+                                } catch (RemoteException ex) {
+                                    Slog.w(TAG, "Failed to send connected event", ex);
+                                }
+                            }
+
+                            @Override
+                            public void onDisconnected(
+                                    @NonNull IOnDeviceSandboxedInferenceService service) {
+                                ensureRemoteIntelligenceServiceInitialized();
+                                mRemoteOnDeviceIntelligenceService.run(
+                                        IOnDeviceIntelligenceService::notifyInferenceServiceDisconnected);
+                            }
+
+                            @Override
+                            public void onBinderDied() {
+                                ensureRemoteIntelligenceServiceInitialized();
+                                mRemoteOnDeviceIntelligenceService.run(
+                                        IOnDeviceIntelligenceService::notifyInferenceServiceDisconnected);
+                            }
+                        });
+            }
+        }
+    }
+
+    private void registerModelLoadingBroadcasts(IOnDeviceSandboxedInferenceService service) {
+        String[] modelBroadcastKeys;
+        try {
+            modelBroadcastKeys = getBroadcastKeys();
+        } catch (Resources.NotFoundException e) {
+            Slog.d(TAG, "Skipping model broadcasts as broadcast intents configured.");
+            return;
+        }
+
+        Bundle bundle = new Bundle();
+        bundle.putBoolean(REGISTER_MODEL_UPDATE_CALLBACK_BUNDLE_KEY, true);
+        try {
+            service.updateProcessingState(bundle, new IProcessingUpdateStatusCallback.Stub() {
+                @Override
+                public void onSuccess(PersistableBundle statusParams) {
+                    Binder.clearCallingIdentity();
+                    synchronized (mLock) {
+                        if (statusParams.containsKey(MODEL_LOADED_BUNDLE_KEY)) {
+                            String modelLoadedBroadcastKey = modelBroadcastKeys[0];
+                            if (modelLoadedBroadcastKey != null
+                                    && !modelLoadedBroadcastKey.isEmpty()) {
+                                final Intent intent = new Intent(modelLoadedBroadcastKey);
+                                intent.setPackage(mBroadcastPackageName);
+                                mContext.sendBroadcast(intent,
+                                        Manifest.permission.USE_ON_DEVICE_INTELLIGENCE);
+                            }
+                        } else if (statusParams.containsKey(MODEL_UNLOADED_BUNDLE_KEY)) {
+                            String modelUnloadedBroadcastKey = modelBroadcastKeys[1];
+                            if (modelUnloadedBroadcastKey != null
+                                    && !modelUnloadedBroadcastKey.isEmpty()) {
+                                final Intent intent = new Intent(modelUnloadedBroadcastKey);
+                                intent.setPackage(mBroadcastPackageName);
+                                mContext.sendBroadcast(intent,
+                                        Manifest.permission.USE_ON_DEVICE_INTELLIGENCE);
+                            }
+                        }
+                    }
+                }
+
+                @Override
+                public void onFailure(int errorCode, String errorMessage) {
+                    Slog.e(TAG, "Failed to register model loading callback with status code",
+                            new OnDeviceIntelligenceException(errorCode, errorMessage));
+                }
+            });
+        } catch (RemoteException e) {
+            Slog.e(TAG, "Failed to register model loading callback with status code", e);
+        }
+    }
+
+    private void registerDeviceConfigChangeListener() {
+        Log.d(TAG, "registerDeviceConfigChangeListener");
+        String configNamespace = getConfigNamespace();
+        if (configNamespace.isEmpty()) {
+            Slog.e(TAG, "config_defaultOnDeviceIntelligenceDeviceConfigNamespace is empty");
+            return;
+        }
+        DeviceConfig.addOnPropertiesChangedListener(
+                configNamespace,
+                mConfigExecutor,
+                mOnPropertiesChangedListener);
+    }
+
+    private String getConfigNamespace() {
+        synchronized (mLock) {
+            if (mTemporaryConfigNamespace != null) {
+                return mTemporaryConfigNamespace;
+            }
+
+            return mContext.getResources().getString(
+                    R.string.config_defaultOnDeviceIntelligenceDeviceConfigNamespace);
+        }
+    }
+
+    private void sendUpdatedConfig(
+            DeviceConfig.Properties props) {
+        Log.d(TAG, "sendUpdatedConfig");
+
+        PersistableBundle persistableBundle = new PersistableBundle();
+        for (String key : props.getKeyset()) {
+            persistableBundle.putString(key, props.getString(key, ""));
+        }
+        Bundle bundle = new Bundle();
+        bundle.putParcelable(DEVICE_CONFIG_UPDATE_BUNDLE_KEY, persistableBundle);
+        ensureRemoteInferenceServiceInitialized();
+        mRemoteInferenceService.run(service -> service.updateProcessingState(bundle,
+                new IProcessingUpdateStatusCallback.Stub() {
+                    @Override
+                    public void onSuccess(PersistableBundle result) {
+                        Slog.d(TAG, "Config update successful." + result);
+                    }
+
+                    @Override
+                    public void onFailure(int errorCode, String errorMessage) {
+                        Slog.e(TAG, "Config update failed with code ["
+                                + String.valueOf(errorCode) + "] and message = " + errorMessage);
+                    }
+                }));
+    }
+
+    @NonNull
+    private IRemoteStorageService.Stub getIRemoteStorageService() {
+        return new IRemoteStorageService.Stub() {
+            @Override
+            public void getReadOnlyFileDescriptor(
+                    String filePath,
+                    AndroidFuture<ParcelFileDescriptor> future) {
+                ensureRemoteIntelligenceServiceInitialized();
+                AndroidFuture<ParcelFileDescriptor> pfdFuture = new AndroidFuture<>();
+                mRemoteOnDeviceIntelligenceService.run(
+                        service -> service.getReadOnlyFileDescriptor(
+                                filePath, pfdFuture));
+                pfdFuture.whenCompleteAsync((pfd, error) -> {
+                    try {
+                        if (error != null) {
+                            future.completeExceptionally(error);
+                        } else {
+                            validatePfdReadOnly(pfd);
+                            future.complete(pfd);
+                        }
+                    } finally {
+                        tryClosePfd(pfd);
+                    }
+                }, callbackExecutor);
+            }
+
+            @Override
+            public void getReadOnlyFeatureFileDescriptorMap(
+                    Feature feature,
+                    RemoteCallback remoteCallback) {
+                ensureRemoteIntelligenceServiceInitialized();
+                mRemoteOnDeviceIntelligenceService.run(
+                        service -> service.getReadOnlyFeatureFileDescriptorMap(
+                                feature,
+                                new RemoteCallback(result -> callbackExecutor.execute(() -> {
+                                    try {
+                                        if (result == null) {
+                                            remoteCallback.sendResult(null);
+                                        }
+                                        for (String key : result.keySet()) {
+                                            ParcelFileDescriptor pfd = result.getParcelable(key,
+                                                    ParcelFileDescriptor.class);
+                                            validatePfdReadOnly(pfd);
+                                        }
+                                        remoteCallback.sendResult(result);
+                                    } finally {
+                                        resourceClosingExecutor.execute(
+                                                () -> BundleUtil.tryCloseResource(result));
+                                    }
+                                }))));
+            }
+        };
+    }
+
+    private void validateServiceElevated(String serviceName, boolean checkIsolated) {
+        try {
+            if (TextUtils.isEmpty(serviceName)) {
+                throw new IllegalStateException(
+                        "Remote service is not configured to complete the request");
+            }
+            ComponentName serviceComponent = ComponentName.unflattenFromString(
+                    serviceName);
+            ServiceInfo serviceInfo = AppGlobals.getPackageManager().getServiceInfo(
+                    serviceComponent,
+                    PackageManager.MATCH_DIRECT_BOOT_AWARE
+                            | PackageManager.MATCH_DIRECT_BOOT_UNAWARE,
+                    UserHandle.SYSTEM.getIdentifier());
+            if (serviceInfo != null) {
+                if (!checkIsolated) {
+                    checkServiceRequiresPermission(serviceInfo,
+                            Manifest.permission.BIND_ON_DEVICE_INTELLIGENCE_SERVICE);
+                    return;
+                }
+
+                checkServiceRequiresPermission(serviceInfo,
+                        Manifest.permission.BIND_ON_DEVICE_SANDBOXED_INFERENCE_SERVICE);
+                if (!isIsolatedService(serviceInfo)) {
+                    throw new SecurityException(
+                            "Call required an isolated service, but the configured service: "
+                                    + serviceName + ", is not isolated");
+                }
+            } else {
+                throw new IllegalStateException(
+                        "Remote service is not configured to complete the request.");
+            }
+        } catch (RemoteException e) {
+            throw new IllegalStateException("Could not fetch service info for remote services", e);
+        }
+    }
+
+    private static void checkServiceRequiresPermission(ServiceInfo serviceInfo,
+            String requiredPermission) {
+        final String permission = serviceInfo.permission;
+        if (!requiredPermission.equals(permission)) {
+            throw new SecurityException(String.format(
+                    "Service %s requires %s permission. Found %s permission",
+                    serviceInfo.getComponentName(),
+                    requiredPermission,
+                    serviceInfo.permission));
+        }
+    }
+
+    private static boolean isIsolatedService(@NonNull ServiceInfo serviceInfo) {
+        return (serviceInfo.flags & ServiceInfo.FLAG_ISOLATED_PROCESS) != 0
+                && (serviceInfo.flags & ServiceInfo.FLAG_EXTERNAL_SERVICE) == 0;
+    }
+
+    private List<InferenceInfo> getLatestInferenceInfo(long startTimeEpochMillis) {
+        return mInferenceInfoStore.getLatestInferenceInfo(startTimeEpochMillis);
+    }
+
+    @Nullable
+    public String getRemoteConfiguredPackageName() {
+        try {
+            String[] serviceNames = getServiceNames();
+            ComponentName componentName = ComponentName.unflattenFromString(serviceNames[1]);
+            if (componentName != null) {
+                return componentName.getPackageName();
+            }
+        } catch (Resources.NotFoundException e) {
+            Slog.e(TAG, "Could not find resource", e);
+        }
+
+        return null;
+    }
+
+
+    protected String[] getServiceNames() throws Resources.NotFoundException {
+        // TODO 329240495 : Consider a small class with explicit field names for the two services
+        synchronized (mLock) {
+            if (mTemporaryServiceNames != null && mTemporaryServiceNames.length == 2) {
+                return mTemporaryServiceNames;
+            }
+        }
+        return new String[]{mContext.getResources().getString(
+                R.string.config_defaultOnDeviceIntelligenceService),
+                mContext.getResources().getString(
+                        R.string.config_defaultOnDeviceSandboxedInferenceService)};
+    }
+
+    protected String[] getBroadcastKeys() throws Resources.NotFoundException {
+        // TODO 329240495 : Consider a small class with explicit field names for the two services
+        synchronized (mLock) {
+            if (mTemporaryBroadcastKeys != null && mTemporaryBroadcastKeys.length == 2) {
+                return mTemporaryBroadcastKeys;
+            }
+        }
+
+        return new String[]{ MODEL_LOADED_BROADCAST_INTENT, MODEL_UNLOADED_BROADCAST_INTENT };
+    }
+
+    @RequiresPermission(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE)
+    public void setTemporaryServices(@NonNull String[] componentNames, int durationMs) {
+        Objects.requireNonNull(componentNames);
+        enforceShellOnly(Binder.getCallingUid(), "setTemporaryServices");
+        mContext.enforceCallingPermission(
+                Manifest.permission.USE_ON_DEVICE_INTELLIGENCE, TAG);
+        synchronized (mLock) {
+            mTemporaryServiceNames = componentNames;
+            if (mRemoteInferenceService != null) {
+                mRemoteInferenceService.unbind();
+                mRemoteInferenceService = null;
+            }
+            if (mRemoteOnDeviceIntelligenceService != null) {
+                mRemoteOnDeviceIntelligenceService.unbind();
+                mRemoteOnDeviceIntelligenceService = null;
+            }
+
+            if (durationMs != -1) {
+                getTemporaryHandler().sendEmptyMessageDelayed(MSG_RESET_TEMPORARY_SERVICE,
+                        durationMs);
+            }
+        }
+    }
+
+    @RequiresPermission(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE)
+    public void setModelBroadcastKeys(@NonNull String[] broadcastKeys, String receiverPackageName,
+            int durationMs) {
+        Objects.requireNonNull(broadcastKeys);
+        enforceShellOnly(Binder.getCallingUid(), "setModelBroadcastKeys");
+        mContext.enforceCallingPermission(
+                Manifest.permission.USE_ON_DEVICE_INTELLIGENCE, TAG);
+        synchronized (mLock) {
+            mTemporaryBroadcastKeys = broadcastKeys;
+            mBroadcastPackageName = receiverPackageName;
+            if (durationMs != -1) {
+                getTemporaryHandler().sendEmptyMessageDelayed(MSG_RESET_BROADCAST_KEYS, durationMs);
+            }
+        }
+    }
+
+    @RequiresPermission(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE)
+    public void setTemporaryDeviceConfigNamespace(@NonNull String configNamespace,
+            int durationMs) {
+        Objects.requireNonNull(configNamespace);
+        enforceShellOnly(Binder.getCallingUid(), "setTemporaryDeviceConfigNamespace");
+        mContext.enforceCallingPermission(
+                Manifest.permission.USE_ON_DEVICE_INTELLIGENCE, TAG);
+        synchronized (mLock) {
+            mTemporaryConfigNamespace = configNamespace;
+            if (durationMs != -1) {
+                getTemporaryHandler().sendEmptyMessageDelayed(MSG_RESET_CONFIG_NAMESPACE,
+                        durationMs);
+            }
+        }
+    }
+
+    /**
+     * Reset the temporary services set in CTS tests, this method is primarily used to only revert
+     * the changes caused by CTS tests.
+     */
+    public void resetTemporaryServices() {
+        synchronized (mLock) {
+            if (mTemporaryHandler != null) {
+                mTemporaryHandler.removeMessages(MSG_RESET_TEMPORARY_SERVICE);
+                mTemporaryHandler = null;
+            }
+
+            mRemoteInferenceService = null;
+            mRemoteOnDeviceIntelligenceService = null;
+            mTemporaryServiceNames = new String[0];
+        }
+    }
+
+    /**
+     * Throws if the caller is not of a shell (or root) UID.
+     *
+     * @param callingUid pass Binder.callingUid().
+     */
+    public static void enforceShellOnly(int callingUid, String message) {
+        if (callingUid == android.os.Process.SHELL_UID
+                || callingUid == android.os.Process.ROOT_UID) {
+            return; // okay
+        }
+
+        throw new SecurityException(message + ": Only shell user can call it");
+    }
+
+    private AndroidFuture<IBinder> wrapCancellationFuture(
+            AndroidFuture future) {
+        if (future == null) {
+            return null;
+        }
+        AndroidFuture<IBinder> cancellationFuture = new AndroidFuture<>();
+        cancellationFuture.whenCompleteAsync((c, e) -> {
+            if (e != null) {
+                Log.e(TAG, "Error forwarding ICancellationSignal to manager layer", e);
+                future.completeExceptionally(e);
+            } else {
+                future.complete(new ICancellationSignal.Stub() {
+                    @Override
+                    public void cancel() throws RemoteException {
+                        ICancellationSignal.Stub.asInterface(c).cancel();
+                    }
+                });
+            }
+        });
+        return cancellationFuture;
+    }
+
+    private AndroidFuture<IBinder> wrapProcessingFuture(
+            AndroidFuture future) {
+        if (future == null) {
+            return null;
+        }
+        AndroidFuture<IBinder> processingSignalFuture = new AndroidFuture<>();
+        processingSignalFuture.whenCompleteAsync((c, e) -> {
+            if (e != null) {
+                future.completeExceptionally(e);
+            } else {
+                future.complete(new IProcessingSignal.Stub() {
+                    @Override
+                    public void sendSignal(PersistableBundle actionParams) throws RemoteException {
+                        IProcessingSignal.Stub.asInterface(c).sendSignal(actionParams);
+                    }
+                });
+            }
+        });
+        return processingSignalFuture;
+    }
+
+    private static void tryClosePfd(ParcelFileDescriptor pfd) {
+        if (pfd != null) {
+            try {
+                pfd.close();
+            } catch (IOException e) {
+                Log.e(TAG, "Failed to close parcel file descriptor ", e);
+            }
+        }
+    }
+
+    private synchronized Handler getTemporaryHandler() {
+        if (mTemporaryHandler == null) {
+            mTemporaryHandler = new Handler(Looper.getMainLooper(), null, true) {
+                @Override
+                public void handleMessage(Message msg) {
+                    synchronized (mLock) {
+                        if (msg.what == MSG_RESET_TEMPORARY_SERVICE) {
+                            resetTemporaryServices();
+                        } else if (msg.what == MSG_RESET_BROADCAST_KEYS) {
+                            mTemporaryBroadcastKeys = null;
+                            mBroadcastPackageName = SYSTEM_PACKAGE;
+                        } else if (msg.what == MSG_RESET_CONFIG_NAMESPACE) {
+                            mTemporaryConfigNamespace = null;
+                        } else {
+                            Slog.wtf(TAG, "invalid handler msg: " + msg);
+                        }
+                    }
+                }
+            };
+        }
+
+        return mTemporaryHandler;
+    }
+
+    private long getIdleTimeoutMs() {
+        return Settings.Secure.getLongForUser(mContext.getContentResolver(),
+                Settings.Secure.ON_DEVICE_INTELLIGENCE_IDLE_TIMEOUT_MS, TimeUnit.HOURS.toMillis(1),
+                mContext.getUserId());
+    }
+
+    private int getRemoteInferenceServiceUid() {
+        synchronized (mLock) {
+            return remoteInferenceServiceUid;
+        }
+    }
+
+    private void setRemoteInferenceServiceUid(int remoteInferenceServiceUid) {
+        synchronized (mLock) {
+            this.remoteInferenceServiceUid = remoteInferenceServiceUid;
+        }
+    }
+}
diff --git a/packages/NeuralNetworks/service/platform/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceShellCommand.java b/packages/NeuralNetworks/service/platform/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceShellCommand.java
new file mode 100644
index 0000000..d2c84fa
--- /dev/null
+++ b/packages/NeuralNetworks/service/platform/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceShellCommand.java
@@ -0,0 +1,144 @@
+/*
+ * Copyright (C) 2022 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.server.ondeviceintelligence;
+
+import android.annotation.NonNull;
+import android.os.Binder;
+import android.os.ShellCommand;
+
+import java.io.PrintWriter;
+import java.util.Objects;
+
+final class OnDeviceIntelligenceShellCommand extends ShellCommand {
+    private static final String TAG = OnDeviceIntelligenceShellCommand.class.getSimpleName();
+
+    @NonNull
+    private final OnDeviceIntelligenceManagerService mService;
+
+    OnDeviceIntelligenceShellCommand(@NonNull OnDeviceIntelligenceManagerService service) {
+        mService = service;
+    }
+
+    @Override
+    public int onCommand(String cmd) {
+        if (cmd == null) {
+            return handleDefaultCommands(cmd);
+        }
+
+        switch (cmd) {
+            case "set-temporary-services":
+                return setTemporaryServices();
+            case "get-services":
+                return getConfiguredServices();
+            case "set-model-broadcasts":
+                return setBroadcastKeys();
+            case "set-deviceconfig-namespace":
+                return setDeviceConfigNamespace();
+            default:
+                return handleDefaultCommands(cmd);
+        }
+    }
+
+    @Override
+    public void onHelp() {
+        PrintWriter pw = getOutPrintWriter();
+        pw.println("OnDeviceIntelligenceShellCommand commands: ");
+        pw.println("  help");
+        pw.println("    Print this help text.");
+        pw.println();
+        pw.println(
+                "  set-temporary-services [IntelligenceServiceComponentName] "
+                        + "[InferenceServiceComponentName] [DURATION]");
+        pw.println("    Temporarily (for DURATION ms) changes the service implementations.");
+        pw.println("    To reset, call without any arguments.");
+
+        pw.println("  get-services To get the names of services that are currently being used.");
+        pw.println(
+                "  set-model-broadcasts [ModelLoadedBroadcastKey] [ModelUnloadedBroadcastKey] "
+                        + "[ReceiverPackageName] "
+                        + "[DURATION] To set the names of broadcast intent keys that are to be "
+                        + "emitted for cts tests.");
+        pw.println(
+                "  set-deviceconfig-namespace [DeviceConfigNamespace] "
+                        + "[DURATION] To set the device config namespace "
+                        + "to use for cts tests.");
+    }
+
+    private int setTemporaryServices() {
+        final PrintWriter out = getOutPrintWriter();
+        final String intelligenceServiceName = getNextArg();
+        final String inferenceServiceName = getNextArg();
+
+        if (getRemainingArgsCount() == 0 && intelligenceServiceName == null
+                && inferenceServiceName == null) {
+            OnDeviceIntelligenceManagerService.enforceShellOnly(Binder.getCallingUid(),
+                    "resetTemporaryServices");
+            mService.resetTemporaryServices();
+            out.println("OnDeviceIntelligenceManagerService temporary reset. ");
+            return 0;
+        }
+
+        Objects.requireNonNull(intelligenceServiceName);
+        Objects.requireNonNull(inferenceServiceName);
+        final int duration = Integer.parseInt(getNextArgRequired());
+        mService.setTemporaryServices(
+                new String[]{intelligenceServiceName, inferenceServiceName},
+                duration);
+        out.println("OnDeviceIntelligenceService temporarily set to " + intelligenceServiceName
+                + " \n and \n OnDeviceTrustedInferenceService set to " + inferenceServiceName
+                + " for " + duration + "ms");
+        return 0;
+    }
+
+    private int getConfiguredServices() {
+        final PrintWriter out = getOutPrintWriter();
+        String[] services = mService.getServiceNames();
+        out.println("OnDeviceIntelligenceService set to :  " + services[0]
+                + " \n and \n OnDeviceTrustedInferenceService set to : " + services[1]);
+        return 0;
+    }
+
+    private int setBroadcastKeys() {
+        final PrintWriter out = getOutPrintWriter();
+        final String modelLoadedKey = getNextArgRequired();
+        final String modelUnloadedKey = getNextArgRequired();
+        final String receiverPackageName = getNextArg();
+
+        final int duration = Integer.parseInt(getNextArgRequired());
+        mService.setModelBroadcastKeys(
+                new String[]{modelLoadedKey, modelUnloadedKey}, receiverPackageName, duration);
+        out.println("OnDeviceIntelligence Model Loading broadcast keys temporarily set to "
+                + modelLoadedKey
+                + " \n and \n OnDeviceTrustedInferenceService set to " + modelUnloadedKey
+                + "\n and Package name set to : " + receiverPackageName
+                + " for " + duration + "ms");
+        return 0;
+    }
+
+    private int setDeviceConfigNamespace() {
+        final PrintWriter out = getOutPrintWriter();
+        final String configNamespace = getNextArg();
+
+        final int duration = Integer.parseInt(getNextArgRequired());
+        mService.setTemporaryDeviceConfigNamespace(configNamespace, duration);
+        out.println("OnDeviceIntelligence DeviceConfig Namespace temporarily set to "
+                + configNamespace
+                + " for " + duration + "ms");
+        return 0;
+    }
+
+}
\ No newline at end of file
diff --git a/packages/NeuralNetworks/service/platform/java/com/android/server/ondeviceintelligence/RemoteOnDeviceIntelligenceService.java b/packages/NeuralNetworks/service/platform/java/com/android/server/ondeviceintelligence/RemoteOnDeviceIntelligenceService.java
new file mode 100644
index 0000000..ac9747a
--- /dev/null
+++ b/packages/NeuralNetworks/service/platform/java/com/android/server/ondeviceintelligence/RemoteOnDeviceIntelligenceService.java
@@ -0,0 +1,66 @@
+/*
+ * 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.server.ondeviceintelligence;
+
+import static android.content.Context.BIND_FOREGROUND_SERVICE;
+import static android.content.Context.BIND_INCLUDE_CAPABILITIES;
+
+import android.content.ComponentName;
+import android.content.Context;
+import android.content.Intent;
+import android.provider.Settings;
+import android.service.ondeviceintelligence.IOnDeviceIntelligenceService;
+import android.service.ondeviceintelligence.OnDeviceIntelligenceService;
+
+import com.android.internal.infra.ServiceConnector;
+
+import java.util.concurrent.TimeUnit;
+
+/**
+ * Manages the connection to the remote on-device intelligence service. Also, handles unbinding
+ * logic set by the service implementation via a Secure Settings flag.
+ */
+public class RemoteOnDeviceIntelligenceService extends
+        ServiceConnector.Impl<IOnDeviceIntelligenceService> {
+    private static final long LONG_TIMEOUT = TimeUnit.HOURS.toMillis(4);
+    private static final String TAG =
+            RemoteOnDeviceIntelligenceService.class.getSimpleName();
+
+    RemoteOnDeviceIntelligenceService(Context context, ComponentName serviceName,
+            int userId) {
+        super(context, new Intent(
+                        OnDeviceIntelligenceService.SERVICE_INTERFACE).setComponent(serviceName),
+                BIND_FOREGROUND_SERVICE | BIND_INCLUDE_CAPABILITIES, userId,
+                IOnDeviceIntelligenceService.Stub::asInterface);
+
+        // Bind right away
+        connect();
+    }
+
+    @Override
+    protected long getRequestTimeoutMs() {
+        return LONG_TIMEOUT;
+    }
+
+    @Override
+    protected long getAutoDisconnectTimeoutMs() {
+        return Settings.Secure.getLongForUser(mContext.getContentResolver(),
+                Settings.Secure.ON_DEVICE_INTELLIGENCE_UNBIND_TIMEOUT_MS,
+                TimeUnit.SECONDS.toMillis(30),
+                mContext.getUserId());
+    }
+}
diff --git a/packages/NeuralNetworks/service/platform/java/com/android/server/ondeviceintelligence/RemoteOnDeviceSandboxedInferenceService.java b/packages/NeuralNetworks/service/platform/java/com/android/server/ondeviceintelligence/RemoteOnDeviceSandboxedInferenceService.java
new file mode 100644
index 0000000..18b1383
--- /dev/null
+++ b/packages/NeuralNetworks/service/platform/java/com/android/server/ondeviceintelligence/RemoteOnDeviceSandboxedInferenceService.java
@@ -0,0 +1,76 @@
+/*
+ * 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.server.ondeviceintelligence;
+
+import static android.content.Context.BIND_FOREGROUND_SERVICE;
+import static android.content.Context.BIND_INCLUDE_CAPABILITIES;
+
+import android.content.ComponentName;
+import android.content.Context;
+import android.content.Intent;
+import android.provider.Settings;
+import android.service.ondeviceintelligence.IOnDeviceSandboxedInferenceService;
+import android.service.ondeviceintelligence.OnDeviceSandboxedInferenceService;
+
+import com.android.internal.infra.ServiceConnector;
+
+import java.util.concurrent.TimeUnit;
+
+
+/**
+ * Manages the connection to the remote on-device sand boxed inference service. Also, handles
+ * unbinding
+ * logic set by the service implementation via a SecureSettings flag.
+ */
+public class RemoteOnDeviceSandboxedInferenceService extends
+        ServiceConnector.Impl<IOnDeviceSandboxedInferenceService> {
+    private static final long LONG_TIMEOUT = TimeUnit.HOURS.toMillis(1);
+
+    /**
+     * Creates an instance of {@link ServiceConnector}
+     *
+     * See {@code protected} methods for optional parameters you can override.
+     *
+     * @param context to be used for {@link Context#bindServiceAsUser binding} and
+     *                {@link Context#unbindService unbinding}
+     * @param userId  to be used for {@link Context#bindServiceAsUser binding}
+     */
+    RemoteOnDeviceSandboxedInferenceService(Context context, ComponentName serviceName,
+            int userId) {
+        super(context, new Intent(
+                        OnDeviceSandboxedInferenceService.SERVICE_INTERFACE).setComponent(serviceName),
+                BIND_FOREGROUND_SERVICE | BIND_INCLUDE_CAPABILITIES, userId,
+                IOnDeviceSandboxedInferenceService.Stub::asInterface);
+
+        // Bind right away
+        connect();
+    }
+
+    @Override
+    protected long getRequestTimeoutMs() {
+        return LONG_TIMEOUT;
+    }
+
+
+    @Override
+    protected long getAutoDisconnectTimeoutMs() {
+        return Settings.Secure.getLongForUser(mContext.getContentResolver(),
+                Settings.Secure.ON_DEVICE_INFERENCE_UNBIND_TIMEOUT_MS,
+                TimeUnit.SECONDS.toMillis(30),
+                mContext.getUserId());
+    }
+}
diff --git a/packages/NeuralNetworks/service/platform/java/com/android/server/ondeviceintelligence/callbacks/ListenableDownloadCallback.java b/packages/NeuralNetworks/service/platform/java/com/android/server/ondeviceintelligence/callbacks/ListenableDownloadCallback.java
new file mode 100644
index 0000000..32f0698
--- /dev/null
+++ b/packages/NeuralNetworks/service/platform/java/com/android/server/ondeviceintelligence/callbacks/ListenableDownloadCallback.java
@@ -0,0 +1,97 @@
+/*
+ * 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.server.ondeviceintelligence.callbacks;
+
+import android.app.ondeviceintelligence.IDownloadCallback;
+import android.os.Handler;
+import android.os.PersistableBundle;
+import android.os.RemoteException;
+
+import com.android.internal.infra.AndroidFuture;
+
+import java.util.concurrent.TimeoutException;
+
+/**
+ * This class extends the {@link IDownloadCallback} and adds a timeout Runnable to the callback
+ * such that, in the case where the callback methods are not invoked, we do not have to wait for
+ * timeout based on {@link #onDownloadCompleted} which might take minutes or hours to complete in
+ * some cases. Instead, in such cases we rely on the remote service sending progress updates and if
+ * there are *no* progress callbacks in the duration of {@link #idleTimeoutMs}, we can assume the
+ * download will not complete and enabling faster cleanup.
+ */
+public class ListenableDownloadCallback extends IDownloadCallback.Stub implements Runnable {
+    private final IDownloadCallback callback;
+    private final Handler handler;
+    private final AndroidFuture future;
+    private final long idleTimeoutMs;
+
+    /**
+     * Constructor to create a ListenableDownloadCallback.
+     *
+     * @param callback      callback to send download updates to caller.
+     * @param handler       handler to schedule timeout runnable.
+     * @param future        future to complete to signal the callback has reached a terminal state.
+     * @param idleTimeoutMs timeout within which download updates should be received.
+     */
+    public ListenableDownloadCallback(IDownloadCallback callback, Handler handler,
+            AndroidFuture future,
+            long idleTimeoutMs) {
+        this.callback = callback;
+        this.handler = handler;
+        this.future = future;
+        this.idleTimeoutMs = idleTimeoutMs;
+        handler.postDelayed(this,
+                idleTimeoutMs); // init the timeout runnable in case no callback is ever invoked
+    }
+
+    @Override
+    public void onDownloadStarted(long bytesToDownload) throws RemoteException {
+        callback.onDownloadStarted(bytesToDownload);
+        handler.removeCallbacks(this);
+        handler.postDelayed(this, idleTimeoutMs);
+    }
+
+    @Override
+    public void onDownloadProgress(long bytesDownloaded) throws RemoteException {
+        callback.onDownloadProgress(bytesDownloaded);
+        handler.removeCallbacks(this); // remove previously queued timeout tasks.
+        handler.postDelayed(this, idleTimeoutMs); // queue fresh timeout task for next update.
+    }
+
+    @Override
+    public void onDownloadFailed(int failureStatus,
+            String errorMessage, PersistableBundle errorParams) throws RemoteException {
+        callback.onDownloadFailed(failureStatus, errorMessage, errorParams);
+        handler.removeCallbacks(this);
+        future.completeExceptionally(new TimeoutException());
+    }
+
+    @Override
+    public void onDownloadCompleted(
+            android.os.PersistableBundle downloadParams) throws RemoteException {
+        callback.onDownloadCompleted(downloadParams);
+        handler.removeCallbacks(this);
+        future.complete(null);
+    }
+
+    @Override
+    public void run() {
+        future.completeExceptionally(
+                new TimeoutException()); // complete the future as we haven't received updates
+        // for download progress.
+    }
+}
\ No newline at end of file
diff --git a/services/Android.bp b/services/Android.bp
index a7cb9bb..7298f14 100644
--- a/services/Android.bp
+++ b/services/Android.bp
@@ -233,8 +233,7 @@
             libs: ["service-ondeviceintelligence.stubs.system_server"],
         },
         release_ondevice_intelligence_platform: {
-            srcs: [":service-ondeviceintelligence-sources"],
-            static_libs: ["modules-utils-backgroundthread"],
+            srcs: [":service-ondeviceintelligence-sources-platform"],
         },
     },
 }
diff --git a/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java b/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java
index 4230cd8..cba606c 100644
--- a/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java
+++ b/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java
@@ -324,7 +324,8 @@
             if (target != imeControlTarget) {
                 // TODO(b/353463205): check if fromUser=false is correct here
                 boolean imeVisible = target.isRequestedVisible(WindowInsets.Type.ime());
-                ImeTracker.Token statsToken = ImeTracker.forLogging().onStart(ImeTracker.TYPE_HIDE,
+                ImeTracker.Token statsToken = ImeTracker.forLogging().onStart(
+                        imeVisible ? ImeTracker.TYPE_SHOW : ImeTracker.TYPE_HIDE,
                         ImeTracker.ORIGIN_SERVER,
                         imeVisible ? SoftInputShowHideReason.SHOW_INPUT_TARGET_CHANGED
                                 : SoftInputShowHideReason.HIDE_INPUT_TARGET_CHANGED,