Sampling capability (#7)

src/mcpdotnet.Tests/IntegrationTests.cs

@@ -341,4 +341,69 @@ public class IntegrationTests
         Assert.Single(result.Contents);
         Assert.NotNull(result.Contents[0].Blob);
     }
+
+    [Fact]
+    public async Task Sampling_Stdio_EverythingServer()
+    {
+        // arrange
+        var config = new McpServerConfig
+        {
+            Id = "everything",
+            Name = "Everything",
+            TransportType = "stdio",
+            TransportOptions = new Dictionary<string, string>
+            {
+                ["command"] = "npx",
+                ["arguments"] = "-y @modelcontextprotocol/server-everything"
+            }
+        };
+
+        var options = new McpClientOptions
+        {
+            ClientInfo = new() { Name = "IntegrationTestClient", Version = "1.0.0" },
+            Capabilities = new ClientCapabilities
+            {
+                Sampling = new()
+            }
+        };
+
+        var factory = new McpClientFactory(
+            new[] { config },
+            options
+        );
+        var client = await factory.GetClientAsync("everything");
+
+        // Set up the sampling handler
+        int samplingHandlerCalls = 0;
+        client.SamplingHandler = async (_) =>
+        {
+            samplingHandlerCalls++;
+            return new CreateMessageResult
+            {
+                Model = "test-model",
+                Role = "assistant",
+                Content = new Content
+                {
+                    Type = "text",
+                    Text = "Test response"
+                }
+            };
+        };
+
+        // Call the server's sampleLLM tool which should trigger our sampling handler
+        var result = await client.CallToolAsync(
+            "sampleLLM",
+            new Dictionary<string, object> 
+            {
+                ["prompt"] = "Test prompt", 
+                ["maxTokens"] = 100 
+            }
+        );
+
+        // assert
+        Assert.NotNull(result);
+        var textContent = Assert.Single(result.Content);
+        Assert.Equal("text", textContent.Type);
+        Assert.Contains("Test response", textContent.Text);
+    }
 }

src/mcpdotnet/Client/IMcpClient.cs

@@ -25,6 +25,11 @@ public interface IMcpClient : IAsyncDisposable
     /// </summary>
     Implementation? ServerInfo { get; }
 
+    /// <summary>
+    /// Gets or sets the handler for the server's sampling requests.
+    /// </summary>
+    Func<CreateMessageRequestParams, Task<CreateMessageResult>>? SamplingHandler { get; set; }
+
     /// <summary>
     /// Establishes a connection to the server.
     /// </summary>

src/mcpdotnet/Client/McpClient.cs

@@ -7,6 +7,8 @@ using McpDotNet.Protocol.Messages;
 using System.Diagnostics;
 using System.Text.Json;
 using McpDotNet.Utils.Json;
+using System.Text;
+using System.Threading;
 
 /// <inheritdoc/>
 internal class McpClient : IMcpClient
@@ -15,6 +17,7 @@ internal class McpClient : IMcpClient
     private readonly McpClientOptions _options;
     private readonly ConcurrentDictionary<RequestId, TaskCompletionSource<IJsonRpcMessage>> _pendingRequests;
     private readonly ConcurrentDictionary<string, List<Func<JsonRpcNotification,Task>>> _notificationHandlers;
+    private readonly Dictionary<string, Func<JsonRpcRequest, Task<object>>> _requestHandlers = new();
     private int _nextRequestId;
     private bool _isInitialized;
     private Task? _messageProcessingTask;
@@ -42,6 +45,17 @@ internal class McpClient : IMcpClient
         _pendingRequests = new();
         _notificationHandlers = new();
         _nextRequestId = 1;
+
+        if (options.Capabilities?.Sampling != null)
+        {
+            SetRequestHandler<CreateMessageRequestParams, CreateMessageResult>("sampling/createMessage",
+                async (request) => {
+                    if (SamplingHandler == null)
+                        throw new McpClientException("Sampling handler not configured");
+
+                    return await SamplingHandler(request);
+                });
+        }
     }
 
     /// <inheritdoc/>
@@ -271,7 +285,7 @@ internal class McpClient : IMcpClient
         {
             await foreach (var message in _transport.MessageReader.ReadAllAsync(cancellationToken).ConfigureAwait(false))
             {
-                await HandleMessageAsync(message).ConfigureAwait(false);
+                await HandleMessageAsync(message, cancellationToken).ConfigureAwait(false);
             }
         }
         catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
@@ -284,10 +298,39 @@ internal class McpClient : IMcpClient
         }
     }
 
-    private async Task HandleMessageAsync(IJsonRpcMessage message)
+    private async Task HandleMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken)
     {
         switch (message)
         {
+            case JsonRpcRequest request:
+                if (_requestHandlers.TryGetValue(request.Method, out var handler))
+                {
+                    try
+                    {
+                        var result = await handler(request);
+                        await _transport.SendMessageAsync(new JsonRpcResponse
+                        {
+                            Id = request.Id,
+                            JsonRpc = "2.0",
+                            Result = result
+                        }, cancellationToken);
+                    }
+                    catch (Exception ex)
+                    {
+                        // Send error response
+                        await _transport.SendMessageAsync(new JsonRpcError
+                        {
+                            Id = request.Id,
+                            JsonRpc = "2.0",
+                            Error = new JsonRpcErrorDetail
+                            {
+                                Code = -32000,  // Implementation defined error
+                                Message = ex.Message
+                            }
+                        }, cancellationToken);
+                    }
+                }
+                break;
             case IJsonRpcMessageWithId messageWithId:
                 if (_pendingRequests.TryRemove(messageWithId.Id, out var tcs))
                 {
@@ -298,11 +341,11 @@ internal class McpClient : IMcpClient
             case JsonRpcNotification notification:
                 if (_notificationHandlers.TryGetValue(notification.Method, out var handlers))
                 {
-                    foreach (var handler in handlers)
+                    foreach (var notificationHandler in handlers)
                     {
                         try
                         {
-                            await handler(notification).ConfigureAwait(false);
+                            await notificationHandler(notification).ConfigureAwait(false);
                         }
                         catch (Exception ex)
                         {
@@ -383,6 +426,18 @@ internal class McpClient : IMcpClient
         }
     }
 
+    public Func<CreateMessageRequestParams, Task<CreateMessageResult>>? SamplingHandler { get; set; }
+
+    private void SetRequestHandler<TRequest, TResponse>(string method, Func<TRequest, Task<TResponse>> handler)
+    {
+        _requestHandlers[method] = async (request) => {
+            // Convert the params JsonElement to our type using the same options
+            var jsonString = JsonSerializer.Serialize(request.Params);
+            var typedRequest = JsonSerializer.Deserialize<TRequest>(jsonString, _jsonOptions);
+            return await handler(typedRequest);
+        };
+    }
+
     private async Task CleanupAsync()
     {
         _cts?.Cancel();
@@ -419,3 +474,4 @@ internal class McpClient : IMcpClient
         GC.SuppressFinalize(this);
     }
 }
+

src/mcpdotnet/Protocol/Messages/JsonRpcResponse.cs

@@ -1,7 +1,7 @@
-namespace McpDotNet.Protocol.Messages;
-
+
 using System.Text.Json.Serialization;
 
+namespace McpDotNet.Protocol.Messages;
 /// <summary>
 /// A successful response message in the JSON-RPC protocol.
 /// </summary>

src/mcpdotnet/Protocol/Transport/StdioTransport.cs

@@ -142,6 +142,7 @@ public sealed class StdioTransport : TransportBase
         try
         {
             var json = JsonSerializer.Serialize(message, _jsonOptions);
+            Debug.WriteLine($"Sending message: {json}");
 
             // Write the message followed by a newline
             await _process!.StandardInput.WriteLineAsync(json.AsMemory(), cancellationToken);

src/mcpdotnet/Protocol/Types/Content.cs

@@ -1,4 +1,6 @@
-namespace McpDotNet.Protocol.Types;
+using System.Text.Json.Serialization;
+
+namespace McpDotNet.Protocol.Types;
 
 /// <summary>
 /// Represents the content of a tool response.
@@ -10,26 +12,32 @@ public class Content
     /// <summary>
     /// The type of content. This determines the structure of the content object. Can be "image", "text", "resource".
     /// </summary>
+
+    [JsonPropertyName("type")]
     public string Type { get; set; } = string.Empty;
 
     /// <summary>
     /// The text content of the message.
     /// </summary>
+    [JsonPropertyName("text")]
     public string? Text { get; set; }
 
     /// <summary>
     /// The base64-encoded image data.
     /// </summary>
+    [JsonPropertyName("data")]
     public string? Data { get; set; }
 
     /// <summary>
     /// The MIME type of the image.
     /// </summary>
+    [JsonPropertyName("mime_type")]
     public string? MimeType { get; set; }
 
     /// <summary>
     /// The resource content of the message (if embedded).
     /// </summary>
+    [JsonPropertyName("resource")]
     public ResourceContents? Resource { get; set; }
 
     /// <summary>

src/mcpdotnet/Protocol/Types/ContextInclusion.cs

@@ -0,0 +1,11 @@
+namespace McpDotNet.Protocol.Types;
+
+/// <summary>
+/// A request to include context from one or more MCP servers (including the caller), to be attached to the prompt.
+/// </summary>
+public enum ContextInclusion
+{
+    None,
+    ThisServer,
+    AllServers
+}

src/mcpdotnet/Protocol/Types/CreateMessageRequestParams.cs

@@ -0,0 +1,52 @@
+namespace McpDotNet.Protocol.Types;
+
+/// <summary>
+/// A request from the server to sample an LLM via the client. 
+/// The client has full discretion over which model to select. 
+/// The client should also inform the user before beginning sampling, to allow them to inspect the request (human in the loop) and decide whether to approve it.
+/// 
+/// While these align with the protocol specification,
+/// clients have full discretion over model selection and should inform users before sampling.
+/// </summary>
+public class CreateMessageRequestParams
+{
+    /// <summary>
+    /// A request to include context from one or more MCP servers (including the caller), to be attached to the prompt. The client MAY ignore this request.
+    /// </summary>
+    public ContextInclusion? IncludeContext { get; init; }
+
+    /// <summary>
+    /// The maximum number of tokens to sample, as requested by the server. The client MAY choose to sample fewer tokens than requested.
+    /// </summary>
+    public int MaxTokens { get; init; }
+
+    /// <summary>
+    /// Messages requested by the server to be included in the prompt.
+    /// </summary>
+    public required IReadOnlyList<SamplingMessage> Messages { get; init; }
+
+    /// <summary>
+    /// Optional metadata to pass through to the LLM provider. The format of this metadata is provider-specific.
+    /// </summary>
+    public object? Metadata { get; init; }
+
+    /// <summary>
+    /// The server's preferences for which model to select. The client MAY ignore these preferences.
+    /// </summary>
+    public ModelPreferences? ModelPreferences { get; init; }
+
+    /// <summary>
+    /// Optional stop sequences that the server wants to use for sampling.
+    /// </summary>
+    public IReadOnlyList<string>? StopSequences { get; init; }
+
+    /// <summary>
+    /// An optional system prompt the server wants to use for sampling. The client MAY modify or omit this prompt.
+    /// </summary>
+    public string? SystemPrompt { get; init; }
+
+    /// <summary>
+    /// The temperature to use for sampling, as requested by the server.
+    /// </summary>
+    public float? Temperature { get; init; }
+}

src/mcpdotnet/Protocol/Types/CreateMessageResult.cs

@@ -0,0 +1,35 @@
+using System.Text.Json.Serialization;
+
+namespace McpDotNet.Protocol.Types;
+
+/// <summary>
+/// The client's response to a sampling/create_message request from the server. 
+/// The client should inform the user before returning the sampled message, to allow them to inspect the response (human in the loop) 
+///  and decide whether to allow the server to see it.
+/// </summary>
+public class CreateMessageResult
+{
+    /// <summary>
+    /// Text or image content of the message.
+    /// </summary>
+    [JsonPropertyName("content")]
+    public required Content Content { get; init; }
+
+    /// <summary>
+    /// The name of the model that generated the message.
+    /// </summary>
+    [JsonPropertyName("model")]
+    public required string Model { get; init; }
+
+    /// <summary>
+    /// The reason why sampling stopped, if known.
+    /// </summary>
+    [JsonPropertyName("stop_reason")]
+    public string? StopReason { get; init; }
+
+    /// <summary>
+    /// The role of the user who generated the message.
+    /// </summary>
+    [JsonPropertyName("role")]
+    public required string Role { get; init; }
+}

src/mcpdotnet/Protocol/Types/ModelHint.cs

@@ -0,0 +1,22 @@
+namespace McpDotNet.Protocol.Types;
+
+/// <summary>
+/// Hints to use for model selection.
+/// Keys not declared here are currently left unspecified by the spec and are up
+/// to the client to interpret.
+/// </summary>
+public class ModelHint
+{
+    /// <summary>
+    /// A hint for a model name.
+    /// 
+    /// The client SHOULD treat this as a substring of a model name; for example:
+    /// - `claude-3-5-sonnet` should match `claude-3-5-sonnet-20241022`
+    /// - `sonnet` should match `claude-3-5-sonnet-20241022`, `claude-3-sonnet-20240229`, etc.
+    /// - `claude` should match any Claude model
+    /// 
+    /// The client MAY also map the string to a different provider's model name or a different model family, as long as it fills a similar niche; for example:
+    /// - `gemini-1.5-flash` could match `claude-3-haiku-20240307`
+    /// </summary>
+    public string? Name { get; init; }
+}

src/mcpdotnet/Protocol/Types/ModelPreferences.cs

@@ -0,0 +1,82 @@
+namespace McpDotNet.Protocol.Types;
+
+/// <summary>
+/// The server's preferences for model selection, requested of the client during sampling.
+/// Because LLMs can vary along multiple dimensions, choosing the \"best\" model is
+/// rarely straightforward.  Different models excel in different areas—some are
+/// faster but less capable, others are more capable but more expensive, and so
+/// on. This interface allows servers to express their priorities across multiple
+/// dimensions to help clients make an appropriate selection for their use case.
+/// 
+/// These preferences are always advisory. The client MAY ignore them. It is also
+/// up to the client to decide how to interpret these preferences and how to
+/// balance them against other considerations.
+/// </summary>
+public class ModelPreferences
+{
+    /// <summary>
+    /// How much to prioritize cost when selecting a model. A value of 0 means cost\nis not important, while a value of 1 means cost is the most important
+    /// factor.
+    /// </summary>
+    public float? CostPriority { get; init; }
+
+    /// <summary>
+    /// Optional hints to use for model selection.
+    /// 
+    /// If multiple hints are specified, the client MUST evaluate them in order
+    /// (such that the first match is taken).
+    /// 
+    /// The client SHOULD prioritize these hints over the numeric priorities, but
+    /// MAY still use the priorities to select from ambiguous matches.
+    /// </summary>
+    public IReadOnlyList<ModelHint>? Hints { get; init; }
+
+    /// <summary>
+    /// How much to prioritize intelligence and capabilities when selecting a
+    /// model. A value of 0 means intelligence is not important, while a value of 1
+    /// means intelligence is the most important factor.
+    /// </summary>
+    public float? SpeedPriority { get; init; }
+
+    /// <summary>
+    /// How much to prioritize sampling speed (latency) when selecting a model. A
+    /// value of 0 means speed is not important, while a value of 1 means speed is
+    /// the most important factor.
+    /// </summary>
+    public float? IntelligencePriority { get; init; }
+
+    /// <summary>
+    /// Validates the model preferences.
+    /// </summary>
+    /// <param name="errorMessage">Error message if object isn't valid</param>
+    /// <returns>True if valid, false if invalid</returns>
+    public bool Validate(out string errorMessage)
+    {
+        bool valid = true;
+        List<string> errors = new();
+        if (CostPriority < 0 || CostPriority > 1)
+        {
+            errors.Add("CostPriority must be between 0 and 1");
+            valid = false;
+        }
+        if (SpeedPriority < 0 || SpeedPriority > 1)
+        {
+            errors.Add("SpeedPriority must be between 0 and 1");
+            valid = false;
+        }
+        if (IntelligencePriority < 0 || IntelligencePriority > 1)
+        {
+            errors.Add("IntelligencePriority must be between 0 and 1");
+            valid = false;
+        }
+        if (!valid)
+        {
+            errorMessage = string.Join(", ", errors);
+        }
+        else
+        {
+            errorMessage = "";
+        }
+        return valid;
+    }
+}

src/mcpdotnet/Protocol/Types/ResourceContents.cs

@@ -1,4 +1,6 @@
-namespace McpDotNet.Protocol.Types;
+using System.Text.Json.Serialization;
+
+namespace McpDotNet.Protocol.Types;
 
 /// <summary>
 /// Represents the content of a resource.
@@ -8,21 +10,25 @@ public class ResourceContents
     /// <summary>
     /// The URI of the resource.
     /// </summary>
+    [JsonPropertyName("uri")]
     public string Uri { get; set; } = string.Empty;
 
     /// <summary>
     /// The type of content.
     /// </summary>
+    [JsonPropertyName("mime_type")]
     public string? MimeType { get; set; }
 
     /// <summary>
     /// The text content of the resource.
     /// </summary>
+    [JsonPropertyName("text")]
     public string? Text { get; set; }
 
 
     /// <summary>
     /// The base64-encoded binary content of the resource.
     /// </summary>
+    [JsonPropertyName("blob")]
     public string? Blob { get; set; }
 }

src/mcpdotnet/Protocol/Types/SamplingMessage.cs

@@ -0,0 +1,18 @@
+namespace McpDotNet.Protocol.Types;
+
+/// <summary>
+/// Describes a message issued to or received from an LLM API.
+/// </summary>
+public class SamplingMessage
+{
+    /// <summary>
+    /// Text or image content of the message.
+    /// </summary>
+    public required Content Content { get; init; }
+
+    /// <summary>
+    /// The role of the message ("user" or "assistant").
+    /// </summary>
+
+    public required Role Role { get; init; }  // "user" or "assistant"
+}

src/mcpdotnet/mcpdotnet.csproj

@@ -9,7 +9,7 @@
   <PropertyGroup>
     <!-- license and package properties -->
     <PackageId>mcpdotnet</PackageId>
-    <Version>0.2.0-alpha1</Version>
+    <Version>0.3.0.1</Version>
     <Authors>PederHP</Authors>
     <Description>.NET client library for the Model Context Protocol (MCP)</Description>
     <PackageProjectUrl>https://github.com/PederHP/mcpdotnet</PackageProjectUrl>