diff --git a/core/java/android/view/textclassifier/TextSelection.java b/core/java/android/view/textclassifier/TextSelection.java
index 5b32649ac55bf..c1913f69546cb 100644
--- a/core/java/android/view/textclassifier/TextSelection.java
+++ b/core/java/android/view/textclassifier/TextSelection.java
@@ -137,6 +137,15 @@ public final class TextSelection implements Parcelable {
         return mExtras;
     }
 
+    /** @hide */
+    public TextSelection.Builder toBuilder() {
+        return new TextSelection.Builder(mStartIndex, mEndIndex)
+                .setId(mId)
+                .setEntityConfidence(mEntityConfidence)
+                .setTextClassification(mTextClassification)
+                .setExtras(mExtras);
+    }
+
     @Override
     public String toString() {
         return String.format(
@@ -188,6 +197,12 @@ public final class TextSelection implements Parcelable {
             return this;
         }
 
+        Builder setEntityConfidence(EntityConfidence scores) {
+            mEntityConfidence.clear();
+            mEntityConfidence.putAll(scores.toMap());
+            return this;
+        }
+
         /**
          * Sets an id for the TextSelection object.
          */
diff --git a/core/tests/coretests/src/android/view/textclassifier/TextSelectionTest.java b/core/tests/coretests/src/android/view/textclassifier/TextSelectionTest.java
index 82788c806fed6..5def552c72b49 100644
--- a/core/tests/coretests/src/android/view/textclassifier/TextSelectionTest.java
+++ b/core/tests/coretests/src/android/view/textclassifier/TextSelectionTest.java
@@ -16,14 +16,24 @@
 
 package android.view.textclassifier;
 
+import static com.google.common.truth.Truth.assertThat;
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 
+import android.app.PendingIntent;
+import android.app.RemoteAction;
+import android.content.Context;
+import android.content.Intent;
+import android.graphics.Bitmap;
+import android.graphics.Color;
+import android.graphics.drawable.Icon;
 import android.os.Bundle;
 import android.os.LocaleList;
 import android.os.Parcel;
 
+import androidx.test.InstrumentationRegistry;
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
@@ -107,4 +117,46 @@ public class TextSelectionTest {
         assertEquals(1, resultSystemTcMetadata.getUserId());
         assertFalse(resultSystemTcMetadata.useDefaultTextClassifier());
     }
+
+    @Test
+    public void testToBuilder() {
+        final Context context = InstrumentationRegistry.getInstrumentation().getContext();
+        final int startIndex = 13;
+        final int endIndex = 37;
+        final String id = "id";
+        final Icon icon1 = generateTestIcon(5, 5, Color.RED);
+
+        final TextClassification classification = new TextClassification.Builder()
+                .addAction(new RemoteAction(icon1, "title1", "desc1",
+                        PendingIntent.getActivity(context, 0, new Intent("action1"), 0)))
+                .setEntityType(TextClassifier.TYPE_ADDRESS, 1.0f)
+                .build();
+        final TextSelection textSelection = new TextSelection.Builder(startIndex, endIndex)
+                .setId(id)
+                .setEntityType(TextClassifier.TYPE_ADDRESS, 1.0f)
+                .setExtras(BUNDLE)
+                .setTextClassification(classification)
+                .build();
+
+        final TextSelection fromBuilder = textSelection.toBuilder().build();
+
+        assertThat(fromBuilder.getId()).isEqualTo(textSelection.getId());
+        assertThat(fromBuilder.getSelectionStartIndex())
+                .isEqualTo(textSelection.getSelectionStartIndex());
+        assertThat(fromBuilder.getSelectionEndIndex())
+                .isEqualTo(textSelection.getSelectionEndIndex());
+        assertThat(fromBuilder.getTextClassification())
+                .isSameInstanceAs(textSelection.getTextClassification());
+        assertThat(fromBuilder.getExtras()).isSameInstanceAs(textSelection.getExtras());
+    }
+
+    private Icon generateTestIcon(int width, int height, int colorValue) {
+        final int numPixels = width * height;
+        final int[] colors = new int[numPixels];
+        for (int i = 0; i < numPixels; ++i) {
+            colors[i] = colorValue;
+        }
+        final Bitmap bitmap = Bitmap.createBitmap(colors, width, height, Bitmap.Config.ARGB_8888);
+        return Icon.createWithBitmap(bitmap);
+    }
 }
diff --git a/services/core/java/com/android/server/textclassifier/TextClassificationManagerService.java b/services/core/java/com/android/server/textclassifier/TextClassificationManagerService.java
index 93ba758d29d62..753b42b245565 100644
--- a/services/core/java/com/android/server/textclassifier/TextClassificationManagerService.java
+++ b/services/core/java/com/android/server/textclassifier/TextClassificationManagerService.java
@@ -192,7 +192,7 @@ public final class TextClassificationManagerService extends ITextClassifierServi
                 request.getSystemTextClassifierMetadata(),
                 /* verifyCallingPackage= */ true,
                 /* attemptToBind= */ true,
-                service -> service.onSuggestSelection(sessionId, request, callback),
+                service -> service.onSuggestSelection(sessionId, request, wrap(callback)),
                 "onSuggestSelection",
                 callback);
     }
@@ -1057,6 +1057,8 @@ public final class TextClassificationManagerService extends ITextClassifierServi
                 rewriteTextClassificationIcons(result);
             } else if (parcelled instanceof ConversationActions) {
                 rewriteConversationActionsIcons(result);
+            } else if (parcelled instanceof TextSelection) {
+                rewriteTextSelectionIcons(result);
             } else {
                 // do nothing.
             }
@@ -1067,10 +1069,32 @@ public final class TextClassificationManagerService extends ITextClassifierServi
             }
         }
 
-        private static void rewriteTextClassificationIcons(Bundle result) {
-            final TextClassification classification = TextClassifierService.getResponse(result);
+        private static void rewriteTextSelectionIcons(Bundle result) {
+            final TextSelection textSelection = TextClassifierService.getResponse(result);
+            if (textSelection.getTextClassification() == null) {
+                return;
+            }
+            TextClassification newTextClassification =
+                    rewriteTextClassificationIcons(textSelection.getTextClassification());
+            if (newTextClassification == null) {
+                return;
+            }
+            TextClassifierService.putResponse(
+                    result,
+                    textSelection.toBuilder()
+                            .setTextClassification(newTextClassification)
+                            .build());
+        }
+
+        /**
+         * Returns a new {@link TextClassification} if any modification is made, {@code null}
+         * otherwise.
+         */
+        @Nullable
+        private static TextClassification rewriteTextClassificationIcons(
+                TextClassification textClassification) {
             boolean rewrite = false;
-            final List<RemoteAction> actions = classification.getActions();
+            final List<RemoteAction> actions = textClassification.getActions();
             final int size = actions.size();
             final List<RemoteAction> validActions = new ArrayList<>(size);
             for (int i = 0; i < size; i++) {
@@ -1084,13 +1108,21 @@ public final class TextClassificationManagerService extends ITextClassifierServi
                 }
                 validActions.add(validAction);
             }
-            if (rewrite) {
-                TextClassifierService.putResponse(
-                        result,
-                        classification.toBuilder()
-                                .clearActions()
-                                .addActions(validActions)
-                                .build());
+            return rewrite
+                    ? textClassification
+                            .toBuilder()
+                            .clearActions()
+                            .addActions(validActions)
+                            .build()
+                    : null;
+        }
+
+        private static void rewriteTextClassificationIcons(Bundle result) {
+            final TextClassification classification = TextClassifierService.getResponse(result);
+            TextClassification newTextClassification = rewriteTextClassificationIcons(
+                    classification);
+            if (newTextClassification != null) {
+                TextClassifierService.putResponse(result, newTextClassification);
             }
         }