From 480979b1b592a4c9a768c1918f994f025f72fd36 Mon Sep 17 00:00:00 2001
From: Copilot <198982749+Copilot@users.noreply.github.com>
Date: Fri, 5 Jun 2026 16:37:23 -0700
Subject: [PATCH] Fix flaky DiagnosticTests.Session_TracksActivities by waiting
 for full server activity predicate (#1495)

---
 .../DiagnosticTests.cs                        | 41 +++++++++++++------
 1 file changed, 28 insertions(+), 13 deletions(-)

diff --git a/tests/ModelContextProtocol.Tests/DiagnosticTests.cs b/tests/ModelContextProtocol.Tests/DiagnosticTests.cs
index bbe7d15..03e00af 100644
--- a/tests/ModelContextProtocol.Tests/DiagnosticTests.cs
+++ b/tests/ModelContextProtocol.Tests/DiagnosticTests.cs
@@ -18,6 +18,18 @@ public class DiagnosticTests
         var activities = new List<Activity>();
         var clientToServerLog = new List<string>();
 
+        // Predicate for the expected server tool-call activity, including all required tags.
+        // Defined here so it can be reused for both the wait and the assertion below.
+        Func<Activity, bool> isExpectedServerToolCall = a =>
+            a.DisplayName == "tools/call DoubleValue" &&
+            a.Kind == ActivityKind.Server &&
+            a.Status == ActivityStatusCode.Unset &&
+            a.Tags.Any(t => t.Key == "gen_ai.tool.name" && t.Value == "DoubleValue") &&
+            a.Tags.Any(t => t.Key == "mcp.method.name" && t.Value == "tools/call") &&
+            a.Tags.Any(t => t.Key == "gen_ai.operation.name" && t.Value == "execute_tool") &&
+            a.Tags.Any(t => t.Key == "mcp.protocol.version" && !string.IsNullOrEmpty(t.Value)) &&
+            a.Tags.Any(t => t.Key == "mcp.session.id" && !string.IsNullOrEmpty(t.Value));
+
         using (var tracerProvider = OpenTelemetry.Sdk.CreateTracerProviderBuilder()
             .AddSource("Experimental.ModelContextProtocol")
             .AddInMemoryExporter(activities)
@@ -36,9 +48,11 @@ public class DiagnosticTests
             // Wait for server-side activities to be exported. The server processes messages
             // via fire-and-forget tasks, so activities may not be immediately available
             // after the client operation completes. Wait for the specific activity we need
-            // rather than a count, as other server activities may be exported first.
-            await WaitForAsync(() => activities.Any(a =>
-                a.DisplayName == "tools/call DoubleValue" && a.Kind == ActivityKind.Server));
+            // (including required tags) rather than just the display name, so that we don't
+            // assert before all tags have been populated.
+            await WaitForAsync(
+                () => activities.Any(isExpectedServerToolCall),
+                failureMessage: "Timed out waiting for the expected server tool-call activity (tools/call DoubleValue) to be exported with required tags.");
         }
 
         Assert.NotEmpty(activities);
@@ -54,13 +68,7 @@ public class DiagnosticTests
         // Per semantic conventions: mcp.protocol.version should be present after initialization
         Assert.Contains(clientToolCall.Tags, t => t.Key == "mcp.protocol.version" && !string.IsNullOrEmpty(t.Value));
 
-        var serverToolCall = Assert.Single(activities, a =>
-            a.Tags.Any(t => t.Key == "gen_ai.tool.name" && t.Value == "DoubleValue") &&
-            a.Tags.Any(t => t.Key == "mcp.method.name" && t.Value == "tools/call") &&
-            a.Tags.Any(t => t.Key == "gen_ai.operation.name" && t.Value == "execute_tool") &&
-            a.DisplayName == "tools/call DoubleValue" &&
-            a.Kind == ActivityKind.Server &&
-            a.Status == ActivityStatusCode.Unset);
+        var serverToolCall = Assert.Single(activities, a => isExpectedServerToolCall(a));
 
         // Per semantic conventions: mcp.protocol.version should be present after initialization
         Assert.Contains(serverToolCall.Tags, t => t.Key == "mcp.protocol.version" && !string.IsNullOrEmpty(t.Value));
@@ -245,12 +253,19 @@ public class DiagnosticTests
         await serverTask;
     }
 
-    private static async Task WaitForAsync(Func<bool> condition, int timeoutMs = 10_000)
+    private static async Task WaitForAsync(Func<bool> condition, int timeoutMs = 10_000, string? failureMessage = null)
     {
         using var cts = new CancellationTokenSource(timeoutMs);
-        while (!condition())
+        try
         {
-            await Task.Delay(10, cts.Token);
+            while (!condition())
+            {
+                await Task.Delay(10, cts.Token);
+            }
+        }
+        catch (TaskCanceledException)
+        {
+            throw new Xunit.Sdk.XunitException(failureMessage ?? $"Condition was not met within {timeoutMs}ms.");
         }
     }
 }