diff --git a/hostsidetests/services/activitymanager/src/android/server/cts/ActivityAndWindowManagersState.java b/hostsidetests/services/activitymanager/src/android/server/cts/ActivityAndWindowManagersState.java
index ee4afbb..a1583d9d 100644
--- a/hostsidetests/services/activitymanager/src/android/server/cts/ActivityAndWindowManagersState.java
+++ b/hostsidetests/services/activitymanager/src/android/server/cts/ActivityAndWindowManagersState.java
@@ -326,6 +326,14 @@
                     return false;
                 }
             }
+
+            for (WindowTask wTask : wStack.mTasks) {
+                if (aStack.getTask(wTask.mTaskId) == null) {
+                    log("Task is in WM but not in AM, waiting for it to settle, taskId="
+                            + wTask.mTaskId);
+                    return false;
+                }
+            }
         }
         return true;
     }
diff --git a/hostsidetests/services/activitymanager/src/android/server/cts/ActivityManagerFreeformStackTests.java b/hostsidetests/services/activitymanager/src/android/server/cts/ActivityManagerFreeformStackTests.java
index f86d9ec..abe067f 100644
--- a/hostsidetests/services/activitymanager/src/android/server/cts/ActivityManagerFreeformStackTests.java
+++ b/hostsidetests/services/activitymanager/src/android/server/cts/ActivityManagerFreeformStackTests.java
@@ -36,7 +36,7 @@
 
         launchActivityInStack(FREEFORM_ACTIVITY, FREEFORM_WORKSPACE_STACK_ID);
 
-        mAmWmState.computeState(mDevice, new String[] {FREEFORM_ACTIVITY});
+        mAmWmState.computeState(mDevice, new String[] {FREEFORM_ACTIVITY, TEST_ACTIVITY});
 
         if (!supportsFreeform()) {
             mAmWmState.assertDoesNotContainStack(
diff --git a/hostsidetests/services/activitymanager/src/android/server/cts/ActivityManagerState.java b/hostsidetests/services/activitymanager/src/android/server/cts/ActivityManagerState.java
index 7e8ee04..d392e00 100644
--- a/hostsidetests/services/activitymanager/src/android/server/cts/ActivityManagerState.java
+++ b/hostsidetests/services/activitymanager/src/android/server/cts/ActivityManagerState.java
@@ -339,6 +339,15 @@
         List<ActivityTask> getTasks() {
             return new ArrayList(mTasks);
         }
+
+        ActivityTask getTask(int taskId) {
+            for (ActivityTask task : mTasks) {
+                if (taskId == task.mTaskId) {
+                    return task;
+                }
+            }
+            return null;
+        }
     }
 
     static class ActivityTask extends ActivityContainer {