Fix SDK test completion races and timeout cleanup

dotnet/test/E2E/HooksE2ETests.cs

@@ -32,13 +32,11 @@ public async Task Should_Invoke_PreToolUse_Hook_When_Model_Runs_A_Tool()
         // Create a file for the model to read
         await File.WriteAllTextAsync(Path.Join(Ctx.WorkDir, "hello.txt"), "Hello from the test!");
 
-        await session.SendAsync(new MessageOptions
+        await TestHelper.SendAndGetFinalAssistantMessageAsync(session, new MessageOptions
         {
             Prompt = "Read the contents of hello.txt and tell me what it says"
         });
 
-        await TestHelper.GetFinalAssistantMessageAsync(session);
-
         // Should have received at least one preToolUse hook call
         Assert.NotEmpty(preToolUseInputs);
 
@@ -68,13 +66,11 @@ public async Task Should_Invoke_PostToolUse_Hook_After_Model_Runs_A_Tool()
         // Create a file for the model to read
         await File.WriteAllTextAsync(Path.Join(Ctx.WorkDir, "world.txt"), "World from the test!");
 
-        await session.SendAsync(new MessageOptions
+        await TestHelper.SendAndGetFinalAssistantMessageAsync(session, new MessageOptions
         {
             Prompt = "Read the contents of world.txt and tell me what it says"
         });
 
-        await TestHelper.GetFinalAssistantMessageAsync(session);
-
         // Should have received at least one postToolUse hook call
         Assert.NotEmpty(postToolUseInputs);
 
@@ -109,13 +105,11 @@ public async Task Should_Invoke_Both_PreToolUse_And_PostToolUse_Hooks_For_Single
 
         await File.WriteAllTextAsync(Path.Join(Ctx.WorkDir, "both.txt"), "Testing both hooks!");
 
-        await session.SendAsync(new MessageOptions
+        await TestHelper.SendAndGetFinalAssistantMessageAsync(session, new MessageOptions
         {
             Prompt = "Read the contents of both.txt"
         });
 
-        await TestHelper.GetFinalAssistantMessageAsync(session);
-
         // Both hooks should have been called
         Assert.NotEmpty(preToolUseInputs);
         Assert.NotEmpty(postToolUseInputs);
@@ -149,13 +143,11 @@ public async Task Should_Deny_Tool_Execution_When_PreToolUse_Returns_Deny()
         var originalContent = "Original content that should not be modified";
         await File.WriteAllTextAsync(Path.Join(Ctx.WorkDir, "protected.txt"), originalContent);
 
-        await session.SendAsync(new MessageOptions
+        var response = await TestHelper.SendAndGetFinalAssistantMessageAsync(session, new MessageOptions
         {
             Prompt = "Edit protected.txt and replace 'Original' with 'Modified'"
         });
 
-        var response = await TestHelper.GetFinalAssistantMessageAsync(session);
-
         // The hook should have been called
         Assert.NotEmpty(preToolUseInputs);
 

dotnet/test/E2E/PermissionE2ETests.cs

@@ -111,13 +111,11 @@ public async Task Should_Deny_Permission_When_Handler_Returns_Denied()
         var testFilePath = Path.Combine(Ctx.WorkDir, "protected.txt");
         await File.WriteAllTextAsync(testFilePath, "protected content");
 
-        await session.SendAsync(new MessageOptions
+        await TestHelper.SendAndGetFinalAssistantMessageAsync(session, new MessageOptions
         {
             Prompt = "Edit protected.txt and replace 'protected' with 'hacked'."
         });
 
-        await TestHelper.GetFinalAssistantMessageAsync(session);
-
         Assert.True(
             userRejectedToolCall,
             "Expected a tool.execution_complete event whose error indicates the user rejected the call.");
@@ -159,13 +157,16 @@ await session.SendAndWaitAsync(new MessageOptions
     public async Task Should_Work_With_Approve_All_Permission_Handler()
     {
         var session = await CreateSessionAsync(new SessionConfig());
+        await AssertApproveAllPermissionHandlerAsync(session, TimeSpan.FromSeconds(120));
+    }
 
-        await session.SendAsync(new MessageOptions
+    internal static async Task AssertApproveAllPermissionHandlerAsync(CopilotSession session, TimeSpan timeout)
+    {
+        var message = await TestHelper.SendAndGetFinalAssistantMessageAsync(session, new MessageOptions
         {
             Prompt = "What is 2+2?"
-        });
+        }, timeout);
 
-        var message = await TestHelper.GetFinalAssistantMessageAsync(session);
         Assert.Contains("4", message?.Data.Content ?? string.Empty);
     }
 
@@ -183,13 +184,11 @@ public async Task Should_Handle_Async_Permission_Handler()
             }
         });
 
-        await session.SendAsync(new MessageOptions
+        await TestHelper.SendAndGetFinalAssistantMessageAsync(session, new MessageOptions
         {
             Prompt = "Run 'echo test' and tell me what happens"
         });
 
-        await TestHelper.GetFinalAssistantMessageAsync(session);
-
         Assert.True(permissionRequestReceived, "Permission request should have been received");
     }
 
@@ -321,13 +320,11 @@ public async Task Should_Receive_ToolCallId_In_Permission_Requests()
             }
         });
 
-        await session.SendAsync(new MessageOptions
+        await TestHelper.SendAndGetFinalAssistantMessageAsync(session, new MessageOptions
         {
             Prompt = "Run 'echo test'"
         });
 
-        await TestHelper.GetFinalAssistantMessageAsync(session);
-
         Assert.True(receivedToolCallId, "Should have received toolCallId in permission request");
     }
 

dotnet/test/E2E/SessionE2ETests.cs

@@ -54,8 +54,7 @@ public async Task Should_Create_A_Session_With_Appended_SystemMessage_Config()
             SystemMessage = new SystemMessageConfig { Mode = SystemMessageMode.Append, Content = systemMessageSuffix }
         });
 
-        await session.SendAsync(new MessageOptions { Prompt = "What is your full name?" });
-        var assistantMessage = await TestHelper.GetFinalAssistantMessageAsync(session);
+        var assistantMessage = await TestHelper.SendAndGetFinalAssistantMessageAsync(session, new MessageOptions { Prompt = "What is your full name?" });
         Assert.NotNull(assistantMessage);
 
         var content = assistantMessage!.Data.Content ?? string.Empty;
@@ -78,8 +77,7 @@ public async Task Should_Create_A_Session_With_Replaced_SystemMessage_Config()
             SystemMessage = new SystemMessageConfig { Mode = SystemMessageMode.Replace, Content = testSystemMessage }
         });
 
-        await session.SendAsync(new MessageOptions { Prompt = "What is your full name?" });
-        var assistantMessage = await TestHelper.GetFinalAssistantMessageAsync(session);
+        var assistantMessage = await TestHelper.SendAndGetFinalAssistantMessageAsync(session, new MessageOptions { Prompt = "What is your full name?" });
         Assert.NotNull(assistantMessage);
 
         var content = assistantMessage!.Data.Content ?? string.Empty;
@@ -219,8 +217,7 @@ public async Task Should_Create_Session_With_Custom_Tool()
             ]
         });
 
-        await session.SendAsync(new MessageOptions { Prompt = "What is the secret number for key ALPHA?" });
-        var assistantMessage = await TestHelper.GetFinalAssistantMessageAsync(session);
+        var assistantMessage = await TestHelper.SendAndGetFinalAssistantMessageAsync(session, new MessageOptions { Prompt = "What is the secret number for key ALPHA?" });
         Assert.NotNull(assistantMessage);
         Assert.Contains("54321", assistantMessage!.Data.Content ?? string.Empty);
     }
@@ -463,10 +460,7 @@ public async Task Should_Receive_Session_Events()
         // Events must be dispatched serially — never more than one handler invocation at a time.
         Assert.Equal(1, maxConcurrent);
 
-        // Verify the assistant response contains the expected answer.
-        // session.idle is ephemeral and not in getEvents(), but we already
-        // confirmed idle via the live event handler above.
-        var assistantMessage = await TestHelper.GetFinalAssistantMessageAsync(session, alreadyIdle: true);
+        var assistantMessage = observedEvents.OfType<AssistantMessageEvent>().LastOrDefault();
         Assert.NotNull(assistantMessage);
         Assert.Contains("300", assistantMessage!.Data.Content);
 
@@ -481,8 +475,17 @@ public async Task Send_Returns_Immediately_While_Events_Stream_In_Background()
             OnPermissionRequest = PermissionHandler.ApproveAll,
         });
         var events = new ConcurrentQueue<string>();
+        AssistantMessageEvent? message = null;
 
-        session.On<SessionEvent>(evt => events.Enqueue(evt.Type));
+        session.On<SessionEvent>(evt =>
+        {
+            events.Enqueue(evt.Type);
+            if (evt is AssistantMessageEvent assistantMessage)
+            {
+                message = assistantMessage;
+            }
+        });
+        var idle = TestHelper.GetNextEventOfTypeAsync<SessionIdleEvent>(session);
 
         // Use a slow command so we can verify SendAsync() returns before completion
         await session.SendAsync(new MessageOptions { Prompt = "Run 'sleep 2 && echo done'" });
@@ -491,7 +494,7 @@ public async Task Send_Returns_Immediately_While_Events_Stream_In_Background()
         Assert.DoesNotContain("session.idle", events);
 
         // Wait for turn to complete
-        var message = await TestHelper.GetFinalAssistantMessageAsync(session);
+        await idle;
 
         Assert.Contains("done", message?.Data.Content ?? string.Empty);
         Assert.Contains("session.idle", events);

dotnet/test/E2E/SystemMessageSectionsE2ETests.cs

@@ -31,8 +31,7 @@ public async Task Should_Use_Replaced_Identity_Section_In_Response()
             }
         });
 
-        await session.SendAsync(new MessageOptions { Prompt = "Who are you?" });
-        var response = await TestHelper.GetFinalAssistantMessageAsync(session);
+        var response = await TestHelper.SendAndGetFinalAssistantMessageAsync(session, new MessageOptions { Prompt = "Who are you?" });
 
         Assert.NotNull(response);
         var content = response.Data.Content.ToLowerInvariant();
@@ -61,8 +60,7 @@ public async Task Should_Use_Replaced_Preamble_Section_In_Response()
             }
         });
 
-        await session.SendAsync(new MessageOptions { Prompt = "Who are you?" });
-        var response = await TestHelper.GetFinalAssistantMessageAsync(session);
+        var response = await TestHelper.SendAndGetFinalAssistantMessageAsync(session, new MessageOptions { Prompt = "Who are you?" });
 
         Assert.NotNull(response);
         var content = response.Data.Content.ToLowerInvariant();

dotnet/test/E2E/SystemMessageTransformE2ETests.cs

@@ -48,13 +48,11 @@ public async Task Should_Invoke_Transform_Callbacks_With_Section_Content()
 
         await File.WriteAllTextAsync(Path.Combine(Ctx.WorkDir, "test.txt"), "Hello transform!");
 
-        await session.SendAsync(new MessageOptions
+        await TestHelper.SendAndGetFinalAssistantMessageAsync(session, new MessageOptions
         {
             Prompt = "Read the contents of test.txt and tell me what it says"
         });
 
-        await TestHelper.GetFinalAssistantMessageAsync(session);
-
         Assert.True(identityCallbackInvoked, "Expected identity transform callback to be invoked");
         Assert.True(toneCallbackInvoked, "Expected tone transform callback to be invoked");
     }
@@ -83,13 +81,11 @@ public async Task Should_Apply_Transform_Modifications_To_Section_Content()
 
         await File.WriteAllTextAsync(Path.Combine(Ctx.WorkDir, "hello.txt"), "Hello!");
 
-        await session.SendAsync(new MessageOptions
+        await TestHelper.SendAndGetFinalAssistantMessageAsync(session, new MessageOptions
         {
             Prompt = "Read the contents of hello.txt"
         });
 
-        await TestHelper.GetFinalAssistantMessageAsync(session);
-
         // Verify the transform result was actually applied to the system message
         var traffic = await Ctx.GetExchangesAsync();
         Assert.NotEmpty(traffic);
@@ -128,13 +124,11 @@ public async Task Should_Work_With_Static_Overrides_And_Transforms_Together()
 
         await File.WriteAllTextAsync(Path.Combine(Ctx.WorkDir, "combo.txt"), "Combo test!");
 
-        await session.SendAsync(new MessageOptions
+        await TestHelper.SendAndGetFinalAssistantMessageAsync(session, new MessageOptions
         {
             Prompt = "Read the contents of combo.txt and tell me what it says"
         });
 
-        await TestHelper.GetFinalAssistantMessageAsync(session);
-
         Assert.True(transformCallbackInvoked, "Expected identity transform callback to be invoked");
     }
 }

dotnet/test/E2E/TelemetryExportE2ETests.cs

@@ -40,8 +40,7 @@ public async Task Should_Export_File_Telemetry_For_Sdk_Interactions()
             OnPermissionRequest = PermissionHandler.ApproveAll,
         });
 
-        await session.SendAsync(new MessageOptions { Prompt = prompt });
-        var assistantMessage = await TestHelper.GetFinalAssistantMessageAsync(session);
+        var assistantMessage = await TestHelper.SendAndGetFinalAssistantMessageAsync(session, new MessageOptions { Prompt = prompt });
         Assert.NotNull(assistantMessage);
         Assert.Contains("TELEMETRY_E2E_DONE", assistantMessage!.Data.Content ?? string.Empty, StringComparison.Ordinal);
 

dotnet/test/E2E/ToolResultsE2ETests.cs

@@ -29,12 +29,11 @@ public async Task Should_Handle_Structured_ToolResultObject_From_Custom_Tool()
             OnPermissionRequest = PermissionHandler.ApproveAll,
         });
 
-        await session.SendAsync(new MessageOptions
+        var assistantMessage = await TestHelper.SendAndGetFinalAssistantMessageAsync(session, new MessageOptions
         {
             Prompt = "What's the weather in Paris?"
         });
 
-        var assistantMessage = await TestHelper.GetFinalAssistantMessageAsync(session);
         Assert.NotNull(assistantMessage);
         Assert.Matches("(?i)sunny|72", assistantMessage!.Data.Content ?? string.Empty);
 
@@ -56,12 +55,11 @@ public async Task Should_Handle_Tool_Result_With_Failure_ResultType()
             OnPermissionRequest = PermissionHandler.ApproveAll,
         });
 
-        await session.SendAsync(new MessageOptions
+        var assistantMessage = await TestHelper.SendAndGetFinalAssistantMessageAsync(session, new MessageOptions
         {
             Prompt = "Check the status of the service using check_status. If it fails, say 'service is down'."
         });
 
-        var assistantMessage = await TestHelper.GetFinalAssistantMessageAsync(session);
         Assert.NotNull(assistantMessage);
         Assert.Contains("service is down", assistantMessage!.Data.Content?.ToLowerInvariant() ?? string.Empty);
 
@@ -84,12 +82,11 @@ public async Task Should_Preserve_ToolTelemetry_And_Not_Stringify_Structured_Res
             OnPermissionRequest = PermissionHandler.ApproveAll,
         });
 
-        await session.SendAsync(new MessageOptions
+        var assistantMessage = await TestHelper.SendAndGetFinalAssistantMessageAsync(session, new MessageOptions
         {
             Prompt = "Analyze the file main.ts for issues."
         });
 
-        var assistantMessage = await TestHelper.GetFinalAssistantMessageAsync(session);
         Assert.NotNull(assistantMessage);
         Assert.Contains("no issues", assistantMessage!.Data.Content?.ToLowerInvariant() ?? string.Empty);