diff --git a/core/java/android/service/textclassifier/TextClassifierService.java b/core/java/android/service/textclassifier/TextClassifierService.java
index 9108c33..5143f18 100644
--- a/core/java/android/service/textclassifier/TextClassifierService.java
+++ b/core/java/android/service/textclassifier/TextClassifierService.java
@@ -51,7 +51,6 @@
 
 import com.android.internal.util.Preconditions;
 
-import java.lang.ref.WeakReference;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 
@@ -431,23 +430,18 @@
      * 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 ITextClassifierCallback mTextClassifierCallback;
 
         private ProxyCallback(ITextClassifierCallback textClassifierCallback) {
-            mTextClassifierCallback =
-                    new WeakReference<>(Preconditions.checkNotNull(textClassifierCallback));
+            mTextClassifierCallback = 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);
+                mTextClassifierCallback.onSuccess(bundle);
             } catch (RemoteException e) {
                 Slog.d(LOG_TAG, "Error calling callback");
             }
@@ -455,13 +449,9 @@
 
         @Override
         public void onFailure(CharSequence error) {
-            ITextClassifierCallback callback = mTextClassifierCallback.get();
-            if (callback == null) {
-                return;
-            }
             try {
                 Slog.w(LOG_TAG, "Request fail: " + error);
-                callback.onFailure();
+                mTextClassifierCallback.onFailure();
             } catch (RemoteException e) {
                 Slog.d(LOG_TAG, "Error calling callback");
             }
diff --git a/core/java/android/widget/SelectionActionModeHelper.java b/core/java/android/widget/SelectionActionModeHelper.java
index 51ca805..d0f8093 100644
--- a/core/java/android/widget/SelectionActionModeHelper.java
+++ b/core/java/android/widget/SelectionActionModeHelper.java
@@ -999,6 +999,7 @@
         }
 
         private void onTimeOut() {
+            Log.d(LOG_TAG, "Timeout in TextClassificationAsyncTask");
             if (getStatus() == Status.RUNNING) {
                 onPostExecute(mTimeOutResultSupplier.get());
             }