Deflake background task removal E2E tests

dotnet/test/E2E/RpcTasksAndHandlersE2ETests.cs

@@ -152,29 +152,38 @@ await TestHelper.WaitForConditionAsync(
             async () =>
             {
                 task = await FindAgentTaskAsync(session, started.AgentId);
-                return task?.LatestResponse?.Contains("TASK_AGENT_DONE", StringComparison.Ordinal) == true
-                    || task?.Result?.Contains("TASK_AGENT_DONE", StringComparison.Ordinal) == true
-                    || task?.Status == GitHub.Copilot.Rpc.TaskStatus.Completed
-                    || task?.Status == GitHub.Copilot.Rpc.TaskStatus.Failed;
+                return task is null
+                    || task.Status == GitHub.Copilot.Rpc.TaskStatus.Completed
+                    || task.Status == GitHub.Copilot.Rpc.TaskStatus.Failed
+                    || task.Status == GitHub.Copilot.Rpc.TaskStatus.Cancelled
+                    || task.Status == GitHub.Copilot.Rpc.TaskStatus.Idle;
             },
             timeout: TimeSpan.FromSeconds(60),
             timeoutMessage: $"Background agent task '{started.AgentId}' did not produce a final observable state.");
 
-        Assert.NotNull(task);
-        Assert.Contains("TASK_AGENT_DONE", task.LatestResponse ?? task.Result ?? string.Empty);
-        await taskCompletionNotification.Task.WaitAsync(TimeSpan.FromSeconds(30));
-
-        if (task.Status == GitHub.Copilot.Rpc.TaskStatus.Idle)
+        if (task is not null)
         {
-            var cancel = await session.Rpc.Tasks.CancelAsync(started.AgentId);
-            Assert.True(cancel.Cancelled);
-        }
+            Assert.Contains("TASK_AGENT_DONE", task.LatestResponse ?? task.Result ?? string.Empty);
+
+            if (task.Status == GitHub.Copilot.Rpc.TaskStatus.Idle)
+            {
+                var cancel = await session.Rpc.Tasks.CancelAsync(started.AgentId);
+                Assert.True(cancel.Cancelled);
+            }
 
-        var remove = await session.Rpc.Tasks.RemoveAsync(started.AgentId);
-        Assert.True(remove.Removed);
+            var remove = await session.Rpc.Tasks.RemoveAsync(started.AgentId);
+            // Completion delivery also removes finished tasks, so this call may lose that race.
+            Assert.True(
+                remove.Removed || taskCompletionNotification.Task.IsCompleted,
+                $"Background agent task '{started.AgentId}' was not removed before its completion notification was delivered.");
+        }
 
         var afterRemove = await session.Rpc.Tasks.ListAsync();
-        Assert.DoesNotContain(afterRemove.Tasks.OfType<TaskInfoAgent>(), t => string.Equals(t.Id, started.AgentId, StringComparison.Ordinal));
+        var taskAfterRemove = afterRemove.Tasks.OfType<TaskInfoAgent>()
+            .SingleOrDefault(t => string.Equals(t.Id, started.AgentId, StringComparison.Ordinal));
+        Assert.Null(taskAfterRemove);
+
+        await taskCompletionNotification.Task.WaitAsync(TimeSpan.FromSeconds(30));
     }
 
     [Fact]

python/e2e/test_rpc_tasks_and_handlers_e2e.py

@@ -445,20 +445,28 @@ def on_event(event):
                 60.0,
                 f"Task {task_id} did not produce a final observable state",
             )
-            assert found_task is not None, f"Task {task_id} disappeared before it completed"
-            assert "TASK_AGENT_DONE" in (found_task.latest_response or found_task.result or "")
-            await asyncio.wait_for(task_completion_notification, timeout=30.0)
-
-            if found_task.status == TaskInfoStatus.IDLE:
-                cancel = await session.rpc.tasks.cancel(TasksCancelRequest(id=task_id))
-                assert cancel.cancelled is True
-
-            # Remove the task
-            remove = await session.rpc.tasks.remove(TasksRemoveRequest(id=task_id))
-            assert remove.removed is True
+            if found_task is not None:
+                assert "TASK_AGENT_DONE" in (found_task.latest_response or found_task.result or "")
+
+                if found_task.status == TaskInfoStatus.IDLE:
+                    cancel = await session.rpc.tasks.cancel(TasksCancelRequest(id=task_id))
+                    assert cancel.cancelled is True
+
+                remove = await session.rpc.tasks.remove(TasksRemoveRequest(id=task_id))
+                # Completion delivery also removes finished tasks, so this call may lose that race.
+                assert remove.removed or task_completion_notification.done(), (
+                    f"Task {task_id} was not removed before its completion "
+                    "notification was delivered"
+                )
 
             after_remove = await session.rpc.tasks.list()
-            assert not any(t.id == task_id for t in (after_remove.tasks or []))
+            task_after_remove = next(
+                (task for task in (after_remove.tasks or []) if task.id == task_id),
+                None,
+            )
+            assert task_after_remove is None
+
+            await asyncio.wait_for(task_completion_notification, timeout=30.0)
         finally:
             unsubscribe()
             await session.disconnect()