From 4ae594ba326b42b4ba95e3fdb62acf3887d6324f Mon Sep 17 00:00:00 2001 From: Stephen Toub Date: Fri, 17 Jul 2026 23:00:17 -0400 Subject: [PATCH 1/2] 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 --- dotnet/test/E2E/McpOAuthE2ETests.cs | 13 +++++++ go/internal/e2e/mcp_oauth_e2e_test.go | 39 ++++++++++++++++--- .../com/github/copilot/McpOAuthE2ETest.java | 16 +++++--- python/e2e/test_mcp_oauth_e2e.py | 14 +++++++ rust/tests/e2e/mcp_oauth.rs | 12 ++++++ 5 files changed, 84 insertions(+), 10 deletions(-) diff --git a/dotnet/test/E2E/McpOAuthE2ETests.cs b/dotnet/test/E2E/McpOAuthE2ETests.cs index f101bbc31b..a29897bdea 100644 --- a/dotnet/test/E2E/McpOAuthE2ETests.cs +++ b/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); diff --git a/go/internal/e2e/mcp_oauth_e2e_test.go b/go/internal/e2e/mcp_oauth_e2e_test.go index 1e08b839c8..95de73eddd 100644 --- a/go/internal/e2e/mcp_oauth_e2e_test.go +++ b/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) } }) diff --git a/java/src/test/java/com/github/copilot/McpOAuthE2ETest.java b/java/src/test/java/com/github/copilot/McpOAuthE2ETest.java index e721468c00..f234337eaf 100644 --- a/java/src/test/java/com/github/copilot/McpOAuthE2ETest.java +++ b/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()); + } } } diff --git a/python/e2e/test_mcp_oauth_e2e.py b/python/e2e/test_mcp_oauth_e2e.py index 8322a3aba1..9d70597c3e 100644 --- a/python/e2e/test_mcp_oauth_e2e.py +++ b/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" diff --git a/rust/tests/e2e/mcp_oauth.rs b/rust/tests/e2e/mcp_oauth.rs index 5a11ef4b7a..fb202536c7 100644 --- a/rust/tests/e2e/mcp_oauth.rs +++ b/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() From 15a1bad2fc4373a4f196d2f6913773cfc517abdb Mon Sep 17 00:00:00 2001 From: Stephen Toub Date: Fri, 17 Jul 2026 23:06:42 -0400 Subject: [PATCH 2/2] 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 --- dotnet/test/E2E/McpOAuthE2ETests.cs | 20 +++++++------------- 1 file changed, 7 insertions(+), 13 deletions(-) diff --git a/dotnet/test/E2E/McpOAuthE2ETests.cs b/dotnet/test/E2E/McpOAuthE2ETests.cs index a29897bdea..1085aba040 100644 --- a/dotnet/test/E2E/McpOAuthE2ETests.cs +++ b/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(TaskCreationOptions.RunContinuationsAsynchronously); await using var session = await CreateSessionAsync(new SessionConfig { OnMcpAuthRequest = request => { - observedRequest = request; + authRequest.TrySetResult(request); return Task.FromResult(McpAuthResult.Cancel()); }, McpServers = new Dictionary @@ -228,19 +228,13 @@ public async Task Should_Cancel_Pending_MCP_OAuth_Request() // 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 + // 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. - 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"); + var observedRequest = await authRequest.Task.WaitAsync(TimeSpan.FromSeconds(60)); - Assert.NotNull(observedRequest); - Assert.NotEmpty(observedRequest!.RequestId); - Assert.Equal(serverName, observedRequest!.ServerName); + Assert.NotEmpty(observedRequest.RequestId); + Assert.Equal(serverName, observedRequest.ServerName); Assert.Equal(McpOauthRequestReason.Initial, observedRequest.Reason); }