Rename project and namespaces

This commit is contained in:
Eirik Tsarpalis 2025-03-19 10:51:27 +02:00
commit c9ed453745
3 changed files with 426 additions and 918 deletions

View file

@ -0,0 +1,382 @@
using System.Collections.Concurrent;
using System.Net;
using System.Text.Json;
using ModelContextProtocol.Protocol.Messages;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Logging.Abstractions;
namespace ModelContextProtocol.Tests.Utils;
public sealed class InMemoryTestSseServer : IAsyncDisposable
{
private readonly HttpListener _listener;
private readonly CancellationTokenSource _cts;
private readonly ILogger<InMemoryTestSseServer> _logger;
private Task? _serverTask;
private readonly TaskCompletionSource _connectionEstablished = new();
// SSE endpoint for GET
private readonly string _endpointPath;
// POST endpoint
private readonly string _messagePath;
// Keep track of all open SSE connections (StreamWriters).
private readonly ConcurrentBag<StreamWriter> _sseClients = [];
public InMemoryTestSseServer(int port = 5000, ILogger<InMemoryTestSseServer>? logger = null)
{
_listener = new HttpListener();
_listener.Prefixes.Add($"http://localhost:{port}/");
_cts = new CancellationTokenSource();
_logger = logger ?? NullLogger<InMemoryTestSseServer>.Instance;
_endpointPath = "/sse";
_messagePath = "/message";
}
/// <summary>
/// This is to be able to use the full URL for the endpoint event.
/// </summary>
public bool UseFullUrlForEndpointEvent { get; set; }
/// <summary>
/// Full URL for the SSE endpoint, e.g. "http://localhost:5000/sse"
/// </summary>
public string SseEndpoint
=> $"http://localhost:{_listener.Prefixes.First().Split(':')[2].TrimEnd('/')}{_endpointPath}";
/// <summary>
/// Full URL for the message endpoint, e.g. "http://localhost:5000/message"
/// </summary>
public string MessageEndpoint
=> $"http://localhost:{_listener.Prefixes.First().Split(':')[2].TrimEnd('/')}{_messagePath}";
/// <summary>
/// Starts the server so it can accept incoming connections and POST requests.
/// </summary>
public async Task StartAsync()
{
_listener.Start();
_serverTask = HandleConnectionsAsync(_cts.Token);
_logger.LogInformation("Test SSE server started on {Endpoint}", SseEndpoint);
await Task.CompletedTask;
}
private async Task HandleConnectionsAsync(CancellationToken ct)
{
try
{
while (!ct.IsCancellationRequested)
{
var context = await _listener.GetContextAsync();
_ = Task.Run(() => HandleRequestAsync(context, ct), ct);
}
}
catch (OperationCanceledException)
{
// Ignore, we are shutting down
}
catch (Exception ex) when (!ct.IsCancellationRequested)
{
_logger.LogError(ex, "Error in SSE server connection handling");
}
}
private async Task HandleRequestAsync(HttpListenerContext context, CancellationToken ct)
{
var request = context.Request;
var response = context.Response;
// Handle SSE endpoint
if (request.HttpMethod.Equals("GET", StringComparison.OrdinalIgnoreCase)
&& request.Url?.AbsolutePath.Equals(_endpointPath, StringComparison.OrdinalIgnoreCase) == true)
{
await HandleSseConnectionAsync(context, ct);
}
// Handle POST /message
else if (request.HttpMethod.Equals("POST", StringComparison.OrdinalIgnoreCase)
&& request.Url?.AbsolutePath.Equals(_messagePath, StringComparison.OrdinalIgnoreCase) == true)
{
await HandlePostMessageAsync(context, ct);
}
else
{
response.StatusCode = 404;
response.Close();
}
}
/// <summary>
/// Handle Server-Sent Events (SSE) connection.
/// Send the initial event: endpoint with the full POST URL.
/// Keep the connection open until the server is disposed or the client disconnects.
/// </summary>
private async Task HandleSseConnectionAsync(HttpListenerContext context, CancellationToken ct)
{
var response = context.Response;
response.ContentType = "text/event-stream";
response.Headers["Cache-Control"] = "no-cache";
// Ensures the response is never chunked away by the framework
response.SendChunked = true;
response.StatusCode = (int)HttpStatusCode.OK;
using var writer = new StreamWriter(response.OutputStream);
_sseClients.Add(writer);
// Immediately send the "endpoint" event with the POST URL
await writer.WriteLineAsync("event: endpoint");
if (UseFullUrlForEndpointEvent)
{
await writer.WriteLineAsync($"data: {MessageEndpoint}");
}
else
{
await writer.WriteLineAsync($"data: {_messagePath}");
}
await writer.WriteLineAsync(); // blank line to end an SSE message
await writer.FlushAsync(ct);
_logger.LogInformation("New SSE client connected.");
_connectionEstablished.TrySetResult(); // Signal connection is ready
try
{
// Keep the connection open by "pinging" or just waiting
// until the client disconnects or the server is canceled.
while (!ct.IsCancellationRequested && response.OutputStream.CanWrite)
{
_logger.LogDebug("SSE connection alive check");
// Optionally do a periodic no-op to keep connection alive:
await writer.WriteLineAsync(": keep-alive");
await writer.FlushAsync(ct);
await Task.Delay(TimeSpan.FromSeconds(5), ct);
}
}
catch (TaskCanceledException)
{
// This is expected on shutdown
_logger.LogInformation("SSE connection cancelled (expected on shutdown)");
}
catch (IOException)
{
// Client likely disconnected
_logger.LogInformation("SSE client disconnected");
}
finally
{
// Remove this writer from bag (we're disposing it anyway)
_sseClients.TryTake(out _);
_logger.LogInformation("SSE client disconnected.");
}
}
// Add method to wait for connection
public Task WaitForConnectionAsync(TimeSpan timeout) =>
_connectionEstablished.Task.WaitAsync(timeout);
/// <summary>
/// Handle POST /message endpoint.
/// Echo the content back to the caller and broadcast it over SSE as well.
/// </summary>
private async Task HandlePostMessageAsync(HttpListenerContext context, CancellationToken cancellationToken)
{
var request = context.Request;
var response = context.Response;
try
{
using var reader = new StreamReader(request.InputStream);
string content = await reader.ReadToEndAsync(cancellationToken);
var jsonRpcNotification = JsonSerializer.Deserialize<JsonRpcNotification>(content);
if (jsonRpcNotification != null && jsonRpcNotification.Method != "initialize")
{
// Test server so just ignore notifications
// Set status code to 202 Accepted
response.StatusCode = 202;
// Write "accepted" as the response body
using var writer = new StreamWriter(response.OutputStream);
await writer.WriteAsync("accepted");
await writer.FlushAsync(cancellationToken);
return;
}
var jsonRpcRequest = JsonSerializer.Deserialize<JsonRpcRequest>(content);
if (jsonRpcRequest != null)
{
if (jsonRpcRequest.Method == "initialize")
{
await HandleInitializationRequest(response, jsonRpcRequest);
}
else
{
// Set status code to 202 Accepted
response.StatusCode = 202;
// Write "accepted" as the response body
using var writer = new StreamWriter(response.OutputStream);
await writer.WriteAsync("accepted");
await writer.FlushAsync(cancellationToken);
// Then process the message and send the response via SSE
if (jsonRpcRequest != null)
{
// Process the request and generate a response
var responseMessage = await ProcessRequestAsync(jsonRpcRequest);
// Send the response via the SSE connection instead of HTTP response
await SendSseMessageAsync(responseMessage);
}
}
}
}
catch (Exception ex)
{
await SendJsonRpcErrorAsync(response, null, -32603, "Internal error", ex.Message);
response.StatusCode = 500;
}
finally
{
response.Close();
}
}
private static Task<JsonRpcResponse> ProcessRequestAsync(JsonRpcRequest jsonRpcRequest)
{
// This is a test server so we just echo back the request
return Task.FromResult(new JsonRpcResponse
{
Id = jsonRpcRequest.Id,
Result = jsonRpcRequest.Params!
});
}
private async Task SendSseMessageAsync(JsonRpcResponse jsonRpcResponse)
{
// This is a test server so we just send to all connected clients
await BroadcastMessageAsync(JsonSerializer.Serialize(jsonRpcResponse));
}
private static async Task SendJsonRpcErrorAsync(HttpListenerResponse response, RequestId? id, int code, string message, string? data = null)
{
var errorResponse = new JsonRpcError
{
Id = id ?? RequestId.FromString("error"),
JsonRpc = "2.0",
Error = new JsonRpcErrorDetail
{
Code = code,
Message = message,
Data = data
}
};
response.StatusCode = 200; // Always 200 for JSON-RPC
response.ContentType = "application/json";
await using var writer = new StreamWriter(response.OutputStream);
await writer.WriteAsync(JsonSerializer.Serialize(errorResponse));
}
private static async Task HandleInitializationRequest(HttpListenerResponse response, JsonRpcRequest jsonRpcRequest)
{
// We don't need to validate the client's initialization request for the test
// Just send back a valid server initialization response
var jsonRpcResponse = new JsonRpcResponse()
{
Id = jsonRpcRequest.Id,
Result = new
{
protocolVersion = "2024-11-05",
capabilities = new
{
experimental = (object?)null,
roots = (object?)null,
sampling = (object?)null
},
serverInfo = new
{
name = "ExampleServer",
version = "1.0.0"
}
}
};
// Echo back to the HTTP response
response.StatusCode = 200;
response.ContentType = "application/json";
await using var writer = new StreamWriter(response.OutputStream);
await writer.WriteAsync(JsonSerializer.Serialize(jsonRpcResponse));
}
/// <summary>
/// Broadcast a message to all currently connected SSE clients.
/// </summary>
/// <param name="message">The raw string to send</param>
public async Task BroadcastMessageAsync(string message)
{
foreach (var client in _sseClients.ToArray()) // ToArray to avoid mutation issues
{
try
{
// SSE requires "event: <name>" + "data: <payload>" + blank line
await client.WriteLineAsync("event: message");
await client.WriteLineAsync($"data: {message}");
await client.WriteLineAsync();
await client.FlushAsync();
}
catch (IOException)
{
// Client may have disconnected. We let them get cleaned up on next iteration.
}
catch (ObjectDisposedException)
{
// Stream is disposed, ignore.
}
}
}
public async ValueTask DisposeAsync()
{
await _cts.CancelAsync();
if (_serverTask != null)
{
try
{
await Task.WhenAny(_serverTask, Task.Delay(2000));
}
catch (TaskCanceledException)
{
// ignore
}
}
_listener.Close();
_cts.Dispose();
_logger.LogInformation("Test SSE server stopped");
}
/// <summary>
/// Send a test notification to all connected SSE clients.
/// </summary>
/// <param name="content"></param>
/// <returns></returns>
public async Task SendTestNotificationAsync(string content)
{
var notification = new JsonRpcNotification
{
JsonRpc = "2.0",
Method = "test/notification",
Params = new { message = content }
};
var serialized = JsonSerializer.Serialize(notification);
await BroadcastMessageAsync(serialized);
}
}

View file

@ -0,0 +1,22 @@
namespace ModelContextProtocol.Tests.Utils;
public class MockHttpHandler : HttpMessageHandler
{
public Func<HttpRequestMessage, Task<HttpResponseMessage>>? RequestHandler { get; set; }
protected async override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
{
if (RequestHandler == null)
throw new InvalidOperationException($"No {nameof(RequestHandler)} was set! Please set handler first and make request afterwards.");
if (cancellationToken.IsCancellationRequested)
throw new OperationCanceledException(cancellationToken);
var result = await RequestHandler.Invoke(request);
if (cancellationToken.IsCancellationRequested)
throw new OperationCanceledException(cancellationToken);
return result;
}
}

View file

@ -0,0 +1,91 @@
using System.Threading.Channels;
using ModelContextProtocol.Protocol.Messages;
using ModelContextProtocol.Protocol.Transport;
namespace ModelContextProtocol.Tests.Utils;
public class TestServerTransport : IServerTransport
{
private readonly Channel<IJsonRpcMessage> _messageChannel;
private bool _isStarted;
public bool IsConnected => _isStarted;
public ChannelReader<IJsonRpcMessage> MessageReader => _messageChannel;
public List<IJsonRpcMessage> SentMessages { get; } = [];
public Action<IJsonRpcMessage>? OnMessageSent { get; set; }
public TestServerTransport()
{
_messageChannel = Channel.CreateUnbounded<IJsonRpcMessage>(new UnboundedChannelOptions
{
SingleReader = true,
SingleWriter = true,
});
}
public ValueTask DisposeAsync() => ValueTask.CompletedTask;
public async Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default)
{
SentMessages.Add(message);
if (message is JsonRpcRequest request)
{
if (request.Method == "roots/list")
await ListRoots(request, cancellationToken);
else if (request.Method == "sampling/createMessage")
await Sampling(request, cancellationToken);
else
await WriteMessageAsync(request, cancellationToken);
}
else if (message is JsonRpcNotification notification)
{
await WriteMessageAsync(notification, cancellationToken);
}
OnMessageSent?.Invoke(message);
}
public Task StartListeningAsync(CancellationToken cancellationToken = default)
{
_isStarted = true;
return Task.CompletedTask;
}
private async Task ListRoots(JsonRpcRequest request, CancellationToken cancellationToken)
{
await WriteMessageAsync(new JsonRpcResponse
{
Id = request.Id,
Result = new ModelContextProtocol.Protocol.Types.ListRootsResult
{
Roots = []
}
}, cancellationToken);
}
private async Task Sampling(JsonRpcRequest request, CancellationToken cancellationToken)
{
await WriteMessageAsync(new JsonRpcResponse
{
Id = request.Id,
Result = new Protocol.Types.CreateMessageResult { Content = new(), Model = "model", Role = "role" }
}, cancellationToken);
}
private async Task Error(JsonRpcRequest request, CancellationToken cancellationToken)
{
await WriteMessageAsync(new JsonRpcError
{
Id = request.Id,
Error = new JsonRpcErrorDetail() { Code = -32601, Message = $"Method '{request.Method}' not supported" }
}, cancellationToken);
}
protected async Task WriteMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default)
{
await _messageChannel.Writer.WriteAsync(message, cancellationToken);
}
}