Add initial Activity / Metric support (#183)

src/ModelContextProtocol/Diagnostics.cs

@@ -0,0 +1,29 @@
+using System.Diagnostics;
+using System.Diagnostics.Metrics;
+
+namespace ModelContextProtocol;
+
+internal static class Diagnostics
+{
+    internal static ActivitySource ActivitySource { get; } = new("ModelContextProtocol");
+
+    /// <summary>
+    /// Follows boundaries from http.server.request.duration/http.client.request.duration
+    /// </summary>
+    internal static InstrumentAdvice<double> ShortSecondsBucketBoundaries { get; } = new()
+    {
+        HistogramBucketBoundaries = [0.005, 0.01, 0.025, 0.05, 0.075, 0.1, 0.25, 0.5, 0.75, 1, 2.5, 5, 7.5, 10],
+    };
+
+    /// <summary>
+    /// Not based on a standard. Larger bucket sizes for longer lasting operations, e.g. HTTP connection duration.
+    /// See https://github.com/open-telemetry/semantic-conventions/issues/336
+    /// </summary>
+    internal static InstrumentAdvice<double> LongSecondsBucketBoundaries { get; } = new()
+    {
+        HistogramBucketBoundaries = [0.01, 0.02, 0.05, 0.1, 0.2, 0.5, 1, 2, 5, 10, 30, 60, 120, 300],
+    };
+
+    internal static Meter Meter { get; } = new("ModelContextProtocol");
+
+}

src/ModelContextProtocol/Protocol/Messages/ProgressToken.cs

@@ -12,15 +12,15 @@ namespace ModelContextProtocol.Protocol.Messages;
 [JsonConverter(typeof(Converter))]
 public readonly struct ProgressToken : IEquatable<ProgressToken>
 {
-    /// <summary>The id, either a string or a boxed long or null.</summary>
-    private readonly object? _id;
+    /// <summary>The token, either a string or a boxed long or null.</summary>
+    private readonly object? _token;
 
     /// <summary>Initializes a new instance of the <see cref="ProgressToken"/> with a specified value.</summary>
     /// <param name="value">The required ID value.</param>
     public ProgressToken(string value)
     {
         Throw.IfNull(value);
-        _id = value;
+        _token = value;
     }
 
     /// <summary>Initializes a new instance of the <see cref="ProgressToken"/> with a specified value.</summary>
@@ -28,28 +28,29 @@ public readonly struct ProgressToken : IEquatable<ProgressToken>
     public ProgressToken(long value)
     {
         // Box the long. Progress tokens are almost always strings in practice, so this should be rare.
-        _id = value;
+        _token = value;
     }
 
-    /// <summary>Gets whether the identifier is uninitialized.</summary>
-    public bool IsDefault => _id is null;
+    /// <summary>Gets the underlying object for this token.</summary>
+    /// <remarks>This will either be a <see cref="string"/>, a boxed <see cref="long"/>, or <see langword="null"/>.</remarks>
+    public object? Token => _token;
 
     /// <inheritdoc />
     public override string? ToString() =>
-        _id is string stringValue ? $"\"{stringValue}\"" :
-        _id is long longValue ? longValue.ToString(CultureInfo.InvariantCulture) :
+        _token is string stringValue ? $"{stringValue}" :
+        _token is long longValue ? longValue.ToString(CultureInfo.InvariantCulture) :
         null;
 
     /// <summary>
     /// Compares this ProgressToken to another ProgressToken.
     /// </summary>
-    public bool Equals(ProgressToken other) => Equals(_id, other._id);
+    public bool Equals(ProgressToken other) => Equals(_token, other._token);
 
     /// <inheritdoc />
     public override bool Equals(object? obj) => obj is ProgressToken other && Equals(other);
 
     /// <inheritdoc />
-    public override int GetHashCode() => _id?.GetHashCode() ?? 0;
+    public override int GetHashCode() => _token?.GetHashCode() ?? 0;
 
     /// <summary>
     /// Compares two ProgressTokens for equality.
@@ -83,7 +84,7 @@ public readonly struct ProgressToken : IEquatable<ProgressToken>
         {
             Throw.IfNull(writer);
 
-            switch (value._id)
+            switch (value._token)
             {
                 case string str:
                     writer.WriteStringValue(str);

src/ModelContextProtocol/Protocol/Messages/RequestId.cs

@@ -31,12 +31,13 @@ public readonly struct RequestId : IEquatable<RequestId>
         _id = value;
     }
 
-    /// <summary>Gets whether the identifier is uninitialized.</summary>
-    public bool IsDefault => _id is null;
+    /// <summary>Gets the underlying object for this id.</summary>
+    /// <remarks>This will either be a <see cref="string"/>, a boxed <see cref="long"/>, or <see langword="null"/>.</remarks>
+    public object? Id => _id;
 
     /// <inheritdoc />
     public override string ToString() =>
-        _id is string stringValue ? $"\"{stringValue}\"" :
+        _id is string stringValue ? stringValue :
         _id is long longValue ? longValue.ToString(CultureInfo.InvariantCulture) :
         string.Empty;
 

src/ModelContextProtocol/Shared/McpJsonRpcEndpoint.cs

@@ -3,6 +3,7 @@ using Microsoft.Extensions.Logging.Abstractions;
 using ModelContextProtocol.Logging;
 using ModelContextProtocol.Protocol.Messages;
 using ModelContextProtocol.Protocol.Transport;
+using ModelContextProtocol.Server;
 using ModelContextProtocol.Utils;
 using System.Diagnostics.CodeAnalysis;
 
@@ -63,7 +64,7 @@ internal abstract class McpJsonRpcEndpoint : IAsyncDisposable
     protected void StartSession(ITransport sessionTransport)
     {
         _sessionCts = new CancellationTokenSource();
-        _session = new McpSession(sessionTransport, EndpointName, _requestHandlers, _notificationHandlers, _logger);
+        _session = new McpSession(this is IMcpServer, sessionTransport, EndpointName, _requestHandlers, _notificationHandlers, _logger);
         MessageProcessingTask = _session.ProcessMessagesAsync(_sessionCts.Token);
     }
 

src/ModelContextProtocol/Shared/McpSession.cs

@@ -4,9 +4,12 @@ using ModelContextProtocol.Client;
 using ModelContextProtocol.Logging;
 using ModelContextProtocol.Protocol.Messages;
 using ModelContextProtocol.Protocol.Transport;
+using ModelContextProtocol.Server;
 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,9 +19,22 @@ namespace ModelContextProtocol.Shared;
 /// </summary>
 internal sealed class McpSession : IDisposable
 {
+    private static readonly Histogram<double> s_clientSessionDuration = Diagnostics.Meter.CreateHistogram<double>(
+        "mcp.client.session.duration", "s", "Measures the duration of a client session.", advice: Diagnostics.LongSecondsBucketBoundaries);
+    private static readonly Histogram<double> s_serverSessionDuration = Diagnostics.Meter.CreateHistogram<double>(
+        "mcp.server.session.duration", "s", "Measures the duration of a server session.", advice: Diagnostics.LongSecondsBucketBoundaries);
+
+    private static readonly Histogram<double> s_serverRequestDuration = Diagnostics.Meter.CreateHistogram<double>(
+        "rpc.server.duration", "s", "Measures the duration of inbound RPC.", advice: Diagnostics.ShortSecondsBucketBoundaries);
+    private static readonly Histogram<double> s_clientRequestDuration = Diagnostics.Meter.CreateHistogram<double>(
+        "rpc.client.duration", "s", "Measures the duration of outbound RPC.", advice: Diagnostics.ShortSecondsBucketBoundaries);
+
+    private readonly bool _isServer;
+    private readonly string _transportKind;
     private readonly ITransport _transport;
     private readonly RequestHandlers _requestHandlers;
     private readonly NotificationHandlers _notificationHandlers;
+    private readonly long _sessionStartingTimestamp = Stopwatch.GetTimestamp();
 
     /// <summary>Collection of requests sent on this session and waiting for responses.</summary>
     private readonly ConcurrentDictionary<RequestId, TaskCompletionSource<IJsonRpcMessage>> _pendingRequests = [];
@@ -36,12 +52,14 @@ internal sealed class McpSession : IDisposable
     /// <summary>
     /// Initializes a new instance of the <see cref="McpSession"/> class.
     /// </summary>
+    /// <param name="isServer">true if this is a server; false if it's a client.</param>
     /// <param name="transport">An MCP transport implementation.</param>
     /// <param name="endpointName">The name of the endpoint for logging and debug purposes.</param>
     /// <param name="requestHandlers">A collection of request handlers.</param>
     /// <param name="notificationHandlers">A collection of notification handlers.</param>
     /// <param name="logger">The logger.</param>
     public McpSession(
+        bool isServer,
         ITransport transport,
         string endpointName,
         RequestHandlers requestHandlers,
@@ -50,6 +68,15 @@ internal sealed class McpSession : IDisposable
     {
         Throw.IfNull(transport);
 
+        _transportKind = transport switch
+        {
+            StdioClientSessionTransport or StdioServerTransport => "stdio",
+            StreamClientSessionTransport or StreamServerTransport => "stream",
+            SseClientSessionTransport or SseResponseStreamTransport => "sse",
+            _ => "unknownTransport"
+        };
+
+        _isServer = isServer;
         _transport = transport;
         EndpointName = endpointName;
         _requestHandlers = requestHandlers;
@@ -121,7 +148,7 @@ internal sealed class McpSession : IDisposable
                                 JsonRpc = "2.0",
                                 Error = new JsonRpcErrorDetail
                                 {
-                                    Code = ErrorCodes.InternalError,
+                                    Code = (ex as McpServerException)?.ErrorCode ?? ErrorCodes.InternalError,
                                     Message = ex.Message
                                 }
                             }, cancellationToken).ConfigureAwait(false);
@@ -152,23 +179,55 @@ internal sealed class McpSession : IDisposable
 
     private async Task HandleMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken)
     {
-        switch (message)
+        Histogram<double> durationMetric = _isServer ? s_serverRequestDuration : s_clientRequestDuration;
+        string method = GetMethodName(message);
+
+        long? startingTimestamp = durationMetric.Enabled ? Stopwatch.GetTimestamp() : null;
+        Activity? activity = Diagnostics.ActivitySource.HasListeners() ?
+            Diagnostics.ActivitySource.StartActivity(CreateActivityName(method)) :
+            null;
+
+        TagList tags = default;
+        bool addTags = activity is { IsAllDataRequested: true } || startingTimestamp is not null;
+        try
         {
-            case JsonRpcRequest request:
-                await HandleRequest(request, cancellationToken).ConfigureAwait(false);
-                break;
+            if (addTags)
+            {
+                AddStandardTags(ref tags, method);
+            }
 
-            case IJsonRpcMessageWithId messageWithId:
-                HandleMessageWithId(message, messageWithId);
-                break;
+            switch (message)
+            {
+                case JsonRpcRequest request:
+                    if (addTags)
+                    {
+                        AddRpcRequestTags(ref tags, activity, request);
+                    }
 
-            case JsonRpcNotification notification:
-                await HandleNotification(notification).ConfigureAwait(false);
-                break;
+                    await HandleRequest(request, cancellationToken).ConfigureAwait(false);
+                    break;
 
-            default:
-                _logger.EndpointHandlerUnexpectedMessageType(EndpointName, message.GetType().Name);
-                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;
+            }
+        }
+        catch (Exception e) when (addTags)
+        {
+            AddExceptionTags(ref tags, e);
+            throw;
+        }
+        finally
+        {
+            FinalizeDiagnostics(activity, startingTimestamp, durationMetric, ref tags);
         }
     }
 
@@ -212,7 +271,7 @@ internal sealed class McpSession : IDisposable
 
     private void HandleMessageWithId(IJsonRpcMessage message, IJsonRpcMessageWithId messageWithId)
     {
-        if (messageWithId.Id.IsDefault)
+        if (messageWithId.Id.Id is null)
         {
             _logger.RequestHasInvalidId(EndpointName);
         }
@@ -229,22 +288,21 @@ internal sealed class McpSession : IDisposable
 
     private async Task HandleRequest(JsonRpcRequest request, CancellationToken cancellationToken)
     {
-        if (_requestHandlers.TryGetValue(request.Method, out var handler))
-        {
-            _logger.RequestHandlerCalled(EndpointName, request.Method);
-            var result = await handler(request, cancellationToken).ConfigureAwait(false);
-            _logger.RequestHandlerCompleted(EndpointName, request.Method);
-            await _transport.SendMessageAsync(new JsonRpcResponse
-            {
-                Id = request.Id,
-                JsonRpc = "2.0",
-                Result = result
-            }, cancellationToken).ConfigureAwait(false);
-        }
-        else
+        if (!_requestHandlers.TryGetValue(request.Method, out var handler))
         {
             _logger.NoHandlerFoundForRequest(EndpointName, request.Method);
+            throw new McpServerException("The method does not exist or is not available.", ErrorCodes.MethodNotFound);
         }
+
+        _logger.RequestHandlerCalled(EndpointName, request.Method);
+        var result = await handler(request, cancellationToken).ConfigureAwait(false);
+        _logger.RequestHandlerCompleted(EndpointName, request.Method);
+        await _transport.SendMessageAsync(new JsonRpcResponse
+        {
+            Id = request.Id,
+            JsonRpc = "2.0",
+            Result = result
+        }, cancellationToken).ConfigureAwait(false);
     }
 
     /// <summary>
@@ -264,17 +322,33 @@ internal sealed class McpSession : IDisposable
             throw new McpClientException("Transport is not connected");
         }
 
+        Histogram<double> durationMetric = _isServer ? s_serverRequestDuration : s_clientRequestDuration;
+        string method = request.Method;
+
+        long? startingTimestamp = durationMetric.Enabled ? Stopwatch.GetTimestamp() : null;
+        using Activity? activity = Diagnostics.ActivitySource.HasListeners() ?
+            Diagnostics.ActivitySource.StartActivity(CreateActivityName(method)) :
+            null;
+
         // Set request ID
-        if (request.Id.IsDefault)
+        if (request.Id.Id is null)
         {
             request.Id = new RequestId($"{_id}-{Interlocked.Increment(ref _nextRequestId)}");
         }
 
+        TagList tags = default;
+        bool addTags = activity is { IsAllDataRequested: true } || startingTimestamp is not null;
+
         var tcs = new TaskCompletionSource<IJsonRpcMessage>(TaskCreationOptions.RunContinuationsAsynchronously);
         _pendingRequests[request.Id] = tcs;
-
         try
         {
+            if (addTags)
+            {
+                AddStandardTags(ref tags, method);
+                AddRpcRequestTags(ref tags, activity, request);
+            }
+
             // Expensive logging, use the logging framework to check if the logger is enabled
             if (_logger.IsEnabled(LogLevel.Debug))
             {
@@ -319,9 +393,15 @@ internal sealed class McpSession : IDisposable
             _logger.RequestInvalidResponseType(EndpointName, request.Method);
             throw new McpClientException("Invalid response type");
         }
+        catch (Exception ex) when (addTags)
+        {
+            AddExceptionTags(ref tags, ex);
+            throw;
+        }
         finally
         {
             _pendingRequests.TryRemove(request.Id, out _);
+            FinalizeDiagnostics(activity, startingTimestamp, durationMetric, ref tags);
         }
     }
 
@@ -335,21 +415,49 @@ internal sealed class McpSession : IDisposable
             throw new McpClientException("Transport is not connected");
         }
 
-        if (_logger.IsEnabled(LogLevel.Debug))
+        Histogram<double> durationMetric = _isServer ? s_serverRequestDuration : s_clientRequestDuration;
+        string method = GetMethodName(message);
+
+        long? startingTimestamp = durationMetric.Enabled ? Stopwatch.GetTimestamp() : null;
+        using Activity? activity = Diagnostics.ActivitySource.HasListeners() ?
+            Diagnostics.ActivitySource.StartActivity(CreateActivityName(method)) :
+            null;
+
+        TagList tags = default;
+        bool addTags = activity is { IsAllDataRequested: true } || startingTimestamp is not null;
+
+        try
         {
-            _logger.SendingMessage(EndpointName, JsonSerializer.Serialize(message, _jsonOptions.GetTypeInfo<IJsonRpcMessage>()));
+            if (addTags)
+            {
+                AddStandardTags(ref tags, method);
+            }
+
+            if (_logger.IsEnabled(LogLevel.Debug))
+            {
+                _logger.SendingMessage(EndpointName, JsonSerializer.Serialize(message, _jsonOptions.GetTypeInfo<IJsonRpcMessage>()));
+            }
+
+            await _transport.SendMessageAsync(message, cancellationToken).ConfigureAwait(false);
+
+            // If the sent notification was a cancellation notification, cancel the pending request's await, as either the
+            // server won't be sending a response, or per the specification, the response should be ignored. There are inherent
+            // race conditions here, so it's possible and allowed for the operation to complete before we get to this point.
+            if (message is JsonRpcNotification { Method: NotificationMethods.CancelledNotification } notification &&
+                GetCancelledNotificationParams(notification.Params) is CancelledNotification cn &&
+                _pendingRequests.TryRemove(cn.RequestId, out var tcs))
+            {
+                tcs.TrySetCanceled(default);
+            }
         }
-
-        await _transport.SendMessageAsync(message, cancellationToken).ConfigureAwait(false);
-
-        // If the sent notification was a cancellation notification, cancel the pending request's await, as either the
-        // server won't be sending a response, or per the specification, the response should be ignored. There are inherent
-        // race conditions here, so it's possible and allowed for the operation to complete before we get to this point.
-        if (message is JsonRpcNotification { Method: NotificationMethods.CancelledNotification } notification &&
-            GetCancelledNotificationParams(notification.Params) is CancelledNotification cn &&
-            _pendingRequests.TryRemove(cn.RequestId, out var tcs))
+        catch (Exception ex) when (addTags)
         {
-            tcs.TrySetCanceled(default);
+            AddExceptionTags(ref tags, ex);
+            throw;
+        }
+        finally
+        {
+            FinalizeDiagnostics(activity, startingTimestamp, durationMetric, ref tags);
         }
     }
 
@@ -380,13 +488,127 @@ internal sealed class McpSession : IDisposable
         }
     }
 
+    private string CreateActivityName(string method) =>
+        $"mcp.{(_isServer ? "server" : "client")}.{_transportKind}/{method}";
+
+    private static string GetMethodName(IJsonRpcMessage message) =>
+        message switch
+        {
+            JsonRpcRequest request => request.Method,
+            JsonRpcNotification notification => notification.Method,
+            _ => "unknownMethod",
+        };
+
+    private void AddStandardTags(ref TagList tags, string method)
+    {
+        tags.Add("session.id", _id);
+        tags.Add("rpc.system", "jsonrpc");
+        tags.Add("rpc.jsonrpc.version", "2.0");
+        tags.Add("rpc.method", method);
+        tags.Add("network.transport", _transportKind);
+
+        // RPC spans convention also includes:
+        // server.address, server.port, client.address, client.port, network.peer.address, network.peer.port, network.type
+    }
+
+    private static void AddRpcRequestTags(ref TagList tags, Activity? activity, JsonRpcRequest request)
+    {
+        tags.Add("rpc.jsonrpc.request_id", request.Id.ToString());
+
+        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)
+                    {
+                        string name = prop.GetString()!;
+                        tags.Add("mcp.request.params.name", name);
+                        if (activity is not null)
+                        {
+                            activity.DisplayName = $"{request.Method}({name})";
+                        }
+                    }
+                    break;
+
+                case RequestMethods.ResourcesRead:
+                    if (je.TryGetProperty("uri", out prop) && prop.ValueKind == JsonValueKind.String)
+                    {
+                        string uri = prop.GetString()!;
+                        tags.Add("mcp.request.params.uri", uri);
+                        if (activity is not null)
+                        {
+                            activity.DisplayName = $"{request.Method}({uri})";
+                        }
+                    }
+                    break;
+            }
+        }
+    }
+
+    private static void AddExceptionTags(ref TagList tags, Exception e)
+    {
+        tags.Add("error.type", e.GetType().FullName);
+        tags.Add("rpc.jsonrpc.error_code",
+            (e as McpClientException)?.ErrorCode is int clientError ? clientError :
+            (e as McpServerException)?.ErrorCode is int serverError ? serverError :
+            e is JsonException ? ErrorCodes.ParseError :
+            ErrorCodes.InternalError);
+    }
+
+    private static void FinalizeDiagnostics(
+        Activity? activity, long? startingTimestamp, Histogram<double> durationMetric, ref TagList tags)
+    {
+        try
+        {
+            if (startingTimestamp is not null)
+            {
+                durationMetric.Record(GetElapsed(startingTimestamp.Value).TotalSeconds, tags);
+            }
+
+            if (activity is { IsAllDataRequested: true })
+            {
+                foreach (var tag in tags)
+                {
+                    activity.AddTag(tag.Key, tag.Value);
+                }
+            }
+        }
+        finally
+        {
+            activity?.Dispose();
+        }
+    }
+
     public void Dispose()
     {
+        Histogram<double> durationMetric = _isServer ? s_serverSessionDuration : s_clientSessionDuration;
+        if (durationMetric.Enabled)
+        {
+            TagList tags = default;
+            tags.Add("session.id", _id);
+            tags.Add("network.transport", _transportKind);
+            durationMetric.Record(GetElapsed(_sessionStartingTimestamp).TotalSeconds, tags);
+        }
+
         // Complete all pending requests with cancellation
         foreach (var entry in _pendingRequests)
         {
             entry.Value.TrySetCanceled();
         }
+
         _pendingRequests.Clear();
     }
+
+#if !NET
+    private static readonly double s_timestampToTicks = TimeSpan.TicksPerSecond / (double)Stopwatch.Frequency;
+#endif
+
+    private static TimeSpan GetElapsed(long startingTimestamp) =>
+#if NET
+        Stopwatch.GetElapsedTime(startingTimestamp);
+#else
+        new((long)(s_timestampToTicks * (Stopwatch.GetTimestamp() - startingTimestamp)));
+#endif
 }