Add support for sampling (#119)

* Add UnitInterval type

* Add test coverage for UnitInterval

* Add support for sampling

* Add integration tests for sampling
This commit is contained in:
Mattt 2025-05-26 04:57:24 -07:00 • committed by GitHub
commit fac4df6b1a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 1584 additions and 0 deletions

126
README.md
View file

@ -189,6 +189,89 @@ for message in messages {
}
```
### Sampling
Sampling allows servers to request LLM completions through the client,
enabling agentic behaviors while maintaining human-in-the-loop control.
Clients register a handler to process incoming sampling requests from servers.
> [!TIP]
> Sampling requests flow from **server to client**,
> not client to server.
> This enables servers to request AI assistance
> while clients maintain control over model access and user approval.
```swift
// Register a sampling handler in the client
await client.withSamplingHandler { parameters in
// Review the sampling request (human-in-the-loop step 1)
print("Server requests completion for: \(parameters.messages)")
// Optionally modify the request based on user input
var messages = parameters.messages
if let systemPrompt = parameters.systemPrompt {
print("System prompt: \(systemPrompt)")
}
// Sample from your LLM (this is where you'd call your AI service)
let completion = try await callYourLLMService(
messages: messages,
maxTokens: parameters.maxTokens,
temperature: parameters.temperature
)
// Review the completion (human-in-the-loop step 2)
print("LLM generated: \(completion)")
// User can approve, modify, or reject the completion here
// Return the result to the server
return CreateSamplingMessage.Result(
model: "your-model-name",
stopReason: .endTurn,
role: .assistant,
content: .text(completion)
)
}
```
The sampling flow follows these steps:
```mermaid
sequenceDiagram
participant S as MCP Server
participant C as MCP Client
participant U as User/Human
participant L as LLM Service
Note over S,L: Server-initiated sampling request
S->>C: sampling/createMessage request
Note right of S: Server needs AI assistance<br/>for decision or content
Note over C,U: Human-in-the-loop review #1
C->>U: Show sampling request
U->>U: Review & optionally modify<br/>messages, system prompt
U->>C: Approve request
Note over C,L: Client handles LLM interaction
C->>L: Send messages to LLM
L->>C: Return completion
Note over C,U: Human-in-the-loop review #2
C->>U: Show LLM completion
U->>U: Review & optionally modify<br/>or reject completion
U->>C: Approve completion
Note over C,S: Return result to server
C->>S: sampling/createMessage response
Note left of C: Contains model used,<br/>stop reason, final content
Note over S: Server continues with<br/>AI-assisted result
```
This human-in-the-loop design ensures that users
maintain control over what the LLM sees and generates,
even when servers initiate the requests.
### Error Handling
Handle common client errors:
@ -504,6 +587,49 @@ server.withMethodHandler(GetPrompt.self) { params in
}
```
### Sampling
Servers can request LLM completions from clients through sampling. This enables agentic behaviors where servers can ask for AI assistance while maintaining human oversight.
> [!NOTE]
> The current implementation provides the correct API design for sampling, but requires bidirectional communication support in the transport layer. This feature will be fully functional when bidirectional transport support is added.
```swift
// Enable sampling capability in server
let server = Server(
name: "MyModelServer",
version: "1.0.0",
capabilities: .init(
sampling: .init(), // Enable sampling capability
tools: .init(listChanged: true)
)
)
// Request sampling from the client (conceptual - requires bidirectional transport)
do {
let result = try await server.requestSampling(
messages: [
Sampling.Message(role: .user, content: .text("Analyze this data and suggest next steps"))
],
systemPrompt: "You are a helpful data analyst",
maxTokens: 150,
temperature: 0.7
)
// Use the LLM completion in your server logic
print("LLM suggested: \(result.content)")
} catch {
print("Sampling request failed: \(error)")
}
```
Sampling enables powerful agentic workflows:
- **Decision-making**: Ask the LLM to choose between options
- **Content generation**: Request drafts for user approval
- **Data analysis**: Get AI insights on complex data
- **Multi-step reasoning**: Chain AI completions with tool calls
#### Initialize Hook
Control client connections with an initialize hook:

View file

@ -0,0 +1,126 @@
/// A value constrained to the range 0.0 to 1.0, inclusive.
///
/// `UnitInterval` represents a normalized value that is guaranteed to be within
/// the unit interval [0, 1]. This type is commonly used for representing
/// priorities in sampling request model preferences.
///
/// The type provides safe initialization that returns `nil` for values outside
/// the valid range, ensuring that all instances contain valid unit interval values.
///
/// - Example:
/// ```swift
/// let zero: UnitInterval = 0 // 0.0
/// let half = UnitInterval(0.5)! // 0.5
/// let one: UnitInterval = 1.0 // 1.0
/// let invalid = UnitInterval(1.5) // nil
/// ```
public struct UnitInterval: Hashable, Sendable {
private let value: Double
/// Creates a unit interval value from a `Double`.
///
/// - Parameter value: A double value that must be in the range 0.0...1.0
/// - Returns: A `UnitInterval` instance if the value is valid, `nil` otherwise
///
/// - Example:
/// ```swift
/// let valid = UnitInterval(0.75) // Optional(0.75)
/// let invalid = UnitInterval(-0.1) // nil
/// let boundary = UnitInterval(1.0) // Optional(1.0)
/// ```
public init?(_ value: Double) {
guard (0...1).contains(value) else { return nil }
self.value = value
}
/// The underlying double value.
///
/// This property provides access to the raw double value that is guaranteed
/// to be within the range [0, 1].
///
/// - Returns: The double value between 0.0 and 1.0, inclusive
public var doubleValue: Double { value }
}
// MARK: - Comparable
extension UnitInterval: Comparable {
public static func < (lhs: UnitInterval, rhs: UnitInterval) -> Bool {
lhs.value < rhs.value
}
}
// MARK: - CustomStringConvertible
extension UnitInterval: CustomStringConvertible {
public var description: String { "\(value)" }
}
// MARK: - Codable
extension UnitInterval: Codable {
public init(from decoder: Decoder) throws {
let container = try decoder.singleValueContainer()
let doubleValue = try container.decode(Double.self)
guard let interval = UnitInterval(doubleValue) else {
throw DecodingError.dataCorrupted(
DecodingError.Context(
codingPath: decoder.codingPath,
debugDescription: "Value \(doubleValue) is not in range 0...1")
)
}
self = interval
}
public func encode(to encoder: Encoder) throws {
var container = encoder.singleValueContainer()
try container.encode(value)
}
}
// MARK: - ExpressibleByFloatLiteral
extension UnitInterval: ExpressibleByFloatLiteral {
/// Creates a unit interval from a floating-point literal.
///
/// This initializer allows you to create `UnitInterval` instances using
/// floating-point literals. The literal value must be in the range [0, 1]
/// or a runtime error will occur.
///
/// - Parameter value: A floating-point literal between 0.0 and 1.0
///
/// - Warning: This initializer will crash if the literal is outside the valid range.
/// Use the failable initializer `init(_:)` for runtime validation.
///
/// - Example:
/// ```swift
/// let quarter: UnitInterval = 0.25
/// let half: UnitInterval = 0.5
/// ```
public init(floatLiteral value: Double) {
self.init(value)!
}
}
// MARK: - ExpressibleByIntegerLiteral
extension UnitInterval: ExpressibleByIntegerLiteral {
/// Creates a unit interval from an integer literal.
///
/// This initializer allows you to create `UnitInterval` instances using
/// integer literals. Only the values 0 and 1 are valid.
///
/// - Parameter value: An integer literal, either 0 or 1
///
/// - Warning: This initializer will crash if the literal is outside the valid range.
/// Use the failable initializer `init(_:)` for runtime validation.
///
/// - Example:
/// ```swift
/// let zero: UnitInterval = 0
/// let one: UnitInterval = 1
/// ```
public init(integerLiteral value: Int) {
self.init(Double(value))!
}
}

View file

@ -583,6 +583,45 @@ public actor Client {
return (content: result.content, isError: result.isError)
}
// MARK: - Sampling
/// Register a handler for sampling requests from servers
///
/// Sampling allows servers to request LLM completions through the client,
/// enabling sophisticated agentic behaviors while maintaining human-in-the-loop control.
///
/// The sampling flow follows these steps:
/// 1. Server sends a `sampling/createMessage` request to the client
/// 2. Client reviews the request and can modify it (via this handler)
/// 3. Client samples from an LLM (via this handler)
/// 4. Client reviews the completion (via this handler)
/// 5. Client returns the result to the server
///
/// - Parameter handler: A closure that processes sampling requests and returns completions
/// - Returns: Self for method chaining
/// - SeeAlso: https://modelcontextprotocol.io/docs/concepts/sampling#how-sampling-works
@discardableResult
public func withSamplingHandler(
_ handler: @escaping @Sendable (CreateSamplingMessage.Parameters) async throws ->
CreateSamplingMessage.Result
) -> Self {
// Note: This would require extending the client architecture to handle incoming requests from servers.
// The current MCP Swift SDK architecture assumes clients only send requests to servers,
// but sampling requires bidirectional communication where servers can send requests to clients.
//
// A full implementation would need:
// 1. Request handlers in the client (similar to how servers handle requests)
// 2. Bidirectional transport support
// 3. Request/response correlation for server-to-client requests
//
// For now, this serves as the correct API design for when bidirectional support is added.
// This would register the handler similar to how servers register method handlers:
// methodHandlers[CreateSamplingMessage.name] = TypedRequestHandler(handler)
return self
}
// MARK: -
private func handleResponse(_ response: Response<AnyMethod>) async {

View file

@ -0,0 +1,196 @@
import Foundation
/// The Model Context Protocol (MCP) allows servers to request LLM completions
/// through the client, enabling sophisticated agentic behaviors while maintaining
/// security and privacy.
///
/// - SeeAlso: https://modelcontextprotocol.io/docs/concepts/sampling#how-sampling-works
public enum Sampling {
/// A message in the conversation history.
public struct Message: Hashable, Codable, Sendable {
/// The message role
public enum Role: String, Hashable, Codable, Sendable {
/// A user message
case user
/// An assistant message
case assistant
}
/// The message role
public let role: Role
/// The message content
public let content: Content
public init(role: Role, content: Content) {
self.role = role
self.content = content
}
/// Content types for sampling messages
public enum Content: Hashable, Codable, Sendable {
/// Text content
case text(String)
/// Image content
case image(data: String, mimeType: String)
private enum CodingKeys: String, CodingKey {
case type, text, data, mimeType
}
public init(from decoder: Decoder) throws {
let container = try decoder.container(keyedBy: CodingKeys.self)
let type = try container.decode(String.self, forKey: .type)
switch type {
case "text":
let text = try container.decode(String.self, forKey: .text)
self = .text(text)
case "image":
let data = try container.decode(String.self, forKey: .data)
let mimeType = try container.decode(String.self, forKey: .mimeType)
self = .image(data: data, mimeType: mimeType)
default:
throw DecodingError.dataCorruptedError(
forKey: .type, in: container,
debugDescription: "Unknown sampling message content type")
}
}
public func encode(to encoder: Encoder) throws {
var container = encoder.container(keyedBy: CodingKeys.self)
switch self {
case .text(let text):
try container.encode("text", forKey: .type)
try container.encode(text, forKey: .text)
case .image(let data, let mimeType):
try container.encode("image", forKey: .type)
try container.encode(data, forKey: .data)
try container.encode(mimeType, forKey: .mimeType)
}
}
}
}
/// Model preferences for sampling requests
public struct ModelPreferences: Hashable, Codable, Sendable {
/// Model hints for selection
public struct Hint: Hashable, Codable, Sendable {
/// Suggested model name/family
public let name: String?
public init(name: String? = nil) {
self.name = name
}
}
/// Array of model name suggestions that clients can use to select an appropriate model
public let hints: [Hint]?
/// Importance of minimizing costs (0-1 normalized)
public let costPriority: UnitInterval?
/// Importance of low latency response (0-1 normalized)
public let speedPriority: UnitInterval?
/// Importance of advanced model capabilities (0-1 normalized)
public let intelligencePriority: UnitInterval?
public init(
hints: [Hint]? = nil,
costPriority: UnitInterval? = nil,
speedPriority: UnitInterval? = nil,
intelligencePriority: UnitInterval? = nil
) {
self.hints = hints
self.costPriority = costPriority
self.speedPriority = speedPriority
self.intelligencePriority = intelligencePriority
}
}
/// Context inclusion options for sampling requests
public enum ContextInclusion: String, Hashable, Codable, Sendable {
/// No additional context
case none
/// Include context from the requesting server
case thisServer
/// Include context from all connected MCP servers
case allServers
}
/// Stop reason for sampling completion
public enum StopReason: String, Hashable, Codable, Sendable {
/// Natural end of turn
case endTurn
/// Hit a stop sequence
case stopSequence
/// Reached maximum tokens
case maxTokens
}
}
/// To request sampling from a client, servers send a `sampling/createMessage` request.
/// - SeeAlso: https://modelcontextprotocol.io/docs/concepts/sampling#how-sampling-works
public enum CreateSamplingMessage: Method {
public static let name = "sampling/createMessage"
public struct Parameters: Hashable, Codable, Sendable {
/// The conversation history to send to the LLM
public let messages: [Sampling.Message]
/// Model selection preferences
public let modelPreferences: Sampling.ModelPreferences?
/// Optional system prompt
public let systemPrompt: String?
/// What MCP context to include
public let includeContext: Sampling.ContextInclusion?
/// Controls randomness (0.0 to 1.0)
public let temperature: Double?
/// Maximum tokens to generate
public let maxTokens: Int
/// Array of sequences that stop generation
public let stopSequences: [String]?
/// Additional provider-specific parameters
public let metadata: [String: Value]?
public init(
messages: [Sampling.Message],
modelPreferences: Sampling.ModelPreferences? = nil,
systemPrompt: String? = nil,
includeContext: Sampling.ContextInclusion? = nil,
temperature: Double? = nil,
maxTokens: Int,
stopSequences: [String]? = nil,
metadata: [String: Value]? = nil
) {
self.messages = messages
self.modelPreferences = modelPreferences
self.systemPrompt = systemPrompt
self.includeContext = includeContext
self.temperature = temperature
self.maxTokens = maxTokens
self.stopSequences = stopSequences
self.metadata = metadata
}
}
public struct Result: Hashable, Codable, Sendable {
/// Name of the model used
public let model: String
/// Why sampling stopped
public let stopReason: Sampling.StopReason?
/// The role of the completion
public let role: Sampling.Message.Role
/// The completion content
public let content: Sampling.Message.Content
public init(
model: String,
stopReason: Sampling.StopReason? = nil,
role: Sampling.Message.Role,
content: Sampling.Message.Content
) {
self.model = model
self.stopReason = stopReason
self.role = role
self.content = content
}
}
}

View file

@ -82,12 +82,19 @@ public actor Server {
public init() {}
}
/// Sampling capabilities
public struct Sampling: Hashable, Codable, Sendable {
public init() {}
}
/// Logging capabilities
public var logging: Logging?
/// Prompts capabilities
public var prompts: Prompts?
/// Resources capabilities
public var resources: Resources?
/// Sampling capabilities
public var sampling: Sampling?
/// Tools capabilities
public var tools: Tools?
@ -95,11 +102,13 @@ public actor Server {
logging: Logging? = nil,
prompts: Prompts? = nil,
resources: Resources? = nil,
sampling: Sampling? = nil,
tools: Tools? = nil
) {
self.logging = logging
self.prompts = prompts
self.resources = resources
self.sampling = sampling
self.tools = tools
}
}
@ -290,6 +299,69 @@ public actor Server {
try await connection.send(notificationData)
}
// MARK: - Sampling
/// Request sampling from the connected client
///
/// Sampling allows servers to request LLM completions through the client,
/// enabling sophisticated agentic behaviors while maintaining human-in-the-loop control.
///
/// The sampling flow follows these steps:
/// 1. Server sends a `sampling/createMessage` request to the client
/// 2. Client reviews the request and can modify it
/// 3. Client samples from an LLM
/// 4. Client reviews the completion
/// 5. Client returns the result to the server
///
/// - Parameters:
/// - messages: The conversation history to send to the LLM
/// - modelPreferences: Model selection preferences
/// - systemPrompt: Optional system prompt
/// - includeContext: What MCP context to include
/// - temperature: Controls randomness (0.0 to 1.0)
/// - maxTokens: Maximum tokens to generate
/// - stopSequences: Array of sequences that stop generation
/// - metadata: Additional provider-specific parameters
/// - Returns: The sampling result containing the model used, stop reason, role, and content
/// - Throws: MCPError if the request fails
/// - SeeAlso: https://modelcontextprotocol.io/docs/concepts/sampling#how-sampling-works
public func requestSampling(
messages: [Sampling.Message],
modelPreferences: Sampling.ModelPreferences? = nil,
systemPrompt: String? = nil,
includeContext: Sampling.ContextInclusion? = nil,
temperature: Double? = nil,
maxTokens: Int,
stopSequences: [String]? = nil,
metadata: [String: Value]? = nil
) async throws -> CreateSamplingMessage.Result {
guard connection != nil else {
throw MCPError.internalError("Server connection not initialized")
}
// Note: This is a conceptual implementation. The actual implementation would require
// bidirectional communication support in the transport layer, allowing servers to
// send requests to clients and receive responses.
_ = CreateSamplingMessage.request(
.init(
messages: messages,
modelPreferences: modelPreferences,
systemPrompt: systemPrompt,
includeContext: includeContext,
temperature: temperature,
maxTokens: maxTokens,
stopSequences: stopSequences,
metadata: metadata
)
)
// This would need to be implemented with proper request/response handling
// similar to how the client sends requests to servers
throw MCPError.internalError(
"Bidirectional sampling requests not yet implemented in transport layer")
}
/// A JSON-RPC batch containing multiple requests and/or notifications
struct Batch: Sendable {
/// An item in a JSON-RPC batch

View file

@ -0,0 +1,806 @@
import Logging
import Testing
import class Foundation.JSONDecoder
import class Foundation.JSONEncoder
@testable import MCP
#if canImport(System)
import System
#else
@preconcurrency import SystemPackage
#endif
@Suite("Sampling Tests")
struct SamplingTests {
@Test("Sampling.Message encoding and decoding")
func testSamplingMessageCoding() throws {
let encoder = JSONEncoder()
let decoder = JSONDecoder()
// Test text content
let textMessage = Sampling.Message(
role: .user,
content: .text("Hello, world!")
)
let textData = try encoder.encode(textMessage)
let decodedTextMessage = try decoder.decode(Sampling.Message.self, from: textData)
#expect(decodedTextMessage.role == .user)
if case .text(let text) = decodedTextMessage.content {
#expect(text == "Hello, world!")
} else {
#expect(Bool(false), "Expected text content")
}
// Test image content
let imageMessage = Sampling.Message(
role: .assistant,
content: .image(data: "base64imagedata", mimeType: "image/png")
)
let imageData = try encoder.encode(imageMessage)
let decodedImageMessage = try decoder.decode(Sampling.Message.self, from: imageData)
#expect(decodedImageMessage.role == .assistant)
if case .image(let data, let mimeType) = decodedImageMessage.content {
#expect(data == "base64imagedata")
#expect(mimeType == "image/png")
} else {
#expect(Bool(false), "Expected image content")
}
}
@Test("ModelPreferences encoding and decoding")
func testModelPreferencesCoding() throws {
let encoder = JSONEncoder()
let decoder = JSONDecoder()
let preferences = Sampling.ModelPreferences(
hints: [
Sampling.ModelPreferences.Hint(name: "claude-4"),
Sampling.ModelPreferences.Hint(name: "gpt-4.1"),
],
costPriority: 0.8,
speedPriority: 0.3,
intelligencePriority: 0.9
)
let data = try encoder.encode(preferences)
let decoded = try decoder.decode(Sampling.ModelPreferences.self, from: data)
#expect(decoded.hints?.count == 2)
#expect(decoded.hints?[0].name == "claude-4")
#expect(decoded.hints?[1].name == "gpt-4.1")
#expect(decoded.costPriority?.doubleValue == 0.8)
#expect(decoded.speedPriority?.doubleValue == 0.3)
#expect(decoded.intelligencePriority?.doubleValue == 0.9)
}
@Test("ContextInclusion encoding and decoding")
func testContextInclusionCoding() throws {
let encoder = JSONEncoder()
let decoder = JSONDecoder()
let contexts: [Sampling.ContextInclusion] = [.none, .thisServer, .allServers]
for context in contexts {
let data = try encoder.encode(context)
let decoded = try decoder.decode(Sampling.ContextInclusion.self, from: data)
#expect(decoded == context)
}
}
@Test("StopReason encoding and decoding")
func testStopReasonCoding() throws {
let encoder = JSONEncoder()
let decoder = JSONDecoder()
let reasons: [Sampling.StopReason] = [.endTurn, .stopSequence, .maxTokens]
for reason in reasons {
let data = try encoder.encode(reason)
let decoded = try decoder.decode(Sampling.StopReason.self, from: data)
#expect(decoded == reason)
}
}
@Test("CreateMessage request parameters")
func testCreateMessageParameters() throws {
let encoder = JSONEncoder()
let decoder = JSONDecoder()
let messages = [
Sampling.Message(role: .user, content: .text("What is the weather like?")),
Sampling.Message(
role: .assistant, content: .text("I need to check the weather for you.")),
]
let modelPreferences = Sampling.ModelPreferences(
hints: [Sampling.ModelPreferences.Hint(name: "claude-4-sonnet")],
costPriority: 0.5,
speedPriority: 0.7,
intelligencePriority: 0.9
)
let parameters = CreateSamplingMessage.Parameters(
messages: messages,
modelPreferences: modelPreferences,
systemPrompt: "You are a helpful weather assistant.",
includeContext: .thisServer,
temperature: 0.7,
maxTokens: 150,
stopSequences: ["END", "STOP"],
metadata: ["provider": "test"]
)
let data = try encoder.encode(parameters)
let decoded = try decoder.decode(CreateSamplingMessage.Parameters.self, from: data)
#expect(decoded.messages.count == 2)
#expect(decoded.messages[0].role == .user)
#expect(decoded.systemPrompt == "You are a helpful weather assistant.")
#expect(decoded.includeContext == .thisServer)
#expect(decoded.temperature == 0.7)
#expect(decoded.maxTokens == 150)
#expect(decoded.stopSequences?.count == 2)
#expect(decoded.stopSequences?[0] == "END")
#expect(decoded.stopSequences?[1] == "STOP")
#expect(decoded.metadata?["provider"]?.stringValue == "test")
}
@Test("CreateMessage result")
func testCreateMessageResult() throws {
let encoder = JSONEncoder()
let decoder = JSONDecoder()
let result = CreateSamplingMessage.Result(
model: "claude-4-sonnet",
stopReason: .endTurn,
role: .assistant,
content: .text("The weather is sunny and 75°F.")
)
let data = try encoder.encode(result)
let decoded = try decoder.decode(CreateSamplingMessage.Result.self, from: data)
#expect(decoded.model == "claude-4-sonnet")
#expect(decoded.stopReason == .endTurn)
#expect(decoded.role == .assistant)
if case .text(let text) = decoded.content {
#expect(text == "The weather is sunny and 75°F.")
} else {
#expect(Bool(false), "Expected text content")
}
}
@Test("CreateMessage request creation")
func testCreateMessageRequest() throws {
let messages = [
Sampling.Message(role: .user, content: .text("Hello"))
]
let request = CreateSamplingMessage.request(
.init(
messages: messages,
maxTokens: 100
)
)
#expect(request.method == "sampling/createMessage")
#expect(request.params.messages.count == 1)
#expect(request.params.maxTokens == 100)
}
@Test("Server capabilities include sampling")
func testServerCapabilitiesIncludeSampling() throws {
let capabilities = Server.Capabilities(
sampling: .init()
)
#expect(capabilities.sampling != nil)
let encoder = JSONEncoder()
let decoder = JSONDecoder()
let data = try encoder.encode(capabilities)
let decoded = try decoder.decode(Server.Capabilities.self, from: data)
#expect(decoded.sampling != nil)
}
@Test("Client capabilities include sampling")
func testClientCapabilitiesIncludeSampling() throws {
let capabilities = Client.Capabilities(
sampling: .init()
)
#expect(capabilities.sampling != nil)
let encoder = JSONEncoder()
let decoder = JSONDecoder()
let data = try encoder.encode(capabilities)
let decoded = try decoder.decode(Client.Capabilities.self, from: data)
#expect(decoded.sampling != nil)
}
@Test("Client sampling handler registration")
func testClientSamplingHandlerRegistration() async throws {
let client = Client(name: "TestClient", version: "1.0")
// Test that sampling handler can be registered
let handlerClient = await client.withSamplingHandler { parameters in
// Mock handler that returns a simple response
return CreateSamplingMessage.Result(
model: "test-model",
stopReason: .endTurn,
role: .assistant,
content: .text("Test response")
)
}
// Should return self for method chaining
#expect(handlerClient === client)
}
@Test("Server sampling request method")
func testServerSamplingRequestMethod() async throws {
let transport = MockTransport()
let server = Server(
name: "TestServer",
version: "1.0",
capabilities: .init(sampling: .init())
)
try await server.start(transport: transport)
// Test that server can attempt to request sampling
let messages = [
Sampling.Message(role: .user, content: .text("Test message"))
]
do {
_ = try await server.requestSampling(
messages: messages,
maxTokens: 100
)
#expect(
Bool(false),
"Should have thrown an error for unimplemented bidirectional communication")
} catch let error as MCPError {
if case .internalError(let message) = error {
#expect(
message?.contains("Bidirectional sampling requests not yet implemented") == true
)
} else {
#expect(Bool(false), "Expected internalError, got \(error)")
}
} catch {
#expect(Bool(false), "Expected MCPError, got \(error)")
}
await server.stop()
}
@Test("Sampling message content JSON format")
func testSamplingMessageContentJSONFormat() throws {
let encoder = JSONEncoder()
encoder.outputFormatting = [.sortedKeys]
// Test text content JSON format
let textContent = Sampling.Message.Content.text("Hello")
let textData = try encoder.encode(textContent)
let textJSON = String(data: textData, encoding: .utf8)!
#expect(textJSON.contains("\"type\":\"text\""))
#expect(textJSON.contains("\"text\":\"Hello\""))
// Test image content JSON format
let imageContent = Sampling.Message.Content.image(data: "base64data", mimeType: "image/png")
let imageData = try encoder.encode(imageContent)
let imageJSON = String(data: imageData, encoding: .utf8)!
#expect(imageJSON.contains("\"type\":\"image\""))
#expect(imageJSON.contains("\"data\":\"base64data\""))
#expect(imageJSON.contains("\"mimeType\":\"image\\/png\""))
}
@Test("UnitInterval in Sampling.ModelPreferences")
func testUnitIntervalInModelPreferences() throws {
// Test that UnitInterval validation works in Sampling.ModelPreferences
let validPreferences = Sampling.ModelPreferences(
costPriority: 0.5,
speedPriority: 1.0,
intelligencePriority: 0.0
)
#expect(validPreferences.costPriority?.doubleValue == 0.5)
#expect(validPreferences.speedPriority?.doubleValue == 1.0)
#expect(validPreferences.intelligencePriority?.doubleValue == 0.0)
// Test JSON encoding/decoding preserves UnitInterval constraints
let encoder = JSONEncoder()
let decoder = JSONDecoder()
let data = try encoder.encode(validPreferences)
let decoded = try decoder.decode(Sampling.ModelPreferences.self, from: data)
#expect(decoded.costPriority?.doubleValue == 0.5)
#expect(decoded.speedPriority?.doubleValue == 1.0)
#expect(decoded.intelligencePriority?.doubleValue == 0.0)
}
}
@Suite("Sampling Integration Tests")
struct SamplingIntegrationTests {
@Test(
.timeLimit(.minutes(1))
)
func testSamplingCapabilitiesNegotiation() async throws {
let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe()
let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe()
var logger = Logger(
label: "mcp.test.sampling",
factory: { StreamLogHandler.standardError(label: $0) })
logger.logLevel = .debug
let serverTransport = StdioTransport(
input: clientToServerRead,
output: serverToClientWrite,
logger: logger
)
let clientTransport = StdioTransport(
input: serverToClientRead,
output: clientToServerWrite,
logger: logger
)
// Server with sampling capability
let server = Server(
name: "SamplingTestServer",
version: "1.0.0",
capabilities: .init(
sampling: .init(), // Enable sampling
tools: .init()
)
)
// Client (capabilities will be set during initialization)
let client = Client(
name: "SamplingTestClient",
version: "1.0"
)
try await server.start(transport: serverTransport)
try await client.connect(transport: clientTransport)
// Test initialization and capability negotiation
let initTask = Task {
let result = try await client.initialize()
#expect(result.serverInfo.name == "SamplingTestServer")
#expect(result.serverInfo.version == "1.0.0")
#expect(
result.capabilities.sampling != nil, "Server should advertise sampling capability")
#expect(result.protocolVersion == Version.latest)
}
try await withThrowingTaskGroup(of: Void.self) { group in
group.addTask {
try await Task.sleep(for: .seconds(1))
initTask.cancel()
throw CancellationError()
}
group.addTask {
try await initTask.value
}
try await group.next()
group.cancelAll()
}
await server.stop()
await client.disconnect()
try? clientToServerRead.close()
try? clientToServerWrite.close()
try? serverToClientRead.close()
try? serverToClientWrite.close()
}
@Test(
.timeLimit(.minutes(1))
)
func testSamplingHandlerRegistration() async throws {
let client = Client(
name: "SamplingHandlerTestClient",
version: "1.0"
)
// Register sampling handler
let handlerClient = await client.withSamplingHandler { parameters in
// Mock LLM response
return CreateSamplingMessage.Result(
model: "test-model-v1",
stopReason: .endTurn,
role: .assistant,
content: .text("This is a test completion from the mock LLM.")
)
}
// Verify method chaining works
#expect(
handlerClient === client, "withSamplingHandler should return self for method chaining")
// Note: We can't test the actual handler invocation without bidirectional transport,
// but we can verify the handler registration doesn't crash and returns correctly
}
@Test(
.timeLimit(.minutes(1))
)
func testServerSamplingRequestAPI() async throws {
let transport = MockTransport()
let server = Server(
name: "SamplingRequestTestServer",
version: "1.0",
capabilities: .init(sampling: .init())
)
try await server.start(transport: transport)
// Test sampling request with comprehensive parameters
let messages = [
Sampling.Message(
role: .user,
content: .text("Analyze the following data and provide insights:")
),
Sampling.Message(
role: .user,
content: .text("Sales data: Q1: $100k, Q2: $150k, Q3: $200k, Q4: $180k")
),
]
let modelPreferences = Sampling.ModelPreferences(
hints: [
Sampling.ModelPreferences.Hint(name: "claude-4-sonnet"),
Sampling.ModelPreferences.Hint(name: "gpt-4.1"),
],
costPriority: 0.3,
speedPriority: 0.7,
intelligencePriority: 0.9
)
// Test that the API accepts all parameters correctly
do {
_ = try await server.requestSampling(
messages: messages,
modelPreferences: modelPreferences,
systemPrompt: "You are a business analyst expert.",
includeContext: .thisServer,
temperature: 0.7,
maxTokens: 500,
stopSequences: ["END_ANALYSIS", "\n\n---"],
metadata: [
"requestId": "test-123",
"priority": "high",
"department": "analytics",
]
)
#expect(Bool(false), "Should throw error for unimplemented bidirectional communication")
} catch let error as MCPError {
if case .internalError(let message) = error {
#expect(
message?.contains("Bidirectional sampling requests not yet implemented")
== true,
"Should indicate bidirectional communication not implemented"
)
} else {
#expect(Bool(false), "Expected internalError, got \(error)")
}
} catch {
#expect(Bool(false), "Expected MCPError, got \(error)")
}
await server.stop()
}
@Test(
.timeLimit(.minutes(1))
)
func testSamplingMessageTypes() async throws {
// Test comprehensive message content types
let textMessage = Sampling.Message(
role: .user,
content: .text("What do you see in this data?")
)
let imageMessage = Sampling.Message(
role: .user,
content: .image(
data:
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==",
mimeType: "image/png"
)
)
// Test encoding/decoding of different message types
let encoder = JSONEncoder()
let decoder = JSONDecoder()
// Test text message
let textData = try encoder.encode(textMessage)
let decodedTextMessage = try decoder.decode(Sampling.Message.self, from: textData)
#expect(decodedTextMessage.role == .user)
if case .text(let text) = decodedTextMessage.content {
#expect(text == "What do you see in this data?")
} else {
#expect(Bool(false), "Expected text content")
}
// Test image message
let imageData = try encoder.encode(imageMessage)
let decodedImageMessage = try decoder.decode(Sampling.Message.self, from: imageData)
#expect(decodedImageMessage.role == .user)
if case .image(let data, let mimeType) = decodedImageMessage.content {
#expect(data.contains("iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ"))
#expect(mimeType == "image/png")
} else {
#expect(Bool(false), "Expected image content")
}
}
@Test(
.timeLimit(.minutes(1))
)
func testSamplingResultTypes() async throws {
// Test different result content types and stop reasons
let textResult = CreateSamplingMessage.Result(
model: "claude-4-sonnet",
stopReason: .endTurn,
role: .assistant,
content: .text(
"Based on the sales data analysis, I can see a strong upward trend through Q3, with a slight decline in Q4. This suggests seasonal factors or market saturation."
)
)
let imageResult = CreateSamplingMessage.Result(
model: "dall-e-3",
stopReason: .maxTokens,
role: .assistant,
content: .image(
data: "generated_chart_base64_data_here",
mimeType: "image/png"
)
)
let stopSequenceResult = CreateSamplingMessage.Result(
model: "gpt-4.1",
stopReason: .stopSequence,
role: .assistant,
content: .text("Analysis complete.\nEND_ANALYSIS")
)
// Test encoding/decoding of different result types
let encoder = JSONEncoder()
let decoder = JSONDecoder()
// Test text result
let textData = try encoder.encode(textResult)
let decodedTextResult = try decoder.decode(
CreateSamplingMessage.Result.self, from: textData)
#expect(decodedTextResult.model == "claude-4-sonnet")
#expect(decodedTextResult.stopReason == .endTurn)
#expect(decodedTextResult.role == .assistant)
// Test image result
let imageData = try encoder.encode(imageResult)
let decodedImageResult = try decoder.decode(
CreateSamplingMessage.Result.self, from: imageData)
#expect(decodedImageResult.model == "dall-e-3")
#expect(decodedImageResult.stopReason == .maxTokens)
// Test stop sequence result
let stopData = try encoder.encode(stopSequenceResult)
let decodedStopResult = try decoder.decode(
CreateSamplingMessage.Result.self, from: stopData)
#expect(decodedStopResult.stopReason == .stopSequence)
}
@Test(
.timeLimit(.minutes(1))
)
func testSamplingErrorHandling() async throws {
let transport = MockTransport()
let server = Server(
name: "ErrorTestServer",
version: "1.0",
capabilities: .init() // No sampling capability
)
try await server.start(transport: transport)
// Test sampling request on server without sampling capability
let messages = [
Sampling.Message(role: .user, content: .text("Test message"))
]
do {
_ = try await server.requestSampling(
messages: messages,
maxTokens: 100
)
#expect(Bool(false), "Should throw error for missing connection")
} catch let error as MCPError {
if case .internalError(let message) = error {
#expect(
message?.contains("Server connection not initialized") == true
|| message?.contains("Bidirectional sampling requests not yet implemented")
== true,
"Should indicate connection or implementation issue"
)
} else {
#expect(Bool(false), "Expected internalError, got \(error)")
}
} catch {
#expect(Bool(false), "Expected MCPError, got \(error)")
}
await server.stop()
}
@Test(
.timeLimit(.minutes(1))
)
func testSamplingParameterValidation() async throws {
// Test parameter validation and edge cases
let validMessages = [
Sampling.Message(role: .user, content: .text("Valid message"))
]
_ = [Sampling.Message]() // Test empty messages array
// Test with valid parameters
let validParams = CreateSamplingMessage.Parameters(
messages: validMessages,
maxTokens: 100
)
#expect(validParams.messages.count == 1)
#expect(validParams.maxTokens == 100)
// Test with comprehensive parameters
let comprehensiveParams = CreateSamplingMessage.Parameters(
messages: validMessages,
modelPreferences: Sampling.ModelPreferences(
hints: [Sampling.ModelPreferences.Hint(name: "claude-4")],
costPriority: 0.5,
speedPriority: 0.8,
intelligencePriority: 0.9
),
systemPrompt: "You are a helpful assistant.",
includeContext: .allServers,
temperature: 0.7,
maxTokens: 500,
stopSequences: ["STOP", "END"],
metadata: [
"sessionId": "test-session-123",
"userId": "user-456",
]
)
#expect(comprehensiveParams.messages.count == 1)
#expect(comprehensiveParams.modelPreferences?.hints?.count == 1)
#expect(comprehensiveParams.systemPrompt == "You are a helpful assistant.")
#expect(comprehensiveParams.includeContext == .allServers)
#expect(comprehensiveParams.temperature == 0.7)
#expect(comprehensiveParams.maxTokens == 500)
#expect(comprehensiveParams.stopSequences?.count == 2)
#expect(comprehensiveParams.metadata?.count == 2)
// Test encoding/decoding of comprehensive parameters
let encoder = JSONEncoder()
let decoder = JSONDecoder()
let data = try encoder.encode(comprehensiveParams)
let decoded = try decoder.decode(CreateSamplingMessage.Parameters.self, from: data)
#expect(decoded.messages.count == 1)
#expect(decoded.modelPreferences?.costPriority?.doubleValue == 0.5)
#expect(decoded.systemPrompt == "You are a helpful assistant.")
#expect(decoded.includeContext == .allServers)
#expect(decoded.temperature == 0.7)
#expect(decoded.maxTokens == 500)
#expect(decoded.stopSequences?[0] == "STOP")
#expect(decoded.metadata?["sessionId"]?.stringValue == "test-session-123")
}
@Test(
.timeLimit(.minutes(1))
)
func testSamplingWorkflowScenarios() async throws {
// Test realistic sampling workflow scenarios
// Scenario 1: Data Analysis Request
let dataAnalysisMessages = [
Sampling.Message(
role: .user,
content: .text("Please analyze the following customer feedback data:")
),
Sampling.Message(
role: .user,
content: .text(
"""
Feedback Summary:
- 85% positive sentiment
- Top complaints: shipping delays (12%), product quality (8%)
- Top praise: customer service (45%), product features (40%)
- NPS Score: 72
""")
),
]
let dataAnalysisParams = CreateSamplingMessage.Parameters(
messages: dataAnalysisMessages,
modelPreferences: Sampling.ModelPreferences(
hints: [Sampling.ModelPreferences.Hint(name: "claude-4-sonnet")],
speedPriority: 0.3,
intelligencePriority: 0.9
),
systemPrompt: "You are an expert business analyst. Provide actionable insights.",
includeContext: .thisServer,
temperature: 0.3, // Lower temperature for analytical tasks
maxTokens: 400,
stopSequences: ["---END---"],
metadata: ["analysisType": "customer-feedback"]
)
// Scenario 2: Creative Content Generation
let creativeMessages = [
Sampling.Message(
role: .user,
content: .text(
"Write a compelling product description for a new smart home device.")
)
]
let creativeParams = CreateSamplingMessage.Parameters(
messages: creativeMessages,
modelPreferences: Sampling.ModelPreferences(
hints: [Sampling.ModelPreferences.Hint(name: "gpt-4.1")],
costPriority: 0.4,
speedPriority: 0.6,
intelligencePriority: 0.8
),
systemPrompt: "You are a creative marketing copywriter.",
temperature: 0.8, // Higher temperature for creativity
maxTokens: 200,
metadata: ["contentType": "marketing-copy"]
)
// Test parameter encoding for both scenarios
let encoder = JSONEncoder()
let analysisData = try encoder.encode(dataAnalysisParams)
let creativeData = try encoder.encode(creativeParams)
// Verify both encode successfully
#expect(analysisData.count > 0)
#expect(creativeData.count > 0)
// Test decoding
let decoder = JSONDecoder()
let decodedAnalysis = try decoder.decode(
CreateSamplingMessage.Parameters.self, from: analysisData)
let decodedCreative = try decoder.decode(
CreateSamplingMessage.Parameters.self, from: creativeData)
#expect(decodedAnalysis.temperature == 0.3)
#expect(decodedCreative.temperature == 0.8)
#expect(decodedAnalysis.modelPreferences?.intelligencePriority?.doubleValue == 0.9)
#expect(decodedCreative.modelPreferences?.costPriority?.doubleValue == 0.4)
}
}

View file

@ -0,0 +1,219 @@
import Testing
import class Foundation.JSONDecoder
import class Foundation.JSONEncoder
@testable import MCP
@Suite("UnitInterval Tests")
struct UnitIntervalTests {
@Test("Valid literal initialization")
func testValidLiteralInitialization() throws {
let zero: UnitInterval = 0.0
#expect(zero.doubleValue == 0.0)
let half: UnitInterval = 0.5
#expect(half.doubleValue == 0.5)
let one: UnitInterval = 1.0
#expect(one.doubleValue == 1.0)
let quarter: UnitInterval = 0.25
#expect(quarter.doubleValue == 0.25)
}
@Test("Valid failable initialization with runtime values")
func testValidFailableInitialization() throws {
// Test with runtime computed values to force use of failable initializer
let values = [0.0, 0.5, 1.0, 0.25]
for value in values {
let computed = value * 1.0 // Force runtime computation
let interval = UnitInterval(computed)
#expect(interval != nil)
#expect(interval!.doubleValue == value)
}
}
@Test("Invalid failable initialization")
func testInvalidFailableInitialization() throws {
// Test with runtime computed values to force use of failable initializer
let invalidValues = [-0.1, 1.1, 100.0, -100.0]
for value in invalidValues {
let computed = value * 1.0 // Force runtime computation
let interval = UnitInterval(computed)
#expect(interval == nil)
}
}
@Test("Boundary and edge case values")
func testBoundaryAndEdgeCaseValues() throws {
// Test exact boundary values
let exactZero = 0.0 * 1.0
let zero = UnitInterval(exactZero)
#expect(zero != nil)
let exactOne = 1.0 * 1.0
let one = UnitInterval(exactOne)
#expect(one != nil)
// Test machine precision boundaries
let justAboveZero = Double.ulpOfOne * 1.0
let aboveZero = UnitInterval(justAboveZero)
#expect(aboveZero != nil)
let justBelowOne = (1.0 - Double.ulpOfOne) * 1.0
let belowOne = UnitInterval(justBelowOne)
#expect(belowOne != nil)
// Test very small positive value
let tinyValue = 1e-10 * 1.0
let tiny = UnitInterval(tinyValue)
#expect(tiny != nil)
#expect(tiny!.doubleValue == 1e-10)
// Test value very close to 1
let almostOneValue = 0.9999999999 * 1.0
let almostOne = UnitInterval(almostOneValue)
#expect(almostOne != nil)
#expect(almostOne!.doubleValue == 0.9999999999)
}
@Test("Float literal initialization")
func testFloatLiteralInitialization() throws {
let zero: UnitInterval = 0.0
#expect(zero.doubleValue == 0.0)
let half: UnitInterval = 0.5
#expect(half.doubleValue == 0.5)
let one: UnitInterval = 1.0
#expect(one.doubleValue == 1.0)
let quarter: UnitInterval = 0.25
#expect(quarter.doubleValue == 0.25)
}
@Test("Integer literal initialization")
func testIntegerLiteralInitialization() throws {
let zero: UnitInterval = 0
#expect(zero.doubleValue == 0.0)
let one: UnitInterval = 1
#expect(one.doubleValue == 1.0)
}
@Test("Comparable conformance")
func testComparable() throws {
let zero: UnitInterval = 0.0
let quarter: UnitInterval = 0.25
let half: UnitInterval = 0.5
let one: UnitInterval = 1.0
#expect(zero < quarter)
#expect(quarter < half)
#expect(half < one)
#expect(zero < one)
#expect(!(quarter < zero))
#expect(!(half < quarter))
#expect(!(one < half))
#expect(zero <= quarter)
#expect(quarter <= half)
#expect(half <= one)
#expect(zero <= zero)
#expect(quarter > zero)
#expect(half > quarter)
#expect(one > half)
#expect(quarter >= zero)
#expect(half >= quarter)
#expect(one >= half)
#expect(one >= one)
}
@Test("Equality and hashing")
func testEqualityAndHashing() throws {
let half1: UnitInterval = 0.5
let half2: UnitInterval = 0.5
let quarter: UnitInterval = 0.25
#expect(half1 == half2)
#expect(half1 != quarter)
#expect(half1.hashValue == half2.hashValue)
}
@Test("String description")
func testStringDescription() throws {
let zero: UnitInterval = 0.0
#expect(zero.description == "0.0")
let half: UnitInterval = 0.5
#expect(half.description == "0.5")
let one: UnitInterval = 1.0
#expect(one.description == "1.0")
let quarter: UnitInterval = 0.25
#expect(quarter.description == "0.25")
}
@Test("JSON encoding and decoding")
func testJSONCodable() throws {
let encoder = JSONEncoder()
let decoder = JSONDecoder()
let original: UnitInterval = 0.75
let data = try encoder.encode(original)
let decoded = try decoder.decode(UnitInterval.self, from: data)
#expect(decoded == original)
#expect(decoded.doubleValue == 0.75)
}
@Test("JSON decoding with invalid values")
func testJSONDecodingInvalidValues() throws {
let decoder = JSONDecoder()
// Test negative value
let negativeJSON = "-0.5".data(using: .utf8)!
#expect(throws: DecodingError.self) {
try decoder.decode(UnitInterval.self, from: negativeJSON)
}
// Test value greater than 1
let tooLargeJSON = "1.5".data(using: .utf8)!
#expect(throws: DecodingError.self) {
try decoder.decode(UnitInterval.self, from: tooLargeJSON)
}
}
@Test("JSON encoding produces expected format")
func testJSONEncodingFormat() throws {
let encoder = JSONEncoder()
let half: UnitInterval = 0.5
let data = try encoder.encode(half)
let jsonString = String(data: data, encoding: .utf8)
#expect(jsonString == "0.5")
}
@Test("Double value property")
func testDoubleValueProperty() throws {
let values = [0.0, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0]
for value in values {
let computed = value * 1.0 // Force runtime computation
if let interval = UnitInterval(computed) {
#expect(interval.doubleValue == value)
} else {
#expect(Bool(false), "UnitInterval(\(value)) should succeed")
}
}
}
}