diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitContainer.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitContainer.java
index 08b7bb8..39cface 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitContainer.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitContainer.java
@@ -201,7 +201,7 @@
             return null;
         }
         return new SplitInfo(primaryActivityStack, secondaryActivityStack,
-                mCurrentSplitAttributes, mToken);
+                mCurrentSplitAttributes, SplitInfo.Token.createFromBinder(mToken));
     }
 
     static boolean shouldFinishPrimaryWithSecondary(@NonNull SplitRule splitRule) {
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
index b2e5b75..2de8cc0 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
@@ -35,6 +35,7 @@
 import static android.window.TaskFragmentTransaction.TYPE_TASK_FRAGMENT_PARENT_INFO_CHANGED;
 import static android.window.TaskFragmentTransaction.TYPE_TASK_FRAGMENT_VANISHED;
 
+import static androidx.window.extensions.embedding.ActivityEmbeddingOptionsProperties.KEY_ACTIVITY_STACK_TOKEN;
 import static androidx.window.extensions.embedding.ActivityEmbeddingOptionsProperties.KEY_OVERLAY_TAG;
 import static androidx.window.extensions.embedding.SplitContainer.getFinishPrimaryWithSecondaryBehavior;
 import static androidx.window.extensions.embedding.SplitContainer.getFinishSecondaryWithPrimaryBehavior;
@@ -111,10 +112,6 @@
     static final boolean ENABLE_SHELL_TRANSITIONS =
             SystemProperties.getBoolean("persist.wm.debug.shell_transit", true);
 
-    // TODO(b/295993745): remove after prebuilt library is updated.
-    private static final String KEY_ACTIVITY_STACK_TOKEN =
-            "androidx.window.extensions.embedding.ActivityStackToken";
-
     @VisibleForTesting
     @GuardedBy("mLock")
     final SplitPresenter mPresenter;
@@ -553,7 +550,7 @@
     }
 
     @Override
-    public void updateActivityStackAttributes(@NonNull IBinder activityStackToken,
+    public void updateActivityStackAttributes(@NonNull ActivityStack.Token activityStackToken,
                                               @NonNull ActivityStackAttributes attributes) {
         if (!Flags.activityEmbeddingOverlayPresentationFlag()) {
             return;
@@ -562,7 +559,7 @@
         Objects.requireNonNull(attributes);
 
         synchronized (mLock) {
-            final TaskFragmentContainer container = getContainer(activityStackToken);
+            final TaskFragmentContainer container = getContainer(activityStackToken.getRawToken());
             if (container == null) {
                 Log.w(TAG, "Cannot find TaskFragmentContainer for token:" + activityStackToken);
                 return;
@@ -582,13 +579,14 @@
 
     @Override
     @Nullable
-    public ParentContainerInfo getParentContainerInfo(@NonNull IBinder activityStackToken) {
+    public ParentContainerInfo getParentContainerInfo(
+            @NonNull ActivityStack.Token activityStackToken) {
         if (!Flags.activityEmbeddingOverlayPresentationFlag()) {
             return null;
         }
         Objects.requireNonNull(activityStackToken);
         synchronized (mLock) {
-            final TaskFragmentContainer container = getContainer(activityStackToken);
+            final TaskFragmentContainer container = getContainer(activityStackToken.getRawToken());
             if (container == null) {
                 return null;
             }
@@ -600,7 +598,7 @@
 
     @Override
     @Nullable
-    public IBinder getActivityStackToken(@NonNull String tag) {
+    public ActivityStack.Token getActivityStackToken(@NonNull String tag) {
         if (!Flags.activityEmbeddingOverlayPresentationFlag()) {
             return null;
         }
@@ -611,7 +609,8 @@
             if (taskFragmentContainer == null) {
                 return null;
             }
-            return taskFragmentContainer.getTaskFragmentToken();
+            return ActivityStack.Token.createFromBinder(taskFragmentContainer
+                    .getTaskFragmentToken());
         }
     }
 
@@ -2760,8 +2759,10 @@
             // TODO(b/232042367): Consolidate the activity create handling so that we can handle
             // cross-process the same as normal.
 
-            IBinder activityStackToken = options.getBinder(KEY_ACTIVITY_STACK_TOKEN);
-            if (activityStackToken != null) {
+            final Bundle bundle = options.getBundle(KEY_ACTIVITY_STACK_TOKEN);
+            if (bundle != null) {
+                final IBinder activityStackToken = ActivityStack.Token.readFromBundle(bundle)
+                        .getRawToken();
                 // Put activityStack token to #KEY_LAUNCH_TASK_FRAGMENT_TOKEN to launch the activity
                 // into the taskFragment associated with the token.
                 options.putBinder(KEY_LAUNCH_TASK_FRAGMENT_TOKEN, activityStackToken);
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentContainer.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentContainer.java
index 6fe8e50..a6bf99d 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentContainer.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentContainer.java
@@ -367,7 +367,8 @@
         if (activities == null) {
             return null;
         }
-        return new ActivityStack(activities, isEmpty(), mToken, mOverlayTag);
+        return new ActivityStack(activities, isEmpty(),
+                ActivityStack.Token.createFromBinder(mToken), mOverlayTag);
     }
 
     /** Adds the activity that will be reparented to this container. */
diff --git a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/OverlayPresentationTest.java b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/OverlayPresentationTest.java
index 34d43ad..28fbadb 100644
--- a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/OverlayPresentationTest.java
+++ b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/OverlayPresentationTest.java
@@ -399,7 +399,8 @@
                         new ActivityStackAttributes.Builder().build()));
 
         assertThrows(NullPointerException.class, () ->
-                mSplitController.updateActivityStackAttributes(new Binder(), null));
+                mSplitController.updateActivityStackAttributes(
+                        ActivityStack.Token.createFromBinder(new Binder()), null));
 
         verify(mSplitPresenter, never()).applyActivityStackAttributes(any(), any(), any(), any());
     }
@@ -408,7 +409,8 @@
     public void testUpdateActivityStackAttributes_nullContainer_earlyReturn() {
         final TaskFragmentContainer container = mSplitController.newContainer(mActivity,
                 mActivity.getTaskId());
-        mSplitController.updateActivityStackAttributes(container.getTaskFragmentToken(),
+        mSplitController.updateActivityStackAttributes(
+                ActivityStack.Token.createFromBinder(container.getTaskFragmentToken()),
                 new ActivityStackAttributes.Builder().build());
 
         verify(mSplitPresenter, never()).applyActivityStackAttributes(any(), any(), any(), any());
@@ -418,7 +420,8 @@
     public void testUpdateActivityStackAttributes_notOverlay_earlyReturn() {
         final TaskFragmentContainer container = createMockTaskFragmentContainer(mActivity);
 
-        mSplitController.updateActivityStackAttributes(container.getTaskFragmentToken(),
+        mSplitController.updateActivityStackAttributes(
+                ActivityStack.Token.createFromBinder(container.getTaskFragmentToken()),
                 new ActivityStackAttributes.Builder().build());
 
         verify(mSplitPresenter, never()).applyActivityStackAttributes(any(), any(), any(), any());
@@ -431,7 +434,8 @@
         final ActivityStackAttributes attrs = new ActivityStackAttributes.Builder().build();
         final IBinder token = container.getTaskFragmentToken();
 
-        mSplitController.updateActivityStackAttributes(token, attrs);
+        mSplitController.updateActivityStackAttributes(ActivityStack.Token.createFromBinder(token),
+                attrs);
 
         verify(mSplitPresenter).applyActivityStackAttributes(any(), eq(container), eq(attrs),
                 any());
diff --git a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitControllerTest.java b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitControllerTest.java
index b60943a..00f8b59 100644
--- a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitControllerTest.java
+++ b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitControllerTest.java
@@ -1437,7 +1437,7 @@
     @Test
     public void testUpdateSplitAttributes_nullParams_throwException() {
         assertThrows(NullPointerException.class,
-                () -> mSplitController.updateSplitAttributes(null, SPLIT_ATTRIBUTES));
+                () -> mSplitController.updateSplitAttributes((IBinder) null, SPLIT_ATTRIBUTES));
 
         final SplitContainer splitContainer = mock(SplitContainer.class);
         final IBinder token = new Binder();