Fix flaky SDK E2E tests

dotnet/test/E2E/PermissionE2ETests.cs

@@ -49,14 +49,14 @@ public async Task Should_Invoke_Permission_Handler_For_Write_Operations()
 
         await File.WriteAllTextAsync(Path.Combine(Ctx.WorkDir, "test.txt"), "original content");
 
-        await session.SendAsync(new MessageOptions
+        var sendTask = session.SendAndWaitAsync(new MessageOptions
         {
             Prompt = "Edit test.txt and replace 'original' with 'modified'"
         });
 
         var readRequest = await readPermissionRequestReceived.Task.WaitAsync(TimeSpan.FromSeconds(30));
         var writeRequest = await writePermissionRequestReceived.Task.WaitAsync(TimeSpan.FromSeconds(30));
-        await TestHelper.GetFinalAssistantMessageAsync(session);
+        await sendTask;
 
         List<PermissionRequest> observedPermissionRequests;
         lock (permissionRequestsLock)

dotnet/test/E2E/PreMcpToolCallHookE2ETests.cs

@@ -15,7 +15,7 @@ namespace GitHub.Copilot.Test.E2E;
 public class PreMcpToolCallHookE2ETests(E2ETestFixture fixture, ITestOutputHelper output)
     : E2ETestBase(fixture, "pre_mcp_tool_call_hook", output)
 {
-    private static string FindTestHarnessDir()
+    private static string FindMetaEchoTestHarnessDir()
     {
         var dir = new DirectoryInfo(AppContext.BaseDirectory);
         while (dir != null)
@@ -42,7 +42,7 @@ private static string FindTestHarnessDir()
     [Fact]
     public async Task Should_Set_Meta_Via_PreMcpToolCall_Hook()
     {
-        var testHarnessDir = FindTestHarnessDir();
+        var testHarnessDir = FindMetaEchoTestHarnessDir();
         var hookInputs = new List<PreMcpToolCallHookInput>();
 
         var session = await CreateSessionAsync(new SessionConfig
@@ -84,7 +84,7 @@ public async Task Should_Set_Meta_Via_PreMcpToolCall_Hook()
     [Fact]
     public async Task Should_Replace_Meta_Via_PreMcpToolCall_Hook()
     {
-        var testHarnessDir = FindTestHarnessDir();
+        var testHarnessDir = FindMetaEchoTestHarnessDir();
         var hookInputs = new List<PreMcpToolCallHookInput>();
 
         var session = await CreateSessionAsync(new SessionConfig
@@ -125,7 +125,7 @@ public async Task Should_Replace_Meta_Via_PreMcpToolCall_Hook()
     [Fact]
     public async Task Should_Remove_Meta_Via_PreMcpToolCall_Hook()
     {
-        var testHarnessDir = FindTestHarnessDir();
+        var testHarnessDir = FindMetaEchoTestHarnessDir();
         var hookInputs = new List<PreMcpToolCallHookInput>();
 
         var session = await CreateSessionAsync(new SessionConfig

dotnet/test/E2E/RpcMcpAndSkillsE2ETests.cs

@@ -2,6 +2,7 @@
  *  Copyright (c) Microsoft Corporation. All rights reserved.
  *--------------------------------------------------------------------------------------------*/
 
+using GitHub.Copilot.Rpc;
 using Xunit;
 using Xunit.Abstractions;
 using RpcSkill = GitHub.Copilot.Rpc.Skill;
@@ -67,21 +68,14 @@ public async Task Should_List_Mcp_Servers_With_Configured_Server()
         const string serverName = "rpc-list-mcp-server";
         var session = await CreateSessionAsync(new SessionConfig
         {
-            McpServers = new Dictionary<string, McpServerConfig>
-            {
-                [serverName] = new McpStdioServerConfig
-                {
-                    Command = "echo",
-                    Args = ["rpc-list-mcp-server"],
-                    Tools = ["*"],
-                },
-            },
+            McpServers = CreateTestMcpServers(serverName),
         });
 
+        await WaitForMcpServerStatusAsync(session, serverName, McpServerStatus.Connected);
         var result = await session.Rpc.Mcp.ListAsync();
 
         var server = Assert.Single(result.Servers, server => string.Equals(server.Name, serverName, StringComparison.Ordinal));
-        Assert.False(string.IsNullOrWhiteSpace(server.Status.Value));
+        Assert.Equal(McpServerStatus.Connected, server.Status);
     }
 
     [Fact]
@@ -143,16 +137,9 @@ public async Task Should_Report_Error_When_Mcp_Oauth_Server_Is_Not_Configured()
     {
         var session = await CreateSessionAsync(new SessionConfig
         {
-            McpServers = new Dictionary<string, McpServerConfig>
-            {
-                ["configured-stdio-server"] = new McpStdioServerConfig
-                {
-                    Command = "echo",
-                    Args = ["configured-stdio-server"],
-                    Tools = ["*"],
-                },
-            },
+            McpServers = CreateTestMcpServers("configured-stdio-server"),
         });
+        await WaitForMcpServerStatusAsync(session, "configured-stdio-server", McpServerStatus.Connected);
 
         await AssertFailureAsync(
             () => session.Rpc.Mcp.Oauth.LoginAsync("missing-server"),
@@ -165,16 +152,9 @@ public async Task Should_Report_Error_When_Mcp_Oauth_Server_Is_Not_Remote()
         const string serverName = "configured-stdio-server";
         var session = await CreateSessionAsync(new SessionConfig
         {
-            McpServers = new Dictionary<string, McpServerConfig>
-            {
-                [serverName] = new McpStdioServerConfig
-                {
-                    Command = "echo",
-                    Args = [serverName],
-                    Tools = ["*"],
-                },
-            },
+            McpServers = CreateTestMcpServers(serverName),
         });
+        await WaitForMcpServerStatusAsync(session, serverName, McpServerStatus.Connected);
 
         await AssertFailureAsync(
             () => session.Rpc.Mcp.Oauth.LoginAsync(serverName, forceReauth: true, clientName: "SDK E2E", callbackSuccessMessage: "Done"),

dotnet/test/E2E/SessionConfigE2ETests.cs

@@ -435,12 +435,17 @@ public async Task Should_Apply_AvailableTools_On_Session_Resume()
             AvailableTools = ["view"],
         });
 
-        await session2.SendAndWaitAsync(new MessageOptions { Prompt = "What is 1+1?" });
-
-        var exchange = Assert.Single(await Ctx.GetExchangesAsync());
-        Assert.Equal(["view"], GetToolNames(exchange));
-
-        await session2.DisposeAsync();
+        try
+        {
+            var exchange = Assert.Single(await SendAndWaitForExchangesAsync(
+                session2,
+                new MessageOptions { Prompt = "What is 1+1?" }));
+            Assert.Equal(["view"], GetToolNames(exchange));
+        }
+        finally
+        {
+            await session2.DisposeAsync();
+        }
     }
 
     [Fact]

dotnet/test/E2E/SessionE2ETests.cs

@@ -130,16 +130,20 @@ public async Task Should_Create_A_Session_With_AvailableTools()
             AvailableTools = ["view", "edit"]
         });
 
-        await session.SendAsync(new MessageOptions { Prompt = "What is 1+1?" });
-        await TestHelper.GetFinalAssistantMessageAsync(session);
-
-        var traffic = await Ctx.GetExchangesAsync();
-        Assert.NotEmpty(traffic);
-
-        var toolNames = GetToolNames(traffic[0]);
-        Assert.Equal(2, toolNames.Count);
-        Assert.Contains("view", toolNames);
-        Assert.Contains("edit", toolNames);
+        try
+        {
+            var traffic = await SendAndWaitForExchangesAsync(
+                session,
+                new MessageOptions { Prompt = "What is 1+1?" });
+            var toolNames = GetToolNames(traffic[0]);
+            Assert.Equal(2, toolNames.Count);
+            Assert.Contains("view", toolNames);
+            Assert.Contains("edit", toolNames);
+        }
+        finally
+        {
+            await session.DisposeAsync();
+        }
     }
 
     [Fact]
@@ -150,16 +154,20 @@ public async Task Should_Create_A_Session_With_ExcludedTools()
             ExcludedTools = ["view"]
         });
 
-        await session.SendAsync(new MessageOptions { Prompt = "What is 1+1?" });
-        await TestHelper.GetFinalAssistantMessageAsync(session);
-
-        var traffic = await Ctx.GetExchangesAsync();
-        Assert.NotEmpty(traffic);
-
-        var toolNames = GetToolNames(traffic[0]);
-        Assert.DoesNotContain("view", toolNames);
-        Assert.Contains("edit", toolNames);
-        Assert.Contains("grep", toolNames);
+        try
+        {
+            var traffic = await SendAndWaitForExchangesAsync(
+                session,
+                new MessageOptions { Prompt = "What is 1+1?" });
+            var toolNames = GetToolNames(traffic[0]);
+            Assert.DoesNotContain("view", toolNames);
+            Assert.Contains("edit", toolNames);
+            Assert.Contains("grep", toolNames);
+        }
+        finally
+        {
+            await session.DisposeAsync();
+        }
     }
 
     [Fact]
@@ -180,14 +188,18 @@ public async Task Should_Create_A_Session_With_DefaultAgent_ExcludedTools()
             },
         });
 
-        await session.SendAsync(new MessageOptions { Prompt = "What is 1+1?" });
-        await TestHelper.GetFinalAssistantMessageAsync(session);
-
-        var traffic = await Ctx.GetExchangesAsync();
-        Assert.NotEmpty(traffic);
-
-        var toolNames = GetToolNames(traffic[0]);
-        Assert.DoesNotContain("secret_tool", toolNames);
+        try
+        {
+            var traffic = await SendAndWaitForExchangesAsync(
+                session,
+                new MessageOptions { Prompt = "What is 1+1?" });
+            var toolNames = GetToolNames(traffic[0]);
+            Assert.DoesNotContain("secret_tool", toolNames);
+        }
+        finally
+        {
+            await session.DisposeAsync();
+        }
     }
 
     [Fact]
@@ -539,11 +551,17 @@ public async Task Should_Create_Session_With_Custom_Config_Dir()
 
         Assert.Matches(@"^[a-f0-9-]+$", session.SessionId);
 
-        // Session should work normally with custom config dir
-        await session.SendAsync(new MessageOptions { Prompt = "What is 1+1?" });
-        var assistantMessage = await TestHelper.GetFinalAssistantMessageAsync(session);
-        Assert.NotNull(assistantMessage);
-        Assert.Contains("2", assistantMessage!.Data.Content);
+        try
+        {
+            // Session should work normally with custom config dir.
+            var assistantMessage = await session.SendAndWaitAsync(new MessageOptions { Prompt = "What is 1+1?" });
+            Assert.NotNull(assistantMessage);
+            Assert.Contains("2", assistantMessage!.Data.Content);
+        }
+        finally
+        {
+            await session.DisposeAsync();
+        }
     }
 
     [Fact]