diff --git a/dotnet/test/E2E/RpcTasksAndHandlersE2ETests.cs b/dotnet/test/E2E/RpcTasksAndHandlersE2ETests.cs index 989fef7c5..640e4f72f 100644 --- a/dotnet/test/E2E/RpcTasksAndHandlersE2ETests.cs +++ b/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(), t => string.Equals(t.Id, started.AgentId, StringComparison.Ordinal)); + var taskAfterRemove = afterRemove.Tasks.OfType() + .SingleOrDefault(t => string.Equals(t.Id, started.AgentId, StringComparison.Ordinal)); + Assert.Null(taskAfterRemove); + + await taskCompletionNotification.Task.WaitAsync(TimeSpan.FromSeconds(30)); } [Fact] diff --git a/python/e2e/test_rpc_tasks_and_handlers_e2e.py b/python/e2e/test_rpc_tasks_and_handlers_e2e.py index 6a99cbb75..f0dd8f757 100644 --- a/python/e2e/test_rpc_tasks_and_handlers_e2e.py +++ b/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()