Skip to content

Commit 7cd6a01

Browse files
Copilotstephentoub
andcommitted
Add Trace-level logging for JSON-RPC payloads in transports
- Add LogTransportSendingMessage and LogTransportSendingMessageSensitive to TransportBase - Update StreamServerTransport to log sent messages at Trace level - Update StreamClientSessionTransport to log sent messages at Trace level - Update SseClientSessionTransport to log sent messages at Trace level - Update StreamableHttpClientSessionTransport to log sent messages at Trace level - Add tests for Trace-level logging of sent and received messages Co-authored-by: stephentoub <2642209+stephentoub@users.noreply.github.com>
1 parent d661159 commit 7cd6a01

6 files changed

Lines changed: 120 additions & 2 deletions

File tree

src/ModelContextProtocol.Core/Client/SseClientSessionTransport.cs

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,15 @@ public override async Task SendMessageAsync(
8080
messageId = messageWithId.Id.ToString();
8181
}
8282

83+
if (_logger.IsEnabled(LogLevel.Trace))
84+
{
85+
LogTransportSendingMessageSensitive(Name, JsonSerializer.Serialize(message, McpJsonUtilities.JsonContext.Default.JsonRpcMessage));
86+
}
87+
else
88+
{
89+
LogTransportSendingMessage(Name, messageId);
90+
}
91+
8392
using var httpRequestMessage = new HttpRequestMessage(HttpMethod.Post, _messageEndpoint);
8493
StreamableHttpClientSessionTransport.CopyAdditionalHeaders(httpRequestMessage.Headers, _options.AdditionalHeaders, sessionId: null, protocolVersion: null);
8594
var response = await _httpClient.SendAsync(httpRequestMessage, message, cancellationToken).ConfigureAwait(false);

src/ModelContextProtocol.Core/Client/StreamClientSessionTransport.cs

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,15 @@ public override async Task SendMessageAsync(JsonRpcMessage message, Cancellation
105105

106106
var json = JsonSerializer.Serialize(message, McpJsonUtilities.JsonContext.Default.JsonRpcMessage);
107107

108+
if (Logger.IsEnabled(LogLevel.Trace))
109+
{
110+
LogTransportSendingMessageSensitive(Name, json);
111+
}
112+
else
113+
{
114+
LogTransportSendingMessage(Name, id);
115+
}
116+
108117
using var _ = await _sendLock.LockAsync(cancellationToken).ConfigureAwait(false);
109118
try
110119
{

src/ModelContextProtocol.Core/Client/StreamableHttpClientSessionTransport.cs

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,21 @@ internal async Task<HttpResponseMessage> SendHttpRequestAsync(JsonRpcMessage mes
7575
$"Call {nameof(McpClient)}.{nameof(McpClient.ResumeSessionAsync)} to resume existing sessions.");
7676
}
7777

78+
string messageId = "(no id)";
79+
if (message is JsonRpcMessageWithId messageWithId)
80+
{
81+
messageId = messageWithId.Id.ToString();
82+
}
83+
84+
if (_logger.IsEnabled(LogLevel.Trace))
85+
{
86+
LogTransportSendingMessageSensitive(Name, JsonSerializer.Serialize(message, McpJsonUtilities.JsonContext.Default.JsonRpcMessage));
87+
}
88+
else
89+
{
90+
LogTransportSendingMessage(Name, messageId);
91+
}
92+
7893
using var sendCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, _connectionCts.Token);
7994
cancellationToken = sendCts.Token;
8095

src/ModelContextProtocol.Core/Protocol/TransportBase.cs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -166,6 +166,12 @@ protected void SetDisconnected(Exception? error = null)
166166
[LoggerMessage(Level = LogLevel.Error, Message = "{EndpointName} transport send failed for message ID '{MessageId}'.")]
167167
private protected partial void LogTransportSendFailed(string endpointName, string messageId, Exception exception);
168168

169+
[LoggerMessage(Level = LogLevel.Debug, Message = "{EndpointName} transport sending message with ID '{MessageId}'.")]
170+
private protected partial void LogTransportSendingMessage(string endpointName, string messageId);
171+
172+
[LoggerMessage(Level = LogLevel.Trace, Message = "{EndpointName} transport sending message. Message: '{Message}'.")]
173+
private protected partial void LogTransportSendingMessageSensitive(string endpointName, string message);
174+
169175
[LoggerMessage(Level = LogLevel.Information, Message = "{EndpointName} transport reading messages.")]
170176
private protected partial void LogTransportEnteringReadMessagesLoop(string endpointName);
171177

src/ModelContextProtocol.Core/Server/StreamServerTransport.cs

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,17 @@ public override async Task SendMessageAsync(JsonRpcMessage message, Cancellation
7474

7575
try
7676
{
77-
await JsonSerializer.SerializeAsync(_outputStream, message, McpJsonUtilities.DefaultOptions.GetTypeInfo(typeof(JsonRpcMessage)), cancellationToken).ConfigureAwait(false);
77+
if (Logger.IsEnabled(LogLevel.Trace))
78+
{
79+
var json = JsonSerializer.Serialize(message, McpJsonUtilities.JsonContext.Default.JsonRpcMessage);
80+
LogTransportSendingMessageSensitive(Name, json);
81+
await _outputStream.WriteAsync(Encoding.UTF8.GetBytes(json), cancellationToken).ConfigureAwait(false);
82+
}
83+
else
84+
{
85+
LogTransportSendingMessage(Name, id);
86+
await JsonSerializer.SerializeAsync(_outputStream, message, McpJsonUtilities.DefaultOptions.GetTypeInfo(typeof(JsonRpcMessage)), cancellationToken).ConfigureAwait(false);
87+
}
7888
await _outputStream.WriteAsync(s_newlineBytes, cancellationToken).ConfigureAwait(false);
7989
await _outputStream.FlushAsync(cancellationToken).ConfigureAwait(false);
8090
}

tests/ModelContextProtocol.Tests/Transport/StdioServerTransportTests.cs

Lines changed: 70 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
1-
using ModelContextProtocol.Protocol;
1+
using Microsoft.Extensions.Logging;
2+
using ModelContextProtocol.Protocol;
23
using ModelContextProtocol.Server;
34
using ModelContextProtocol.Tests.Utils;
45
using System.IO.Pipelines;
@@ -193,4 +194,72 @@ public async Task SendMessageAsync_Should_Preserve_Unicode_Characters()
193194
Assert.True(magnifyingGlassFound, "Magnifying glass emoji not found in result");
194195
Assert.True(rocketFound, "Rocket emoji not found in result");
195196
}
197+
198+
[Fact]
199+
public async Task SendMessageAsync_Should_Log_At_Trace_Level()
200+
{
201+
// Arrange
202+
var mockLoggerProvider = new MockLoggerProvider();
203+
using var traceLoggerFactory = Microsoft.Extensions.Logging.LoggerFactory.Create(builder =>
204+
{
205+
builder.AddProvider(mockLoggerProvider);
206+
builder.SetMinimumLevel(LogLevel.Trace);
207+
});
208+
using var output = new MemoryStream();
209+
210+
await using var transport = new StreamServerTransport(
211+
new Pipe().Reader.AsStream(),
212+
output,
213+
loggerFactory: traceLoggerFactory);
214+
215+
// Act
216+
var message = new JsonRpcRequest { Method = "test", Id = new RequestId(44) };
217+
await transport.SendMessageAsync(message, TestContext.Current.CancellationToken);
218+
219+
// Assert
220+
var traceLogMessages = mockLoggerProvider.LogMessages
221+
.Where(x => x.LogLevel == LogLevel.Trace && x.Message.Contains("transport sending message"))
222+
.ToList();
223+
224+
Assert.NotEmpty(traceLogMessages);
225+
Assert.Contains(traceLogMessages, x => x.Message.Contains("\"method\":\"test\"") && x.Message.Contains("\"id\":44"));
226+
}
227+
228+
[Fact]
229+
public async Task ReadMessagesAsync_Should_Log_Received_At_Trace_Level()
230+
{
231+
// Arrange
232+
var mockLoggerProvider = new MockLoggerProvider();
233+
using var traceLoggerFactory = Microsoft.Extensions.Logging.LoggerFactory.Create(builder =>
234+
{
235+
builder.AddProvider(mockLoggerProvider);
236+
builder.SetMinimumLevel(LogLevel.Trace);
237+
});
238+
239+
var message = new JsonRpcRequest { Method = "test", Id = new RequestId(99) };
240+
var json = JsonSerializer.Serialize(message, McpJsonUtilities.DefaultOptions);
241+
242+
Pipe pipe = new();
243+
using var input = pipe.Reader.AsStream();
244+
245+
await using var transport = new StreamServerTransport(
246+
input,
247+
Stream.Null,
248+
loggerFactory: traceLoggerFactory);
249+
250+
// Act
251+
await pipe.Writer.WriteAsync(Encoding.UTF8.GetBytes($"{json}\n"), TestContext.Current.CancellationToken);
252+
253+
// Wait for the message to be processed
254+
var canRead = await transport.MessageReader.WaitToReadAsync(TestContext.Current.CancellationToken);
255+
Assert.True(canRead, "Nothing to read here from transport message reader");
256+
257+
// Assert
258+
var traceLogMessages = mockLoggerProvider.LogMessages
259+
.Where(x => x.LogLevel == LogLevel.Trace && x.Message.Contains("transport received message"))
260+
.ToList();
261+
262+
Assert.NotEmpty(traceLogMessages);
263+
Assert.Contains(traceLogMessages, x => x.Message.Contains("\"method\":\"test\"") && x.Message.Contains("\"id\":99"));
264+
}
196265
}

0 commit comments

Comments
 (0)