Add initial Activity / Metric support

- Creates an Activity around every outgoing send / incoming receive handling
- Creates a histogram for each incoming operation, tagged with request/notification details

Author: stephentoub

Directory.Packages.props

@@ -33,6 +33,8 @@
     <PackageVersion Include="Microsoft.Extensions.Logging.Console" Version="$(MicrosoftExtensionsVersion)" />
     <PackageVersion Include="Microsoft.NET.Test.Sdk" Version="17.12.0" />
     <PackageVersion Include="Moq" Version="4.20.72" />
+    <PackageVersion Include="OpenTelemetry" Version="1.11.2" />
+    <PackageVersion Include="OpenTelemetry.Exporter.InMemory" Version="1.11.2" />
     <PackageVersion Include="Serilog.Extensions.Hosting" Version="9.0.0" />
     <PackageVersion Include="Serilog.Extensions.Logging" Version="9.0.0" />
     <PackageVersion Include="Serilog.Sinks.Console" Version="6.0.0" />

src/ModelContextProtocol/Shared/McpSession.cs

@@ -7,6 +7,8 @@
 using ModelContextProtocol.Utils;
 using ModelContextProtocol.Utils.Json;
 using System.Collections.Concurrent;
+using System.Diagnostics;
+using System.Diagnostics.Metrics;
 using System.Text.Json;
 
 namespace ModelContextProtocol.Shared;
@@ -16,6 +18,13 @@ namespace ModelContextProtocol.Shared;
 /// </summary>
 internal sealed class McpSession : IDisposable
 {
+    private static readonly ActivitySource s_activitySource = new("ModelContextProtocol");
+    private static readonly Meter s_meter = new("ModelContextProtocol");
+    private static readonly Histogram<double> s_operationDurationHistogram = s_meter.CreateHistogram<double>(
+        "modelcontextprotocol.operation.duration",
+        "s",
+        "Measures the duration of an operation in seconds.");
+
     private readonly ITransport _transport;
     private readonly RequestHandlers _requestHandlers;
     private readonly NotificationHandlers _notificationHandlers;
@@ -152,23 +161,89 @@ await _transport.SendMessageAsync(new JsonRpcError
 
     private async Task HandleMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken)
     {
-        switch (message)
+        using Activity? activity = s_activitySource.StartActivity("HandlingMessage");
+        TagList tags = default;
+
+        Stopwatch? measuring = s_operationDurationHistogram.Enabled ? Stopwatch.StartNew() : null;
+        bool addTags = activity is not null || measuring is not null;
+
+        try
         {
-            case JsonRpcRequest request:
-                await HandleRequest(request, cancellationToken).ConfigureAwait(false);
-                break;
+            if (addTags)
+            {
+                tags.Add("mcp.session.id", _id);
+            }
+
+            switch (message)
+            {
+                case JsonRpcRequest request:
+                    if (addTags)
+                    {
+                        tags.Add("mcp.request.id", request.Id.ToString());
+                        tags.Add("mcp.request.method", request.Method);
+
+                        if (request.Params is JsonElement je)
+                        {
+                            switch (request.Method)
+                            {
+                                case RequestMethods.ToolsCall:
+                                case RequestMethods.PromptsGet:
+                                    if (je.TryGetProperty("name", out var prop) && prop.ValueKind == JsonValueKind.String)
+                                    {
+                                        tags.Add("mcp.request.params.name", prop.GetString());
+                                    }
+                                    break;
+
+                                case RequestMethods.ResourcesRead:
+                                    if (je.TryGetProperty("uri", out prop) && prop.ValueKind == JsonValueKind.String)
+                                    {
+                                        tags.Add("mcp.request.params.uri", prop.GetString());
+                                    }
+                                    break;
+                            }
+                        }
+                    }
+
+                    await HandleRequest(request, cancellationToken).ConfigureAwait(false);
+                    break;
 
-            case IJsonRpcMessageWithId messageWithId:
-                HandleMessageWithId(message, messageWithId);
-                break;
+                case JsonRpcNotification notification:
+                    if (addTags)
+                    {
+                        tags.Add("mcp.notification.method", notification.Method);
+                    }
+
+                    await HandleNotification(notification).ConfigureAwait(false);
+                    break;
 
-            case JsonRpcNotification notification:
-                await HandleNotification(notification).ConfigureAwait(false);
-                break;
+                case IJsonRpcMessageWithId messageWithId:
+                    HandleMessageWithId(message, messageWithId);
+                    break;
 
-            default:
-                _logger.EndpointHandlerUnexpectedMessageType(EndpointName, message.GetType().Name);
-                break;
+                default:
+                    _logger.EndpointHandlerUnexpectedMessageType(EndpointName, message.GetType().Name);
+                    break;
+            }
+        }
+        catch (Exception e) when (addTags)
+        {
+            tags.Add("error.type", e.GetType().FullName);
+            throw;
+        }
+        finally
+        {
+            if (activity is not null)
+            {
+                foreach (var tag in tags)
+                {
+                    activity.AddTag(tag.Key, tag.Value);
+                }
+            }
+
+            if (measuring is not null)
+            {
+                s_operationDurationHistogram.Record(measuring.Elapsed.TotalSeconds, tags);
+            }
         }
     }
 
@@ -264,12 +339,20 @@ public async Task<TResult> SendRequestAsync<TResult>(JsonRpcRequest request, Can
             throw new McpClientException("Transport is not connected");
         }
 
+        using Activity? activity = s_activitySource.StartActivity("SendingRequest");
+
         // Set request ID
         if (request.Id.IsDefault)
         {
             request.Id = new RequestId($"{_id}-{Interlocked.Increment(ref _nextRequestId)}");
         }
 
+        if (activity is not null)
+        {
+            activity.SetTag("mcp.request.id", request.Id.ToString());
+            activity.SetTag("mcp.request.method", request.Method);
+        }
+
         var tcs = new TaskCompletionSource<IJsonRpcMessage>(TaskCreationOptions.RunContinuationsAsynchronously);
         _pendingRequests[request.Id] = tcs;
 
@@ -319,6 +402,11 @@ public async Task<TResult> SendRequestAsync<TResult>(JsonRpcRequest request, Can
             _logger.RequestInvalidResponseType(EndpointName, request.Method);
             throw new McpClientException("Invalid response type");
         }
+        catch (Exception ex) when (activity is not null)
+        {
+            activity.AddTag("error.type", ex.GetType().FullName);
+            throw;
+        }
         finally
         {
             _pendingRequests.TryRemove(request.Id, out _);
@@ -335,6 +423,8 @@ public async Task SendMessageAsync(IJsonRpcMessage message, CancellationToken ca
             throw new McpClientException("Transport is not connected");
         }
 
+        using Activity? activity = s_activitySource.StartActivity("SendingMessage");
+
         if (_logger.IsEnabled(LogLevel.Debug))
         {
             _logger.SendingMessage(EndpointName, JsonSerializer.Serialize(message, _jsonOptions.GetTypeInfo<IJsonRpcMessage>()));

tests/ModelContextProtocol.Tests/DiagnosticTests.cs

@@ -0,0 +1,79 @@
+using ModelContextProtocol.Client;
+using ModelContextProtocol.Protocol.Transport;
+using ModelContextProtocol.Protocol.Types;
+using ModelContextProtocol.Server;
+using OpenTelemetry.Trace;
+using System.Diagnostics;
+using System.IO.Pipelines;
+
+namespace ModelContextProtocol.Tests;
+
+[Collection(nameof(DisableParallelization))]
+public class DiagnosticTests
+{
+    [Fact]
+    public async Task Session_TracksActivities()
+    {
+        var activities = new List<Activity>();
+
+        using (var tracerProvider = OpenTelemetry.Sdk.CreateTracerProviderBuilder()
+            .AddSource("ModelContextProtocol")
+            .AddInMemoryExporter(activities)
+            .Build())
+        {
+            await RunConnected(async (client, server) =>
+            {
+                var tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
+                Assert.NotNull(tools);
+                Assert.NotEmpty(tools);
+
+                var tool = tools.First(t => t.Name == "DoubleValue");
+                await tool.InvokeAsync(new Dictionary<string, object?>() { ["amount"] = 42 }, TestContext.Current.CancellationToken);
+            });
+        }
+
+        Assert.NotEmpty(activities);
+
+        Activity toolCallActivity = activities.First(a =>
+            a.Tags.Any(t => t.Key == "mcp.request.method" && t.Value == "tools/call"));
+        Assert.Equal("DoubleValue", toolCallActivity.Tags.First(t => t.Key == "mcp.request.params.name").Value);
+    }
+
+    private static async Task RunConnected(Func<IMcpClient, IMcpServer, Task> action)
+    {
+        Pipe clientToServerPipe = new(), serverToClientPipe = new();
+        StreamServerTransport serverTransport = new(clientToServerPipe.Reader.AsStream(), serverToClientPipe.Writer.AsStream());
+        StreamClientTransport clientTransport = new(clientToServerPipe.Writer.AsStream(), serverToClientPipe.Reader.AsStream());
+
+        Task serverTask;
+
+        await using (IMcpServer server = McpServerFactory.Create(serverTransport, new()
+            {
+                ServerInfo = new Implementation { Name = "TestServer", Version = "1.0.0" },
+                Capabilities = new()
+                {
+                    Tools = new()
+                    {
+                        ToolCollection = [McpServerTool.Create((int amount) => amount * 2, new() { Name = "DoubleValue", Description = "Doubles the value." })],
+                    }
+                }
+            }))
+        {
+            serverTask = server.RunAsync(TestContext.Current.CancellationToken);
+
+            await using (IMcpClient client = await McpClientFactory.CreateAsync(new()
+                {
+                    Id = "TestServer",
+                    Name = "TestServer",
+                    TransportType = TransportTypes.StdIo,
+                },
+                createTransportFunc: (_, __) => clientTransport,
+                cancellationToken: TestContext.Current.CancellationToken))
+            {
+                await action(client, server);
+            }
+        }
+
+        await serverTask;
+    }
+}

tests/ModelContextProtocol.Tests/ModelContextProtocol.Tests.csproj

@@ -25,6 +25,8 @@
     <PackageReference Include="Microsoft.Extensions.Logging.Console" />
     <PackageReference Include="Microsoft.NET.Test.Sdk" />
     <PackageReference Include="Moq" />
+    <PackageReference Include="OpenTelemetry" />
+    <PackageReference Include="OpenTelemetry.Exporter.InMemory" />
     <PackageReference Include="System.Linq.AsyncEnumerable" />
     <PackageReference Include="xunit.v3" />
     <PackageReference Include="xunit.runner.visualstudio">

tests/ModelContextProtocol.Tests/TestAttributes.cs

@@ -1,2 +1,9 @@
-// Uncomment to disable parallel test execution
-//[assembly: CollectionBehavior(DisableTestParallelization = true)]
+// Uncomment to disable parallel test execution for the whole assembly
+//[assembly: CollectionBehavior(DisableTestParallelization = true)]
+
+/// <summary>
+/// Enables test classes to individually be attributed as [Collection(nameof(DisableParallelization))]
+/// to have those tests run non-concurrently with any other tests.
+/// </summary>
+[CollectionDefinition(nameof(DisableParallelization), DisableParallelization = true)]
+public sealed class DisableParallelization;