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

Author: stephentoub

dotnet/test/E2E/McpOAuthE2ETests.cs

@@ -225,6 +225,19 @@ public async Task Should_Cancel_Pending_MCP_OAuth_Request()
 
         await WaitForMcpServerStatusAsync(session, serverName, McpServerStatus.NeedsAuth);
 
+        // 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 null 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.
+        await TestHelper.WaitForConditionAsync(
+            () => Task.FromResult(observedRequest is not null),
+            timeout: TimeSpan.FromSeconds(60),
+            pollInterval: TimeSpan.FromMilliseconds(200),
+            timeoutMessage: $"{serverName} OAuth request reaching the host callback");
+
         Assert.NotNull(observedRequest);
         Assert.NotEmpty(observedRequest!.RequestId);
         Assert.Equal(serverName, observedRequest!.ServerName);

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()