diff --git a/services/core/java/com/android/server/wm/TaskOrganizerController.java b/services/core/java/com/android/server/wm/TaskOrganizerController.java
index 8edcd2fb77655..8e95ca7ef9614 100644
--- a/services/core/java/com/android/server/wm/TaskOrganizerController.java
+++ b/services/core/java/com/android/server/wm/TaskOrganizerController.java
@@ -78,7 +78,9 @@ class TaskOrganizerController extends ITaskOrganizerController.Stub {
             synchronized (mGlobalLock) {
                 final TaskOrganizerState state = mTaskOrganizerStates.remove(
                         mTaskOrganizer.asBinder());
-                state.dispose();
+                if (state != null) {
+                    state.dispose();
+                }
             }
         }
     };
@@ -222,9 +224,20 @@ class TaskOrganizerController extends ITaskOrganizerController.Stub {
 
     @Override
     public void unregisterTaskOrganizer(ITaskOrganizer organizer) {
-        final TaskOrganizerState state = mTaskOrganizerStates.remove(organizer.asBinder());
-        state.unlinkDeath();
-        state.dispose();
+        enforceStackPermission("unregisterTaskOrganizer()");
+        final long origId = Binder.clearCallingIdentity();
+        try {
+            synchronized (mGlobalLock) {
+                final TaskOrganizerState state = mTaskOrganizerStates.remove(organizer.asBinder());
+                if (state == null) {
+                    return;
+                }
+                state.unlinkDeath();
+                state.dispose();
+            }
+        } finally {
+            Binder.restoreCallingIdentity(origId);
+        }
     }
 
     ITaskOrganizer getTaskOrganizer(int windowingMode) {