Refactoring and some fixes around TextClassifierService

This change is mainly to address the comment
"TextClassifierService.onDetectLanguage receives a callback but
immediately calls callback.onSuccess before method returns, that is
unexpected behavior. Method should either have a return value or post
callback. Same for onSuggestConversationActions" from API  council.

1. Like other services that allows developers to extend,
   APIs in TextClassifierService are now delivered in the main thread.

2. The default implementation of onDetectLanguage and
   onSuggestConversationActions now run local textclassifier in a worker
   thread. So it won't block the main thread and make sure the callback
   is called after the function returns.

3. Refactor the code so that fewer boilerplate codes around the callback.
   We are now using ITextClassifierCallback to replace all the existing
   I*Callback. This allows us to write template code easily.

BUG: 120841922

Test: Enable system textclassifier tests and run atest TextClassifierTest.java
Test: Manual. Try out smart selection and related features.

Change-Id: I33b2cc80f968fb8b5d5007a153b6e50041e57ce6

11 files changed, 157 insertions, 359 deletions

diff --git a/Android.bp b/Android.bp
index 628894016f39..0f53600172d0 100644
--- a/Android.bp
+++ b/Android.bp
@@ -352,12 +352,8 @@ java_defaults {
         "core/java/android/service/chooser/IChooserTargetResult.aidl",
         "core/java/android/service/resolver/IResolverRankerService.aidl",
         "core/java/android/service/resolver/IResolverRankerResult.aidl",
-        "core/java/android/service/textclassifier/IConversationActionsCallback.aidl",
-        "core/java/android/service/textclassifier/ITextClassificationCallback.aidl",
+        "core/java/android/service/textclassifier/ITextClassifierCallback.aidl",
         "core/java/android/service/textclassifier/ITextClassifierService.aidl",
-        "core/java/android/service/textclassifier/ITextLanguageCallback.aidl",
-        "core/java/android/service/textclassifier/ITextLinksCallback.aidl",
-        "core/java/android/service/textclassifier/ITextSelectionCallback.aidl",
         "core/java/android/service/attention/IAttentionService.aidl",
         "core/java/android/service/attention/IAttentionCallback.aidl",
         "core/java/android/view/accessibility/IAccessibilityInteractionConnection.aidl",
diff --git a/api/system-current.txt b/api/system-current.txt
index e7879277234b..2699d4a8fae0 100644
--- a/api/system-current.txt
+++ b/api/system-current.txt
@@ -6724,15 +6724,15 @@ package android.service.textclassifier {
     method public static android.view.textclassifier.TextClassifier getDefaultTextClassifierImplementation(@NonNull android.content.Context);
     method @Deprecated public final android.view.textclassifier.TextClassifier getLocalTextClassifier();
     method @Nullable public final android.os.IBinder onBind(android.content.Intent);
-    method public abstract void onClassifyText(@Nullable android.view.textclassifier.TextClassificationSessionId, @NonNull android.view.textclassifier.TextClassification.Request, @NonNull android.os.CancellationSignal, @NonNull android.service.textclassifier.TextClassifierService.Callback<android.view.textclassifier.TextClassification>);
-    method public void onCreateTextClassificationSession(@NonNull android.view.textclassifier.TextClassificationContext, @NonNull android.view.textclassifier.TextClassificationSessionId);
-    method public void onDestroyTextClassificationSession(@NonNull android.view.textclassifier.TextClassificationSessionId);
-    method public void onDetectLanguage(@Nullable android.view.textclassifier.TextClassificationSessionId, @NonNull android.view.textclassifier.TextLanguage.Request, @NonNull android.os.CancellationSignal, @NonNull android.service.textclassifier.TextClassifierService.Callback<android.view.textclassifier.TextLanguage>);
-    method public abstract void onGenerateLinks(@Nullable android.view.textclassifier.TextClassificationSessionId, @NonNull android.view.textclassifier.TextLinks.Request, @NonNull android.os.CancellationSignal, @NonNull android.service.textclassifier.TextClassifierService.Callback<android.view.textclassifier.TextLinks>);
-    method @Deprecated public void onSelectionEvent(@Nullable android.view.textclassifier.TextClassificationSessionId, @NonNull android.view.textclassifier.SelectionEvent);
-    method public void onSuggestConversationActions(@Nullable android.view.textclassifier.TextClassificationSessionId, @NonNull android.view.textclassifier.ConversationActions.Request, @NonNull android.os.CancellationSignal, @NonNull android.service.textclassifier.TextClassifierService.Callback<android.view.textclassifier.ConversationActions>);
-    method public abstract void onSuggestSelection(@Nullable android.view.textclassifier.TextClassificationSessionId, @NonNull android.view.textclassifier.TextSelection.Request, @NonNull android.os.CancellationSignal, @NonNull android.service.textclassifier.TextClassifierService.Callback<android.view.textclassifier.TextSelection>);
-    method public void onTextClassifierEvent(@Nullable android.view.textclassifier.TextClassificationSessionId, @NonNull android.view.textclassifier.TextClassifierEvent);
+    method @MainThread public abstract void onClassifyText(@Nullable android.view.textclassifier.TextClassificationSessionId, @NonNull android.view.textclassifier.TextClassification.Request, @NonNull android.os.CancellationSignal, @NonNull android.service.textclassifier.TextClassifierService.Callback<android.view.textclassifier.TextClassification>);
+    method @MainThread public void onCreateTextClassificationSession(@NonNull android.view.textclassifier.TextClassificationContext, @NonNull android.view.textclassifier.TextClassificationSessionId);
+    method @MainThread public void onDestroyTextClassificationSession(@NonNull android.view.textclassifier.TextClassificationSessionId);
+    method @MainThread public void onDetectLanguage(@Nullable android.view.textclassifier.TextClassificationSessionId, @NonNull android.view.textclassifier.TextLanguage.Request, @NonNull android.os.CancellationSignal, @NonNull android.service.textclassifier.TextClassifierService.Callback<android.view.textclassifier.TextLanguage>);
+    method @MainThread public abstract void onGenerateLinks(@Nullable android.view.textclassifier.TextClassificationSessionId, @NonNull android.view.textclassifier.TextLinks.Request, @NonNull android.os.CancellationSignal, @NonNull android.service.textclassifier.TextClassifierService.Callback<android.view.textclassifier.TextLinks>);
+    method @Deprecated @MainThread public void onSelectionEvent(@Nullable android.view.textclassifier.TextClassificationSessionId, @NonNull android.view.textclassifier.SelectionEvent);
+    method @MainThread public void onSuggestConversationActions(@Nullable android.view.textclassifier.TextClassificationSessionId, @NonNull android.view.textclassifier.ConversationActions.Request, @NonNull android.os.CancellationSignal, @NonNull android.service.textclassifier.TextClassifierService.Callback<android.view.textclassifier.ConversationActions>);
+    method @MainThread public abstract void onSuggestSelection(@Nullable android.view.textclassifier.TextClassificationSessionId, @NonNull android.view.textclassifier.TextSelection.Request, @NonNull android.os.CancellationSignal, @NonNull android.service.textclassifier.TextClassifierService.Callback<android.view.textclassifier.TextSelection>);
+    method @MainThread public void onTextClassifierEvent(@Nullable android.view.textclassifier.TextClassificationSessionId, @NonNull android.view.textclassifier.TextClassifierEvent);
     field public static final String SERVICE_INTERFACE = "android.service.textclassifier.TextClassifierService";
   }
 
diff --git a/core/java/android/service/textclassifier/IConversationActionsCallback.aidl b/core/java/android/service/textclassifier/IConversationActionsCallback.aidl
deleted file mode 100644
index c35d4246e924..000000000000
--- a/core/java/android/service/textclassifier/IConversationActionsCallback.aidl
+++ /dev/null
@@ -1,28 +0,0 @@
-/*
- * Copyright (C) 2018 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package android.service.textclassifier;
-
-import android.view.textclassifier.ConversationActions;
-
-/**
- * Callback for a ConversationActions request.
- * @hide
- */
-oneway interface IConversationActionsCallback {
-    void onSuccess(in ConversationActions conversationActions);
-    void onFailure();
-}
\ No newline at end of file
diff --git a/core/java/android/service/textclassifier/ITextClassificationCallback.aidl b/core/java/android/service/textclassifier/ITextClassificationCallback.aidl
deleted file mode 100644
index 10bfe6324509..000000000000
--- a/core/java/android/service/textclassifier/ITextClassificationCallback.aidl
+++ /dev/null
@@ -1,28 +0,0 @@
-/*
- * Copyright (C) 2018 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package android.service.textclassifier;
-
-import android.view.textclassifier.TextClassification;
-
-/**
- * Callback for a TextClassification request.
- * @hide
- */
-oneway interface ITextClassificationCallback {
-    void onSuccess(in TextClassification classification);
-    void onFailure();
-}
diff --git a/core/java/android/service/textclassifier/ITextSelectionCallback.aidl b/core/java/android/service/textclassifier/ITextClassifierCallback.aidl
index 1b4c4d10d427..2926734086a5 100644
--- a/core/java/android/service/textclassifier/ITextSelectionCallback.aidl
+++ b/core/java/android/service/textclassifier/ITextClassifierCallback.aidl
@@ -16,13 +16,14 @@
 
 package android.service.textclassifier;
 
+import android.os.Bundle;
 import android.view.textclassifier.TextSelection;
 
 /**
- * Callback for a TextSelection request.
+ * Callback for all requests from SystemTextClassifier.
  * @hide
  */
-oneway interface ITextSelectionCallback {
-    void onSuccess(in TextSelection selection);
+oneway interface ITextClassifierCallback {
+    void onSuccess(in Bundle result);
     void onFailure();
 }
\ No newline at end of file
diff --git a/core/java/android/service/textclassifier/ITextClassifierService.aidl b/core/java/android/service/textclassifier/ITextClassifierService.aidl
index 794179415fc0..2f8d67b6ccee 100644
--- a/core/java/android/service/textclassifier/ITextClassifierService.aidl
+++ b/core/java/android/service/textclassifier/ITextClassifierService.aidl
@@ -16,11 +16,7 @@
 
 package android.service.textclassifier;
 
-import android.service.textclassifier.IConversationActionsCallback;
-import android.service.textclassifier.ITextClassificationCallback;
-import android.service.textclassifier.ITextLanguageCallback;
-import android.service.textclassifier.ITextLinksCallback;
-import android.service.textclassifier.ITextSelectionCallback;
+import android.service.textclassifier.ITextClassifierCallback;
 import android.view.textclassifier.ConversationActions;
 import android.view.textclassifier.SelectionEvent;
 import android.view.textclassifier.TextClassification;
@@ -41,17 +37,17 @@ oneway interface ITextClassifierService {
     void onSuggestSelection(
             in TextClassificationSessionId sessionId,
             in TextSelection.Request request,
-            in ITextSelectionCallback callback);
+            in ITextClassifierCallback callback);
 
     void onClassifyText(
             in TextClassificationSessionId sessionId,
             in TextClassification.Request request,
-            in ITextClassificationCallback callback);
+            in ITextClassifierCallback callback);
 
     void onGenerateLinks(
             in TextClassificationSessionId sessionId,
             in TextLinks.Request request,
-            in ITextLinksCallback callback);
+            in ITextClassifierCallback callback);
 
     // TODO: Remove
     void onSelectionEvent(
@@ -72,10 +68,10 @@ oneway interface ITextClassifierService {
     void onDetectLanguage(
             in TextClassificationSessionId sessionId,
             in TextLanguage.Request request,
-            in ITextLanguageCallback callback);
+            in ITextClassifierCallback callback);
 
     void onSuggestConversationActions(
             in TextClassificationSessionId sessionId,
             in ConversationActions.Request request,
-            in IConversationActionsCallback callback);
+            in ITextClassifierCallback callback);
 }
diff --git a/core/java/android/service/textclassifier/ITextLanguageCallback.aidl b/core/java/android/service/textclassifier/ITextLanguageCallback.aidl
deleted file mode 100644
index 263d99afca5c..000000000000
--- a/core/java/android/service/textclassifier/ITextLanguageCallback.aidl
+++ /dev/null
@@ -1,28 +0,0 @@
-/*
- * Copyright (C) 2018 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package android.service.textclassifier;
-
-import android.view.textclassifier.TextLanguage;
-
-/**
- * Callback for a TextLanguage request.
- * @hide
- */
-oneway interface ITextLanguageCallback {
-    void onSuccess(in TextLanguage textLanguage);
-    void onFailure();
-}
\ No newline at end of file
diff --git a/core/java/android/service/textclassifier/ITextLinksCallback.aidl b/core/java/android/service/textclassifier/ITextLinksCallback.aidl
deleted file mode 100644
index a9e0dde56773..000000000000
--- a/core/java/android/service/textclassifier/ITextLinksCallback.aidl
+++ /dev/null
@@ -1,28 +0,0 @@
-/*
- * Copyright (C) 2018 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package android.service.textclassifier;
-
-import android.view.textclassifier.TextLinks;
-
-/**
- * Callback for a TextLinks request.
- * @hide
- */
-oneway interface ITextLinksCallback {
-    void onSuccess(in TextLinks links);
-    void onFailure();
-}
\ No newline at end of file
diff --git a/core/java/android/service/textclassifier/TextClassifierService.java b/core/java/android/service/textclassifier/TextClassifierService.java
index 235b8e8c9176..4088ce8ea4b0 100644
--- a/core/java/android/service/textclassifier/TextClassifierService.java
+++ b/core/java/android/service/textclassifier/TextClassifierService.java
@@ -17,6 +17,7 @@
 package android.service.textclassifier;
 
 import android.Manifest;
+import android.annotation.MainThread;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.SystemApi;
@@ -27,8 +28,12 @@ import android.content.Intent;
 import android.content.pm.PackageManager;
 import android.content.pm.ResolveInfo;
 import android.content.pm.ServiceInfo;
+import android.os.Bundle;
 import android.os.CancellationSignal;
+import android.os.Handler;
 import android.os.IBinder;
+import android.os.Looper;
+import android.os.Parcelable;
 import android.os.RemoteException;
 import android.text.TextUtils;
 import android.util.Slog;
@@ -46,6 +51,10 @@ import android.view.textclassifier.TextSelection;
 
 import com.android.internal.util.Preconditions;
 
+import java.lang.ref.WeakReference;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+
 /**
  * Abstract base class for the TextClassifier service.
  *
@@ -68,6 +77,11 @@ import com.android.internal.util.Preconditions;
  *     </intent-filter>
  * </service>}</pre>
  *
+ * <p>From {@link android.os.Build.VERSION_CODES#Q} onward, all callbacks are called on the main
+ * thread. Prior to Q, there is no guarantee on what thread the callback will happen. You should
+ * make sure the callbacks are executed in your desired thread by using a executor, a handler or
+ * something else along the line.
+ *
  * @see TextClassifier
  * @hide
  */
@@ -85,201 +99,102 @@ public abstract class TextClassifierService extends Service {
     public static final String SERVICE_INTERFACE =
             "android.service.textclassifier.TextClassifierService";
 
+    /** @hide **/
+    private static final String KEY_RESULT = "key_result";
+
+    private final Handler mMainThreadHandler = new Handler(Looper.getMainLooper(), null, true);
+    private final ExecutorService mSingleThreadExecutor = Executors.newSingleThreadExecutor();
+
     private final ITextClassifierService.Stub mBinder = new ITextClassifierService.Stub() {
 
         // TODO(b/72533911): Implement cancellation signal
         @NonNull private final CancellationSignal mCancellationSignal = new CancellationSignal();
 
-        /** {@inheritDoc} */
         @Override
         public void onSuggestSelection(
                 TextClassificationSessionId sessionId,
-                TextSelection.Request request, ITextSelectionCallback callback) {
+                TextSelection.Request request, ITextClassifierCallback callback) {
             Preconditions.checkNotNull(request);
             Preconditions.checkNotNull(callback);
-            TextClassifierService.this.onSuggestSelection(
-                    sessionId, request, mCancellationSignal,
-                    new Callback<TextSelection>() {
-                        @Override
-                        public void onSuccess(TextSelection result) {
-                            try {
-                                callback.onSuccess(result);
-                            } catch (RemoteException e) {
-                                Slog.d(LOG_TAG, "Error calling callback");
-                            }
-                        }
-
-                        @Override
-                        public void onFailure(CharSequence error) {
-                            try {
-                                if (callback.asBinder().isBinderAlive()) {
-                                    callback.onFailure();
-                                }
-                            } catch (RemoteException e) {
-                                Slog.d(LOG_TAG, "Error calling callback");
-                            }
-                        }
-                    });
+            mMainThreadHandler.post(() -> TextClassifierService.this.onSuggestSelection(
+                    sessionId, request, mCancellationSignal, new ProxyCallback<>(callback)));
+
         }
 
-        /** {@inheritDoc} */
         @Override
         public void onClassifyText(
                 TextClassificationSessionId sessionId,
-                TextClassification.Request request, ITextClassificationCallback callback) {
+                TextClassification.Request request, ITextClassifierCallback callback) {
             Preconditions.checkNotNull(request);
             Preconditions.checkNotNull(callback);
-            TextClassifierService.this.onClassifyText(
-                    sessionId, request, mCancellationSignal,
-                    new Callback<TextClassification>() {
-                        @Override
-                        public void onSuccess(TextClassification result) {
-                            try {
-                                callback.onSuccess(result);
-                            } catch (RemoteException e) {
-                                Slog.d(LOG_TAG, "Error calling callback");
-                            }
-                        }
-
-                        @Override
-                        public void onFailure(CharSequence error) {
-                            try {
-                                callback.onFailure();
-                            } catch (RemoteException e) {
-                                Slog.d(LOG_TAG, "Error calling callback");
-                            }
-                        }
-                    });
+            mMainThreadHandler.post(() -> TextClassifierService.this.onClassifyText(
+                    sessionId, request, mCancellationSignal, new ProxyCallback<>(callback)));
         }
 
-        /** {@inheritDoc} */
         @Override
         public void onGenerateLinks(
                 TextClassificationSessionId sessionId,
-                TextLinks.Request request, ITextLinksCallback callback) {
+                TextLinks.Request request, ITextClassifierCallback callback) {
             Preconditions.checkNotNull(request);
             Preconditions.checkNotNull(callback);
-            TextClassifierService.this.onGenerateLinks(
-                    sessionId, request,
-                    mCancellationSignal,
-                    new Callback<TextLinks>() {
-                        @Override
-                        public void onSuccess(TextLinks result) {
-                            try {
-                                callback.onSuccess(result);
-                            } catch (RemoteException e) {
-                                Slog.d(LOG_TAG, "Error calling callback");
-                            }
-                        }
-
-                        @Override
-                        public void onFailure(CharSequence error) {
-                            try {
-                                callback.onFailure();
-                            } catch (RemoteException e) {
-                                Slog.d(LOG_TAG, "Error calling callback");
-                            }
-                        }
-                    });
+            mMainThreadHandler.post(() -> TextClassifierService.this.onGenerateLinks(
+                    sessionId, request, mCancellationSignal, new ProxyCallback<>(callback)));
         }
 
-        /** {@inheritDoc} */
         @Override
         public void onSelectionEvent(
                 TextClassificationSessionId sessionId,
                 SelectionEvent event) {
             Preconditions.checkNotNull(event);
-            TextClassifierService.this.onSelectionEvent(sessionId, event);
+            mMainThreadHandler.post(
+                    () -> TextClassifierService.this.onSelectionEvent(sessionId, event));
         }
 
-        /** {@inheritDoc} */
         @Override
         public void onTextClassifierEvent(
                 TextClassificationSessionId sessionId,
                 TextClassifierEvent event) {
             Preconditions.checkNotNull(event);
-            TextClassifierService.this.onTextClassifierEvent(sessionId, event);
+            mMainThreadHandler.post(
+                    () -> TextClassifierService.this.onTextClassifierEvent(sessionId, event));
         }
 
-        /** {@inheritDoc} */
         @Override
         public void onDetectLanguage(
                 TextClassificationSessionId sessionId,
                 TextLanguage.Request request,
-                ITextLanguageCallback callback) {
+                ITextClassifierCallback callback) {
             Preconditions.checkNotNull(request);
             Preconditions.checkNotNull(callback);
-            TextClassifierService.this.onDetectLanguage(
-                    sessionId,
-                    request,
-                    mCancellationSignal,
-                    new Callback<TextLanguage>() {
-                        @Override
-                        public void onSuccess(TextLanguage result) {
-                            try {
-                                callback.onSuccess(result);
-                            } catch (RemoteException e) {
-                                Slog.d(LOG_TAG, "Error calling callback");
-                            }
-                        }
-
-                        @Override
-                        public void onFailure(CharSequence error) {
-                            try {
-                                callback.onFailure();
-                            } catch (RemoteException e) {
-                                Slog.d(LOG_TAG, "Error calling callback");
-                            }
-                        };
-                    });
+            mMainThreadHandler.post(() -> TextClassifierService.this.onDetectLanguage(
+                    sessionId, request, mCancellationSignal, new ProxyCallback<>(callback)));
         }
 
-        /** {@inheritDoc} */
         @Override
         public void onSuggestConversationActions(
                 TextClassificationSessionId sessionId,
                 ConversationActions.Request request,
-                IConversationActionsCallback callback) {
+                ITextClassifierCallback callback) {
             Preconditions.checkNotNull(request);
             Preconditions.checkNotNull(callback);
-            TextClassifierService.this.onSuggestConversationActions(
-                    sessionId,
-                    request,
-                    mCancellationSignal,
-                    new Callback<ConversationActions>() {
-                        @Override
-                        public void onSuccess(ConversationActions result) {
-                            try {
-                                callback.onSuccess(result);
-                            } catch (RemoteException e) {
-                                Slog.d(LOG_TAG, "Error calling callback");
-                            }
-                        }
-
-                        @Override
-                        public void onFailure(CharSequence error) {
-                            try {
-                                callback.onFailure();
-                            } catch (RemoteException e) {
-                                Slog.d(LOG_TAG, "Error calling callback");
-                            }
-                        }
-                    });
+            mMainThreadHandler.post(() -> TextClassifierService.this.onSuggestConversationActions(
+                    sessionId, request, mCancellationSignal, new ProxyCallback<>(callback)));
         }
 
-        /** {@inheritDoc} */
         @Override
         public void onCreateTextClassificationSession(
                 TextClassificationContext context, TextClassificationSessionId sessionId) {
             Preconditions.checkNotNull(context);
             Preconditions.checkNotNull(sessionId);
-            TextClassifierService.this.onCreateTextClassificationSession(context, sessionId);
+            mMainThreadHandler.post(
+                    () -> TextClassifierService.this.onCreateTextClassificationSession(
+                            context, sessionId));
         }
 
-        /** {@inheritDoc} */
         @Override
         public void onDestroyTextClassificationSession(TextClassificationSessionId sessionId) {
-            TextClassifierService.this.onDestroyTextClassificationSession(sessionId);
+            mMainThreadHandler.post(
+                    () -> TextClassifierService.this.onDestroyTextClassificationSession(sessionId));
         }
     };
 
@@ -301,6 +216,7 @@ public abstract class TextClassifierService extends Service {
      * @param cancellationSignal object to watch for canceling the current operation
      * @param callback the callback to return the result to
      */
+    @MainThread
     public abstract void onSuggestSelection(
             @Nullable TextClassificationSessionId sessionId,
             @NonNull TextSelection.Request request,
@@ -316,6 +232,7 @@ public abstract class TextClassifierService extends Service {
      * @param cancellationSignal object to watch for canceling the current operation
      * @param callback the callback to return the result to
      */
+    @MainThread
     public abstract void onClassifyText(
             @Nullable TextClassificationSessionId sessionId,
             @NonNull TextClassification.Request request,
@@ -331,6 +248,7 @@ public abstract class TextClassifierService extends Service {
      * @param cancellationSignal object to watch for canceling the current operation
      * @param callback the callback to return the result to
      */
+    @MainThread
     public abstract void onGenerateLinks(
             @Nullable TextClassificationSessionId sessionId,
             @NonNull TextLinks.Request request,
@@ -345,12 +263,14 @@ public abstract class TextClassifierService extends Service {
      * @param cancellationSignal object to watch for canceling the current operation
      * @param callback the callback to return the result to
      */
+    @MainThread
     public void onDetectLanguage(
             @Nullable TextClassificationSessionId sessionId,
             @NonNull TextLanguage.Request request,
             @NonNull CancellationSignal cancellationSignal,
             @NonNull Callback<TextLanguage> callback) {
-        callback.onSuccess(getLocalTextClassifier().detectLanguage(request));
+        mSingleThreadExecutor.submit(() ->
+                callback.onSuccess(getLocalTextClassifier().detectLanguage(request)));
     }
 
     /**
@@ -361,12 +281,14 @@ public abstract class TextClassifierService extends Service {
      * @param cancellationSignal object to watch for canceling the current operation
      * @param callback the callback to return the result to
      */
+    @MainThread
     public void onSuggestConversationActions(
             @Nullable TextClassificationSessionId sessionId,
             @NonNull ConversationActions.Request request,
             @NonNull CancellationSignal cancellationSignal,
             @NonNull Callback<ConversationActions> callback) {
-        callback.onSuccess(getLocalTextClassifier().suggestConversationActions(request));
+        mSingleThreadExecutor.submit(() ->
+                callback.onSuccess(getLocalTextClassifier().suggestConversationActions(request)));
     }
 
     /**
@@ -383,6 +305,7 @@ public abstract class TextClassifierService extends Service {
      *      instead
      */
     @Deprecated
+    @MainThread
     public void onSelectionEvent(
             @Nullable TextClassificationSessionId sessionId, @NonNull SelectionEvent event) {}
 
@@ -396,6 +319,7 @@ public abstract class TextClassifierService extends Service {
      * @param sessionId the session id
      * @param event the TextClassifier event
      */
+    @MainThread
     public void onTextClassifierEvent(
             @Nullable TextClassificationSessionId sessionId, @NonNull TextClassifierEvent event) {}
 
@@ -405,6 +329,7 @@ public abstract class TextClassifierService extends Service {
      * @param context the text classification context
      * @param sessionId the session's Id
      */
+    @MainThread
     public void onCreateTextClassificationSession(
             @NonNull TextClassificationContext context,
             @NonNull TextClassificationSessionId sessionId) {}
@@ -414,6 +339,7 @@ public abstract class TextClassifierService extends Service {
      *
      * @param sessionId the id of the session to destroy
      */
+    @MainThread
     public void onDestroyTextClassificationSession(
             @NonNull TextClassificationSessionId sessionId) {}
 
@@ -441,6 +367,11 @@ public abstract class TextClassifierService extends Service {
         return TextClassifier.NO_OP;
     }
 
+    /** @hide **/
+    public static <T extends Parcelable> T getResponse(Bundle bundle) {
+        return bundle.getParcelable(KEY_RESULT);
+    }
+
     /**
      * Callbacks for TextClassifierService results.
      *
@@ -494,4 +425,44 @@ public abstract class TextClassifierService extends Service {
                 si.permission));
         return null;
     }
+
+    /**
+     * Forwards the callback result to a wrapped binder callback.
+     */
+    private static final class ProxyCallback<T extends Parcelable> implements Callback<T> {
+        private WeakReference<ITextClassifierCallback> mTextClassifierCallback;
+
+        private ProxyCallback(ITextClassifierCallback textClassifierCallback) {
+            mTextClassifierCallback =
+                    new WeakReference<>(Preconditions.checkNotNull(textClassifierCallback));
+        }
+
+        @Override
+        public void onSuccess(T result) {
+            ITextClassifierCallback callback = mTextClassifierCallback.get();
+            if (callback == null) {
+                return;
+            }
+            try {
+                Bundle bundle = new Bundle(1);
+                bundle.putParcelable(KEY_RESULT, result);
+                callback.onSuccess(bundle);
+            } catch (RemoteException e) {
+                Slog.d(LOG_TAG, "Error calling callback");
+            }
+        }
+
+        @Override
+        public void onFailure(CharSequence error) {
+            ITextClassifierCallback callback = mTextClassifierCallback.get();
+            if (callback == null) {
+                return;
+            }
+            try {
+                callback.onFailure();
+            } catch (RemoteException e) {
+                Slog.d(LOG_TAG, "Error calling callback");
+            }
+        }
+    }
 }
diff --git a/core/java/android/view/textclassifier/SystemTextClassifier.java b/core/java/android/view/textclassifier/SystemTextClassifier.java
index 8b370f543ccc..8f8766e3f783 100644
--- a/core/java/android/view/textclassifier/SystemTextClassifier.java
+++ b/core/java/android/view/textclassifier/SystemTextClassifier.java
@@ -20,15 +20,14 @@ import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.WorkerThread;
 import android.content.Context;
+import android.os.Bundle;
 import android.os.Looper;
+import android.os.Parcelable;
 import android.os.RemoteException;
 import android.os.ServiceManager;
-import android.service.textclassifier.IConversationActionsCallback;
-import android.service.textclassifier.ITextClassificationCallback;
+import android.service.textclassifier.ITextClassifierCallback;
 import android.service.textclassifier.ITextClassifierService;
-import android.service.textclassifier.ITextLanguageCallback;
-import android.service.textclassifier.ITextLinksCallback;
-import android.service.textclassifier.ITextSelectionCallback;
+import android.service.textclassifier.TextClassifierService;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.annotations.VisibleForTesting.Visibility;
@@ -73,9 +72,10 @@ public final class SystemTextClassifier implements TextClassifier {
         Utils.checkMainThread();
         try {
             request.setCallingPackageName(mPackageName);
-            final TextSelectionCallback callback = new TextSelectionCallback();
+            final BlockingCallback<TextSelection> callback =
+                    new BlockingCallback<>("textselection");
             mManagerService.onSuggestSelection(mSessionId, request, callback);
-            final TextSelection selection = callback.mReceiver.get();
+            final TextSelection selection = callback.get();
             if (selection != null) {
                 return selection;
             }
@@ -95,9 +95,10 @@ public final class SystemTextClassifier implements TextClassifier {
         Utils.checkMainThread();
         try {
             request.setCallingPackageName(mPackageName);
-            final TextClassificationCallback callback = new TextClassificationCallback();
+            final BlockingCallback<TextClassification> callback =
+                    new BlockingCallback<>("textclassification");
             mManagerService.onClassifyText(mSessionId, request, callback);
-            final TextClassification classification = callback.mReceiver.get();
+            final TextClassification classification = callback.get();
             if (classification != null) {
                 return classification;
             }
@@ -122,9 +123,10 @@ public final class SystemTextClassifier implements TextClassifier {
 
         try {
             request.setCallingPackageName(mPackageName);
-            final TextLinksCallback callback = new TextLinksCallback();
+            final BlockingCallback<TextLinks> callback =
+                    new BlockingCallback<>("textlinks");
             mManagerService.onGenerateLinks(mSessionId, request, callback);
-            final TextLinks links = callback.mReceiver.get();
+            final TextLinks links = callback.get();
             if (links != null) {
                 return links;
             }
@@ -165,9 +167,10 @@ public final class SystemTextClassifier implements TextClassifier {
 
         try {
             request.setCallingPackageName(mPackageName);
-            final TextLanguageCallback callback = new TextLanguageCallback();
+            final BlockingCallback<TextLanguage> callback =
+                    new BlockingCallback<>("textlanguage");
             mManagerService.onDetectLanguage(mSessionId, request, callback);
-            final TextLanguage textLanguage = callback.mReceiver.get();
+            final TextLanguage textLanguage = callback.get();
             if (textLanguage != null) {
                 return textLanguage;
             }
@@ -184,9 +187,10 @@ public final class SystemTextClassifier implements TextClassifier {
 
         try {
             request.setCallingPackageName(mPackageName);
-            final ConversationActionsCallback callback = new ConversationActionsCallback();
+            final BlockingCallback<ConversationActions> callback =
+                    new BlockingCallback<>("conversation-actions");
             mManagerService.onSuggestConversationActions(mSessionId, request, callback);
-            final ConversationActions conversationActions = callback.mReceiver.get();
+            final ConversationActions conversationActions = callback.get();
             if (conversationActions != null) {
                 return conversationActions;
             }
@@ -245,82 +249,28 @@ public final class SystemTextClassifier implements TextClassifier {
         }
     }
 
-    private static final class TextSelectionCallback extends ITextSelectionCallback.Stub {
+    private static final class BlockingCallback<T extends Parcelable>
+            extends ITextClassifierCallback.Stub {
+        private final ResponseReceiver<T> mReceiver;
 
-        final ResponseReceiver<TextSelection> mReceiver = new ResponseReceiver<>("textselection");
-
-        @Override
-        public void onSuccess(TextSelection selection) {
-            mReceiver.onSuccess(selection);
+        BlockingCallback(String name) {
+            mReceiver = new ResponseReceiver<>(name);
         }
 
         @Override
-        public void onFailure() {
-            mReceiver.onFailure();
-        }
-    }
-
-    private static final class TextClassificationCallback extends ITextClassificationCallback.Stub {
-
-        final ResponseReceiver<TextClassification> mReceiver =
-                new ResponseReceiver<>("textclassification");
-
-        @Override
-        public void onSuccess(TextClassification classification) {
-            mReceiver.onSuccess(classification);
+        public void onSuccess(Bundle result) {
+            mReceiver.onSuccess(TextClassifierService.getResponse(result));
         }
 
         @Override
         public void onFailure() {
             mReceiver.onFailure();
         }
-    }
 
-    private static final class TextLinksCallback extends ITextLinksCallback.Stub {
-
-        final ResponseReceiver<TextLinks> mReceiver = new ResponseReceiver<>("textlinks");
-
-        @Override
-        public void onSuccess(TextLinks links) {
-            mReceiver.onSuccess(links);
-        }
-
-        @Override
-        public void onFailure() {
-            mReceiver.onFailure();
-        }
-    }
-
-    private static final class TextLanguageCallback extends ITextLanguageCallback.Stub {
-
-        final ResponseReceiver<TextLanguage> mReceiver = new ResponseReceiver<>("textlanguage");
-
-        @Override
-        public void onSuccess(TextLanguage textLanguage) {
-            mReceiver.onSuccess(textLanguage);
-        }
-
-        @Override
-        public void onFailure() {
-            mReceiver.onFailure();
-        }
-    }
-
-    private static final class ConversationActionsCallback
-            extends IConversationActionsCallback.Stub {
-
-        final ResponseReceiver<ConversationActions> mReceiver =
-                new ResponseReceiver<>("conversationaction");
-
-        @Override
-        public void onSuccess(ConversationActions conversationActions) {
-            mReceiver.onSuccess(conversationActions);
+        public T get() {
+            return mReceiver.get();
         }
 
-        @Override
-        public void onFailure() {
-            mReceiver.onFailure();
-        }
     }
 
     private static final class ResponseReceiver<T> {
diff --git a/services/core/java/com/android/server/textclassifier/TextClassificationManagerService.java b/services/core/java/com/android/server/textclassifier/TextClassificationManagerService.java
index ef771406805b..a5d291f94751 100644
--- a/services/core/java/com/android/server/textclassifier/TextClassificationManagerService.java
+++ b/services/core/java/com/android/server/textclassifier/TextClassificationManagerService.java
@@ -27,12 +27,8 @@ import android.os.Binder;
 import android.os.IBinder;
 import android.os.RemoteException;
 import android.os.UserHandle;
-import android.service.textclassifier.IConversationActionsCallback;
-import android.service.textclassifier.ITextClassificationCallback;
+import android.service.textclassifier.ITextClassifierCallback;
 import android.service.textclassifier.ITextClassifierService;
-import android.service.textclassifier.ITextLanguageCallback;
-import android.service.textclassifier.ITextLinksCallback;
-import android.service.textclassifier.ITextSelectionCallback;
 import android.service.textclassifier.TextClassifierService;
 import android.util.Slog;
 import android.util.SparseArray;
@@ -132,7 +128,7 @@ public final class TextClassificationManagerService extends ITextClassifierServi
     @Override
     public void onSuggestSelection(
             TextClassificationSessionId sessionId,
-            TextSelection.Request request, ITextSelectionCallback callback)
+            TextSelection.Request request, ITextClassifierCallback callback)
             throws RemoteException {
         Preconditions.checkNotNull(request);
         Preconditions.checkNotNull(callback);
@@ -155,7 +151,7 @@ public final class TextClassificationManagerService extends ITextClassifierServi
     @Override
     public void onClassifyText(
             TextClassificationSessionId sessionId,
-            TextClassification.Request request, ITextClassificationCallback callback)
+            TextClassification.Request request, ITextClassifierCallback callback)
             throws RemoteException {
         Preconditions.checkNotNull(request);
         Preconditions.checkNotNull(callback);
@@ -178,7 +174,7 @@ public final class TextClassificationManagerService extends ITextClassifierServi
     @Override
     public void onGenerateLinks(
             TextClassificationSessionId sessionId,
-            TextLinks.Request request, ITextLinksCallback callback)
+            TextLinks.Request request, ITextClassifierCallback callback)
             throws RemoteException {
         Preconditions.checkNotNull(request);
         Preconditions.checkNotNull(callback);
@@ -241,7 +237,7 @@ public final class TextClassificationManagerService extends ITextClassifierServi
     public void onDetectLanguage(
             TextClassificationSessionId sessionId,
             TextLanguage.Request request,
-            ITextLanguageCallback callback) throws RemoteException {
+            ITextClassifierCallback callback) throws RemoteException {
         Preconditions.checkNotNull(request);
         Preconditions.checkNotNull(callback);
         validateInput(mContext, request.getCallingPackageName());
@@ -264,7 +260,7 @@ public final class TextClassificationManagerService extends ITextClassifierServi
     public void onSuggestConversationActions(
             TextClassificationSessionId sessionId,
             ConversationActions.Request request,
-            IConversationActionsCallback callback) throws RemoteException {
+            ITextClassifierCallback callback) throws RemoteException {
         Preconditions.checkNotNull(request);
         Preconditions.checkNotNull(callback);
         validateInput(mContext, request.getCallingPackageName());