Fix flaky E2E tests

Make permission handler error coverage assert deterministic replayed tool results instead of waiting for final assistant text, and ensure background-agent tests wait for the completion notification before cleanup. Normalize equivalent replay proxy notification wording across SDK suites.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: stephentoub

dotnet/test/E2E/PermissionE2ETests.cs

@@ -226,24 +226,46 @@ await session2.SendAndWaitAsync(new MessageOptions
     [Fact]
     public async Task Should_Handle_Permission_Handler_Errors_Gracefully()
     {
+        var permissionRequestReceived =
+            new TaskCompletionSource<PermissionRequest>(TaskCreationOptions.RunContinuationsAsynchronously);
         var session = await CreateSessionAsync(new SessionConfig
         {
             OnPermissionRequest = (request, invocation) =>
             {
-                // Simulate an error in the handler
+                permissionRequestReceived.TrySetResult(request);
                 throw new InvalidOperationException("Handler error");
             }
         });
 
-        await session.SendAsync(new MessageOptions
+        try
         {
-            Prompt = "Run 'echo test'. If you can't, say 'failed'."
-        });
-
-        var message = await TestHelper.GetFinalAssistantMessageAsync(session);
-
-        // Should handle the error and deny permission
-        Assert.Matches("fail|cannot|unable|permission", message?.Data.Content?.ToLowerInvariant() ?? string.Empty);
+            var exchanges = await SendAndWaitForExchangesAsync(
+                session,
+                new MessageOptions
+                {
+                    Prompt = "Run 'echo test'. If you can't, say 'failed'."
+                },
+                minimumCount: 2);
+
+            var permissionRequest = await permissionRequestReceived.Task.WaitAsync(TimeSpan.FromSeconds(30));
+            Assert.IsType<PermissionRequestShell>(permissionRequest);
+
+            var toolResultMessage = exchanges
+                .SelectMany(exchange => exchange.Request.Messages)
+                .LastOrDefault(message =>
+                    message.Role == "tool" &&
+                    message.StringContent?.Contains("Permission denied", StringComparison.OrdinalIgnoreCase) == true);
+
+            Assert.NotNull(toolResultMessage);
+            Assert.Contains(
+                "could not request permission",
+                toolResultMessage.StringContent ?? string.Empty,
+                StringComparison.OrdinalIgnoreCase);
+        }
+        finally
+        {
+            await session.DisposeAsync();
+        }
     }
 
     [Fact]

dotnet/test/E2E/RpcTasksAndHandlersE2ETests.cs

@@ -80,6 +80,22 @@ public async Task Should_Start_Background_Agent_And_Report_Task_Details()
         });
         Assert.Contains("TASK_AGENT_READY", ready?.Data.Content ?? string.Empty, StringComparison.Ordinal);
 
+        var taskCompletionNotification =
+            new TaskCompletionSource<AssistantMessageEvent>(TaskCreationOptions.RunContinuationsAsynchronously);
+        using var subscription = session.On<SessionEvent>(evt =>
+        {
+            switch (evt)
+            {
+                case AssistantMessageEvent assistantMessage
+                    when assistantMessage.Data.Content?.Contains("TASK_AGENT_DONE", StringComparison.Ordinal) == true:
+                    taskCompletionNotification.TrySetResult(assistantMessage);
+                    break;
+                case SessionErrorEvent error:
+                    taskCompletionNotification.TrySetException(new Exception(error.Data.Message ?? "session error"));
+                    break;
+            }
+        });
+
         var started = await session.Rpc.Tasks.StartAgentAsync(
             agentType: "general-purpose",
             prompt: "Reply with TASK_AGENT_DONE exactly.",
@@ -123,6 +139,7 @@ await TestHelper.WaitForConditionAsync(
 
         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)
         {

python/e2e/test_rpc_tasks_and_handlers_e2e.py

@@ -29,6 +29,7 @@
     UIElicitationResponseAction,
     UIHandlePendingElicitationRequest,
 )
+from copilot.generated.session_events import AssistantMessageData, SessionErrorData
 from copilot.session import PermissionHandler
 
 from .testharness import E2ETestContext
@@ -211,6 +212,21 @@ async def test_should_start_background_agent_and_report_task_details(self, ctx:
         session = await ctx.client.create_session(
             on_permission_request=PermissionHandler.approve_all,
         )
+        task_completion_notification = asyncio.get_running_loop().create_future()
+
+        def on_event(event):
+            if isinstance(event.data, AssistantMessageData) and "TASK_AGENT_DONE" in (
+                event.data.content or ""
+            ):
+                if not task_completion_notification.done():
+                    task_completion_notification.set_result(event)
+            elif isinstance(event.data, SessionErrorData):
+                if not task_completion_notification.done():
+                    task_completion_notification.set_exception(
+                        RuntimeError(event.data.message or "session error")
+                    )
+
+        unsubscribe = session.on(on_event)
         try:
             ready = await session.send_and_wait(
                 "Reply with TASK_AGENT_READY exactly.",
@@ -262,6 +278,7 @@ async def test_should_start_background_agent_and_report_task_details(self, ctx:
             )
             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))
@@ -274,4 +291,5 @@ async def test_should_start_background_agent_and_report_task_details(self, ctx:
             after_remove = await session.rpc.tasks.list()
             assert not any(t.id == task_id for t in (after_remove.tasks or []))
         finally:
+            unsubscribe()
             await session.disconnect()

test/harness/replayingCapiProxy.test.ts

@@ -324,6 +324,38 @@ describe("ReplayingCapiProxy", () => {
     expect(result.conversations[0].messages[0].content).toBe("What is 2+2?");
   });
 
+  test("normalizes task completion notification wording", async () => {
+    const unreadNotification = [
+      "<system_notification>",
+      'Agent "sdk-background-agent" (general-purpose) has completed successfully. Use read_agent with agent_id "sdk-background-agent" to retrieve unread results.',
+      "</system_notification>",
+    ].join("\n");
+    const fullNotification = [
+      "<system_notification>",
+      'Agent "sdk-background-agent" (general-purpose) has completed successfully. Use read_agent with agent_id "sdk-background-agent" to retrieve the full results.',
+      "</system_notification>",
+    ].join("\n");
+
+    const requestBody = JSON.stringify({
+      messages: [
+        {
+          role: "user",
+          content: unreadNotification,
+        },
+      ],
+    });
+    const responseBody = JSON.stringify({
+      choices: [{ message: { role: "assistant", content: "Done" } }],
+    });
+
+    const outputPath = await createProxy([
+      { url: "/chat/completions", requestBody, responseBody },
+    ]);
+
+    const result = await readYamlOutput(outputPath);
+    expect(result.conversations[0].messages[0].content).toBe(fullNotification);
+  });
+
   test("strips agent_instructions from user messages", async () => {
     const requestBody = JSON.stringify({
       messages: [

test/harness/replayingCapiProxy.ts

@@ -1034,6 +1034,10 @@ function transformOpenAIRequestMessage(
 
 function normalizeUserMessage(content: string): string {
   return normalizeSkillContextFrontmatter(content)
+    .replace(
+      /Use read_agent with agent_id "([^"]+)" to retrieve unread results\./g,
+      'Use read_agent with agent_id "$1" to retrieve the full results.',
+    )
     .replace(/<current_datetime>.*?<\/current_datetime>/g, "")
     .replace(/<reminder>[\s\S]*?<\/reminder>/g, "")
     .replace(/<system_reminder>[\s\S]*?<\/system_reminder>/g, "")