95 lines
3.7 KiB
C#
95 lines
3.7 KiB
C#
using ModelContextProtocol.Protocol;
|
|
using ModelContextProtocol.Server;
|
|
using System.Collections.Concurrent;
|
|
using System.Text.Json.Nodes;
|
|
using Xunit;
|
|
|
|
namespace ModelContextProtocol.Tests.Utils;
|
|
|
|
/// <summary>
|
|
/// Tracks MRTR protocol mode via incoming and outgoing message filters.
|
|
/// Used by MRTR tests to verify the correct protocol mode (MRTR vs legacy) was used.
|
|
/// </summary>
|
|
internal sealed class ServerMessageTracker
|
|
{
|
|
private static readonly HashSet<string> LegacyMrtrMethods =
|
|
[
|
|
RequestMethods.ElicitationCreate,
|
|
RequestMethods.SamplingCreateMessage,
|
|
RequestMethods.RootsList,
|
|
];
|
|
|
|
private readonly ConcurrentBag<string> _legacyRequestMethods = [];
|
|
private int _mrtrRetryCount;
|
|
private int _incompleteResultCount;
|
|
|
|
/// <summary>
|
|
/// Adds incoming and outgoing message filters to track MRTR protocol usage.
|
|
/// Call this in <c>services.Configure<McpServerOptions></c> or <c>AddMcpServer</c> callbacks.
|
|
/// </summary>
|
|
public void AddFilters(McpMessageFilters messageFilters)
|
|
{
|
|
// Track outgoing legacy JSON-RPC requests and InputRequiredResult responses.
|
|
messageFilters.OutgoingFilters.Add(next => async (context, cancellationToken) =>
|
|
{
|
|
if (context.JsonRpcMessage is JsonRpcRequest request && LegacyMrtrMethods.Contains(request.Method))
|
|
{
|
|
_legacyRequestMethods.Add(request.Method);
|
|
}
|
|
else if (context.JsonRpcMessage is JsonRpcResponse response &&
|
|
response.Result is JsonObject resultObj &&
|
|
resultObj.TryGetPropertyValue("resultType", out var resultTypeNode) &&
|
|
resultTypeNode?.GetValue<string>() == "input_required")
|
|
{
|
|
Interlocked.Increment(ref _incompleteResultCount);
|
|
}
|
|
|
|
await next(context, cancellationToken);
|
|
});
|
|
|
|
// Track incoming MRTR retries (requests with inputResponses or requestState in params).
|
|
messageFilters.IncomingFilters.Add(next => async (context, cancellationToken) =>
|
|
{
|
|
if (context.JsonRpcMessage is JsonRpcRequest request &&
|
|
request.Params is JsonObject paramsObj &&
|
|
(paramsObj.ContainsKey("inputResponses") || paramsObj.ContainsKey("requestState")))
|
|
{
|
|
Interlocked.Increment(ref _mrtrRetryCount);
|
|
}
|
|
|
|
await next(context, cancellationToken);
|
|
});
|
|
}
|
|
|
|
/// <summary>
|
|
/// Asserts that MRTR was used: at least one InputRequiredResult response was sent
|
|
/// and no legacy JSON-RPC requests (elicitation/create, sampling/createMessage, roots/list) were sent.
|
|
/// </summary>
|
|
public void AssertMrtrUsed()
|
|
{
|
|
Assert.True(_incompleteResultCount > 0,
|
|
"Expected at least one InputRequiredResult response (MRTR mode), but none were detected.");
|
|
Assert.Empty(_legacyRequestMethods);
|
|
}
|
|
|
|
/// <summary>
|
|
/// Asserts that MRTR was used at least once (at least one InputRequiredResult response was sent),
|
|
/// independent of whether the session also issued any legacy server-to-client requests.
|
|
/// </summary>
|
|
public void AssertMrtrUsedAtLeastOnce()
|
|
{
|
|
Assert.True(_incompleteResultCount > 0,
|
|
"Expected at least one InputRequiredResult response (MRTR mode), but none were detected.");
|
|
}
|
|
|
|
/// <summary>
|
|
/// Asserts that legacy mode was used: at least one legacy JSON-RPC request was sent
|
|
/// and no MRTR retries or InputRequiredResult responses were detected.
|
|
/// </summary>
|
|
public void AssertMrtrNotUsed()
|
|
{
|
|
Assert.NotEmpty(_legacyRequestMethods);
|
|
Assert.Equal(0, _mrtrRetryCount);
|
|
Assert.Equal(0, _incompleteResultCount);
|
|
}
|
|
}
|