Harden MCP OAuth cancel E2E tests against create/interest race (#2029)

* Harden MCP OAuth cancel E2E tests against create/interest race

The MCP OAuth "cancel" E2E tests sampled the host auth callback the instant
the server reached `needs-auth`, which is racy: `session.create` kicks off the
MCP connection, but the SDK only registers its `mcp.oauth_required` event
interest after create returns. When the server's initial 401 wins that race,
the runtime records `needs-auth` without invoking the host callback, so the
callback observation was intermittently empty (e.g. the flaky C# CI leg in
copilot-agent-runtime).

Wait for the callback to be invoked (bounded, reusing each suite's existing
wait-for-condition helper) instead of sampling it immediately. A later runtime
auth retry fires the callback with the same `initial` reason, so the assertions
stay valid. Applied uniformly to C#, Go, Python, Java, and Rust; the Go change
also guards the observed request with a mutex to fix a latent data race.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Copilot-Session: f7849882-a2a2-4587-a602-da7718889a8c

* Use TaskCompletionSource for thread-safe callback wait in C# cancel test

Addresses review feedback: the previous poll read observedRequest (written by the callback thread) from the test continuation without synchronization. Switch to the TaskCompletionSource pattern already used by the direct-RPC test in this file, so the callback result is handed off safely and awaited with a timeout.

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

Copilot-Session: f7849882-a2a2-4587-a602-da7718889a8c

---------

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

Author: stephentoub

dotnet/test/E2E/McpOAuthE2ETests.cs

@@ -204,13 +204,13 @@ public async Task Should_Cancel_Pending_MCP_OAuth_Request()
     {
         await using var oauthServer = await OAuthMcpServer.StartAsync(ExpectedToken);
         var serverName = "oauth-cancelled-mcp";
-        McpAuthContext? observedRequest = null;
+        var authRequest = new TaskCompletionSource<McpAuthContext>(TaskCreationOptions.RunContinuationsAsynchronously);
 
         await using var session = await CreateSessionAsync(new SessionConfig
         {
             OnMcpAuthRequest = request =>
             {
-                observedRequest = request;
+                authRequest.TrySetResult(request);
                 return Task.FromResult<McpAuthResult?>(McpAuthResult.Cancel());
             },
             McpServers = new Dictionary<string, McpServerConfig>
@@ -225,9 +225,16 @@ public async Task Should_Cancel_Pending_MCP_OAuth_Request()
 
         await WaitForMcpServerStatusAsync(session, serverName, McpServerStatus.NeedsAuth);
 
-        Assert.NotNull(observedRequest);
-        Assert.NotEmpty(observedRequest!.RequestId);
-        Assert.Equal(serverName, observedRequest!.ServerName);
+        // The MCP connection is kicked off by session.create, but the SDK only registers its
+        // `mcp.oauth_required` event interest once create returns. If the server's initial 401
+        // wins that race, the runtime records `needs-auth` WITHOUT invoking the host callback,
+        // so the callback fires only on a later auth retry (now that interest is registered),
+        // with the same `Initial` reason. Await the callback rather than sampling it the instant
+        // `needs-auth` first appears, which is what made this test flaky.
+        var observedRequest = await authRequest.Task.WaitAsync(TimeSpan.FromSeconds(60));
+
+        Assert.NotEmpty(observedRequest.RequestId);
+        Assert.Equal(serverName, observedRequest.ServerName);
         Assert.Equal(McpOauthRequestReason.Initial, observedRequest.Reason);
     }
 

go/internal/e2e/mcp_oauth_e2e_test.go

@@ -197,12 +197,17 @@ func TestMCPOAuthE2E(t *testing.T) {
 	t.Run("cancel pending MCP OAuth request", func(t *testing.T) {
 		baseURL := startOAuthMCPServer(t)
 		serverName := "oauth-cancelled-mcp"
+		var mu sync.Mutex
 		var observedRequest copilot.MCPAuthRequest
+		var observed bool
 
 		session, err := client.CreateSession(t.Context(), &copilot.SessionConfig{
 			OnPermissionRequest: copilot.PermissionHandler.ApproveAll,
 			OnMCPAuthRequest: func(request copilot.MCPAuthRequest, _ copilot.MCPAuthInvocation) (*copilot.MCPAuthResult, error) {
+				mu.Lock()
 				observedRequest = request
+				observed = true
+				mu.Unlock()
 				return copilot.MCPAuthResultCancelled(), nil
 			},
 			MCPServers: map[string]copilot.MCPServerConfig{
@@ -218,11 +223,35 @@ func TestMCPOAuthE2E(t *testing.T) {
 		t.Cleanup(func() { session.Disconnect() })
 
 		waitForMCPServerStatus(t, session, serverName, rpc.MCPServerStatusNeedsAuth)
-		if observedRequest.ServerName != serverName {
-			t.Fatalf("Expected serverName %q, got %q", serverName, observedRequest.ServerName)
-		}
-		if observedRequest.Reason != copilot.MCPOauthRequestReasonInitial {
-			t.Fatalf("Unexpected auth request reason: %q", observedRequest.Reason)
+
+		// The MCP connection is kicked off by session.create, but the SDK only registers its
+		// `mcp.oauth_required` event interest once create returns. If the server's initial 401
+		// wins that race, the runtime records `needs-auth` WITHOUT invoking the host callback,
+		// so `observedRequest` is briefly unset even after `needs-auth` is observed. A later
+		// auth retry (now that interest is registered) invokes the callback with the same
+		// `Initial` reason. Wait for the callback rather than sampling it the instant
+		// `needs-auth` first appears, which is what made this test flaky.
+		var request copilot.MCPAuthRequest
+		deadline := time.Now().Add(60 * time.Second)
+		for {
+			mu.Lock()
+			got := observed
+			request = observedRequest
+			mu.Unlock()
+			if got {
+				break
+			}
+			if time.Now().After(deadline) {
+				t.Fatalf("%s OAuth request did not reach the host callback", serverName)
+			}
+			time.Sleep(200 * time.Millisecond)
+		}
+
+		if request.ServerName != serverName {
+			t.Fatalf("Expected serverName %q, got %q", serverName, request.ServerName)
+		}
+		if request.Reason != copilot.MCPOauthRequestReasonInitial {
+			t.Fatalf("Unexpected auth request reason: %q", request.Reason)
 		}
 	})
 

java/src/test/java/com/github/copilot/McpOAuthE2ETest.java

@@ -188,12 +188,18 @@ void testShouldCancelPendingMcpOauthRequest() throws Exception {
                                             .setUrl(oauthServer.url() + "/mcp").setTools(List.of("*")))))
                             .get()) {
                 waitForMcpServerStatus(session, serverName, McpServerStatus.NEEDS_AUTH, observedRequest);
-            }
 
-            var request = observedRequest.get();
-            assertNotNull(request, "MCP auth handler should be invoked");
-            assertEquals(serverName, request.serverName());
-            assertEquals(McpOauthRequestReason.INITIAL, request.reason());
+                // Race: session.create kicks off the MCP connection, but the SDK
+                // registers its `mcp.oauth_required` interest only after create
+                // returns. If the initial 401 wins, the runtime records
+                // `needs-auth` without invoking the host callback. A later auth
+                // retry (interest now registered) fires the callback with the same
+                // INITIAL reason. Wait for the callback instead of sampling it the
+                // instant `needs-auth` appears, which is what made this test flaky.
+                var request = waitForAuthRequest(observedRequest);
+                assertEquals(serverName, request.serverName());
+                assertEquals(McpOauthRequestReason.INITIAL, request.reason());
+            }
         }
     }
 

python/e2e/test_mcp_oauth_e2e.py

@@ -326,6 +326,20 @@ def on_mcp_auth_request(request, _invocation):
             ) as session:
                 await _wait_for_mcp_server_status(session, server_name, McpServerStatus.NEEDS_AUTH)
 
+                # The MCP connection is kicked off by session.create, but the SDK only registers
+                # its `mcp.oauth_required` event interest once create returns. If the server's
+                # initial 401 wins that race, the runtime records `needs-auth` WITHOUT invoking
+                # the host callback, so `observed_request` is briefly None even after `needs-auth`
+                # is observed. A later auth retry (now that interest is registered) invokes the
+                # callback with the same `initial` reason. Wait for the callback rather than
+                # sampling it the instant `needs-auth` first appears, which made this test flaky.
+                await wait_for_condition(
+                    lambda: observed_request is not None,
+                    timeout=60.0,
+                    poll_interval=0.2,
+                    timeout_message=f"{server_name} OAuth request did not reach the host callback",
+                )
+
             assert observed_request is not None
             assert observed_request["serverName"] == server_name
             assert observed_request["reason"] == "initial"

rust/tests/e2e/mcp_oauth.rs

@@ -209,6 +209,18 @@ async fn should_cancel_pending_mcp_oauth_request() {
 
             wait_for_mcp_server_status(&session, server_name, McpServerStatus::NeedsAuth).await;
 
+            // The MCP connection is kicked off by session.create, but the SDK only registers its
+            // `mcp.oauth_required` event interest once create returns. If the server's initial 401
+            // wins that race, the runtime records `needs-auth` WITHOUT invoking the host callback,
+            // so `handler.request` is briefly `None` even after `needs-auth` is observed. A later
+            // auth retry (now that interest is registered) invokes the callback with the same
+            // `Initial` reason. Wait for the callback rather than sampling it the instant
+            // `needs-auth` first appears, which is what made this test flaky.
+            wait_for_condition("MCP OAuth request reaching the host callback", || async {
+                handler.request.lock().is_some()
+            })
+            .await;
+
             let request = handler
                 .request
                 .lock()