Address feedback

Author: stephentoub

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/RequestId.cs

@@ -34,6 +34,10 @@ public RequestId(long 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 :

src/ModelContextProtocol/Shared/McpSession.cs

@@ -19,30 +19,15 @@ namespace ModelContextProtocol.Shared;
 /// </summary>
 internal sealed class McpSession : IDisposable
 {
-    private static readonly ActivitySource s_activitySource = new("ModelContextProtocol");
+    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 InstrumentAdvice<double> s_shortSecondsBucketBoundaries = new()
-    {
-        // Follows boundaries from http.server.request.duration/http.client.request.duration
-        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],
-    };
-
-    private static readonly InstrumentAdvice<double> s_longSecondsBucketBoundaries = new()
-    {
-        // 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
-        HistogramBucketBoundaries = [0.01, 0.02, 0.05, 0.1, 0.2, 0.5, 1, 2, 5, 10, 30, 60, 120, 300],
-    };
-
-    private static readonly Meter s_meter = new("ModelContextProtocol");
-    private static readonly Histogram<double> s_clientSessionDuration = s_meter.CreateHistogram<double>(
-        "mcp.client.session.duration", "s", "Measures the duration of a client session.", advice: s_longSecondsBucketBoundaries);
-    private static readonly Histogram<double> s_serverSessionDuration = s_meter.CreateHistogram<double>(
-        "mcp.server.session.duration", "s", "Measures the duration of a server session.", advice: s_longSecondsBucketBoundaries);
-    private static readonly Histogram<double> s_serverRequestDuration = s_meter.CreateHistogram<double>(
-        "rpc.server.duration", "s", "Measures the duration of inbound RPC.", advice: s_shortSecondsBucketBoundaries);
-    private static readonly Histogram<double> s_clientRequestDuration = s_meter.CreateHistogram<double>(
-        "rpc.client.duration", "s", "Measures the duration of outbound RPC.", advice: s_shortSecondsBucketBoundaries);
+    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;
@@ -198,12 +183,12 @@ private async Task HandleMessageAsync(IJsonRpcMessage message, CancellationToken
         string method = GetMethodName(message);
 
         long? startingTimestamp = durationMetric.Enabled ? Stopwatch.GetTimestamp() : null;
-        Activity? activity = s_activitySource.HasListeners() ?
-            s_activitySource.StartActivity(CreateActivityName(method)) :
+        Activity? activity = Diagnostics.ActivitySource.HasListeners() ?
+            Diagnostics.ActivitySource.StartActivity(CreateActivityName(method)) :
             null;
 
         TagList tags = default;
-        bool addTags = activity is not null || startingTimestamp is not null;
+        bool addTags = activity is { IsAllDataRequested: true } || startingTimestamp is not null;
         try
         {
             if (addTags)
@@ -341,8 +326,8 @@ public async Task<TResult> SendRequestAsync<TResult>(JsonRpcRequest request, Can
         string method = request.Method;
 
         long? startingTimestamp = durationMetric.Enabled ? Stopwatch.GetTimestamp() : null;
-        using Activity? activity = s_activitySource.HasListeners() ?
-            s_activitySource.StartActivity(CreateActivityName(method)) :
+        using Activity? activity = Diagnostics.ActivitySource.HasListeners() ?
+            Diagnostics.ActivitySource.StartActivity(CreateActivityName(method)) :
             null;
 
         // Set request ID
@@ -352,7 +337,7 @@ public async Task<TResult> SendRequestAsync<TResult>(JsonRpcRequest request, Can
         }
 
         TagList tags = default;
-        bool addTags = activity is not null || startingTimestamp is not null;
+        bool addTags = activity is { IsAllDataRequested: true } || startingTimestamp is not null;
 
         var tcs = new TaskCompletionSource<IJsonRpcMessage>(TaskCreationOptions.RunContinuationsAsynchronously);
         _pendingRequests[request.Id] = tcs;
@@ -434,12 +419,12 @@ public async Task SendMessageAsync(IJsonRpcMessage message, CancellationToken ca
         string method = GetMethodName(message);
 
         long? startingTimestamp = durationMetric.Enabled ? Stopwatch.GetTimestamp() : null;
-        using Activity? activity = s_activitySource.HasListeners() ?
-            s_activitySource.StartActivity(CreateActivityName(method)) :
+        using Activity? activity = Diagnostics.ActivitySource.HasListeners() ?
+            Diagnostics.ActivitySource.StartActivity(CreateActivityName(method)) :
             null;
 
         TagList tags = default;
-        bool addTags = activity is not null || startingTimestamp is not null;
+        bool addTags = activity is { IsAllDataRequested: true } || startingTimestamp is not null;
 
         try
         {
@@ -516,6 +501,7 @@ private string GetMethodName(IJsonRpcMessage message) =>
 
     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);
@@ -581,7 +567,7 @@ private static void FinalizeDiagnostics(
                 durationMetric.Record(GetElapsed(startingTimestamp.Value).TotalSeconds, tags);
             }
 
-            if (activity is not null)
+            if (activity is { IsAllDataRequested: true })
             {
                 foreach (var tag in tags)
                 {
@@ -600,7 +586,10 @@ public void Dispose()
         Histogram<double> durationMetric = _isServer ? s_serverSessionDuration : s_clientSessionDuration;
         if (durationMetric.Enabled)
         {
-            durationMetric.Record(GetElapsed(_sessionStartingTimestamp).TotalSeconds);
+            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

tests/ModelContextProtocol.Tests/Protocol/RequestIdTests.cs

@@ -0,0 +1,38 @@
+using ModelContextProtocol.Protocol.Messages;
+using System.Text.Json;
+
+namespace ModelContextProtocol.Tests.Protocol;
+
+public class RequestIdTests
+{
+    [Fact]
+    public void StringCtor_Roundtrips()
+    {
+        RequestId id = new("test-id");
+        Assert.Equal("test-id", id.ToString());
+        Assert.Equal("\"test-id\"", JsonSerializer.Serialize(id));
+        Assert.Same("test-id", id.Id);
+
+        Assert.True(id.Equals(new("test-id")));
+        Assert.False(id.Equals(new("tEst-id")));
+        Assert.Equal("test-id".GetHashCode(), id.GetHashCode());
+
+        Assert.Equal(id, JsonSerializer.Deserialize<RequestId>(JsonSerializer.Serialize(id)));
+    }
+
+    [Fact]
+    public void Int64Ctor_Roundtrips()
+    {
+        RequestId id = new(42);
+        Assert.Equal("42", id.ToString());
+        Assert.Equal("42", JsonSerializer.Serialize(id));
+        Assert.Equal(42, Assert.IsType<long>(id.Id));
+
+        Assert.True(id.Equals(new(42)));
+        Assert.False(id.Equals(new(43)));
+        Assert.False(id.Equals(new("42")));
+        Assert.Equal(42L.GetHashCode(), id.GetHashCode());
+
+        Assert.Equal(id, JsonSerializer.Deserialize<RequestId>(JsonSerializer.Serialize(id)));
+    }
+}