Fix Session_TracksActivities test (#1200)

Author: stephentoub

tests/ModelContextProtocol.Tests/DiagnosticTests.cs

@@ -32,6 +32,11 @@ await RunConnected(async (client, server) =>
                 var tool = tools.First(t => t.Name == "DoubleValue");
                 await tool.InvokeAsync(new() { ["amount"] = 42 }, TestContext.Current.CancellationToken);
             }, clientToServerLog);
+
+            // 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.
+            await WaitForAsync(() => activities.Count(a => a.Kind == ActivityKind.Server) >= 4);
         }
 
         Assert.NotEmpty(activities);
@@ -97,6 +102,9 @@ await RunConnected(async (client, server) =>
                 await client.CallToolAsync("Throw", cancellationToken: TestContext.Current.CancellationToken);
                 await Assert.ThrowsAsync<McpProtocolException>(async () => await client.CallToolAsync("does-not-exist", cancellationToken: TestContext.Current.CancellationToken));
             }, []);
+
+            // Wait for server-side activities to be exported.
+            await WaitForAsync(() => activities.Count(a => a.Kind == ActivityKind.Server) >= 4);
         }
 
         Assert.NotEmpty(activities);
@@ -164,6 +172,9 @@ await RunConnected(async (client, server) =>
                     .First(t => t.Name == "DoubleValue");
                 await tool.InvokeAsync(new() { ["amount"] = 42 }, TestContext.Current.CancellationToken);
             }, []);
+
+            // Wait for server-side activities to be exported.
+            await WaitForAsync(() => activities.Count(a => a.Kind == ActivityKind.Server) >= 3);
         }
 
         // The outer activity should have MCP-specific attributes added to it
@@ -215,6 +226,15 @@ private static async Task RunConnected(Func<McpClient, McpServer, Task> action,
 
         await serverTask;
     }
+
+    private static async Task WaitForAsync(Func<bool> condition, int timeoutMs = 10_000)
+    {
+        using var cts = new CancellationTokenSource(timeoutMs);
+        while (!condition())
+        {
+            await Task.Delay(10, cts.Token);
+        }
+    }
 }
 
 public class LoggingStream : Stream