diff --git a/core/java/android/os/CancellationSignalBeamer.java b/core/java/android/os/CancellationSignalBeamer.java
index b4247831ddc5c..5c0b22172674f 100644
--- a/core/java/android/os/CancellationSignalBeamer.java
+++ b/core/java/android/os/CancellationSignalBeamer.java
@@ -18,6 +18,7 @@ package android.os;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.annotation.SuppressLint;
 import android.system.SystemCleaner;
 import android.util.Pair;
 import android.view.inputmethod.CancellableHandwritingGesture;
@@ -137,7 +138,7 @@ public class CancellationSignalBeamer {
          * MUST be forwarded to {@link Receiver#cancel} with proper ordering. See
          * {@link CancellationSignalBeamer} for details.
          */
-        public abstract void onCancel(IBinder token);
+        public abstract void onCancel(@NonNull IBinder token);
 
         /**
          * A {@link #beam}ed {@link CancellationSignal} was GC'd.
@@ -145,7 +146,7 @@ public class CancellationSignalBeamer {
          * MUST be forwarded to {@link Receiver#forget} with proper ordering. See
          * {@link CancellationSignalBeamer} for details.
          */
-        public abstract void onForget(IBinder token);
+        public abstract void onForget(@NonNull IBinder token);
 
         private static final ThreadLocal<Pair<Sender, ArrayList<CloseableToken>>> sScope =
                 new ThreadLocal<>();
@@ -159,7 +160,8 @@ public class CancellationSignalBeamer {
          *  try-with-resources. {@code null} if {@code cs} was {@code null} or if
          *  {@link HandwritingGesture} isn't {@link CancellableHandwritingGesture cancellable}.
          */
-        public MustClose beamScopeIfNeeded(HandwritingGesture gesture) {
+        @NonNull
+        public MustClose beamScopeIfNeeded(@NonNull HandwritingGesture gesture) {
             if (!(gesture instanceof CancellableHandwritingGesture)) {
                 return null;
             }
@@ -189,7 +191,8 @@ public class CancellationSignalBeamer {
          * @param cs {@link CancellationSignal} for which token should be returned.
          * @return {@link IBinder} token.
          */
-        public static IBinder beamFromScope(CancellationSignal cs) {
+        @NonNull
+        public static IBinder beamFromScope(@NonNull CancellationSignal cs) {
             var state = sScope.get();
             if (state != null) {
                 var token = state.first.beam(cs);
@@ -291,6 +294,7 @@ public class CancellationSignalBeamer {
          * @return a {@link CancellationSignal} linked to the given token.
          */
         @Nullable
+        @SuppressLint("VisiblySynchronized")
         public CancellationSignal unbeam(@Nullable IBinder token) {
             if (token == null) {
                 return null;
@@ -327,6 +331,7 @@ public class CancellationSignalBeamer {
          *
          * @param token the token to forget. No-op if {@code null}.
          */
+        @SuppressLint("VisiblySynchronized")
         public void forget(@Nullable IBinder token) {
             synchronized (this) {
                 if (mTokenMap.remove(token) != null) {
@@ -347,6 +352,7 @@ public class CancellationSignalBeamer {
          *
          * @param token the token to forget. No-op if {@code null}.
          */
+        @SuppressLint("VisiblySynchronized")
         public void cancel(@Nullable IBinder token) {
             CancellationSignal cs;
             synchronized (this) {
diff --git a/core/java/android/view/inputmethod/CancellableHandwritingGesture.java b/core/java/android/view/inputmethod/CancellableHandwritingGesture.java
index 3e7974b0a6b8c..09c2b90068177 100644
--- a/core/java/android/view/inputmethod/CancellableHandwritingGesture.java
+++ b/core/java/android/view/inputmethod/CancellableHandwritingGesture.java
@@ -17,6 +17,7 @@
 package android.view.inputmethod;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.annotation.TestApi;
 import android.os.CancellationSignal;
 import android.os.CancellationSignalBeamer;
@@ -28,10 +29,13 @@ import android.os.IBinder;
  */
 @TestApi
 public abstract class CancellableHandwritingGesture extends HandwritingGesture {
+    @NonNull
     CancellationSignal mCancellationSignal;
 
+    @Nullable
     IBinder mCancellationSignalToken;
 
+
     /**
      * Set {@link CancellationSignal} for testing only.
      * @hide
@@ -41,13 +45,20 @@ public abstract class CancellableHandwritingGesture extends HandwritingGesture {
         mCancellationSignal = cancellationSignal;
     }
 
+    @NonNull
     CancellationSignal getCancellationSignal() {
         return mCancellationSignal;
     }
 
-    void unbeamCancellationSignal(CancellationSignalBeamer.Receiver receiver) {
-        mCancellationSignal = receiver.unbeam(mCancellationSignalToken);
-        mCancellationSignalToken = null;
+    /**
+     * Unbeam cancellation token.
+     * @hide
+     */
+    public void unbeamCancellationSignal(@NonNull CancellationSignalBeamer.Receiver receiver) {
+        if (mCancellationSignalToken != null) {
+            mCancellationSignal = receiver.unbeam(mCancellationSignalToken);
+            mCancellationSignalToken = null;
+        }
     }
 
 }
diff --git a/core/tests/coretests/src/android/view/inputmethod/InsertModeGestureTest.java b/core/tests/coretests/src/android/view/inputmethod/InsertModeGestureTest.java
index 11ddba110f7a1..a94f8772fcd08 100644
--- a/core/tests/coretests/src/android/view/inputmethod/InsertModeGestureTest.java
+++ b/core/tests/coretests/src/android/view/inputmethod/InsertModeGestureTest.java
@@ -21,6 +21,7 @@ import static org.junit.Assert.assertNotNull;
 
 import android.graphics.PointF;
 import android.os.CancellationSignal;
+import android.os.CancellationSignalBeamer;
 import android.platform.test.annotations.Presubmit;
 
 import androidx.test.filters.SmallTest;
@@ -54,4 +55,15 @@ public class InsertModeGestureTest {
         assertEquals(FALLBACK_TEXT, gesture.getFallbackText());
         assertEquals(CANCELLATION_SIGNAL, gesture.getCancellationSignal());
     }
+
+    @Test
+    public void testCancellationSignal() {
+        var cs = CANCELLATION_SIGNAL;
+        var gesture = new InsertModeGesture.Builder().setInsertionPoint(INSERTION_POINT)
+                .setCancellationSignal(CANCELLATION_SIGNAL)
+                .setFallbackText(FALLBACK_TEXT).build();
+        gesture.unbeamCancellationSignal(
+                new CancellationSignalBeamer.Receiver(true /* cancelOnSenderDeath */));
+        assertEquals(gesture.getCancellationSignal(), cs);
+    }
 }