Merge "Expose InferenceInfo via Hidden API for V to use in Settings App." into main

10 files changed, 433 insertions, 8 deletions

diff --git a/core/java/android/app/ondeviceintelligence/IOnDeviceIntelligenceManager.aidl b/core/java/android/app/ondeviceintelligence/IOnDeviceIntelligenceManager.aidl
index 470b1ec29105..1977a3923578 100644
--- a/core/java/android/app/ondeviceintelligence/IOnDeviceIntelligenceManager.aidl
+++ b/core/java/android/app/ondeviceintelligence/IOnDeviceIntelligenceManager.aidl
@@ -24,6 +24,8 @@
  import android.os.Bundle;
  import android.app.ondeviceintelligence.Feature;
  import android.app.ondeviceintelligence.FeatureDetails;
+ import android.app.ondeviceintelligence.InferenceInfo;
+ import java.util.List;
  import android.app.ondeviceintelligence.IDownloadCallback;
  import android.app.ondeviceintelligence.IListFeaturesCallback;
  import android.app.ondeviceintelligence.IFeatureCallback;
@@ -72,4 +74,6 @@ interface IOnDeviceIntelligenceManager {
                     in IStreamingResponseCallback streamingCallback) = 8;
 
       String getRemoteServicePackageName() = 9;
+
+      List<InferenceInfo> getLatestInferenceInfo(long startTimeEpochMillis) = 10;
  }
diff --git a/core/java/android/app/ondeviceintelligence/InferenceInfo.aidl b/core/java/android/app/ondeviceintelligence/InferenceInfo.aidl
new file mode 100644
index 000000000000..6d70fc4577a2
--- /dev/null
+++ b/core/java/android/app/ondeviceintelligence/InferenceInfo.aidl
@@ -0,0 +1,22 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.app.ondeviceintelligence;
+
+/**
+  * @hide
+  */
+parcelable InferenceInfo;
diff --git a/core/java/android/app/ondeviceintelligence/InferenceInfo.java b/core/java/android/app/ondeviceintelligence/InferenceInfo.java
new file mode 100644
index 000000000000..5557a81cbe40
--- /dev/null
+++ b/core/java/android/app/ondeviceintelligence/InferenceInfo.java
@@ -0,0 +1,211 @@
+/*
+ * 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 android.os.Parcel;
+import android.os.Parcelable;
+
+/**
+ * This class represents the information related to an inference event to track the resource usage
+ * as a function of inference time.
+ *
+ * @hide
+ */
+public class InferenceInfo implements Parcelable {
+
+    /**
+     * Uid for the caller app.
+     */
+    private final int uid;
+
+    /**
+     * Inference start time (milliseconds from the epoch time).
+     */
+    private final long startTimeMs;
+
+    /**
+     * Inference end time (milliseconds from the epoch time).
+     */
+    private final long endTimeMs;
+
+    /**
+     * Suspended time in milliseconds.
+     */
+    private final long suspendedTimeMs;
+
+    /**
+     * Constructs an InferenceInfo object with the specified parameters.
+     *
+     * @param uid             Uid for the caller app.
+     * @param startTimeMs     Inference start time (milliseconds from the epoch time).
+     * @param endTimeMs       Inference end time (milliseconds from the epoch time).
+     * @param suspendedTimeMs Suspended time in milliseconds.
+     */
+    public InferenceInfo(int uid, long startTimeMs, long endTimeMs,
+            long suspendedTimeMs) {
+        this.uid = uid;
+        this.startTimeMs = startTimeMs;
+        this.endTimeMs = endTimeMs;
+        this.suspendedTimeMs = suspendedTimeMs;
+    }
+
+    /**
+     * Constructs an InferenceInfo object from a Parcel.
+     *
+     * @param in The Parcel to read the object's data from.
+     */
+    protected InferenceInfo(Parcel in) {
+        uid = in.readInt();
+        startTimeMs = in.readLong();
+        endTimeMs = in.readLong();
+        suspendedTimeMs = in.readLong();
+    }
+
+
+    /**
+     * Writes the object's data to the provided Parcel.
+     *
+     * @param dest The Parcel to write the object's data to.
+     * @param flags Additional flags about how the object should be written.
+     */
+    @Override
+    public void writeToParcel(Parcel dest, int flags) {
+        dest.writeInt(uid);
+        dest.writeLong(startTimeMs);
+        dest.writeLong(endTimeMs);
+        dest.writeLong(suspendedTimeMs);
+    }
+
+    /**
+     * Returns the UID for the caller app.
+     *
+     * @return the UID for the caller app.
+     */
+    public int getUid() {
+        return uid;
+    }
+
+    /**
+     * Returns the inference start time in milliseconds from the epoch time.
+     *
+     * @return the inference start time in milliseconds from the epoch time.
+     */
+    public long getStartTimeMs() {
+        return startTimeMs;
+    }
+
+    /**
+     * Returns the inference end time in milliseconds from the epoch time.
+     *
+     * @return the inference end time in milliseconds from the epoch time.
+     */
+    public long getEndTimeMs() {
+        return endTimeMs;
+    }
+
+    /**
+     * Returns the suspended time in milliseconds.
+     *
+     * @return the suspended time in milliseconds.
+     */
+    public long getSuspendedTimeMs() {
+        return suspendedTimeMs;
+    }
+
+    @Override
+    public int describeContents() {
+        return 0;
+    }
+
+
+    public static final @android.annotation.NonNull Parcelable.Creator<InferenceInfo> CREATOR
+            = new Parcelable.Creator<InferenceInfo>() {
+        @Override
+        public InferenceInfo[] newArray(int size) {
+            return new InferenceInfo[size];
+        }
+
+        @Override
+        public InferenceInfo createFromParcel(@android.annotation.NonNull Parcel in) {
+            return new InferenceInfo(in);
+        }
+    };
+
+    /**
+     * Builder class for creating instances of {@link InferenceInfo}.
+     */
+    public static class Builder {
+        private int uid;
+        private long startTimeMs;
+        private long endTimeMs;
+        private long suspendedTimeMs;
+
+        /**
+         * Sets the UID for the caller app.
+         *
+         * @param uid the UID for the caller app.
+         * @return the Builder instance.
+         */
+        public Builder setUid(int uid) {
+            this.uid = uid;
+            return this;
+        }
+
+        /**
+         * Sets the inference start time in milliseconds from the epoch time.
+         *
+         * @param startTimeMs the inference start time in milliseconds from the epoch time.
+         * @return the Builder instance.
+         */
+        public Builder setStartTimeMs(long startTimeMs) {
+            this.startTimeMs = startTimeMs;
+            return this;
+        }
+
+        /**
+         * Sets the inference end time in milliseconds from the epoch time.
+         *
+         * @param endTimeMs the inference end time in milliseconds from the epoch time.
+         * @return the Builder instance.
+         */
+        public Builder setEndTimeMs(long endTimeMs) {
+            this.endTimeMs = endTimeMs;
+            return this;
+        }
+
+        /**
+         * Sets the suspended time in milliseconds.
+         *
+         * @param suspendedTimeMs the suspended time in milliseconds.
+         * @return the Builder instance.
+         */
+        public Builder setSuspendedTimeMs(long suspendedTimeMs) {
+            this.suspendedTimeMs = suspendedTimeMs;
+            return this;
+        }
+
+        /**
+         * Builds and returns an instance of {@link InferenceInfo}.
+         *
+         * @return an instance of {@link InferenceInfo}.
+         */
+        public InferenceInfo build() {
+            return new InferenceInfo(uid, startTimeMs, endTimeMs,
+                    suspendedTimeMs);
+        }
+    }
+}
diff --git a/core/java/android/app/ondeviceintelligence/OnDeviceIntelligenceManager.java b/core/java/android/app/ondeviceintelligence/OnDeviceIntelligenceManager.java
index a37f51bb6944..937a9cde1097 100644
--- a/core/java/android/app/ondeviceintelligence/OnDeviceIntelligenceManager.java
+++ b/core/java/android/app/ondeviceintelligence/OnDeviceIntelligenceManager.java
@@ -496,6 +496,24 @@ public final class OnDeviceIntelligenceManager {
         }
     }
 
+    /**
+     * This is primarily intended to be used to attribute/blame on-device intelligence power usage,
+     * via the configured remote implementation, to its actual caller.
+     *
+     * @param startTimeEpochMillis epoch millis used to filter the InferenceInfo events.
+     * @return InferenceInfo events since the passed in startTimeEpochMillis.
+     *
+     * @hide
+     */
+    @RequiresPermission(Manifest.permission.DUMP)
+    public List<InferenceInfo> getLatestInferenceInfo(long startTimeEpochMillis) {
+        try {
+            return mService.getLatestInferenceInfo(startTimeEpochMillis);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
 
     /** Request inference with provided Bundle and Params. */
     public static final int REQUEST_TYPE_INFERENCE = 0;
diff --git a/core/java/android/service/ondeviceintelligence/OnDeviceSandboxedInferenceService.java b/core/java/android/service/ondeviceintelligence/OnDeviceSandboxedInferenceService.java
index a77e07662d23..f123a962def9 100644
--- a/core/java/android/service/ondeviceintelligence/OnDeviceSandboxedInferenceService.java
+++ b/core/java/android/service/ondeviceintelligence/OnDeviceSandboxedInferenceService.java
@@ -101,6 +101,11 @@ public abstract class OnDeviceSandboxedInferenceService extends Service {
     private static final String TAG = OnDeviceSandboxedInferenceService.class.getSimpleName();
 
     /**
+     * @hide
+     */
+    public static final String INFERENCE_INFO_BUNDLE_KEY = "inference_info";
+
+    /**
      * The {@link Intent} that must be declared as handled by the service. To be supported, the
      * service must also require the
      * {@link android.Manifest.permission#BIND_ON_DEVICE_SANDBOXED_INFERENCE_SERVICE}
diff --git a/proto/src/ondeviceintelligence/inference_info.proto b/proto/src/ondeviceintelligence/inference_info.proto
new file mode 100644
index 000000000000..a6f4f4fa4eba
--- /dev/null
+++ b/proto/src/ondeviceintelligence/inference_info.proto
@@ -0,0 +1,34 @@
+/*
+ * 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.
+ */
+
+syntax = "proto2";
+
+package android.ondeviceintelligence;
+
+option java_package = "com.android.server.ondeviceintelligence";
+option java_multiple_files = true;
+
+
+message InferenceInfo {
+  // Uid for the caller app.
+  optional int32 uid = 1;
+  // Inference start time(milliseconds from the epoch time).
+  optional int64 start_time_ms = 2;
+  // Inference end time(milliseconds from the epoch time).
+  optional int64 end_time_ms = 3;
+  // Suspended time in milliseconds.
+  optional int64 suspended_time_ms = 4;
+}
\ No newline at end of file
diff --git a/services/core/java/com/android/server/ondeviceintelligence/BundleUtil.java b/services/core/java/com/android/server/ondeviceintelligence/BundleUtil.java
index 96ab2ccc8611..7dd8f2fdcecb 100644
--- a/services/core/java/com/android/server/ondeviceintelligence/BundleUtil.java
+++ b/services/core/java/com/android/server/ondeviceintelligence/BundleUtil.java
@@ -188,7 +188,8 @@ public class BundleUtil {
     public static IStreamingResponseCallback wrapWithValidation(
             IStreamingResponseCallback streamingResponseCallback,
             Executor resourceClosingExecutor,
-            AndroidFuture future) {
+            AndroidFuture future,
+            InferenceInfoStore inferenceInfoStore) {
         return new IStreamingResponseCallback.Stub() {
             @Override
             public void onNewContent(Bundle processedResult) throws RemoteException {
@@ -207,6 +208,7 @@ public class BundleUtil {
                     sanitizeResponseParams(resultBundle);
                     streamingResponseCallback.onSuccess(resultBundle);
                 } finally {
+                    inferenceInfoStore.addInferenceInfoFromBundle(resultBundle);
                     resourceClosingExecutor.execute(() -> tryCloseResource(resultBundle));
                     future.complete(null);
                 }
@@ -216,6 +218,7 @@ public class BundleUtil {
             public void onFailure(int errorCode, String errorMessage,
                     PersistableBundle errorParams) throws RemoteException {
                 streamingResponseCallback.onFailure(errorCode, errorMessage, errorParams);
+                inferenceInfoStore.addInferenceInfoFromBundle(errorParams);
                 future.completeExceptionally(new TimeoutException());
             }
 
@@ -245,7 +248,8 @@ public class BundleUtil {
 
     public static IResponseCallback wrapWithValidation(IResponseCallback responseCallback,
             Executor resourceClosingExecutor,
-            AndroidFuture future) {
+            AndroidFuture future,
+            InferenceInfoStore inferenceInfoStore) {
         return new IResponseCallback.Stub() {
             @Override
             public void onSuccess(Bundle resultBundle)
@@ -254,6 +258,7 @@ public class BundleUtil {
                     sanitizeResponseParams(resultBundle);
                     responseCallback.onSuccess(resultBundle);
                 } finally {
+                    inferenceInfoStore.addInferenceInfoFromBundle(resultBundle);
                     resourceClosingExecutor.execute(() -> tryCloseResource(resultBundle));
                     future.complete(null);
                 }
@@ -263,6 +268,7 @@ public class BundleUtil {
             public void onFailure(int errorCode, String errorMessage,
                     PersistableBundle errorParams) throws RemoteException {
                 responseCallback.onFailure(errorCode, errorMessage, errorParams);
+                inferenceInfoStore.addInferenceInfoFromBundle(errorParams);
                 future.completeExceptionally(new TimeoutException());
             }
 
@@ -291,11 +297,13 @@ public class BundleUtil {
 
 
     public static ITokenInfoCallback wrapWithValidation(ITokenInfoCallback responseCallback,
-            AndroidFuture future) {
+            AndroidFuture future,
+            InferenceInfoStore inferenceInfoStore) {
         return new ITokenInfoCallback.Stub() {
             @Override
             public void onSuccess(TokenInfo tokenInfo) throws RemoteException {
                 responseCallback.onSuccess(tokenInfo);
+                inferenceInfoStore.addInferenceInfoFromBundle(tokenInfo.getInfoParams());
                 future.complete(null);
             }
 
@@ -303,6 +311,7 @@ public class BundleUtil {
             public void onFailure(int errorCode, String errorMessage, PersistableBundle errorParams)
                     throws RemoteException {
                 responseCallback.onFailure(errorCode, errorMessage, errorParams);
+                inferenceInfoStore.addInferenceInfoFromBundle(errorParams);
                 future.completeExceptionally(new TimeoutException());
             }
         };
diff --git a/services/core/java/com/android/server/ondeviceintelligence/InferenceInfoStore.java b/services/core/java/com/android/server/ondeviceintelligence/InferenceInfoStore.java
new file mode 100644
index 000000000000..6578853909c9
--- /dev/null
+++ b/services/core/java/com/android/server/ondeviceintelligence/InferenceInfoStore.java
@@ -0,0 +1,99 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.ondeviceintelligence;
+
+import android.app.ondeviceintelligence.InferenceInfo;
+import android.os.Bundle;
+import android.os.PersistableBundle;
+import android.service.ondeviceintelligence.OnDeviceSandboxedInferenceService;
+import android.util.Slog;
+
+import java.io.IOException;
+import java.util.Base64;
+import java.util.Comparator;
+import java.util.List;
+import java.util.TreeSet;
+
+public class InferenceInfoStore {
+    private static final String TAG = "InferenceInfoStore";
+    private final TreeSet<InferenceInfo> inferenceInfos;
+    private final long maxAgeMs;
+
+    public InferenceInfoStore(long maxAgeMs) {
+        this.maxAgeMs = maxAgeMs;
+        this.inferenceInfos = new TreeSet<>(
+                Comparator.comparingLong(InferenceInfo::getStartTimeMs));
+    }
+
+    public List<InferenceInfo> getLatestInferenceInfo(long startTimeEpochMillis) {
+        return inferenceInfos.stream().filter(
+                info -> info.getStartTimeMs() > startTimeEpochMillis).toList();
+    }
+
+    public void addInferenceInfoFromBundle(PersistableBundle pb) {
+        if (!pb.containsKey(OnDeviceSandboxedInferenceService.INFERENCE_INFO_BUNDLE_KEY)) {
+            return;
+        }
+
+        try {
+            String infoBytesBase64String = pb.getString(
+                    OnDeviceSandboxedInferenceService.INFERENCE_INFO_BUNDLE_KEY);
+            if (infoBytesBase64String != null) {
+                byte[] infoBytes = Base64.getDecoder().decode(infoBytesBase64String);
+                com.android.server.ondeviceintelligence.nano.InferenceInfo inferenceInfo =
+                        com.android.server.ondeviceintelligence.nano.InferenceInfo.parseFrom(
+                                infoBytes);
+                add(inferenceInfo);
+            }
+        } catch (IOException e) {
+            Slog.e(TAG, "Unable to parse InferenceInfo from the received bytes.");
+        }
+    }
+
+    public void addInferenceInfoFromBundle(Bundle b) {
+        if (!b.containsKey(OnDeviceSandboxedInferenceService.INFERENCE_INFO_BUNDLE_KEY)) {
+            return;
+        }
+
+        try {
+            byte[] infoBytes = b.getByteArray(
+                    OnDeviceSandboxedInferenceService.INFERENCE_INFO_BUNDLE_KEY);
+            if (infoBytes != null) {
+                com.android.server.ondeviceintelligence.nano.InferenceInfo inferenceInfo =
+                        com.android.server.ondeviceintelligence.nano.InferenceInfo.parseFrom(
+                                infoBytes);
+                add(inferenceInfo);
+            }
+        } catch (IOException e) {
+            Slog.e(TAG, "Unable to parse InferenceInfo from the received bytes.");
+        }
+    }
+
+    private synchronized void add(com.android.server.ondeviceintelligence.nano.InferenceInfo info) {
+        while (System.currentTimeMillis() - inferenceInfos.first().getStartTimeMs() > maxAgeMs) {
+            inferenceInfos.pollFirst();
+        }
+        inferenceInfos.add(toInferenceInfo(info));
+    }
+
+    private static InferenceInfo toInferenceInfo(
+            com.android.server.ondeviceintelligence.nano.InferenceInfo info) {
+        return new InferenceInfo.Builder().setUid(info.uid).setStartTimeMs(
+                info.startTimeMs).setEndTimeMs(info.endTimeMs).setSuspendedTimeMs(
+                info.suspendedTimeMs).build();
+    }
+}
\ No newline at end of file
diff --git a/services/core/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerInternal.java b/services/core/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerInternal.java
index 07af8d042420..1450dc0803d6 100644
--- a/services/core/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerInternal.java
+++ b/services/core/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerInternal.java
@@ -17,5 +17,10 @@
 package com.android.server.ondeviceintelligence;
 
 public interface OnDeviceIntelligenceManagerInternal {
+    /**
+     * Gets the uid for the process that is currently hosting the
+     * {@link android.service.ondeviceintelligence.OnDeviceSandboxedInferenceService} registered on
+     * the device.
+     */
     int getInferenceServiceUid();
 }
\ No newline at end of file
diff --git a/services/core/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerService.java b/services/core/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerService.java
index 1a43bc27bdfc..9ef2e12e55c5 100644
--- a/services/core/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerService.java
+++ b/services/core/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerService.java
@@ -44,6 +44,7 @@ import android.app.ondeviceintelligence.IProcessingSignal;
 import android.app.ondeviceintelligence.IResponseCallback;
 import android.app.ondeviceintelligence.IStreamingResponseCallback;
 import android.app.ondeviceintelligence.ITokenInfoCallback;
+import android.app.ondeviceintelligence.InferenceInfo;
 import android.app.ondeviceintelligence.OnDeviceIntelligenceException;
 import android.content.ComponentName;
 import android.content.Context;
@@ -127,6 +128,7 @@ public class OnDeviceIntelligenceManagerService extends SystemService {
     private static final String NAMESPACE_ON_DEVICE_INTELLIGENCE = "ondeviceintelligence";
 
     private static final String SYSTEM_PACKAGE = "android";
+    private static final long MAX_AGE_MS = TimeUnit.HOURS.toMillis(3);
 
 
     private final Executor resourceClosingExecutor = Executors.newCachedThreadPool();
@@ -138,7 +140,7 @@ public class OnDeviceIntelligenceManagerService extends SystemService {
     private final Context mContext;
     protected final Object mLock = new Object();
 
-
+    private final InferenceInfoStore mInferenceInfoStore;
     private RemoteOnDeviceSandboxedInferenceService mRemoteInferenceService;
     private RemoteOnDeviceIntelligenceService mRemoteOnDeviceIntelligenceService;
     volatile boolean mIsServiceEnabled;
@@ -170,6 +172,7 @@ public class OnDeviceIntelligenceManagerService extends SystemService {
         super(context);
         mContext = context;
         mTemporaryServiceNames = new String[0];
+        mInferenceInfoStore = new InferenceInfoStore(MAX_AGE_MS);
     }
 
     @Override
@@ -223,6 +226,14 @@ public class OnDeviceIntelligenceManagerService extends SystemService {
             }
 
             @Override
+            public List<InferenceInfo> getLatestInferenceInfo(long startTimeEpochMillis) {
+                mContext.enforceCallingPermission(
+                        Manifest.permission.DUMP, TAG);
+                return OnDeviceIntelligenceManagerService.this.getLatestInferenceInfo(
+                        startTimeEpochMillis);
+            }
+
+            @Override
             public void getVersion(RemoteCallback remoteCallback) {
                 Slog.i(TAG, "OnDeviceIntelligenceManagerInternal getVersion");
                 Objects.requireNonNull(remoteCallback);
@@ -434,7 +445,8 @@ public class OnDeviceIntelligenceManagerService extends SystemService {
                                 service.requestTokenInfo(callerUid, feature,
                                         request,
                                         wrapCancellationFuture(cancellationSignalFuture),
-                                        wrapWithValidation(tokenInfoCallback, future));
+                                        wrapWithValidation(tokenInfoCallback, future,
+                                                mInferenceInfoStore));
                                 return future.orTimeout(getIdleTimeoutMs(), TimeUnit.MILLISECONDS);
                             });
                     result.whenCompleteAsync((c, e) -> BundleUtil.tryCloseResource(request),
@@ -480,7 +492,8 @@ public class OnDeviceIntelligenceManagerService extends SystemService {
                                         wrapCancellationFuture(cancellationSignalFuture),
                                         wrapProcessingFuture(processingSignalFuture),
                                         wrapWithValidation(responseCallback,
-                                                resourceClosingExecutor, future));
+                                                resourceClosingExecutor, future,
+                                                mInferenceInfoStore));
                                 return future.orTimeout(getIdleTimeoutMs(), TimeUnit.MILLISECONDS);
                             });
                     result.whenCompleteAsync((c, e) -> BundleUtil.tryCloseResource(request),
@@ -525,7 +538,8 @@ public class OnDeviceIntelligenceManagerService extends SystemService {
                                         wrapCancellationFuture(cancellationSignalFuture),
                                         wrapProcessingFuture(processingSignalFuture),
                                         wrapWithValidation(streamingCallback,
-                                                resourceClosingExecutor, future));
+                                                resourceClosingExecutor, future,
+                                                mInferenceInfoStore));
                                 return future.orTimeout(getIdleTimeoutMs(), TimeUnit.MILLISECONDS);
                             });
                     result.whenCompleteAsync((c, e) -> BundleUtil.tryCloseResource(request),
@@ -846,6 +860,10 @@ public class OnDeviceIntelligenceManagerService extends SystemService {
                 && (serviceInfo.flags & ServiceInfo.FLAG_EXTERNAL_SERVICE) == 0;
     }
 
+    private List<InferenceInfo> getLatestInferenceInfo(long startTimeEpochMillis) {
+        return mInferenceInfoStore.getLatestInferenceInfo(startTimeEpochMillis);
+    }
+
     @Nullable
     public String getRemoteConfiguredPackageName() {
         try {
@@ -1066,7 +1084,7 @@ public class OnDeviceIntelligenceManagerService extends SystemService {
     }
 
     private void setRemoteInferenceServiceUid(int remoteInferenceServiceUid) {
-        synchronized (mLock){
+        synchronized (mLock) {
             this.remoteInferenceServiceUid = remoteInferenceServiceUid;
         }
     }