diff --git a/AconfigFlags.bp b/AconfigFlags.bp
index 38ade2f..5e87859 100644
--- a/AconfigFlags.bp
+++ b/AconfigFlags.bp
@@ -808,21 +808,6 @@
     defaults: ["framework-minus-apex-aconfig-java-defaults"],
 }
 
-// OnDeviceIntelligence
-aconfig_declarations {
-    name: "android.app.ondeviceintelligence-aconfig",
-    exportable: true,
-    package: "android.app.ondeviceintelligence.flags",
-    container: "system",
-    srcs: ["core/java/android/app/ondeviceintelligence/flags/ondevice_intelligence.aconfig"],
-}
-
-java_aconfig_library {
-    name: "android.app.ondeviceintelligence-aconfig-java",
-    aconfig_declarations: "android.app.ondeviceintelligence-aconfig",
-    defaults: ["framework-minus-apex-aconfig-java-defaults"],
-}
-
 // Permissions
 aconfig_declarations {
     name: "android.permission.flags-aconfig",
diff --git a/Android.bp b/Android.bp
index a525583b8..529da53 100644
--- a/Android.bp
+++ b/Android.bp
@@ -446,6 +446,9 @@
         default: [
             "framework-platformcrashrecovery.impl",
         ],
+    }) + select(release_flag("RELEASE_ONDEVICE_INTELLIGENCE_MODULE"), {
+        true: [],
+        default: ["framework-ondeviceintelligence-platform.impl"],
     }),
     sdk_version: "core_platform",
     installable: false,
@@ -489,6 +492,7 @@
     apex_available: ["//apex_available:platform"],
     visibility: [
         "//frameworks/base:__subpackages__",
+        "//packages/modules/NeuralNetworks:__subpackages__",
     ],
     compile_dex: false,
     headers_only: true,
@@ -584,6 +588,9 @@
         default: [
             "framework-platformcrashrecovery-compat-config",
         ],
+    }) + select(release_flag("RELEASE_ONDEVICE_INTELLIGENCE_MODULE"), {
+        true: [],
+        default: ["framework-ondeviceintelligence-platform-compat-config"],
     }),
 }
 
diff --git a/api/Android.bp b/api/Android.bp
index 7326203..14c2766 100644
--- a/api/Android.bp
+++ b/api/Android.bp
@@ -105,6 +105,13 @@
         default: [
             "framework-platformcrashrecovery",
         ],
+    }) + select(release_flag("RELEASE_ONDEVICE_INTELLIGENCE_MODULE"), {
+        true: [
+            "framework-ondeviceintelligence",
+        ],
+        default: [
+            "framework-ondeviceintelligence-platform",
+        ],
     }) + select(release_flag("RELEASE_RANGING_STACK"), {
         true: [
             "framework-ranging",
@@ -119,7 +126,12 @@
         "service-permission",
         "service-rkp",
         "service-sdksandbox",
-    ] + select(soong_config_variable("ANDROID", "release_crashrecovery_module"), {
+    ] + select(release_flag("RELEASE_ONDEVICE_INTELLIGENCE_MODULE"), {
+        true: [
+            "service-ondeviceintelligence",
+        ],
+        default: [],
+    }) + select(soong_config_variable("ANDROID", "release_crashrecovery_module"), {
         "true": [
             "service-crashrecovery",
         ],
@@ -478,6 +490,7 @@
         "//frameworks/base/location",
         "//frameworks/base/packages/CrashRecovery/framework",
         "//frameworks/base/nfc",
+        "//packages/modules/NeuralNetworks:__subpackages__",
     ],
     plugins: ["error_prone_android_framework"],
     errorprone: {
diff --git a/api/api.go b/api/api.go
index 5ca24de..e4d783e 100644
--- a/api/api.go
+++ b/api/api.go
@@ -29,6 +29,7 @@
 const virtualization = "framework-virtualization"
 const location = "framework-location"
 const platformCrashrecovery = "framework-platformcrashrecovery"
+const ondeviceintelligence = "framework-ondeviceintelligence-platform"
 
 var core_libraries_modules = []string{art, conscrypt, i18n}
 
@@ -40,7 +41,7 @@
 // APIs.
 // In addition, the modules in this list are allowed to contribute to test APIs
 // stubs.
-var non_updatable_modules = []string{virtualization, location, platformCrashrecovery}
+var non_updatable_modules = []string{virtualization, location, platformCrashrecovery, ondeviceintelligence}
 
 // The intention behind this soong plugin is to generate a number of "merged"
 // API-related modules that would otherwise require a large amount of very
diff --git a/boot/Android.bp b/boot/Android.bp
index 6eead42..eaa984a 100644
--- a/boot/Android.bp
+++ b/boot/Android.bp
@@ -31,6 +31,7 @@
         "car_bootclasspath_fragment",
         "nfc_apex_bootclasspath_fragment",
         "release_crashrecovery_module",
+        "release_ondevice_intelligence_module",
         "release_package_profiling_module",
     ],
     properties: [
@@ -176,6 +177,15 @@
                 },
             ],
         },
+        release_ondevice_intelligence_module: {
+            fragments: [
+                // only used when ondeviceintelligence is moved to neuralnetworks module
+                {
+                    apex: "com.android.neuralnetworks",
+                    module: "com.android.ondeviceintelligence-bootclasspath-fragment",
+                },
+            ],
+        },
         release_package_profiling_module: {
             fragments: [
                 // only used if profiling is enabled.
diff --git a/core/api/system-current.txt b/core/api/system-current.txt
index edb30bd..71a6467 100644
--- a/core/api/system-current.txt
+++ b/core/api/system-current.txt
@@ -2276,149 +2276,6 @@
 
 }
 
-package android.app.ondeviceintelligence {
-
-  @FlaggedApi("android.app.ondeviceintelligence.flags.enable_on_device_intelligence") public interface DownloadCallback {
-    method public void onDownloadCompleted(@NonNull android.os.PersistableBundle);
-    method public void onDownloadFailed(int, @Nullable String, @NonNull android.os.PersistableBundle);
-    method public default void onDownloadProgress(long);
-    method public default void onDownloadStarted(long);
-    field public static final int DOWNLOAD_FAILURE_STATUS_DOWNLOADING = 3; // 0x3
-    field public static final int DOWNLOAD_FAILURE_STATUS_NETWORK_FAILURE = 2; // 0x2
-    field public static final int DOWNLOAD_FAILURE_STATUS_NOT_ENOUGH_DISK_SPACE = 1; // 0x1
-    field public static final int DOWNLOAD_FAILURE_STATUS_UNAVAILABLE = 4; // 0x4
-    field public static final int DOWNLOAD_FAILURE_STATUS_UNKNOWN = 0; // 0x0
-  }
-
-  @FlaggedApi("android.app.ondeviceintelligence.flags.enable_on_device_intelligence") public final class Feature implements android.os.Parcelable {
-    method public int describeContents();
-    method @NonNull public android.os.PersistableBundle getFeatureParams();
-    method public int getId();
-    method @Nullable public String getModelName();
-    method @Nullable public String getName();
-    method public int getType();
-    method public int getVariant();
-    method public void writeToParcel(@NonNull android.os.Parcel, int);
-    field @NonNull public static final android.os.Parcelable.Creator<android.app.ondeviceintelligence.Feature> CREATOR;
-  }
-
-  public static final class Feature.Builder {
-    ctor public Feature.Builder(int);
-    method @NonNull public android.app.ondeviceintelligence.Feature build();
-    method @NonNull public android.app.ondeviceintelligence.Feature.Builder setFeatureParams(@NonNull android.os.PersistableBundle);
-    method @NonNull public android.app.ondeviceintelligence.Feature.Builder setModelName(@NonNull String);
-    method @NonNull public android.app.ondeviceintelligence.Feature.Builder setName(@NonNull String);
-    method @NonNull public android.app.ondeviceintelligence.Feature.Builder setType(int);
-    method @NonNull public android.app.ondeviceintelligence.Feature.Builder setVariant(int);
-  }
-
-  @FlaggedApi("android.app.ondeviceintelligence.flags.enable_on_device_intelligence") public final class FeatureDetails implements android.os.Parcelable {
-    ctor public FeatureDetails(int, @NonNull android.os.PersistableBundle);
-    ctor public FeatureDetails(int);
-    method public int describeContents();
-    method @NonNull public android.os.PersistableBundle getFeatureDetailParams();
-    method public int getFeatureStatus();
-    method public void writeToParcel(@NonNull android.os.Parcel, int);
-    field @NonNull public static final android.os.Parcelable.Creator<android.app.ondeviceintelligence.FeatureDetails> CREATOR;
-    field public static final int FEATURE_STATUS_AVAILABLE = 3; // 0x3
-    field public static final int FEATURE_STATUS_DOWNLOADABLE = 1; // 0x1
-    field public static final int FEATURE_STATUS_DOWNLOADING = 2; // 0x2
-    field public static final int FEATURE_STATUS_SERVICE_UNAVAILABLE = 4; // 0x4
-    field public static final int FEATURE_STATUS_UNAVAILABLE = 0; // 0x0
-  }
-
-  @FlaggedApi("android.app.ondeviceintelligence.flags.enable_on_device_intelligence_module") public final class InferenceInfo implements android.os.Parcelable {
-    method public int describeContents();
-    method public long getEndTimeMillis();
-    method public long getStartTimeMillis();
-    method public long getSuspendedTimeMillis();
-    method public int getUid();
-    method public void writeToParcel(@NonNull android.os.Parcel, int);
-    field @NonNull public static final android.os.Parcelable.Creator<android.app.ondeviceintelligence.InferenceInfo> CREATOR;
-  }
-
-  public static final class InferenceInfo.Builder {
-    ctor public InferenceInfo.Builder(int);
-    method @NonNull public android.app.ondeviceintelligence.InferenceInfo build();
-    method @NonNull public android.app.ondeviceintelligence.InferenceInfo.Builder setEndTimeMillis(long);
-    method @NonNull public android.app.ondeviceintelligence.InferenceInfo.Builder setStartTimeMillis(long);
-    method @NonNull public android.app.ondeviceintelligence.InferenceInfo.Builder setSuspendedTimeMillis(long);
-  }
-
-  @FlaggedApi("android.app.ondeviceintelligence.flags.enable_on_device_intelligence") public class OnDeviceIntelligenceException extends java.lang.Exception {
-    ctor public OnDeviceIntelligenceException(int, @NonNull String, @NonNull android.os.PersistableBundle);
-    ctor public OnDeviceIntelligenceException(int, @NonNull android.os.PersistableBundle);
-    ctor public OnDeviceIntelligenceException(int, @NonNull String);
-    ctor public OnDeviceIntelligenceException(int);
-    method public int getErrorCode();
-    method @NonNull public android.os.PersistableBundle getErrorParams();
-    field public static final int ON_DEVICE_INTELLIGENCE_SERVICE_UNAVAILABLE = 100; // 0x64
-    field public static final int PROCESSING_ERROR_BAD_DATA = 2; // 0x2
-    field public static final int PROCESSING_ERROR_BAD_REQUEST = 3; // 0x3
-    field public static final int PROCESSING_ERROR_BUSY = 9; // 0x9
-    field public static final int PROCESSING_ERROR_CANCELLED = 7; // 0x7
-    field public static final int PROCESSING_ERROR_COMPUTE_ERROR = 5; // 0x5
-    field public static final int PROCESSING_ERROR_INTERNAL = 14; // 0xe
-    field public static final int PROCESSING_ERROR_IPC_ERROR = 6; // 0x6
-    field public static final int PROCESSING_ERROR_NOT_AVAILABLE = 8; // 0x8
-    field public static final int PROCESSING_ERROR_REQUEST_NOT_SAFE = 4; // 0x4
-    field public static final int PROCESSING_ERROR_REQUEST_TOO_LARGE = 12; // 0xc
-    field public static final int PROCESSING_ERROR_RESPONSE_NOT_SAFE = 11; // 0xb
-    field public static final int PROCESSING_ERROR_SAFETY_ERROR = 10; // 0xa
-    field public static final int PROCESSING_ERROR_SERVICE_UNAVAILABLE = 15; // 0xf
-    field public static final int PROCESSING_ERROR_SUSPENDED = 13; // 0xd
-    field public static final int PROCESSING_ERROR_UNKNOWN = 1; // 0x1
-    field public static final int PROCESSING_UPDATE_STATUS_CONNECTION_FAILED = 200; // 0xc8
-  }
-
-  @FlaggedApi("android.app.ondeviceintelligence.flags.enable_on_device_intelligence") public final class OnDeviceIntelligenceManager {
-    method @RequiresPermission(android.Manifest.permission.USE_ON_DEVICE_INTELLIGENCE) public void getFeature(int, @NonNull java.util.concurrent.Executor, @NonNull android.os.OutcomeReceiver<android.app.ondeviceintelligence.Feature,android.app.ondeviceintelligence.OnDeviceIntelligenceException>);
-    method @RequiresPermission(android.Manifest.permission.USE_ON_DEVICE_INTELLIGENCE) public void getFeatureDetails(@NonNull android.app.ondeviceintelligence.Feature, @NonNull java.util.concurrent.Executor, @NonNull android.os.OutcomeReceiver<android.app.ondeviceintelligence.FeatureDetails,android.app.ondeviceintelligence.OnDeviceIntelligenceException>);
-    method @FlaggedApi("android.app.ondeviceintelligence.flags.enable_on_device_intelligence_module") @NonNull @RequiresPermission(android.Manifest.permission.DUMP) public java.util.List<android.app.ondeviceintelligence.InferenceInfo> getLatestInferenceInfo(long);
-    method @Nullable @RequiresPermission(android.Manifest.permission.USE_ON_DEVICE_INTELLIGENCE) public String getRemoteServicePackageName();
-    method @RequiresPermission(android.Manifest.permission.USE_ON_DEVICE_INTELLIGENCE) public void getVersion(@NonNull java.util.concurrent.Executor, @NonNull java.util.function.LongConsumer);
-    method @RequiresPermission(android.Manifest.permission.USE_ON_DEVICE_INTELLIGENCE) public void listFeatures(@NonNull java.util.concurrent.Executor, @NonNull android.os.OutcomeReceiver<java.util.List<android.app.ondeviceintelligence.Feature>,android.app.ondeviceintelligence.OnDeviceIntelligenceException>);
-    method @RequiresPermission(android.Manifest.permission.USE_ON_DEVICE_INTELLIGENCE) public void processRequest(@NonNull android.app.ondeviceintelligence.Feature, @NonNull android.os.Bundle, int, @Nullable android.os.CancellationSignal, @Nullable android.app.ondeviceintelligence.ProcessingSignal, @NonNull java.util.concurrent.Executor, @NonNull android.app.ondeviceintelligence.ProcessingCallback);
-    method @RequiresPermission(android.Manifest.permission.USE_ON_DEVICE_INTELLIGENCE) public void processRequestStreaming(@NonNull android.app.ondeviceintelligence.Feature, @NonNull android.os.Bundle, int, @Nullable android.os.CancellationSignal, @Nullable android.app.ondeviceintelligence.ProcessingSignal, @NonNull java.util.concurrent.Executor, @NonNull android.app.ondeviceintelligence.StreamingProcessingCallback);
-    method @RequiresPermission(android.Manifest.permission.USE_ON_DEVICE_INTELLIGENCE) public void requestFeatureDownload(@NonNull android.app.ondeviceintelligence.Feature, @Nullable android.os.CancellationSignal, @NonNull java.util.concurrent.Executor, @NonNull android.app.ondeviceintelligence.DownloadCallback);
-    method @RequiresPermission(android.Manifest.permission.USE_ON_DEVICE_INTELLIGENCE) public void requestTokenInfo(@NonNull android.app.ondeviceintelligence.Feature, @NonNull android.os.Bundle, @Nullable android.os.CancellationSignal, @NonNull java.util.concurrent.Executor, @NonNull android.os.OutcomeReceiver<android.app.ondeviceintelligence.TokenInfo,android.app.ondeviceintelligence.OnDeviceIntelligenceException>);
-    field public static final int REQUEST_TYPE_EMBEDDINGS = 2; // 0x2
-    field public static final int REQUEST_TYPE_INFERENCE = 0; // 0x0
-    field public static final int REQUEST_TYPE_PREPARE = 1; // 0x1
-  }
-
-  @FlaggedApi("android.app.ondeviceintelligence.flags.enable_on_device_intelligence") public interface ProcessingCallback {
-    method public default void onDataAugmentRequest(@NonNull android.os.Bundle, @NonNull java.util.function.Consumer<android.os.Bundle>);
-    method public void onError(@NonNull android.app.ondeviceintelligence.OnDeviceIntelligenceException);
-    method public void onResult(@NonNull android.os.Bundle);
-  }
-
-  @FlaggedApi("android.app.ondeviceintelligence.flags.enable_on_device_intelligence") public final class ProcessingSignal {
-    ctor public ProcessingSignal();
-    method public void sendSignal(@NonNull android.os.PersistableBundle);
-    method public void setOnProcessingSignalCallback(@NonNull java.util.concurrent.Executor, @Nullable android.app.ondeviceintelligence.ProcessingSignal.OnProcessingSignalCallback);
-  }
-
-  public static interface ProcessingSignal.OnProcessingSignalCallback {
-    method public void onSignalReceived(@NonNull android.os.PersistableBundle);
-  }
-
-  @FlaggedApi("android.app.ondeviceintelligence.flags.enable_on_device_intelligence") public interface StreamingProcessingCallback extends android.app.ondeviceintelligence.ProcessingCallback {
-    method public void onPartialResult(@NonNull android.os.Bundle);
-  }
-
-  @FlaggedApi("android.app.ondeviceintelligence.flags.enable_on_device_intelligence") public final class TokenInfo implements android.os.Parcelable {
-    ctor public TokenInfo(long, @NonNull android.os.PersistableBundle);
-    ctor public TokenInfo(long);
-    method public int describeContents();
-    method public long getCount();
-    method @NonNull public android.os.PersistableBundle getInfoParams();
-    method public void writeToParcel(@NonNull android.os.Parcel, int);
-    field @NonNull public static final android.os.Parcelable.Creator<android.app.ondeviceintelligence.TokenInfo> CREATOR;
-  }
-
-}
-
 package android.app.people {
 
   public final class PeopleManager {
@@ -13791,39 +13648,6 @@
 
 }
 
-package android.service.ondeviceintelligence {
-
-  @FlaggedApi("android.app.ondeviceintelligence.flags.enable_on_device_intelligence") public abstract class OnDeviceIntelligenceService extends android.app.Service {
-    ctor public OnDeviceIntelligenceService();
-    method @Nullable public final android.os.IBinder onBind(@NonNull android.content.Intent);
-    method public abstract void onDownloadFeature(int, @NonNull android.app.ondeviceintelligence.Feature, @Nullable android.os.CancellationSignal, @NonNull android.app.ondeviceintelligence.DownloadCallback);
-    method public abstract void onGetFeature(int, int, @NonNull android.os.OutcomeReceiver<android.app.ondeviceintelligence.Feature,android.app.ondeviceintelligence.OnDeviceIntelligenceException>);
-    method public abstract void onGetFeatureDetails(int, @NonNull android.app.ondeviceintelligence.Feature, @NonNull android.os.OutcomeReceiver<android.app.ondeviceintelligence.FeatureDetails,android.app.ondeviceintelligence.OnDeviceIntelligenceException>);
-    method public abstract void onGetReadOnlyFeatureFileDescriptorMap(@NonNull android.app.ondeviceintelligence.Feature, @NonNull java.util.function.Consumer<java.util.Map<java.lang.String,android.os.ParcelFileDescriptor>>);
-    method public abstract void onGetVersion(@NonNull java.util.function.LongConsumer);
-    method public abstract void onInferenceServiceConnected();
-    method public abstract void onInferenceServiceDisconnected();
-    method public abstract void onListFeatures(int, @NonNull android.os.OutcomeReceiver<java.util.List<android.app.ondeviceintelligence.Feature>,android.app.ondeviceintelligence.OnDeviceIntelligenceException>);
-    method public final void updateProcessingState(@NonNull android.os.Bundle, @NonNull java.util.concurrent.Executor, @NonNull android.os.OutcomeReceiver<android.os.PersistableBundle,android.app.ondeviceintelligence.OnDeviceIntelligenceException>);
-    field public static final String SERVICE_INTERFACE = "android.service.ondeviceintelligence.OnDeviceIntelligenceService";
-  }
-
-  @FlaggedApi("android.app.ondeviceintelligence.flags.enable_on_device_intelligence") public abstract class OnDeviceSandboxedInferenceService extends android.app.Service {
-    ctor public OnDeviceSandboxedInferenceService();
-    method public final void fetchFeatureFileDescriptorMap(@NonNull android.app.ondeviceintelligence.Feature, @NonNull java.util.concurrent.Executor, @NonNull java.util.function.Consumer<java.util.Map<java.lang.String,android.os.ParcelFileDescriptor>>);
-    method @NonNull public java.util.concurrent.Executor getCallbackExecutor();
-    method public final void getReadOnlyFileDescriptor(@NonNull String, @NonNull java.util.concurrent.Executor, @NonNull java.util.function.Consumer<android.os.ParcelFileDescriptor>) throws java.io.FileNotFoundException;
-    method @Nullable public final android.os.IBinder onBind(@NonNull android.content.Intent);
-    method @NonNull public abstract void onProcessRequest(int, @NonNull android.app.ondeviceintelligence.Feature, @NonNull android.os.Bundle, int, @Nullable android.os.CancellationSignal, @Nullable android.app.ondeviceintelligence.ProcessingSignal, @NonNull android.app.ondeviceintelligence.ProcessingCallback);
-    method @NonNull public abstract void onProcessRequestStreaming(int, @NonNull android.app.ondeviceintelligence.Feature, @NonNull android.os.Bundle, int, @Nullable android.os.CancellationSignal, @Nullable android.app.ondeviceintelligence.ProcessingSignal, @NonNull android.app.ondeviceintelligence.StreamingProcessingCallback);
-    method @NonNull public abstract void onTokenInfoRequest(int, @NonNull android.app.ondeviceintelligence.Feature, @NonNull android.os.Bundle, @Nullable android.os.CancellationSignal, @NonNull android.os.OutcomeReceiver<android.app.ondeviceintelligence.TokenInfo,android.app.ondeviceintelligence.OnDeviceIntelligenceException>);
-    method public abstract void onUpdateProcessingState(@NonNull android.os.Bundle, @NonNull android.os.OutcomeReceiver<android.os.PersistableBundle,android.app.ondeviceintelligence.OnDeviceIntelligenceException>);
-    method public final java.io.FileInputStream openFileInput(@NonNull String) throws java.io.FileNotFoundException;
-    field public static final String SERVICE_INTERFACE = "android.service.ondeviceintelligence.OnDeviceSandboxedInferenceService";
-  }
-
-}
-
 package android.service.persistentdata {
 
   @FlaggedApi("android.security.frp_enforcement") public class PersistentDataBlockManager {
diff --git a/core/api/test-current.txt b/core/api/test-current.txt
index 0a806c7..6230a59 100644
--- a/core/api/test-current.txt
+++ b/core/api/test-current.txt
@@ -3247,14 +3247,6 @@
 
 }
 
-package android.service.ondeviceintelligence {
-
-  @FlaggedApi("android.app.ondeviceintelligence.flags.enable_on_device_intelligence") public abstract class OnDeviceIntelligenceService extends android.app.Service {
-    method public void onReady();
-  }
-
-}
-
 package android.service.quickaccesswallet {
 
   public interface QuickAccessWalletClient extends java.io.Closeable {
diff --git a/core/java/android/app/SystemServiceRegistry.java b/core/java/android/app/SystemServiceRegistry.java
index f357836e..248e0433 100644
--- a/core/java/android/app/SystemServiceRegistry.java
+++ b/core/java/android/app/SystemServiceRegistry.java
@@ -42,8 +42,7 @@
 import android.app.contextualsearch.ContextualSearchManager;
 import android.app.ecm.EnhancedConfirmationFrameworkInitializer;
 import android.app.job.JobSchedulerFrameworkInitializer;
-import android.app.ondeviceintelligence.IOnDeviceIntelligenceManager;
-import android.app.ondeviceintelligence.OnDeviceIntelligenceManager;
+import android.app.ondeviceintelligence.OnDeviceIntelligenceFrameworkInitializer;
 import android.app.people.PeopleManager;
 import android.app.prediction.AppPredictionManager;
 import android.app.role.RoleFrameworkInitializer;
@@ -1692,19 +1691,6 @@
                         throw new ServiceNotFoundException(Context.WEARABLE_SENSING_SERVICE);
                     }});
 
-        registerService(Context.ON_DEVICE_INTELLIGENCE_SERVICE, OnDeviceIntelligenceManager.class,
-                new CachedServiceFetcher<OnDeviceIntelligenceManager>() {
-                    @Override
-                    public OnDeviceIntelligenceManager createService(ContextImpl ctx)
-                            throws ServiceNotFoundException {
-                        IBinder iBinder = ServiceManager.getServiceOrThrow(
-                                Context.ON_DEVICE_INTELLIGENCE_SERVICE);
-                        IOnDeviceIntelligenceManager manager =
-                                IOnDeviceIntelligenceManager.Stub.asInterface(iBinder);
-                        return new OnDeviceIntelligenceManager(ctx.getOuterContext(), manager);
-                    }
-                });
-
         registerService(Context.GRAMMATICAL_INFLECTION_SERVICE, GrammaticalInflectionManager.class,
                 new CachedServiceFetcher<GrammaticalInflectionManager>() {
                     @Override
@@ -1849,6 +1835,7 @@
             ConnectivityFrameworkInitializerTiramisu.registerServiceWrappers();
             NearbyFrameworkInitializer.registerServiceWrappers();
             OnDevicePersonalizationFrameworkInitializer.registerServiceWrappers();
+            OnDeviceIntelligenceFrameworkInitializer.registerServiceWrappers();
             DeviceLockFrameworkInitializer.registerServiceWrappers();
             VirtualizationFrameworkInitializer.registerServiceWrappers();
             ConnectivityFrameworkInitializerBaklava.registerServiceWrappers();
diff --git a/packages/NeuralNetworks/framework/Android.bp b/packages/NeuralNetworks/framework/Android.bp
new file mode 100644
index 0000000..6f45daa
--- /dev/null
+++ b/packages/NeuralNetworks/framework/Android.bp
@@ -0,0 +1,30 @@
+// 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 {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+filegroup {
+    name: "framework-ondeviceintelligence-sources",
+    srcs: [
+        "java/**/*.aidl",
+        "java/**/*.java",
+    ],
+    path: "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/java/android/app/ondeviceintelligence/DownloadCallback.java
index 30c6e19..95fb288 100644
--- a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/DownloadCallback.java
+++ b/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/DownloadCallback.java
@@ -23,8 +23,7 @@
 import android.annotation.Nullable;
 import android.annotation.SystemApi;
 import android.os.PersistableBundle;
-
-import androidx.annotation.IntDef;
+import android.annotation.IntDef;
 
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
@@ -53,14 +52,14 @@
 
     /**
      * Sent when feature download has been initiated already, hence no need to request download
-     * again. Caller can query {@link OnDeviceIntelligenceManager#getFeatureStatus} to check if
+     * again. Caller can query {@link OnDeviceIntelligenceManager#getFeatureDetails} to check if
      * download has been completed.
      */
     int DOWNLOAD_FAILURE_STATUS_DOWNLOADING = 3;
 
     /**
      * Sent when feature download did not start due to errors (e.g. remote exception of features not
-     * available). Caller can query {@link OnDeviceIntelligenceManager#getFeatureStatus} to check
+     * available). Caller can query {@link OnDeviceIntelligenceManager#getFeatureDetails} to check
      * if feature-status is {@link FeatureDetails#FEATURE_STATUS_DOWNLOADABLE}.
      */
     int DOWNLOAD_FAILURE_STATUS_UNAVAILABLE = 4;
@@ -72,7 +71,7 @@
             DOWNLOAD_FAILURE_STATUS_NETWORK_FAILURE,
             DOWNLOAD_FAILURE_STATUS_DOWNLOADING,
             DOWNLOAD_FAILURE_STATUS_UNAVAILABLE
-    }, open = true)
+    })
     @Retention(RetentionPolicy.SOURCE)
     @interface DownloadFailureStatus {
     }
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/Feature.aidl b/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/Feature.aidl
index 18494d7..47cfb4a 100644
--- a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/Feature.aidl
+++ b/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/Feature.aidl
@@ -19,4 +19,4 @@
 /**
   * @hide
   */
-parcelable Feature;
+@JavaOnlyStableParcelable parcelable Feature;
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/Feature.java b/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/Feature.java
index bcc56073..88f4de29 100644
--- a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/Feature.java
+++ b/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/Feature.java
@@ -26,6 +26,8 @@
 import android.os.Parcelable;
 import android.os.PersistableBundle;
 
+import java.util.Objects;
+
 /**
  * Represents a typical feature associated with on-device intelligence.
  *
@@ -56,9 +58,8 @@
         this.mModelName = modelName;
         this.mType = type;
         this.mVariant = variant;
-        this.mFeatureParams = featureParams;
-        com.android.internal.util.AnnotationValidations.validate(
-                NonNull.class, null, mFeatureParams);
+        this.mFeatureParams = Objects.requireNonNull(featureParams,
+                "featureParams should be non-null.");
     }
 
     /** Returns the unique and immutable identifier of this feature. */
@@ -167,8 +168,6 @@
         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
@@ -200,6 +199,7 @@
 
         /**
          * Provides a builder instance to create a feature for given id.
+         *
          * @param id the unique identifier for the feature.
          */
         public Builder(int id) {
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/FeatureDetails.aidl b/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/FeatureDetails.aidl
index 0589bf8..c5b3532 100644
--- a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/FeatureDetails.aidl
+++ b/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/FeatureDetails.aidl
@@ -19,4 +19,4 @@
 /**
   * @hide
   */
-parcelable FeatureDetails;
+@JavaOnlyStableParcelable parcelable FeatureDetails;
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/FeatureDetails.java b/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/FeatureDetails.java
index 44930f2..063cfb8 100644
--- a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/FeatureDetails.java
+++ b/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/FeatureDetails.java
@@ -19,18 +19,18 @@
 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 androidx.annotation.IntDef;
-
 import java.lang.annotation.ElementType;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
 import java.text.MessageFormat;
+import java.util.Objects;
 
 /**
  * Represents a status of a requested {@link Feature}.
@@ -69,7 +69,7 @@
             FEATURE_STATUS_DOWNLOADING,
             FEATURE_STATUS_AVAILABLE,
             FEATURE_STATUS_SERVICE_UNAVAILABLE
-    }, open = true)
+    })
     @Target({ElementType.TYPE_USE, ElementType.METHOD, ElementType.PARAMETER, ElementType.FIELD})
     @Retention(RetentionPolicy.SOURCE)
     public @interface Status {
@@ -79,18 +79,12 @@
             @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);
+        this.mFeatureDetailParams = Objects.requireNonNull(featureDetailParams);
     }
 
     public FeatureDetails(
             @Status int featureStatus) {
         this.mFeatureStatus = featureStatus;
-        com.android.internal.util.AnnotationValidations.validate(
-                Status.class, null, mFeatureStatus);
         this.mFeatureDetailParams = new PersistableBundle();
     }
 
@@ -155,11 +149,7 @@
                 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);
     }
 
 
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/ICancellationSignal.aidl b/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/ICancellationSignal.aidl
new file mode 100644
index 0000000..1fe201f
--- /dev/null
+++ b/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/ICancellationSignal.aidl
@@ -0,0 +1,24 @@
+/*
+ * 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
+ */
+oneway interface ICancellationSignal {
+    void cancel();
+}
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IOnDeviceIntelligenceManager.aidl b/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IOnDeviceIntelligenceManager.aidl
index 1977a39..fac5ec6 100644
--- a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IOnDeviceIntelligenceManager.aidl
+++ b/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IOnDeviceIntelligenceManager.aidl
@@ -16,8 +16,8 @@
 
  package android.app.ondeviceintelligence;
 
- import com.android.internal.infra.AndroidFuture;
- import android.os.ICancellationSignal;
+ import com.android.modules.utils.AndroidFuture;
+ import android.app.ondeviceintelligence.ICancellationSignal;
  import android.os.ParcelFileDescriptor;
  import android.os.PersistableBundle;
  import android.os.RemoteCallback;
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IRemoteCallback.aidl b/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IRemoteCallback.aidl
new file mode 100644
index 0000000..6f07693
--- /dev/null
+++ b/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/IRemoteCallback.aidl
@@ -0,0 +1,24 @@
+/*
+* Copyright 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 android.os.Bundle;
+
+/* @hide */
+oneway interface IRemoteCallback {
+    void sendResult(in Bundle data);
+}
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/InferenceInfo.aidl b/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/InferenceInfo.aidl
index 6d70fc4..6f63254 100644
--- a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/InferenceInfo.aidl
+++ b/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/InferenceInfo.aidl
@@ -19,4 +19,4 @@
 /**
   * @hide
   */
-parcelable InferenceInfo;
+@JavaOnlyStableParcelable parcelable InferenceInfo;
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/OnDeviceIntelligenceException.java b/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/OnDeviceIntelligenceException.java
index 03ff563a..2881c9d 100644
--- a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/OnDeviceIntelligenceException.java
+++ b/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/OnDeviceIntelligenceException.java
@@ -20,13 +20,14 @@
 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.PersistableBundle;
 
-import androidx.annotation.IntDef;
-
 import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
 
 /**
@@ -124,8 +125,9 @@
                     PROCESSING_ERROR_SERVICE_UNAVAILABLE,
                     ON_DEVICE_INTELLIGENCE_SERVICE_UNAVAILABLE,
                     PROCESSING_UPDATE_STATUS_CONNECTION_FAILED
-            }, open = true)
+            })
     @Target({ElementType.TYPE_PARAMETER, ElementType.TYPE_USE})
+    @Retention(RetentionPolicy.SOURCE)
     @interface OnDeviceIntelligenceError {
     }
 
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/OnDeviceIntelligenceFrameworkInitializer.java b/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/OnDeviceIntelligenceFrameworkInitializer.java
new file mode 100644
index 0000000..7d35dd7
--- /dev/null
+++ b/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/OnDeviceIntelligenceFrameworkInitializer.java
@@ -0,0 +1,54 @@
+/*
+ * 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.SystemApi;
+import android.app.SystemServiceRegistry;
+import android.content.Context;
+
+/**
+ * Class for performing registration for OnDeviceIntelligence service.
+ *
+ * @hide
+ */
+@SystemApi(client = SystemApi.Client.MODULE_LIBRARIES)
+@FlaggedApi(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
+public class OnDeviceIntelligenceFrameworkInitializer {
+    private OnDeviceIntelligenceFrameworkInitializer() {
+    }
+
+    /**
+     * Called by {@link SystemServiceRegistry}'s static initializer and registers
+     * OnDeviceIntelligence service to {@link Context}, so that {@link Context#getSystemService} can
+     * return them.
+     *
+     * @throws IllegalStateException if this is called from anywhere besides {@link
+     *                               SystemServiceRegistry}
+     */
+    public static void registerServiceWrappers() {
+        SystemServiceRegistry.registerContextAwareService(Context.ON_DEVICE_INTELLIGENCE_SERVICE,
+                OnDeviceIntelligenceManager.class,
+                (context, serviceBinder) -> {
+                    IOnDeviceIntelligenceManager manager =
+                            IOnDeviceIntelligenceManager.Stub.asInterface(serviceBinder);
+                    return new OnDeviceIntelligenceManager(context, manager);
+                });
+    }
+}
\ No newline at end of file
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/OnDeviceIntelligenceManager.java b/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/OnDeviceIntelligenceManager.java
index 91651e3..dc0665a 100644
--- a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/OnDeviceIntelligenceManager.java
+++ b/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/OnDeviceIntelligenceManager.java
@@ -23,18 +23,18 @@
 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.app.ondeviceintelligence.utils.BinderUtils;
 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;
@@ -42,9 +42,7 @@
 import android.system.OsConstants;
 import android.util.Log;
 
-import androidx.annotation.IntDef;
-
-import com.android.internal.infra.AndroidFuture;
+import com.android.modules.utils.AndroidFuture;
 
 import java.lang.annotation.ElementType;
 import java.lang.annotation.Retention;
@@ -80,10 +78,39 @@
     public static final String AUGMENT_REQUEST_CONTENT_BUNDLE_KEY =
             "AugmentRequestContentBundleKey";
 
+    /**
+     * Timeout to be used for unbinding to the configured remote {@link
+     * android.service.ondeviceintelligence.OnDeviceIntelligenceService} if there are no requests in
+     * the queue. A value of -1 represents to never unbind.
+     *
+     * @hide
+     */
+    public static final String ON_DEVICE_INTELLIGENCE_UNBIND_TIMEOUT_MS =
+        "on_device_intelligence_unbind_timeout_ms";
+
+    /**
+     * Timeout that represents maximum idle time before which a callback should be populated.
+     *
+     * @hide
+     */
+    public static final String ON_DEVICE_INTELLIGENCE_IDLE_TIMEOUT_MS =
+        "on_device_intelligence_idle_timeout_ms";
+
+    /**
+     * Timeout to be used for unbinding to the configured remote {@link
+     * android.service.ondeviceintelligence.OnDeviceSandboxedInferenceService} if there are no
+     * requests in the queue. A value of -1 represents to never unbind.
+     *
+     * @hide
+     */
+    public static final String ON_DEVICE_INFERENCE_UNBIND_TIMEOUT_MS =
+        "on_device_inference_unbind_timeout_ms";
+
     private static final String TAG = "OnDeviceIntelligence";
     private final Context mContext;
     private final IOnDeviceIntelligenceManager mService;
 
+
     /**
      * @hide
      */
@@ -105,11 +132,11 @@
         try {
             RemoteCallback callback = new RemoteCallback(result -> {
                 if (result == null) {
-                    Binder.withCleanCallingIdentity(
+                    BinderUtils.withCleanCallingIdentity(
                             () -> callbackExecutor.execute(() -> versionConsumer.accept(0)));
                 }
                 long version = result.getLong(API_VERSION_BUNDLE_KEY);
-                Binder.withCleanCallingIdentity(
+                BinderUtils.withCleanCallingIdentity(
                         () -> callbackExecutor.execute(() -> versionConsumer.accept(version)));
             });
             mService.getVersion(callback);
@@ -151,14 +178,14 @@
                     new IFeatureCallback.Stub() {
                         @Override
                         public void onSuccess(Feature result) {
-                            Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                            BinderUtils.withCleanCallingIdentity(() -> callbackExecutor.execute(
                                     () -> featureReceiver.onResult(result)));
                         }
 
                         @Override
                         public void onFailure(int errorCode, String errorMessage,
                                 PersistableBundle errorParams) {
-                            Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                            BinderUtils.withCleanCallingIdentity(() -> callbackExecutor.execute(
                                     () -> featureReceiver.onError(
                                             new OnDeviceIntelligenceException(
                                                     errorCode, errorMessage, errorParams))));
@@ -185,14 +212,14 @@
                     new IListFeaturesCallback.Stub() {
                         @Override
                         public void onSuccess(List<Feature> result) {
-                            Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                            BinderUtils.withCleanCallingIdentity(() -> callbackExecutor.execute(
                                     () -> featureListReceiver.onResult(result)));
                         }
 
                         @Override
                         public void onFailure(int errorCode, String errorMessage,
                                 PersistableBundle errorParams) {
-                            Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                            BinderUtils.withCleanCallingIdentity(() -> callbackExecutor.execute(
                                     () -> featureListReceiver.onError(
                                             new OnDeviceIntelligenceException(
                                                     errorCode, errorMessage, errorParams))));
@@ -223,14 +250,14 @@
 
                 @Override
                 public void onSuccess(FeatureDetails result) {
-                    Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                    BinderUtils.withCleanCallingIdentity(() -> callbackExecutor.execute(
                             () -> featureDetailsReceiver.onResult(result)));
                 }
 
                 @Override
                 public void onFailure(int errorCode, String errorMessage,
                         PersistableBundle errorParams) {
-                    Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                    BinderUtils.withCleanCallingIdentity(() -> callbackExecutor.execute(
                             () -> featureDetailsReceiver.onError(
                                     new OnDeviceIntelligenceException(errorCode,
                                             errorMessage, errorParams))));
@@ -268,27 +295,27 @@
 
                 @Override
                 public void onDownloadStarted(long bytesToDownload) {
-                    Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                    BinderUtils.withCleanCallingIdentity(() -> callbackExecutor.execute(
                             () -> callback.onDownloadStarted(bytesToDownload)));
                 }
 
                 @Override
                 public void onDownloadProgress(long bytesDownloaded) {
-                    Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                    BinderUtils.withCleanCallingIdentity(() -> callbackExecutor.execute(
                             () -> callback.onDownloadProgress(bytesDownloaded)));
                 }
 
                 @Override
                 public void onDownloadFailed(int failureStatus, String errorMessage,
                         PersistableBundle errorParams) {
-                    Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                    BinderUtils.withCleanCallingIdentity(() -> callbackExecutor.execute(
                             () -> callback.onDownloadFailed(failureStatus, errorMessage,
                                     errorParams)));
                 }
 
                 @Override
                 public void onDownloadCompleted(PersistableBundle downloadParams) {
-                    Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                    BinderUtils.withCleanCallingIdentity(() -> callbackExecutor.execute(
                             () -> callback.onDownloadCompleted(downloadParams)));
                 }
             };
@@ -325,14 +352,14 @@
             ITokenInfoCallback callback = new ITokenInfoCallback.Stub() {
                 @Override
                 public void onSuccess(TokenInfo tokenInfo) {
-                    Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                    BinderUtils.withCleanCallingIdentity(() -> callbackExecutor.execute(
                             () -> outcomeReceiver.onResult(tokenInfo)));
                 }
 
                 @Override
                 public void onFailure(int errorCode, String errorMessage,
                         PersistableBundle errorParams) {
-                    Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                    BinderUtils.withCleanCallingIdentity(() -> callbackExecutor.execute(
                             () -> outcomeReceiver.onError(
                                     new OnDeviceIntelligenceException(
                                             errorCode, errorMessage, errorParams))));
@@ -377,7 +404,7 @@
             IResponseCallback callback = new IResponseCallback.Stub() {
                 @Override
                 public void onSuccess(@InferenceParams Bundle result) {
-                    Binder.withCleanCallingIdentity(() -> {
+                    BinderUtils.withCleanCallingIdentity(() -> {
                         callbackExecutor.execute(() -> processingCallback.onResult(result));
                     });
                 }
@@ -385,7 +412,7 @@
                 @Override
                 public void onFailure(int errorCode, String errorMessage,
                         PersistableBundle errorParams) {
-                    Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                    BinderUtils.withCleanCallingIdentity(() -> callbackExecutor.execute(
                             () -> processingCallback.onError(
                                     new OnDeviceIntelligenceException(
                                             errorCode, errorMessage, errorParams))));
@@ -394,7 +421,7 @@
                 @Override
                 public void onDataAugmentRequest(@NonNull @InferenceParams Bundle request,
                         @NonNull RemoteCallback contentCallback) {
-                    Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                    BinderUtils.withCleanCallingIdentity(() -> callbackExecutor.execute(
                             () -> processingCallback.onDataAugmentRequest(request, result -> {
                                 Bundle bundle = new Bundle();
                                 bundle.putParcelable(AUGMENT_REQUEST_CONTENT_BUNDLE_KEY, result);
@@ -447,7 +474,7 @@
             IStreamingResponseCallback callback = new IStreamingResponseCallback.Stub() {
                 @Override
                 public void onNewContent(@InferenceParams Bundle result) {
-                    Binder.withCleanCallingIdentity(() -> {
+                    BinderUtils.withCleanCallingIdentity(() -> {
                         callbackExecutor.execute(
                                 () -> streamingProcessingCallback.onPartialResult(result));
                     });
@@ -455,7 +482,7 @@
 
                 @Override
                 public void onSuccess(@InferenceParams Bundle result) {
-                    Binder.withCleanCallingIdentity(() -> {
+                    BinderUtils.withCleanCallingIdentity(() -> {
                         callbackExecutor.execute(
                                 () -> streamingProcessingCallback.onResult(result));
                     });
@@ -464,7 +491,7 @@
                 @Override
                 public void onFailure(int errorCode, String errorMessage,
                         PersistableBundle errorParams) {
-                    Binder.withCleanCallingIdentity(() -> {
+                    BinderUtils.withCleanCallingIdentity(() -> {
                         callbackExecutor.execute(
                                 () -> streamingProcessingCallback.onError(
                                         new OnDeviceIntelligenceException(
@@ -476,7 +503,7 @@
                 @Override
                 public void onDataAugmentRequest(@NonNull @InferenceParams Bundle content,
                         @NonNull RemoteCallback contentCallback) {
-                    Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                    BinderUtils.withCleanCallingIdentity(() -> callbackExecutor.execute(
                             () -> streamingProcessingCallback.onDataAugmentRequest(content,
                                     contentResponse -> {
                                         Bundle bundle = new Bundle();
@@ -537,7 +564,7 @@
             REQUEST_TYPE_INFERENCE,
             REQUEST_TYPE_PREPARE,
             REQUEST_TYPE_EMBEDDINGS
-    }, open = true)
+    })
     @Target({ElementType.TYPE_USE, ElementType.METHOD, ElementType.PARAMETER,
             ElementType.FIELD})
     @Retention(RetentionPolicy.SOURCE)
@@ -614,8 +641,17 @@
                     if (error != null || cancellationTransport == null) {
                         Log.e(TAG, "Unable to receive the remote cancellation signal.", error);
                     } else {
-                        cancellationSignal.setRemote(
-                                ICancellationSignal.Stub.asInterface(cancellationTransport));
+                        ICancellationSignal remoteCancellationSignal =
+                                ICancellationSignal.Stub.asInterface(cancellationTransport);
+                        cancellationSignal.setOnCancelListener(
+                                () -> {
+                                    try {
+                                        remoteCancellationSignal.cancel();
+                                    } catch (RemoteException e) {
+                                        Log.w(TAG, "Unable to propagate cancellation signal.",
+                                                e);
+                                    }
+                                });
                     }
                 }, callbackExecutor);
         return cancellationFuture;
@@ -638,6 +674,4 @@
                 }, executor);
         return processingSignalFuture;
     }
-
-
-}
+}
\ No newline at end of file
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/TokenInfo.aidl b/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/TokenInfo.aidl
index 2c19c1e..599b337 100644
--- a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/TokenInfo.aidl
+++ b/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/TokenInfo.aidl
@@ -19,4 +19,4 @@
 /**
   * @hide
   */
-parcelable TokenInfo;
+@JavaOnlyStableParcelable parcelable TokenInfo;
diff --git a/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/utils/BinderUtils.java b/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/utils/BinderUtils.java
new file mode 100644
index 0000000..2916f03
--- /dev/null
+++ b/packages/NeuralNetworks/framework/java/android/app/ondeviceintelligence/utils/BinderUtils.java
@@ -0,0 +1,96 @@
+/*
+ * 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.utils;
+
+import android.annotation.NonNull;
+import android.os.Binder;
+
+import java.util.function.Supplier;
+
+/**
+ * Collection of utilities for {@link Binder} and related classes.
+ * @hide
+ */
+public class BinderUtils {
+    /**
+     * Convenience method for running the provided action enclosed in
+     * {@link Binder#clearCallingIdentity}/{@link Binder#restoreCallingIdentity}
+     *
+     * Any exception thrown by the given action will be caught and rethrown after the call to
+     * {@link Binder#restoreCallingIdentity}
+     *
+     * Note that this is copied from Binder#withCleanCallingIdentity with minor changes
+     * since it is not public.
+     *
+     * @hide
+     */
+    public static final <T extends Exception> void withCleanCallingIdentity(
+            @NonNull ThrowingRunnable<T> action) throws T {
+        final long callingIdentity = Binder.clearCallingIdentity();
+        try {
+            action.run();
+        } finally {
+            Binder.restoreCallingIdentity(callingIdentity);
+        }
+    }
+
+    /**
+     * Like a Runnable, but declared to throw an exception.
+     *
+     * @param <T> The exception class which is declared to be thrown.
+     */
+    @FunctionalInterface
+    public interface ThrowingRunnable<T extends Exception> {
+        /** @see java.lang.Runnable */
+        void run() throws T;
+    }
+
+    /**
+     * Convenience method for running the provided action enclosed in
+     * {@link Binder#clearCallingIdentity}/{@link Binder#restoreCallingIdentity} returning the
+     * result.
+     *
+     * <p>Any exception thrown by the given action will be caught and rethrown after
+     * the call to {@link Binder#restoreCallingIdentity}.
+     *
+     * Note that this is copied from Binder#withCleanCallingIdentity with minor changes
+     * since it is not public.
+     *
+     * @hide
+     */
+    public static final <T, E extends Exception> T withCleanCallingIdentity(
+            @NonNull ThrowingSupplier<T, E> action) throws E {
+        final long callingIdentity = Binder.clearCallingIdentity();
+        try {
+            return action.get();
+        } finally {
+            Binder.restoreCallingIdentity(callingIdentity);
+        }
+    }
+
+    /**
+     * An equivalent of {@link Supplier}
+     *
+     * @param <T> The class which is declared to be returned.
+     * @param <E> The exception class which is declared to be thrown.
+     */
+    @FunctionalInterface
+    public interface ThrowingSupplier<T, E extends Exception> {
+        /** @see java.util.function.Supplier */
+        T get() throws E;
+    }
+}
\ No newline at end of file
diff --git a/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/IOnDeviceIntelligenceService.aidl b/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/IOnDeviceIntelligenceService.aidl
index 45c4350..cba18c1 100644
--- a/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/IOnDeviceIntelligenceService.aidl
+++ b/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/IOnDeviceIntelligenceService.aidl
@@ -18,14 +18,14 @@
 
 import android.os.PersistableBundle;
 import android.os.ParcelFileDescriptor;
-import android.os.ICancellationSignal;
+import android.app.ondeviceintelligence.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 com.android.modules.utils.AndroidFuture;
 import android.service.ondeviceintelligence.IRemoteProcessingService;
 
 
diff --git a/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/IOnDeviceSandboxedInferenceService.aidl b/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/IOnDeviceSandboxedInferenceService.aidl
index 1af3b0f..504fdd9 100644
--- a/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/IOnDeviceSandboxedInferenceService.aidl
+++ b/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/IOnDeviceSandboxedInferenceService.aidl
@@ -21,11 +21,11 @@
 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.app.ondeviceintelligence.IRemoteCallback;
+import android.app.ondeviceintelligence.ICancellationSignal;
 import android.os.PersistableBundle;
 import android.os.Bundle;
-import com.android.internal.infra.AndroidFuture;
+import com.android.modules.utils.AndroidFuture;
 import android.service.ondeviceintelligence.IRemoteStorageService;
 import android.service.ondeviceintelligence.IProcessingUpdateStatusCallback;
 
diff --git a/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/IRemoteStorageService.aidl b/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/IRemoteStorageService.aidl
index a6f49e1..253df89 100644
--- a/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/IRemoteStorageService.aidl
+++ b/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/IRemoteStorageService.aidl
@@ -20,7 +20,7 @@
 import android.os.ParcelFileDescriptor;
 import android.os.RemoteCallback;
 
-import com.android.internal.infra.AndroidFuture;
+import com.android.modules.utils.AndroidFuture;
 
 /**
  * Interface for a concrete implementation to provide access to storage read access
diff --git a/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/OnDeviceIntelligenceService.java b/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/OnDeviceIntelligenceService.java
index d82fe1c..6907e2b 100644
--- a/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/OnDeviceIntelligenceService.java
+++ b/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/OnDeviceIntelligenceService.java
@@ -18,8 +18,6 @@
 
 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;
@@ -27,11 +25,11 @@
 import android.annotation.Nullable;
 import android.annotation.SdkConstant;
 import android.annotation.SystemApi;
-import android.annotation.TestApi;
 import android.app.Service;
 import android.app.ondeviceintelligence.DownloadCallback;
 import android.app.ondeviceintelligence.Feature;
 import android.app.ondeviceintelligence.FeatureDetails;
+import android.app.ondeviceintelligence.ICancellationSignal;
 import android.app.ondeviceintelligence.IDownloadCallback;
 import android.app.ondeviceintelligence.IFeatureCallback;
 import android.app.ondeviceintelligence.IFeatureDetailsCallback;
@@ -39,14 +37,14 @@
 import android.app.ondeviceintelligence.OnDeviceIntelligenceException;
 import android.app.ondeviceintelligence.OnDeviceIntelligenceManager;
 import android.app.ondeviceintelligence.OnDeviceIntelligenceManager.StateParams;
+import android.app.ondeviceintelligence.utils.BinderUtils;
 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.Message;
 import android.os.OutcomeReceiver;
 import android.os.ParcelFileDescriptor;
 import android.os.PersistableBundle;
@@ -55,10 +53,11 @@
 import android.util.Log;
 import android.util.Slog;
 
-import com.android.internal.infra.AndroidFuture;
+import com.android.modules.utils.AndroidFuture;
 
 import java.io.File;
 import java.io.FileNotFoundException;
+import java.io.IOException;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
@@ -92,6 +91,18 @@
 @SystemApi
 @FlaggedApi(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
 public abstract class OnDeviceIntelligenceService extends Service {
+    private static final int MSG_ON_READY = 1;
+    private static final int MSG_GET_VERSION = 2;
+    private static final int MSG_LIST_FEATURES = 3;
+    private static final int MSG_GET_FEATURE = 4;
+    private static final int MSG_GET_FEATURE_DETAILS = 5;
+    private static final int MSG_DOWNLOAD_FEATURE = 6;
+    private static final int MSG_GET_READ_ONLY_FILE_DESCRIPTOR = 7;
+    private static final int MSG_GET_READ_ONLY_FEATURE_FILE_DESCRIPTOR_MAP = 8;
+    private static final int MSG_REGISTER_REMOTE_SERVICES = 9;
+    private static final int MSG_INFERENCE_SERVICE_CONNECTED = 10;
+    private static final int MSG_INFERENCE_SERVICE_DISCONNECTED = 11;
+
     private static final String TAG = OnDeviceIntelligenceService.class.getSimpleName();
 
     private volatile IRemoteProcessingService mRemoteProcessingService;
@@ -101,19 +112,71 @@
     @Override
     public void onCreate() {
         super.onCreate();
-        mHandler = new Handler(Looper.getMainLooper(), null /* callback */, true /* async */);
+        mHandler = new Handler(Looper.getMainLooper()) {
+            @Override
+            public void handleMessage(@NonNull Message msg) {
+                switch (msg.what) {
+                    case MSG_ON_READY:
+                        OnDeviceIntelligenceService.this.onReady();
+                        break;
+                    case MSG_GET_VERSION:
+                        OnDeviceIntelligenceService.this.onGetVersion(
+                                (LongConsumer) msg.obj);
+                        break;
+                    case MSG_LIST_FEATURES:
+                        OnDeviceIntelligenceService.this.onListFeatures(
+                                msg.arg1,
+                                (OutcomeReceiver<List<Feature>, OnDeviceIntelligenceException>) msg.obj);
+                        break;
+                    case MSG_GET_FEATURE:
+                        GetFeatureParams params = (GetFeatureParams) msg.obj;
+                        OnDeviceIntelligenceService.this.onGetFeature(
+                                msg.arg1,
+                                msg.arg2,
+                                params.callback);
+                        break;
+                    case MSG_GET_FEATURE_DETAILS:
+                        FeatureDetailsParams detailsParams = (FeatureDetailsParams) msg.obj;
+                        OnDeviceIntelligenceService.this.onGetFeatureDetails(
+                                msg.arg1,
+                                detailsParams.feature,
+                                detailsParams.callback);
+                        break;
+                    case MSG_DOWNLOAD_FEATURE:
+                        DownloadParams downloadParams = (DownloadParams) msg.obj;
+                        OnDeviceIntelligenceService.this.onDownloadFeature(
+                                msg.arg1,
+                                downloadParams.feature,
+                                downloadParams.cancellationSignal,
+                                downloadParams.callback);
+                        break;
+                    case MSG_GET_READ_ONLY_FILE_DESCRIPTOR:
+                        FileDescriptorParams fdParams = (FileDescriptorParams) msg.obj;
+                        OnDeviceIntelligenceService.this.onGetReadOnlyFileDescriptor(
+                                fdParams.fileName,
+                                fdParams.future);
+                        break;
+                    case MSG_GET_READ_ONLY_FEATURE_FILE_DESCRIPTOR_MAP:
+                        FeatureFileDescriptorParams ffdParams =
+                                (FeatureFileDescriptorParams) msg.obj;
+                        OnDeviceIntelligenceService.this.onGetReadOnlyFeatureFileDescriptorMap(
+                                ffdParams.feature,
+                                ffdParams.consumer);
+                        break;
+                    case MSG_REGISTER_REMOTE_SERVICES:
+                        mRemoteProcessingService = (IRemoteProcessingService) msg.obj;
+                        break;
+                    case MSG_INFERENCE_SERVICE_CONNECTED:
+                        OnDeviceIntelligenceService.this.onInferenceServiceConnected();
+                        break;
+                    case MSG_INFERENCE_SERVICE_DISCONNECTED:
+                        OnDeviceIntelligenceService.this.onInferenceServiceDisconnected();
+                        break;
+                }
+            }
+        };
     }
 
-    /**
-     * 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
@@ -126,45 +189,37 @@
                 /** {@inheritDoc} */
                 @Override
                 public void ready() {
-                    mHandler.executeOrSendMessage(
-                            obtainMessage(OnDeviceIntelligenceService::onReady,
-                                    OnDeviceIntelligenceService.this));
+                    mHandler.sendEmptyMessage(MSG_ON_READY);
                 }
 
                 @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);
-                                    }));
+                    Message msg = Message.obtain(mHandler, MSG_GET_VERSION,
+                            (LongConsumer) (l -> {
+                                Bundle b = new Bundle();
+                                b.putLong(OnDeviceIntelligenceManager.API_VERSION_BUNDLE_KEY, l);
+                                remoteCallback.sendResult(b);
+                            }));
+                    mHandler.sendMessage(msg);
                 }
 
                 @Override
                 public void listFeatures(int callerUid,
                         IListFeaturesCallback listFeaturesCallback) {
                     Objects.requireNonNull(listFeaturesCallback);
-                    mHandler.executeOrSendMessage(
-                            obtainMessage(
-                                    OnDeviceIntelligenceService::onListFeatures,
-                                    OnDeviceIntelligenceService.this, callerUid,
-                                    wrapListFeaturesCallback(listFeaturesCallback)));
+                    Message msg = Message.obtain(mHandler, MSG_LIST_FEATURES,
+                            callerUid, 0, wrapListFeaturesCallback(listFeaturesCallback));
+                    mHandler.sendMessage(msg);
                 }
 
                 @Override
                 public void getFeature(int callerUid, int id, IFeatureCallback featureCallback) {
                     Objects.requireNonNull(featureCallback);
-                    mHandler.executeOrSendMessage(
-                            obtainMessage(
-                                    OnDeviceIntelligenceService::onGetFeature,
-                                    OnDeviceIntelligenceService.this, callerUid,
-                                    id, wrapFeatureCallback(featureCallback)));
+                    Message msg = Message.obtain(mHandler, MSG_GET_FEATURE,
+                            callerUid, id,
+                            new GetFeatureParams(wrapFeatureCallback(featureCallback)));
+                    mHandler.sendMessage(msg);
                 }
 
 
@@ -173,11 +228,11 @@
                         IFeatureDetailsCallback featureDetailsCallback) {
                     Objects.requireNonNull(feature);
                     Objects.requireNonNull(featureDetailsCallback);
-                    mHandler.executeOrSendMessage(
-                            obtainMessage(
-                                    OnDeviceIntelligenceService::onGetFeatureDetails,
-                                    OnDeviceIntelligenceService.this, callerUid,
-                                    feature, wrapFeatureDetailsCallback(featureDetailsCallback)));
+                    Message msg = Message.obtain(mHandler, MSG_GET_FEATURE_DETAILS,
+                            new FeatureDetailsParams(feature,
+                                    wrapFeatureDetailsCallback(featureDetailsCallback)));
+                    msg.arg1 = callerUid;
+                    mHandler.sendMessage(msg);
                 }
 
                 @Override
@@ -186,18 +241,24 @@
                         IDownloadCallback downloadCallback) {
                     Objects.requireNonNull(feature);
                     Objects.requireNonNull(downloadCallback);
-                    ICancellationSignal transport = null;
+
+                    CancellationSignal cancellationSignal = new CancellationSignal();
                     if (cancellationSignalFuture != null) {
-                        transport = CancellationSignal.createTransport();
+                        ICancellationSignal transport = new ICancellationSignal.Stub() {
+                            @Override
+                            public void cancel() {
+                                cancellationSignal.cancel();
+                            }
+                        };
                         cancellationSignalFuture.complete(transport);
                     }
-                    mHandler.executeOrSendMessage(
-                            obtainMessage(
-                                    OnDeviceIntelligenceService::onDownloadFeature,
-                                    OnDeviceIntelligenceService.this, callerUid,
-                                    feature,
-                                    CancellationSignal.fromTransport(transport),
+
+                    Message msg = Message.obtain(mHandler, MSG_DOWNLOAD_FEATURE,
+                            new DownloadParams(feature,
+                                    cancellationSignalFuture != null ? cancellationSignal : null,
                                     wrapDownloadCallback(downloadCallback)));
+                    msg.arg1 = callerUid;
+                    mHandler.sendMessage(msg);
                 }
 
                 @Override
@@ -205,11 +266,9 @@
                         AndroidFuture<ParcelFileDescriptor> future) {
                     Objects.requireNonNull(fileName);
                     Objects.requireNonNull(future);
-                    mHandler.executeOrSendMessage(
-                            obtainMessage(
-                                    OnDeviceIntelligenceService::onGetReadOnlyFileDescriptor,
-                                    OnDeviceIntelligenceService.this, fileName,
-                                    future));
+                    Message msg = Message.obtain(mHandler, MSG_GET_READ_ONLY_FILE_DESCRIPTOR,
+                            new FileDescriptorParams(fileName, future));
+                    mHandler.sendMessage(msg);
                 }
 
                 @Override
@@ -217,16 +276,15 @@
                         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());
-                                    }));
+                    Message msg = Message.obtain(mHandler,
+                            MSG_GET_READ_ONLY_FEATURE_FILE_DESCRIPTOR_MAP,
+                            new FeatureFileDescriptorParams(feature, parcelFileDescriptorMap -> {
+                                Bundle bundle = new Bundle();
+                                parcelFileDescriptorMap.forEach(bundle::putParcelable);
+                                remoteCallback.sendResult(bundle);
+                                tryClosePfds(parcelFileDescriptorMap.values());
+                            }));
+                    mHandler.sendMessage(msg);
                 }
 
                 @Override
@@ -237,18 +295,12 @@
 
                 @Override
                 public void notifyInferenceServiceConnected() {
-                    mHandler.executeOrSendMessage(
-                            obtainMessage(
-                                    OnDeviceIntelligenceService::onInferenceServiceConnected,
-                                    OnDeviceIntelligenceService.this));
+                    mHandler.sendEmptyMessage(MSG_INFERENCE_SERVICE_CONNECTED);
                 }
 
                 @Override
                 public void notifyInferenceServiceDisconnected() {
-                    mHandler.executeOrSendMessage(
-                            obtainMessage(
-                                    OnDeviceIntelligenceService::onInferenceServiceDisconnected,
-                                    OnDeviceIntelligenceService.this));
+                    mHandler.sendEmptyMessage(MSG_INFERENCE_SERVICE_DISCONNECTED);
                 }
             };
         }
@@ -257,13 +309,77 @@
     }
 
     /**
+     * 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";
+
+    // Parameter holder classes
+    private static class GetFeatureParams {
+        final OutcomeReceiver<Feature, OnDeviceIntelligenceException> callback;
+
+        GetFeatureParams(OutcomeReceiver<Feature, OnDeviceIntelligenceException> callback) {
+            this.callback = callback;
+        }
+    }
+
+    private static class FeatureDetailsParams {
+        final Feature feature;
+        final OutcomeReceiver<FeatureDetails, OnDeviceIntelligenceException> callback;
+
+        FeatureDetailsParams(Feature feature,
+                OutcomeReceiver<FeatureDetails, OnDeviceIntelligenceException> callback) {
+            this.feature = feature;
+            this.callback = callback;
+        }
+    }
+
+    private static class DownloadParams {
+        final Feature feature;
+        final CancellationSignal cancellationSignal;
+        final DownloadCallback callback;
+
+        DownloadParams(Feature feature, CancellationSignal cancellationSignal,
+                DownloadCallback callback) {
+            this.feature = feature;
+            this.cancellationSignal = cancellationSignal;
+            this.callback = callback;
+        }
+    }
+
+    private static class FileDescriptorParams {
+        final String fileName;
+        final AndroidFuture<ParcelFileDescriptor> future;
+
+        FileDescriptorParams(String fileName, AndroidFuture<ParcelFileDescriptor> future) {
+            this.fileName = fileName;
+            this.future = future;
+        }
+    }
+
+    private static class FeatureFileDescriptorParams {
+        final Feature feature;
+        final Consumer<Map<String, ParcelFileDescriptor>> consumer;
+
+        FeatureFileDescriptorParams(Feature feature,
+                Consumer<Map<String, ParcelFileDescriptor>> consumer) {
+            this.feature = feature;
+            this.consumer = consumer;
+        }
+    }
+
+    /**
      * 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.
      *
      * @hide
      */
-    @TestApi
+    @SystemApi
     public void onReady() {
     }
 
@@ -306,7 +422,7 @@
                     new IProcessingUpdateStatusCallback.Stub() {
                         @Override
                         public void onSuccess(PersistableBundle result) {
-                            Binder.withCleanCallingIdentity(() -> {
+                            BinderUtils.withCleanCallingIdentity(() -> {
                                 callbackExecutor.execute(
                                         () -> statusReceiver.onResult(result));
                             });
@@ -314,7 +430,7 @@
 
                         @Override
                         public void onFailure(int errorCode, String errorMessage) {
-                            Binder.withCleanCallingIdentity(() -> callbackExecutor.execute(
+                            BinderUtils.withCleanCallingIdentity(() -> callbackExecutor.execute(
                                     () -> statusReceiver.onError(
                                             new OnDeviceIntelligenceException(
                                                     errorCode, errorMessage))));
@@ -459,7 +575,7 @@
     private void onGetReadOnlyFileDescriptor(@NonNull String fileName,
             @NonNull AndroidFuture<ParcelFileDescriptor> future) {
         Slog.v(TAG, "onGetReadOnlyFileDescriptor " + fileName);
-        Binder.withCleanCallingIdentity(() -> {
+        BinderUtils.withCleanCallingIdentity(() -> {
             Slog.v(TAG,
                     "onGetReadOnlyFileDescriptor: " + fileName + " under internal app storage.");
             File f = new File(getBaseContext().getFilesDir(), fileName);
@@ -476,7 +592,11 @@
             } finally {
                 future.complete(pfd);
                 if (pfd != null) {
-                    pfd.close();
+                    try {
+                        pfd.close();
+                    } catch (IOException e) {
+                        Log.w(TAG, "Error closing FD", e);
+                    }
                 }
             }
         });
diff --git a/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/OnDeviceSandboxedInferenceService.java b/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/OnDeviceSandboxedInferenceService.java
index a6aee247..315dbaf 100644
--- a/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/OnDeviceSandboxedInferenceService.java
+++ b/packages/NeuralNetworks/framework/java/android/service/ondeviceintelligence/OnDeviceSandboxedInferenceService.java
@@ -19,8 +19,6 @@
 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;
@@ -31,7 +29,9 @@
 import android.annotation.SystemApi;
 import android.app.Service;
 import android.app.ondeviceintelligence.Feature;
+import android.app.ondeviceintelligence.ICancellationSignal;
 import android.app.ondeviceintelligence.IProcessingSignal;
+import android.app.ondeviceintelligence.IRemoteCallback;
 import android.app.ondeviceintelligence.IResponseCallback;
 import android.app.ondeviceintelligence.IStreamingResponseCallback;
 import android.app.ondeviceintelligence.ITokenInfoCallback;
@@ -48,11 +48,9 @@
 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.Message;
 import android.os.OutcomeReceiver;
 import android.os.ParcelFileDescriptor;
 import android.os.PersistableBundle;
@@ -61,7 +59,8 @@
 import android.util.Log;
 import android.util.Slog;
 
-import com.android.internal.infra.AndroidFuture;
+import com.android.modules.utils.AndroidFuture;
+import com.android.modules.utils.HandlerExecutor;
 
 import java.io.FileInputStream;
 import java.io.FileNotFoundException;
@@ -74,7 +73,7 @@
 
 /**
  * Abstract base class for performing inference in a isolated process. This service exposes its
- * methods via {@link OnDeviceIntelligenceManager}.
+ * methods via {@link android.app.ondeviceintelligence.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
@@ -100,6 +99,12 @@
 public abstract class OnDeviceSandboxedInferenceService extends Service {
     private static final String TAG = OnDeviceSandboxedInferenceService.class.getSimpleName();
 
+    private static final int MSG_TOKEN_INFO_REQUEST = 1;
+    private static final int MSG_PROCESS_REQUEST_STREAMING = 2;
+    private static final int MSG_PROCESS_REQUEST = 3;
+    private static final int MSG_UPDATE_PROCESSING_STATE = 4;
+
+
     /**
      * @hide
      */
@@ -133,12 +138,12 @@
      * @hide
      */
     public static final String MODEL_LOADED_BROADCAST_INTENT =
-        "android.service.ondeviceintelligence.MODEL_LOADED";
+            "android.service.ondeviceintelligence.MODEL_LOADED";
     /**
      * @hide
      */
     public static final String MODEL_UNLOADED_BROADCAST_INTENT =
-        "android.service.ondeviceintelligence.MODEL_UNLOADED";
+            "android.service.ondeviceintelligence.MODEL_UNLOADED";
 
     /**
      * @hide
@@ -152,12 +157,115 @@
     @Override
     public void onCreate() {
         super.onCreate();
-        mHandler = new Handler(Looper.getMainLooper(), null /* callback */, true /* async */);
+        mHandler = new Handler(Looper.getMainLooper()) {
+            @Override
+            public void handleMessage(Message msg) {
+                switch (msg.what) {
+                    case MSG_TOKEN_INFO_REQUEST:
+                        TokenInfoParams params = (TokenInfoParams) msg.obj;
+                        OnDeviceSandboxedInferenceService.this.onTokenInfoRequest(
+                                msg.arg1,
+                                params.feature,
+                                params.request,
+                                params.cancellationSignal,
+                                params.callback);
+                        break;
+                    case MSG_PROCESS_REQUEST_STREAMING:
+                        StreamingRequestParams streamParams = (StreamingRequestParams) msg.obj;
+                        OnDeviceSandboxedInferenceService.this.onProcessRequestStreaming(
+                                msg.arg1,
+                                streamParams.feature,
+                                streamParams.request,
+                                msg.arg2,
+                                streamParams.cancellationSignal,
+                                streamParams.processingSignal,
+                                streamParams.callback);
+                        break;
+                    case MSG_PROCESS_REQUEST:
+                        RequestParams requestParams = (RequestParams) msg.obj;
+                        OnDeviceSandboxedInferenceService.this.onProcessRequest(
+                                msg.arg1,
+                                requestParams.feature,
+                                requestParams.request,
+                                msg.arg2,
+                                requestParams.cancellationSignal,
+                                requestParams.processingSignal,
+                                requestParams.callback);
+                        break;
+                    case MSG_UPDATE_PROCESSING_STATE:
+                        UpdateStateParams stateParams = (UpdateStateParams) msg.obj;
+                        OnDeviceSandboxedInferenceService.this.onUpdateProcessingState(
+                                stateParams.processingState,
+                                stateParams.callback);
+                        break;
+                }
+            }
+        };
     }
 
-    /**
-     * @hide
-     */
+    // Parameter holder classes
+    private static class TokenInfoParams {
+        final Feature feature;
+        final Bundle request;
+        final CancellationSignal cancellationSignal;
+        final OutcomeReceiver<TokenInfo, OnDeviceIntelligenceException> callback;
+
+        TokenInfoParams(Feature feature, Bundle request, CancellationSignal cancellationSignal,
+                OutcomeReceiver<TokenInfo, OnDeviceIntelligenceException> callback) {
+            this.feature = feature;
+            this.request = request;
+            this.cancellationSignal = cancellationSignal;
+            this.callback = callback;
+        }
+    }
+
+    private static class StreamingRequestParams {
+        final Feature feature;
+        final Bundle request;
+        final CancellationSignal cancellationSignal;
+        final ProcessingSignal processingSignal;
+        final StreamingProcessingCallback callback;
+
+        StreamingRequestParams(Feature feature, Bundle request,
+                CancellationSignal cancellationSignal, ProcessingSignal processingSignal,
+                StreamingProcessingCallback callback) {
+            this.feature = feature;
+            this.request = request;
+            this.cancellationSignal = cancellationSignal;
+            this.processingSignal = processingSignal;
+            this.callback = callback;
+        }
+    }
+
+    private static class RequestParams {
+        final Feature feature;
+        final Bundle request;
+        final CancellationSignal cancellationSignal;
+        final ProcessingSignal processingSignal;
+        final ProcessingCallback callback;
+
+        RequestParams(Feature feature, Bundle request,
+                CancellationSignal cancellationSignal, ProcessingSignal processingSignal,
+                ProcessingCallback callback) {
+            this.feature = feature;
+            this.request = request;
+            this.cancellationSignal = cancellationSignal;
+            this.processingSignal = processingSignal;
+            this.callback = callback;
+        }
+    }
+
+    private static class UpdateStateParams {
+        final Bundle processingState;
+        final OutcomeReceiver<PersistableBundle, OnDeviceIntelligenceException> callback;
+
+        UpdateStateParams(Bundle processingState,
+                OutcomeReceiver<PersistableBundle, OnDeviceIntelligenceException> callback) {
+            this.processingState = processingState;
+            this.callback = callback;
+        }
+    }
+
     @Nullable
     @Override
     public final IBinder onBind(@NonNull Intent intent) {
@@ -168,8 +276,7 @@
                         IRemoteCallback remoteCallback) throws RemoteException {
                     Objects.requireNonNull(storageService);
                     mRemoteStorageService = storageService;
-                    remoteCallback.sendResult(
-                            Bundle.EMPTY); //to notify caller uid to system-server.
+                    remoteCallback.sendResult(Bundle.EMPTY);
                 }
 
                 @Override
@@ -178,34 +285,42 @@
                         ITokenInfoCallback tokenInfoCallback) {
                     Objects.requireNonNull(feature);
                     Objects.requireNonNull(tokenInfoCallback);
-                    ICancellationSignal transport = null;
+                    CancellationSignal cancellationSignal = new CancellationSignal();
                     if (cancellationSignalFuture != null) {
-                        transport = CancellationSignal.createTransport();
+                        ICancellationSignal transport = new ICancellationSignal.Stub() {
+                            @Override
+                            public void cancel() {
+                                cancellationSignal.cancel();
+                            }
+                        };
                         cancellationSignalFuture.complete(transport);
                     }
 
-                    mHandler.executeOrSendMessage(
-                            obtainMessage(
-                                    OnDeviceSandboxedInferenceService::onTokenInfoRequest,
-                                    OnDeviceSandboxedInferenceService.this,
-                                    callerUid, feature,
-                                    request,
-                                    CancellationSignal.fromTransport(transport),
+                    Message msg = Message.obtain(mHandler, MSG_TOKEN_INFO_REQUEST,
+                            callerUid, 0,
+                            new TokenInfoParams(feature, request,
+                                    cancellationSignalFuture != null ? cancellationSignal : null,
                                     wrapTokenInfoCallback(tokenInfoCallback)));
+                    mHandler.sendMessage(msg);
                 }
 
                 @Override
-                public void processRequestStreaming(int callerUid, Feature feature, Bundle request,
-                        int requestType,
+                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;
+                    CancellationSignal cancellationSignal = new CancellationSignal();
                     if (cancellationSignalFuture != null) {
-                        transport = CancellationSignal.createTransport();
+                        ICancellationSignal transport = new ICancellationSignal.Stub() {
+                            @Override
+                            public void cancel() {
+                                cancellationSignal.cancel();
+                            }
+                        };
                         cancellationSignalFuture.complete(transport);
                     }
                     IProcessingSignal processingSignalTransport = null;
@@ -214,30 +329,32 @@
                         processingSignalFuture.complete(processingSignalTransport);
                     }
 
-
-                    mHandler.executeOrSendMessage(
-                            obtainMessage(
-                                    OnDeviceSandboxedInferenceService::onProcessRequestStreaming,
-                                    OnDeviceSandboxedInferenceService.this, callerUid,
-                                    feature,
-                                    request,
-                                    requestType,
-                                    CancellationSignal.fromTransport(transport),
+                    Message msg = Message.obtain(mHandler, MSG_PROCESS_REQUEST_STREAMING,
+                            callerUid, requestType,
+                            new StreamingRequestParams(feature, request,
+                                    cancellationSignalFuture != null ? cancellationSignal : null,
                                     ProcessingSignal.fromTransport(processingSignalTransport),
                                     wrapStreamingResponseCallback(callback)));
+                    mHandler.sendMessage(msg);
                 }
 
                 @Override
-                public void processRequest(int callerUid, Feature feature, Bundle request,
-                        int requestType,
+                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;
+
+                    CancellationSignal cancellationSignal = new CancellationSignal();
                     if (cancellationSignalFuture != null) {
-                        transport = CancellationSignal.createTransport();
+                        ICancellationSignal transport = new ICancellationSignal.Stub() {
+                            @Override
+                            public void cancel() {
+                                cancellationSignal.cancel();
+                            }
+                        };
                         cancellationSignalFuture.complete(transport);
                     }
                     IProcessingSignal processingSignalTransport = null;
@@ -245,14 +362,14 @@
                         processingSignalTransport = ProcessingSignal.createTransport();
                         processingSignalFuture.complete(processingSignalTransport);
                     }
-                    mHandler.executeOrSendMessage(
-                            obtainMessage(
-                                    OnDeviceSandboxedInferenceService::onProcessRequest,
-                                    OnDeviceSandboxedInferenceService.this, callerUid, feature,
-                                    request, requestType,
-                                    CancellationSignal.fromTransport(transport),
+
+                    Message msg = Message.obtain(mHandler, MSG_PROCESS_REQUEST,
+                            callerUid, requestType,
+                            new RequestParams(feature, request,
+                                    cancellationSignalFuture != null ? cancellationSignal : null,
                                     ProcessingSignal.fromTransport(processingSignalTransport),
                                     wrapResponseCallback(callback)));
+                    mHandler.sendMessage(msg);
                 }
 
                 @Override
@@ -260,11 +377,11 @@
                         IProcessingUpdateStatusCallback callback) {
                     Objects.requireNonNull(processingState);
                     Objects.requireNonNull(callback);
-                    mHandler.executeOrSendMessage(
-                            obtainMessage(
-                                    OnDeviceSandboxedInferenceService::onUpdateProcessingState,
-                                    OnDeviceSandboxedInferenceService.this, processingState,
+
+                    Message msg = Message.obtain(mHandler, MSG_UPDATE_PROCESSING_STATE,
+                            new UpdateStateParams(processingState,
                                     wrapOutcomeReceiver(callback)));
+                    mHandler.sendMessage(msg);
                 }
             };
         }
@@ -471,7 +588,7 @@
             IResponseCallback callback) {
         return new ProcessingCallback() {
             @Override
-            public void onResult(@androidx.annotation.NonNull Bundle result) {
+            public void onResult(@NonNull Bundle result) {
                 try {
                     callback.onSuccess(result);
                 } catch (RemoteException e) {
@@ -507,7 +624,7 @@
             IStreamingResponseCallback callback) {
         return new StreamingProcessingCallback() {
             @Override
-            public void onPartialResult(@androidx.annotation.NonNull Bundle partialResult) {
+            public void onPartialResult(@NonNull Bundle partialResult) {
                 try {
                     callback.onNewContent(partialResult);
                 } catch (RemoteException e) {
@@ -516,7 +633,7 @@
             }
 
             @Override
-            public void onResult(@androidx.annotation.NonNull Bundle result) {
+            public void onResult(@NonNull Bundle result) {
                 try {
                     callback.onSuccess(result);
                 } catch (RemoteException e) {
@@ -549,7 +666,7 @@
     }
 
     private RemoteCallback wrapRemoteCallback(
-            @androidx.annotation.NonNull Consumer<Bundle> contentCallback) {
+            @NonNull Consumer<Bundle> contentCallback) {
         return new RemoteCallback(
                 result -> {
                     if (result != null) {
@@ -604,7 +721,7 @@
 
             @Override
             public void onError(
-                    @androidx.annotation.NonNull OnDeviceIntelligenceException error) {
+                    @NonNull OnDeviceIntelligenceException error) {
                 try {
                     callback.onFailure(error.getErrorCode(), error.getMessage());
                 } catch (RemoteException e) {
diff --git a/packages/NeuralNetworks/service/Android.bp b/packages/NeuralNetworks/service/Android.bp
new file mode 100644
index 0000000..05c603f
--- /dev/null
+++ b/packages/NeuralNetworks/service/Android.bp
@@ -0,0 +1,29 @@
+// 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 {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+filegroup {
+    name: "service-ondeviceintelligence-sources",
+    srcs: [
+        "java/**/*.java",
+    ],
+    path: "java",
+    visibility: [
+        "//frameworks/base:__subpackages__",
+        "//packages/modules/NeuralNetworks:__subpackages__",
+    ],
+}
diff --git a/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/BundleUtil.java b/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/BundleUtil.java
index 7dd8f2f..53ef9e8 100644
--- a/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/BundleUtil.java
+++ b/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/BundleUtil.java
@@ -42,7 +42,7 @@
 import android.system.Os;
 import android.util.Log;
 
-import com.android.internal.infra.AndroidFuture;
+import com.android.modules.utils.AndroidFuture;
 
 import java.util.concurrent.Executor;
 import java.util.concurrent.TimeoutException;
@@ -50,6 +50,8 @@
 /**
  * Util methods for ensuring the Bundle passed in various methods are read-only and restricted to
  * some known types.
+ *
+ * @hide
  */
 public class BundleUtil {
     private static final String TAG = "BundleUtil";
@@ -76,7 +78,7 @@
                  * {@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);
+                bundle.putParcelable(key, null);
                 continue;
             }
             if (canMarshall(obj) || obj instanceof CursorWindow) {
@@ -122,7 +124,7 @@
                  * {@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);
+                bundle.putParcelable(key, null);
                 continue;
             }
             if (canMarshall(obj)) {
@@ -167,7 +169,7 @@
                  * {@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);
+                bundle.putParcelable(key, null);
                 continue;
             }
             if (canMarshall(obj)) {
@@ -317,9 +319,13 @@
         };
     }
 
-    private static boolean canMarshall(Object obj) {
-        return obj instanceof byte[] || obj instanceof PersistableBundle
-                || PersistableBundle.isValidType(obj);
+    private static boolean canMarshall(Object value) {
+        return (value instanceof byte[]) || (value instanceof Integer) || (value instanceof Long) ||
+                (value instanceof Double) || (value instanceof String) ||
+                (value instanceof int[]) || (value instanceof long[]) ||
+                (value instanceof double[]) || (value instanceof String[]) ||
+                (value instanceof PersistableBundle) || (value == null) ||
+                (value instanceof Boolean) || (value instanceof boolean[]);
     }
 
     private static void ensureValidBundle(Bundle bundle) {
@@ -364,7 +370,7 @@
             }
         } catch (ErrnoException e) {
             throw new BadParcelableException(
-                    "Invalid File descriptor passed in the Bundle.", e);
+                    "Invalid File descriptor passed in the Bundle.");
         }
     }
 
diff --git a/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/InferenceInfoStore.java b/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/InferenceInfoStore.java
index bef3f80..e8a1b322 100644
--- a/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/InferenceInfoStore.java
+++ b/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/InferenceInfoStore.java
@@ -28,6 +28,9 @@
 import java.util.List;
 import java.util.TreeSet;
 
+/**
+ * @hide
+ */
 public class InferenceInfoStore {
     private static final String TAG = "InferenceInfoStore";
     private final TreeSet<InferenceInfo> inferenceInfos;
@@ -98,4 +101,4 @@
                 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/OnDeviceIntelligenceManagerInternal.java b/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerLocal.java
similarity index 61%
rename from packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerInternal.java
rename to packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerLocal.java
index 1450dc0..6badc53 100644
--- a/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerInternal.java
+++ b/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerLocal.java
@@ -16,7 +16,21 @@
 
 package com.android.server.ondeviceintelligence;
 
-public interface OnDeviceIntelligenceManagerInternal {
+import static android.app.ondeviceintelligence.flags.Flags.FLAG_ENABLE_ON_DEVICE_INTELLIGENCE_MODULE;
+
+import android.annotation.FlaggedApi;
+import android.annotation.SystemApi;
+import android.annotation.SystemApi.Client;
+
+/**
+ * Exposes APIs to {@code system_server} components outside of the module boundaries.
+ * <p> This API should be access using {@link com.android.server.LocalManagerRegistry}. </p>
+ *
+ * @hide
+ */
+@SystemApi(client = Client.SYSTEM_SERVER)
+@FlaggedApi(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE_MODULE)
+public interface OnDeviceIntelligenceManagerLocal {
     /**
      * Gets the uid for the process that is currently hosting the
      * {@link android.service.ondeviceintelligence.OnDeviceSandboxedInferenceService} registered on
diff --git a/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerService.java b/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerService.java
index b0d69e6..607ec1c 100644
--- a/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerService.java
+++ b/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerService.java
@@ -16,38 +16,42 @@
 
 package com.android.server.ondeviceintelligence;
 
+import static android.app.ondeviceintelligence.flags.Flags.enableOnDeviceIntelligenceModule;
+
+import static android.app.ondeviceintelligence.OnDeviceIntelligenceManager.ON_DEVICE_INTELLIGENCE_IDLE_TIMEOUT_MS;
 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_LOADED_BUNDLE_KEY;
 import static android.service.ondeviceintelligence.OnDeviceSandboxedInferenceService.MODEL_UNLOADED_BROADCAST_INTENT;
+import static android.service.ondeviceintelligence.OnDeviceSandboxedInferenceService.MODEL_UNLOADED_BUNDLE_KEY;
 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.validatePfdReadOnly;
 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.ICancellationSignal;
 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.IRemoteCallback;
 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.app.ondeviceintelligence.utils.BinderUtils;
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
@@ -58,16 +62,12 @@
 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;
@@ -82,17 +82,14 @@
 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.LocalServices;
+import com.android.modules.utils.AndroidFuture;
+import com.android.modules.utils.BackgroundThread;
+import com.android.modules.utils.ServiceConnector;
+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;
@@ -182,9 +179,11 @@
     public void onStart() {
         publishBinderService(
                 Context.ON_DEVICE_INTELLIGENCE_SERVICE, getOnDeviceIntelligenceManagerService(),
-                /* allowIsolated = */true);
-        LocalServices.addService(OnDeviceIntelligenceManagerInternal.class,
-                this::getRemoteInferenceServiceUid);
+                /* allowIsolated = */ true);
+        if (enableOnDeviceIntelligenceModule()) {
+            LocalManagerRegistry.addManager(OnDeviceIntelligenceManagerLocal.class,
+                    this::getRemoteInferenceServiceUid);
+        }
     }
 
     @Override
@@ -203,10 +202,10 @@
     public void onUserUnlocked(@NonNull TargetUser user) {
         Slog.d(TAG, "onUserUnlocked: " + user.getUserHandle());
         //connect to remote services(if available) during boot.
-        if(user.getUserHandle().equals(UserHandle.SYSTEM)) {
+        if (user.getUserHandle().equals(UserHandle.SYSTEM)) {
             try {
-                ensureRemoteInferenceServiceInitialized();
-                ensureRemoteIntelligenceServiceInitialized();
+                ensureRemoteInferenceServiceInitialized(/* throwServiceIfInvalid */ false);
+                ensureRemoteIntelligenceServiceInitialized(/* throwServiceIfInvalid */ false);
             } catch (Exception e) {
                 Slog.w(TAG, "Couldn't pre-start remote ondeviceintelligence services.", e);
             }
@@ -251,7 +250,7 @@
                     remoteCallback.sendResult(null);
                     return;
                 }
-                ensureRemoteIntelligenceServiceInitialized();
+                ensureRemoteIntelligenceServiceInitialized(/* throwServiceIfInvalid */ true);
                 mRemoteOnDeviceIntelligenceService.postAsync(
                         service -> {
                             AndroidFuture future = new AndroidFuture();
@@ -279,7 +278,7 @@
                             PersistableBundle.EMPTY);
                     return;
                 }
-                ensureRemoteIntelligenceServiceInitialized();
+                ensureRemoteIntelligenceServiceInitialized(/* throwServiceIfInvalid */ true);
                 int callerUid = Binder.getCallingUid();
                 mRemoteOnDeviceIntelligenceService.postAsync(
                         service -> {
@@ -317,7 +316,7 @@
                             PersistableBundle.EMPTY);
                     return;
                 }
-                ensureRemoteIntelligenceServiceInitialized();
+                ensureRemoteIntelligenceServiceInitialized(/* throwServiceIfInvalid */ true);
                 int callerUid = Binder.getCallingUid();
                 mRemoteOnDeviceIntelligenceService.postAsync(
                         service -> {
@@ -361,7 +360,7 @@
                             PersistableBundle.EMPTY);
                     return;
                 }
-                ensureRemoteIntelligenceServiceInitialized();
+                ensureRemoteIntelligenceServiceInitialized(/* throwServiceIfInvalid */ true);
                 int callerUid = Binder.getCallingUid();
                 mRemoteOnDeviceIntelligenceService.postAsync(
                         service -> {
@@ -404,7 +403,7 @@
                             "OnDeviceIntelligenceManagerService is unavailable",
                             PersistableBundle.EMPTY);
                 }
-                ensureRemoteIntelligenceServiceInitialized();
+                ensureRemoteIntelligenceServiceInitialized(/* throwServiceIfInvalid */ true);
                 int callerUid = Binder.getCallingUid();
                 mRemoteOnDeviceIntelligenceService.postAsync(
                         service -> {
@@ -444,7 +443,7 @@
                                 "OnDeviceIntelligenceManagerService is unavailable",
                                 PersistableBundle.EMPTY);
                     }
-                    ensureRemoteInferenceServiceInitialized();
+                    ensureRemoteInferenceServiceInitialized(/* throwServiceIfInvalid */ true);
                     int callerUid = Binder.getCallingUid();
                     result = mRemoteInferenceService.postAsync(
                             service -> {
@@ -488,7 +487,7 @@
                                 "OnDeviceIntelligenceManagerService is unavailable",
                                 PersistableBundle.EMPTY);
                     }
-                    ensureRemoteInferenceServiceInitialized();
+                    ensureRemoteInferenceServiceInitialized(/* throwServiceIfInvalid */ true);
                     int callerUid = Binder.getCallingUid();
                     result = mRemoteInferenceService.postAsync(
                             service -> {
@@ -534,7 +533,7 @@
                                 "OnDeviceIntelligenceManagerService is unavailable",
                                 PersistableBundle.EMPTY);
                     }
-                    ensureRemoteInferenceServiceInitialized();
+                    ensureRemoteInferenceServiceInitialized(/* throwServiceIfInvalid */ true);
                     int callerUid = Binder.getCallingUid();
                     result = mRemoteInferenceService.postAsync(
                             service -> {
@@ -559,20 +558,31 @@
             }
 
             @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);
+            public int handleShellCommand(@NonNull ParcelFileDescriptor in,
+                    @NonNull ParcelFileDescriptor out, @NonNull ParcelFileDescriptor err,
+                    @NonNull String[] args) {
+                return new com.android.server.ondeviceintelligence.OnDeviceIntelligenceShellCommand(
+                        OnDeviceIntelligenceManagerService.this).exec(
+                        this,
+                        in.getFileDescriptor(),
+                        out.getFileDescriptor(),
+                        err.getFileDescriptor(),
+                        args);
             }
         };
     }
 
-    private void ensureRemoteIntelligenceServiceInitialized() {
+    private boolean ensureRemoteIntelligenceServiceInitialized(boolean throwIfServiceInvalid) {
         synchronized (mLock) {
             if (mRemoteOnDeviceIntelligenceService == null) {
                 String serviceName = getServiceNames()[0];
-                Binder.withCleanCallingIdentity(() -> validateServiceElevated(serviceName, false));
-                mRemoteOnDeviceIntelligenceService = new RemoteOnDeviceIntelligenceService(mContext,
+                if (!BinderUtils.withCleanCallingIdentity(
+                        () -> validateServiceElevated(serviceName, false,
+                                throwIfServiceInvalid))) {
+                    return false;
+                }
+                mRemoteOnDeviceIntelligenceService = new RemoteOnDeviceIntelligenceService(
+                        mContext,
                         ComponentName.unflattenFromString(serviceName),
                         UserHandle.SYSTEM.getIdentifier());
                 mRemoteOnDeviceIntelligenceService.setServiceLifecycleCallbacks(
@@ -591,6 +601,7 @@
                         });
             }
         }
+        return true;
     }
 
     @NonNull
@@ -604,13 +615,21 @@
                     AndroidFuture<Void> result = null;
                     try {
                         sanitizeStateParams(processingState);
-                        ensureRemoteInferenceServiceInitialized();
-                        result = mRemoteInferenceService.post(
-                                service -> service.updateProcessingState(
-                                        processingState, callback));
-                        result.whenCompleteAsync(
-                                (c, e) -> BundleUtil.tryCloseResource(processingState),
-                                resourceClosingExecutor);
+                        if (ensureRemoteInferenceServiceInitialized(/* throwServiceIfInvalid */
+                                false)) {
+                            result = mRemoteInferenceService.post(
+                                    service -> service.updateProcessingState(
+                                            processingState, callback));
+                            result.whenCompleteAsync(
+                                    (c, e) -> BundleUtil.tryCloseResource(processingState),
+                                    resourceClosingExecutor);
+                        } else {
+                            callback.onFailure(
+                                    OnDeviceIntelligenceException.PROCESSING_ERROR_SERVICE_UNAVAILABLE,
+                                    "Remote service cannot be initialized.");
+                        }
+                    } catch (RemoteException e) {
+                        Slog.w("Failed to invoke updateProcessingState", e);
                     } finally {
                         if (result == null) {
                             resourceClosingExecutor.execute(
@@ -622,11 +641,14 @@
         };
     }
 
-    private void ensureRemoteInferenceServiceInitialized() {
+    private boolean ensureRemoteInferenceServiceInitialized(boolean throwIfServiceInvalid) {
         synchronized (mLock) {
             if (mRemoteInferenceService == null) {
                 String serviceName = getServiceNames()[1];
-                Binder.withCleanCallingIdentity(() -> validateServiceElevated(serviceName, true));
+                if (!BinderUtils.withCleanCallingIdentity(
+                        () -> validateServiceElevated(serviceName, true, throwIfServiceInvalid))) {
+                    return false;
+                }
                 mRemoteInferenceService = new RemoteOnDeviceSandboxedInferenceService(mContext,
                         ComponentName.unflattenFromString(serviceName),
                         UserHandle.SYSTEM.getIdentifier());
@@ -636,7 +658,11 @@
                             public void onConnected(
                                     @NonNull IOnDeviceSandboxedInferenceService service) {
                                 try {
-                                    ensureRemoteIntelligenceServiceInitialized();
+                                    if (!ensureRemoteIntelligenceServiceInitialized(
+                                            /* throwServiceIfInvalid */
+                                            false)) {
+                                        return;
+                                    }
                                     service.registerRemoteStorageService(
                                             getIRemoteStorageService(), new IRemoteCallback.Stub() {
                                                 @Override
@@ -659,20 +685,29 @@
                             @Override
                             public void onDisconnected(
                                     @NonNull IOnDeviceSandboxedInferenceService service) {
-                                ensureRemoteIntelligenceServiceInitialized();
+                                if (!ensureRemoteIntelligenceServiceInitialized(
+                                        /* throwServiceIfInvalid */
+                                        false)) {
+                                    return;
+                                }
                                 mRemoteOnDeviceIntelligenceService.run(
                                         IOnDeviceIntelligenceService::notifyInferenceServiceDisconnected);
                             }
 
                             @Override
                             public void onBinderDied() {
-                                ensureRemoteIntelligenceServiceInitialized();
+                                if (!ensureRemoteIntelligenceServiceInitialized(
+                                        /* throwServiceIfInvalid */
+                                        false)) {
+                                    return;
+                                }
                                 mRemoteOnDeviceIntelligenceService.run(
                                         IOnDeviceIntelligenceService::notifyInferenceServiceDisconnected);
                             }
                         });
             }
         }
+        return true;
     }
 
     private void registerModelLoadingBroadcasts(IOnDeviceSandboxedInferenceService service) {
@@ -743,9 +778,13 @@
             if (mTemporaryConfigNamespace != null) {
                 return mTemporaryConfigNamespace;
             }
-
-            return mContext.getResources().getString(
-                    R.string.config_defaultOnDeviceIntelligenceDeviceConfigNamespace);
+            return mContext.getResources()
+                    .getString(
+                            mContext.getResources()
+                                    .getIdentifier(
+                                            "config_defaultOnDeviceIntelligenceDeviceConfigNamespace",
+                                            "string",
+                                            "android"));
         }
     }
 
@@ -759,7 +798,11 @@
         }
         Bundle bundle = new Bundle();
         bundle.putParcelable(DEVICE_CONFIG_UPDATE_BUNDLE_KEY, persistableBundle);
-        ensureRemoteInferenceServiceInitialized();
+        if (!ensureRemoteIntelligenceServiceInitialized(
+                /* throwServiceIfInvalid */
+                false)) {
+            return;
+        }
         mRemoteInferenceService.run(service -> service.updateProcessingState(bundle,
                 new IProcessingUpdateStatusCallback.Stub() {
                     @Override
@@ -782,7 +825,13 @@
             public void getReadOnlyFileDescriptor(
                     String filePath,
                     AndroidFuture<ParcelFileDescriptor> future) {
-                ensureRemoteIntelligenceServiceInitialized();
+                if (!ensureRemoteIntelligenceServiceInitialized(
+                        /* throwServiceIfInvalid */
+                        false)) {
+                    future.completeExceptionally(new OnDeviceIntelligenceException(
+                            OnDeviceIntelligenceException.PROCESSING_ERROR_NOT_AVAILABLE));
+                    return;
+                }
                 AndroidFuture<ParcelFileDescriptor> pfdFuture = new AndroidFuture<>();
                 mRemoteOnDeviceIntelligenceService.run(
                         service -> service.getReadOnlyFileDescriptor(
@@ -805,7 +854,7 @@
             public void getReadOnlyFeatureFileDescriptorMap(
                     Feature feature,
                     RemoteCallback remoteCallback) {
-                ensureRemoteIntelligenceServiceInitialized();
+                ensureRemoteIntelligenceServiceInitialized(/* throwServiceIfInvalid */ true);
                 mRemoteOnDeviceIntelligenceService.run(
                         service -> service.getReadOnlyFeatureFileDescriptorMap(
                                 feature,
@@ -829,40 +878,48 @@
         };
     }
 
-    private void validateServiceElevated(String serviceName, boolean checkIsolated) {
+    private boolean validateServiceElevated(String serviceName, boolean checkIsolated,
+            boolean throwIfServiceInvalid) {
         try {
             if (TextUtils.isEmpty(serviceName)) {
-                throw new IllegalStateException(
-                        "Remote service is not configured to complete the request");
+                if (throwIfServiceInvalid) {
+                    throw new IllegalStateException(
+                            "Remote service is not configured to complete the request");
+                }
+                return false;
             }
             ComponentName serviceComponent = ComponentName.unflattenFromString(
                     serviceName);
-            ServiceInfo serviceInfo = AppGlobals.getPackageManager().getServiceInfo(
+            ServiceInfo serviceInfo = mContext.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;
-                }
-
+                            | PackageManager.MATCH_DIRECT_BOOT_UNAWARE);
+            if (!checkIsolated) {
                 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 {
+                        Manifest.permission.BIND_ON_DEVICE_INTELLIGENCE_SERVICE);
+                return true;
+            }
+
+            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");
+            }
+        } catch (PackageManager.NameNotFoundException e) {
+            if (throwIfServiceInvalid) {
                 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);
+            return false;
+        } catch (SecurityException e) {
+            if (throwIfServiceInvalid) {
+                throw e;
+            }
+            return false;
         }
+        return true;
     }
 
     private static void checkServiceRequiresPermission(ServiceInfo serviceInfo,
@@ -870,8 +927,8 @@
         final String permission = serviceInfo.permission;
         if (!requiredPermission.equals(permission)) {
             throw new SecurityException(String.format(
-                    "Service %s requires %s permission. Found %s permission",
-                    serviceInfo.getComponentName(),
+                    "%s requires %s permission. Found %s permission",
+                    serviceInfo,
                     requiredPermission,
                     serviceInfo.permission));
         }
@@ -909,10 +966,22 @@
                 return mTemporaryServiceNames;
             }
         }
-        return new String[]{mContext.getResources().getString(
-                R.string.config_defaultOnDeviceIntelligenceService),
-                mContext.getResources().getString(
-                        R.string.config_defaultOnDeviceSandboxedInferenceService)};
+        return new String[]{
+                mContext.getResources()
+                        .getString(
+                        mContext.getResources()
+                                .getIdentifier(
+                                        "config_defaultOnDeviceIntelligenceService",
+                                        "string",
+                                        "android")),
+                mContext.getResources()
+                        .getString(
+                        mContext.getResources()
+                                .getIdentifier(
+                                        "config_defaultOnDeviceSandboxedInferenceService",
+                                        "string",
+                                        "android"))
+        };
     }
 
     protected String[] getBroadcastKeys() throws Resources.NotFoundException {
@@ -923,7 +992,7 @@
             }
         }
 
-        return new String[]{ MODEL_LOADED_BROADCAST_INTENT, MODEL_UNLOADED_BROADCAST_INTENT };
+        return new String[]{MODEL_LOADED_BROADCAST_INTENT, MODEL_UNLOADED_BROADCAST_INTENT};
     }
 
     @RequiresPermission(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE)
@@ -1068,7 +1137,7 @@
 
     private synchronized Handler getTemporaryHandler() {
         if (mTemporaryHandler == null) {
-            mTemporaryHandler = new Handler(Looper.getMainLooper(), null, true) {
+            mTemporaryHandler = new Handler(Looper.getMainLooper()) {
                 @Override
                 public void handleMessage(Message msg) {
                     synchronized (mLock) {
@@ -1090,10 +1159,13 @@
         return mTemporaryHandler;
     }
 
+    // Using #getLong here as the timeout settings are only applicable to the services running in
+    // SYSTEM user only.
+    @SuppressWarnings("NonUserGetterCalled")
     private long getIdleTimeoutMs() {
-        return Settings.Secure.getLongForUser(mContext.getContentResolver(),
-                Settings.Secure.ON_DEVICE_INTELLIGENCE_IDLE_TIMEOUT_MS, TimeUnit.HOURS.toMillis(1),
-                mContext.getUserId());
+        return Settings.Secure.getLong(mContext.getContentResolver(),
+                ON_DEVICE_INTELLIGENCE_IDLE_TIMEOUT_MS,
+                TimeUnit.HOURS.toMillis(1));
     }
 
     private int getRemoteInferenceServiceUid() {
diff --git a/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceShellCommand.java b/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceShellCommand.java
index d2c84fa..c641de8 100644
--- a/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceShellCommand.java
+++ b/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceShellCommand.java
@@ -18,12 +18,16 @@
 
 import android.annotation.NonNull;
 import android.os.Binder;
-import android.os.ShellCommand;
+
+import com.android.modules.utils.BasicShellCommandHandler;
 
 import java.io.PrintWriter;
 import java.util.Objects;
 
-final class OnDeviceIntelligenceShellCommand extends ShellCommand {
+/**
+ * @hide
+ */
+final class OnDeviceIntelligenceShellCommand extends BasicShellCommandHandler {
     private static final String TAG = OnDeviceIntelligenceShellCommand.class.getSimpleName();
 
     @NonNull
diff --git a/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/RemoteOnDeviceIntelligenceService.java b/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/RemoteOnDeviceIntelligenceService.java
index ac9747a..0c43a30 100644
--- a/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/RemoteOnDeviceIntelligenceService.java
+++ b/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/RemoteOnDeviceIntelligenceService.java
@@ -16,6 +16,7 @@
 
 package com.android.server.ondeviceintelligence;
 
+import static android.app.ondeviceintelligence.OnDeviceIntelligenceManager.ON_DEVICE_INTELLIGENCE_UNBIND_TIMEOUT_MS;
 import static android.content.Context.BIND_FOREGROUND_SERVICE;
 import static android.content.Context.BIND_INCLUDE_CAPABILITIES;
 
@@ -26,13 +27,15 @@
 import android.service.ondeviceintelligence.IOnDeviceIntelligenceService;
 import android.service.ondeviceintelligence.OnDeviceIntelligenceService;
 
-import com.android.internal.infra.ServiceConnector;
+import com.android.modules.utils.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.
+ *
+ * @hide
  */
 public class RemoteOnDeviceIntelligenceService extends
         ServiceConnector.Impl<IOnDeviceIntelligenceService> {
@@ -56,11 +59,13 @@
         return LONG_TIMEOUT;
     }
 
+    // Using #getLong here as the timeout settings are only applicable to the services running in
+    // SYSTEM user only.
     @Override
+    @SuppressWarnings("NonUserGetterCalled")
     protected long getAutoDisconnectTimeoutMs() {
-        return Settings.Secure.getLongForUser(mContext.getContentResolver(),
-                Settings.Secure.ON_DEVICE_INTELLIGENCE_UNBIND_TIMEOUT_MS,
-                TimeUnit.SECONDS.toMillis(30),
-                mContext.getUserId());
+        return Settings.Secure.getLong(mContext.getContentResolver(),
+                ON_DEVICE_INTELLIGENCE_UNBIND_TIMEOUT_MS,
+                TimeUnit.SECONDS.toMillis(30));
     }
 }
diff --git a/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/RemoteOnDeviceSandboxedInferenceService.java b/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/RemoteOnDeviceSandboxedInferenceService.java
index 18b1383..8c5d5a7 100644
--- a/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/RemoteOnDeviceSandboxedInferenceService.java
+++ b/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/RemoteOnDeviceSandboxedInferenceService.java
@@ -16,6 +16,7 @@
 
 package com.android.server.ondeviceintelligence;
 
+import static android.app.ondeviceintelligence.OnDeviceIntelligenceManager.ON_DEVICE_INFERENCE_UNBIND_TIMEOUT_MS;
 import static android.content.Context.BIND_FOREGROUND_SERVICE;
 import static android.content.Context.BIND_INCLUDE_CAPABILITIES;
 
@@ -26,7 +27,7 @@
 import android.service.ondeviceintelligence.IOnDeviceSandboxedInferenceService;
 import android.service.ondeviceintelligence.OnDeviceSandboxedInferenceService;
 
-import com.android.internal.infra.ServiceConnector;
+import com.android.modules.utils.ServiceConnector;
 
 import java.util.concurrent.TimeUnit;
 
@@ -35,6 +36,8 @@
  * 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.
+ *
+ * @hide
  */
 public class RemoteOnDeviceSandboxedInferenceService extends
         ServiceConnector.Impl<IOnDeviceSandboxedInferenceService> {
@@ -65,12 +68,13 @@
         return LONG_TIMEOUT;
     }
 
-
+    // Using #getLong here as the timeout settings are only applicable to the services running in
+    // SYSTEM user only.
     @Override
+    @SuppressWarnings("NonUserGetterCalled")
     protected long getAutoDisconnectTimeoutMs() {
-        return Settings.Secure.getLongForUser(mContext.getContentResolver(),
-                Settings.Secure.ON_DEVICE_INFERENCE_UNBIND_TIMEOUT_MS,
-                TimeUnit.SECONDS.toMillis(30),
-                mContext.getUserId());
+        return Settings.Secure.getLong(mContext.getContentResolver(),
+                ON_DEVICE_INFERENCE_UNBIND_TIMEOUT_MS,
+                TimeUnit.SECONDS.toMillis(30));
     }
 }
diff --git a/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/callbacks/ListenableDownloadCallback.java b/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/callbacks/ListenableDownloadCallback.java
index 32f0698..249bcd3 100644
--- a/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/callbacks/ListenableDownloadCallback.java
+++ b/packages/NeuralNetworks/service/java/com/android/server/ondeviceintelligence/callbacks/ListenableDownloadCallback.java
@@ -21,7 +21,7 @@
 import android.os.PersistableBundle;
 import android.os.RemoteException;
 
-import com.android.internal.infra.AndroidFuture;
+import com.android.modules.utils.AndroidFuture;
 
 import java.util.concurrent.TimeoutException;
 
@@ -32,6 +32,8 @@
  * 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.
+ *
+ * @hide
  */
 public class ListenableDownloadCallback extends IDownloadCallback.Stub implements Runnable {
     private final IDownloadCallback callback;
diff --git a/proto/Android.bp b/proto/Android.bp
index a5e1335..feaa6d2 100644
--- a/proto/Android.bp
+++ b/proto/Android.bp
@@ -25,6 +25,10 @@
             static_libs: ["libprotobuf-java-nano"],
         },
     },
+    apex_available: [
+        "com.android.neuralnetworks",
+        "//apex_available:platform",
+    ],
 }
 
 java_library_static {
diff --git a/services/Android.bp b/services/Android.bp
index fc0bb33..a7cb9bb 100644
--- a/services/Android.bp
+++ b/services/Android.bp
@@ -210,6 +210,35 @@
     },
 }
 
+soong_config_module_type {
+    name: "ondeviceintelligence_module_java_defaults",
+    module_type: "java_defaults",
+    config_namespace: "ANDROID",
+    bool_variables: [
+        "release_ondevice_intelligence_module",
+        "release_ondevice_intelligence_platform",
+    ],
+    properties: [
+        "libs",
+        "srcs",
+        "static_libs",
+    ],
+}
+
+// Conditionally add ondeviceintelligence stubs library
+ondeviceintelligence_module_java_defaults {
+    name: "ondeviceintelligence_conditionally",
+    soong_config_variables: {
+        release_ondevice_intelligence_module: {
+            libs: ["service-ondeviceintelligence.stubs.system_server"],
+        },
+        release_ondevice_intelligence_platform: {
+            srcs: [":service-ondeviceintelligence-sources"],
+            static_libs: ["modules-utils-backgroundthread"],
+        },
+    },
+}
+
 // merge all required services into one jar
 // ============================================================
 soong_config_module_type {
@@ -236,6 +265,7 @@
         "services_java_defaults",
         "art_profile_java_defaults",
         "services_crashrecovery_stubs_conditionally",
+        "ondeviceintelligence_conditionally",
     ],
     installable: true,
 
diff --git a/services/core/Android.bp b/services/core/Android.bp
index ffa259b..371306f 100644
--- a/services/core/Android.bp
+++ b/services/core/Android.bp
@@ -127,6 +127,7 @@
         "android.hardware.power-java_shared",
         "latest_android_hardware_broadcastradio_java_static",
         "services_crashrecovery_stubs_conditionally",
+        "ondeviceintelligence_conditionally",
     ],
     srcs: [
         ":android.hardware.tv.hdmi.connection-V1-java-source",
diff --git a/services/core/java/com/android/server/pm/ComputerEngine.java b/services/core/java/com/android/server/pm/ComputerEngine.java
index 58b1e49..be2f58d 100644
--- a/services/core/java/com/android/server/pm/ComputerEngine.java
+++ b/services/core/java/com/android/server/pm/ComputerEngine.java
@@ -138,7 +138,8 @@
 import com.android.internal.util.IndentingPrintWriter;
 import com.android.internal.util.Preconditions;
 import com.android.modules.utils.TypedXmlSerializer;
-import com.android.server.ondeviceintelligence.OnDeviceIntelligenceManagerInternal;
+import com.android.server.LocalManagerRegistry;
+import com.android.server.ondeviceintelligence.OnDeviceIntelligenceManagerLocal;
 import com.android.server.pm.dex.DexManager;
 import com.android.server.pm.dex.PackageDexUsage;
 import com.android.server.pm.parsing.PackageInfoUtils;
@@ -5851,10 +5852,10 @@
         if (isHotword) {
             return true;
         }
-        OnDeviceIntelligenceManagerInternal onDeviceIntelligenceManagerInternal =
-                mInjector.getLocalService(OnDeviceIntelligenceManagerInternal.class);
-        return onDeviceIntelligenceManagerInternal != null
-                && uid == onDeviceIntelligenceManagerInternal.getInferenceServiceUid();
+        OnDeviceIntelligenceManagerLocal onDeviceIntelligenceManagerLocal =
+                LocalManagerRegistry.getManager(OnDeviceIntelligenceManagerLocal.class);
+        return onDeviceIntelligenceManagerLocal != null
+                && uid == onDeviceIntelligenceManagerLocal.getInferenceServiceUid();
     }
 
     @Nullable
diff --git a/services/java/com/android/server/SystemServer.java b/services/java/com/android/server/SystemServer.java
index aa63c4a..65315af 100644
--- a/services/java/com/android/server/SystemServer.java
+++ b/services/java/com/android/server/SystemServer.java
@@ -201,7 +201,6 @@
 import com.android.server.notification.NotificationManagerService;
 import com.android.server.oemlock.OemLockService;
 import com.android.server.om.OverlayManagerService;
-import com.android.server.ondeviceintelligence.OnDeviceIntelligenceManagerService;
 import com.android.server.os.BugreportManagerService;
 import com.android.server.os.DeviceIdentifiersPolicyService;
 import com.android.server.os.NativeTombstoneManagerService;
@@ -392,6 +391,8 @@
             "com.android.server.sdksandbox.SdkSandboxManagerService$Lifecycle";
     private static final String AD_SERVICES_MANAGER_SERVICE_CLASS =
             "com.android.server.adservices.AdServicesManagerService$Lifecycle";
+    private static final String ON_DEVICE_INTELLIGENCE_MANAGER_SERVICE_CLASS =
+            "com.android.server.ondeviceintelligence.OnDeviceIntelligenceManagerService";
     private static final String ON_DEVICE_PERSONALIZATION_SYSTEM_SERVICE_CLASS =
             "com.android.server.ondevicepersonalization."
                     + "OnDevicePersonalizationSystemService$Lifecycle";
@@ -3453,7 +3454,7 @@
 
     private void startOnDeviceIntelligenceService(TimingsTraceAndSlog t) {
         t.traceBegin("startOnDeviceIntelligenceManagerService");
-        mSystemServiceManager.startService(OnDeviceIntelligenceManagerService.class);
+        mSystemServiceManager.startService(ON_DEVICE_INTELLIGENCE_MANAGER_SERVICE_CLASS);
         t.traceEnd();
     }
 
diff --git a/services/proguard.flags b/services/proguard.flags
index 977bd19..0e1f68e 100644
--- a/services/proguard.flags
+++ b/services/proguard.flags
@@ -44,6 +44,9 @@
 -keep,allowoptimization,allowaccessmodification class com.android.server.input.NativeInputManagerService$NativeImpl { *; }
 -keep,allowoptimization,allowaccessmodification class com.android.server.ThreadPriorityBooster { *; }
 
+# allow invoking start-service using class name in both apex and services jar.
+-keep,allowoptimization,allowaccessmodification class com.android.server.ondeviceintelligence.OnDeviceIntelligenceManagerService { *; }
+
 # Keep all aconfig Flag class as they might be statically referenced by other packages
 # An merge or inlining could lead to missing dependencies that cause run time errors
 -keepclassmembernames class android.**.Flags, com.android.**.Flags { public *; }