Backfills trace injection test and migrates from request.params to request.params._meta (#360)
* Adds test to show where trace IDs are encoded Signed-off-by: Adrian Cole <adrian.cole@elastic.co> * update impl to params._meta Signed-off-by: Adrian Cole <adrian.cole@elastic.co> --------- Signed-off-by: Adrian Cole <adrian.cole@elastic.co>
This commit is contained in:
parent
2eeb61f4fa
commit
17a94272d9
2 changed files with 54 additions and 16 deletions
|
|
@ -50,22 +50,19 @@ internal static class Diagnostics
|
|||
fieldValues = null;
|
||||
fieldValue = null;
|
||||
|
||||
JsonNode? parameters = null;
|
||||
JsonNode? meta = null;
|
||||
switch (message)
|
||||
{
|
||||
case JsonRpcRequest request:
|
||||
parameters = request.Params;
|
||||
meta = request.Params?["_meta"];
|
||||
break;
|
||||
|
||||
case JsonRpcNotification notification:
|
||||
parameters = notification.Params;
|
||||
break;
|
||||
|
||||
default:
|
||||
meta = notification.Params?["_meta"];
|
||||
break;
|
||||
}
|
||||
|
||||
if (parameters?[fieldName] is JsonValue value && value.GetValueKind() == JsonValueKind.String)
|
||||
if (meta?[fieldName] is JsonValue value && value.GetValueKind() == JsonValueKind.String)
|
||||
{
|
||||
fieldValue = value.GetValue<string>();
|
||||
}
|
||||
|
|
@ -89,14 +86,17 @@ internal static class Diagnostics
|
|||
case JsonRpcNotification notification:
|
||||
parameters = notification.Params;
|
||||
break;
|
||||
|
||||
default:
|
||||
break;
|
||||
}
|
||||
|
||||
if (parameters is JsonObject jsonObject && jsonObject[key] == null)
|
||||
// Replace any params._meta with the current value
|
||||
if (parameters is JsonObject jsonObject)
|
||||
{
|
||||
jsonObject[key] = value;
|
||||
if (jsonObject["_meta"] is not JsonObject meta)
|
||||
{
|
||||
meta = new JsonObject();
|
||||
jsonObject["_meta"] = meta;
|
||||
}
|
||||
meta[key] = value;
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -4,6 +4,8 @@ using ModelContextProtocol.Server;
|
|||
using OpenTelemetry.Trace;
|
||||
using System.Diagnostics;
|
||||
using System.IO.Pipelines;
|
||||
using System.Text;
|
||||
using System.Text.Json;
|
||||
|
||||
namespace ModelContextProtocol.Tests;
|
||||
|
||||
|
|
@ -14,6 +16,7 @@ public class DiagnosticTests
|
|||
public async Task Session_TracksActivities()
|
||||
{
|
||||
var activities = new List<Activity>();
|
||||
var clientToServerLog = new List<string>();
|
||||
|
||||
using (var tracerProvider = OpenTelemetry.Sdk.CreateTracerProviderBuilder()
|
||||
.AddSource("Experimental.ModelContextProtocol")
|
||||
|
|
@ -28,7 +31,7 @@ public class DiagnosticTests
|
|||
|
||||
var tool = tools.First(t => t.Name == "DoubleValue");
|
||||
await tool.InvokeAsync(new() { ["amount"] = 42 }, TestContext.Current.CancellationToken);
|
||||
});
|
||||
}, clientToServerLog);
|
||||
}
|
||||
|
||||
Assert.NotEmpty(activities);
|
||||
|
|
@ -64,6 +67,11 @@ public class DiagnosticTests
|
|||
|
||||
Assert.Equal(clientListToolsCall.SpanId, serverListToolsCall.ParentSpanId);
|
||||
Assert.Equal(clientListToolsCall.TraceId, serverListToolsCall.TraceId);
|
||||
|
||||
// Validate that the client trace context encoded to request.params._meta[traceparent]
|
||||
using var listToolsJson = JsonDocument.Parse(clientToServerLog.First(s => s.Contains("\"method\":\"tools/list\"")));
|
||||
var metaJson = listToolsJson.RootElement.GetProperty("params").GetProperty("_meta").GetRawText();
|
||||
Assert.Equal($$"""{"traceparent":"00-{{clientListToolsCall.TraceId}}-{{clientListToolsCall.SpanId}}-01"}""", metaJson);
|
||||
}
|
||||
|
||||
[Fact]
|
||||
|
|
@ -80,7 +88,7 @@ public class DiagnosticTests
|
|||
{
|
||||
await client.CallToolAsync("Throw", cancellationToken: TestContext.Current.CancellationToken);
|
||||
await Assert.ThrowsAsync<McpException>(() => client.CallToolAsync("does-not-exist", cancellationToken: TestContext.Current.CancellationToken));
|
||||
});
|
||||
}, new List<string>());
|
||||
}
|
||||
|
||||
Assert.NotEmpty(activities);
|
||||
|
|
@ -120,11 +128,12 @@ public class DiagnosticTests
|
|||
Assert.Equal("-32602", doesNotExistToolClient.Tags.Single(t => t.Key == "rpc.jsonrpc.error_code").Value);
|
||||
}
|
||||
|
||||
private static async Task RunConnected(Func<IMcpClient, IMcpServer, Task> action)
|
||||
private static async Task RunConnected(Func<IMcpClient, IMcpServer, Task> action, List<string> clientToServerLog)
|
||||
{
|
||||
Pipe clientToServerPipe = new(), serverToClientPipe = new();
|
||||
StreamServerTransport serverTransport = new(clientToServerPipe.Reader.AsStream(), serverToClientPipe.Writer.AsStream());
|
||||
StreamClientTransport clientTransport = new(clientToServerPipe.Writer.AsStream(), serverToClientPipe.Reader.AsStream());
|
||||
StreamClientTransport clientTransport = new(new LoggingStream(
|
||||
clientToServerPipe.Writer.AsStream(), clientToServerLog.Add), serverToClientPipe.Reader.AsStream());
|
||||
|
||||
Task serverTask;
|
||||
|
||||
|
|
@ -155,3 +164,32 @@ public class DiagnosticTests
|
|||
await serverTask;
|
||||
}
|
||||
}
|
||||
|
||||
public class LoggingStream : Stream
|
||||
{
|
||||
private readonly Stream _innerStream;
|
||||
private readonly Action<string> _logAction;
|
||||
|
||||
public LoggingStream(Stream innerStream, Action<string> logAction)
|
||||
{
|
||||
_innerStream = innerStream ?? throw new ArgumentNullException(nameof(innerStream));
|
||||
_logAction = logAction ?? throw new ArgumentNullException(nameof(logAction));
|
||||
}
|
||||
|
||||
public override void Write(byte[] buffer, int offset, int count)
|
||||
{
|
||||
var data = Encoding.UTF8.GetString(buffer, offset, count);
|
||||
_logAction(data);
|
||||
_innerStream.Write(buffer, offset, count);
|
||||
}
|
||||
|
||||
public override bool CanRead => _innerStream.CanRead;
|
||||
public override bool CanSeek => _innerStream.CanSeek;
|
||||
public override bool CanWrite => _innerStream.CanWrite;
|
||||
public override long Length => _innerStream.Length;
|
||||
public override long Position { get => _innerStream.Position; set => _innerStream.Position = value; }
|
||||
public override void Flush() => _innerStream.Flush();
|
||||
public override int Read(byte[] buffer, int offset, int count) => _innerStream.Read(buffer, offset, count);
|
||||
public override long Seek(long offset, SeekOrigin origin) => _innerStream.Seek(offset, origin);
|
||||
public override void SetLength(long value) => _innerStream.SetLength(value);
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue