From 521ed2a9678671b60a1316a9ab0446cc9b3459b8 Mon Sep 17 00:00:00 2001 From: Anthony DePasquale Date: Wed, 31 Dec 2025 13:23:25 +0100 Subject: [PATCH 1/8] Fix concurrency errors in client tests --- Tests/MCPTests/ClientTests.swift | 34 ++++++++++++++++++++++++-------- 1 file changed, 26 insertions(+), 8 deletions(-) diff --git a/Tests/MCPTests/ClientTests.swift b/Tests/MCPTests/ClientTests.swift index 6fcc87a4..80096cc1 100644 --- a/Tests/MCPTests/ClientTests.swift +++ b/Tests/MCPTests/ClientTests.swift @@ -345,12 +345,22 @@ struct ClientTests { let request1 = Ping.request() let request2 = Ping.request() - var resultTask1: Task? - var resultTask2: Task? + + // Use an actor to safely capture the tasks from the closure + actor TaskHolder { + var task1: Task? + var task2: Task? + func set(task1: Task, task2: Task) { + self.task1 = task1 + self.task2 = task2 + } + } + let holder = TaskHolder() try await client.withBatch { batch in - resultTask1 = try await batch.addRequest(request1) - resultTask2 = try await batch.addRequest(request2) + let t1 = try await batch.addRequest(request1) + let t2 = try await batch.addRequest(request2) + await holder.set(task1: t1, task2: t2) } // Check if batch message was sent (after initialize and initialized notification) @@ -381,7 +391,7 @@ struct ClientTests { try await transport.queue(batch: [anyResponse1, anyResponse2]) // Wait for results and verify - guard let task1 = resultTask1, let task2 = resultTask2 else { + guard let task1 = await holder.task1, let task2 = await holder.task2 else { #expect(Bool(false), "Result tasks not created") return } @@ -426,11 +436,18 @@ struct ClientTests { let request1 = Ping.request() // Success let request2 = Ping.request() // Error - var resultTasks: [Task] = [] + // Use an actor to safely capture the tasks from the closure + actor TasksHolder { + var tasks: [Task] = [] + func append(_ task: Task) { + tasks.append(task) + } + } + let holder = TasksHolder() try await client.withBatch { batch in - resultTasks.append(try await batch.addRequest(request1)) - resultTasks.append(try await batch.addRequest(request2)) + await holder.append(try await batch.addRequest(request1)) + await holder.append(try await batch.addRequest(request2)) } // Check if batch message was sent (after initialize and initialized notification) @@ -447,6 +464,7 @@ struct ClientTests { try await transport.queue(batch: [anyResponse1, anyResponse2]) // Wait for results and verify + let resultTasks = await holder.tasks #expect(resultTasks.count == 2) guard resultTasks.count == 2 else { #expect(Bool(false), "Expected 2 result tasks") From d30394e73f0357ea4c2ad94e31098a23eea56e6c Mon Sep 17 00:00:00 2001 From: Paul Berman Date: Fri, 23 May 2025 11:50:36 -0500 Subject: [PATCH 2/8] Append handlers in-situ. --- Sources/MCP/Client/Client.swift | 3 +-- Sources/MCP/Server/Server.swift | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/Sources/MCP/Client/Client.swift b/Sources/MCP/Client/Client.swift index 696ffd14..a2deb3dc 100644 --- a/Sources/MCP/Client/Client.swift +++ b/Sources/MCP/Client/Client.swift @@ -272,8 +272,7 @@ public actor Client { _ type: N.Type, handler: @escaping @Sendable (Message) async throws -> Void ) async -> Self { - let handlers = notificationHandlers[N.name, default: []] - notificationHandlers[N.name] = handlers + [TypedNotificationHandler(handler)] + notificationHandlers[N.name, default: []].append(TypedNotificationHandler(handler)) return self } diff --git a/Sources/MCP/Server/Server.swift b/Sources/MCP/Server/Server.swift index 6ba1e27b..f2c975af 100644 --- a/Sources/MCP/Server/Server.swift +++ b/Sources/MCP/Server/Server.swift @@ -276,8 +276,7 @@ public actor Server { _ type: N.Type, handler: @escaping @Sendable (Message) async throws -> Void ) -> Self { - let handlers = notificationHandlers[N.name, default: []] - notificationHandlers[N.name] = handlers + [TypedNotificationHandler(handler)] + notificationHandlers[N.name, default: []].append(TypedNotificationHandler(handler)) return self } From a1a9db0ee719a20f46db4fa3e9f5fdd7a3cbff33 Mon Sep 17 00:00:00 2001 From: Anthony DePasquale Date: Sat, 3 Jan 2026 21:58:59 +0100 Subject: [PATCH 3/8] Add missing functionality and many fixes --- .gitignore | 2 +- .../HummingbirdIntegration/Package.resolved | 231 ++ Examples/HummingbirdIntegration/Package.swift | 24 + .../HummingbirdIntegration/Sources/main.swift | 333 ++ Examples/README.md | 190 ++ Examples/VaporIntegration/Package.resolved | 267 ++ Examples/VaporIntegration/Package.swift | 24 + Examples/VaporIntegration/Sources/main.swift | 279 ++ Package.resolved | 4 +- README.md | 147 +- Sources/MCP/Base/Annotations.swift | 40 + Sources/MCP/Base/Error.swift | 516 +++- Sources/MCP/Base/HTTPHeader.swift | 45 + Sources/MCP/Base/Icon.swift | 51 + Sources/MCP/Base/Lifecycle.swift | 85 +- Sources/MCP/Base/Messages.swift | 113 +- Sources/MCP/Base/Progress.swift | 298 ++ .../MCP/Base/{ID.swift => RequestId.swift} | 16 +- Sources/MCP/Base/Transport.swift | 38 + .../Base/Transports/HTTPClientTransport.swift | 639 +++- .../HTTPServerTransport+Types.swift | 304 ++ .../Base/Transports/HTTPServerTransport.swift | 1163 +++++++ .../Base/Transports/InMemoryEventStore.swift | 277 ++ .../Base/Transports/NetworkTransport.swift | 277 +- .../MCP/Base/Transports/StdioTransport.swift | 19 +- .../Base/Utilities/ExtraFieldsCoding.swift | 112 + Sources/MCP/Base/Utilities/Ping.swift | 13 + Sources/MCP/Base/Versioning.swift | 38 +- Sources/MCP/Client/Client+Tasks.swift | 256 ++ Sources/MCP/Client/Client.swift | 2008 +++++++++++- Sources/MCP/Client/Elicitation.swift | 789 +++++ .../ExperimentalClientFeatures.swift | 330 ++ .../Tasks/ClientTaskSupport.swift | 231 ++ Sources/MCP/Client/Roots.swift | 135 + Sources/MCP/Client/Sampling.swift | 805 ++++- Sources/MCP/Server/Completions.swift | 338 ++ .../ExperimentalServerFeatures.swift | 217 ++ .../Tasks/ServerTaskContext.swift | 936 ++++++ .../Experimental/Tasks/TaskContext.swift | 276 ++ .../Experimental/Tasks/TaskMessageQueue.swift | 416 +++ .../Tasks/TaskResultHandler.swift | 136 + .../Server/Experimental/Tasks/TaskStore.swift | 333 ++ .../Experimental/Tasks/TaskSupport.swift | 287 ++ .../MCP/Server/Experimental/Tasks/Tasks.swift | 1013 ++++++ Sources/MCP/Server/Logging.swift | 100 + Sources/MCP/Server/Prompts.swift | 252 +- Sources/MCP/Server/Resources.swift | 359 ++- Sources/MCP/Server/Server.swift | 1210 +++++++- Sources/MCP/Server/SessionManager.swift | 221 ++ Sources/MCP/Server/ToolNameValidation.swift | 111 + Sources/MCP/Server/Tools.swift | 289 +- Tests/MCPTests/AdditionalServerTests.swift | 677 ++++ Tests/MCPTests/CancellationTests.swift | 1848 +++++++++++ Tests/MCPTests/CapabilitiesTests.swift | 919 ++++++ Tests/MCPTests/ClientReconnectionTests.swift | 238 ++ Tests/MCPTests/ClientTests.swift | 154 + Tests/MCPTests/CompletionTests.swift | 1035 +++++++ Tests/MCPTests/ElicitationTests.swift | 2460 +++++++++++++++ Tests/MCPTests/ErrorTests.swift | 417 +++ Tests/MCPTests/FullRoundtripTests.swift | 342 ++ Tests/MCPTests/HTTPClientTransportTests.swift | 1831 +++++++++-- Tests/MCPTests/HTTPIntegrationTests.swift | 536 ++++ Tests/MCPTests/Helpers/MockTransport.swift | 45 + Tests/MCPTests/Helpers/TestPayloads.swift | 199 ++ Tests/MCPTests/IDTests.swift | 12 +- Tests/MCPTests/InMemoryEventStoreTests.swift | 464 +++ Tests/MCPTests/NotificationTests.swift | 239 ++ Tests/MCPTests/PrimingEventsTests.swift | 371 +++ Tests/MCPTests/ProgressTests.swift | 2737 +++++++++++++++++ Tests/MCPTests/PromptTests.swift | 762 ++++- Tests/MCPTests/RequestTests.swift | 2 +- .../MCPTests/ResourceSubscriptionTests.swift | 446 +++ Tests/MCPTests/ResourceTests.swift | 806 ++++- Tests/MCPTests/ResponseTests.swift | 16 +- Tests/MCPTests/ResumabilityTests.swift | 684 ++++ Tests/MCPTests/RootsTests.swift | 622 ++++ Tests/MCPTests/RoundtripTests.swift | 16 +- Tests/MCPTests/SamplingTests.swift | 1311 ++++++-- Tests/MCPTests/ServerTests.swift | 49 +- Tests/MCPTests/SessionManagerTests.swift | 239 ++ Tests/MCPTests/StdioTransportTests.swift | 313 ++ .../StreamableHTTPServerTransportTests.swift | 2489 +++++++++++++++ Tests/MCPTests/TaskTests.swift | 1187 +++++++ Tests/MCPTests/ToolTests.swift | 845 ++++- Tests/MCPTests/TransportSwitchingTests.swift | 441 +++ Tests/MCPTests/VersioningTests.swift | 14 +- 86 files changed, 39972 insertions(+), 1387 deletions(-) create mode 100644 Examples/HummingbirdIntegration/Package.resolved create mode 100644 Examples/HummingbirdIntegration/Package.swift create mode 100644 Examples/HummingbirdIntegration/Sources/main.swift create mode 100644 Examples/README.md create mode 100644 Examples/VaporIntegration/Package.resolved create mode 100644 Examples/VaporIntegration/Package.swift create mode 100644 Examples/VaporIntegration/Sources/main.swift create mode 100644 Sources/MCP/Base/Annotations.swift create mode 100644 Sources/MCP/Base/HTTPHeader.swift create mode 100644 Sources/MCP/Base/Icon.swift create mode 100644 Sources/MCP/Base/Progress.swift rename Sources/MCP/Base/{ID.swift => RequestId.swift} (79%) create mode 100644 Sources/MCP/Base/Transports/HTTPServerTransport+Types.swift create mode 100644 Sources/MCP/Base/Transports/HTTPServerTransport.swift create mode 100644 Sources/MCP/Base/Transports/InMemoryEventStore.swift create mode 100644 Sources/MCP/Base/Utilities/ExtraFieldsCoding.swift create mode 100644 Sources/MCP/Client/Client+Tasks.swift create mode 100644 Sources/MCP/Client/Elicitation.swift create mode 100644 Sources/MCP/Client/Experimental/ExperimentalClientFeatures.swift create mode 100644 Sources/MCP/Client/Experimental/Tasks/ClientTaskSupport.swift create mode 100644 Sources/MCP/Client/Roots.swift create mode 100644 Sources/MCP/Server/Completions.swift create mode 100644 Sources/MCP/Server/Experimental/ExperimentalServerFeatures.swift create mode 100644 Sources/MCP/Server/Experimental/Tasks/ServerTaskContext.swift create mode 100644 Sources/MCP/Server/Experimental/Tasks/TaskContext.swift create mode 100644 Sources/MCP/Server/Experimental/Tasks/TaskMessageQueue.swift create mode 100644 Sources/MCP/Server/Experimental/Tasks/TaskResultHandler.swift create mode 100644 Sources/MCP/Server/Experimental/Tasks/TaskStore.swift create mode 100644 Sources/MCP/Server/Experimental/Tasks/TaskSupport.swift create mode 100644 Sources/MCP/Server/Experimental/Tasks/Tasks.swift create mode 100644 Sources/MCP/Server/Logging.swift create mode 100644 Sources/MCP/Server/SessionManager.swift create mode 100644 Sources/MCP/Server/ToolNameValidation.swift create mode 100644 Tests/MCPTests/AdditionalServerTests.swift create mode 100644 Tests/MCPTests/CancellationTests.swift create mode 100644 Tests/MCPTests/CapabilitiesTests.swift create mode 100644 Tests/MCPTests/ClientReconnectionTests.swift create mode 100644 Tests/MCPTests/CompletionTests.swift create mode 100644 Tests/MCPTests/ElicitationTests.swift create mode 100644 Tests/MCPTests/ErrorTests.swift create mode 100644 Tests/MCPTests/FullRoundtripTests.swift create mode 100644 Tests/MCPTests/HTTPIntegrationTests.swift create mode 100644 Tests/MCPTests/Helpers/TestPayloads.swift create mode 100644 Tests/MCPTests/InMemoryEventStoreTests.swift create mode 100644 Tests/MCPTests/PrimingEventsTests.swift create mode 100644 Tests/MCPTests/ProgressTests.swift create mode 100644 Tests/MCPTests/ResourceSubscriptionTests.swift create mode 100644 Tests/MCPTests/ResumabilityTests.swift create mode 100644 Tests/MCPTests/RootsTests.swift create mode 100644 Tests/MCPTests/SessionManagerTests.swift create mode 100644 Tests/MCPTests/StreamableHTTPServerTransportTests.swift create mode 100644 Tests/MCPTests/TaskTests.swift create mode 100644 Tests/MCPTests/TransportSwitchingTests.swift diff --git a/.gitignore b/.gitignore index 2d0abd04..77b66726 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,5 @@ .DS_Store -/.build +.build /Packages xcuserdata/ DerivedData/ diff --git a/Examples/HummingbirdIntegration/Package.resolved b/Examples/HummingbirdIntegration/Package.resolved new file mode 100644 index 00000000..9d6543f4 --- /dev/null +++ b/Examples/HummingbirdIntegration/Package.resolved @@ -0,0 +1,231 @@ +{ + "originHash" : "6e195967fb65405cea73fd5bebe27163ba86af50fcdda39e0266ea0df2f70c7c", + "pins" : [ + { + "identity" : "async-http-client", + "kind" : "remoteSourceControl", + "location" : "https://github.com/swift-server/async-http-client.git", + "state" : { + "revision" : "5dd84c7bb48b348751d7bbe7ba94a17bafdcef37", + "version" : "1.30.2" + } + }, + { + "identity" : "eventsource", + "kind" : "remoteSourceControl", + "location" : "https://github.com/mattt/eventsource.git", + "state" : { + "revision" : "ca2a9d90cbe49e09b92f4b6ebd922c03ebea51d0", + "version" : "1.3.0" + } + }, + { + "identity" : "hummingbird", + "kind" : "remoteSourceControl", + "location" : "https://github.com/hummingbird-project/hummingbird.git", + "state" : { + "revision" : "e98b27919198f7578e8e349f701655b13635fb7c", + "version" : "2.18.3" + } + }, + { + "identity" : "swift-algorithms", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-algorithms.git", + "state" : { + "revision" : "87e50f483c54e6efd60e885f7f5aa946cee68023", + "version" : "1.2.1" + } + }, + { + "identity" : "swift-asn1", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-asn1.git", + "state" : { + "revision" : "810496cf121e525d660cd0ea89a758740476b85f", + "version" : "1.5.1" + } + }, + { + "identity" : "swift-async-algorithms", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-async-algorithms.git", + "state" : { + "revision" : "6c050d5ef8e1aa6342528460db614e9770d7f804", + "version" : "1.1.1" + } + }, + { + "identity" : "swift-atomics", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-atomics.git", + "state" : { + "revision" : "b601256eab081c0f92f059e12818ac1d4f178ff7", + "version" : "1.3.0" + } + }, + { + "identity" : "swift-certificates", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-certificates.git", + "state" : { + "revision" : "133a347911b6ad0fc8fe3bf46ca90c66cff97130", + "version" : "1.17.0" + } + }, + { + "identity" : "swift-collections", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-collections.git", + "state" : { + "revision" : "7b847a3b7008b2dc2f47ca3110d8c782fb2e5c7e", + "version" : "1.3.0" + } + }, + { + "identity" : "swift-configuration", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-configuration.git", + "state" : { + "revision" : "3528deb75256d7dcbb0d71fa75077caae0a8c749", + "version" : "1.0.0" + } + }, + { + "identity" : "swift-crypto", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-crypto.git", + "state" : { + "revision" : "6f70fa9eab24c1fd982af18c281c4525d05e3095", + "version" : "4.2.0" + } + }, + { + "identity" : "swift-distributed-tracing", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-distributed-tracing.git", + "state" : { + "revision" : "baa932c1336f7894145cbaafcd34ce2dd0b77c97", + "version" : "1.3.1" + } + }, + { + "identity" : "swift-http-structured-headers", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-http-structured-headers.git", + "state" : { + "revision" : "76d7627bd88b47bf5a0f8497dd244885960dde0b", + "version" : "1.6.0" + } + }, + { + "identity" : "swift-http-types", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-http-types.git", + "state" : { + "revision" : "45eb0224913ea070ec4fba17291b9e7ecf4749ca", + "version" : "1.5.1" + } + }, + { + "identity" : "swift-log", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-log.git", + "state" : { + "revision" : "bc386b95f2a16ccd0150a8235e7c69eab2b866ca", + "version" : "1.8.0" + } + }, + { + "identity" : "swift-metrics", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-metrics.git", + "state" : { + "revision" : "0743a9364382629da3bf5677b46a2c4b1ce5d2a6", + "version" : "2.7.1" + } + }, + { + "identity" : "swift-nio", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-nio.git", + "state" : { + "revision" : "a1605a3303a28e14d822dec8aaa53da8a9490461", + "version" : "2.92.0" + } + }, + { + "identity" : "swift-nio-extras", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-nio-extras.git", + "state" : { + "revision" : "1c90641b02b6ab47c6d0db2063a12198b04e83e2", + "version" : "1.31.2" + } + }, + { + "identity" : "swift-nio-http2", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-nio-http2.git", + "state" : { + "revision" : "c2ba4cfbb83f307c66f5a6df6bb43e3c88dfbf80", + "version" : "1.39.0" + } + }, + { + "identity" : "swift-nio-ssl", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-nio-ssl.git", + "state" : { + "revision" : "173cc69a058623525a58ae6710e2f5727c663793", + "version" : "2.36.0" + } + }, + { + "identity" : "swift-nio-transport-services", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-nio-transport-services.git", + "state" : { + "revision" : "60c3e187154421171721c1a38e800b390680fb5d", + "version" : "1.26.0" + } + }, + { + "identity" : "swift-numerics", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-numerics.git", + "state" : { + "revision" : "0c0290ff6b24942dadb83a929ffaaa1481df04a2", + "version" : "1.1.1" + } + }, + { + "identity" : "swift-service-context", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-service-context.git", + "state" : { + "revision" : "1983448fefc717a2bc2ebde5490fe99873c5b8a6", + "version" : "1.2.1" + } + }, + { + "identity" : "swift-service-lifecycle", + "kind" : "remoteSourceControl", + "location" : "https://github.com/swift-server/swift-service-lifecycle.git", + "state" : { + "revision" : "1de37290c0ab3c5a96028e0f02911b672fd42348", + "version" : "2.9.1" + } + }, + { + "identity" : "swift-system", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-system.git", + "state" : { + "revision" : "395a77f0aa927f0ff73941d7ac35f2b46d47c9db", + "version" : "1.6.3" + } + } + ], + "version" : 3 +} diff --git a/Examples/HummingbirdIntegration/Package.swift b/Examples/HummingbirdIntegration/Package.swift new file mode 100644 index 00000000..c7af3739 --- /dev/null +++ b/Examples/HummingbirdIntegration/Package.swift @@ -0,0 +1,24 @@ +// swift-tools-version: 6.0 + +import PackageDescription + +let package = Package( + name: "HummingbirdMCPExample", + platforms: [ + .macOS(.v14) + ], + dependencies: [ + .package(url: "https://github.com/hummingbird-project/hummingbird.git", from: "2.0.0"), + .package(path: "../.."), // MCP Swift SDK + ], + targets: [ + .executableTarget( + name: "HummingbirdMCPExample", + dependencies: [ + .product(name: "Hummingbird", package: "hummingbird"), + .product(name: "MCP", package: "mcp-swift-sdk"), + ], + path: "Sources" + ) + ] +) diff --git a/Examples/HummingbirdIntegration/Sources/main.swift b/Examples/HummingbirdIntegration/Sources/main.swift new file mode 100644 index 00000000..0ea7243c --- /dev/null +++ b/Examples/HummingbirdIntegration/Sources/main.swift @@ -0,0 +1,333 @@ +/// Hummingbird MCP Server Example +/// +/// This example demonstrates how to integrate an MCP server with the Hummingbird web framework. +/// It follows the TypeScript SDK's pattern from `examples/server/src/simpleStreamableHttp.ts`. +/// +/// ## Architecture +/// +/// - ONE `Server` instance is shared across all HTTP clients +/// - Each client session gets its own `HTTPServerTransport` +/// - The `SessionManager` actor manages transport instances by session ID +/// - Request capture in the Server ensures responses route to the correct client +/// +/// ## Endpoints +/// +/// - `POST /mcp` - Handle JSON-RPC requests (initialize, tools/list, tools/call, etc.) +/// - `GET /mcp` - Server-Sent Events stream for server-initiated notifications +/// - `DELETE /mcp` - Terminate a session +/// +/// ## Running +/// +/// ```bash +/// cd Examples/HummingbirdIntegration +/// swift run +/// ``` +/// +/// The server will listen on http://localhost:3000/mcp + +import Foundation +import HTTPTypes +import Hummingbird +import Logging +import MCP + +// MARK: - Server Setup + +/// Create the MCP server (ONE instance for all clients) +let mcpServer = MCP.Server( + name: "hummingbird-mcp-example", + version: "1.0.0", + capabilities: .init(tools: .init()) +) + +/// Register tool handlers +func setUpToolHandlers() async { + // Register tool list handler + await mcpServer.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool( + name: "echo", + description: "Echoes back the input message", + inputSchema: [ + "type": "object", + "properties": [ + "message": ["type": "string", "description": "The message to echo"] + ], + "required": ["message"] + ] + ), + Tool( + name: "add", + description: "Adds two numbers", + inputSchema: [ + "type": "object", + "properties": [ + "a": ["type": "number", "description": "First number"], + "b": ["type": "number", "description": "Second number"] + ], + "required": ["a", "b"] + ] + ), + ]) + } + + // Register tool call handler + await mcpServer.withRequestHandler(CallTool.self) { request, _ in + switch request.name { + case "echo": + let message = request.arguments?["message"]?.stringValue ?? "No message provided" + return CallTool.Result(content: [.text(message)]) + + case "add": + let a = request.arguments?["a"]?.doubleValue ?? 0 + let b = request.arguments?["b"]?.doubleValue ?? 0 + return CallTool.Result(content: [.text("Result: \(a + b)")]) + + default: + return CallTool.Result(content: [.text("Unknown tool: \(request.name)")], isError: true) + } + } +} + +// MARK: - Session Management + +/// Session manager for tracking active sessions +let sessionManager = SessionManager(maxSessions: 100) + +/// Logger for the example +let logger = Logger(label: "mcp.example.hummingbird") + +// MARK: - Request Context + +struct MCPRequestContext: RequestContext { + var coreContext: CoreRequestContextStorage + + init(source: Source) { + self.coreContext = .init(source: source) + } +} + +// MARK: - HTTP Handlers + +/// Handle POST /mcp requests +func handlePost(request: Request, context: MCPRequestContext) async throws -> Response { + // Get session ID from header (if present) + let sessionId = request.headers[HTTPField.Name(HTTPHeader.sessionId)!] + + // Read request body + let body = try await request.body.collect(upTo: .max) + let data = Data(buffer: body) + + // Check if this is an initialize request + let isInitializeRequest = String(data: data, encoding: .utf8)?.contains("\"method\":\"initialize\"") ?? false + + // Get or create transport + let transport: HTTPServerTransport + + if let sid = sessionId, let existing = await sessionManager.transport(forSessionId: sid) { + // Reuse existing transport for this session + transport = existing + } else if isInitializeRequest { + // Check capacity + guard await sessionManager.canAddSession() else { + return Response( + status: .serviceUnavailable, + headers: [.retryAfter: "60"], + body: .init(byteBuffer: .init(string: "Server at capacity")) + ) + } + + // Generate session ID upfront so we can store the transport + let newSessionId = UUID().uuidString + + // Create new transport with session callbacks + let newTransport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { newSessionId }, + onSessionInitialized: { sessionId in + logger.info("Session initialized: \(sessionId)") + }, + onSessionClosed: { sessionId in + await sessionManager.remove(sessionId) + logger.info("Session closed: \(sessionId)") + } + ) + ) + + // Store the transport immediately (we know the session ID) + await sessionManager.store(newTransport, forSessionId: newSessionId) + transport = newTransport + + // Connect transport to server + try await mcpServer.start(transport: transport) + } else if sessionId != nil { + // Client sent a session ID that no longer exists + return Response( + status: .notFound, + body: .init(byteBuffer: .init(string: "Session expired. Try reconnecting.")) + ) + } else { + // No session ID and not an initialize request + return Response( + status: .badRequest, + body: .init(byteBuffer: .init(string: "Missing \(HTTPHeader.sessionId) header")) + ) + } + + // Create the MCP HTTP request for the transport + let mcpRequest = MCP.HTTPRequest( + method: "POST", + headers: extractHeaders(from: request), + body: data + ) + + // Handle the request + let mcpResponse = await transport.handleRequest(mcpRequest) + + // Build response + return buildResponse(from: mcpResponse) +} + +/// Handle GET /mcp requests (SSE stream for server-initiated notifications) +func handleGet(request: Request, context: MCPRequestContext) async throws -> Response { + guard let sessionId = request.headers[HTTPField.Name(HTTPHeader.sessionId)!], + let transport = await sessionManager.transport(forSessionId: sessionId) + else { + return Response( + status: .badRequest, + body: .init(byteBuffer: .init(string: "Invalid or missing session ID")) + ) + } + + let mcpRequest = MCP.HTTPRequest( + method: "GET", + headers: extractHeaders(from: request) + ) + + let mcpResponse = await transport.handleRequest(mcpRequest) + + return buildResponse(from: mcpResponse) +} + +/// Handle DELETE /mcp requests (session termination) +func handleDelete(request: Request, context: MCPRequestContext) async throws -> Response { + guard let sessionId = request.headers[HTTPField.Name(HTTPHeader.sessionId)!], + let transport = await sessionManager.transport(forSessionId: sessionId) + else { + return Response( + status: .notFound, + body: .init(byteBuffer: .init(string: "Session not found")) + ) + } + + let mcpRequest = MCP.HTTPRequest( + method: "DELETE", + headers: extractHeaders(from: request) + ) + + let mcpResponse = await transport.handleRequest(mcpRequest) + + return Response(status: .init(code: mcpResponse.statusCode)) +} + +// MARK: - Helper Functions + +/// Extract headers from Hummingbird request to dictionary +func extractHeaders(from request: Request) -> [String: String] { + var headers: [String: String] = [:] + for field in request.headers { + headers[field.name.rawName] = field.value + } + return headers +} + +/// Build a Hummingbird Response from an MCP HTTPResponse +func buildResponse(from mcpResponse: MCP.HTTPResponse) -> Response { + var responseHeaders = HTTPFields() + for (key, value) in mcpResponse.headers { + if let name = HTTPField.Name(key) { + responseHeaders[name] = value + } + } + + let status = HTTPResponse.Status(code: mcpResponse.statusCode) + + if let stream = mcpResponse.stream { + // SSE response - stream the events + let responseBody = ResponseBody(asyncSequence: SSEResponseSequence(stream: stream)) + return Response( + status: status, + headers: responseHeaders, + body: responseBody + ) + } else if let body = mcpResponse.body { + // JSON response + return Response( + status: status, + headers: responseHeaders, + body: .init(byteBuffer: .init(data: body)) + ) + } else { + // No content (e.g., 202 Accepted for notifications) + return Response( + status: status, + headers: responseHeaders + ) + } +} + +/// Async sequence wrapper for SSE stream +struct SSEResponseSequence: AsyncSequence, Sendable { + typealias Element = ByteBuffer + + let stream: AsyncThrowingStream + + struct AsyncIterator: AsyncIteratorProtocol { + var iterator: AsyncThrowingStream.AsyncIterator + + mutating func next() async throws -> ByteBuffer? { + guard let data = try await iterator.next() else { + return nil + } + return ByteBuffer(data: data) + } + } + + func makeAsyncIterator() -> AsyncIterator { + AsyncIterator(iterator: stream.makeAsyncIterator()) + } +} + +// MARK: - Main + +@main +struct HummingbirdMCPExample { + static func main() async throws { + // Set up tool handlers + await setUpToolHandlers() + + // Create router + let router = Router(context: MCPRequestContext.self) + + // MCP endpoints + router.post("/mcp", use: handlePost) + router.get("/mcp", use: handleGet) + router.delete("/mcp", use: handleDelete) + + // Health check + router.get("/health") { _, _ in + Response(status: .ok, body: .init(byteBuffer: .init(string: "OK"))) + } + + // Create and run application + let app = Application( + router: router, + configuration: .init(address: .hostname("localhost", port: 3000)) + ) + + logger.info("Starting MCP server on http://localhost:3000/mcp") + logger.info("Available tools: echo, add") + + try await app.run() + } +} diff --git a/Examples/README.md b/Examples/README.md new file mode 100644 index 00000000..62f8563e --- /dev/null +++ b/Examples/README.md @@ -0,0 +1,190 @@ +# MCP Swift SDK Examples + +This directory contains example integrations showing how to build HTTP-based MCP servers with popular Swift web frameworks. + +## Available Examples + +### [HummingbirdIntegration](./HummingbirdIntegration) + +Integration with [Hummingbird](https://github.com/hummingbird-project/hummingbird), a lightweight, flexible Swift web framework. + +```bash +cd HummingbirdIntegration +swift run +# Server starts on http://localhost:3000/mcp +``` + +### [VaporIntegration](./VaporIntegration) + +Integration with [Vapor](https://vapor.codes/), a popular full-featured Swift web framework. + +```bash +cd VaporIntegration +swift run +# Server starts on http://localhost:8080/mcp +``` + +## Architecture Pattern + +Both examples follow the same architecture pattern from the TypeScript SDK: + +1. **One Server instance** is shared across all HTTP clients +2. **Each client session** gets its own `HTTPServerTransport` +3. **SessionManager** tracks active transports by session ID +4. **Request capture** ensures responses route to the correct client + +``` + ┌─────────────────────────────┐ + │ MCP Server (shared) │ + │ - Tool handlers │ + │ - Resource handlers │ + └─────────────┬───────────────┘ + │ + ┌─────────────────────────┼─────────────────────────┐ + │ │ │ + ▼ ▼ ▼ +┌───────────────┐ ┌───────────────┐ ┌───────────────┐ +│ Transport A │ │ Transport B │ │ Transport C │ +│ (session-1) │ │ (session-2) │ │ (session-3) │ +└───────┬───────┘ └───────┬───────┘ └───────┬───────┘ + │ │ │ + ▼ ▼ ▼ + Client A Client B Client C +``` + +## HTTP Endpoints + +Both examples implement the standard MCP HTTP endpoints: + +| Endpoint | Method | Description | +|----------|--------|-------------| +| `/mcp` | POST | Handle JSON-RPC requests (initialize, tools/list, tools/call, etc.) | +| `/mcp` | GET | Server-Sent Events stream for server-initiated notifications | +| `/mcp` | DELETE | Terminate a session | +| `/health` | GET | Health check endpoint | + +## Testing with curl + +### Initialize a session + +```bash +curl -X POST http://localhost:3000/mcp \ + -H "Content-Type: application/json" \ + -H "Accept: application/json, text/event-stream" \ + -d '{ + "jsonrpc": "2.0", + "method": "initialize", + "id": "1", + "params": { + "protocolVersion": "2024-11-05", + "capabilities": {}, + "clientInfo": {"name": "curl-test", "version": "1.0"} + } + }' +``` + +Save the `Mcp-Session-Id` header from the response for subsequent requests. + +### List tools + +```bash +curl -X POST http://localhost:3000/mcp \ + -H "Content-Type: application/json" \ + -H "Accept: application/json, text/event-stream" \ + -H "Mcp-Session-Id: YOUR_SESSION_ID" \ + -H "Mcp-Protocol-Version: 2024-11-05" \ + -d '{"jsonrpc": "2.0", "method": "tools/list", "id": "2"}' +``` + +### Call a tool + +```bash +curl -X POST http://localhost:3000/mcp \ + -H "Content-Type: application/json" \ + -H "Accept: application/json, text/event-stream" \ + -H "Mcp-Session-Id: YOUR_SESSION_ID" \ + -H "Mcp-Protocol-Version: 2024-11-05" \ + -d '{ + "jsonrpc": "2.0", + "method": "tools/call", + "id": "3", + "params": { + "name": "echo", + "arguments": {"message": "Hello, MCP!"} + } + }' +``` + +### Terminate session + +```bash +curl -X DELETE http://localhost:3000/mcp \ + -H "Mcp-Session-Id: YOUR_SESSION_ID" \ + -H "Mcp-Protocol-Version: 2024-11-05" +``` + +## Key Components + +### SessionManager + +The `SessionManager` actor provides thread-safe session storage: + +```swift +let sessionManager = SessionManager(maxSessions: 100) + +// Store a transport +await sessionManager.store(transport, forSessionId: sessionId) + +// Get a transport +if let transport = await sessionManager.transport(forSessionId: sessionId) { + // Use transport +} + +// Remove a transport +await sessionManager.remove(sessionId) + +// Cleanup stale sessions +await sessionManager.cleanupStaleSessions(olderThan: .seconds(3600)) +``` + +### HTTPServerTransport + +The transport handles HTTP request/response multiplexing: + +```swift +let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { UUID().uuidString }, + onSessionInitialized: { sessionId in + // Called when session is initialized + await sessionManager.store(transport, forSessionId: sessionId) + }, + onSessionClosed: { sessionId in + // Called when session is terminated (DELETE request) + await sessionManager.remove(sessionId) + } + ) +) +``` + +## Stateless Mode + +For simpler deployments that don't need session persistence, you can run in stateless mode by omitting the `sessionIdGenerator`: + +```swift +// Stateless mode - no session tracking +let transport = HTTPServerTransport() +``` + +In stateless mode: +- No `Mcp-Session-Id` header is returned or required +- Each request is independent +- Server-initiated notifications are not supported (no GET endpoint) + +## Production Considerations + +1. **Session cleanup**: Implement periodic cleanup of stale sessions +2. **Connection limits**: Set `maxSessions` to prevent resource exhaustion +3. **Load balancing**: Use sticky sessions or shared session storage +4. **TLS**: Always use HTTPS in production +5. **Authentication**: Add authentication middleware as needed diff --git a/Examples/VaporIntegration/Package.resolved b/Examples/VaporIntegration/Package.resolved new file mode 100644 index 00000000..42c1b906 --- /dev/null +++ b/Examples/VaporIntegration/Package.resolved @@ -0,0 +1,267 @@ +{ + "originHash" : "28076ce7563165da628d782493ffb5aecf1e96b39274d05ddf9ab98ec415554d", + "pins" : [ + { + "identity" : "async-http-client", + "kind" : "remoteSourceControl", + "location" : "https://github.com/swift-server/async-http-client.git", + "state" : { + "revision" : "5dd84c7bb48b348751d7bbe7ba94a17bafdcef37", + "version" : "1.30.2" + } + }, + { + "identity" : "async-kit", + "kind" : "remoteSourceControl", + "location" : "https://github.com/vapor/async-kit.git", + "state" : { + "revision" : "6f3615ccf2ac3c2ae0c8087d527546e9544a43dd", + "version" : "1.21.0" + } + }, + { + "identity" : "console-kit", + "kind" : "remoteSourceControl", + "location" : "https://github.com/vapor/console-kit.git", + "state" : { + "revision" : "742f624a998cba2a9e653d9b1e91ad3f3a5dff6b", + "version" : "4.15.2" + } + }, + { + "identity" : "eventsource", + "kind" : "remoteSourceControl", + "location" : "https://github.com/mattt/eventsource.git", + "state" : { + "revision" : "ca2a9d90cbe49e09b92f4b6ebd922c03ebea51d0", + "version" : "1.3.0" + } + }, + { + "identity" : "multipart-kit", + "kind" : "remoteSourceControl", + "location" : "https://github.com/vapor/multipart-kit.git", + "state" : { + "revision" : "3498e60218e6003894ff95192d756e238c01f44e", + "version" : "4.7.1" + } + }, + { + "identity" : "routing-kit", + "kind" : "remoteSourceControl", + "location" : "https://github.com/vapor/routing-kit.git", + "state" : { + "revision" : "1a10ccea61e4248effd23b6e814999ce7bdf0ee0", + "version" : "4.9.3" + } + }, + { + "identity" : "swift-algorithms", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-algorithms.git", + "state" : { + "revision" : "87e50f483c54e6efd60e885f7f5aa946cee68023", + "version" : "1.2.1" + } + }, + { + "identity" : "swift-asn1", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-asn1.git", + "state" : { + "revision" : "810496cf121e525d660cd0ea89a758740476b85f", + "version" : "1.5.1" + } + }, + { + "identity" : "swift-async-algorithms", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-async-algorithms.git", + "state" : { + "revision" : "6c050d5ef8e1aa6342528460db614e9770d7f804", + "version" : "1.1.1" + } + }, + { + "identity" : "swift-atomics", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-atomics.git", + "state" : { + "revision" : "b601256eab081c0f92f059e12818ac1d4f178ff7", + "version" : "1.3.0" + } + }, + { + "identity" : "swift-certificates", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-certificates.git", + "state" : { + "revision" : "133a347911b6ad0fc8fe3bf46ca90c66cff97130", + "version" : "1.17.0" + } + }, + { + "identity" : "swift-collections", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-collections.git", + "state" : { + "revision" : "7b847a3b7008b2dc2f47ca3110d8c782fb2e5c7e", + "version" : "1.3.0" + } + }, + { + "identity" : "swift-crypto", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-crypto.git", + "state" : { + "revision" : "6f70fa9eab24c1fd982af18c281c4525d05e3095", + "version" : "4.2.0" + } + }, + { + "identity" : "swift-distributed-tracing", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-distributed-tracing.git", + "state" : { + "revision" : "baa932c1336f7894145cbaafcd34ce2dd0b77c97", + "version" : "1.3.1" + } + }, + { + "identity" : "swift-http-structured-headers", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-http-structured-headers.git", + "state" : { + "revision" : "76d7627bd88b47bf5a0f8497dd244885960dde0b", + "version" : "1.6.0" + } + }, + { + "identity" : "swift-http-types", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-http-types.git", + "state" : { + "revision" : "45eb0224913ea070ec4fba17291b9e7ecf4749ca", + "version" : "1.5.1" + } + }, + { + "identity" : "swift-log", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-log.git", + "state" : { + "revision" : "bc386b95f2a16ccd0150a8235e7c69eab2b866ca", + "version" : "1.8.0" + } + }, + { + "identity" : "swift-metrics", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-metrics.git", + "state" : { + "revision" : "0743a9364382629da3bf5677b46a2c4b1ce5d2a6", + "version" : "2.7.1" + } + }, + { + "identity" : "swift-nio", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-nio.git", + "state" : { + "revision" : "a1605a3303a28e14d822dec8aaa53da8a9490461", + "version" : "2.92.0" + } + }, + { + "identity" : "swift-nio-extras", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-nio-extras.git", + "state" : { + "revision" : "1c90641b02b6ab47c6d0db2063a12198b04e83e2", + "version" : "1.31.2" + } + }, + { + "identity" : "swift-nio-http2", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-nio-http2.git", + "state" : { + "revision" : "c2ba4cfbb83f307c66f5a6df6bb43e3c88dfbf80", + "version" : "1.39.0" + } + }, + { + "identity" : "swift-nio-ssl", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-nio-ssl.git", + "state" : { + "revision" : "173cc69a058623525a58ae6710e2f5727c663793", + "version" : "2.36.0" + } + }, + { + "identity" : "swift-nio-transport-services", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-nio-transport-services.git", + "state" : { + "revision" : "60c3e187154421171721c1a38e800b390680fb5d", + "version" : "1.26.0" + } + }, + { + "identity" : "swift-numerics", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-numerics.git", + "state" : { + "revision" : "0c0290ff6b24942dadb83a929ffaaa1481df04a2", + "version" : "1.1.1" + } + }, + { + "identity" : "swift-service-context", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-service-context.git", + "state" : { + "revision" : "1983448fefc717a2bc2ebde5490fe99873c5b8a6", + "version" : "1.2.1" + } + }, + { + "identity" : "swift-service-lifecycle", + "kind" : "remoteSourceControl", + "location" : "https://github.com/swift-server/swift-service-lifecycle.git", + "state" : { + "revision" : "1de37290c0ab3c5a96028e0f02911b672fd42348", + "version" : "2.9.1" + } + }, + { + "identity" : "swift-system", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-system.git", + "state" : { + "revision" : "395a77f0aa927f0ff73941d7ac35f2b46d47c9db", + "version" : "1.6.3" + } + }, + { + "identity" : "vapor", + "kind" : "remoteSourceControl", + "location" : "https://github.com/vapor/vapor.git", + "state" : { + "revision" : "f7090db27390ebc4cadbff06d76fe8ce79d6ece6", + "version" : "4.120.0" + } + }, + { + "identity" : "websocket-kit", + "kind" : "remoteSourceControl", + "location" : "https://github.com/vapor/websocket-kit.git", + "state" : { + "revision" : "8666c92dbbb3c8eefc8008c9c8dcf50bfd302167", + "version" : "2.16.1" + } + } + ], + "version" : 3 +} diff --git a/Examples/VaporIntegration/Package.swift b/Examples/VaporIntegration/Package.swift new file mode 100644 index 00000000..3bd2dd3c --- /dev/null +++ b/Examples/VaporIntegration/Package.swift @@ -0,0 +1,24 @@ +// swift-tools-version: 6.0 + +import PackageDescription + +let package = Package( + name: "VaporMCPExample", + platforms: [ + .macOS(.v14) + ], + dependencies: [ + .package(url: "https://github.com/vapor/vapor.git", from: "4.0.0"), + .package(path: "../.."), // MCP Swift SDK + ], + targets: [ + .executableTarget( + name: "VaporMCPExample", + dependencies: [ + .product(name: "Vapor", package: "vapor"), + .product(name: "MCP", package: "mcp-swift-sdk"), + ], + path: "Sources" + ) + ] +) diff --git a/Examples/VaporIntegration/Sources/main.swift b/Examples/VaporIntegration/Sources/main.swift new file mode 100644 index 00000000..fd6985e6 --- /dev/null +++ b/Examples/VaporIntegration/Sources/main.swift @@ -0,0 +1,279 @@ +/// Vapor MCP Server Example +/// +/// This example demonstrates how to integrate an MCP server with the Vapor web framework. +/// It follows the TypeScript SDK's pattern from `examples/server/src/simpleStreamableHttp.ts`. +/// +/// ## Architecture +/// +/// - ONE `Server` instance is shared across all HTTP clients +/// - Each client session gets its own `HTTPServerTransport` +/// - The `SessionManager` actor manages transport instances by session ID +/// - Request capture in the Server ensures responses route to the correct client +/// +/// ## Endpoints +/// +/// - `POST /mcp` - Handle JSON-RPC requests (initialize, tools/list, tools/call, etc.) +/// - `GET /mcp` - Server-Sent Events stream for server-initiated notifications +/// - `DELETE /mcp` - Terminate a session +/// +/// ## Running +/// +/// ```bash +/// cd Examples/VaporIntegration +/// swift run +/// ``` +/// +/// The server will listen on http://localhost:8080/mcp + +import Foundation +import MCP +import Vapor + +// MARK: - Server Setup + +/// Create the MCP server (ONE instance for all clients) +let mcpServer = MCP.Server( + name: "vapor-mcp-example", + version: "1.0.0", + capabilities: .init(tools: .init()) +) + +/// Register tool handlers +func setUpToolHandlers() async { + // Register tool list handler + await mcpServer.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool( + name: "echo", + description: "Echoes back the input message", + inputSchema: [ + "type": "object", + "properties": [ + "message": ["type": "string", "description": "The message to echo"] + ], + "required": ["message"] + ] + ), + Tool( + name: "add", + description: "Adds two numbers", + inputSchema: [ + "type": "object", + "properties": [ + "a": ["type": "number", "description": "First number"], + "b": ["type": "number", "description": "Second number"] + ], + "required": ["a", "b"] + ] + ), + ]) + } + + // Register tool call handler + await mcpServer.withRequestHandler(CallTool.self) { request, _ in + switch request.name { + case "echo": + let message = request.arguments?["message"]?.stringValue ?? "No message provided" + return CallTool.Result(content: [.text(message)]) + + case "add": + let a = request.arguments?["a"]?.doubleValue ?? 0 + let b = request.arguments?["b"]?.doubleValue ?? 0 + return CallTool.Result(content: [.text("Result: \(a + b)")]) + + default: + return CallTool.Result(content: [.text("Unknown tool: \(request.name)")], isError: true) + } + } +} + +// MARK: - Session Management + +/// Session manager for tracking active sessions +let sessionManager = SessionManager(maxSessions: 100) + +// MARK: - HTTP Handlers + +/// Handle POST /mcp requests +func handlePost(_ req: Vapor.Request) async throws -> Vapor.Response { + // Get session ID from header (if present) + let sessionId = req.headers.first(name: HTTPHeader.sessionId) + + // Read request body + guard let bodyData = req.body.data else { + throw Abort(.badRequest, reason: "Missing request body") + } + let data = Data(buffer: bodyData) + + // Check if this is an initialize request + let isInitializeRequest = String(data: data, encoding: .utf8)?.contains("\"method\":\"initialize\"") ?? false + + // Get or create transport + let transport: HTTPServerTransport + + if let sid = sessionId, let existing = await sessionManager.transport(forSessionId: sid) { + // Reuse existing transport for this session + transport = existing + } else if isInitializeRequest { + // Check capacity + guard await sessionManager.canAddSession() else { + throw Abort(.serviceUnavailable, reason: "Server at capacity") + } + + // Generate session ID upfront so we can store the transport + let newSessionId = UUID().uuidString + + // Create new transport with session callbacks + let newTransport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { newSessionId }, + onSessionInitialized: { sessionId in + req.logger.info("Session initialized: \(sessionId)") + }, + onSessionClosed: { sessionId in + await sessionManager.remove(sessionId) + req.logger.info("Session closed: \(sessionId)") + } + ) + ) + + // Store the transport immediately (we know the session ID) + await sessionManager.store(newTransport, forSessionId: newSessionId) + transport = newTransport + + // Connect transport to server + try await mcpServer.start(transport: transport) + } else if sessionId != nil { + // Client sent a session ID that no longer exists + throw Abort(.notFound, reason: "Session expired. Try reconnecting.") + } else { + // No session ID and not an initialize request + throw Abort(.badRequest, reason: "Missing \(HTTPHeader.sessionId) header") + } + + // Create the MCP HTTP request for the transport + let mcpRequest = MCP.HTTPRequest( + method: "POST", + headers: extractHeaders(from: req), + body: data + ) + + // Handle the request + let mcpResponse = await transport.handleRequest(mcpRequest) + + // Build response + return buildVaporResponse(from: mcpResponse, for: req) +} + +/// Handle GET /mcp requests (SSE stream for server-initiated notifications) +func handleGet(_ req: Vapor.Request) async throws -> Vapor.Response { + guard let sessionId = req.headers.first(name: HTTPHeader.sessionId), + let transport = await sessionManager.transport(forSessionId: sessionId) + else { + throw Abort(.badRequest, reason: "Invalid or missing session ID") + } + + let mcpRequest = MCP.HTTPRequest( + method: "GET", + headers: extractHeaders(from: req) + ) + + let mcpResponse = await transport.handleRequest(mcpRequest) + + return buildVaporResponse(from: mcpResponse, for: req) +} + +/// Handle DELETE /mcp requests (session termination) +func handleDelete(_ req: Vapor.Request) async throws -> Vapor.Response { + guard let sessionId = req.headers.first(name: HTTPHeader.sessionId), + let transport = await sessionManager.transport(forSessionId: sessionId) + else { + throw Abort(.notFound, reason: "Session not found") + } + + let mcpRequest = MCP.HTTPRequest( + method: "DELETE", + headers: extractHeaders(from: req) + ) + + let mcpResponse = await transport.handleRequest(mcpRequest) + + return Vapor.Response(status: .init(statusCode: mcpResponse.statusCode)) +} + +// MARK: - Helper Functions + +/// Extract headers from Vapor request to dictionary +func extractHeaders(from req: Vapor.Request) -> [String: String] { + var headers: [String: String] = [:] + for (name, value) in req.headers { + headers[name] = value + } + return headers +} + +/// Build a Vapor Response from an MCP HTTPResponse +func buildVaporResponse(from mcpResponse: MCP.HTTPResponse, for req: Vapor.Request) -> Vapor.Response { + var headers = HTTPHeaders() + for (key, value) in mcpResponse.headers { + headers.add(name: key, value: value) + } + + let status = HTTPResponseStatus(statusCode: mcpResponse.statusCode) + + if let stream = mcpResponse.stream { + // SSE response - create streaming body + let response = Vapor.Response(status: status, headers: headers) + response.body = .init(asyncStream: { writer in + do { + for try await data in stream { + try await writer.write(.buffer(.init(data: data))) + } + try await writer.write(.end) + } catch { + req.logger.error("SSE stream error: \(error)") + } + }) + return response + } else if let body = mcpResponse.body { + // JSON response + return Vapor.Response( + status: status, + headers: headers, + body: .init(data: body) + ) + } else { + // No content (e.g., 202 Accepted for notifications) + return Vapor.Response(status: status, headers: headers) + } +} + +// MARK: - Main + +@main +struct VaporMCPExample { + static func main() async throws { + // Set up tool handlers + await setUpToolHandlers() + + // Create Vapor application + let env = try Environment.detect() + let app = try await Application.make(env) + + // Configure routes + app.post("mcp", use: handlePost) + app.get("mcp", use: handleGet) + app.delete("mcp", use: handleDelete) + + // Health check + app.get("health") { _ in + "OK" + } + + app.logger.info("Starting MCP server on http://localhost:8080/mcp") + app.logger.info("Available tools: echo, add") + + try await app.execute() + try await app.asyncShutdown() + } +} diff --git a/Package.resolved b/Package.resolved index 5e9023c5..fb776dd5 100644 --- a/Package.resolved +++ b/Package.resolved @@ -1,10 +1,10 @@ { - "originHash" : "08de61941b7919a65e36c0e34f8c1c41995469b86a39122158b75b4a68c4527d", + "originHash" : "371f3dfcfa1201fc8d50e924ad31f9ebc4f90242924df1275958ac79df15dc12", "pins" : [ { "identity" : "eventsource", "kind" : "remoteSourceControl", - "location" : "https://github.com/loopwork-ai/eventsource.git", + "location" : "https://github.com/mattt/eventsource.git", "state" : { "revision" : "e83f076811f32757305b8bf69ac92d05626ffdd7", "version" : "1.1.0" diff --git a/README.md b/README.md index ac23a4c3..3061811c 100644 --- a/README.md +++ b/README.md @@ -428,33 +428,37 @@ Register tool handlers to respond to client tool calls: ```swift // Register a tool list handler -await server.withMethodHandler(ListTools.self) { _ in +await server.withRequestHandler(ListTools.self) { _ in let tools = [ Tool( name: "weather", description: "Get current weather for a location", - inputSchema: .object([ - "properties": .object([ - "location": .string("City name or coordinates"), - "units": .string("Units of measurement, e.g., metric, imperial") - ]) - ]) + inputSchema: [ + "type": "object", + "properties": [ + "location": ["type": "string", "description": "City name or coordinates"], + "units": ["type": "string", "description": "Units of measurement (metric or imperial)"] + ], + "required": ["location"] + ] ), Tool( name: "calculator", description: "Perform calculations", - inputSchema: .object([ - "properties": .object([ - "expression": .string("Mathematical expression to evaluate") - ]) - ]) + inputSchema: [ + "type": "object", + "properties": [ + "expression": ["type": "string", "description": "Mathematical expression to evaluate"] + ], + "required": ["expression"] + ] ) ] return .init(tools: tools) } // Register a tool call handler -await server.withMethodHandler(CallTool.self) { params in +await server.withRequestHandler(CallTool.self) { params in switch params.name { case "weather": let location = params.arguments?["location"]?.stringValue ?? "Unknown" @@ -479,13 +483,116 @@ await server.withMethodHandler(CallTool.self) { params in } ``` +### Progress Notifications + +For long-running operations, servers can send progress notifications to keep clients informed. Clients include a `progressToken` in the request's `_meta` field, and servers use that token when sending progress updates. + +The `RequestHandlerContext` provides convenience methods for sending notifications: + +```swift +// Register a tool call handler that reports progress +await server.withRequestHandler(CallTool.self) { params, context in + switch params.name { + case "long-running-task": + // Extract the progress token from the request metadata + // Clients that want progress updates will include this in _meta.progressToken + let progressToken = params._meta?.progressToken + + // Helper to send progress if client requested it + func reportProgress(_ progress: Double, _ message: String) async throws { + if let token = progressToken { + try await context.sendProgress( + token: token, + progress: progress, + total: 100.0, + message: message + ) + } + } + + try await reportProgress(0.0, "Starting task...") + + // Perform first part of work + await performFirstStep() + + try await reportProgress(50.0, "Halfway complete...") + + // Perform second part of work + await performSecondStep() + + try await reportProgress(100.0, "Task complete!") + + return .init(content: [.text("Task completed successfully")], isError: false) + + default: + return .init(content: [.text("Unknown tool")], isError: true) + } +} +``` + +The `RequestHandlerContext` provides these convenience methods: + +| Method | Description | +|--------|-------------| +| `sendProgress(token:progress:total:message:)` | Send progress updates for long-running operations | +| `sendLogMessage(level:logger:data:)` | Send log messages to the client | +| `sendResourceListChanged()` | Notify client that available resources have changed | +| `sendResourceUpdated(uri:)` | Notify client that a specific resource was updated | +| `sendToolListChanged()` | Notify client that available tools have changed | +| `sendPromptListChanged()` | Notify client that available prompts have changed | + +You can also use `sendMessage()` for full control over notification parameters: + +```swift +try await context.sendMessage(ProgressNotification.message(.init( + progressToken: .integer(42), // Tokens can be strings or integers + progress: 75.0, + total: 100.0, + message: "Processing item 3 of 4..." +))) +``` + +#### Handling Progress Notifications (Client) + +Clients register handlers to receive progress notifications, and include a `progressToken` in requests where they want progress updates: + +```swift +// Register a progress notification handler +await client.onNotification(ProgressNotification.self) { message in + let params = message.params + + // Match the token to know which request this progress is for + guard params.progressToken == .string("my-task-token") else { return } + + // Calculate percentage if total is known + if let total = params.total, total > 0 { + let percentage = (params.progress / total) * 100 + print("Progress: \(percentage)%") + } + + // Show human-readable message if available + if let progressMessage = params.message { + print("Status: \(progressMessage)") + } +} + +// Make a request with a progress token to receive updates +let result = try await client.request( + CallTool.request(.init( + name: "long-running-task", + arguments: [:], + _meta: RequestMeta(progressToken: "my-task-token") + )) +) +``` + ### Resources Implement resource handlers for data access: ```swift // Register a resource list handler -await server.withMethodHandler(ListResources.self) { params in +await server.withRequestHandler(ListResources.self) { params in let resources = [ Resource( name: "Knowledge Base Articles", @@ -502,7 +609,7 @@ await server.withMethodHandler(ListResources.self) { params in } // Register a resource read handler -await server.withMethodHandler(ReadResource.self) { params in +await server.withRequestHandler(ReadResource.self) { params in switch params.uri { case "resource://knowledge-base/articles": return .init(contents: [Resource.Content.text("# Knowledge Base\n\nThis is the content of the knowledge base...", uri: params.uri)]) @@ -528,7 +635,7 @@ await server.withMethodHandler(ReadResource.self) { params in } // Register a resource subscribe handler -await server.withMethodHandler(ResourceSubscribe.self) { params in +await server.withRequestHandler(ResourceSubscribe.self) { params in // Store subscription for later notifications. // Client identity for multi-client scenarios needs to be managed by the server application, // potentially using information from the initialize handshake if the server handles one client post-init. @@ -544,7 +651,7 @@ Implement prompt handlers: ```swift // Register a prompt list handler -await server.withMethodHandler(ListPrompts.self) { params in +await server.withRequestHandler(ListPrompts.self) { params in let prompts = [ Prompt( name: "interview", @@ -568,7 +675,7 @@ await server.withMethodHandler(ListPrompts.self) { params in } // Register a prompt get handler -await server.withMethodHandler(GetPrompt.self) { params in +await server.withRequestHandler(GetPrompt.self) { params in switch params.name { case "interview": let position = params.arguments?["position"]?.stringValue ?? "Software Engineer" @@ -724,14 +831,14 @@ let server = Server( ) // Add handlers directly to the server -await server.withMethodHandler(ListTools.self) { _ in +await server.withRequestHandler(ListTools.self) { _ in // Your implementation return .init(tools: [ Tool(name: "example", description: "An example tool") ]) } -await server.withMethodHandler(CallTool.self) { params in +await server.withRequestHandler(CallTool.self) { params in // Your implementation return .init(content: [.text("Tool result")], isError: false) } diff --git a/Sources/MCP/Base/Annotations.swift b/Sources/MCP/Base/Annotations.swift new file mode 100644 index 00000000..2edf9769 --- /dev/null +++ b/Sources/MCP/Base/Annotations.swift @@ -0,0 +1,40 @@ +import Foundation + +/// The sender or recipient of messages and data in a conversation. +public enum Role: String, Hashable, Codable, Sendable { + /// A user message + case user + /// An assistant message + case assistant +} + +/// Optional annotations for content, used to inform how objects are used or displayed. +/// +/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/basic/ +public struct Annotations: Hashable, Codable, Sendable { + // TODO: Deprecate in a future version + /// Backwards compatibility alias for top-level `Role`. + public typealias Role = MCP.Role + + /// Describes who the intended audience of this object or data is. + /// It can include multiple entries to indicate content useful for multiple audiences. + public var audience: [Role]? + + /// Describes how important this data is for operating the server. + /// A value of 1 means "most important" (effectively required), + /// while 0 means "least important" (entirely optional). + public var priority: Double? + + /// The moment the resource was last modified, as an ISO 8601 formatted string. + public var lastModified: String? + + public init( + audience: [Role]? = nil, + priority: Double? = nil, + lastModified: String? = nil + ) { + self.audience = audience + self.priority = priority + self.lastModified = lastModified + } +} diff --git a/Sources/MCP/Base/Error.swift b/Sources/MCP/Base/Error.swift index 0c461a46..c2301696 100644 --- a/Sources/MCP/Base/Error.swift +++ b/Sources/MCP/Base/Error.swift @@ -6,33 +6,231 @@ import Foundation @preconcurrency import SystemPackage #endif +// MARK: - Error Codes + +/// JSON-RPC and MCP error codes. +/// +/// Error codes are organized by source: +/// - Standard JSON-RPC 2.0 error codes (-32700 to -32600) +/// - MCP specification error codes (-32002, -32042) +/// - SDK-specific error codes (-32000, -32001, -32003) +public enum ErrorCode { + // MARK: Standard JSON-RPC 2.0 Errors + + /// Parse error: Invalid JSON was received by the server. + public static let parseError: Int = -32700 + + /// Invalid request: The JSON sent is not a valid Request object. + public static let invalidRequest: Int = -32600 + + /// Method not found: The method does not exist or is not available. + public static let methodNotFound: Int = -32601 + + /// Invalid params: Invalid method parameter(s). + public static let invalidParams: Int = -32602 + + /// Internal error: Internal JSON-RPC error. + public static let internalError: Int = -32603 + + // MARK: MCP Specification Errors + + /// Resource not found: The requested resource does not exist. + /// + /// Defined in MCP specification (resources.mdx). + public static let resourceNotFound: Int = -32002 + + /// URL elicitation required: The request requires URL-mode elicitation(s) to be completed. + /// + /// Defined in MCP specification (schema). + public static let urlElicitationRequired: Int = -32042 + + // MARK: SDK-Specific Errors + + /// Connection closed: The connection to the server was closed. + /// + /// Not defined in MCP spec. SDK-specific, matches TypeScript SDK. + public static let connectionClosed: Int = -32000 + + /// Request timeout: The server did not respond within the timeout period. + /// + /// Not defined in MCP spec. SDK-specific, matches TypeScript SDK. + public static let requestTimeout: Int = -32001 + + /// Transport error: An error occurred in the transport layer. + /// + /// Not defined in MCP spec. SDK-specific for Swift. + public static let transportError: Int = -32003 + + /// Request cancelled: The request was cancelled before completion. + /// + /// Not defined in MCP spec. SDK-specific for Swift. + public static let requestCancelled: Int = -32004 +} + +// MARK: - MCPError + /// A model context protocol error. public enum MCPError: Swift.Error, Sendable { - // Standard JSON-RPC 2.0 errors (-32700 to -32603) - case parseError(String?) // -32700 - case invalidRequest(String?) // -32600 - case methodNotFound(String?) // -32601 - case invalidParams(String?) // -32602 - case internalError(String?) // -32603 + // Standard JSON-RPC 2.0 errors + case parseError(String?) + case invalidRequest(String?) + case methodNotFound(String?) + case invalidParams(String?) + case internalError(String?) + + // MCP-specific errors + /// The requested resource was not found. + /// Defined in MCP specification (resources.mdx) with error code -32002. + case resourceNotFound(uri: String?) + + /// URL elicitation is required before the request can proceed. + /// Servers throw this from tool handlers when URL-mode elicitation(s) must be completed. + case urlElicitationRequired(message: String, elicitations: [ElicitRequestURLParams]) // Server errors (-32000 to -32099) case serverError(code: Int, message: String) + /// Server error with additional data payload. + case serverErrorWithData(code: Int, message: String, data: Value) - // Transport specific errors + // Transport and connection errors case connectionClosed case transportError(Swift.Error) + // Request timeout + /// Request timed out waiting for a response. + case requestTimeout(timeout: Duration, message: String?) + + // Request cancellation + /// Request was cancelled before completion. + case requestCancelled(reason: String?) + /// The JSON-RPC 2.0 error code public var code: Int { switch self { - case .parseError: return -32700 - case .invalidRequest: return -32600 - case .methodNotFound: return -32601 - case .invalidParams: return -32602 - case .internalError: return -32603 + case .parseError: return ErrorCode.parseError + case .invalidRequest: return ErrorCode.invalidRequest + case .methodNotFound: return ErrorCode.methodNotFound + case .invalidParams: return ErrorCode.invalidParams + case .internalError: return ErrorCode.internalError + case .resourceNotFound: return ErrorCode.resourceNotFound + case .urlElicitationRequired: return ErrorCode.urlElicitationRequired case .serverError(let code, _): return code - case .connectionClosed: return -32000 - case .transportError: return -32001 + case .serverErrorWithData(let code, _, _): return code + case .connectionClosed: return ErrorCode.connectionClosed + case .transportError: return ErrorCode.transportError + case .requestTimeout: return ErrorCode.requestTimeout + case .requestCancelled: return ErrorCode.requestCancelled + } + } + + /// Creates a URL elicitation required error. + /// + /// Use this when a tool handler needs the client to complete URL-mode elicitation(s) + /// before the request can proceed. + /// + /// Example: + /// ```swift + /// throw MCPError.urlElicitationRequired( + /// elicitations: [ + /// ElicitRequestURLParams( + /// message: "Please authorize access", + /// elicitationId: "auth-123", + /// url: "https://example.com/oauth" + /// ) + /// ] + /// ) + /// ``` + public static func urlElicitationRequired( + elicitations: [ElicitRequestURLParams], + message: String? = nil + ) -> MCPError { + let msg = message ?? "URL elicitation\(elicitations.count > 1 ? "s" : "") required" + return .urlElicitationRequired(message: msg, elicitations: elicitations) + } + + /// Attempts to extract elicitations from an error if it's a URL elicitation required error. + public var elicitations: [ElicitRequestURLParams]? { + if case .urlElicitationRequired(_, let elicitations) = self { + return elicitations + } + return nil + } + + /// The raw error message for wire format serialization. + /// + /// This returns the message suitable for JSON-RPC 2.0 error format, without + /// any additional prefixes or formatting that `errorDescription` might add. + /// Use this for serialization; use `errorDescription` for human-readable display. + public var message: String { + switch self { + case .parseError(let detail): + return detail ?? "Invalid JSON" + case .invalidRequest(let detail): + return detail ?? "Invalid Request" + case .methodNotFound(let detail): + return detail ?? "Method not found" + case .invalidParams(let detail): + return detail ?? "Invalid params" + case .internalError(let detail): + return detail ?? "Internal error" + case .resourceNotFound(let uri): + return uri.map { "Resource not found: \($0)" } ?? "Resource not found" + case .urlElicitationRequired(let message, _): + return message + case .serverError(_, let message): + return message + case .serverErrorWithData(_, let message, _): + return message + case .connectionClosed: + return "Connection closed" + case .transportError(let error): + return error.localizedDescription + case .requestTimeout(let timeout, let message): + return message ?? "Request timed out after \(timeout)" + case .requestCancelled(let reason): + return reason ?? "Request cancelled" + } + } + + /// The error data payload for wire format serialization. + /// + /// This returns the additional data to include in the JSON-RPC 2.0 error, + /// following MCP specification requirements for specific error types. + public var data: Value? { + switch self { + case .parseError, .invalidRequest, .methodNotFound, .invalidParams, .internalError: + // Standard JSON-RPC errors don't require data + return nil + case .resourceNotFound(let uri): + // Resource not found includes the URI in data per MCP spec + if let uri { + return .object(["uri": .string(uri)]) + } + return nil + case .urlElicitationRequired(_, let elicitations): + // URL elicitation required includes elicitations in data per MCP spec + do { + let encoded = try JSONEncoder().encode(ElicitationRequiredErrorData(elicitations: elicitations)) + return try JSONDecoder().decode(Value.self, from: encoded) + } catch { + return nil + } + case .serverError: + return nil + case .serverErrorWithData(_, _, let data): + return data + case .connectionClosed: + return nil + case .transportError(let error): + return .object(["error": .string(error.localizedDescription)]) + case .requestTimeout(let timeout, _): + let timeoutMs = Int(timeout.components.seconds * 1000 + timeout.components.attoseconds / 1_000_000_000_000_000) + return .object(["timeout": .int(timeoutMs)]) + case .requestCancelled(let reason): + if let reason { + return .object(["reason": .string(reason)]) + } + return nil } } @@ -66,12 +264,26 @@ extension MCPError: LocalizedError { return "Invalid params" + (detail.map { ": \($0)" } ?? "") case .internalError(let detail): return "Internal error" + (detail.map { ": \($0)" } ?? "") + case .resourceNotFound(let uri): + return "Resource not found" + (uri.map { ": \($0)" } ?? "") + case .urlElicitationRequired(let message, _): + return message case .serverError(_, let message): return "Server error: \(message)" + case .serverErrorWithData(_, let message, _): + return "Server error: \(message)" case .connectionClosed: return "Connection closed" case .transportError(let error): return "Transport error: \(error.localizedDescription)" + case .requestTimeout(let timeout, let message): + if let message { + return "Request timeout: \(message)" + } else { + return "Request timed out after \(timeout)" + } + case .requestCancelled(let reason): + return "Request cancelled" + (reason.map { ": \($0)" } ?? "") } } @@ -87,12 +299,20 @@ extension MCPError: LocalizedError { return "Invalid method parameter(s)" case .internalError: return "Internal JSON-RPC error" - case .serverError: + case .resourceNotFound: + return "The requested resource does not exist" + case .urlElicitationRequired: + return "The request requires URL-mode elicitation(s) to be completed first" + case .serverError, .serverErrorWithData: return "Server-defined error occurred" case .connectionClosed: return "The connection to the server was closed" case .transportError(let error): return (error as? LocalizedError)?.failureReason ?? error.localizedDescription + case .requestTimeout: + return "The server did not respond within the timeout period" + case .requestCancelled: + return "The request was cancelled before it could complete" } } @@ -106,8 +326,16 @@ extension MCPError: LocalizedError { return "Check the method name and ensure it is supported by the server" case .invalidParams: return "Verify the parameters match the method's expected parameters" + case .resourceNotFound: + return "Verify the resource URI is correct and the resource exists" + case .urlElicitationRequired: + return "Complete the required URL elicitation(s) and retry the request" case .connectionClosed: return "Try reconnecting to the server" + case .requestTimeout: + return "Try increasing the timeout or check if the server is responding" + case .requestCancelled: + return "Retry the request if needed" default: return nil } @@ -139,28 +367,11 @@ extension MCPError: Codable { public func encode(to encoder: Encoder) throws { var container = encoder.container(keyedBy: CodingKeys.self) try container.encode(code, forKey: .code) - try container.encode(errorDescription ?? "Unknown error", forKey: .message) + try container.encode(message, forKey: .message) - // Encode additional data if available - switch self { - case .parseError(let detail), - .invalidRequest(let detail), - .methodNotFound(let detail), - .invalidParams(let detail), - .internalError(let detail): - if let detail = detail { - try container.encode(["detail": detail], forKey: .data) - } - case .serverError(_, _): - // No additional data for server errors - break - case .connectionClosed: - break - case .transportError(let error): - try container.encode( - ["error": error.localizedDescription], - forKey: .data - ) + // Encode data if available + if let data = self.data { + try container.encode(data, forKey: .data) } } @@ -168,35 +379,61 @@ extension MCPError: Codable { let container = try decoder.container(keyedBy: CodingKeys.self) let code = try container.decode(Int.self, forKey: .code) let message = try container.decode(String.self, forKey: .message) - let data = try container.decodeIfPresent([String: Value].self, forKey: .data) - // Helper to extract detail from data, falling back to message if needed - let unwrapDetail: (String?) -> String? = { fallback in - guard let detailValue = data?["detail"] else { return fallback } - if case .string(let str) = detailValue { return str } - return fallback + // Try to decode data as a generic Value first + let dataValue = try container.decodeIfPresent(Value.self, forKey: .data) + + // Helper to check if message is the default for a given error type. + // If it's the default, we use nil as the detail; otherwise we use the custom message. + func customDetailOrNil(ifNotDefault defaultMessage: String) -> String? { + message == defaultMessage ? nil : message } switch code { - case -32700: - self = .parseError(unwrapDetail(message)) - case -32600: - self = .invalidRequest(unwrapDetail(message)) - case -32601: - self = .methodNotFound(unwrapDetail(message)) - case -32602: - self = .invalidParams(unwrapDetail(message)) - case -32603: - self = .internalError(unwrapDetail(nil)) - case -32000: + case ErrorCode.parseError: + self = .parseError(customDetailOrNil(ifNotDefault: "Invalid JSON")) + case ErrorCode.invalidRequest: + self = .invalidRequest(customDetailOrNil(ifNotDefault: "Invalid Request")) + case ErrorCode.methodNotFound: + self = .methodNotFound(customDetailOrNil(ifNotDefault: "Method not found")) + case ErrorCode.invalidParams: + self = .invalidParams(customDetailOrNil(ifNotDefault: "Invalid params")) + case ErrorCode.internalError: + self = .internalError(customDetailOrNil(ifNotDefault: "Internal error")) + case ErrorCode.resourceNotFound: + // Extract URI from data if present + var uri: String? = nil + if case .object(let dict) = dataValue, + case .string(let u) = dict["uri"] { + uri = u + } + self = .resourceNotFound(uri: uri) + case ErrorCode.urlElicitationRequired: + // Try to decode elicitations from data + if let errorData = try? container.decode(ElicitationRequiredErrorData.self, forKey: .data) { + self = .urlElicitationRequired(message: message, elicitations: errorData.elicitations) + } else { + // Fall back to server error if data doesn't match expected format + self = .serverError(code: code, message: message) + } + case ErrorCode.connectionClosed: self = .connectionClosed - case -32001: + case ErrorCode.requestTimeout: + // Extract timeout from data if present + var timeoutMs = 60000 // Default 60 seconds + if case .object(let dict) = dataValue, + let timeoutValue = dict["timeout"], + case .int(let t) = timeoutValue { + timeoutMs = t + } + self = .requestTimeout(timeout: .milliseconds(timeoutMs), message: message) + case ErrorCode.transportError: // Extract underlying error string if present - let underlyingErrorString = - data?["error"].flatMap { val -> String? in - if case .string(let str) = val { return str } - return nil - } ?? message + var underlyingErrorString = message + if case .object(let dict) = dataValue, + case .string(let str) = dict["error"] { + underlyingErrorString = str + } self = .transportError( NSError( domain: "org.jsonrpc.error", @@ -204,8 +441,123 @@ extension MCPError: Codable { userInfo: [NSLocalizedDescriptionKey: underlyingErrorString] ) ) + case ErrorCode.requestCancelled: + // Extract reason from data if present + var reason: String? = nil + if case .object(let dict) = dataValue, + case .string(let r) = dict["reason"] { + reason = r + } else if message != "Request cancelled" { + reason = message + } + self = .requestCancelled(reason: reason) default: - self = .serverError(code: code, message: message) + // Preserve data if present + if let dataValue { + self = .serverErrorWithData(code: code, message: message, data: dataValue) + } else { + self = .serverError(code: code, message: message) + } + } + } + + /// Reconstructs an MCPError from error code, message, and optional data. + /// + /// This is useful for clients receiving error responses and wanting to + /// work with typed error values. + /// + /// - Parameters: + /// - code: The JSON-RPC error code + /// - message: The error message + /// - data: Optional additional error data + /// - Returns: The appropriate MCPError type + public static func fromError(code: Int, message: String, data: Value? = nil) -> MCPError { + // Helper to check if message is the default for a given error type. + // If it's the default, we use nil as the detail; otherwise we use the custom message. + func customDetailOrNil(ifNotDefault defaultMessage: String) -> String? { + message == defaultMessage ? nil : message + } + + switch code { + case ErrorCode.parseError: + return .parseError(customDetailOrNil(ifNotDefault: "Invalid JSON")) + case ErrorCode.invalidRequest: + return .invalidRequest(customDetailOrNil(ifNotDefault: "Invalid Request")) + case ErrorCode.methodNotFound: + return .methodNotFound(customDetailOrNil(ifNotDefault: "Method not found")) + case ErrorCode.invalidParams: + return .invalidParams(customDetailOrNil(ifNotDefault: "Invalid params")) + case ErrorCode.internalError: + return .internalError(customDetailOrNil(ifNotDefault: "Internal error")) + case ErrorCode.resourceNotFound: + // Extract URI from data if present + var uri: String? = nil + if case .object(let dict) = data, + case .string(let u) = dict["uri"] { + uri = u + } + return .resourceNotFound(uri: uri) + case ErrorCode.urlElicitationRequired: + // Try to extract elicitations from data + if case .object(let dict) = data, + case .array(let elicitationsArray) = dict["elicitations"] { + // Decode each elicitation + var elicitations: [ElicitRequestURLParams] = [] + for item in elicitationsArray { + if case .object = item { + // Re-encode and decode to get proper type + if let jsonData = try? JSONEncoder().encode(item), + let params = try? JSONDecoder().decode(ElicitRequestURLParams.self, from: jsonData) { + elicitations.append(params) + } + } + } + if !elicitations.isEmpty { + return .urlElicitationRequired(message: message, elicitations: elicitations) + } + } + // Fall back to server error if we can't parse elicitations + return .serverError(code: code, message: message) + case ErrorCode.connectionClosed: + return .connectionClosed + case ErrorCode.requestTimeout: + // Extract timeout from data if present + var timeoutMs = 60000 // Default 60 seconds + if case .object(let dict) = data, + let timeoutValue = dict["timeout"], + case .int(let t) = timeoutValue { + timeoutMs = t + } + return .requestTimeout(timeout: .milliseconds(timeoutMs), message: message) + case ErrorCode.transportError: + // Extract underlying error string if present + var underlyingErrorString = message + if case .object(let dict) = data, + case .string(let str) = dict["error"] { + underlyingErrorString = str + } + return .transportError( + NSError( + domain: "org.jsonrpc.error", + code: code, + userInfo: [NSLocalizedDescriptionKey: underlyingErrorString] + ) + ) + case ErrorCode.requestCancelled: + // Extract reason from data if present + var reason: String? = nil + if case .object(let dict) = data, + case .string(let r) = dict["reason"] { + reason = r + } else if message != "Request cancelled" { + reason = message + } + return .requestCancelled(reason: reason) + default: + if let data { + return .serverErrorWithData(code: code, message: message, data: data) + } + return .serverError(code: code, message: message) } } } @@ -214,7 +566,36 @@ extension MCPError: Codable { extension MCPError: Equatable { public static func == (lhs: MCPError, rhs: MCPError) -> Bool { - lhs.code == rhs.code + switch (lhs, rhs) { + case (.parseError(let l), .parseError(let r)): + return l == r + case (.invalidRequest(let l), .invalidRequest(let r)): + return l == r + case (.methodNotFound(let l), .methodNotFound(let r)): + return l == r + case (.invalidParams(let l), .invalidParams(let r)): + return l == r + case (.internalError(let l), .internalError(let r)): + return l == r + case (.resourceNotFound(let l), .resourceNotFound(let r)): + return l == r + case (.urlElicitationRequired(let lMsg, let lElicit), .urlElicitationRequired(let rMsg, let rElicit)): + return lMsg == rMsg && lElicit == rElicit + case (.serverError(let lCode, let lMsg), .serverError(let rCode, let rMsg)): + return lCode == rCode && lMsg == rMsg + case (.serverErrorWithData(let lCode, let lMsg, let lData), .serverErrorWithData(let rCode, let rMsg, let rData)): + return lCode == rCode && lMsg == rMsg && lData == rData + case (.connectionClosed, .connectionClosed): + return true + case (.transportError(let l), .transportError(let r)): + return l.localizedDescription == r.localizedDescription + case (.requestTimeout(let lTimeout, let lMsg), .requestTimeout(let rTimeout, let rMsg)): + return lTimeout == rTimeout && lMsg == rMsg + case (.requestCancelled(let lReason), .requestCancelled(let rReason)): + return lReason == rReason + default: + return false + } } } @@ -234,12 +615,25 @@ extension MCPError: Hashable { hasher.combine(detail) case .internalError(let detail): hasher.combine(detail) + case .resourceNotFound(let uri): + hasher.combine(uri) + case .urlElicitationRequired(let message, let elicitations): + hasher.combine(message) + hasher.combine(elicitations) case .serverError(_, let message): hasher.combine(message) + case .serverErrorWithData(_, let message, let data): + hasher.combine(message) + hasher.combine(data) case .connectionClosed: break case .transportError(let error): hasher.combine(error.localizedDescription) + case .requestTimeout(let timeout, let message): + hasher.combine(timeout) + hasher.combine(message) + case .requestCancelled(let reason): + hasher.combine(reason) } } } diff --git a/Sources/MCP/Base/HTTPHeader.swift b/Sources/MCP/Base/HTTPHeader.swift new file mode 100644 index 00000000..cafbf614 --- /dev/null +++ b/Sources/MCP/Base/HTTPHeader.swift @@ -0,0 +1,45 @@ +/// HTTP header names used by the MCP protocol. +public enum HTTPHeader { + // MARK: - MCP Protocol Headers + + /// Session identifier for Streamable HTTP transport. + /// + /// Servers may return this header during initialization. Clients must + /// include it in all subsequent requests when present. + public static let sessionId = "mcp-session-id" + + /// Protocol version indicating the MCP version in use. + /// + /// Required for protocol versions >= 2025-06-18. Clients should send + /// the version negotiated during initialization. + public static let protocolVersion = "mcp-protocol-version" + + // MARK: - Standard HTTP Headers + + /// Content type of the request or response body. + public static let contentType = "content-type" + + /// Caching directives for the response. + public static let cacheControl = "cache-control" + + /// Connection management options. + public static let connection = "connection" + + /// Media types acceptable for the response. + public static let accept = "accept" + + /// Last event ID for SSE stream resumability. + public static let lastEventId = "last-event-id" + + /// Allowed HTTP methods for a resource (used in 405 responses). + public static let allow = "allow" + + /// Host header for the request target. + public static let host = "host" + + /// Origin header indicating where a request originated from. + public static let origin = "origin" + + /// Authorization header for bearer tokens and other auth schemes. + public static let authorization = "authorization" +} diff --git a/Sources/MCP/Base/Icon.swift b/Sources/MCP/Base/Icon.swift new file mode 100644 index 00000000..ee7ad651 --- /dev/null +++ b/Sources/MCP/Base/Icon.swift @@ -0,0 +1,51 @@ +import Foundation + +/// Icon metadata for representing visual icons for tools, resources, prompts, and implementations. +/// +/// Icons can be provided as HTTP/HTTPS URLs or data URIs (base64-encoded images). +/// +/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2025-11-25/ +public struct Icon: Hashable, Codable, Sendable { + /// URL or data URI for the icon. + /// + /// Can be an HTTP/HTTPS URL or a data URI (e.g., `data:image/png;base64,...`). + public let src: String + + /// Optional MIME type for the icon. + /// + /// Useful when the MIME type cannot be inferred from the `src` URL. + public let mimeType: String? + + /// Optional array of strings that specify sizes at which the icon can be used. + /// + /// Each string should be in WxH format (e.g., `"48x48"`, `"96x96"`) or `"any"` for + /// scalable formats like SVG. + /// + /// If not provided, the client should assume that the icon can be used at any size. + public let sizes: [String]? + + /// Optional specifier for the theme this icon is designed for. + /// + /// If not provided, the client should assume the icon can be used with any theme. + public let theme: Theme? + + /// The theme an icon is designed for. + public enum Theme: String, Hashable, Codable, Sendable { + /// Icon designed for use with a light background. + case light + /// Icon designed for use with a dark background. + case dark + } + + public init( + src: String, + mimeType: String? = nil, + sizes: [String]? = nil, + theme: Theme? = nil + ) { + self.src = src + self.mimeType = mimeType + self.sizes = sizes + self.theme = theme + } +} diff --git a/Sources/MCP/Base/Lifecycle.swift b/Sources/MCP/Base/Lifecycle.swift index 7d3e7119..e3131a6a 100644 --- a/Sources/MCP/Base/Lifecycle.swift +++ b/Sources/MCP/Base/Lifecycle.swift @@ -12,19 +12,23 @@ public enum Initialize: Method { public let protocolVersion: String public let capabilities: Client.Capabilities public let clientInfo: Client.Info + /// Request metadata including progress token. + public var _meta: RequestMeta? public init( protocolVersion: String = Version.latest, capabilities: Client.Capabilities, - clientInfo: Client.Info + clientInfo: Client.Info, + _meta: RequestMeta? = nil ) { self.protocolVersion = protocolVersion self.capabilities = capabilities self.clientInfo = clientInfo + self._meta = _meta } private enum CodingKeys: String, CodingKey { - case protocolVersion, capabilities, clientInfo + case protocolVersion, capabilities, clientInfo, _meta } public init(from decoder: Decoder) throws { @@ -38,14 +42,61 @@ public enum Initialize: Method { clientInfo = try container.decodeIfPresent(Client.Info.self, forKey: .clientInfo) ?? .init(name: "unknown", version: "0.0.0") + _meta = try container.decodeIfPresent(RequestMeta.self, forKey: ._meta) } } - public struct Result: Hashable, Codable, Sendable { + public struct Result: ResultWithExtraFields { + public typealias ResultCodingKeys = CodingKeys + public let protocolVersion: String public let capabilities: Server.Capabilities public let serverInfo: Server.Info public let instructions: String? + /// Reserved for clients and servers to attach additional metadata. + public var _meta: [String: Value]? + /// Additional fields not defined in the schema (for forward compatibility). + public var extraFields: [String: Value]? + + public init( + protocolVersion: String, + capabilities: Server.Capabilities, + serverInfo: Server.Info, + instructions: String? = nil, + _meta: [String: Value]? = nil, + extraFields: [String: Value]? = nil + ) { + self.protocolVersion = protocolVersion + self.capabilities = capabilities + self.serverInfo = serverInfo + self.instructions = instructions + self._meta = _meta + self.extraFields = extraFields + } + + public enum CodingKeys: String, CodingKey, CaseIterable { + case protocolVersion, capabilities, serverInfo, instructions, _meta + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + protocolVersion = try container.decode(String.self, forKey: .protocolVersion) + capabilities = try container.decode(Server.Capabilities.self, forKey: .capabilities) + serverInfo = try container.decode(Server.Info.self, forKey: .serverInfo) + instructions = try container.decodeIfPresent(String.self, forKey: .instructions) + _meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + extraFields = try Self.decodeExtraFields(from: decoder) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode(protocolVersion, forKey: .protocolVersion) + try container.encode(capabilities, forKey: .capabilities) + try container.encode(serverInfo, forKey: .serverInfo) + try container.encodeIfPresent(instructions, forKey: .instructions) + try container.encodeIfPresent(_meta, forKey: ._meta) + try encodeExtraFields(to: encoder) + } } } @@ -53,4 +104,32 @@ public enum Initialize: Method { /// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2024-11-05/basic/lifecycle/#initialization public struct InitializedNotification: Notification { public static let name: String = "notifications/initialized" + + public typealias Parameters = NotificationParams +} + +/// Notification sent when an operation is cancelled. +/// +/// This can be used by either client or server to indicate that an +/// ongoing operation should be terminated. +/// +/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2025-11-25/ +public struct CancelledNotification: Notification { + public static let name: String = "notifications/cancelled" + + public struct Parameters: Hashable, Codable, Sendable { + /// The ID of the request to cancel. + /// Optional in protocol version 2025-11-25 and later. + public var requestId: RequestId? + /// The reason for cancellation. + public var reason: String? + /// Reserved for additional metadata. + public var _meta: [String: Value]? + + public init(requestId: RequestId? = nil, reason: String? = nil, _meta: [String: Value]? = nil) { + self.requestId = requestId + self.reason = reason + self._meta = _meta + } + } } diff --git a/Sources/MCP/Base/Messages.swift b/Sources/MCP/Base/Messages.swift index b9058f7e..bc568023 100644 --- a/Sources/MCP/Base/Messages.swift +++ b/Sources/MCP/Base/Messages.swift @@ -11,6 +11,22 @@ public struct Empty: NotRequired, Hashable, Codable, Sendable { public init() {} } +/// Base notification parameters with optional metadata. +/// +/// Used by notifications that have no additional parameters beyond `_meta`. +public struct NotificationParams: NotRequired, Hashable, Codable, Sendable { + /// Reserved for additional metadata. + public var _meta: [String: Value]? + + public init() { + self._meta = nil + } + + public init(_meta: [String: Value]?) { + self._meta = _meta + } +} + extension Value: NotRequired { public init() { self = .null @@ -37,30 +53,37 @@ struct AnyMethod: Method, Sendable { } extension Method where Parameters == Empty { - public static func request(id: ID = .random) -> Request { + public static func request(id: RequestId = .random) -> Request { Request(id: id, method: name, params: Empty()) } } +extension Method where Parameters: NotRequired { + /// Create a request with default parameters. + public static func request(id: RequestId = .random) -> Request { + Request(id: id, method: name, params: Parameters()) + } +} + extension Method where Result == Empty { - public static func response(id: ID) -> Response { + public static func response(id: RequestId) -> Response { Response(id: id, result: Empty()) } } extension Method { /// Create a request with the given parameters. - public static func request(id: ID = .random, _ parameters: Self.Parameters) -> Request { + public static func request(id: RequestId = .random, _ parameters: Self.Parameters) -> Request { Request(id: id, method: name, params: parameters) } /// Create a response with the given result. - public static func response(id: ID, result: Self.Result) -> Response { + public static func response(id: RequestId, result: Self.Result) -> Response { Response(id: id, result: result) } /// Create a response with the given error. - public static func response(id: ID, error: MCPError) -> Response { + public static func response(id: RequestId, error: MCPError) -> Response { Response(id: id, error: error) } } @@ -70,13 +93,13 @@ extension Method { /// A request message. public struct Request: Hashable, Identifiable, Codable, Sendable { /// The request ID. - public let id: ID + public let id: RequestId /// The method name. public let method: String /// The request parameters. public let params: M.Parameters - init(id: ID = .random, method: String, params: M.Parameters) { + init(id: RequestId = .random, method: String, params: M.Parameters) { self.id = id self.method = method self.params = params @@ -151,21 +174,21 @@ extension AnyRequest { /// A box for request handlers that can be type-erased class RequestHandlerBox: @unchecked Sendable { - func callAsFunction(_ request: AnyRequest) async throws -> AnyResponse { + func callAsFunction(_ request: AnyRequest, context: Server.RequestHandlerContext) async throws -> AnyResponse { fatalError("Must override") } } /// A typed request handler that can be used to handle requests of a specific type final class TypedRequestHandler: RequestHandlerBox, @unchecked Sendable { - private let _handle: @Sendable (Request) async throws -> Response + private let _handle: @Sendable (Request, Server.RequestHandlerContext) async throws -> Response - init(_ handler: @escaping @Sendable (Request) async throws -> Response) { + init(_ handler: @escaping @Sendable (Request, Server.RequestHandlerContext) async throws -> Response) { self._handle = handler super.init() } - override func callAsFunction(_ request: AnyRequest) async throws -> AnyResponse { + override func callAsFunction(_ request: AnyRequest, context: Server.RequestHandlerContext) async throws -> AnyResponse { let encoder = JSONEncoder() let decoder = JSONDecoder() @@ -174,7 +197,7 @@ final class TypedRequestHandler: RequestHandlerBox, @unchecked Sendab let request = try decoder.decode(Request.self, from: data) // Handle with concrete type - let response = try await _handle(request) + let response = try await _handle(request, context) // Convert result to AnyMethod response switch response.result { @@ -193,16 +216,16 @@ final class TypedRequestHandler: RequestHandlerBox, @unchecked Sendab /// A response message. public struct Response: Hashable, Identifiable, Codable, Sendable { /// The response ID. - public let id: ID + public let id: RequestId /// The response result. public let result: Swift.Result - public init(id: ID, result: M.Result) { + public init(id: RequestId, result: M.Result) { self.id = id self.result = .success(result) } - public init(id: ID, error: MCPError) { + public init(id: RequestId, error: MCPError) { self.id = id self.result = .failure(error) } @@ -291,8 +314,17 @@ extension AnyNotification { } } +/// Protocol for type-erased notification messages. +/// +/// This protocol allows sending notification messages with parameters through +/// a type-erased interface. `Message` conforms to this protocol. +public protocol NotificationMessageProtocol: Sendable, Encodable { + /// The notification method name. + var method: String { get } +} + /// A message that can be used to send notifications. -public struct Message: Hashable, Codable, Sendable { +public struct Message: NotificationMessageProtocol, Hashable, Codable, Sendable { /// The method name. public let method: String /// The notification parameters. @@ -365,6 +397,13 @@ extension Notification where Parameters == Empty { } } +extension Notification where Parameters == NotificationParams { + /// Create a message with default parameters (no metadata). + public static func message() -> Message { + Message(method: name, params: NotificationParams()) + } +} + extension Notification { /// Create a message with the given parameters. public static func message(_ parameters: Parameters) -> Message { @@ -396,3 +435,45 @@ final class TypedNotificationHandler: NotificationHandlerBox, try await _handle(typedNotification) } } + +// MARK: - Client Request Handlers + +/// A box for client request handlers that can be type-erased +class ClientRequestHandlerBox: @unchecked Sendable { + func callAsFunction(_ request: AnyRequest, context: Client.RequestHandlerContext) async throws -> AnyResponse { + fatalError("Must override") + } +} + +/// A typed client request handler that can be used to handle requests of a specific type +final class TypedClientRequestHandler: ClientRequestHandlerBox, @unchecked Sendable { + private let _handle: @Sendable (M.Parameters, Client.RequestHandlerContext) async throws -> M.Result + + init(_ handler: @escaping @Sendable (M.Parameters, Client.RequestHandlerContext) async throws -> M.Result) { + self._handle = handler + super.init() + } + + override func callAsFunction(_ request: AnyRequest, context: Client.RequestHandlerContext) async throws -> AnyResponse { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + // Create a concrete request from the type-erased one + let data = try encoder.encode(request) + let typedRequest = try decoder.decode(Request.self, from: data) + + // Handle with concrete type + do { + let result = try await _handle(typedRequest.params, context) + + // Convert result to AnyMethod response + let resultData = try encoder.encode(result) + let resultValue = try decoder.decode(Value.self, from: resultData) + return Response(id: typedRequest.id, result: resultValue) + } catch let error as MCPError { + return Response(id: typedRequest.id, error: error) + } catch { + return Response(id: typedRequest.id, error: MCPError.internalError(error.localizedDescription)) + } + } +} diff --git a/Sources/MCP/Base/Progress.swift b/Sources/MCP/Base/Progress.swift new file mode 100644 index 00000000..41835731 --- /dev/null +++ b/Sources/MCP/Base/Progress.swift @@ -0,0 +1,298 @@ +/// Progress tracking for long-running operations. +/// +/// Clients can include a `progressToken` in request metadata (`_meta.progressToken`) +/// to receive progress notifications during operation execution. +/// +/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2025-11-25/utilities/progress/ + +/// Metadata that can be attached to any request via the `_meta` field. +/// +/// This is used primarily for progress tracking, but can also carry +/// arbitrary additional metadata. +public struct RequestMeta: Hashable, Codable, Sendable { + /// If specified, the caller is requesting out-of-band progress notifications + /// for this request. The value is an opaque token that will be attached to + /// any subsequent progress notifications. + public var progressToken: ProgressToken? + + /// Additional metadata fields. + public var additionalFields: [String: Value]? + + public init( + progressToken: ProgressToken? = nil, + additionalFields: [String: Value]? = nil + ) { + self.progressToken = progressToken + self.additionalFields = additionalFields + } + + private enum CodingKeys: String, CodingKey { + case progressToken + } + + private struct DynamicCodingKey: CodingKey { + var stringValue: String + var intValue: Int? { nil } + + init?(stringValue: String) { + self.stringValue = stringValue + } + + init?(intValue: Int) { + return nil + } + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + progressToken = try container.decodeIfPresent(ProgressToken.self, forKey: .progressToken) + + // Decode additional fields + let dynamicContainer = try decoder.container(keyedBy: DynamicCodingKey.self) + var extra: [String: Value] = [:] + for key in dynamicContainer.allKeys { + if key.stringValue == CodingKeys.progressToken.stringValue { + continue + } + if let value = try? dynamicContainer.decode(Value.self, forKey: key) { + extra[key.stringValue] = value + } + } + additionalFields = extra.isEmpty ? nil : extra + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encodeIfPresent(progressToken, forKey: .progressToken) + + // Encode additional fields + if let additional = additionalFields { + var dynamicContainer = encoder.container(keyedBy: DynamicCodingKey.self) + for (key, value) in additional { + if let codingKey = DynamicCodingKey(stringValue: key) { + try dynamicContainer.encode(value, forKey: codingKey) + } + } + } + } +} + +/// A token used to associate progress notifications with a specific request. +/// +/// Progress tokens can be either strings or integers. +public enum ProgressToken: Hashable, Sendable { + case string(String) + case integer(Int) +} + +extension ProgressToken: Codable { + public init(from decoder: Decoder) throws { + let container = try decoder.singleValueContainer() + if let intValue = try? container.decode(Int.self) { + self = .integer(intValue) + } else if let stringValue = try? container.decode(String.self) { + self = .string(stringValue) + } else { + throw DecodingError.typeMismatch( + ProgressToken.self, + DecodingError.Context( + codingPath: decoder.codingPath, + debugDescription: "Expected string or integer for ProgressToken" + ) + ) + } + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.singleValueContainer() + switch self { + case .string(let value): + try container.encode(value) + case .integer(let value): + try container.encode(value) + } + } +} + +extension ProgressToken: ExpressibleByStringLiteral { + public init(stringLiteral value: String) { + self = .string(value) + } +} + +extension ProgressToken: ExpressibleByIntegerLiteral { + public init(integerLiteral value: Int) { + self = .integer(value) + } +} + +/// Notification sent to report progress on a long-running operation. +/// +/// Servers send progress notifications to inform clients about the status +/// of operations that may take significant time to complete. +/// +/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2025-11-25/utilities/progress/ +public struct ProgressNotification: Notification { + public static let name: String = "notifications/progress" + + public struct Parameters: Hashable, Codable, Sendable { + /// The progress token from the original request's `_meta.progressToken`. + public let progressToken: ProgressToken + + /// The current progress value. Should increase monotonically. + public let progress: Double + + /// The total progress value, if known. + public let total: Double? + + /// An optional human-readable message describing the current progress. + public let message: String? + + /// Reserved for additional metadata. + public var _meta: [String: Value]? + + public init( + progressToken: ProgressToken, + progress: Double, + total: Double? = nil, + message: String? = nil, + _meta: [String: Value]? = nil + ) { + self.progressToken = progressToken + self.progress = progress + self.total = total + self.message = message + self._meta = _meta + } + } +} + +// MARK: - Progress Callback + +/// Progress information received during a long-running operation. +/// +/// This struct is passed to progress callbacks when using `send(_:onProgress:)`. +public struct Progress: Sendable, Hashable { + /// The current progress value. Increases monotonically. + public let value: Double + + /// The total progress value, if known. + public let total: Double? + + /// An optional human-readable message describing current progress. + public let message: String? + + public init(value: Double, total: Double? = nil, message: String? = nil) { + self.value = value + self.total = total + self.message = message + } +} + +/// A callback invoked when a progress notification is received. +/// +/// This is used by the client to receive progress updates for specific requests +/// when using `send(_:onProgress:)`. +public typealias ProgressCallback = @Sendable (Progress) async -> Void + +// MARK: - Progress Tracker (Server-Side) + +/// An actor for tracking and sending cumulative progress during a request. +/// +/// This follows the Python SDK's `ProgressContext` pattern, providing a convenient +/// way to track cumulative progress and send notifications without manually +/// tracking the current value. +/// +/// ## Example +/// +/// ```swift +/// server.withRequestHandler(CallTool.self) { request, context in +/// guard let token = request._meta?.progressToken else { +/// return CallTool.Result(content: [.text("Done")]) +/// } +/// +/// let tracker = ProgressTracker(token: token, total: 100, context: context) +/// +/// try await tracker.advance(by: 25, message: "Loading...") +/// try await tracker.advance(by: 50, message: "Processing...") +/// try await tracker.advance(by: 25, message: "Completing...") +/// +/// return CallTool.Result(content: [.text("Done")]) +/// } +/// ``` +public actor ProgressTracker { + /// The progress token from the request. + public let token: ProgressToken + + /// The total progress value, if known. + public let total: Double? + + /// The request handler context for sending notifications. + private let context: Server.RequestHandlerContext + + /// The current cumulative progress value. + public private(set) var current: Double = 0 + + /// Creates a new progress tracker. + /// + /// - Parameters: + /// - token: The progress token from the request's `_meta.progressToken` + /// - total: The total progress value, if known + /// - context: The request handler context for sending notifications + public init( + token: ProgressToken, + total: Double? = nil, + context: Server.RequestHandlerContext + ) { + self.token = token + self.total = total + self.context = context + } + + /// Advance progress by the given amount and send a notification. + /// + /// - Parameters: + /// - amount: The amount to add to the current progress + /// - message: An optional human-readable message describing current progress + public func advance(by amount: Double, message: String? = nil) async throws { + current += amount + try await context.sendProgress( + token: token, + progress: current, + total: total, + message: message + ) + } + + /// Set progress to a specific value and send a notification. + /// + /// Use this when you want to set progress to an absolute value rather than + /// incrementing. The progress value should still increase monotonically. + /// + /// - Parameters: + /// - value: The new progress value + /// - message: An optional human-readable message describing current progress + public func set(to value: Double, message: String? = nil) async throws { + current = value + try await context.sendProgress( + token: token, + progress: current, + total: total, + message: message + ) + } + + /// Send a progress notification without changing the current value. + /// + /// Use this to update the message without changing the progress value. + /// + /// - Parameter message: A human-readable message describing current progress + public func update(message: String) async throws { + try await context.sendProgress( + token: token, + progress: current, + total: total, + message: message + ) + } +} diff --git a/Sources/MCP/Base/ID.swift b/Sources/MCP/Base/RequestId.swift similarity index 79% rename from Sources/MCP/Base/ID.swift rename to Sources/MCP/Base/RequestId.swift index 271b6150..cf4c1b08 100644 --- a/Sources/MCP/Base/ID.swift +++ b/Sources/MCP/Base/RequestId.swift @@ -1,7 +1,7 @@ import struct Foundation.UUID /// A unique identifier for a request. -public enum ID: Hashable, Sendable { +public enum RequestId: Hashable, Sendable { /// A string ID. case string(String) @@ -9,14 +9,18 @@ public enum ID: Hashable, Sendable { case number(Int) /// Generates a random string ID. - public static var random: ID { + public static var random: RequestId { return .string(UUID().uuidString) } } +/// Backwards compatibility alias for `RequestId`. +@available(*, deprecated, renamed: "RequestId") +public typealias ID = RequestId + // MARK: - ExpressibleByStringLiteral -extension ID: ExpressibleByStringLiteral { +extension RequestId: ExpressibleByStringLiteral { public init(stringLiteral value: String) { self = .string(value) } @@ -24,7 +28,7 @@ extension ID: ExpressibleByStringLiteral { // MARK: - ExpressibleByIntegerLiteral -extension ID: ExpressibleByIntegerLiteral { +extension RequestId: ExpressibleByIntegerLiteral { public init(integerLiteral value: Int) { self = .number(value) } @@ -32,7 +36,7 @@ extension ID: ExpressibleByIntegerLiteral { // MARK: - CustomStringConvertible -extension ID: CustomStringConvertible { +extension RequestId: CustomStringConvertible { public var description: String { switch self { case .string(let str): return str @@ -43,7 +47,7 @@ extension ID: CustomStringConvertible { // MARK: - Codable -extension ID: Codable { +extension RequestId: Codable { public init(from decoder: Decoder) throws { let container = try decoder.singleValueContainer() if let string = try? container.decode(String.self) { diff --git a/Sources/MCP/Base/Transport.swift b/Sources/MCP/Base/Transport.swift index 4e4350fd..2c9057a2 100644 --- a/Sources/MCP/Base/Transport.swift +++ b/Sources/MCP/Base/Transport.swift @@ -6,6 +6,15 @@ import struct Foundation.Data public protocol Transport: Actor { var logger: Logger { get } + /// The session identifier for this transport connection. + /// + /// For HTTP transports supporting multiple concurrent clients, each client + /// session has a unique identifier. This enables per-session features like + /// independent log levels for each client. + /// + /// For simple transports (stdio, single-connection), this returns `nil`. + var sessionId: String? { get } + /// Establishes connection with the transport func connect() async throws @@ -15,6 +24,35 @@ public protocol Transport: Actor { /// Sends data func send(_ data: Data) async throws + /// Sends data with an optional related request ID for response routing. + /// + /// For transports that support multiplexing (like HTTP), the `relatedRequestId` + /// parameter enables routing responses back to the correct client connection. + /// + /// For simple transports (stdio, single-connection), this can be ignored. + /// + /// - Parameters: + /// - data: The data to send + /// - relatedRequestId: The ID of the request this message relates to (for response routing) + func send(_ data: Data, relatedRequestId: RequestId?) async throws + /// Receives data in an async sequence func receive() -> AsyncThrowingStream } + +// MARK: - Default Implementation + +extension Transport { + /// Default implementation returns `nil` for simple transports. + /// + /// HTTP transports override this to return their session identifier. + public var sessionId: String? { nil } + + /// Default implementation that ignores the request ID. + /// + /// Simple transports (stdio, single-connection) don't need request ID routing, + /// so they can use this default implementation that delegates to `send(_:)`. + public func send(_ data: Data, relatedRequestId: RequestId?) async throws { + try await send(data) + } +} diff --git a/Sources/MCP/Base/Transports/HTTPClientTransport.swift b/Sources/MCP/Base/Transports/HTTPClientTransport.swift index 11a4455e..dde41daf 100644 --- a/Sources/MCP/Base/Transports/HTTPClientTransport.swift +++ b/Sources/MCP/Base/Transports/HTTPClientTransport.swift @@ -9,6 +9,44 @@ import Logging import FoundationNetworking #endif +/// Configuration options for reconnection behavior of the HTTPClientTransport. +/// +/// These options control how the transport handles SSE stream disconnections +/// and reconnection attempts. +public struct HTTPReconnectionOptions: Sendable { + /// Initial delay between reconnection attempts in seconds. + /// Default is 1.0 second. + public var initialReconnectionDelay: TimeInterval + + /// Maximum delay between reconnection attempts in seconds. + /// Default is 30.0 seconds. + public var maxReconnectionDelay: TimeInterval + + /// Factor by which the reconnection delay increases after each attempt. + /// Default is 1.5. + public var reconnectionDelayGrowFactor: Double + + /// Maximum number of reconnection attempts before giving up. + /// Default is 2. + public var maxRetries: Int + + /// Creates reconnection options with default values. + public init( + initialReconnectionDelay: TimeInterval = 1.0, + maxReconnectionDelay: TimeInterval = 30.0, + reconnectionDelayGrowFactor: Double = 1.5, + maxRetries: Int = 2 + ) { + self.initialReconnectionDelay = initialReconnectionDelay + self.maxReconnectionDelay = maxReconnectionDelay + self.reconnectionDelayGrowFactor = reconnectionDelayGrowFactor + self.maxRetries = maxRetries + } + + /// Default reconnection options. + public static let `default` = HTTPReconnectionOptions() +} + /// An implementation of the MCP Streamable HTTP transport protocol for clients. /// /// This transport implements the [Streamable HTTP transport](https://modelcontextprotocol.io/specification/2025-03-26/basic/transports#streamable-http) @@ -56,6 +94,9 @@ public actor HTTPClientTransport: Transport { /// The session ID assigned by the server, used for maintaining state across requests public private(set) var sessionID: String? + + /// The negotiated protocol version, set after initialization + public private(set) var protocolVersion: String? private let streaming: Bool private var streamingTask: Task? @@ -65,6 +106,9 @@ public actor HTTPClientTransport: Transport { /// Maximum time to wait for a session ID before proceeding with SSE connection public let sseInitializationTimeout: TimeInterval + /// Configuration for reconnection behavior + public nonisolated let reconnectionOptions: HTTPReconnectionOptions + /// Closure to modify requests before they are sent private let requestModifier: (URLRequest) -> URLRequest @@ -72,8 +116,30 @@ public actor HTTPClientTransport: Transport { private let messageStream: AsyncThrowingStream private let messageContinuation: AsyncThrowingStream.Continuation - private var initialSessionIDSignalTask: Task? - private var initialSessionIDContinuation: CheckedContinuation? + /// Stream for signaling when session ID is set + private var sessionIDSignalStream: AsyncStream? + private var sessionIDSignalContinuation: AsyncStream.Continuation? + + // MARK: - Reconnection State + + /// The last event ID received from the server, used for resumability + private var lastEventId: String? + + /// Server-provided retry delay in seconds (from SSE retry: field) + private var serverRetryDelay: TimeInterval? + + /// Current reconnection attempt count + private var reconnectionAttempt: Int = 0 + + /// Callback invoked when a new resumption token (event ID) is received + public var onResumptionToken: ((String) -> Void)? + + /// Sets the callback invoked when a new resumption token (event ID) is received. + /// + /// - Parameter callback: The callback to invoke with the event ID + public func setOnResumptionToken(_ callback: ((String) -> Void)?) { + onResumptionToken = callback + } /// Creates a new HTTP transport client with the specified endpoint /// @@ -82,6 +148,7 @@ public actor HTTPClientTransport: Transport { /// - configuration: URLSession configuration to use for HTTP requests /// - streaming: Whether to enable SSE streaming mode (default: true) /// - sseInitializationTimeout: Maximum time to wait for session ID before proceeding with SSE (default: 10 seconds) + /// - reconnectionOptions: Configuration for reconnection behavior (default: .default) /// - requestModifier: Optional closure to customize requests before they are sent (default: no modification) /// - logger: Optional logger instance for transport events public init( @@ -89,6 +156,7 @@ public actor HTTPClientTransport: Transport { configuration: URLSessionConfiguration = .default, streaming: Bool = true, sseInitializationTimeout: TimeInterval = 10, + reconnectionOptions: HTTPReconnectionOptions = .default, requestModifier: @escaping (URLRequest) -> URLRequest = { $0 }, logger: Logger? = nil ) { @@ -97,6 +165,7 @@ public actor HTTPClientTransport: Transport { session: URLSession(configuration: configuration), streaming: streaming, sseInitializationTimeout: sseInitializationTimeout, + reconnectionOptions: reconnectionOptions, requestModifier: requestModifier, logger: logger ) @@ -107,6 +176,7 @@ public actor HTTPClientTransport: Transport { session: URLSession, streaming: Bool = false, sseInitializationTimeout: TimeInterval = 10, + reconnectionOptions: HTTPReconnectionOptions = .default, requestModifier: @escaping (URLRequest) -> URLRequest = { $0 }, logger: Logger? = nil ) { @@ -114,6 +184,7 @@ public actor HTTPClientTransport: Transport { self.session = session self.streaming = streaming self.sseInitializationTimeout = sseInitializationTimeout + self.reconnectionOptions = reconnectionOptions self.requestModifier = requestModifier // Create message stream @@ -129,21 +200,19 @@ public actor HTTPClientTransport: Transport { ) } - // Setup the initial session ID signal - private func setupInitialSessionIDSignal() { - self.initialSessionIDSignalTask = Task { - await withCheckedContinuation { continuation in - self.initialSessionIDContinuation = continuation - // This task will suspend here until continuation.resume() is called - } - } + // Setup the initial session ID signal stream + private func setUpInitialSessionIDSignal() { + let (stream, continuation) = AsyncStream.makeStream() + self.sessionIDSignalStream = stream + self.sessionIDSignalContinuation = continuation } // Trigger the initial session ID signal when a session ID is established private func triggerInitialSessionIDSignal() { - if let continuation = self.initialSessionIDContinuation { - continuation.resume() - self.initialSessionIDContinuation = nil // Consume the continuation + if let continuation = self.sessionIDSignalContinuation { + continuation.yield(()) + continuation.finish() + self.sessionIDSignalContinuation = nil // Consume the continuation logger.trace("Initial session ID signal triggered for SSE task.") } } @@ -158,7 +227,7 @@ public actor HTTPClientTransport: Transport { isConnected = true // Setup initial session ID signal - setupInitialSessionIDSignal() + setUpInitialSessionIDSignal() if streaming { // Start listening to server events @@ -186,16 +255,76 @@ public actor HTTPClientTransport: Transport { // Clean up message stream messageContinuation.finish() - // Cancel the initial session ID signal task if active - initialSessionIDSignalTask?.cancel() - initialSessionIDSignalTask = nil - // Resume the continuation if it's still pending to avoid leaks - initialSessionIDContinuation?.resume() - initialSessionIDContinuation = nil + // Finish the session ID signal stream if it's still pending + sessionIDSignalContinuation?.finish() + sessionIDSignalContinuation = nil + sessionIDSignalStream = nil logger.debug("HTTP clienttransport disconnected") } + /// Terminates the current session by sending a DELETE request to the server. + /// + /// Clients that no longer need a particular session (e.g., because the user is + /// leaving the client application) SHOULD send an HTTP DELETE to the MCP endpoint + /// with the `Mcp-Session-Id` header to explicitly terminate the session. + /// + /// This allows the server to clean up any resources associated with the session. + /// + /// - Note: The server MAY respond with HTTP 405 Method Not Allowed, indicating + /// that the server does not allow clients to terminate sessions. This is handled + /// gracefully and does not throw an error. + /// + /// - Throws: MCPError if the DELETE request fails for reasons other than 405. + public func terminateSession() async throws { + guard let sessionID else { + // No session to terminate + return + } + + var request = URLRequest(url: endpoint) + request.httpMethod = "DELETE" + + // Add session ID header + request.addValue(sessionID, forHTTPHeaderField: HTTPHeader.sessionId) + + // Add protocol version if available + if let protocolVersion { + request.addValue(protocolVersion, forHTTPHeaderField: HTTPHeader.protocolVersion) + } + + // Apply request modifier (for auth headers, etc.) + request = requestModifier(request) + + logger.debug("Terminating session", metadata: ["sessionID": "\(sessionID)"]) + + let (_, response) = try await session.data(for: request) + + guard let httpResponse = response as? HTTPURLResponse else { + throw MCPError.internalError("Invalid HTTP response") + } + + switch httpResponse.statusCode { + case 200, 204: + // Success - session terminated + self.sessionID = nil + logger.debug("Session terminated successfully") + + case 405: + // Server does not support session termination - this is OK per spec + logger.debug("Server does not support session termination (405)") + + case 404: + // Session already expired or doesn't exist + self.sessionID = nil + logger.debug("Session not found (already expired)") + + default: + throw MCPError.internalError( + "Failed to terminate session: HTTP \(httpResponse.statusCode)") + } + } + /// Sends data through an HTTP POST request /// /// This sends a JSON-RPC message to the server via HTTP POST and processes @@ -206,22 +335,41 @@ public actor HTTPClientTransport: Transport { /// - Processing different response types (JSON vs SSE) /// - Handling HTTP error codes according to the specification /// + /// ## Implementation Note + /// + /// This method signature differs from TypeScript and Python SDKs which receive + /// typed `JSONRPCMessage` objects instead of raw `Data`. Swift parses the JSON + /// internally to determine message type (request vs notification) for proper + /// content-type validation per the MCP spec. + /// + /// This design avoids breaking changes to the `Transport` protocol. A future + /// revision could consider changing the protocol to receive typed messages + /// for better alignment with other SDKs. + /// /// - Parameter data: The JSON-RPC message to send /// - Throws: MCPError for transport failures or server errors public func send(_ data: Data) async throws { + // Determine if message is a request (has both "method" and "id") + // Per MCP spec, only requests require content-type validation + let expectsContentType = isRequest(data) guard isConnected else { throw MCPError.internalError("Transport not connected") } var request = URLRequest(url: endpoint) request.httpMethod = "POST" - request.addValue("application/json, text/event-stream", forHTTPHeaderField: "Accept") - request.addValue("application/json", forHTTPHeaderField: "Content-Type") + request.addValue("application/json, text/event-stream", forHTTPHeaderField: HTTPHeader.accept) + request.addValue("application/json", forHTTPHeaderField: HTTPHeader.contentType) request.httpBody = data // Add session ID if available - if let sessionID = sessionID { - request.addValue(sessionID, forHTTPHeaderField: "Mcp-Session-Id") + if let sessionID { + request.addValue(sessionID, forHTTPHeaderField: HTTPHeader.sessionId) + } + + // Add protocol version if available (required after initialization) + if let protocolVersion { + request.addValue(protocolVersion, forHTTPHeaderField: HTTPHeader.protocolVersion) } // Apply request modifier @@ -230,26 +378,98 @@ public actor HTTPClientTransport: Transport { #if os(Linux) // Linux implementation using data(for:) instead of bytes(for:) let (responseData, response) = try await session.data(for: request) - try await processResponse(response: response, data: responseData) + try await processResponse(response: response, data: responseData, expectsContentType: expectsContentType) #else // macOS and other platforms with bytes(for:) support let (responseStream, response) = try await session.bytes(for: request) - try await processResponse(response: response, stream: responseStream) + try await processResponse(response: response, stream: responseStream, expectsContentType: expectsContentType) #endif } + /// Checks if the given data represents a JSON-RPC request. + /// + /// Per JSON-RPC 2.0 spec, a request has both "method" and "id" fields. + /// Notifications have "method" but no "id". Responses have "id" but no "method". + /// + /// This is used to determine content-type validation behavior per MCP spec: + /// - Requests: Server MUST return `application/json` or `text/event-stream` + /// - Notifications: Server MUST return 202 Accepted with no body + private func isRequest(_ data: Data) -> Bool { + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { + return false + } + // A request has both "method" and "id" fields + return json["method"] != nil && json["id"] != nil + } + + /// Result of processing a JSON-RPC message for response detection and optional ID remapping. + private struct ProcessedMessage { + /// Whether the message is a JSON-RPC response (success or error) + let isResponse: Bool + /// The message data, potentially with ID remapped + let data: Data + } + + /// Processes a JSON-RPC message, detecting if it's a response and optionally remapping its ID. + /// + /// Per JSON-RPC 2.0 spec: + /// - A successful response has "id" and "result" fields, but no "method" + /// - An error response has "id" and "error" fields, but no "method" + /// + /// This combines response detection with ID remapping for efficiency (single parse). + /// ID remapping is used during stream resumption to ensure responses match the + /// original pending request, aligning with Python SDK behavior. + /// + /// Note: This implementation handles both success AND error responses, which aligns + /// with Python SDK but is more complete than TypeScript SDK. TypeScript's streamableHttp.ts + /// only checks `isJSONRPCResultResponse` (success only), missing error response handling. + /// TODO: File PR to fix TypeScript SDK - streamableHttp.ts line 364 should also handle + /// `isJSONRPCErrorResponse` for both `receivedResponse` flag and ID remapping. + /// + /// - Parameters: + /// - data: The raw JSON-RPC message data + /// - originalRequestId: Optional ID to remap response IDs to (for stream resumption) + /// - Returns: ProcessedMessage with isResponse flag and potentially remapped data + private func processJSONRPCMessage(_ data: Data, originalRequestId: RequestId?) -> ProcessedMessage { + guard var json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { + return ProcessedMessage(isResponse: false, data: data) + } + + // Check if it's a response (has id + result/error, no method) + let hasId = json["id"] != nil + let hasResult = json["result"] != nil + let hasError = json["error"] != nil + let hasMethod = json["method"] != nil + let isResponse = hasId && (hasResult || hasError) && !hasMethod + + // If it's a response and we have an original request ID, remap the ID + if isResponse, let originalId = originalRequestId { + switch originalId { + case .string(let s): json["id"] = s + case .number(let n): json["id"] = n + } + + // Re-encode with remapped ID + if let remappedData = try? JSONSerialization.data(withJSONObject: json) { + return ProcessedMessage(isResponse: true, data: remappedData) + } + } + + return ProcessedMessage(isResponse: isResponse, data: data) + } + #if os(Linux) // Process response with data payload (Linux) - private func processResponse(response: URLResponse, data: Data) async throws { + private func processResponse(response: URLResponse, data: Data, expectsContentType: Bool) async throws { guard let httpResponse = response as? HTTPURLResponse else { throw MCPError.internalError("Invalid HTTP response") } // Process the response based on content type and status code - let contentType = httpResponse.value(forHTTPHeaderField: "Content-Type") ?? "" + let contentType = httpResponse.value(forHTTPHeaderField: HTTPHeader.contentType) ?? "" // Extract session ID if present - if let newSessionID = httpResponse.value(forHTTPHeaderField: "Mcp-Session-Id") { + if let newSessionID = httpResponse.value(forHTTPHeaderField: HTTPHeader.sessionId) { let wasSessionIDNil = (self.sessionID == nil) self.sessionID = newSessionID if wasSessionIDNil { @@ -262,20 +482,22 @@ public actor HTTPClientTransport: Transport { try processHTTPResponse(httpResponse, contentType: contentType) guard case 200..<300 = httpResponse.statusCode else { return } - // For JSON responses, yield the data + // Process response based on content type if contentType.contains("text/event-stream") { logger.warning("SSE responses aren't fully supported on Linux") messageContinuation.yield(data) } else if contentType.contains("application/json") { logger.trace("Received JSON response", metadata: ["size": "\(data.count)"]) messageContinuation.yield(data) - } else { - logger.warning("Unexpected content type: \(contentType)") + } else if expectsContentType && !data.isEmpty { + // Per MCP spec: requests MUST receive application/json or text/event-stream + // Notifications expect 202 Accepted with no body, so unexpected content-type is ignored + throw MCPError.internalError("Unexpected content type: \(contentType)") } } #else // Process response with byte stream (macOS, iOS, etc.) - private func processResponse(response: URLResponse, stream: URLSession.AsyncBytes) + private func processResponse(response: URLResponse, stream: URLSession.AsyncBytes, expectsContentType: Bool) async throws { guard let httpResponse = response as? HTTPURLResponse else { @@ -283,10 +505,10 @@ public actor HTTPClientTransport: Transport { } // Process the response based on content type and status code - let contentType = httpResponse.value(forHTTPHeaderField: "Content-Type") ?? "" + let contentType = httpResponse.value(forHTTPHeaderField: HTTPHeader.contentType) ?? "" // Extract session ID if present - if let newSessionID = httpResponse.value(forHTTPHeaderField: "Mcp-Session-Id") { + if let newSessionID = httpResponse.value(forHTTPHeaderField: HTTPHeader.sessionId) { let wasSessionIDNil = (self.sessionID == nil) self.sessionID = newSessionID if wasSessionIDNil { @@ -300,9 +522,10 @@ public actor HTTPClientTransport: Transport { guard case 200..<300 = httpResponse.statusCode else { return } if contentType.contains("text/event-stream") { - // For SSE, processing happens via the stream + // For SSE response from POST, isReconnectable is false initially + // but can become reconnectable after receiving a priming event logger.trace("Received SSE response, processing in streaming task") - try await self.processSSE(stream) + try await self.processSSE(stream, isReconnectable: false) } else if contentType.contains("application/json") { // For JSON responses, collect and deliver the data var buffer = Data() @@ -312,12 +535,25 @@ public actor HTTPClientTransport: Transport { logger.trace("Received JSON response", metadata: ["size": "\(buffer.count)"]) messageContinuation.yield(buffer) } else { - logger.warning("Unexpected content type: \(contentType)") + // Collect data to check if response has content + var buffer = Data() + for try await byte in stream { + buffer.append(byte) + } + // Per MCP spec: requests MUST receive application/json or text/event-stream + // Notifications expect 202 Accepted with no body, so unexpected content-type is ignored + if expectsContentType && !buffer.isEmpty { + throw MCPError.internalError("Unexpected content type: \(contentType)") + } } } #endif // Common HTTP response handling for all platforms + // + // Note: The MCP spec recommends auto-detecting legacy SSE servers by falling back + // to GET on 400/404/405 errors. We don't implement this, consistent with the + // TypeScript and Python SDKs which provide separate transports instead. private func processHTTPResponse(_ response: HTTPURLResponse, contentType: String) throws { // Handle status codes according to HTTP semantics switch response.statusCode { @@ -336,6 +572,10 @@ public actor HTTPClientTransport: Transport { case 404: // If we get a 404 with a session ID, it means our session is invalid + // TODO: Consider Python's approach - send JSON-RPC error through stream + // with request ID (code -32600) before throwing. This gives pending requests + // proper error responses. Options: (1) catch in send() and yield error, + // (2) use RequestContext pattern like Python. Both are spec-compliant. if sessionID != nil { logger.warning("Session has expired") sessionID = nil @@ -381,6 +621,18 @@ public actor HTTPClientTransport: Transport { return messageStream } + /// Sets the protocol version to include in request headers. + /// + /// This should be called after initialization when the protocol version is negotiated. + /// HTTP transports must include the `Mcp-Protocol-Version` header in all requests + /// after initialization. + /// + /// - Parameter version: The negotiated protocol version (e.g., "2024-11-05") + public func setProtocolVersion(_ version: String) { + self.protocolVersion = version + logger.debug("Protocol version set", metadata: ["version": "\(version)"]) + } + // MARK: - SSE /// Starts listening for server events using SSE @@ -405,29 +657,24 @@ public actor HTTPClientTransport: Transport { guard isConnected else { return } // Wait for the initial session ID signal, but only if sessionID isn't already set - if self.sessionID == nil, let signalTask = self.initialSessionIDSignalTask { + if self.sessionID == nil, let signalStream = self.sessionIDSignalStream { logger.trace("SSE streaming task waiting for initial sessionID signal...") - // Race the signalTask against a timeout - let timeoutTask = Task { - try? await Task.sleep(for: .seconds(self.sseInitializationTimeout)) - return false - } - - let signalCompletionTask = Task { - await signalTask.value - return true // Indicates signal received - } - - // Use TaskGroup to race the two tasks + // Race the stream against a timeout using TaskGroup var signalReceived = false do { signalReceived = try await withThrowingTaskGroup(of: Bool.self) { group in group.addTask { - await signalCompletionTask.value + // Wait for signal from stream + for await _ in signalStream { + return true + } + return false // Stream finished without yielding } group.addTask { - await timeoutTask.value + // Timeout task + try await Task.sleep(for: .seconds(self.sseInitializationTimeout)) + return false } // Take the first result and cancel the other task @@ -441,9 +688,6 @@ public actor HTTPClientTransport: Transport { logger.error("Error while waiting for session ID signal: \(error)") } - // Clean up tasks - timeoutTask.cancel() - if signalReceived { logger.trace("SSE streaming task proceeding after initial sessionID signal.") } else { @@ -461,39 +705,102 @@ public actor HTTPClientTransport: Transport { ) } - // Retry loop for connection drops + // Retry loop for connection drops with exponential backoff while isConnected && !Task.isCancelled { do { try await connectToEventStream() + // Reset attempt counter on successful connection + reconnectionAttempt = 0 } catch { if !Task.isCancelled { logger.error("SSE connection error: \(error)") - // Wait before retrying - try? await Task.sleep(for: .seconds(1)) + + // Check if we've exceeded max retries + if reconnectionAttempt >= reconnectionOptions.maxRetries { + logger.error( + "Maximum reconnection attempts exceeded", + metadata: ["maxRetries": "\(reconnectionOptions.maxRetries)"] + ) + break + } + + // Calculate delay with exponential backoff + let delay = getNextReconnectionDelay() + reconnectionAttempt += 1 + + logger.debug( + "Scheduling reconnection", + metadata: [ + "attempt": "\(reconnectionAttempt)", + "delay": "\(delay)s", + ] + ) + + try? await Task.sleep(for: .seconds(delay)) } } } #endif } + /// Calculates the next reconnection delay using exponential backoff + /// + /// Uses server-provided retry value if available, otherwise falls back + /// to exponential backoff based on current attempt count. + /// + /// - Returns: Time to wait in seconds before next reconnection attempt + private func getNextReconnectionDelay() -> TimeInterval { + // Use server-provided retry value if available + if let serverDelay = serverRetryDelay { + return serverDelay + } + + // Fall back to exponential backoff + let initialDelay = reconnectionOptions.initialReconnectionDelay + let growFactor = reconnectionOptions.reconnectionDelayGrowFactor + let maxDelay = reconnectionOptions.maxReconnectionDelay + + // Calculate delay with exponential growth, capped at maximum + let delay = initialDelay * pow(growFactor, Double(reconnectionAttempt)) + return min(delay, maxDelay) + } + #if !os(Linux) /// Establishes an SSE connection to the server /// /// This initiates a GET request to the server endpoint with appropriate /// headers to establish an SSE stream according to the MCP specification. /// + /// - Parameters: + /// - resumptionToken: Optional event ID to resume from (sent as Last-Event-ID header) + /// - originalRequestId: Optional request ID to remap response IDs to (for stream resumption) /// - Throws: MCPError for connection failures or server errors - private func connectToEventStream() async throws { + private func connectToEventStream( + resumptionToken: String? = nil, + originalRequestId: RequestId? = nil + ) async throws { guard isConnected else { return } var request = URLRequest(url: endpoint) request.httpMethod = "GET" - request.addValue("text/event-stream", forHTTPHeaderField: "Accept") - request.addValue("no-cache", forHTTPHeaderField: "Cache-Control") + request.addValue("text/event-stream", forHTTPHeaderField: HTTPHeader.accept) + request.addValue("no-cache", forHTTPHeaderField: HTTPHeader.cacheControl) // Add session ID if available - if let sessionID = sessionID { - request.addValue(sessionID, forHTTPHeaderField: "Mcp-Session-Id") + if let sessionID { + request.addValue(sessionID, forHTTPHeaderField: HTTPHeader.sessionId) + } + + // Add protocol version if available + if let protocolVersion { + request.addValue(protocolVersion, forHTTPHeaderField: HTTPHeader.protocolVersion) + } + + // Add Last-Event-ID for resumability (use provided token or stored lastEventId) + let eventIdToSend = resumptionToken ?? lastEventId + if let eventId = eventIdToSend { + request.addValue(eventId, forHTTPHeaderField: HTTPHeader.lastEventId) + logger.debug("Resuming SSE stream", metadata: ["lastEventId": "\(eventId)"]) } // Apply request modifier @@ -501,6 +808,9 @@ public actor HTTPClientTransport: Transport { logger.debug("Starting SSE connection") + // Reset reconnection attempt on new connection + reconnectionAttempt = 0 + // Create URLSession task for SSE let (stream, response) = try await session.bytes(for: request) @@ -520,7 +830,7 @@ public actor HTTPClientTransport: Transport { } // Extract session ID if present - if let newSessionID = httpResponse.value(forHTTPHeaderField: "Mcp-Session-Id") { + if let newSessionID = httpResponse.value(forHTTPHeaderField: HTTPHeader.sessionId) { let wasSessionIDNil = (self.sessionID == nil) self.sessionID = newSessionID if wasSessionIDNil { @@ -531,14 +841,32 @@ public actor HTTPClientTransport: Transport { logger.debug("Session ID received", metadata: ["sessionID": "\(newSessionID)"]) } - try await self.processSSE(stream) + try await self.processSSE(stream, isReconnectable: true, originalRequestId: originalRequestId) } /// Processes an SSE byte stream, extracting events and delivering them /// - /// - Parameter stream: The URLSession.AsyncBytes stream to process + /// This method tracks event IDs for resumability and handles the retry directive + /// from the server to adjust reconnection timing. + /// + /// - Parameters: + /// - stream: The URLSession.AsyncBytes stream to process + /// - isReconnectable: Whether this stream should automatically reconnect on disconnect + /// - originalRequestId: Optional request ID to remap response IDs to (for stream resumption) /// - Throws: Error for stream processing failures - private func processSSE(_ stream: URLSession.AsyncBytes) async throws { + private func processSSE( + _ stream: URLSession.AsyncBytes, + isReconnectable: Bool, + originalRequestId: RequestId? = nil + ) async throws { + // Track whether we've received a priming event (event with ID) + // Per spec, server SHOULD send a priming event with ID before closing + var hasPrimingEvent = false + + // Track whether we've received a response - if so, no need to reconnect + // Reconnection is for when server disconnects BEFORE sending response + var receivedResponse = false + do { for try await event in stream.events { // Check if task has been cancelled @@ -549,18 +877,187 @@ public actor HTTPClientTransport: Transport { metadata: [ "type": "\(event.event ?? "message")", "id": "\(event.id ?? "none")", + "retry": "\(event.retry.map(String.init) ?? "none")", ] ) + // Update last event ID if provided + if let eventId = event.id { + lastEventId = eventId + // Mark that we've received a priming event - stream is now resumable + hasPrimingEvent = true + // Notify callback + onResumptionToken?(eventId) + } + + // Handle server-provided retry directive (in milliseconds, convert to seconds) + if let retryMs = event.retry { + serverRetryDelay = TimeInterval(retryMs) / 1000.0 + logger.debug( + "Server retry directive received", + metadata: ["retryMs": "\(retryMs)"] + ) + } + + // Skip events with no data (priming events, keep-alives) + if event.data.isEmpty { + continue + } + // Convert the event data to Data and yield it to the message stream - if !event.data.isEmpty, let data = event.data.data(using: .utf8) { - messageContinuation.yield(data) + if let data = event.data.data(using: .utf8) { + // Process the message: detect if it's a response and optionally remap ID + // Per MCP spec, reconnection should only stop after receiving + // the response to the original request + let processed = processJSONRPCMessage(data, originalRequestId: originalRequestId) + if processed.isResponse { + receivedResponse = true + } + messageContinuation.yield(processed.data) + } + } + + // Stream ended gracefully - check if we need to reconnect + // Reconnect if: already reconnectable (GET stream) OR received a priming event + // BUT don't reconnect if we already received a response - the request is complete + let canResume = isReconnectable || hasPrimingEvent + let needsReconnect = canResume && !receivedResponse + + if needsReconnect && isConnected && !Task.isCancelled { + logger.debug( + "SSE stream ended gracefully, will reconnect", + metadata: ["lastEventId": "\(lastEventId ?? "none")"] + ) + + // For GET streams (isReconnectable=true), the outer loop in + // startListeningForServerEvents handles reconnection. + // For POST SSE responses that received a priming event, we need to + // schedule reconnection via GET (per MCP spec: "Resumption is always via HTTP GET"). + if !isReconnectable && hasPrimingEvent { + schedulePostSSEReconnection() } } } catch { logger.error("Error processing SSE events: \(error)") - throw error + + // For GET streams, the outer loop will handle reconnection with exponential backoff. + // For POST SSE responses with a priming event, schedule reconnection via GET. + if !isReconnectable && hasPrimingEvent && !receivedResponse && isConnected + && !Task.isCancelled + { + schedulePostSSEReconnection() + } else { + throw error + } } } + + /// Schedules reconnection for a POST SSE response that was interrupted. + /// + /// Per MCP spec, resumption is always via HTTP GET with Last-Event-ID header. + /// This method spawns a task that handles reconnection with exponential backoff. + private func schedulePostSSEReconnection() { + guard let eventId = lastEventId else { + logger.warning("Cannot schedule POST SSE reconnection without lastEventId") + return + } + + // Reset reconnection attempt counter for this new reconnection sequence + reconnectionAttempt = 0 + + Task { [weak self] in + guard let self else { return } + + let maxRetries = self.reconnectionOptions.maxRetries + + while await self.isConnected && !Task.isCancelled { + let attempt = await self.reconnectionAttempt + + if attempt >= maxRetries { + self.logger.error( + "POST SSE reconnection: max attempts exceeded", + metadata: ["maxRetries": "\(maxRetries)"] + ) + return + } + + // Calculate delay with exponential backoff + let delay = await self.getNextReconnectionDelay() + await self.incrementReconnectionAttempt() + + self.logger.debug( + "POST SSE reconnection: scheduling attempt", + metadata: [ + "attempt": "\(attempt + 1)", + "delay": "\(delay)s", + "lastEventId": "\(eventId)", + ] + ) + + try? await Task.sleep(for: .seconds(delay)) + + // Check again after sleep + guard await self.isConnected && !Task.isCancelled else { return } + + do { + try await self.connectToEventStream(resumptionToken: eventId) + // Success - connectToEventStream handles SSE processing + // Reset attempt counter on success + await self.resetReconnectionAttempt() + return + } catch { + self.logger.error( + "POST SSE reconnection failed: \(error)", + metadata: ["attempt": "\(attempt + 1)"] + ) + // Continue to next iteration for retry + } + } + } + } + + /// Increments the reconnection attempt counter. + private func incrementReconnectionAttempt() { + reconnectionAttempt += 1 + } + + /// Resets the reconnection attempt counter. + private func resetReconnectionAttempt() { + reconnectionAttempt = 0 + } #endif + + // MARK: - Public Resumption API + + /// Resumes an SSE stream from a previous event ID. + /// + /// Opens a GET SSE connection with the Last-Event-ID header to replay missed events. + /// This is useful for clients that need to reconnect after a disconnection and want + /// to resume from where they left off. + /// + /// When `originalRequestId` is provided, any JSON-RPC response received on the + /// resumed stream will have its ID remapped to match the original request. This + /// ensures the response is correctly matched to the pending request in the client, + /// even if the server sends a different ID during replay. This behavior aligns + /// with the TypeScript and Python MCP SDK implementations. + /// + /// - Parameters: + /// - lastEventId: The event ID to resume from (sent as Last-Event-ID header) + /// - originalRequestId: Optional request ID to remap response IDs to + /// - Throws: MCPError if the connection fails + public func resumeStream(from lastEventId: String, forRequestId originalRequestId: RequestId? = nil) async throws { + #if os(Linux) + logger.warning("resumeStream is not supported on Linux (SSE not available)") + #else + try await connectToEventStream(resumptionToken: lastEventId, originalRequestId: originalRequestId) + #endif + } + + /// The last event ID received from the server. + /// + /// This can be used to persist the event ID and resume the stream later + /// using `resumeStream(from:)`. + public var lastReceivedEventId: String? { + lastEventId + } } diff --git a/Sources/MCP/Base/Transports/HTTPServerTransport+Types.swift b/Sources/MCP/Base/Transports/HTTPServerTransport+Types.swift new file mode 100644 index 00000000..6adb29e0 --- /dev/null +++ b/Sources/MCP/Base/Transports/HTTPServerTransport+Types.swift @@ -0,0 +1,304 @@ +import Foundation + +// Types extracted from HTTPServerTransport.swift +// - Options +// - SecuritySettings +// - EventStore protocol +// - HTTPRequest +// - HTTPResponse + +/// Configuration options for HTTPServerTransport +public struct HTTPServerTransportOptions: Sendable { + /// Function that generates a session ID for the transport. + /// The session ID SHOULD be globally unique and cryptographically secure + /// (e.g., a securely generated UUID, a JWT, or a cryptographic hash). + /// + /// If not provided, session management is disabled (stateless mode). + public var sessionIdGenerator: (@Sendable () -> String)? + + /// Called when the server initializes a new session. + /// This is called when the server receives an initialize request and generates a session ID. + /// Useful for tracking multiple MCP sessions. + public var onSessionInitialized: (@Sendable (String) async -> Void)? + + /// Called when the server closes a session (DELETE request). + /// Useful for cleaning up resources associated with the session. + public var onSessionClosed: (@Sendable (String) async -> Void)? + + /// If true, the server will return JSON responses instead of starting an SSE stream. + /// This can be useful for simple request/response scenarios without streaming. + /// Default is false (SSE streams are preferred). + public var enableJsonResponse: Bool + + /// Event store for resumability support. + /// If provided, resumability will be enabled, allowing clients to reconnect and resume messages. + public var eventStore: EventStore? + + /// Retry interval in milliseconds to suggest to clients in SSE retry field. + /// When set, the server will send a retry field in SSE priming events to control + /// client reconnection timing for polling behavior. + public var retryInterval: Int? + + /// Security settings for DNS rebinding protection. + /// + /// When nil, no security validation is performed. + /// Use `TransportSecuritySettings.forLocalhost(port:)` for localhost-bound servers. + /// + /// See `TransportSecuritySettings` documentation for details on DNS rebinding attacks + /// and the rationale for protection. + public var security: TransportSecuritySettings? + + public init( + sessionIdGenerator: (@Sendable () -> String)? = nil, + onSessionInitialized: (@Sendable (String) async -> Void)? = nil, + onSessionClosed: (@Sendable (String) async -> Void)? = nil, + enableJsonResponse: Bool = false, + eventStore: EventStore? = nil, + retryInterval: Int? = nil, + security: TransportSecuritySettings? = nil + ) { + self.sessionIdGenerator = sessionIdGenerator + self.onSessionInitialized = onSessionInitialized + self.onSessionClosed = onSessionClosed + self.enableJsonResponse = enableJsonResponse + self.eventStore = eventStore + self.retryInterval = retryInterval + self.security = security + } +} + +/// Security settings for DNS rebinding protection. +/// +/// DNS rebinding is an attack where a malicious website can bypass same-origin policy +/// by manipulating DNS responses, potentially allowing browser-based attackers to +/// interact with local MCP servers. This is particularly dangerous for servers +/// bound to localhost. +/// +/// ## How Protection Works +/// +/// When enabled, the transport validates: +/// 1. **Host header**: Must match an allowed host pattern (prevents DNS rebinding) +/// 2. **Origin header**: If present (browser requests), must match an allowed origin +/// +/// ## Usage +/// +/// ```swift +/// // Auto-enabled for localhost (recommended) +/// let settings = TransportSecuritySettings.forLocalhost(port: 8080) +/// +/// // Or manually configure +/// let settings = TransportSecuritySettings( +/// enableDnsRebindingProtection: true, +/// allowedHosts: ["myserver.local:8080"], +/// allowedOrigins: ["http://myserver.local:8080"] +/// ) +/// ``` +public struct TransportSecuritySettings: Sendable { + /// Whether to validate Host and Origin headers for DNS rebinding protection. + public var enableDnsRebindingProtection: Bool + + /// Allowed Host header values. Supports wildcard port patterns like "127.0.0.1:*". + /// When protection is enabled, requests with Host headers not matching any pattern are rejected. + public var allowedHosts: [String] + + /// Allowed Origin header values. Supports wildcard port patterns like "http://localhost:*". + /// When protection is enabled and an Origin header is present, it must match one of these patterns. + /// Requests without an Origin header are allowed (non-browser clients). + public var allowedOrigins: [String] + + public init( + enableDnsRebindingProtection: Bool = false, + allowedHosts: [String] = [], + allowedOrigins: [String] = [] + ) { + self.enableDnsRebindingProtection = enableDnsRebindingProtection + self.allowedHosts = allowedHosts + self.allowedOrigins = allowedOrigins + } + + /// Creates security settings for a localhost-bound server. + /// + /// - Parameter port: The port number (use "*" pattern if port varies) + /// - Returns: Security settings with protection enabled for all localhost variants + public static func forLocalhost(port: Int? = nil) -> TransportSecuritySettings { + let portPattern = port.map { String($0) } ?? "*" + return TransportSecuritySettings( + enableDnsRebindingProtection: true, + allowedHosts: [ + "127.0.0.1:\(portPattern)", + "localhost:\(portPattern)", + "[::1]:\(portPattern)", + ], + allowedOrigins: [ + "http://127.0.0.1:\(portPattern)", + "http://localhost:\(portPattern)", + "http://[::1]:\(portPattern)", + ] + ) + } + + /// Creates security settings appropriate for the given bind address. + /// + /// Auto-enables DNS rebinding protection for localhost addresses, + /// returns nil for other addresses (no protection needed for remote bindings). + /// + /// - Parameters: + /// - host: The host address the server is binding to + /// - port: The port number + /// - Returns: Security settings if protection should be enabled, nil otherwise + public static func forBindAddress(host: String, port: Int) -> TransportSecuritySettings? { + let localhostAddresses = ["127.0.0.1", "localhost", "::1"] + if localhostAddresses.contains(host) { + return forLocalhost(port: port) + } + return nil + } +} + +/// Protocol for storing and replaying SSE events for resumability support. +/// +/// Implementations should store events durably and support replaying them +/// when clients reconnect with a Last-Event-ID header. +/// +/// ## Priming Events +/// +/// Priming events are stored with empty `Data()` as the message. These events +/// establish the initial event ID for a stream but should **not** be replayed +/// as regular messages. During replay, implementations should skip events with +/// empty message data and only replay actual JSON-RPC messages. +public protocol EventStore: Sendable { + /// Stores an event and returns its unique ID. + /// + /// - Parameters: + /// - streamId: The stream this event belongs to + /// - message: The JSON-RPC message data. Empty `Data()` indicates a priming event + /// which should be skipped during replay. + /// - Returns: A unique event ID for this event + func storeEvent(streamId: String, message: Data) async throws -> String + + /// Gets the stream ID associated with an event ID. + /// - Parameter eventId: The event ID to look up + /// - Returns: The stream ID, or nil if not found + func streamIdForEventId(_ eventId: String) async -> String? + + /// Replays events after the given event ID. + /// + /// Implementations should skip priming events (empty message data) during replay. + /// Only actual JSON-RPC messages should be sent to the callback. + /// + /// - Parameters: + /// - lastEventId: The last event ID the client received + /// - send: Callback to send each replayed event (eventId, message) + /// - Returns: The stream ID for continued event delivery + func replayEventsAfter( + _ lastEventId: String, + send: @escaping @Sendable (String, Data) async throws -> Void + ) async throws -> String +} + +/// HTTP response returned by `HTTPServerTransport.handleRequest(_:)`. +/// +/// This struct represents the result of processing an MCP request. It can contain either: +/// - A simple JSON response with `body` data (for non-streaming responses) +/// - An SSE stream for streaming responses (for long-running operations or server-initiated messages) +/// +/// ## Usage with HTTP Frameworks +/// +/// When integrating with an HTTP framework like Vapor or Hummingbird, convert this response +/// to the framework's native response type: +/// +/// ```swift +/// // Vapor example +/// func handleMCP(req: Request) async throws -> Response { +/// let httpRequest = HTTPRequest( +/// method: req.method.rawValue, +/// headers: Dictionary(req.headers.map { ($0.name, $0.value) }) { _, last in last }, +/// body: req.body.data +/// ) +/// let response = await transport.handleRequest(httpRequest) +/// +/// if let stream = response.stream { +/// // Return SSE response +/// return Response(status: .init(statusCode: response.statusCode), body: .init(asyncSequence: stream)) +/// } else { +/// // Return JSON response +/// return Response(status: .init(statusCode: response.statusCode), body: .init(data: response.body ?? Data())) +/// } +/// } +/// ``` +public struct HTTPResponse: Sendable { + /// The HTTP status code for the response (e.g., 200, 400, 404). + public let statusCode: Int + /// HTTP headers to include in the response (e.g., Content-Type, Mcp-Session-Id). + public let headers: [String: String] + /// Response body data for non-streaming responses. Nil for SSE streaming responses. + public let body: Data? + /// SSE stream for streaming responses. Nil for simple JSON responses. + /// When present, the caller should stream this data to the client as Server-Sent Events. + public let stream: AsyncThrowingStream? + + public init( + statusCode: Int, + headers: [String: String] = [:], + body: Data? = nil, + stream: AsyncThrowingStream? = nil + ) { + self.statusCode = statusCode + self.headers = headers + self.body = body + self.stream = stream + } +} + +/// HTTP request abstraction for framework-agnostic handling. +/// +/// This struct provides a common interface for HTTP requests that can be populated from +/// any HTTP server framework (Vapor, Hummingbird, SwiftNIO, etc.). The +/// `HTTPServerTransport` uses this abstraction to process MCP requests +/// without being coupled to a specific framework. +/// +/// ## Usage +/// +/// Convert your framework's request type to `HTTPRequest` before passing to the transport: +/// +/// ```swift +/// // Vapor example +/// let httpRequest = HTTPRequest( +/// method: req.method.rawValue, +/// headers: Dictionary(req.headers.map { ($0.name, $0.value) }) { _, last in last }, +/// body: req.body.data +/// ) +/// +/// // Hummingbird example +/// let httpRequest = HTTPRequest( +/// method: String(describing: request.method), +/// headers: Dictionary(request.headers.map { ($0.name.rawName, $0.value) }) { _, last in last }, +/// body: request.body.buffer?.getData(at: 0, length: request.body.buffer?.readableBytes ?? 0) +/// ) +/// ``` +public struct HTTPRequest: Sendable { + /// The HTTP method (e.g., "GET", "POST", "DELETE"). + public let method: String + /// Request headers as a case-sensitive dictionary. + /// Use the `header(_:)` method for case-insensitive header lookup. + public let headers: [String: String] + /// The request body data, if present. + public let body: Data? + + public init(method: String, headers: [String: String] = [:], body: Data? = nil) { + self.method = method + self.headers = headers + self.body = body + } + + /// Get a header value (case-insensitive) + public func header(_ name: String) -> String? { + let lowercased = name.lowercased() + for (key, value) in headers { + if key.lowercased() == lowercased { + return value + } + } + return nil + } +} diff --git a/Sources/MCP/Base/Transports/HTTPServerTransport.swift b/Sources/MCP/Base/Transports/HTTPServerTransport.swift new file mode 100644 index 00000000..69f9c872 --- /dev/null +++ b/Sources/MCP/Base/Transports/HTTPServerTransport.swift @@ -0,0 +1,1163 @@ +import Foundation +import Logging + +#if canImport(FoundationNetworking) + import FoundationNetworking +#endif + +/// Internal stream state for managing SSE connections +private struct StreamState: Sendable { + /// Continuation for pushing SSE data + let continuation: AsyncThrowingStream.Continuation + + /// Cleanup function to close stream and remove mapping + let cleanup: @Sendable () -> Void +} + +/// Internal stream state for JSON response mode +private struct JsonStreamState: Sendable { + /// Continuation for yielding the HTTP response + let continuation: AsyncThrowingStream.Continuation +} + +/// Server transport for Streamable HTTP: implements the MCP Streamable HTTP transport specification. +/// +/// This transport can be integrated with any HTTP server framework (Hummingbird, Vapor, etc.) +/// by passing incoming requests to `handleRequest`. +/// +/// Usage example: +/// ```swift +/// // Stateful mode - server manages session IDs +/// let transport = HTTPServerTransport( +/// options: .init( +/// sessionIdGenerator: { UUID().uuidString }, +/// onSessionInitialized: { sessionId in +/// await sessions.store(sessionId, transport: transport) +/// }, +/// onSessionClosed: { sessionId in +/// await sessions.remove(sessionId) +/// } +/// ) +/// ) +/// +/// // Stateless mode +/// let statelessTransport = HTTPServerTransport() +/// +/// // In your HTTP handler: +/// let response = try await transport.handleRequest(httpRequest) +/// ``` +/// +/// In stateful mode: +/// - Session ID is generated and included in response headers +/// - Requests with invalid session IDs are rejected with 404 Not Found +/// - Non-initialization requests without a session ID are rejected with 400 Bad Request +/// +/// In stateless mode: +/// - No Session ID is included in responses +/// - No session validation is performed +public actor HTTPServerTransport: Transport { + /// Logger for transport events + public nonisolated let logger: Logger + + /// The session ID for this transport (nil in stateless mode) + public private(set) var sessionId: String? + + /// Whether this transport has been initialized + private var initialized = false + + /// Whether this session has been terminated (via DELETE) + private var terminated = false + + /// Whether the transport has been started + private var started = false + + /// The negotiated protocol version, set after initialization + private var negotiatedProtocolVersion: String? + + // Configuration + private let options: HTTPServerTransportOptions + + // Stream multiplexing (matching TypeScript's three maps pattern) + private var streamMapping: [String: StreamState] = [:] + private var jsonStreamMapping: [String: JsonStreamState] = [:] + private var requestToStreamMapping: [RequestId: String] = [:] + private var requestResponseMap: [RequestId: Data] = [:] + + // Standalone SSE stream ID for GET requests + private let standaloneSseStreamId = "_GET_stream" + + // Server receive stream (messages from HTTP clients go here) + private let serverStream: AsyncThrowingStream + private let serverContinuation: AsyncThrowingStream.Continuation + + /// Closure called when the transport is closed + public var onClose: (@Sendable () async -> Void)? + + /// Creates a new HTTPServerTransport. + /// + /// - Parameters: + /// - options: Transport configuration options + /// - logger: Optional logger instance + public init( + options: HTTPServerTransportOptions = .init(), + logger: Logger? = nil + ) { + self.options = options + self.logger = + logger + ?? Logger( + label: "mcp.transport.http.server", + factory: { _ in SwiftLogNoOpLogHandler() } + ) + + // Create server receive stream + var continuation: AsyncThrowingStream.Continuation! + self.serverStream = AsyncThrowingStream { continuation = $0 } + self.serverContinuation = continuation + } + + // MARK: - Transport Protocol + + /// Starts the transport. + /// This is required by the Transport interface but is a no-op for HTTP transports + /// as connections are managed per-request. + public func connect() async throws { + guard !started else { + throw MCPError.internalError("Transport already started") + } + started = true + } + + /// Disconnects and closes the transport. + public func disconnect() async { + await close() + } + + /// Sends data to the appropriate client connection. + /// + /// For responses, the request ID is extracted from the message. + /// For notifications during tool execution, use `send(_:relatedRequestId:)`. + public func send(_ data: Data) async throws { + try await send(data, relatedRequestId: nil) + } + + /// Sends data with an optional related request ID for response routing. + /// + /// - Parameters: + /// - data: The data to send + /// - relatedRequestId: The ID of the request this message relates to + public func send(_ data: Data, relatedRequestId: RequestId?) async throws { + var requestId = relatedRequestId + + // For responses, extract the ID from the message + if requestId == nil { + requestId = extractResponseId(from: data) + } + + // If no request ID, send to standalone SSE stream + if requestId == nil { + // Generate and store event ID if event store is provided + var eventId: String? + if let eventStore = options.eventStore { + eventId = try await eventStore.storeEvent(streamId: standaloneSseStreamId, message: data) + } + + if let streamState = streamMapping[standaloneSseStreamId] { + let sseData = formatSSEEvent(data: data, eventId: eventId) + streamState.continuation.yield(sseData) + } + return + } + + guard let requestId = requestId else { return } + + // Get the stream for this request + guard let streamId = requestToStreamMapping[requestId] else { + logger.debug("No stream found for request \(requestId) - client may have disconnected") + return + } + + // Check if using JSON response mode + if let jsonState = jsonStreamMapping[streamId] { + // Store the response + requestResponseMap[requestId] = data + try await checkBatchCompletion(streamId: streamId, jsonState: jsonState) + return + } + + // SSE streaming mode + if let streamState = streamMapping[streamId] { + // Generate event ID if event store is provided + var eventId: String? + if let eventStore = options.eventStore { + eventId = try await eventStore.storeEvent(streamId: streamId, message: data) + } + + let sseData = formatSSEEvent(data: data, eventId: eventId) + streamState.continuation.yield(sseData) + + // Track response for batch completion + let isResponse = isJSONRPCResponse(data) + if isResponse { + requestResponseMap[requestId] = data + try await checkStreamCompletion(streamId: streamId) + } + } + } + + /// Returns the stream of messages from HTTP clients. + public func receive() -> AsyncThrowingStream { + return serverStream + } + + // MARK: - HTTP Request Handling + + /// Handles an incoming HTTP request. + /// + /// This method routes the request based on HTTP method: + /// - POST: Handle JSON-RPC messages + /// - GET: Establish SSE stream for server-initiated notifications + /// - DELETE: Terminate the session + /// + /// - Parameter request: The incoming HTTP request + /// - Returns: An HTTP response + public func handleRequest(_ request: HTTPRequest) async -> HTTPResponse { + // Check if transport has been terminated (applies to all modes) + // Per spec: server MUST respond to requests after termination with 404 Not Found + if terminated { + return createJsonErrorResponse( + status: 404, + code: ErrorCode.connectionClosed, + message: "Session has been terminated" + ) + } + + // Validate security headers (DNS rebinding protection) + if let error = validateSecurityHeaders(request) { + return error + } + + switch request.method.uppercased() { + case "POST": + return await handlePostRequest(request) + case "GET": + return await handleGetRequest(request) + case "DELETE": + return await handleDeleteRequest(request) + default: + return createJsonErrorResponse( + status: 405, + code: ErrorCode.connectionClosed, + message: "Method not allowed", + extraHeaders: [HTTPHeader.allow: "GET, POST, DELETE"] + ) + } + } + + // MARK: - POST Request Handling + + private func handlePostRequest(_ request: HTTPRequest) async -> HTTPResponse { + // Validate Accept header + // Per spec: Client must accept both application/json and text/event-stream for SSE mode. + // However, when JSON response mode is enabled, only application/json is required. + let acceptHeader = request.header(HTTPHeader.accept) ?? "" + if options.enableJsonResponse { + // JSON response mode only requires application/json + guard acceptHeader.contains("application/json") else { + return createJsonErrorResponse( + status: 406, + code: ErrorCode.connectionClosed, + message: "Not Acceptable: Client must accept application/json" + ) + } + } else { + // SSE mode requires both content types + guard acceptHeader.contains("application/json") && acceptHeader.contains("text/event-stream") else { + return createJsonErrorResponse( + status: 406, + code: ErrorCode.connectionClosed, + message: "Not Acceptable: Client must accept both application/json and text/event-stream" + ) + } + } + + // Validate Content-Type + let contentType = request.header(HTTPHeader.contentType) ?? "" + guard contentType.contains("application/json") else { + return createJsonErrorResponse( + status: 415, + code: ErrorCode.connectionClosed, + message: "Unsupported Media Type: Content-Type must be application/json" + ) + } + + // Parse the request body + guard let body = request.body, !body.isEmpty else { + return createJsonErrorResponse( + status: 400, + code: ErrorCode.parseError, + message: "Parse error: Empty request body" + ) + } + + // Try to parse as JSON-RPC message(s) + let messages: [[String: Any]] + do { + if let parsed = try JSONSerialization.jsonObject(with: body) as? [[String: Any]] { + messages = parsed + } else if let single = try JSONSerialization.jsonObject(with: body) as? [String: Any] { + messages = [single] + } else { + return createJsonErrorResponse( + status: 400, + code: ErrorCode.parseError, + message: "Parse error: Invalid JSON-RPC message" + ) + } + } catch { + return createJsonErrorResponse( + status: 400, + code: ErrorCode.parseError, + message: "Parse error: Invalid JSON" + ) + } + + // Validate JSON-RPC format - all messages must have "jsonrpc": "2.0" + for message in messages { + guard let jsonrpc = message["jsonrpc"] as? String, jsonrpc == "2.0" else { + return createJsonErrorResponse( + status: 400, + code: ErrorCode.invalidRequest, + message: "Invalid Request: Missing or invalid jsonrpc version" + ) + } + } + + // Check for initialization request + let isInitializationRequest = messages.contains { isInitializeRequest($0) } + + // Check for batch requests (protocol version conditional) + let isBatchRequest = messages.count > 1 + + if isInitializationRequest { + // Check if already initialized in stateful mode + if initialized && sessionId != nil { + return createJsonErrorResponse( + status: 400, + code: ErrorCode.invalidRequest, + message: "Invalid Request: Server already initialized" + ) + } + + // Only one initialize request allowed + if isBatchRequest { + return createJsonErrorResponse( + status: 400, + code: ErrorCode.invalidRequest, + message: "Invalid Request: Only one initialization request is allowed" + ) + } + + // Extract and store the protocol version from the initialize request + let clientProtocolVersion = extractProtocolVersionFromInitialize(messages) + negotiatedProtocolVersion = clientProtocolVersion + + // Generate session ID if in stateful mode + if let generator = options.sessionIdGenerator { + let generatedId = generator() + + // Validate session ID per spec: must be visible ASCII (0x21-0x7E) + if !isValidSessionId(generatedId) { + logger.error( + "Generated session ID contains invalid characters", + metadata: ["sessionId": "\(generatedId)"] + ) + return createJsonErrorResponse( + status: 500, + code: ErrorCode.connectionClosed, + message: "Internal error: Invalid session ID generated" + ) + } + + sessionId = generatedId + initialized = true + + // Fire session initialized callback BEFORE dispatching to server + if let sessionId, let callback = options.onSessionInitialized { + await callback(sessionId) + } + } else { + initialized = true + } + } else { + // Validate session for non-initialization requests + if let error = validateSession(request) { + return error + } + + // Validate protocol version + if let error = validateProtocolVersion(request) { + return error + } + + // Reject batch requests for protocol version >= 2025-06-18 + // Batching was removed from the spec starting with 2025-06-18 + if isBatchRequest { + let protocolVersion = request.header(HTTPHeader.protocolVersion) ?? Version.defaultNegotiated + if protocolVersion >= Version.v2025_06_18 { + return createJsonErrorResponse( + status: 400, + code: ErrorCode.invalidRequest, + message: "Invalid Request: Batch requests not supported in protocol version \(protocolVersion)" + ) + } + } + } + + // Check if messages contain any requests (vs just notifications) + let hasRequests = messages.contains { isJSONRPCRequest($0) } + + if !hasRequests { + // Only notifications - yield to server and return 202 + serverContinuation.yield(body) + return HTTPResponse(statusCode: 202, headers: sessionHeaders()) + } + + // Extract request IDs + let requestIds = extractRequestIds(from: messages) + let streamId = UUID().uuidString + + // Map request IDs to this stream + for id in requestIds { + requestToStreamMapping[id] = streamId + } + + // Check if using JSON response mode + if options.enableJsonResponse { + return await handleJsonResponseMode(streamId: streamId, requestIds: requestIds, body: body) + } + + // SSE streaming mode + return await handleSSEStreamingMode( + streamId: streamId, + requestIds: requestIds, + body: body, + request: request, + messages: messages + ) + } + + private func handleJsonResponseMode( + streamId: String, + requestIds: [RequestId], + body: Data + ) async -> HTTPResponse { + // Create stream for receiving the response + let (stream, continuation) = AsyncThrowingStream.makeStream() + + let state = JsonStreamState(continuation: continuation) + jsonStreamMapping[streamId] = state + + // Yield the message to the server + serverContinuation.yield(body) + + // Wait for response - this is cancellation-aware unlike withCheckedContinuation + do { + for try await response in stream { + return response + } + } catch { + // Stream was finished with error (e.g., transport closed) + logger.debug("JSON response stream ended with error: \(error)") + } + + // Stream closed without yielding a response - return error + return createJsonErrorResponse( + status: 503, + code: ErrorCode.connectionClosed, + message: "Service Unavailable: No response received" + ) + } + + private func handleSSEStreamingMode( + streamId: String, + requestIds: [RequestId], + body: Data, + request: HTTPRequest, + messages: [[String: Any]] + ) async -> HTTPResponse { + let (stream, streamContinuation) = AsyncThrowingStream.makeStream() + + // Clean up mapping when stream terminates (e.g., client disconnect) + streamContinuation.onTermination = { @Sendable [weak self] _ in + Task { await self?.cleanUpStreamMapping(for: streamId) } + } + + let cleanup: @Sendable () -> Void = { + streamContinuation.finish() + } + + let state = StreamState( + continuation: streamContinuation, + cleanup: cleanup + ) + + streamMapping[streamId] = state + + // Use negotiated protocol version if available, otherwise extract from request + let protocolVersion = negotiatedProtocolVersion ?? extractProtocolVersion(from: messages, request: request) + + // Write priming event if appropriate + await writePrimingEvent(streamId: streamId, continuation: streamContinuation, protocolVersion: protocolVersion) + + // Yield the message to the server + serverContinuation.yield(body) + + var headers = sessionHeaders() + headers[HTTPHeader.contentType] = "text/event-stream" + headers[HTTPHeader.cacheControl] = "no-cache, no-transform" + headers[HTTPHeader.connection] = "keep-alive" + + return HTTPResponse(statusCode: 200, headers: headers, stream: stream) + } + + // MARK: - GET Request Handling + + private func handleGetRequest(_ request: HTTPRequest) async -> HTTPResponse { + // Validate Accept header + let acceptHeader = request.header(HTTPHeader.accept) ?? "" + guard acceptHeader.contains("text/event-stream") else { + return createJsonErrorResponse( + status: 406, + code: ErrorCode.connectionClosed, + message: "Not Acceptable: Client must accept text/event-stream" + ) + } + + // Validate session + if let error = validateSession(request) { + return error + } + + // Validate protocol version + if let error = validateProtocolVersion(request) { + return error + } + + // Handle resumability + if let eventStore = options.eventStore, + let lastEventId = request.header(HTTPHeader.lastEventId) + { + return await replayEvents(lastEventId: lastEventId, eventStore: eventStore, request: request) + } + + // Check if there's already an active standalone SSE stream + if streamMapping[standaloneSseStreamId] != nil { + return createJsonErrorResponse( + status: 409, + code: ErrorCode.connectionClosed, + message: "Conflict: Only one SSE stream is allowed per session" + ) + } + + let (stream, streamContinuation) = AsyncThrowingStream.makeStream() + + // Clean up mapping when stream terminates (e.g., client disconnect) + let streamId = standaloneSseStreamId + streamContinuation.onTermination = { @Sendable [weak self] _ in + Task { await self?.cleanUpStreamMapping(for: streamId) } + } + + let cleanup: @Sendable () -> Void = { + streamContinuation.finish() + } + + streamMapping[standaloneSseStreamId] = StreamState( + continuation: streamContinuation, + cleanup: cleanup + ) + + // Write priming event for resumability (use negotiated version or header) + let protocolVersion = negotiatedProtocolVersion ?? request.header(HTTPHeader.protocolVersion) ?? Version.defaultNegotiated + await writePrimingEvent(streamId: standaloneSseStreamId, continuation: streamContinuation, protocolVersion: protocolVersion) + + var headers = sessionHeaders() + headers[HTTPHeader.contentType] = "text/event-stream" + headers[HTTPHeader.cacheControl] = "no-cache, no-transform" + headers[HTTPHeader.connection] = "keep-alive" + + return HTTPResponse(statusCode: 200, headers: headers, stream: stream) + } + + // MARK: - DELETE Request Handling + + private func handleDeleteRequest(_ request: HTTPRequest) async -> HTTPResponse { + // DELETE is only valid in stateful mode (when session management is enabled) + // In stateless mode, there's no session to terminate + guard options.sessionIdGenerator != nil else { + return createJsonErrorResponse( + status: 405, + code: ErrorCode.connectionClosed, + message: "Method Not Allowed: Session management is not enabled", + extraHeaders: [HTTPHeader.allow: "GET, POST"] + ) + } + + // Validate session + if let error = validateSession(request) { + return error + } + + // Validate protocol version + if let error = validateProtocolVersion(request) { + return error + } + + // Fire session closed callback + if let sessionId, let callback = options.onSessionClosed { + await callback(sessionId) + } + + await close() + + return HTTPResponse(statusCode: 200, headers: sessionHeaders()) + } + + // MARK: - Close + + /// Closes the transport and all active streams. + public func close() async { + // Mark session as terminated so subsequent requests are rejected with 404 + terminated = true + + // Close all SSE streams and remove mappings synchronously + for (streamId, state) in streamMapping { + state.cleanup() + streamMapping.removeValue(forKey: streamId) + } + + // Finish all pending JSON response streams. + // The for-await loop in handleJsonResponseMode will exit and return a 503 error. + for (streamId, state) in jsonStreamMapping { + state.continuation.finish() + jsonStreamMapping.removeValue(forKey: streamId) + } + + // Clear request mappings + requestToStreamMapping.removeAll() + requestResponseMap.removeAll() + + // Finish the server stream + serverContinuation.finish() + + await onClose?() + } + + // MARK: - Stream Control + + /// Closes an SSE stream for a specific request, triggering client reconnection. + /// + /// Use this to implement polling behavior during long-running operations - + /// the client will reconnect after the retry interval specified in the priming event. + /// + /// - Parameter requestId: The ID of the request whose stream should be closed + public func closeSSEStream(for requestId: RequestId) { + guard let streamId = requestToStreamMapping[requestId] else { return } + + if let stream = streamMapping.removeValue(forKey: streamId) { + stream.cleanup() + } + } + + /// Closes the standalone GET SSE stream, triggering client reconnection. + /// + /// Use this to implement polling behavior for server-initiated notifications. + public func closeStandaloneSSEStream() { + if let stream = streamMapping.removeValue(forKey: standaloneSseStreamId) { + stream.cleanup() + } + } + + /// Removes a stream from the mapping without calling cleanup. + /// Used by onTermination handlers when the stream has already terminated. + private func cleanUpStreamMapping(for streamId: String) { + streamMapping.removeValue(forKey: streamId) + } + + // MARK: - Helper Methods + + private func sessionHeaders() -> [String: String] { + var headers: [String: String] = [:] + if let sessionId { + headers[HTTPHeader.sessionId] = sessionId + } + return headers + } + + // MARK: - Security Validation + + /// Validates Host and Origin headers for DNS rebinding protection. + /// + /// DNS rebinding attacks allow malicious websites to bypass browser same-origin policy + /// by manipulating DNS responses. This is particularly dangerous for localhost servers + /// as browsers may allow requests from attacker-controlled pages to local services. + private func validateSecurityHeaders(_ request: HTTPRequest) -> HTTPResponse? { + guard let security = options.security, + security.enableDnsRebindingProtection + else { + return nil + } + + // Validate Host header (required when protection is enabled) + let hostHeader = request.header(HTTPHeader.host) + if hostHeader == nil { + logger.warning("DNS rebinding protection: Missing Host header") + // Use 421 Misdirected Request for Host header issues + return createJsonErrorResponse( + status: 421, + code: ErrorCode.connectionClosed, + message: "Misdirected Request: Missing Host header" + ) + } + + let hostMatches = security.allowedHosts.contains { pattern in + matchesHostPattern(hostHeader!, pattern: pattern) + } + + if !hostMatches { + logger.warning( + "DNS rebinding protection: Host header rejected", + metadata: ["host": "\(hostHeader!)"] + ) + // Use 421 Misdirected Request for Host header issues + return createJsonErrorResponse( + status: 421, + code: ErrorCode.connectionClosed, + message: "Misdirected Request: Host header not allowed" + ) + } + + // Validate Origin header (only if present - non-browser clients won't send it) + if let originHeader = request.header(HTTPHeader.origin) { + let originMatches = security.allowedOrigins.contains { pattern in + matchesOriginPattern(originHeader, pattern: pattern) + } + + if !originMatches { + logger.warning( + "DNS rebinding protection: Origin header rejected", + metadata: ["origin": "\(originHeader)"] + ) + return createJsonErrorResponse( + status: 403, + code: ErrorCode.connectionClosed, + message: "Forbidden: Origin not allowed" + ) + } + } + + return nil + } + + /// Matches a host value against a pattern that may contain port wildcards. + /// + /// Patterns like "localhost:*" match "localhost:8080", "localhost:3000", etc. + private func matchesHostPattern(_ host: String, pattern: String) -> Bool { + if pattern.hasSuffix(":*") { + let patternHost = String(pattern.dropLast(2)) // Remove ":*" + // Host must start with pattern host and have a port + if host.hasPrefix(patternHost + ":") { + let portPart = host.dropFirst(patternHost.count + 1) + // Verify the rest is a valid port (digits only) + return !portPart.isEmpty && portPart.allSatisfy { $0.isNumber } + } + return false + } + // Exact match + return host == pattern + } + + /// Matches an origin value against a pattern that may contain port wildcards. + /// + /// Patterns like "http://localhost:*" match "http://localhost:8080", etc. + private func matchesOriginPattern(_ origin: String, pattern: String) -> Bool { + if pattern.hasSuffix(":*") { + let patternPrefix = String(pattern.dropLast(2)) // Remove ":*" + // Origin must start with pattern prefix and have a port + if origin.hasPrefix(patternPrefix + ":") { + let portPart = origin.dropFirst(patternPrefix.count + 1) + // Verify the rest is a valid port (digits only, possibly followed by path) + let portString = portPart.prefix(while: { $0.isNumber }) + return !portString.isEmpty + } + return false + } + // Exact match or origin is prefix (origin may have trailing path) + return origin == pattern || origin.hasPrefix(pattern + "/") + } + + // MARK: - Session Validation + + private func validateSession(_ request: HTTPRequest) -> HTTPResponse? { + // Check initialization status first - applies to BOTH stateful and stateless modes + // Per MCP spec, clients should not send requests before initialization + guard initialized else { + return createJsonErrorResponse( + status: 400, + code: ErrorCode.connectionClosed, + message: "Bad Request: Server not initialized" + ) + } + + // If no session ID generator, we're in stateless mode - skip session ID validation + guard options.sessionIdGenerator != nil else { + return nil + } + + // If session was terminated (via DELETE), reject with 404 + if terminated { + return createJsonErrorResponse( + status: 404, + code: ErrorCode.connectionClosed, + message: "Session has been terminated" + ) + } + + let requestSessionId = request.header(HTTPHeader.sessionId) + + // Non-initialization requests must include session ID + guard let requestSessionId else { + return createJsonErrorResponse( + status: 400, + code: ErrorCode.connectionClosed, + message: "Bad Request: \(HTTPHeader.sessionId) header is required" + ) + } + + // Session ID must match + guard requestSessionId == sessionId else { + return createJsonErrorResponse( + status: 404, + code: ErrorCode.connectionClosed, + message: "Session not found" + ) + } + + return nil + } + + /// Supported protocol versions for header validation. + /// This is more lenient than Version.supported - it accepts headers from clients + /// even if we don't fully implement all features of that version yet. + private static let supportedProtocolVersions = [ + Version.v2024_11_05, + Version.v2025_03_26, + Version.v2025_06_18, + Version.v2025_11_25, + ] + + private func validateProtocolVersion(_ request: HTTPRequest) -> HTTPResponse? { + let protocolVersion = request.header(HTTPHeader.protocolVersion) + + // If header is present, validate it + if let version = protocolVersion { + guard Self.supportedProtocolVersions.contains(version) else { + return createJsonErrorResponse( + status: 400, + code: ErrorCode.connectionClosed, + message: + "Bad Request: Unsupported protocol version: \(version) (supported: \(Self.supportedProtocolVersions.joined(separator: ", ")))" + ) + } + } + + return nil + } + + private func isInitializeRequest(_ message: [String: Any]) -> Bool { + guard let method = message["method"] as? String else { return false } + return method == "initialize" + } + + /// Validates that a session ID contains only visible ASCII characters (0x21-0x7E). + /// + /// Per MCP spec: "Session IDs MUST be visible ASCII characters only." + /// This range includes printable characters from '!' (0x21) to '~' (0x7E), + /// excluding space (0x20) and control characters. + private func isValidSessionId(_ sessionId: String) -> Bool { + guard !sessionId.isEmpty else { return false } + return sessionId.utf8.allSatisfy { byte in + byte >= 0x21 && byte <= 0x7E + } + } + + /// Extracts the protocol version from an initialize request's params. + private func extractProtocolVersionFromInitialize(_ messages: [[String: Any]]) -> String { + for message in messages where isInitializeRequest(message) { + if let params = message["params"] as? [String: Any], + let version = params["protocolVersion"] as? String + { + return version + } + } + return Version.defaultNegotiated // Default per spec + } + + private func isJSONRPCRequest(_ message: [String: Any]) -> Bool { + return message["method"] != nil && message["id"] != nil + } + + private func isJSONRPCResponse(_ data: Data) -> Bool { + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { + return false + } + return json["result"] != nil || json["error"] != nil + } + + private func extractRequestIds(from messages: [[String: Any]]) -> [RequestId] { + var ids: [RequestId] = [] + for message in messages { + guard message["method"] != nil else { continue } + if let stringId = message["id"] as? String { + ids.append(.string(stringId)) + } else if let intId = message["id"] as? Int { + ids.append(.number(intId)) + } + } + return ids + } + + private func extractResponseId(from data: Data) -> RequestId? { + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { + return nil + } + + // Check if it's a response (has result or error) + guard json["result"] != nil || json["error"] != nil else { + return nil + } + + if let stringId = json["id"] as? String { + return .string(stringId) + } else if let intId = json["id"] as? Int { + return .number(intId) + } + return nil + } + + private func extractProtocolVersion(from messages: [[String: Any]], request: HTTPRequest) -> String { + // For initialize requests, get from request params + for message in messages where isInitializeRequest(message) { + if let params = message["params"] as? [String: Any], + let version = params["protocolVersion"] as? String + { + return version + } + } + + // For other requests, get from header + return request.header(HTTPHeader.protocolVersion) ?? Version.defaultNegotiated + } + + private func formatSSEEvent(data: Data, eventId: String?) -> Data { + Self.formatSSEEventStatic(data: data, eventId: eventId) + } + + /// Static version of formatSSEEvent for use in Sendable closures + private static func formatSSEEventStatic(data: Data, eventId: String?) -> Data { + var event = "event: message\n" + if let eventId { + event += "id: \(eventId)\n" + } + if let jsonString = String(data: data, encoding: .utf8) { + event += "data: \(jsonString)\n\n" + } + return Data(event.utf8) + } + + private func writePrimingEvent( + streamId: String, + continuation: AsyncThrowingStream.Continuation, + protocolVersion: String + ) async { + // Only write priming events if event store is configured + guard let eventStore = options.eventStore else { return } + + // Priming events have empty data which older clients cannot handle + // Only send to clients with protocol version >= 2025-11-25 + guard protocolVersion >= Version.v2025_11_25 else { return } + + do { + let primingEventId = try await eventStore.storeEvent(streamId: streamId, message: Data()) + + var primingEvent = "id: \(primingEventId)\n" + if let retryInterval = options.retryInterval { + primingEvent += "retry: \(retryInterval)\n" + } + primingEvent += "data: \n\n" + + continuation.yield(Data(primingEvent.utf8)) + } catch { + logger.error("Failed to write priming event: \(error)") + } + } + + private func replayEvents(lastEventId: String, eventStore: EventStore, request: HTTPRequest) async -> HTTPResponse { + // Get stream ID for this event + guard let streamId = await eventStore.streamIdForEventId(lastEventId) else { + return createJsonErrorResponse( + status: 400, + code: ErrorCode.connectionClosed, + message: "Invalid event ID format" + ) + } + + // Check for conflict + if streamMapping[streamId] != nil { + return createJsonErrorResponse( + status: 409, + code: ErrorCode.connectionClosed, + message: "Conflict: Stream already has an active connection" + ) + } + + let (stream, streamContinuation) = AsyncThrowingStream.makeStream() + + // Capture continuation by value for Sendable closure + let capturedContinuation = streamContinuation + + do { + // Replay events - use static method for SSE formatting + let replayedStreamId = try await eventStore.replayEventsAfter(lastEventId) { eventId, message in + let sseData = Self.formatSSEEventStatic(data: message, eventId: eventId) + capturedContinuation.yield(sseData) + } + + // Clean up mapping when stream terminates (e.g., client disconnect) + streamContinuation.onTermination = { @Sendable [weak self] _ in + Task { await self?.cleanUpStreamMapping(for: replayedStreamId) } + } + + let cleanup: @Sendable () -> Void = { + capturedContinuation.finish() + } + + streamMapping[replayedStreamId] = StreamState( + continuation: streamContinuation, + cleanup: cleanup + ) + + // Write a new priming event after replay so clients can resume again + // if they disconnect during this stream. Use the replayed stream ID. + let protocolVersion = negotiatedProtocolVersion ?? request.header(HTTPHeader.protocolVersion) ?? Version.defaultNegotiated + await writePrimingEvent(streamId: replayedStreamId, continuation: streamContinuation, protocolVersion: protocolVersion) + + var headers = sessionHeaders() + headers[HTTPHeader.contentType] = "text/event-stream" + headers[HTTPHeader.cacheControl] = "no-cache, no-transform" + headers[HTTPHeader.connection] = "keep-alive" + + return HTTPResponse(statusCode: 200, headers: headers, stream: stream) + } catch { + logger.error("Error replaying events: \(error)") + streamContinuation.finish() + return createJsonErrorResponse( + status: 500, + code: ErrorCode.connectionClosed, + message: "Error replaying events" + ) + } + } + + private func checkBatchCompletion(streamId: String, jsonState: JsonStreamState) async throws { + // Find all request IDs using this stream + let relatedIds = requestToStreamMapping.filter { $0.value == streamId }.map { $0.key } + + // Check if all requests have responses + let allComplete = relatedIds.allSatisfy { requestResponseMap[$0] != nil } + + guard allComplete else { return } + + // Gather responses + let responses = relatedIds.compactMap { requestResponseMap[$0] } + + // Build JSON response + var responseData: Data + if responses.count == 1, let singleResponse = responses.first { + responseData = singleResponse + } else { + // Combine into array + var jsonArray = Data("[".utf8) + for (index, response) in responses.enumerated() { + if index > 0 { + jsonArray.append(contentsOf: ",".utf8) + } + jsonArray.append(response) + } + jsonArray.append(contentsOf: "]".utf8) + responseData = jsonArray + } + + // Clean up + for id in relatedIds { + requestResponseMap.removeValue(forKey: id) + requestToStreamMapping.removeValue(forKey: id) + } + jsonStreamMapping.removeValue(forKey: streamId) + + var headers = sessionHeaders() + headers[HTTPHeader.contentType] = "application/json" + + // Yield the response to the stream and finish + jsonState.continuation.yield(HTTPResponse(statusCode: 200, headers: headers, body: responseData)) + jsonState.continuation.finish() + } + + private func checkStreamCompletion(streamId: String) async throws { + // Find all request IDs using this stream + let relatedIds = requestToStreamMapping.filter { $0.value == streamId }.map { $0.key } + + // Check if all requests have responses + let allComplete = relatedIds.allSatisfy { requestResponseMap[$0] != nil } + + guard allComplete else { return } + + // Close the stream + if let state = streamMapping[streamId] { + state.cleanup() + } + + // Clean up + for id in relatedIds { + requestResponseMap.removeValue(forKey: id) + requestToStreamMapping.removeValue(forKey: id) + } + streamMapping.removeValue(forKey: streamId) + } + + private func createJsonErrorResponse( + status: Int, + code: Int, + message: String, + extraHeaders: [String: String] = [:] + ) -> HTTPResponse { + let error: [String: Any] = [ + "jsonrpc": "2.0", + "error": [ + "code": code, + "message": message, + ], + "id": NSNull(), + ] + + let body = (try? JSONSerialization.data(withJSONObject: error)) ?? Data() + + var headers = sessionHeaders() + headers[HTTPHeader.contentType] = "application/json" + for (key, value) in extraHeaders { + headers[key] = value + } + + return HTTPResponse(statusCode: status, headers: headers, body: body) + } +} diff --git a/Sources/MCP/Base/Transports/InMemoryEventStore.swift b/Sources/MCP/Base/Transports/InMemoryEventStore.swift new file mode 100644 index 00000000..74359554 --- /dev/null +++ b/Sources/MCP/Base/Transports/InMemoryEventStore.swift @@ -0,0 +1,277 @@ +import Foundation + +/// Simple in-memory implementation of the `EventStore` protocol for resumability support. +/// +/// This implementation is primarily intended for examples and testing. For production use, +/// consider implementing a persistent storage solution (e.g., using a database or cache). +/// +/// ## How It Works +/// +/// The `InMemoryEventStore` generates event IDs that encode the stream ID, allowing +/// events to be replayed for a specific stream when a client reconnects with a +/// `Last-Event-ID` header. +/// +/// Event ID format: `{streamId}_{timestamp}_{random}` +/// +/// ## Memory Management +/// +/// The event store automatically limits memory usage by keeping only the last N events +/// per stream (configurable via `maxEventsPerStream`). When a stream exceeds its limit, +/// the oldest events are automatically evicted. +/// +/// You can also manually manage memory using: +/// - `cleanup(olderThan:)`: Remove events older than a specified duration +/// - `removeEvents(forStream:)`: Remove all events for a specific stream +/// - `clear()`: Remove all events +/// +/// ## Example Usage +/// +/// ```swift +/// // Default: 100 events per stream +/// let eventStore = InMemoryEventStore() +/// +/// // Custom limit: 500 events per stream +/// let eventStore = InMemoryEventStore(maxEventsPerStream: 500) +/// +/// let transport = HTTPServerTransport( +/// options: .init( +/// sessionIdGenerator: { UUID().uuidString }, +/// eventStore: eventStore +/// ) +/// ) +/// ``` +/// +/// ## Limitations +/// +/// - **Not persistent**: Events are lost when the process restarts +/// - **Single process**: Cannot be shared across multiple server instances +/// +/// For production deployments, implement `EventStore` with a persistent backend like +/// Redis, PostgreSQL, or another appropriate storage system. +public actor InMemoryEventStore: EventStore { + /// Maximum number of events to keep per stream. + /// When a stream exceeds this limit, the oldest events are automatically evicted. + public nonisolated let maxEventsPerStream: Int + + /// Per-stream event storage, maintaining chronological order within each stream + private var streams: [String: [StoredEvent]] = [:] + + /// Event ID to event entry lookup for quick access + private var eventIndex: [String: StoredEvent] = [:] + + private struct StoredEvent { + let eventId: String + let streamId: String + let message: Data + let timestamp: Date + } + + /// Creates a new in-memory event store. + /// + /// - Parameter maxEventsPerStream: Maximum number of events to keep per stream. + /// When a stream exceeds this limit, the oldest events are automatically evicted. + /// Default is 100 events per stream. + public init(maxEventsPerStream: Int = 100) { + precondition(maxEventsPerStream > 0, "maxEventsPerStream must be positive") + self.maxEventsPerStream = maxEventsPerStream + } + + // MARK: - EventStore Protocol + + /// Stores an event and returns its unique ID. + /// + /// If the stream has reached its maximum event limit (`maxEventsPerStream`), + /// the oldest event in that stream is automatically evicted. + /// + /// - Parameters: + /// - streamId: The stream this event belongs to + /// - message: The JSON-RPC message data. Empty `Data()` indicates a priming event. + /// - Returns: A unique event ID for this event + public func storeEvent(streamId: String, message: Data) async throws -> String { + let eventId = generateEventId(streamId: streamId) + let event = StoredEvent( + eventId: eventId, + streamId: streamId, + message: message, + timestamp: Date() + ) + + // Get or create the event list for this stream + var streamEvents = streams[streamId] ?? [] + + // If stream is at capacity, evict the oldest event + if streamEvents.count >= maxEventsPerStream { + let oldestEvent = streamEvents.removeFirst() + eventIndex.removeValue(forKey: oldestEvent.eventId) + } + + // Add the new event + streamEvents.append(event) + streams[streamId] = streamEvents + eventIndex[eventId] = event + + return eventId + } + + /// Gets the stream ID associated with an event ID. + /// + /// - Parameter eventId: The event ID to look up + /// - Returns: The stream ID, or nil if not found + public func streamIdForEventId(_ eventId: String) async -> String? { + // Try to get from stored event first (fast O(1) lookup) + if let event = eventIndex[eventId] { + return event.streamId + } + // Fall back to parsing from event ID format + return extractStreamId(from: eventId) + } + + /// Replays events after the given event ID. + /// + /// Events are replayed in chronological order, only including events from the same stream. + /// Priming events (empty message data) are skipped during replay. + /// + /// - Parameters: + /// - lastEventId: The last event ID the client received + /// - send: Callback to send each replayed event (eventId, message) + /// - Returns: The stream ID for continued event delivery + /// - Throws: `EventStoreError.eventNotFound` if the event ID doesn't exist + public func replayEventsAfter( + _ lastEventId: String, + send: @escaping @Sendable (String, Data) async throws -> Void + ) async throws -> String { + // Look up the event in our index + guard let lastEvent = eventIndex[lastEventId] else { + throw EventStoreError.eventNotFound(lastEventId) + } + + let streamId = lastEvent.streamId + + // Get events for this stream + guard let streamEvents = streams[streamId] else { + // Stream exists in index but not in streams - should not happen + throw EventStoreError.eventNotFound(lastEventId) + } + + // Find the position of the last event and replay everything after it + var foundLastEvent = false + for event in streamEvents { + if foundLastEvent { + // Skip priming events (empty message data) + if !event.message.isEmpty { + try await send(event.eventId, event.message) + } + } else if event.eventId == lastEventId { + foundLastEvent = true + } + } + + return streamId + } + + // MARK: - Cleanup + + /// Removes events older than the specified duration. + /// + /// Call this periodically to remove stale events and free memory. + /// + /// - Parameter age: Events older than this duration will be removed + /// - Returns: The number of events removed + @discardableResult + public func cleanUp(olderThan age: Duration) -> Int { + let cutoff = Date().addingTimeInterval(-age.timeInterval) + var removed = 0 + + for (streamId, events) in streams { + var remaining: [StoredEvent] = [] + for event in events { + if event.timestamp < cutoff { + eventIndex.removeValue(forKey: event.eventId) + removed += 1 + } else { + remaining.append(event) + } + } + if remaining.isEmpty { + streams.removeValue(forKey: streamId) + } else { + streams[streamId] = remaining + } + } + + return removed + } + + /// Removes all events for a specific stream. + /// + /// - Parameter streamId: The stream ID whose events should be removed + /// - Returns: The number of events removed + @discardableResult + public func removeEvents(forStream streamId: String) -> Int { + guard let events = streams.removeValue(forKey: streamId) else { + return 0 + } + + for event in events { + eventIndex.removeValue(forKey: event.eventId) + } + + return events.count + } + + /// The total number of stored events across all streams. + public var eventCount: Int { + eventIndex.count + } + + /// The number of active streams. + public var streamCount: Int { + streams.count + } + + /// Removes all events from all streams. + public func clear() { + streams.removeAll() + eventIndex.removeAll() + } + + // MARK: - Private Helpers + + /// Generates a unique event ID that encodes the stream ID. + /// + /// Format: `{streamId}_{timestamp}_{random}` + private func generateEventId(streamId: String) -> String { + let timestamp = Int(Date().timeIntervalSince1970 * 1000) + let random = String(format: "%08x", UInt32.random(in: 0...UInt32.max)) + return "\(streamId)_\(timestamp)_\(random)" + } + + /// Extracts the stream ID from an event ID. + /// + /// Handles event IDs in format: `{streamId}_{timestamp}_{random}` + private func extractStreamId(from eventId: String) -> String? { + // Find the last two underscores (timestamp and random parts) + let parts = eventId.split(separator: "_", omittingEmptySubsequences: false) + guard parts.count >= 3 else { return nil } + + // Reconstruct stream ID (everything before the last two parts) + // This handles stream IDs that contain underscores + let streamIdParts = parts.dropLast(2) + guard !streamIdParts.isEmpty else { return nil } + + return streamIdParts.joined(separator: "_") + } +} + +/// Errors that can occur when working with the event store. +public enum EventStoreError: Error, CustomStringConvertible { + /// The specified event ID was not found. + case eventNotFound(String) + + public var description: String { + switch self { + case .eventNotFound(let eventId): + return "Event not found: \(eventId)" + } + } +} diff --git a/Sources/MCP/Base/Transports/NetworkTransport.swift b/Sources/MCP/Base/Transports/NetworkTransport.swift index b34af418..4a27cd18 100644 --- a/Sources/MCP/Base/Transports/NetworkTransport.swift +++ b/Sources/MCP/Base/Transports/NetworkTransport.swift @@ -242,9 +242,6 @@ import Logging private let messageStream: AsyncThrowingStream private let messageContinuation: AsyncThrowingStream.Continuation - // Track connection state for continuations - private var connectionContinuationResumed = false - // Connection is marked nonisolated(unsafe) to allow access from closures private nonisolated(unsafe) var connection: NetworkConnectionProtocol @@ -315,70 +312,67 @@ import Logging // Reset state for fresh connection isStopping = false - reconnectAttempt = 0 - // Reset continuation state - connectionContinuationResumed = false + try await waitForConnectionReady() + isConnected = true + logger.debug("Network transport connected successfully") - // Wait for connection to be ready - try await withCheckedThrowingContinuation { - [weak self] (continuation: CheckedContinuation) in - guard let self = self else { - continuation.resume(throwing: MCPError.internalError("Transport deallocated")) - return - } + // Start the receive loop after connection is established + Task { await receiveLoop() } - connection.stateUpdateHandler = { [weak self] state in - guard let self = self else { return } - - Task { @MainActor in - switch state { - case .ready: - await self.handleConnectionReady(continuation: continuation) - case .failed(let error): - await self.handleConnectionFailed( - error: error, continuation: continuation) - case .cancelled: - await self.handleConnectionCancelled(continuation: continuation) - case .waiting(let error): - self.logger.debug("Connection waiting: \(error)") - case .preparing: - self.logger.debug("Connection preparing...") - case .setup: - self.logger.debug("Connection setup...") - @unknown default: - self.logger.warning("Unknown connection state") - } + // Start heartbeat task if enabled + if heartbeatConfig.enabled { + startHeartbeat() + } + } + + /// Waits for the connection to reach ready state using AsyncStream + /// + /// This safely bridges NWConnection's callback-based state updates to async/await. + /// The stream finishes on terminal states (ready, failed, cancelled) to ensure + /// the continuation is resumed exactly once. + /// + /// - Throws: Error if the connection fails or is cancelled + private func waitForConnectionReady() async throws { + let stateStream = AsyncStream { continuation in + connection.stateUpdateHandler = { state in + continuation.yield(state) + // Finish stream on terminal states + switch state { + case .ready, .failed, .cancelled: + continuation.finish() + default: + break } } - connection.start(queue: .main) + continuation.onTermination = { [weak self] _ in + self?.connection.stateUpdateHandler = nil + } } - } - /// Handles when the connection reaches the ready state - /// - /// - Parameter continuation: The continuation to resume when connection is ready - private func handleConnectionReady(continuation: CheckedContinuation) - async - { - if !connectionContinuationResumed { - connectionContinuationResumed = true - isConnected = true - - // Reset reconnect attempt counter on successful connection - reconnectAttempt = 0 - logger.debug("Network transport connected successfully") - continuation.resume() - - // Start the receive loop after connection is established - Task { await self.receiveLoop() } - - // Start heartbeat task if enabled - if heartbeatConfig.enabled { - startHeartbeat() + connection.start(queue: .main) + + for await state in stateStream { + switch state { + case .ready: + return + case .failed(let error): + throw error + case .cancelled: + throw MCPError.internalError("Connection cancelled") + case .waiting(let error): + logger.debug("Connection waiting: \(error)") + case .preparing: + logger.debug("Connection preparing...") + case .setup: + logger.debug("Connection setup...") + @unknown default: + logger.warning("Unknown connection state") } } + + throw MCPError.internalError("Connection stream ended unexpectedly") } /// Starts a task to periodically send heartbeats to check connection health @@ -443,87 +437,6 @@ import Logging logger.trace("Heartbeat sent") } - /// Handles connection failure - /// - /// - Parameters: - /// - error: The error that caused the connection to fail - /// - continuation: The continuation to resume with the error - private func handleConnectionFailed( - error: Swift.Error, continuation: CheckedContinuation - ) async { - if !connectionContinuationResumed { - connectionContinuationResumed = true - logger.error("Connection failed: \(error)") - - await handleReconnection( - error: error, - continuation: continuation, - context: "failure" - ) - } - } - - /// Handles connection cancellation - /// - /// - Parameter continuation: The continuation to resume with cancellation error - private func handleConnectionCancelled(continuation: CheckedContinuation) - async - { - if !connectionContinuationResumed { - connectionContinuationResumed = true - logger.warning("Connection cancelled") - - await handleReconnection( - error: MCPError.internalError("Connection cancelled"), - continuation: continuation, - context: "cancellation" - ) - } - } - - /// Common reconnection handling logic - /// - /// - Parameters: - /// - error: The error that triggered the reconnection - /// - continuation: The continuation to resume with the error - /// - context: The context of the reconnection (for logging) - private func handleReconnection( - error: Swift.Error, - continuation: CheckedContinuation, - context: String - ) async { - if !isStopping, - reconnectionConfig.enabled, - reconnectAttempt < reconnectionConfig.maxAttempts - { - // Try to reconnect with exponential backoff - reconnectAttempt += 1 - logger.debug( - "Attempting reconnection after \(context) (\(reconnectAttempt)/\(reconnectionConfig.maxAttempts))..." - ) - - // Calculate backoff delay - let delay = reconnectionConfig.backoffDelay(for: reconnectAttempt) - - // Schedule reconnection attempt after delay - Task { - try? await Task.sleep(for: .seconds(delay)) - if !isStopping { - // Cancel the current connection before attempting to reconnect. - self.connection.cancel() - // Resume original continuation with error; outer logic or a new call to connect() will handle retry. - continuation.resume(throwing: error) - } else { - continuation.resume(throwing: error) // Stopping, so fail. - } - } - } else { - // Not configured to reconnect, exceeded max attempts, or stopping - self.connection.cancel() // Ensure connection is cancelled - continuation.resume(throwing: error) - } - } - /// Disconnects from the transport /// /// This cancels the NWConnection, finalizes the message stream, @@ -560,9 +473,6 @@ import Logging var messageWithNewline = message messageWithNewline.append(UInt8(ascii: "\n")) - // Use a local actor-isolated variable to track continuation state - var sendContinuationResumed = false - try await withCheckedThrowingContinuation { [weak self] (continuation: CheckedContinuation) in guard let self = self else { @@ -577,47 +487,25 @@ import Logging completion: .contentProcessed { [weak self] error in guard let self = self else { return } - Task { @MainActor in - if !sendContinuationResumed { - sendContinuationResumed = true - if let error = error { - self.logger.error("Send error: \(error)") - - // Check if we should attempt to reconnect on send failure - let isStopping = await self.isStopping // Await actor-isolated property - if !isStopping && self.reconnectionConfig.enabled { - let isConnected = await self.isConnected - if isConnected { - if error.isConnectionLost { - self.logger.warning( - "Connection appears broken, will attempt to reconnect..." - ) - - // Schedule connection restart - Task { [weak self] in // Operate on self's executor - guard let self = self else { return } - - await self.setIsConnected(false) - - try? await Task.sleep(for: .milliseconds(500)) - - let currentIsStopping = await self.isStopping - if !currentIsStopping { - // Cancel the connection, then attempt to reconnect fully. - self.connection.cancel() - try? await self.connect() - } - } - } - } - } + if let error = error { + self.logger.error("Send error: \(error)") - continuation.resume( - throwing: MCPError.internalError("Send error: \(error)")) - } else { - continuation.resume() + // Schedule reconnection attempt if connection lost + if error.isConnectionLost && self.reconnectionConfig.enabled { + Task { + await self.setIsConnected(false) + try? await Task.sleep(for: .milliseconds(500)) + if await !self.isStopping { + self.connection.cancel() + try? await self.connect() + } } } + + continuation.resume( + throwing: MCPError.internalError("Send error: \(error)")) + } else { + continuation.resume() } }) } @@ -796,9 +684,7 @@ import Logging /// - Returns: The received data chunk /// - Throws: Network errors or transport failures private func receiveData() async throws -> Data { - var receiveContinuationResumed = false - - return try await withCheckedThrowingContinuation { + try await withCheckedThrowingContinuation { [weak self] (continuation: CheckedContinuation) in guard let self = self else { continuation.resume(throwing: MCPError.internalError("Transport deallocated")) @@ -807,22 +693,17 @@ import Logging let maxLength = bufferConfig.maxReceiveBufferSize ?? Int.max connection.receive(minimumIncompleteLength: 1, maximumLength: maxLength) { - content, _, isComplete, error in - Task { @MainActor in - if !receiveContinuationResumed { - receiveContinuationResumed = true - if let error = error { - continuation.resume(throwing: MCPError.transportError(error)) - } else if let content = content { - continuation.resume(returning: content) - } else if isComplete { - self.logger.trace("Connection completed by peer") - continuation.resume(throwing: MCPError.connectionClosed) - } else { - // EOF: Resume with empty data instead of throwing an error - continuation.resume(returning: Data()) - } - } + [weak self] content, _, isComplete, error in + if let error = error { + continuation.resume(throwing: MCPError.transportError(error)) + } else if let content = content { + continuation.resume(returning: content) + } else if isComplete { + self?.logger.trace("Connection completed by peer") + continuation.resume(throwing: MCPError.connectionClosed) + } else { + // EOF: Resume with empty data instead of throwing an error + continuation.resume(returning: Data()) } } } diff --git a/Sources/MCP/Base/Transports/StdioTransport.swift b/Sources/MCP/Base/Transports/StdioTransport.swift index 84bfd93a..b6ae1fae 100644 --- a/Sources/MCP/Base/Transports/StdioTransport.swift +++ b/Sources/MCP/Base/Transports/StdioTransport.swift @@ -20,7 +20,7 @@ import struct Foundation.Data #if canImport(Darwin) || canImport(Glibc) || canImport(Musl) /// An implementation of the MCP stdio transport protocol. /// - /// This transport implements the [stdio transport](https://modelcontextprotocol.io/specification/2025-03-26/basic/transports#stdio) + /// This transport implements the [stdio transport](https://modelcontextprotocol.io/specification/2025-11-25/basic/transports#stdio) /// specification from the Model Context Protocol. /// /// The stdio transport works by: @@ -133,6 +133,14 @@ import struct Foundation.Data /// This method runs in the background while the transport is connected, /// parsing complete messages delimited by newlines and yielding them /// to the message stream. + /// + /// - Note: This implementation uses synchronous `FileDescriptor.read()` with + /// non-blocking mode and EAGAIN retry. This works but requires polling. + /// + /// - TODO: Consider refactoring to use proper async I/O (e.g., Dispatch I/O or + /// DispatchSource for read/write events) to eliminate polling. This would + /// remove the need for EAGAIN handling and improve efficiency. The same + /// applies to the `send()` method's write loop. private func readLoop() async { let bufferSize = 4096 var buffer = [UInt8](repeating: 0, count: bufferSize) @@ -151,11 +159,16 @@ import struct Foundation.Data pendingData.append(Data(buffer[.. + ) throws -> [String: Value]? { + let container = try decoder.container(keyedBy: AnyCodingKey.self) + var extra: [String: Value] = [:] + + for key in container.allKeys where !knownKeys.contains(key.stringValue) { + if let value = try? container.decode(Value.self, forKey: key) { + extra[key.stringValue] = value + } + } + + return extra.isEmpty ? nil : extra + } +} + +/// Helpers for encoding extra fields to an encoder. +public enum ExtraFieldsEncoder { + /// Encodes extra fields to a keyed container. + /// + /// - Parameters: + /// - extraFields: The extra fields to encode (can be nil) + /// - encoder: The encoder to write to + public static func encode( + _ extraFields: [String: Value]?, + to encoder: Encoder + ) throws { + guard let extra = extraFields else { return } + var container = encoder.container(keyedBy: AnyCodingKey.self) + for (key, value) in extra { + try container.encode(value, forKey: AnyCodingKey(key)) + } + } +} + +// MARK: - Protocol for Types with Extra Fields + +// TODO: Consider using a Swift macro to further reduce code duplication. +// A macro like `@ExtraFieldsCodable` could auto-generate the full +// `init(from:)` and `encode(to:)` implementations, eliminating the need +// for conforming types to manually write custom Codable implementations. + +/// Protocol for result types that support forward-compatible extra fields. +/// +/// Conforming types gain helper methods for encoding/decoding extra fields. +/// +/// Usage: +/// ```swift +/// extension ListResources.Result: ResultWithExtraFields { +/// public typealias ResultCodingKeys = CodingKeys +/// } +/// +/// // In init(from decoder:): +/// extraFields = try Self.decodeExtraFields(from: decoder) +/// +/// // In encode(to encoder:): +/// try encodeExtraFields(to: encoder) +/// ``` +public protocol ResultWithExtraFields: Codable, Hashable, Sendable { + associatedtype ResultCodingKeys: CodingKey & CaseIterable & RawRepresentable + where ResultCodingKeys.RawValue == String + + /// Additional fields not defined in the schema (for forward compatibility). + var extraFields: [String: Value]? { get set } +} + +extension ResultWithExtraFields { + /// Decodes extra fields from the decoder (static, callable from init). + public static func decodeExtraFields(from decoder: Decoder) throws -> [String: Value]? { + try ExtraFieldsDecoder.decode( + from: decoder, + knownKeys: Set(ResultCodingKeys.allCases.map { $0.rawValue }) + ) + } + + /// Encodes extra fields to the encoder. + public func encodeExtraFields(to encoder: Encoder) throws { + try ExtraFieldsEncoder.encode(extraFields, to: encoder) + } +} diff --git a/Sources/MCP/Base/Utilities/Ping.swift b/Sources/MCP/Base/Utilities/Ping.swift index e526ad4e..b2131abe 100644 --- a/Sources/MCP/Base/Utilities/Ping.swift +++ b/Sources/MCP/Base/Utilities/Ping.swift @@ -2,4 +2,17 @@ /// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2024-11-05/basic/utilities/ping public enum Ping: Method { public static let name: String = "ping" + + public struct Parameters: NotRequired, Hashable, Codable, Sendable { + /// Request metadata including progress token. + public var _meta: RequestMeta? + + public init() { + self._meta = nil + } + + public init(_meta: RequestMeta?) { + self._meta = _meta + } + } } diff --git a/Sources/MCP/Base/Versioning.swift b/Sources/MCP/Base/Versioning.swift index 05c77a00..4ac0d447 100644 --- a/Sources/MCP/Base/Versioning.swift +++ b/Sources/MCP/Base/Versioning.swift @@ -4,16 +4,42 @@ import Foundation /// following the format YYYY-MM-DD, to indicate /// the last date backwards incompatible changes were made. /// -/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2025-03-26/ +/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2025-11-25/ public enum Version { - /// All protocol versions supported by this implementation, ordered from newest to oldest. - static let supported: Set = [ - "2025-03-26", - "2024-11-05", + // MARK: - Version Constants + + /// Protocol version 2025-11-25: Tasks, icons, URL elicitation, sampling tools, tool execution + public static let v2025_11_25 = "2025-11-25" + + /// Protocol version 2025-06-18: Elicitation, structured output, title fields, resource links + public static let v2025_06_18 = "2025-06-18" + + /// Protocol version 2025-03-26: JSON-RPC batching + public static let v2025_03_26 = "2025-03-26" + + /// Protocol version 2024-11-05: Initial stable release + public static let v2024_11_05 = "2024-11-05" + + // MARK: - Computed Properties + + /// All protocol versions supported by this implementation. + public static let supported: Set = [ + v2025_11_25, + v2025_06_18, + v2025_03_26, + v2024_11_05, ] /// The latest protocol version supported by this implementation. - public static let latest = supported.max()! + public static let latest = v2025_11_25 + + /// The default protocol version assumed when no `MCP-Protocol-Version` header is received. + /// + /// Per the spec: "For backwards compatibility, if the server does _not_ receive an + /// `MCP-Protocol-Version` header, and has no other way to identify the version - for example, + /// by relying on the protocol version negotiated during initialization - the server **SHOULD** + /// assume protocol version `2025-03-26`." + public static let defaultNegotiated = v2025_03_26 /// Negotiates the protocol version based on the client's request and server's capabilities. /// - Parameter clientRequestedVersion: The protocol version requested by the client. diff --git a/Sources/MCP/Client/Client+Tasks.swift b/Sources/MCP/Client/Client+Tasks.swift new file mode 100644 index 00000000..bfcefc87 --- /dev/null +++ b/Sources/MCP/Client/Client+Tasks.swift @@ -0,0 +1,256 @@ +import Foundation + +extension Client { + // MARK: - Tasks (Experimental) + // Note: These methods are internal. Access via client.experimental.* + + func getTask(taskId: String) async throws -> GetTask.Result { + try validateServerCapability(\.tasks, "Tasks") + let request = GetTask.request(.init(taskId: taskId)) + return try await send(request) + } + + func listTasks(cursor: String? = nil) async throws -> (tasks: [MCPTask], nextCursor: String?) { + try validateServerCapability(\.tasks, "Tasks") + let request: Request + if let cursor { + request = ListTasks.request(.init(cursor: cursor)) + } else { + request = ListTasks.request(.init()) + } + let result = try await send(request) + return (tasks: result.tasks, nextCursor: result.nextCursor) + } + + func cancelTask(taskId: String) async throws -> CancelTask.Result { + try validateServerCapability(\.tasks, "Tasks") + let request = CancelTask.request(.init(taskId: taskId)) + return try await send(request) + } + + func getTaskResult(taskId: String) async throws -> GetTaskPayload.Result { + try validateServerCapability(\.tasks, "Tasks") + let request = GetTaskPayload.request(.init(taskId: taskId)) + return try await send(request) + } + + /// Get the task result decoded as a specific type. + /// + /// This method retrieves the task result and decodes the `extraFields` as the specified type. + /// The `extraFields` contain the actual result payload (e.g., CallTool.Result fields). + func getTaskResultAs(taskId: String, type: T.Type) async throws -> T { + let result = try await getTaskResult(taskId: taskId) + + // The result's extraFields contain the actual result payload + // We need to encode them back to JSON and decode as the target type + guard let extraFields = result.extraFields else { + throw MCPError.invalidParams("Task result has no payload") + } + + // Convert extraFields to the target type + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + // Encode the extraFields as JSON + let jsonData = try encoder.encode(extraFields) + + // Decode as the target type + return try decoder.decode(T.self, from: jsonData) + } + + func callToolAsTask( + name: String, + arguments: [String: Value]? = nil, + ttl: Int? = nil + ) async throws -> CreateTaskResult { + try validateServerCapability(\.tasks, "Tasks") + try validateServerCapability(\.tools, "Tools") + + let taskMetadata = TaskMetadata(ttl: ttl) + let request = CallTool.request(.init( + name: name, + arguments: arguments, + task: taskMetadata + )) + + // The server should return CreateTaskResult for task-augmented requests + // We need to decode as CreateTaskResult instead of CallTool.Result + guard let connection = connection else { + throw MCPError.internalError("Client connection not initialized") + } + + let requestData = try encoder.encode(request) + + // Create stream for receiving the response + let (stream, continuation) = AsyncThrowingStream.makeStream() + + let requestId = request.id + continuation.onTermination = { @Sendable [weak self] _ in + Task { await self?.cleanUpPendingRequest(id: requestId) } + } + + addPendingRequest(id: request.id, continuation: continuation) + + do { + try await connection.send(requestData) + } catch { + if removePendingRequest(id: request.id) != nil { + continuation.finish(throwing: error) + } + throw error + } + + for try await result in stream { + return result + } + + throw MCPError.internalError("No response received") + } + + func pollTask(taskId: String) -> AsyncThrowingStream { + AsyncThrowingStream { continuation in + let pollingTask = Task { + do { + while !Task.isCancelled { + let task = try await self.getTask(taskId: taskId) + continuation.yield(task) + + if isTerminalStatus(task.status) { + continuation.finish() + return + } + + // Wait based on pollInterval (default 1 second) + let intervalMs = task.pollInterval ?? 1000 + try await Task.sleep(for: .milliseconds(intervalMs)) + } + // Task was cancelled + continuation.finish(throwing: CancellationError()) + } catch { + continuation.finish(throwing: error) + } + } + + // Cancel the polling task when the stream is terminated + continuation.onTermination = { _ in + pollingTask.cancel() + } + } + } + + func pollUntilTerminal(taskId: String) async throws -> GetTask.Result { + for try await status in pollTask(taskId: taskId) { + if isTerminalStatus(status.status) { + return status + } + } + // This shouldn't happen, but handle it gracefully + throw MCPError.internalError("Task polling ended unexpectedly") + } + + func callToolAsTaskAndWait( + name: String, + arguments: [String: Value]? = nil, + ttl: Int? = nil + ) async throws -> (content: [Tool.Content], isError: Bool?) { + // Start the task + let createResult = try await callToolAsTask(name: name, arguments: arguments, ttl: ttl) + let taskId = createResult.task.taskId + + // Wait for the result (uses blocking getTaskResult) + let payloadResult = try await getTaskResult(taskId: taskId) + + // Decode the result as CallTool.Result + // Per MCP spec, the result fields are flattened directly in the response (via extraFields) + guard let extraFields = payloadResult.extraFields else { + throw MCPError.internalError("Task completed but no result available") + } + + // Convert extraFields back to Value for decoding + let resultValue = Value.object(extraFields) + let resultData = try encoder.encode(resultValue) + let toolResult = try decoder.decode(CallTool.Result.self, from: resultData) + return (content: toolResult.content, isError: toolResult.isError) + } + + func callToolStream( + name: String, + arguments: [String: Value]? = nil, + ttl: Int? = nil + ) -> AsyncThrowingStream { + AsyncThrowingStream { continuation in + let streamTask = Task { + do { + // Step 1: Create the task + let createResult = try await self.callToolAsTask(name: name, arguments: arguments, ttl: ttl) + let task = createResult.task + continuation.yield(.taskCreated(task)) + + // Step 2: Poll for status updates until terminal + var lastStatus = task.status + var finalTask = task + + while !isTerminalStatus(lastStatus) { + // Wait based on pollInterval (default 1 second) + let intervalMs = finalTask.pollInterval ?? 1000 + try await Task.sleep(for: .milliseconds(intervalMs)) + + // Get updated status + let statusResult = try await self.getTask(taskId: task.taskId) + finalTask = MCPTask( + taskId: statusResult.taskId, + status: statusResult.status, + ttl: statusResult.ttl, + createdAt: statusResult.createdAt, + lastUpdatedAt: statusResult.lastUpdatedAt, + pollInterval: statusResult.pollInterval, + statusMessage: statusResult.statusMessage + ) + + // Only yield if status or message changed + if statusResult.status != lastStatus || statusResult.statusMessage != nil { + continuation.yield(.taskStatus(finalTask)) + } + lastStatus = statusResult.status + } + + // Step 3: Get the final result + if finalTask.status == .completed { + let payloadResult = try await self.getTaskResult(taskId: task.taskId) + + // Decode the result as CallTool.Result + if let extraFields = payloadResult.extraFields { + let resultValue = Value.object(extraFields) + let resultData = try self.encoder.encode(resultValue) + let toolResult = try self.decoder.decode(CallTool.Result.self, from: resultData) + continuation.yield(.result(toolResult)) + } else { + // No result available - return empty result + continuation.yield(.result(CallTool.Result(content: []))) + } + } else if finalTask.status == .failed { + let error = MCPError.internalError(finalTask.statusMessage ?? "Task failed") + continuation.yield(.error(error)) + } else if finalTask.status == .cancelled { + let error = MCPError.internalError("Task was cancelled") + continuation.yield(.error(error)) + } + + continuation.finish() + } catch let error as MCPError { + continuation.yield(.error(error)) + continuation.finish() + } catch { + let mcpError = MCPError.internalError(error.localizedDescription) + continuation.yield(.error(mcpError)) + continuation.finish() + } + } + + // Cancel the stream task if the stream is terminated + continuation.onTermination = { _ in + streamTask.cancel() + } + } + } +} diff --git a/Sources/MCP/Client/Client.swift b/Sources/MCP/Client/Client.swift index a2deb3dc..9f11a3af 100644 --- a/Sources/MCP/Client/Client.swift +++ b/Sources/MCP/Client/Client.swift @@ -36,10 +36,30 @@ public actor Client { public var name: String /// The client version public var version: String + /// A human-readable title for the client, intended for UI display. + /// If not provided, the `name` should be used for display. + public var title: String? + /// An optional human-readable description of what this implementation does. + public var description: String? + /// Optional icons representing this implementation. + public var icons: [Icon]? + /// An optional URL of the website for this implementation. + public var websiteUrl: String? - public init(name: String, version: String) { + public init( + name: String, + version: String, + title: String? = nil, + description: String? = nil, + icons: [Icon]? = nil, + websiteUrl: String? = nil + ) { self.name = name self.version = version + self.title = title + self.description = description + self.icons = icons + self.websiteUrl = websiteUrl } } @@ -57,24 +77,198 @@ public actor Client { /// The sampling capabilities public struct Sampling: Hashable, Codable, Sendable { - public init() {} + /// Context capability for sampling requests. + /// + /// When declared, indicates the client supports the `includeContext` parameter + /// with values "thisServer" and "allServers". If not declared, servers should + /// only use `includeContext: "none"` (or omit it). + public struct Context: Hashable, Codable, Sendable { + public init() {} + } + + /// Tools capability for sampling requests + public struct Tools: Hashable, Codable, Sendable { + public init() {} + } + + /// Whether the client supports includeContext parameter + public var context: Context? + /// Whether the client supports tools in sampling requests + public var tools: Tools? + + public init(context: Context? = nil, tools: Tools? = nil) { + self.context = context + self.tools = tools + } + } + + /// The elicitation capabilities + public struct Elicitation: Hashable, Codable, Sendable { + /// Form mode capabilities + public struct Form: Hashable, Codable, Sendable { + /// Whether the client applies schema defaults to missing fields. + public var applyDefaults: Bool? + + public init(applyDefaults: Bool? = nil) { + self.applyDefaults = applyDefaults + } + } + + /// URL mode capabilities (for out-of-band flows like OAuth) + public struct URL: Hashable, Codable, Sendable { + public init() {} + } + + /// Form mode capabilities + public var form: Form? + /// URL mode capabilities + public var url: URL? + + public init(form: Form? = nil, url: URL? = nil) { + self.form = form + self.url = url + } } /// Whether the client supports sampling public var sampling: Sampling? - /// Experimental features supported by the client - public var experimental: [String: String]? + /// Whether the client supports elicitation (user input requests) + public var elicitation: Elicitation? + /// Experimental, non-standard capabilities that the client supports. + public var experimental: [String: [String: Value]]? /// Whether the client supports roots public var roots: Capabilities.Roots? + /// Task capabilities (experimental, for bidirectional task support) + public var tasks: Tasks? public init( sampling: Sampling? = nil, - experimental: [String: String]? = nil, - roots: Capabilities.Roots? = nil + elicitation: Elicitation? = nil, + experimental: [String: [String: Value]]? = nil, + roots: Capabilities.Roots? = nil, + tasks: Tasks? = nil ) { self.sampling = sampling + self.elicitation = elicitation self.experimental = experimental self.roots = roots + self.tasks = tasks + } + } + + /// Context provided to client request handlers. + /// + /// This context is passed to handlers for server→client requests (e.g., sampling, + /// elicitation, roots) and provides: + /// - Cancellation checking via `isCancelled` and `checkCancellation()` + /// - Notification sending to the server + /// - Progress reporting convenience methods + /// + /// ## Example + /// + /// ```swift + /// client.withRequestHandler(CreateSamplingMessage.self) { params, context in + /// // Check for cancellation periodically + /// try context.checkCancellation() + /// + /// // Report progress back to server + /// try await context.sendProgressNotification( + /// token: progressToken, + /// progress: 50.0, + /// total: 100.0, + /// message: "Processing..." + /// ) + /// + /// return result + /// } + /// ``` + public struct RequestHandlerContext: Sendable { + /// Send a notification to the server. + /// + /// Use this to send notifications from within a request handler. + let sendNotification: @Sendable (any NotificationMessageProtocol) async throws -> Void + + // MARK: - Convenience Methods + + /// Send a progress notification to the server. + /// + /// Use this to report progress on long-running operations initiated by + /// server→client requests. + /// + /// - Parameters: + /// - token: The progress token from the request's `_meta.progressToken` + /// - progress: The current progress value (should increase monotonically) + /// - total: The total progress value, if known + /// - message: An optional human-readable message describing current progress + public func sendProgressNotification( + token: ProgressToken, + progress: Double, + total: Double? = nil, + message: String? = nil + ) async throws { + try await sendNotification(ProgressNotification.message(.init( + progressToken: token, + progress: progress, + total: total, + message: message + ))) + } + + // MARK: - Cancellation Checking + + /// Whether the request has been cancelled. + /// + /// Check this property periodically during long-running operations + /// to respond to cancellation requests from the server. + /// + /// This returns `true` when: + /// - The server sends a `CancelledNotification` for this request + /// - The client is disconnecting + /// + /// When cancelled, the handler should clean up resources and return + /// or throw an error. Per MCP spec, responses are not sent for cancelled requests. + /// + /// ## Example + /// + /// ```swift + /// client.withRequestHandler(CreateSamplingMessage.self) { params, context in + /// for chunk in largeInput { + /// // Check cancellation periodically + /// guard !context.isCancelled else { + /// throw CancellationError() + /// } + /// try await process(chunk) + /// } + /// return result + /// } + /// ``` + public var isCancelled: Bool { + Task.isCancelled + } + + /// Check if the request has been cancelled and throw if so. + /// + /// Call this method periodically during long-running operations. + /// If the request has been cancelled, this throws `CancellationError`. + /// + /// This is equivalent to checking `isCancelled` and throwing manually, + /// but provides a more idiomatic Swift concurrency pattern. + /// + /// ## Example + /// + /// ```swift + /// client.withRequestHandler(CreateSamplingMessage.self) { params, context in + /// for chunk in largeInput { + /// try context.checkCancellation() // Throws if cancelled + /// try await process(chunk) + /// } + /// return result + /// } + /// ``` + /// + /// - Throws: `CancellationError` if the request has been cancelled. + public func checkCancellation() throws { + try Task.checkCancellation() } } @@ -99,6 +293,19 @@ public actor Client { /// The client configuration public var configuration: Configuration + /// Experimental APIs for tasks and other features. + /// + /// Access experimental features via this property: + /// ```swift + /// let result = try await client.experimental.tasks.callToolAsTask(name: "tool", arguments: [:]) + /// let status = try await client.experimental.tasks.getTask(result.task.taskId) + /// ``` + /// + /// - Warning: These APIs are experimental and may change without notice. + public var experimental: ExperimentalClientFeatures { + ExperimentalClientFeatures(client: self) + } + /// The server capabilities private var serverCapabilities: Server.Capabilities? /// The server version @@ -108,55 +315,197 @@ public actor Client { /// A dictionary of type-erased notification handlers, keyed by method name private var notificationHandlers: [String: [NotificationHandlerBox]] = [:] + /// A dictionary of type-erased request handlers for server→client requests, keyed by method name + private var requestHandlers: [String: ClientRequestHandlerBox] = [:] + /// Task-augmented sampling handler (called when request has `task` field) + private var taskAugmentedSamplingHandler: ExperimentalClientTaskHandlers.TaskAugmentedSamplingHandler? + /// Task-augmented elicitation handler (called when request has `task` field) + private var taskAugmentedElicitationHandler: ExperimentalClientTaskHandlers.TaskAugmentedElicitationHandler? /// The task for the message handling loop private var task: Task? + /// In-flight server request handler Tasks, tracked by request ID. + /// Used for protocol-level cancellation when CancelledNotification is received. + private var inFlightServerRequestTasks: [RequestId: Task] = [:] + /// An error indicating a type mismatch when decoding a pending request private struct TypeMismatchError: Swift.Error {} - /// A pending request with a continuation for the result - private struct PendingRequest { - let continuation: CheckedContinuation - } - - /// A type-erased pending request + /// A type-erased pending request using AsyncThrowingStream for cancellation-aware waiting. private struct AnyPendingRequest { - private let _resume: (Result) -> Void + private let _yield: (Result) -> Void + private let _finish: () -> Void - init(_ request: PendingRequest) { - _resume = { result in + init( + continuation: AsyncThrowingStream.Continuation + ) { + _yield = { result in switch result { case .success(let value): if let typedValue = value as? T { - request.continuation.resume(returning: typedValue) + continuation.yield(typedValue) + continuation.finish() } else if let value = value as? Value, let data = try? JSONEncoder().encode(value), let decoded = try? JSONDecoder().decode(T.self, from: data) { - request.continuation.resume(returning: decoded) + continuation.yield(decoded) + continuation.finish() } else { - request.continuation.resume(throwing: TypeMismatchError()) + continuation.finish(throwing: TypeMismatchError()) } case .failure(let error): - request.continuation.resume(throwing: error) + continuation.finish(throwing: error) } } + _finish = { + continuation.finish() + } } + func resume(returning value: Any) { - _resume(.success(value)) + _yield(.success(value)) } func resume(throwing error: Swift.Error) { - _resume(.failure(error)) + _yield(.failure(error)) + } + + func finish() { + _finish() } } /// A dictionary of type-erased pending requests, keyed by request ID - private var pendingRequests: [ID: AnyPendingRequest] = [:] + private var pendingRequests: [RequestId: AnyPendingRequest] = [:] + /// Progress callbacks for requests, keyed by progress token. + /// Used to invoke callbacks when progress notifications are received. + private var progressCallbacks: [ProgressToken: ProgressCallback] = [:] + /// Timeout controllers for requests with progress-aware timeouts. + /// Used to reset timeouts when progress notifications are received. + private var timeoutControllers: [ProgressToken: TimeoutController] = [:] + /// Mapping from request ID to progress token. + /// Used to detect task-augmented responses and keep progress handlers alive. + private var requestProgressTokens: [RequestId: ProgressToken] = [:] + /// Mapping from task ID to progress token. + /// Keeps progress handlers alive for task-augmented requests until the task completes. + /// Per MCP spec 2025-11-25: "For task-augmented requests, the progressToken provided + /// in the original request MUST continue to be used for progress notifications + /// throughout the task's lifetime, even after the CreateTaskResult has been returned." + private var taskProgressTokens: [String: ProgressToken] = [:] // Add reusable JSON encoder/decoder private let encoder = JSONEncoder() private let decoder = JSONDecoder() + /// Controls timeout behavior for a single request, supporting reset on progress. + /// + /// This actor manages the timeout state for requests that use `resetTimeoutOnProgress`. + /// When progress is received, calling `signalProgress()` resets the timeout clock. + private actor TimeoutController { + /// The per-interval timeout duration. + let timeout: Duration + /// Whether to reset timeout when progress is received. + let resetOnProgress: Bool + /// Maximum total time to wait regardless of progress. + let maxTotalTimeout: Duration? + /// The start time of the request (for maxTotalTimeout tracking). + let startTime: ContinuousClock.Instant + /// The current deadline (updated when progress is received). + private var deadline: ContinuousClock.Instant + /// Whether the controller has been cancelled. + private var isCancelled = false + /// Continuation for signaling progress. + private var progressContinuation: AsyncStream.Continuation? + + init(timeout: Duration, resetOnProgress: Bool, maxTotalTimeout: Duration?) { + self.timeout = timeout + self.resetOnProgress = resetOnProgress + self.maxTotalTimeout = maxTotalTimeout + self.startTime = ContinuousClock.now + self.deadline = ContinuousClock.now.advanced(by: timeout) + } + + /// Signal that progress was received, resetting the timeout. + func signalProgress() { + guard resetOnProgress, !isCancelled else { return } + deadline = ContinuousClock.now.advanced(by: timeout) + progressContinuation?.yield() + } + + /// Cancel the timeout controller. + func cancel() { + isCancelled = true + progressContinuation?.finish() + } + + /// Wait until the timeout expires. + /// + /// If `resetOnProgress` is true, the timeout resets each time `signalProgress()` is called. + /// If `maxTotalTimeout` is set, the wait will end when that limit is exceeded. + /// + /// - Throws: `MCPError.requestTimeout` when the timeout expires. + func waitForTimeout() async throws { + let clock = ContinuousClock() + + // Create a stream for progress signals + let (progressStream, continuation) = AsyncStream.makeStream() + self.progressContinuation = continuation + + while !isCancelled { + // Check maxTotalTimeout + if let maxTotal = maxTotalTimeout { + let elapsed = clock.now - startTime + if elapsed >= maxTotal { + throw MCPError.requestTimeout( + timeout: maxTotal, + message: "Request exceeded maximum total timeout" + ) + } + } + + // Calculate time until deadline + let now = clock.now + let timeUntilDeadline = deadline - now + + if timeUntilDeadline <= .zero { + throw MCPError.requestTimeout( + timeout: timeout, + message: "Request timed out" + ) + } + + // Wait for either timeout or progress signal + do { + try await withThrowingTaskGroup(of: Void.self) { group in + // Timeout task + group.addTask { + try await Task.sleep(for: timeUntilDeadline) + } + + // Progress signal task (if reset is enabled) + if resetOnProgress { + group.addTask { + for await _ in progressStream { + // Progress received, exit to recalculate deadline + return + } + } + } + + // Wait for whichever completes first + _ = try await group.next() + group.cancelAll() + } + } catch is CancellationError { + return // Task was cancelled, exit gracefully + } + + // If we get here after a progress signal, loop to recalculate deadline + // If we get here after timeout, the next iteration will throw + } + } + } + public init( name: String, version: String, @@ -167,6 +516,30 @@ public actor Client { self.configuration = configuration } + /// Set the client capabilities. + /// + /// This should be called before `connect()` to configure what capabilities + /// the client will advertise to the server during initialization. + /// + /// - Parameter capabilities: The capabilities to set. + public func setCapabilities(_ capabilities: Capabilities) { + self.capabilities = capabilities + } + + /// Returns the server capabilities received during initialization. + /// + /// Use this method to check what capabilities the server supports after + /// successfully connecting. This can be useful for: + /// - Conditionally enabling features based on server support + /// - Logging or debugging connection details + /// - Building adaptive clients that work with various server implementations + /// + /// - Returns: The server's capabilities, or `nil` if the client has not + /// been initialized yet (i.e., `connect()` has not been called or failed). + public func getServerCapabilities() -> Server.Capabilities? { + return serverCapabilities + } + /// Connect to the server using the given transport @discardableResult public func connect(transport: any Transport) async throws -> Initialize.Result { @@ -177,48 +550,81 @@ public actor Client { "Client connected", metadata: ["name": "\(name)", "version": "\(version)"]) // Start message handling loop + // + // The receive loop: + // - Calls receive() once to get the stream + // - Iterates until the stream ends or throws + // - Cleans up pending requests on exit + // + // EAGAIN is handled by the transport layer internally. task = Task { guard let connection = self.connection else { return } - repeat { - // Check for cancellation before starting the iteration - if Task.isCancelled { break } - do { - let stream = await connection.receive() - for try await data in stream { - if Task.isCancelled { break } // Check inside loop too - - // Attempt to decode data - // Try decoding as a batch response first - if let batchResponse = try? decoder.decode([AnyResponse].self, from: data) { - await handleBatchResponse(batchResponse) - } else if let response = try? decoder.decode(AnyResponse.self, from: data) { - await handleResponse(response) - } else if let message = try? decoder.decode(AnyMessage.self, from: data) { - await handleMessage(message) - } else { - var metadata: Logger.Metadata = [:] - if let string = String(data: data, encoding: .utf8) { - metadata["message"] = .string(string) + defer { + // When the receive loop exits unexpectedly (transport closed without + // disconnect() being called), clean up pending requests. + Task { + await self.cleanupPendingRequestsOnUnexpectedDisconnect() + } + } + + do { + let stream = await connection.receive() + for try await data in stream { + if Task.isCancelled { break } + + // Attempt to decode data + // Try decoding as a batch response first + if let batchResponse = try? decoder.decode([AnyResponse].self, from: data) { + await handleBatchResponse(batchResponse) + } else if let response = try? decoder.decode(AnyResponse.self, from: data) { + await handleResponse(response) + } else if let request = try? decoder.decode(AnyRequest.self, from: data) { + // Handle incoming request from server (bidirectional communication) + // Spawn in a separate task to avoid blocking the message loop. + // This allows client request handlers to make nested requests + // back to the server if needed. + let requestId = request.id + let handlerTask = Task { [weak self] in + guard let self else { return } + defer { + Task { await self.removeInFlightServerRequest(requestId) } } - await logger?.warning( - "Unexpected message received by client (not single/batch response or notification)", - metadata: metadata - ) + await self.handleIncomingRequest(request) + } + trackInFlightServerRequest(requestId, task: handlerTask) + } else if let message = try? decoder.decode(AnyMessage.self, from: data) { + await handleMessage(message) + } else { + var metadata: Logger.Metadata = [:] + if let string = String(data: data, encoding: .utf8) { + metadata["message"] = .string(string) } + await logger?.warning( + "Unexpected message received by client (not single/batch response, request, or notification)", + metadata: metadata + ) } - } catch let error where MCPError.isResourceTemporarilyUnavailable(error) { - try? await Task.sleep(for: .milliseconds(10)) - continue - } catch { - await logger?.error( - "Error in message handling loop", metadata: ["error": "\(error)"]) - break } - } while true + await logger?.debug("Client receive stream ended") + } catch { + await logger?.error( + "Error in message handling loop", metadata: ["error": "\(error)"]) + } await self.logger?.debug("Client message handling loop task is terminating.") } + // Register default handler for CancelledNotification (protocol-level cancellation) + _ = await onNotification(CancelledNotification.self) { [weak self] message in + guard let self else { return } + guard let requestId = message.params.requestId else { + // Per protocol 2025-11-25+, requestId is optional. + // If not provided, we cannot cancel a specific request. + return + } + await self.cancelInFlightServerRequest(requestId, reason: message.params.reason) + } + // Automatically initialize after connecting return try await _initialize() } @@ -227,6 +633,16 @@ public actor Client { public func disconnect() async { await logger?.debug("Initiating client disconnect...") + // Cancel all in-flight server request handlers + for (requestId, handlerTask) in inFlightServerRequestTasks { + handlerTask.cancel() + await logger?.debug( + "Cancelled in-flight server request during disconnect", + metadata: ["id": "\(requestId)"] + ) + } + inFlightServerRequestTasks.removeAll() + // Part 1: Inside actor - Grab state and clear internal references let taskToCancel = self.task let connectionToDisconnect = self.connection @@ -236,11 +652,17 @@ public actor Client { self.connection = nil self.pendingRequests = [:] // Use empty dictionary literal + // Clear all progress-related state + progressCallbacks.removeAll() + timeoutControllers.removeAll() + requestProgressTokens.removeAll() + taskProgressTokens.removeAll() + // Part 2: Outside actor - Resume continuations, disconnect transport, await task - // Resume continuations first + // Resume pending request continuations with connection closed error for (_, request) in pendingRequestsToCancel { - request.resume(throwing: MCPError.internalError("Client disconnected")) + request.resume(throwing: MCPError.connectionClosed) } await logger?.debug("Pending requests cancelled.") @@ -264,6 +686,54 @@ public actor Client { await logger?.debug("Client disconnect complete.") } + /// Cleans up pending requests when the receive loop exits unexpectedly. + /// + /// This is called from the receive loop's defer block when the transport closes + /// without `disconnect()` being called (e.g., server process exits). We only + /// clean up requests that haven't already been handled by `disconnect()`. + private func cleanupPendingRequestsOnUnexpectedDisconnect() async { + guard !pendingRequests.isEmpty else { return } + + await logger?.debug( + "Cleaning up pending requests after unexpected disconnect", + metadata: ["count": "\(pendingRequests.count)"]) + + for (_, request) in pendingRequests { + request.resume(throwing: MCPError.connectionClosed) + } + pendingRequests.removeAll() + } + + // MARK: - In-Flight Server Request Tracking (Protocol-Level Cancellation) + + /// Track an in-flight server request handler Task. + private func trackInFlightServerRequest(_ requestId: RequestId, task: Task) { + inFlightServerRequestTasks[requestId] = task + } + + /// Remove an in-flight server request handler Task. + private func removeInFlightServerRequest(_ requestId: RequestId) { + inFlightServerRequestTasks.removeValue(forKey: requestId) + } + + /// Cancel an in-flight server request handler Task. + /// + /// Called when a CancelledNotification is received for a specific requestId. + /// Per MCP spec, if the request is unknown or already completed, this is a no-op. + private func cancelInFlightServerRequest(_ requestId: RequestId, reason: String?) async { + if let task = inFlightServerRequestTasks[requestId] { + task.cancel() + await logger?.debug( + "Cancelled in-flight server request", + metadata: [ + "id": "\(requestId)", + "reason": "\(reason ?? "none")", + ] + ) + } + // Per spec: MAY ignore if request is unknown - no error needed + } + // MARK: - Registration /// Register a handler for a notification @@ -282,61 +752,682 @@ public actor Client { throw MCPError.internalError("Client connection not initialized") } - let notificationData = try encoder.encode(notification) - try await connection.send(notificationData) - } + let notificationData = try encoder.encode(notification) + try await connection.send(notificationData) + } + + /// Send a progress notification to the server. + /// + /// This is a convenience method for sending progress notifications from the client + /// to the server. This enables bidirectional progress reporting where clients can + /// inform servers about their own progress (e.g., during client-side processing). + /// + /// ## Example + /// + /// ```swift + /// // Client reports its own progress to the server + /// try await client.sendProgressNotification( + /// token: .string("client-task-123"), + /// progress: 50.0, + /// total: 100.0, + /// message: "Processing client-side data..." + /// ) + /// ``` + /// + /// - Parameters: + /// - token: The progress token to associate with this notification + /// - progress: The current progress value (should increase monotonically) + /// - total: The total progress value, if known + /// - message: An optional human-readable message describing current progress + public func sendProgressNotification( + token: ProgressToken, + progress: Double, + total: Double? = nil, + message: String? = nil + ) async throws { + try await notify(ProgressNotification.message(.init( + progressToken: token, + progress: progress, + total: total, + message: message + ))) + } + + /// Send a notification that the list of available roots has changed. + /// + /// Servers that receive this notification should request an updated + /// list of roots via the roots/list request. + /// + /// - Throws: `MCPError.invalidRequest` if the client has not declared + /// the `roots.listChanged` capability. + public func sendRootsChanged() async throws { + guard capabilities.roots?.listChanged == true else { + throw MCPError.invalidRequest( + "Client does not support roots.listChanged capability") + } + try await notify(RootsListChangedNotification.message(.init())) + } + + /// Register a handler for server→client requests. + /// + /// This enables bidirectional communication where the server can send requests + /// to the client (e.g., sampling, roots, elicitation). + /// + /// - Parameters: + /// - type: The method type to handle + /// - handler: The handler function that receives parameters and returns a result + /// - Returns: Self for chaining + @discardableResult + public func withRequestHandler( + _ type: M.Type, + handler: @escaping @Sendable (M.Parameters) async throws -> M.Result + ) -> Self { + requestHandlers[M.name] = TypedClientRequestHandler(handler) + return self + } + + /// Register a handler for `roots/list` requests from the server. + /// + /// When the server requests the list of roots, this handler will be called + /// to provide the available filesystem directories. + /// + /// - Important: The client must have declared `roots` capability during initialization. + /// + /// - Parameter handler: A closure that returns the list of available roots. + /// - Returns: Self for chaining. + /// - Precondition: `capabilities.roots` must be non-nil. + @discardableResult + public func withRootsHandler( + _ handler: @escaping @Sendable () async throws -> [Root] + ) -> Self { + precondition( + capabilities.roots != nil, + "Cannot register roots handler: Client does not have roots capability" + ) + return withRequestHandler(ListRoots.self) { _ in + ListRoots.Result(roots: try await handler()) + } + } + + /// Register a handler for `sampling/createMessage` requests from the server. + /// + /// When the server requests a sampling completion, this handler will be called + /// to generate the LLM response. + /// + /// The handler receives parameters that may or may not include tools. Check `params.hasTools` + /// to determine if tool use is enabled for this request. + /// + /// - Important: The client must have declared `sampling` capability during initialization. + /// + /// ## Example + /// + /// ```swift + /// client.withSamplingHandler { params in + /// // Call your LLM with the messages + /// let response = try await llm.complete( + /// messages: params.messages, + /// tools: params.tools, // May be nil + /// maxTokens: params.maxTokens + /// ) + /// + /// return ClientSamplingRequest.Result( + /// model: "gpt-4", + /// stopReason: .endTurn, + /// role: .assistant, + /// content: .text(response.text) + /// ) + /// } + /// ``` + /// + /// - Parameter handler: A closure that receives sampling parameters and returns the result. + /// - Returns: Self for chaining. + /// - Precondition: `capabilities.sampling` must be non-nil. + @discardableResult + public func withSamplingHandler( + _ handler: @escaping @Sendable (ClientSamplingRequest.Parameters) async throws -> ClientSamplingRequest.Result + ) -> Self { + precondition( + capabilities.sampling != nil, + "Cannot register sampling handler: Client does not have sampling capability" + ) + return withRequestHandler(ClientSamplingRequest.self, handler: handler) + } + + /// Register a handler for `elicitation/create` requests from the server. + /// + /// When the server requests user input via elicitation, this handler will be called + /// to collect the input and return the result. + /// + /// - Important: The client must have declared `elicitation` capability during initialization. + /// + /// - Parameter handler: A closure that receives elicitation parameters and returns the result. + /// - Returns: Self for chaining. + /// - Precondition: `capabilities.elicitation` must be non-nil. + @discardableResult + public func withElicitationHandler( + _ handler: @escaping @Sendable (Elicit.Parameters) async throws -> Elicit.Result + ) -> Self { + precondition( + capabilities.elicitation != nil, + "Cannot register elicitation handler: Client does not have elicitation capability" + ) + return withRequestHandler(Elicit.self, handler: handler) + } + + /// Internal method to set a request handler box directly. + /// + /// This is used by task-augmented handlers that need to return different result types + /// based on whether the request has a `task` field. + /// + /// - Important: This is an internal API that may change without notice. + internal func _setRequestHandler(method: String, handler: ClientRequestHandlerBox) { + requestHandlers[method] = handler + } + + /// Internal method to get an existing request handler box. + /// + /// This is used to retrieve the existing handler before wrapping it with + /// a task-aware handler that preserves the normal handler as a fallback. + /// + /// - Important: This is an internal API that may change without notice. + internal func _getRequestHandler(method: String) -> ClientRequestHandlerBox? { + requestHandlers[method] + } + + /// Internal method to set the task-augmented sampling handler. + /// + /// This handler is called when the server sends a `sampling/createMessage` request + /// with a `task` field. The handler should return `CreateTaskResult` instead of + /// the normal sampling result. + /// + /// - Important: This is an internal API that may change without notice. + internal func _setTaskAugmentedSamplingHandler( + _ handler: @escaping ExperimentalClientTaskHandlers.TaskAugmentedSamplingHandler + ) { + taskAugmentedSamplingHandler = handler + } + + /// Internal method to set the task-augmented elicitation handler. + /// + /// This handler is called when the server sends an `elicitation/create` request + /// with a `task` field. The handler should return `CreateTaskResult` instead of + /// the normal elicitation result. + /// + /// - Important: This is an internal API that may change without notice. + internal func _setTaskAugmentedElicitationHandler( + _ handler: @escaping ExperimentalClientTaskHandlers.TaskAugmentedElicitationHandler + ) { + taskAugmentedElicitationHandler = handler + } + + // MARK: - Request Options + + /// Options that can be given per request. + /// + /// Similar to TypeScript SDK's `RequestOptions`, this allows configuring + /// timeout behavior for individual requests, including progress-aware timeouts. + public struct RequestOptions: Sendable { + /// The default request timeout (60 seconds), matching TypeScript SDK. + public static let defaultTimeout: Duration = .seconds(60) + + /// A timeout for this request. + /// + /// If exceeded, the request will be cancelled and an `MCPError.requestTimeout` + /// will be thrown. A `CancelledNotification` will also be sent to the server. + /// + /// If `nil`, no timeout is applied (the request can wait indefinitely). + /// Default is `nil` to match existing behavior. + public var timeout: Duration? + + /// If `true`, receiving a progress notification resets the timeout clock. + /// + /// This is useful for long-running operations that send periodic progress updates. + /// As long as the server keeps sending progress, the request won't time out. + /// + /// When combined with `maxTotalTimeout`, this allows both: + /// - Per-interval timeout that resets on progress + /// - Overall hard limit that prevents infinite waiting + /// + /// Default is `false`. + /// + /// - Note: Only effective when `timeout` is set and the request uses `onProgress`. + public var resetTimeoutOnProgress: Bool + + /// Maximum total time to wait for the request, regardless of progress. + /// + /// When `resetTimeoutOnProgress` is `true`, this provides a hard upper limit + /// on the total wait time. Even if progress notifications keep arriving, + /// the request will be cancelled if this limit is exceeded. + /// + /// If `nil`, there's no maximum total timeout (only the regular `timeout` + /// applies, potentially reset by progress). + /// + /// - Note: Only effective when both `timeout` and `resetTimeoutOnProgress` are set. + public var maxTotalTimeout: Duration? + + /// Creates request options with the specified configuration. + /// + /// - Parameters: + /// - timeout: The timeout duration, or `nil` for no timeout. + /// - resetTimeoutOnProgress: Whether to reset the timeout when progress is received. + /// - maxTotalTimeout: Maximum total time to wait regardless of progress. + public init( + timeout: Duration? = nil, + resetTimeoutOnProgress: Bool = false, + maxTotalTimeout: Duration? = nil + ) { + self.timeout = timeout + self.resetTimeoutOnProgress = resetTimeoutOnProgress + self.maxTotalTimeout = maxTotalTimeout + } + + /// Request options with the default timeout (60 seconds). + public static let withDefaultTimeout = RequestOptions(timeout: defaultTimeout) + + /// Request options with no timeout. + public static let noTimeout = RequestOptions(timeout: nil) + } + + // MARK: - Requests + + /// Send a request and receive its response. + /// + /// This method sends a request without a timeout. For timeout support, + /// use `send(_:options:)` instead. + public func send(_ request: Request) async throws -> M.Result { + try await send(request, options: nil) + } + + /// Send a request and receive its response with options. + /// + /// - Parameters: + /// - request: The request to send. + /// - options: Options for this request, including timeout configuration. + /// - Returns: The response result. + /// - Throws: `MCPError.requestTimeout` if the timeout is exceeded. + public func send( + _ request: Request, + options: RequestOptions? + ) async throws -> M.Result { + guard let connection = connection else { + throw MCPError.internalError("Client connection not initialized") + } + + let requestData = try encoder.encode(request) + + // Create stream for receiving the response + let (stream, continuation) = AsyncThrowingStream.makeStream() + + // Track whether we've timed out (for the onTermination handler) + let requestId = request.id + let timeout = options?.timeout + + // Clean up pending request if caller cancels (e.g., task cancelled or timeout) + // and send CancelledNotification to server per MCP spec + continuation.onTermination = { @Sendable [weak self] termination in + Task { + guard let self else { return } + await self.cleanupPendingRequest(id: requestId) + + // Per MCP spec: send notifications/cancelled when cancelling a request + // Only send if the stream was cancelled (not finished normally) + if case .cancelled = termination { + let reason = if let timeout { + "Request timed out after \(timeout)" + } else { + "Client cancelled the request" + } + await self.sendCancellationNotification( + requestId: requestId, + reason: reason + ) + } + } + } + + // Add the pending request before attempting to send + addPendingRequest(id: request.id, continuation: continuation) + + // Send the request data + do { + try await connection.send(requestData) + } catch { + // If send fails, remove the pending request and rethrow + if removePendingRequest(id: request.id) != nil { + continuation.finish(throwing: error) + } + throw error + } + + // Wait for response with optional timeout + if let timeout { + // Use withTimeout pattern for cancellation-aware timeout + return try await withThrowingTaskGroup(of: M.Result.self) { group in + // Add the main task that waits for the response + group.addTask { + for try await result in stream { + return result + } + throw MCPError.internalError("No response received") + } + + // Add the timeout task + group.addTask { + try await Task.sleep(for: timeout) + throw MCPError.requestTimeout(timeout: timeout, message: "Request timed out") + } + + // Return whichever completes first + guard let result = try await group.next() else { + throw MCPError.internalError("No response received") + } + + // Cancel the other task + group.cancelAll() + + return result + } + } else { + // No timeout - wait indefinitely for response + for try await result in stream { + return result + } + + // Stream closed without yielding a response + throw MCPError.internalError("No response received") + } + } + + /// Send a request with a progress callback. + /// + /// This method automatically sets up progress tracking by: + /// 1. Generating a unique progress token based on the request ID + /// 2. Injecting the token into the request's `_meta.progressToken` + /// 3. Invoking the callback when progress notifications are received + /// + /// The callback is automatically cleaned up when the request completes. + /// + /// ## Example + /// + /// ```swift + /// let result = try await client.send( + /// CallTool.request(.init(name: "slow_operation", arguments: ["steps": 5])), + /// onProgress: { progress in + /// print("Progress: \(progress.value)/\(progress.total ?? 0) - \(progress.message ?? "")") + /// } + /// ) + /// ``` + /// + /// - Parameters: + /// - request: The request to send + /// - onProgress: A callback invoked when progress notifications are received + /// - Returns: The response result + public func send( + _ request: Request, + onProgress: @escaping ProgressCallback + ) async throws -> M.Result { + try await send(request, options: nil, onProgress: onProgress) + } + + /// Send a request with options and a progress callback. + /// + /// - Parameters: + /// - request: The request to send. + /// - options: Options for this request, including timeout configuration. + /// - onProgress: A callback invoked when progress notifications are received. + /// - Returns: The response result. + /// - Throws: `MCPError.requestTimeout` if the timeout is exceeded. + public func send( + _ request: Request, + options: RequestOptions?, + onProgress: @escaping ProgressCallback + ) async throws -> M.Result { + guard let connection = connection else { + throw MCPError.internalError("Client connection not initialized") + } + + // Generate a progress token from the request ID + let progressToken: ProgressToken = switch request.id { + case .number(let n): .integer(n) + case .string(let s): .string(s) + } + + // Encode the request, inject progressToken into _meta, then re-encode + let requestData = try encoder.encode(request) + var requestDict = try decoder.decode([String: Value].self, from: requestData) + + // Ensure params exists and inject _meta.progressToken + var params = requestDict["params"]?.objectValue ?? [:] + var meta = params["_meta"]?.objectValue ?? [:] + meta["progressToken"] = switch progressToken { + case .string(let s): .string(s) + case .integer(let n): .int(n) + } + params["_meta"] = .object(meta) + requestDict["params"] = .object(params) + + let modifiedRequestData = try encoder.encode(requestDict) + + // Register the progress callback and track the request → token mapping + // (used to detect task-augmented responses and keep progress handlers alive) + progressCallbacks[progressToken] = onProgress + requestProgressTokens[request.id] = progressToken + + // Create timeout controller if resetTimeoutOnProgress is enabled + let timeoutController: TimeoutController? + if let timeout = options?.timeout, options?.resetTimeoutOnProgress == true { + let controller = TimeoutController( + timeout: timeout, + resetOnProgress: true, + maxTotalTimeout: options?.maxTotalTimeout + ) + timeoutControllers[progressToken] = controller + timeoutController = controller + } else { + timeoutController = nil + } + + // Create stream for receiving the response + let (stream, continuation) = AsyncThrowingStream.makeStream() + + let requestId = request.id + let timeout = options?.timeout + continuation.onTermination = { @Sendable [weak self] termination in + Task { + guard let self else { return } + await self.cleanupPendingRequest(id: requestId) + await self.removeRequestProgressToken(id: requestId) + await self.removeProgressCallback(token: progressToken) + await self.removeTimeoutController(token: progressToken) + + if case .cancelled = termination { + let reason = if let timeout { + "Request timed out after \(timeout)" + } else { + "Client cancelled the request" + } + await self.sendCancellationNotification( + requestId: requestId, + reason: reason + ) + } + } + } + + // Add the pending request before attempting to send + addPendingRequest(id: request.id, continuation: continuation) + + // Send the modified request data + do { + try await connection.send(modifiedRequestData) + } catch { + if removePendingRequest(id: request.id) != nil { + continuation.finish(throwing: error) + } + removeRequestProgressToken(id: request.id) + removeProgressCallback(token: progressToken) + removeTimeoutController(token: progressToken) + throw error + } + + // Wait for response with optional timeout + if let timeout { + // Use TimeoutController if resetTimeoutOnProgress is enabled + if let controller = timeoutController { + return try await withThrowingTaskGroup(of: M.Result.self) { group in + group.addTask { + for try await result in stream { + return result + } + throw MCPError.internalError("No response received") + } - // MARK: - Requests + group.addTask { + try await controller.waitForTimeout() + throw MCPError.internalError("Unreachable - timeout should throw") + } - /// Send a request and receive its response - public func send(_ request: Request) async throws -> M.Result { - guard let connection = connection else { - throw MCPError.internalError("Client connection not initialized") - } + guard let result = try await group.next() else { + throw MCPError.internalError("No response received") + } - let requestData = try encoder.encode(request) + group.cancelAll() + await controller.cancel() + removeProgressCallback(token: progressToken) + removeTimeoutController(token: progressToken) + return result + } + } else { + // Simple timeout without progress-aware reset + return try await withThrowingTaskGroup(of: M.Result.self) { group in + group.addTask { + for try await result in stream { + return result + } + throw MCPError.internalError("No response received") + } - // Store the pending request first - return try await withCheckedThrowingContinuation { continuation in - Task { - // Add the pending request before attempting to send - self.addPendingRequest( - id: request.id, - continuation: continuation, - type: M.Result.self - ) + group.addTask { + try await Task.sleep(for: timeout) + throw MCPError.requestTimeout(timeout: timeout, message: "Request timed out") + } - // Send the request data - do { - // Use the existing connection send - try await connection.send(requestData) - } catch { - // If send fails, try to remove the pending request. - // Resume with the send error only if we successfully removed the request, - // indicating the response handler hasn't processed it yet. - if self.removePendingRequest(id: request.id) != nil { - continuation.resume(throwing: error) + guard let result = try await group.next() else { + throw MCPError.internalError("No response received") } - // Otherwise, the request was already removed by the response handler - // or by disconnect, so the continuation was already resumed. - // Do nothing here. + + group.cancelAll() + removeProgressCallback(token: progressToken) + return result } } + } else { + for try await result in stream { + removeProgressCallback(token: progressToken) + removeTimeoutController(token: progressToken) + return result + } + + removeProgressCallback(token: progressToken) + removeTimeoutController(token: progressToken) + throw MCPError.internalError("No response received") + } + } + + /// Remove a progress callback for the given token. + /// + /// If the token is being tracked for a task (task-augmented response), the callback + /// is NOT removed. This keeps progress handlers alive until the task completes. + private func removeProgressCallback(token: ProgressToken) { + // Check if this token is being tracked for a task + // If so, don't remove the callback - it needs to stay alive until task completes + let isTaskProgressToken = taskProgressTokens.values.contains(token) + if isTaskProgressToken { + return } + progressCallbacks.removeValue(forKey: token) + } + + /// Remove a timeout controller for the given token. + /// + /// If the token is being tracked for a task (task-augmented response), the controller + /// is NOT removed. This keeps timeout tracking alive until the task completes. + private func removeTimeoutController(token: ProgressToken) { + // Check if this token is being tracked for a task + // If so, don't remove the controller - it needs to stay alive until task completes + let isTaskProgressToken = taskProgressTokens.values.contains(token) + if isTaskProgressToken { + return + } + timeoutControllers.removeValue(forKey: token) + } + + /// Remove the request → progress token mapping for the given request ID. + private func removeRequestProgressToken(id: RequestId) { + requestProgressTokens.removeValue(forKey: id) } private func addPendingRequest( - id: ID, - continuation: CheckedContinuation, - type: T.Type // Keep type for AnyPendingRequest internal logic + id: RequestId, + continuation: AsyncThrowingStream.Continuation ) { - pendingRequests[id] = AnyPendingRequest(PendingRequest(continuation: continuation)) + pendingRequests[id] = AnyPendingRequest(continuation: continuation) } - private func removePendingRequest(id: ID) -> AnyPendingRequest? { + private func removePendingRequest(id: RequestId) -> AnyPendingRequest? { return pendingRequests.removeValue(forKey: id) } + /// Removes a pending request without returning it. + /// Used by onTermination handlers when the request has been cancelled. + private func cleanupPendingRequest(id: RequestId) { + pendingRequests.removeValue(forKey: id) + } + + /// Send a CancelledNotification to the server for a cancelled request. + /// + /// Per MCP spec: "When a party wants to cancel an in-progress request, it sends + /// a `notifications/cancelled` notification containing the ID of the request to cancel." + /// + /// This is called when a client Task waiting for a response is cancelled. + /// The notification is sent on a best-effort basis - failures are logged but not thrown. + private func sendCancellationNotification(requestId: RequestId, reason: String?) async { + guard let connection = connection else { + await logger?.debug( + "Cannot send cancellation notification - connection is nil", + metadata: ["requestId": "\(requestId)"] + ) + return + } + + let notification = CancelledNotification.message(.init( + requestId: requestId, + reason: reason + )) + + do { + let notificationData = try encoder.encode(notification) + try await connection.send(notificationData) + await logger?.debug( + "Sent cancellation notification", + metadata: [ + "requestId": "\(requestId)", + "reason": "\(reason ?? "none")", + ] + ) + } catch { + // Log but don't throw - cancellation notification is best-effort + // per MCP spec's fire-and-forget nature of notifications + await logger?.debug( + "Failed to send cancellation notification", + metadata: [ + "requestId": "\(requestId)", + "error": "\(error)", + ] + ) + } + } + // MARK: - Batching /// A batch of requests. @@ -359,21 +1450,38 @@ public actor Client { > { requests.append(try AnyRequest(request)) - // Return a Task that registers the pending request and awaits its result. - // The continuation is resumed when the response arrives. - return Task { - try await withCheckedThrowingContinuation { continuation in - // We are already inside a Task, but need another Task - // to bridge to the client actor's context. - Task { - await client.addPendingRequest( - id: request.id, - continuation: continuation, - type: M.Result.self + // Create stream for receiving the response + let (stream, continuation) = AsyncThrowingStream.makeStream() + + // Clean up pending request if caller cancels (e.g., task cancelled) + // and send CancelledNotification to server per MCP spec + let requestId = request.id + continuation.onTermination = { @Sendable [weak client] termination in + Task { + guard let client else { return } + await client.cleanupPendingRequest(id: requestId) + + // Per MCP spec: send notifications/cancelled when cancelling a request + // Only send if the stream was cancelled (not finished normally) + if case .cancelled = termination { + await client.sendCancellationNotification( + requestId: requestId, + reason: "Client cancelled the batch request" ) } } } + + // Register the pending request + await client.addPendingRequest(id: request.id, continuation: continuation) + + // Return a Task that waits for the response via the stream + return Task { + for try await result in stream { + return result + } + throw MCPError.internalError("No response received") + } } } @@ -509,10 +1617,25 @@ public actor Client { let result = try await send(request) + // Per MCP spec: "If the client does not support the version in the + // server's response, it SHOULD disconnect." + guard Version.supported.contains(result.protocolVersion) else { + await disconnect() + throw MCPError.invalidRequest( + "Server responded with unsupported protocol version: \(result.protocolVersion). " + + "Supported versions: \(Version.supported.sorted().joined(separator: ", "))" + ) + } + self.serverCapabilities = result.capabilities self.serverVersion = result.protocolVersion self.instructions = result.instructions + // HTTP transports must set the protocol version in headers after initialization + if let httpTransport = connection as? HTTPClientTransport { + await httpTransport.setProtocolVersion(result.protocolVersion) + } + try await notify(InitializedNotification.message()) return result @@ -525,7 +1648,7 @@ public actor Client { // MARK: - Prompts - public func getPrompt(name: String, arguments: [String: Value]? = nil) async throws + public func getPrompt(name: String, arguments: [String: String]? = nil) async throws -> (description: String?, messages: [Prompt.Message]) { try validateServerCapability(\.prompts, "Prompts") @@ -577,6 +1700,12 @@ public actor Client { _ = try await send(request) } + public func unsubscribeFromResource(uri: String) async throws { + try validateServerCapability(\.resources?.subscribe, "Resource subscription") + let request = ResourceUnsubscribe.request(.init(uri: uri)) + _ = try await send(request) + } + public func listResourceTemplates(cursor: String? = nil) async throws -> ( templates: [Resource.Template], nextCursor: String? ) { @@ -608,51 +1737,305 @@ public actor Client { } public func callTool(name: String, arguments: [String: Value]? = nil) async throws -> ( - content: [Tool.Content], isError: Bool? + content: [Tool.Content], structuredContent: Value?, isError: Bool? ) { try validateServerCapability(\.tools, "Tools") let request = CallTool.request(.init(name: name, arguments: arguments)) let result = try await send(request) - return (content: result.content, isError: result.isError) + // TODO: Add client-side output validation against the tool's outputSchema. + // TypeScript and Python SDKs cache tool outputSchemas from listTools() and + // validate structuredContent when receiving tool results. + return (content: result.content, structuredContent: result.structuredContent, isError: result.isError) } - // MARK: - Sampling + // MARK: - Completions - /// Register a handler for sampling requests from servers + /// Request completion suggestions from the server. /// - /// Sampling allows servers to request LLM completions through the client, - /// enabling sophisticated agentic behaviors while maintaining human-in-the-loop control. + /// Completions provide autocomplete suggestions for prompt arguments or resource + /// template URI parameters. /// - /// 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 + /// - Parameters: + /// - ref: A reference to the prompt or resource template to get completions for. + /// - argument: The argument being completed, including its name and partial value. + /// - context: Optional additional context with previously-resolved argument values. + /// - Returns: The completion suggestions from the server. + public func complete( + ref: CompletionReference, + argument: CompletionArgument, + context: CompletionContext? = nil + ) async throws -> CompletionSuggestions { + try validateServerCapability(\.completions, "Completions") + let request = Complete.request(.init(ref: ref, argument: argument, context: context)) + let result = try await send(request) + return result.completion + } + + // MARK: - Logging + + /// Set the minimum log level for messages from 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. + /// After calling this method, the server should only send log messages + /// at the specified level or higher (more severe). + /// + /// - Parameter level: The minimum log level to receive. + public func setLoggingLevel(_ level: LoggingLevel) async throws { + try validateServerCapability(\.logging, "Logging") + let request = SetLoggingLevel.request(.init(level: level)) + _ = try await send(request) + } - // This would register the handler similar to how servers register method handlers: - // methodHandlers[CreateSamplingMessage.name] = TypedRequestHandler(handler) + // MARK: - Tasks (Experimental) + // Note: These methods are internal. Access via client.experimental.* - return self + func getTask(taskId: String) async throws -> GetTask.Result { + try validateServerCapability(\.tasks, "Tasks") + let request = GetTask.request(.init(taskId: taskId)) + return try await send(request) + } + + func listTasks(cursor: String? = nil) async throws -> (tasks: [MCPTask], nextCursor: String?) { + try validateServerCapability(\.tasks, "Tasks") + let request: Request + if let cursor { + request = ListTasks.request(.init(cursor: cursor)) + } else { + request = ListTasks.request(.init()) + } + let result = try await send(request) + return (tasks: result.tasks, nextCursor: result.nextCursor) + } + + func cancelTask(taskId: String) async throws -> CancelTask.Result { + try validateServerCapability(\.tasks, "Tasks") + let request = CancelTask.request(.init(taskId: taskId)) + return try await send(request) + } + + func getTaskResult(taskId: String) async throws -> GetTaskPayload.Result { + try validateServerCapability(\.tasks, "Tasks") + let request = GetTaskPayload.request(.init(taskId: taskId)) + return try await send(request) + } + + /// Get the task result decoded as a specific type. + /// + /// This method retrieves the task result and decodes the `extraFields` as the specified type. + /// The `extraFields` contain the actual result payload (e.g., CallTool.Result fields). + func getTaskResultAs(taskId: String, type: T.Type) async throws -> T { + let result = try await getTaskResult(taskId: taskId) + + // The result's extraFields contain the actual result payload + // We need to encode them back to JSON and decode as the target type + guard let extraFields = result.extraFields else { + throw MCPError.invalidParams("Task result has no payload") + } + + // Convert extraFields to the target type + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + // Encode the extraFields as JSON + let jsonData = try encoder.encode(extraFields) + + // Decode as the target type + return try decoder.decode(T.self, from: jsonData) + } + + func callToolAsTask( + name: String, + arguments: [String: Value]? = nil, + ttl: Int? = nil + ) async throws -> CreateTaskResult { + try validateServerCapability(\.tasks, "Tasks") + try validateServerCapability(\.tools, "Tools") + + let taskMetadata = TaskMetadata(ttl: ttl) + let request = CallTool.request(.init( + name: name, + arguments: arguments, + task: taskMetadata + )) + + // The server should return CreateTaskResult for task-augmented requests + // We need to decode as CreateTaskResult instead of CallTool.Result + guard let connection = connection else { + throw MCPError.internalError("Client connection not initialized") + } + + let requestData = try encoder.encode(request) + + // Create stream for receiving the response + let (stream, continuation) = AsyncThrowingStream.makeStream() + + let requestId = request.id + continuation.onTermination = { @Sendable [weak self] _ in + Task { await self?.cleanupPendingRequest(id: requestId) } + } + + addPendingRequest(id: request.id, continuation: continuation) + + do { + try await connection.send(requestData) + } catch { + if removePendingRequest(id: request.id) != nil { + continuation.finish(throwing: error) + } + throw error + } + + for try await result in stream { + return result + } + + throw MCPError.internalError("No response received") + } + + func pollTask(taskId: String) -> AsyncThrowingStream { + AsyncThrowingStream { continuation in + let pollingTask = Task { + do { + while !Task.isCancelled { + let task = try await self.getTask(taskId: taskId) + continuation.yield(task) + + if isTerminalStatus(task.status) { + continuation.finish() + return + } + + // Wait based on pollInterval (default 1 second) + let intervalMs = task.pollInterval ?? 1000 + try await Task.sleep(for: .milliseconds(intervalMs)) + } + // Task was cancelled + continuation.finish(throwing: CancellationError()) + } catch { + continuation.finish(throwing: error) + } + } + + // Cancel the polling task when the stream is terminated + continuation.onTermination = { _ in + pollingTask.cancel() + } + } + } + + func pollUntilTerminal(taskId: String) async throws -> GetTask.Result { + for try await status in pollTask(taskId: taskId) { + if isTerminalStatus(status.status) { + return status + } + } + // This shouldn't happen, but handle it gracefully + throw MCPError.internalError("Task polling ended unexpectedly") + } + + func callToolAsTaskAndWait( + name: String, + arguments: [String: Value]? = nil, + ttl: Int? = nil + ) async throws -> (content: [Tool.Content], isError: Bool?) { + // Start the task + let createResult = try await callToolAsTask(name: name, arguments: arguments, ttl: ttl) + let taskId = createResult.task.taskId + + // Wait for the result (uses blocking getTaskResult) + let payloadResult = try await getTaskResult(taskId: taskId) + + // Decode the result as CallTool.Result + // Per MCP spec, the result fields are flattened directly in the response (via extraFields) + guard let extraFields = payloadResult.extraFields else { + throw MCPError.internalError("Task completed but no result available") + } + + // Convert extraFields back to Value for decoding + let resultValue = Value.object(extraFields) + let resultData = try encoder.encode(resultValue) + let toolResult = try decoder.decode(CallTool.Result.self, from: resultData) + return (content: toolResult.content, isError: toolResult.isError) + } + + func callToolStream( + name: String, + arguments: [String: Value]? = nil, + ttl: Int? = nil + ) -> AsyncThrowingStream { + AsyncThrowingStream { continuation in + let streamTask = Task { + do { + // Step 1: Create the task + let createResult = try await self.callToolAsTask(name: name, arguments: arguments, ttl: ttl) + let task = createResult.task + continuation.yield(.taskCreated(task)) + + // Step 2: Poll for status updates until terminal + var lastStatus = task.status + var finalTask = task + + while !isTerminalStatus(lastStatus) { + // Wait based on pollInterval (default 1 second) + let intervalMs = finalTask.pollInterval ?? 1000 + try await Task.sleep(for: .milliseconds(intervalMs)) + + // Get updated status + let statusResult = try await self.getTask(taskId: task.taskId) + finalTask = MCPTask( + taskId: statusResult.taskId, + status: statusResult.status, + ttl: statusResult.ttl, + createdAt: statusResult.createdAt, + lastUpdatedAt: statusResult.lastUpdatedAt, + pollInterval: statusResult.pollInterval, + statusMessage: statusResult.statusMessage + ) + + // Only yield if status or message changed + if statusResult.status != lastStatus || statusResult.statusMessage != nil { + continuation.yield(.taskStatus(finalTask)) + } + lastStatus = statusResult.status + } + + // Step 3: Get the final result + if finalTask.status == .completed { + let payloadResult = try await self.getTaskResult(taskId: task.taskId) + + // Decode the result as CallTool.Result + if let extraFields = payloadResult.extraFields { + let resultValue = Value.object(extraFields) + let resultData = try self.encoder.encode(resultValue) + let toolResult = try self.decoder.decode(CallTool.Result.self, from: resultData) + continuation.yield(.result(toolResult)) + } else { + // No result available - return empty result + continuation.yield(.result(CallTool.Result(content: []))) + } + } else if finalTask.status == .failed { + let error = MCPError.internalError(finalTask.statusMessage ?? "Task failed") + continuation.yield(.error(error)) + } else if finalTask.status == .cancelled { + let error = MCPError.internalError("Task was cancelled") + continuation.yield(.error(error)) + } + + continuation.finish() + } catch let error as MCPError { + continuation.yield(.error(error)) + continuation.finish() + } catch { + let mcpError = MCPError.internalError(error.localizedDescription) + continuation.yield(.error(mcpError)) + continuation.finish() + } + } + + // Cancel the stream task if the stream is terminated + continuation.onTermination = { _ in + streamTask.cancel() + } + } } // MARK: - @@ -662,6 +2045,14 @@ public actor Client { "Processing response", metadata: ["id": "\(response.id)"]) + // Check for task-augmented response BEFORE resuming the request. + // Per MCP spec 2025-11-25: progress tokens continue for task lifetime. + // If this is a CreateTaskResult, we need to keep the progress handler alive. + if case .success(let value) = response.result, + case .object(let resultObject) = value { + checkForTaskResponse(response: response, value: resultObject) + } + // Attempt to remove the pending request using the response ID. // Resume with the response only if it hadn't yet been removed. if let removedRequest = self.removePendingRequest(id: response.id) { @@ -682,11 +2073,94 @@ public actor Client { } } + /// Check if a response is a task-augmented response (CreateTaskResult). + /// + /// If the response contains a `task` object with `taskId`, this is a task-augmented + /// response. Per MCP spec, progress notifications can continue until the task reaches + /// terminal status, so we migrate the progress handler from request tracking to task tracking. + /// + /// This matches the TypeScript SDK pattern where task progress tokens are kept alive + /// until the task completes. + private func checkForTaskResponse(response: Response, value: [String: Value]) { + // Check if we have a progress token for this request + guard let progressToken = requestProgressTokens[response.id] else { return } + + // Check if response has task.taskId (CreateTaskResult pattern) + // This mirrors TypeScript's check: result.task?.taskId + guard let taskValue = value["task"], + case .object(let taskObject) = taskValue, + let taskIdValue = taskObject["taskId"], + case .string(let taskId) = taskIdValue else { + // Not a task response - clean up request tracking + // (the progress callback itself is cleaned up in send() after receiving result) + requestProgressTokens.removeValue(forKey: response.id) + return + } + + // This is a task-augmented response! + // Migrate progress token from request tracking to task tracking. + // This keeps the progress handler alive until the task completes. + taskProgressTokens[taskId] = progressToken + requestProgressTokens.removeValue(forKey: response.id) + + Task { + await logger?.debug( + "Keeping progress handler alive for task", + metadata: [ + "taskId": "\(taskId)", + "progressToken": "\(progressToken)", + ] + ) + } + } + + /// Clean up the progress handler for a completed task. + /// + /// Call this method when a task reaches terminal status (completed, failed, cancelled) + /// to remove the progress callback and timeout controller. + /// + /// ## Example + /// + /// ```swift + /// // Register task status notification handler + /// await client.onNotification(TaskStatusNotification.self) { message in + /// if message.params.status.isTerminal { + /// await client.cleanupTaskProgressHandler(taskId: message.params.taskId) + /// } + /// } + /// ``` + /// + /// - Parameter taskId: The ID of the task that completed. + public func cleanupTaskProgressHandler(taskId: String) { + guard let progressToken = taskProgressTokens.removeValue(forKey: taskId) else { return } + + progressCallbacks.removeValue(forKey: progressToken) + timeoutControllers.removeValue(forKey: progressToken) + + Task { + await logger?.debug( + "Cleaned up progress handler for completed task", + metadata: ["taskId": "\(taskId)"] + ) + } + } + private func handleMessage(_ message: Message) async { await logger?.trace( "Processing notification", metadata: ["method": "\(message.method)"]) + // Check if this is a progress notification and invoke any registered callback + if message.method == ProgressNotification.name { + await handleProgressNotification(message) + } + + // Check if this is a task status notification and clean up progress handlers + // for terminal task statuses (per MCP spec, progress tokens are valid until terminal status) + if message.method == TaskStatusNotification.name { + await handleTaskStatusNotification(message) + } + // Find notification handlers for this method guard let handlers = notificationHandlers[message.method] else { return } @@ -705,6 +2179,264 @@ public actor Client { } } + /// Handle a progress notification by invoking any registered callback. + private func handleProgressNotification(_ message: Message) async { + do { + // Decode as ProgressNotification.Parameters + let paramsData = try encoder.encode(message.params) + let params = try decoder.decode(ProgressNotification.Parameters.self, from: paramsData) + + // Look up the callback for this token + guard let callback = progressCallbacks[params.progressToken] else { + // TypeScript SDK logs an error for unknown progress tokens + await logger?.warning( + "Received progress notification for unknown token", + metadata: ["progressToken": "\(params.progressToken)"]) + return + } + + // Signal the timeout controller if one exists for this token + // This allows resetTimeoutOnProgress to work + if let timeoutController = timeoutControllers[params.progressToken] { + await timeoutController.signalProgress() + } + + // Invoke the callback + let progress = Progress( + value: params.progress, + total: params.total, + message: params.message + ) + await callback(progress) + } catch { + await logger?.warning( + "Failed to decode progress notification", + metadata: ["error": "\(error)"]) + } + } + + /// Handle a task status notification by cleaning up progress handlers for terminal tasks. + /// + /// Per MCP spec 2025-11-25: progress tokens continue throughout task lifetime until terminal status. + /// This method automatically cleans up progress handlers when a task reaches completed, failed, or cancelled. + private func handleTaskStatusNotification(_ message: Message) async { + do { + // Decode as TaskStatusNotification.Parameters + let paramsData = try encoder.encode(message.params) + let params = try decoder.decode(TaskStatusNotification.Parameters.self, from: paramsData) + + // If the task reached a terminal status, clean up its progress handler + if params.status.isTerminal { + cleanupTaskProgressHandler(taskId: params.taskId) + } + } catch { + // Don't log errors for task status notifications - they may not be task-related + // and the user may not have registered a handler for them + } + } + + /// Handle an incoming request from the server (bidirectional communication). + /// + /// This enables server→client requests such as sampling, roots, and elicitation. + /// + /// ## Task-Augmented Request Handling + /// + /// For `sampling/createMessage` and `elicitation/create` requests, this method + /// checks for a `task` field in the request params. If present, it routes to + /// the task-augmented handler (which returns `CreateTaskResult`) instead of + /// the normal handler. + /// + /// This follows the Python SDK pattern of storing task-augmented handlers + /// separately and checking at dispatch time, rather than the TypeScript pattern + /// of wrapping handlers at registration time. The Python pattern was chosen + /// because: + /// - It allows handlers to be registered in any order without losing task-awareness + /// - It keeps task logic separate from normal handler logic + /// - It's more explicit about which handler is called for which request type + private func handleIncomingRequest(_ request: Request) async { + await logger?.trace( + "Processing incoming request from server", + metadata: [ + "method": "\(request.method)", + "id": "\(request.id)", + ]) + + // Validate elicitation mode against client capabilities + // Per spec: Client MUST return -32602 if server requests unsupported mode + if request.method == Elicit.name { + if let modeError = await validateElicitationMode(request) { + await sendResponse(modeError) + return + } + } + + // Check for task-augmented sampling/elicitation requests first + // This matches the Python SDK pattern where task detection happens at dispatch time + if let taskResponse = await handleTaskAugmentedRequest(request) { + await sendResponse(taskResponse) + return + } + + // Find handler for method name + guard let handler = requestHandlers[request.method] else { + await logger?.warning( + "No handler registered for server request", + metadata: ["method": "\(request.method)"]) + + // Send error response + let response = AnyMethod.response( + id: request.id, + error: MCPError.methodNotFound("Client has no handler for: \(request.method)") + ) + await sendResponse(response) + return + } + + // Execute the handler and send response + do { + let response = try await handler(request) + + // Check cancellation before sending response (per MCP spec: + // "Receivers of a cancellation notification SHOULD... Not send a response + // for the cancelled request") + if Task.isCancelled { + await logger?.debug( + "Server request cancelled, suppressing response", + metadata: ["id": "\(request.id)"] + ) + return + } + + await sendResponse(response) + } catch { + // Also check cancellation on error path - don't send error response if cancelled + if Task.isCancelled { + await logger?.debug( + "Server request cancelled during error handling, suppressing response", + metadata: ["id": "\(request.id)"] + ) + return + } + + await logger?.error( + "Error handling server request", + metadata: [ + "method": "\(request.method)", + "error": "\(error)", + ]) + let errorResponse = AnyMethod.response( + id: request.id, + error: (error as? MCPError) ?? MCPError.internalError(error.localizedDescription) + ) + await sendResponse(errorResponse) + } + } + + /// Validate that an elicitation request uses a mode supported by client capabilities. + /// + /// Per MCP spec: Client MUST return -32602 (Invalid params) if server sends + /// an elicitation/create request with a mode not declared in client capabilities. + /// + /// - Parameter request: The incoming elicitation request + /// - Returns: An error response if mode is unsupported, nil if valid + private func validateElicitationMode(_ request: Request) async -> Response? { + do { + let paramsData = try encoder.encode(request.params) + let params = try decoder.decode(Elicit.Parameters.self, from: paramsData) + + switch params { + case .form: + // Form mode requires form capability + if capabilities.elicitation?.form == nil { + return Response( + id: request.id, + error: .invalidParams("Client does not support form elicitation mode") + ) + } + case .url: + // URL mode requires url capability + if capabilities.elicitation?.url == nil { + return Response( + id: request.id, + error: .invalidParams("Client does not support URL elicitation mode") + ) + } + } + } catch { + // If we can't decode the params, let the normal handler deal with it + await logger?.warning( + "Failed to decode elicitation params for mode validation", + metadata: ["error": "\(error)"]) + } + + return nil + } + + /// Check if a request is task-augmented and handle it if so. + /// + /// - Parameter request: The incoming request + /// - Returns: A response if the request was task-augmented and handled, nil otherwise + private func handleTaskAugmentedRequest(_ request: Request) async -> Response? { + do { + // Check for task-augmented sampling request + if request.method == CreateSamplingMessage.name, + let taskHandler = taskAugmentedSamplingHandler { + let paramsData = try encoder.encode(request.params) + let params = try decoder.decode(CreateSamplingMessage.Parameters.self, from: paramsData) + + if let taskMetadata = params.task { + let result = try await taskHandler(params, taskMetadata) + let resultData = try encoder.encode(result) + let resultValue = try decoder.decode(Value.self, from: resultData) + return Response(id: request.id, result: resultValue) + } + } + + // Check for task-augmented elicitation request + if request.method == Elicit.name, + let taskHandler = taskAugmentedElicitationHandler { + let paramsData = try encoder.encode(request.params) + let params = try decoder.decode(Elicit.Parameters.self, from: paramsData) + + let taskMetadata: TaskMetadata? = switch params { + case .form(let formParams): formParams.task + case .url(let urlParams): urlParams.task + } + + if let taskMetadata { + let result = try await taskHandler(params, taskMetadata) + let resultData = try encoder.encode(result) + let resultValue = try decoder.decode(Value.self, from: resultData) + return Response(id: request.id, result: resultValue) + } + } + } catch let error as MCPError { + return Response(id: request.id, error: error) + } catch { + return Response(id: request.id, error: MCPError.internalError(error.localizedDescription)) + } + + // Not a task-augmented request + return nil + } + + /// Send a response back to the server. + private func sendResponse(_ response: Response) async { + guard let connection = connection else { + await logger?.warning("Cannot send response - client not connected") + return + } + + do { + let responseData = try encoder.encode(response) + try await connection.send(responseData) + } catch { + await logger?.error( + "Failed to send response to server", + metadata: ["error": "\(error)"]) + } + } + // MARK: - /// Validate the server capabilities. diff --git a/Sources/MCP/Client/Elicitation.swift b/Sources/MCP/Client/Elicitation.swift new file mode 100644 index 00000000..06de8862 --- /dev/null +++ b/Sources/MCP/Client/Elicitation.swift @@ -0,0 +1,789 @@ +import Foundation + +/// Elicitation allows servers to request additional information from users +/// through the client. This enables interactive workflows where the server +/// needs user input during an operation. +/// +/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/client/elicitation/ + +// MARK: - Schema Types + +/// Format constraints for string fields in elicitation forms. +/// +/// These formats provide validation hints to the client for user input fields. +/// The client may use these to provide appropriate input controls or validation. +public enum StringSchemaFormat: String, Hashable, Codable, Sendable { + /// Email address format (e.g., "user@example.com") + case email + /// URI format (e.g., "https://example.com/path") + case uri + /// Date format (ISO 8601 date, e.g., "2024-01-15") + case date + /// Date-time format (ISO 8601 date-time, e.g., "2024-01-15T10:30:00Z") + case dateTime = "date-time" +} + +/// Schema definition for string input fields in elicitation forms. +/// +/// Use this to request text input from the user, optionally with validation constraints +/// like minimum/maximum length, pattern (regex), or format requirements. +/// +/// ```swift +/// // Simple text field +/// let nameField = StringSchema(title: "Name", description: "Your full name") +/// +/// // Email field with validation +/// let emailField = StringSchema( +/// title: "Email", +/// format: .email, +/// defaultValue: "user@example.com" +/// ) +/// +/// // Field with regex pattern validation +/// let zipCode = StringSchema( +/// title: "ZIP Code", +/// pattern: "^[0-9]{5}$" +/// ) +/// ``` +public struct StringSchema: Hashable, Codable, Sendable { + public let type: String + public var title: String? + public var description: String? + public var minLength: Int? + public var maxLength: Int? + public var pattern: String? + public var format: StringSchemaFormat? + public var defaultValue: String? + + private enum CodingKeys: String, CodingKey { + case type, title, description, minLength, maxLength, pattern, format + case defaultValue = "default" + } + + public init( + title: String? = nil, + description: String? = nil, + minLength: Int? = nil, + maxLength: Int? = nil, + pattern: String? = nil, + format: StringSchemaFormat? = nil, + defaultValue: String? = nil + ) { + self.type = "string" + self.title = title + self.description = description + self.minLength = minLength + self.maxLength = maxLength + self.pattern = pattern + self.format = format + self.defaultValue = defaultValue + } +} + +/// Schema definition for numeric input fields in elicitation forms. +/// +/// Use this to request number or integer input from the user, optionally with +/// minimum/maximum constraints. +/// +/// ```swift +/// // Integer field for age +/// let ageField = NumberSchema(isInteger: true, title: "Age", minimum: 0, maximum: 150) +/// +/// // Decimal field for price +/// let priceField = NumberSchema(title: "Price", minimum: 0.0) +/// ``` +public struct NumberSchema: Hashable, Codable, Sendable { + public let type: String + public var title: String? + public var description: String? + public var minimum: Double? + public var maximum: Double? + public var defaultValue: Double? + + private enum CodingKeys: String, CodingKey { + case type, title, description, minimum, maximum + case defaultValue = "default" + } + + public init( + isInteger: Bool = false, + title: String? = nil, + description: String? = nil, + minimum: Double? = nil, + maximum: Double? = nil, + defaultValue: Double? = nil + ) { + self.type = isInteger ? "integer" : "number" + self.title = title + self.description = description + self.minimum = minimum + self.maximum = maximum + self.defaultValue = defaultValue + } +} + +/// Schema definition for boolean (checkbox/toggle) fields in elicitation forms. +/// +/// Use this to request a yes/no or true/false choice from the user. +/// +/// ```swift +/// let agreeField = BooleanSchema( +/// title: "I agree to the terms", +/// description: "You must accept the terms to continue", +/// defaultValue: false +/// ) +/// ``` +public struct BooleanSchema: Hashable, Codable, Sendable { + public let type: String + public var title: String? + public var description: String? + public var defaultValue: Bool? + + private enum CodingKeys: String, CodingKey { + case type, title, description + case defaultValue = "default" + } + + public init( + title: String? = nil, + description: String? = nil, + defaultValue: Bool? = nil + ) { + self.type = "boolean" + self.title = title + self.description = description + self.defaultValue = defaultValue + } +} + +/// An option in a titled enum with a value and display label. +public struct TitledEnumOption: Hashable, Codable, Sendable { + /// The constant value for this option. + public let const: String + /// The display label for this option. + public let title: String + + public init(const: String, title: String) { + self.const = const + self.title = title + } +} + +/// Schema definition for single-select enum fields without display titles. +public struct UntitledEnumSchema: Hashable, Codable, Sendable { + public let type: String + public var title: String? + public var description: String? + public var enumValues: [String] + public var defaultValue: String? + + private enum CodingKeys: String, CodingKey { + case type, title, description + case enumValues = "enum" + case defaultValue = "default" + } + + public init( + title: String? = nil, + description: String? = nil, + enumValues: [String], + defaultValue: String? = nil + ) { + self.type = "string" + self.title = title + self.description = description + self.enumValues = enumValues + self.defaultValue = defaultValue + } +} + +/// Schema definition for single-select enum fields with display titles. +public struct TitledEnumSchema: Hashable, Codable, Sendable { + public let type: String + public var title: String? + public var description: String? + public var oneOf: [TitledEnumOption] + public var defaultValue: String? + + private enum CodingKeys: String, CodingKey { + case type, title, description, oneOf + case defaultValue = "default" + } + + public init( + title: String? = nil, + description: String? = nil, + oneOf: [TitledEnumOption], + defaultValue: String? = nil + ) { + self.type = "string" + self.title = title + self.description = description + self.oneOf = oneOf + self.defaultValue = defaultValue + } +} + +/// Schema definition for legacy enum fields with enumNames (non-standard). +public struct LegacyTitledEnumSchema: Hashable, Codable, Sendable { + public let type: String + public var title: String? + public var description: String? + public var enumValues: [String] + public var enumNames: [String]? + public var defaultValue: String? + + private enum CodingKeys: String, CodingKey { + case type, title, description, enumNames + case enumValues = "enum" + case defaultValue = "default" + } + + public init( + title: String? = nil, + description: String? = nil, + enumValues: [String], + enumNames: [String]? = nil, + defaultValue: String? = nil + ) { + self.type = "string" + self.title = title + self.description = description + self.enumValues = enumValues + self.enumNames = enumNames + self.defaultValue = defaultValue + } +} + +// MARK: - Multi-Select Enum Schemas + +/// Items definition for untitled multi-select enum. +public struct UntitledMultiSelectItems: Hashable, Codable, Sendable { + public let type: String + public var enumValues: [String] + + private enum CodingKeys: String, CodingKey { + case type + case enumValues = "enum" + } + + public init(enumValues: [String]) { + self.type = "string" + self.enumValues = enumValues + } +} + +/// Schema definition for multi-select enum fields without display titles. +public struct UntitledMultiSelectEnumSchema: Hashable, Codable, Sendable { + public let type: String + public var title: String? + public var description: String? + public var minItems: Int? + public var maxItems: Int? + public var items: UntitledMultiSelectItems + public var defaultValue: [String]? + + private enum CodingKeys: String, CodingKey { + case type, title, description, minItems, maxItems, items + case defaultValue = "default" + } + + public init( + title: String? = nil, + description: String? = nil, + minItems: Int? = nil, + maxItems: Int? = nil, + enumValues: [String], + defaultValue: [String]? = nil + ) { + self.type = "array" + self.title = title + self.description = description + self.minItems = minItems + self.maxItems = maxItems + self.items = UntitledMultiSelectItems(enumValues: enumValues) + self.defaultValue = defaultValue + } +} + +/// Items definition for titled multi-select enum. +public struct TitledMultiSelectItems: Hashable, Codable, Sendable { + public var anyOf: [TitledEnumOption] + + public init(anyOf: [TitledEnumOption]) { + self.anyOf = anyOf + } +} + +/// Schema definition for multi-select enum fields with display titles. +public struct TitledMultiSelectEnumSchema: Hashable, Codable, Sendable { + public let type: String + public var title: String? + public var description: String? + public var minItems: Int? + public var maxItems: Int? + public var items: TitledMultiSelectItems + public var defaultValue: [String]? + + private enum CodingKeys: String, CodingKey { + case type, title, description, minItems, maxItems, items + case defaultValue = "default" + } + + public init( + title: String? = nil, + description: String? = nil, + minItems: Int? = nil, + maxItems: Int? = nil, + options: [TitledEnumOption], + defaultValue: [String]? = nil + ) { + self.type = "array" + self.title = title + self.description = description + self.minItems = minItems + self.maxItems = maxItems + self.items = TitledMultiSelectItems(anyOf: options) + self.defaultValue = defaultValue + } +} + +/// A primitive schema definition for form fields in elicitation requests. +/// +/// This enum represents all the field types that can be used in an elicitation form. +/// Each case corresponds to a specific input type with its own validation and display options. +/// +/// ## Supported Field Types +/// +/// - **string**: Text input with optional format validation (email, URI, date) +/// - **number**: Numeric input (integer or decimal) with optional min/max +/// - **boolean**: Checkbox or toggle for true/false values +/// - **untitledEnum**: Single-select dropdown with simple string values +/// - **titledEnum**: Single-select dropdown with separate values and display labels +/// - **untitledMultiSelect**: Multi-select list with simple string values +/// - **titledMultiSelect**: Multi-select list with separate values and display labels +/// +/// ## Example +/// +/// ```swift +/// let schema = ElicitationSchema(properties: [ +/// "name": .string(StringSchema(title: "Name")), +/// "age": .number(NumberSchema(isInteger: true, title: "Age", minimum: 0)), +/// "agree": .boolean(BooleanSchema(title: "Accept terms")), +/// "color": .untitledEnum(UntitledEnumSchema(title: "Color", enumValues: ["red", "green", "blue"])) +/// ]) +/// ``` +public enum PrimitiveSchemaDefinition: Hashable, Sendable { + case string(StringSchema) + case number(NumberSchema) + case boolean(BooleanSchema) + case untitledEnum(UntitledEnumSchema) + case titledEnum(TitledEnumSchema) + case legacyTitledEnum(LegacyTitledEnumSchema) + case untitledMultiSelect(UntitledMultiSelectEnumSchema) + case titledMultiSelect(TitledMultiSelectEnumSchema) +} + +extension PrimitiveSchemaDefinition: Codable { + private enum CodingKeys: String, CodingKey { + case type, oneOf, enumValues = "enum", enumNames, items + } + + private enum ItemsCodingKeys: String, CodingKey { + case enumValues = "enum", anyOf + } + + 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 "string": + // Check if it's an enum type (has enum or oneOf) + if container.contains(.oneOf) { + let schema = try TitledEnumSchema(from: decoder) + self = .titledEnum(schema) + } else if container.contains(.enumValues) { + // Check for enumNames (legacy format) + if container.contains(.enumNames) { + let schema = try LegacyTitledEnumSchema(from: decoder) + self = .legacyTitledEnum(schema) + } else { + let schema = try UntitledEnumSchema(from: decoder) + self = .untitledEnum(schema) + } + } else { + let schema = try StringSchema(from: decoder) + self = .string(schema) + } + case "array": + // Multi-select enum - check items for anyOf (titled) or enum (untitled) + if container.contains(.items) { + let itemsContainer = try container.nestedContainer( + keyedBy: ItemsCodingKeys.self, forKey: .items) + if itemsContainer.contains(.anyOf) { + let schema = try TitledMultiSelectEnumSchema(from: decoder) + self = .titledMultiSelect(schema) + } else { + let schema = try UntitledMultiSelectEnumSchema(from: decoder) + self = .untitledMultiSelect(schema) + } + } else { + throw DecodingError.dataCorruptedError( + forKey: .type, in: container, + debugDescription: "Array type must have items property") + } + case "number", "integer": + let schema = try NumberSchema(from: decoder) + self = .number(schema) + case "boolean": + let schema = try BooleanSchema(from: decoder) + self = .boolean(schema) + default: + throw DecodingError.dataCorruptedError( + forKey: .type, in: container, + debugDescription: "Unknown primitive schema type: \(type)") + } + } + + public func encode(to encoder: Encoder) throws { + switch self { + case .string(let schema): + try schema.encode(to: encoder) + case .number(let schema): + try schema.encode(to: encoder) + case .boolean(let schema): + try schema.encode(to: encoder) + case .untitledEnum(let schema): + try schema.encode(to: encoder) + case .titledEnum(let schema): + try schema.encode(to: encoder) + case .legacyTitledEnum(let schema): + try schema.encode(to: encoder) + case .untitledMultiSelect(let schema): + try schema.encode(to: encoder) + case .titledMultiSelect(let schema): + try schema.encode(to: encoder) + } + } +} + +// MARK: - Elicitation Request + +/// Parameters for a form-mode elicitation request. +public struct ElicitRequestFormParams: Hashable, Codable, Sendable { + /// The elicitation mode (optional, defaults to "form"). + public var mode: String? + /// The message to present to the user describing what information is being requested. + public var message: String + /// A restricted subset of JSON Schema defining the form fields. + public var requestedSchema: ElicitationSchema + /// Request metadata including progress token. + public var _meta: RequestMeta? + /// Task augmentation metadata. If present, the receiver should run the elicitation + /// as a background task and return `CreateTaskResult` instead of `ElicitResult`. + public var task: TaskMetadata? + + public init( + mode: String? = nil, + message: String, + requestedSchema: ElicitationSchema, + _meta: RequestMeta? = nil, + task: TaskMetadata? = nil + ) { + self.mode = mode + self.message = message + self.requestedSchema = requestedSchema + self._meta = _meta + self.task = task + } +} + +/// The schema for an elicitation form, defining the fields and their types. +public struct ElicitationSchema: Hashable, Codable, Sendable { + /// The JSON Schema dialect (optional). + public var schema: String? + /// Must be "object". + public let type: String + /// The form field definitions. + public var properties: [String: PrimitiveSchemaDefinition] + /// The list of required field names. + public var required: [String]? + + private enum CodingKeys: String, CodingKey { + case schema = "$schema" + case type, properties, required + } + + public init( + schema: String? = nil, + properties: [String: PrimitiveSchemaDefinition], + required: [String]? = nil + ) { + self.schema = schema + self.type = "object" + self.properties = properties + self.required = required + } +} + +// MARK: - Elicitation Result + +/// The action taken by the user in response to an elicitation request. +public enum ElicitAction: String, Hashable, Codable, Sendable { + /// User submitted the form/confirmed the action. + case accept + /// User explicitly declined the action. + case decline + /// User dismissed without making an explicit choice. + case cancel +} + +/// A value that can be returned in elicitation form content. +public enum ElicitValue: Hashable, Sendable { + case string(String) + case int(Int) + case double(Double) + case bool(Bool) + case strings([String]) +} + +extension ElicitValue: Codable { + public init(from decoder: Decoder) throws { + let container = try decoder.singleValueContainer() + + // Try to decode as each type + if let value = try? container.decode(Bool.self) { + self = .bool(value) + } else if let value = try? container.decode(Int.self) { + self = .int(value) + } else if let value = try? container.decode(Double.self) { + self = .double(value) + } else if let value = try? container.decode(String.self) { + self = .string(value) + } else if let value = try? container.decode([String].self) { + self = .strings(value) + } else { + throw DecodingError.typeMismatch( + ElicitValue.self, + DecodingError.Context( + codingPath: decoder.codingPath, + debugDescription: "Expected String, Int, Double, Bool, or [String]")) + } + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.singleValueContainer() + switch self { + case .string(let value): + try container.encode(value) + case .int(let value): + try container.encode(value) + case .double(let value): + try container.encode(value) + case .bool(let value): + try container.encode(value) + case .strings(let value): + try container.encode(value) + } + } +} + +/// The result of an elicitation request. +public struct ElicitResult: ResultWithExtraFields { + public typealias ResultCodingKeys = CodingKeys + + /// The user action in response to the elicitation. + public var action: ElicitAction + /// The submitted form data, only present when action is "accept". + public var content: [String: ElicitValue]? + /// Reserved for clients and servers to attach additional metadata. + public var _meta: [String: Value]? + /// Additional fields not defined in the schema (for forward compatibility). + public var extraFields: [String: Value]? + + public init( + action: ElicitAction, + content: [String: ElicitValue]? = nil, + _meta: [String: Value]? = nil, + extraFields: [String: Value]? = nil + ) { + self.action = action + self.content = content + self._meta = _meta + self.extraFields = extraFields + } + + public enum CodingKeys: String, CodingKey, CaseIterable { + case action, content, _meta + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + action = try container.decode(ElicitAction.self, forKey: .action) + content = try container.decodeIfPresent([String: ElicitValue].self, forKey: .content) + _meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + extraFields = try Self.decodeExtraFields(from: decoder) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode(action, forKey: .action) + try container.encodeIfPresent(content, forKey: .content) + try container.encodeIfPresent(_meta, forKey: ._meta) + try encodeExtraFields(to: encoder) + } +} + +// MARK: - URL Mode Elicitation + +/// Parameters for a URL-mode elicitation request. +/// +/// URL mode is used for out-of-band flows like OAuth or credential collection, +/// where the user needs to navigate to an external URL. +public struct ElicitRequestURLParams: Hashable, Codable, Sendable { + /// The elicitation mode (must be "url"). + public let mode: String + /// The message to present to the user explaining why the interaction is needed. + public var message: String + /// The ID of the elicitation, which must be unique within the context of the server. + /// The client MUST treat this ID as an opaque value. + public var elicitationId: String + /// The URL that the user should navigate to. + public var url: String + /// Request metadata including progress token. + public var _meta: RequestMeta? + /// Task augmentation metadata. If present, the receiver should run the elicitation + /// as a background task and return `CreateTaskResult` instead of `ElicitResult`. + public var task: TaskMetadata? + + public init( + message: String, + elicitationId: String, + url: String, + _meta: RequestMeta? = nil, + task: TaskMetadata? = nil + ) { + self.mode = "url" + self.message = message + self.elicitationId = elicitationId + self.url = url + self._meta = _meta + self.task = task + } +} + +/// Parameters for elicitation requests (either form or URL mode). +public enum ElicitRequestParams: Hashable, Sendable { + case form(ElicitRequestFormParams) + case url(ElicitRequestURLParams) +} + +extension ElicitRequestParams: Codable { + private enum CodingKeys: String, CodingKey { + case mode + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + let mode = try container.decodeIfPresent(String.self, forKey: .mode) ?? "form" + + switch mode { + case "url": + let params = try ElicitRequestURLParams(from: decoder) + self = .url(params) + default: + let params = try ElicitRequestFormParams(from: decoder) + self = .form(params) + } + } + + public func encode(to encoder: Encoder) throws { + switch self { + case .form(let params): + try params.encode(to: encoder) + case .url(let params): + try params.encode(to: encoder) + } + } +} + +/// Notification from the server to the client, informing it of completion +/// of an out-of-band (URL mode) elicitation request. +/// +/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2025-11-25/ +public struct ElicitationCompleteNotification: Notification { + public static let name = "notifications/elicitation/complete" + + public struct Parameters: Hashable, Codable, Sendable { + /// The ID of the elicitation that completed. + public var elicitationId: String + /// Reserved for additional metadata. + public var _meta: [String: Value]? + + public init(elicitationId: String, _meta: [String: Value]? = nil) { + self.elicitationId = elicitationId + self._meta = _meta + } + } +} + +// MARK: - Error Codes + +/// Error code indicating URL elicitation is required. +/// +/// This error is returned when a server requires the client to perform +/// URL-mode elicitation but the client doesn't support it. +/// +/// - Note: Prefer using `ErrorCode.urlElicitationRequired` or +/// throwing `MCPError.urlElicitationRequired(elicitations:)` directly. +@available(*, deprecated, renamed: "ErrorCode.urlElicitationRequired") +public let URLElicitationRequiredErrorCode: Int = ErrorCode.urlElicitationRequired + +/// Error data for `URLElicitationRequiredError`. +/// +/// Servers return this when a request cannot be processed until one or more +/// URL mode elicitations are completed. The error response includes this data +/// in the `data` field with the error code `-32042`. +/// +/// Example error response: +/// ```json +/// { +/// "jsonrpc": "2.0", +/// "id": 2, +/// "error": { +/// "code": -32042, +/// "message": "This request requires more information.", +/// "data": { +/// "elicitations": [ +/// { +/// "mode": "url", +/// "elicitationId": "...", +/// "url": "https://example.com/...", +/// "message": "..." +/// } +/// ] +/// } +/// } +/// } +/// ``` +public struct ElicitationRequiredErrorData: Hashable, Codable, Sendable { + /// List of URL mode elicitations that must be completed. + public var elicitations: [ElicitRequestURLParams] + + public init(elicitations: [ElicitRequestURLParams]) { + self.elicitations = elicitations + } +} + +// MARK: - Method + +/// Server requests additional information from the user via the client. +/// +/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/client/elicitation/ +public enum Elicit: Method { + public static let name = "elicitation/create" + + public typealias Parameters = ElicitRequestParams + public typealias Result = ElicitResult +} diff --git a/Sources/MCP/Client/Experimental/ExperimentalClientFeatures.swift b/Sources/MCP/Client/Experimental/ExperimentalClientFeatures.swift new file mode 100644 index 00000000..4c341ea3 --- /dev/null +++ b/Sources/MCP/Client/Experimental/ExperimentalClientFeatures.swift @@ -0,0 +1,330 @@ +import Foundation + +// MARK: - Task Stream Message + +/// Messages yielded by streaming task operations. +/// +/// Similar to TypeScript SDK's `ResponseMessage`, this enum represents the different +/// types of messages that can occur during task-augmented tool execution. +/// +/// ## Example +/// +/// ```swift +/// for try await message in await client.experimental.tasks.callToolStream(name: "myTool") { +/// switch message { +/// case .taskCreated(let task): +/// print("Task started: \(task.taskId)") +/// case .taskStatus(let task): +/// print("Status: \(task.status)") +/// case .result(let result): +/// print("Tool completed with \(result.content.count) content blocks") +/// case .error(let error): +/// print("Error: \(error.localizedDescription)") +/// } +/// } +/// ``` +public enum TaskStreamMessage: Sendable { + /// A task has been created. This is always the first message for task-augmented requests. + case taskCreated(MCPTask) + + /// The task status has changed. Yielded when polling detects a status update. + case taskStatus(MCPTask) + + /// The task completed successfully with a result. + case result(CallTool.Result) + + /// The task or request encountered an error. + case error(MCPError) +} + +// MARK: - Experimental Client Features + +/// Experimental APIs for MCP clients. +/// +/// Access via `client.experimental.tasks`: +/// ```swift +/// // Call a tool as a task +/// let createResult = try await client.experimental.tasks.callToolAsTask( +/// name: "long_running_tool", +/// arguments: ["input": .string("data")] +/// ) +/// +/// // Get task status +/// let status = try await client.experimental.tasks.getTask(createResult.task.taskId) +/// +/// // Get task result when complete +/// let result = try await client.experimental.tasks.getTaskResult(taskId) +/// +/// // Poll for completion +/// for try await status in await client.experimental.tasks.pollTask(taskId) { +/// print("Status: \(status.status)") +/// } +/// ``` +/// +/// - Warning: These APIs are experimental and may change without notice. +public struct ExperimentalClientFeatures: Sendable { + private let client: Client + + init(client: Client) { + self.client = client + } + + /// Task-related experimental APIs. + public var tasks: ExperimentalClientTasks { + ExperimentalClientTasks(client: client) + } +} + +/// Experimental task APIs for MCP clients. +/// +/// - Warning: These APIs are experimental and may change without notice. +public struct ExperimentalClientTasks: Sendable { + private let client: Client + + init(client: Client) { + self.client = client + } + + /// Get the current status of a task. + /// + /// - Parameter taskId: The task identifier + /// - Returns: The task status information + /// - Throws: MCPError if the server doesn't support tasks or if the task is not found + public func getTask(_ taskId: String) async throws -> GetTask.Result { + try await client.getTask(taskId: taskId) + } + + /// List all tasks. + /// + /// - Parameter cursor: Optional pagination cursor + /// - Returns: Tuple of (tasks, nextCursor). nextCursor is nil if no more pages. + /// - Throws: MCPError if the server doesn't support tasks + public func listTasks(cursor: String? = nil) async throws -> (tasks: [MCPTask], nextCursor: String?) { + try await client.listTasks(cursor: cursor) + } + + /// Cancel a running task. + /// + /// - Parameter taskId: The task identifier + /// - Returns: The updated task status after cancellation + /// - Throws: MCPError if the server doesn't support tasks or if the task is not found + public func cancelTask(_ taskId: String) async throws -> CancelTask.Result { + try await client.cancelTask(taskId: taskId) + } + + /// Get the result payload of a completed task. + /// + /// The result type depends on the original request that created the task + /// (e.g., a tool call result for a task created from `tools/call`). + /// + /// - Note: For non-terminal tasks, this will block until the task completes. + /// The server implements long-polling behavior. + /// + /// - Parameter taskId: The task identifier + /// - Returns: The task result payload + /// - Throws: MCPError if the server doesn't support tasks or if the task is not found + public func getTaskResult(_ taskId: String) async throws -> GetTaskPayload.Result { + try await client.getTaskResult(taskId: taskId) + } + + /// Get the result payload of a completed task, decoded as a specific type. + /// + /// This is a convenience method that retrieves the task result and decodes it + /// as the expected result type. Use this when you know the type of result + /// the task will produce. + /// + /// - Note: For non-terminal tasks, this will block until the task completes. + /// The server implements long-polling behavior. + /// + /// ## Example + /// + /// ```swift + /// // Start a task-augmented tool call + /// let createResult = try await client.experimental.tasks.callToolAsTask( + /// name: "long_running_tool", + /// arguments: ["input": .string("data")] + /// ) + /// + /// // Get the result decoded as CallTool.Result + /// let result: CallTool.Result = try await client.experimental.tasks.getTaskResult( + /// createResult.task.taskId, + /// as: CallTool.Result.self + /// ) + /// print("Tool returned \(result.content.count) content blocks") + /// ``` + /// + /// - Parameters: + /// - taskId: The task identifier + /// - type: The type to decode the result as + /// - Returns: The task result decoded as the specified type + /// - Throws: MCPError if the server doesn't support tasks, if the task is not found, + /// or DecodingError if the result cannot be decoded as the specified type + public func getTaskResult(_ taskId: String, as type: T.Type) async throws -> T { + try await client.getTaskResultAs(taskId: taskId, type: type) + } + + /// Get the tool result of a completed task. + /// + /// This is a convenience method specifically for tasks created from `callToolAsTask()`. + /// It retrieves and decodes the result as `CallTool.Result`. + /// + /// - Note: For non-terminal tasks, this will block until the task completes. + /// The server implements long-polling behavior. + /// + /// ## Example + /// + /// ```swift + /// let createResult = try await client.experimental.tasks.callToolAsTask( + /// name: "process_data", + /// arguments: ["input": .string("data")] + /// ) + /// + /// let toolResult = try await client.experimental.tasks.getToolResult(createResult.task.taskId) + /// for content in toolResult.content { + /// // Process content blocks + /// } + /// ``` + /// + /// - Parameter taskId: The task identifier + /// - Returns: The tool call result + /// - Throws: MCPError if the server doesn't support tasks or if the task is not found + public func getToolResult(_ taskId: String) async throws -> CallTool.Result { + try await client.getTaskResultAs(taskId: taskId, type: CallTool.Result.self) + } + + /// Call a tool as a task, returning immediately with a task reference. + /// + /// This is the recommended way to call tools that may take a long time to complete. + /// Instead of waiting for the result, this method returns a `CreateTaskResult` + /// containing the task ID. You can then poll for the result using `getTaskResult()`. + /// + /// ## Example + /// + /// ```swift + /// // Start the task + /// let createResult = try await client.experimental.tasks.callToolAsTask( + /// name: "long_running_tool", + /// arguments: ["input": .string("data")], + /// ttl: 60000 // Keep results for 60 seconds + /// ) + /// print("Task started: \(createResult.task.taskId)") + /// + /// // Poll for result + /// let result = try await client.experimental.tasks.getTaskResult(createResult.task.taskId) + /// ``` + /// + /// - Parameters: + /// - name: The name of the tool to call + /// - arguments: Optional arguments for the tool + /// - ttl: Optional time-to-live in milliseconds for the task result + /// - Returns: The created task information + /// - Throws: MCPError if the server doesn't support tasks or the request fails + public func callToolAsTask( + name: String, + arguments: [String: Value]? = nil, + ttl: Int? = nil + ) async throws -> CreateTaskResult { + try await client.callToolAsTask(name: name, arguments: arguments, ttl: ttl) + } + + /// Poll a task until it reaches a terminal state. + /// + /// This method repeatedly polls the task status until it reaches a terminal + /// state (completed, failed, or cancelled). It yields each status update as + /// it occurs. + /// + /// The polling respects the server's suggested `pollInterval` if provided, + /// otherwise defaults to 1 second. + /// + /// ## Example + /// + /// ```swift + /// for try await status in await client.experimental.tasks.pollTask(taskId) { + /// print("Status: \(status.status)") + /// if status.status == .inputRequired { + /// // Handle user input request + /// } + /// } + /// // Task is now terminal - get the result + /// let result = try await client.experimental.tasks.getTaskResult(taskId) + /// ``` + /// + /// - Parameter taskId: The task identifier + /// - Returns: An async stream that yields task status updates until terminal + /// - Throws: MCPError if polling fails + public func pollTask(_ taskId: String) async -> AsyncThrowingStream { + await client.pollTask(taskId: taskId) + } + + /// Wait for a task to reach a terminal state. + /// + /// This is a convenience method that polls the task and returns only + /// when it has completed, failed, or been cancelled. + /// + /// - Parameter taskId: The task identifier + /// - Returns: The final task status + /// - Throws: MCPError if polling fails + public func pollUntilTerminal(_ taskId: String) async throws -> GetTask.Result { + try await client.pollUntilTerminal(taskId: taskId) + } + + /// Call a tool as a task and wait for the result. + /// + /// This is a convenience method that combines `callToolAsTask()` and + /// `getTaskResult()`. It starts the task and waits for the result. + /// + /// - Parameters: + /// - name: The name of the tool to call + /// - arguments: Optional arguments for the tool + /// - ttl: Optional time-to-live in milliseconds for the task result + /// - Returns: The tool result (same as `callTool()`) + /// - Throws: MCPError if the request fails or the task fails + public func callToolAsTaskAndWait( + name: String, + arguments: [String: Value]? = nil, + ttl: Int? = nil + ) async throws -> (content: [Tool.Content], isError: Bool?) { + try await client.callToolAsTaskAndWait(name: name, arguments: arguments, ttl: ttl) + } + + /// Call a tool as a task and stream status updates until completion. + /// + /// This method provides streaming access to tool execution, allowing you to + /// observe intermediate task status updates for long-running tool calls. + /// It combines `callToolAsTask()`, `pollTask()`, and `getTaskResult()` into + /// a single stream that yields all events. + /// + /// The stream is guaranteed to end with either a `.result` or `.error` message. + /// + /// This is similar to TypeScript SDK's `callToolStream` method. + /// + /// ## Example + /// + /// ```swift + /// for try await message in await client.experimental.tasks.callToolStream(name: "myTool") { + /// switch message { + /// case .taskCreated(let task): + /// print("Task started: \(task.taskId)") + /// case .taskStatus(let task): + /// print("Status: \(task.status), message: \(task.statusMessage ?? "none")") + /// case .result(let result): + /// print("Tool completed with \(result.content.count) content blocks") + /// case .error(let error): + /// print("Error: \(error.localizedDescription)") + /// } + /// } + /// ``` + /// + /// - Parameters: + /// - name: The name of the tool to call + /// - arguments: Optional arguments for the tool + /// - ttl: Optional time-to-live in milliseconds for the task result + /// - Returns: An async stream that yields `TaskStreamMessage` values + public func callToolStream( + name: String, + arguments: [String: Value]? = nil, + ttl: Int? = nil + ) async -> AsyncThrowingStream { + await client.callToolStream(name: name, arguments: arguments, ttl: ttl) + } +} diff --git a/Sources/MCP/Client/Experimental/Tasks/ClientTaskSupport.swift b/Sources/MCP/Client/Experimental/Tasks/ClientTaskSupport.swift new file mode 100644 index 00000000..b7820dcf --- /dev/null +++ b/Sources/MCP/Client/Experimental/Tasks/ClientTaskSupport.swift @@ -0,0 +1,231 @@ +import Foundation + +// MARK: - Client Task Handlers + +/// Container for client-side task handlers. +/// +/// This allows clients to handle task requests from servers (bidirectional task support). +/// When a server initiates a task on the client, these handlers process the requests. +/// +/// - Important: This is an experimental API that may change without notice. +/// +/// ## Example +/// +/// ```swift +/// let handlers = ExperimentalClientTaskHandlers( +/// getTask: { taskId in +/// // Return task status +/// return GetTask.Result(task: myTask) +/// }, +/// listTasks: { cursor in +/// // Return list of tasks +/// return ListTasks.Result(tasks: myTasks) +/// } +/// ) +/// +/// let client = Client(name: "MyClient", version: "1.0") +/// client.enableTaskHandlers(handlers) +/// ``` +public struct ExperimentalClientTaskHandlers: Sendable { + /// Handler for `tasks/get` requests from the server. + public typealias GetTaskHandler = @Sendable (String) async throws -> GetTask.Result + + /// Handler for `tasks/list` requests from the server. + public typealias ListTasksHandler = @Sendable (String?) async throws -> ListTasks.Result + + /// Handler for `tasks/cancel` requests from the server. + public typealias CancelTaskHandler = @Sendable (String) async throws -> CancelTask.Result + + /// Handler for `tasks/result` requests from the server. + public typealias GetTaskPayloadHandler = @Sendable (String) async throws -> GetTaskPayload.Result + + /// Handler for task-augmented sampling requests from the server. + /// + /// This is called when the server sends a `sampling/createMessage` request with a `task` field, + /// indicating the client should run the sampling as a background task. + public typealias TaskAugmentedSamplingHandler = @Sendable (CreateSamplingMessage.Parameters, TaskMetadata) async throws -> CreateTaskResult + + /// Handler for task-augmented elicitation requests from the server. + /// + /// This is called when the server sends an `elicitation/create` request with a `task` field, + /// indicating the client should run the elicitation as a background task. + public typealias TaskAugmentedElicitationHandler = @Sendable (Elicit.Parameters, TaskMetadata) async throws -> CreateTaskResult + + /// Handler for `tasks/get` requests. + public var getTask: GetTaskHandler? + + /// Handler for `tasks/list` requests. + public var listTasks: ListTasksHandler? + + /// Handler for `tasks/cancel` requests. + public var cancelTask: CancelTaskHandler? + + /// Handler for `tasks/result` requests. + public var getTaskPayload: GetTaskPayloadHandler? + + /// Handler for task-augmented sampling requests. + public var taskAugmentedSampling: TaskAugmentedSamplingHandler? + + /// Handler for task-augmented elicitation requests. + public var taskAugmentedElicitation: TaskAugmentedElicitationHandler? + + /// Create empty task handlers. + public init() {} + + /// Create task handlers with specific implementations. + public init( + getTask: GetTaskHandler? = nil, + listTasks: ListTasksHandler? = nil, + cancelTask: CancelTaskHandler? = nil, + getTaskPayload: GetTaskPayloadHandler? = nil, + taskAugmentedSampling: TaskAugmentedSamplingHandler? = nil, + taskAugmentedElicitation: TaskAugmentedElicitationHandler? = nil + ) { + self.getTask = getTask + self.listTasks = listTasks + self.cancelTask = cancelTask + self.getTaskPayload = getTaskPayload + self.taskAugmentedSampling = taskAugmentedSampling + self.taskAugmentedElicitation = taskAugmentedElicitation + } + + /// Build the client tasks capability based on which handlers are implemented. + /// + /// - Returns: The capability declaration, or nil if no handlers are set + public func buildCapability() -> Client.Capabilities.Tasks? { + // Check if any handlers are set + let hasTaskHandlers = getTask != nil || listTasks != nil || cancelTask != nil || getTaskPayload != nil + let hasAugmentedHandlers = taskAugmentedSampling != nil || taskAugmentedElicitation != nil + + guard hasTaskHandlers || hasAugmentedHandlers else { + return nil + } + + var requests: Client.Capabilities.Tasks.Requests? + if hasAugmentedHandlers { + requests = .init( + sampling: taskAugmentedSampling != nil ? .init(createMessage: .init()) : nil, + elicitation: taskAugmentedElicitation != nil ? .init(create: .init()) : nil + ) + } + + return .init( + list: listTasks != nil ? .init() : nil, + cancel: cancelTask != nil ? .init() : nil, + requests: requests + ) + } +} + +// MARK: - Client Task Support + +/// Configuration for client-side task support. +/// +/// This enables clients to run tasks initiated by servers (bidirectional task support). +/// The client can optionally provide its own task store and message queue for +/// tracking and managing tasks. +/// +/// - Important: This is an experimental API that may change without notice. +public final class ClientTaskSupport: Sendable { + /// The task store for persisting task state. + public let store: any TaskStore + + /// The message queue for side-channel communication. + public let queue: any TaskMessageQueue + + /// The task handlers. + public let handlers: ExperimentalClientTaskHandlers + + /// Create client task support with custom store and queue. + /// + /// - Parameters: + /// - store: The task store implementation + /// - queue: The message queue implementation + /// - handlers: The task handlers + public init( + store: any TaskStore, + queue: any TaskMessageQueue, + handlers: ExperimentalClientTaskHandlers + ) { + self.store = store + self.queue = queue + self.handlers = handlers + } + + /// Create in-memory client task support. + /// + /// - Parameter handlers: The task handlers + /// - Returns: ClientTaskSupport configured with in-memory store and queue + public static func inMemory(handlers: ExperimentalClientTaskHandlers = .init()) -> ClientTaskSupport { + ClientTaskSupport( + store: InMemoryTaskStore(), + queue: InMemoryTaskMessageQueue(), + handlers: handlers + ) + } +} + +// MARK: - Client Extension + +extension Client { + /// Enable task handlers on this client. + /// + /// This registers handlers for task requests from the server, enabling + /// bidirectional task support where the server can initiate tasks on the client. + /// + /// This method also integrates task-augmented sampling and elicitation handlers + /// that are called when the server sends requests with a `task` field, expecting + /// `CreateTaskResult` instead of the normal result. + /// + /// - Important: This is an experimental API that may change without notice. + /// + /// - Parameter taskSupport: The client task support configuration + /// - Returns: Self for chaining + @discardableResult + public func enableTaskHandlers(_ taskSupport: ClientTaskSupport) -> Self { + let handlers = taskSupport.handlers + + // Update capabilities based on handlers + if let tasksCap = handlers.buildCapability() { + capabilities.tasks = tasksCap + } + + // Register handlers for task requests from server + if let getTaskHandler = handlers.getTask { + withRequestHandler(GetTask.self) { params, _ in + try await getTaskHandler(params.taskId) + } + } + + if let listTasksHandler = handlers.listTasks { + withRequestHandler(ListTasks.self) { params, _ in + try await listTasksHandler(params.cursor) + } + } + + if let cancelTaskHandler = handlers.cancelTask { + withRequestHandler(CancelTask.self) { params, _ in + try await cancelTaskHandler(params.taskId) + } + } + + if let getTaskPayloadHandler = handlers.getTaskPayload { + withRequestHandler(GetTaskPayload.self) { params, _ in + try await getTaskPayloadHandler(params.taskId) + } + } + + // Register task-augmented sampling/elicitation handlers + // These are stored separately and checked at dispatch time (Python SDK pattern) + // This ensures handlers can be registered in any order without losing task-awareness + if let taskAugmentedSampling = handlers.taskAugmentedSampling { + _setTaskAugmentedSamplingHandler(taskAugmentedSampling) + } + + if let taskAugmentedElicitation = handlers.taskAugmentedElicitation { + _setTaskAugmentedElicitationHandler(taskAugmentedElicitation) + } + + return self + } +} diff --git a/Sources/MCP/Client/Roots.swift b/Sources/MCP/Client/Roots.swift new file mode 100644 index 00000000..96b425cd --- /dev/null +++ b/Sources/MCP/Client/Roots.swift @@ -0,0 +1,135 @@ +/// Roots represent filesystem directories that the client has access to. +/// +/// Servers can request the list of roots from clients to understand +/// the scope of files they can work with. +/// +/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2025-11-25/client/roots/ + +/// A root directory that the client has access to. +/// +/// Roots allow clients to inform servers about which parts of the +/// filesystem are available for operations. +public struct Root: Hashable, Codable, Sendable { + /// The prefix required for all root URIs. + public static let requiredURIPrefix = "file://" + + /// The URI of the root. Must be a `file://` URI. + public let uri: String + + /// An optional human-readable name for the root. + public let name: String? + + /// Reserved for additional metadata. + public var _meta: [String: Value]? + + /// Creates a new root with the specified URI. + /// + /// - Parameters: + /// - uri: The URI of the root. Must start with `file://`. + /// - name: An optional human-readable name for the root. + /// - _meta: Optional metadata for the root. + /// - Precondition: `uri` must start with `file://`. + public init( + uri: String, + name: String? = nil, + _meta: [String: Value]? = nil + ) { + precondition( + uri.hasPrefix(Self.requiredURIPrefix), + "Root URI must start with '\(Self.requiredURIPrefix)', got: \(uri)" + ) + self.uri = uri + self.name = name + self._meta = _meta + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + let uri = try container.decode(String.self, forKey: .uri) + guard uri.hasPrefix(Self.requiredURIPrefix) else { + throw DecodingError.dataCorruptedError( + forKey: .uri, + in: container, + debugDescription: "Root URI must start with '\(Self.requiredURIPrefix)', got: \(uri)" + ) + } + self.uri = uri + self.name = try container.decodeIfPresent(String.self, forKey: .name) + self._meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + } + + private enum CodingKeys: String, CodingKey { + case uri, name, _meta + } +} + +/// Request from server to client to list available filesystem roots. +/// +/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2025-11-25/client/roots/ +public enum ListRoots: Method { + public static let name: String = "roots/list" + + public struct Parameters: NotRequired, Hashable, Codable, Sendable { + /// Request metadata including progress token. + public var _meta: RequestMeta? + + public init() { + self._meta = nil + } + + public init(_meta: RequestMeta?) { + self._meta = _meta + } + } + + public struct Result: ResultWithExtraFields { + public typealias ResultCodingKeys = CodingKeys + + /// The list of available roots. + public let roots: [Root] + /// Reserved for clients and servers to attach additional metadata. + public var _meta: [String: Value]? + /// Additional fields not defined in the schema (for forward compatibility). + public var extraFields: [String: Value]? + + public init( + roots: [Root], + _meta: [String: Value]? = nil, + extraFields: [String: Value]? = nil + ) { + self.roots = roots + self._meta = _meta + self.extraFields = extraFields + } + + public enum CodingKeys: String, CodingKey, CaseIterable { + case roots, _meta + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + roots = try container.decode([Root].self, forKey: .roots) + _meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + extraFields = try Self.decodeExtraFields(from: decoder) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode(roots, forKey: .roots) + try container.encodeIfPresent(_meta, forKey: ._meta) + try encodeExtraFields(to: encoder) + } + } +} + +/// Notification sent by clients when the list of available roots changes. +/// +/// Servers that receive this notification should request an updated +/// list of roots via `ListRoots`. +/// +/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2025-11-25/client/roots/ +public struct RootsListChangedNotification: Notification { + public static let name: String = "notifications/roots/list_changed" + + public typealias Parameters = NotificationParams +} diff --git a/Sources/MCP/Client/Sampling.swift b/Sources/MCP/Client/Sampling.swift index 46563985..4169471c 100644 --- a/Sources/MCP/Client/Sampling.swift +++ b/Sources/MCP/Client/Sampling.swift @@ -1,120 +1,241 @@ import Foundation +/// Controls how the model uses tools during sampling. +/// +/// This allows servers to influence whether the model should use tools in its response. +/// The client may ignore this preference if it doesn't support the requested mode. +public struct ToolChoice: Hashable, Codable, Sendable { + /// How tools should be used during sampling. + public enum Mode: String, Hashable, Codable, Sendable { + /// Model decides whether to use tools (default). + case auto + /// Model MUST use at least one tool before completing. + case required + /// Model MUST NOT use any tools. + case none + } + + /// The tool choice mode. If nil, defaults to `.auto`. + public var mode: Mode? + + public init(mode: Mode? = nil) { + self.mode = mode + } +} + +/// Stop reason for sampling completion. +/// +/// This is an open string type to allow for provider-specific stop reasons. +/// Standard values are: `endTurn`, `stopSequence`, `maxTokens`, `toolUse`. +public struct StopReason: RawRepresentable, Hashable, Codable, Sendable { + public let rawValue: String + + public init(rawValue: String) { + self.rawValue = rawValue + } + + /// Natural end of turn + public static let endTurn = StopReason(rawValue: "endTurn") + /// Hit a stop sequence + public static let stopSequence = StopReason(rawValue: "stopSequence") + /// Reached maximum tokens + public static let maxTokens = StopReason(rawValue: "maxTokens") + /// Model decided to use a tool + public static let toolUse = StopReason(rawValue: "toolUse") +} + +/// Model preferences for sampling requests. +public struct ModelPreferences: Hashable, Codable, Sendable { + /// A hint suggesting a model name or family. + public struct Hint: Hashable, Codable, Sendable { + public let name: String? + public init(name: String? = nil) { self.name = name } + } + + public let hints: [Hint]? + public let costPriority: UnitInterval? + public let speedPriority: UnitInterval? + 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 + } +} + +// MARK: - Sampling Namespace + /// 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 - } + public typealias Role = MCP.Role - /// The message role public let role: Role - /// The message content - public let content: Content + public let content: [ContentBlock] + public var _meta: [String: Value]? - /// Creates a message with the specified role and content - @available( - *, deprecated, message: "Use static factory methods .user(_:) or .assistant(_:) instead" - ) - public init(role: Role, content: Content) { + public init(role: Role, content: ContentBlock, _meta: [String: Value]? = nil) { self.role = role - self.content = content + self.content = [content] + self._meta = _meta } - /// Private initializer for convenience methods to avoid deprecation warnings - private init(_role role: Role, _content content: Content) { + public init(role: Role, content: [ContentBlock], _meta: [String: Value]? = nil) { self.role = role self.content = content + self._meta = _meta } - /// Creates a user message with the specified content - public static func user(_ content: Content) -> Message { - return Message(_role: .user, _content: content) + public static func user(_ content: ContentBlock) -> Message { + Message(role: .user, content: content) } - /// Creates an assistant message with the specified content - public static func assistant(_ content: Content) -> Message { - return Message(_role: .assistant, _content: content) + public static func user(_ content: [ContentBlock]) -> Message { + Message(role: .user, content: content) } - /// Content types for sampling messages - public enum Content: Hashable, Sendable { - /// Text content - case text(String) - /// Image content - case image(data: String, mimeType: String) + public static func assistant(_ content: ContentBlock) -> Message { + Message(role: .assistant, content: content) } - } - /// 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 static func assistant(_ content: [ContentBlock]) -> Message { + Message(role: .assistant, content: content) + } - public init(name: String? = nil) { - self.name = name + private enum CodingKeys: String, CodingKey { + case role, content, _meta + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + role = try container.decode(Role.self, forKey: .role) + + // Content can be a single block or an array of blocks + if var arrayContainer = try? container.nestedUnkeyedContainer(forKey: .content) { + var blocks: [ContentBlock] = [] + while !arrayContainer.isAtEnd { + blocks.append(try arrayContainer.decode(ContentBlock.self)) + } + content = blocks + } else { + content = [try container.decode(ContentBlock.self, forKey: .content)] } + + _meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) } - /// 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 func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode(role, forKey: .role) - 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 + // Encode as single object if one block, array if multiple + if content.count == 1, let block = content.first { + try container.encode(block, forKey: .content) + } else { + try container.encode(content, forKey: .content) + } + + try container.encodeIfPresent(_meta, forKey: ._meta) + } + + /// Content block types for sampling messages. + public enum ContentBlock: Hashable, Sendable { + case text(String, annotations: Annotations?, _meta: [String: Value]?) + case image(data: String, mimeType: String, annotations: Annotations?, _meta: [String: Value]?) + case audio(data: String, mimeType: String, annotations: Annotations?, _meta: [String: Value]?) + case toolUse(ToolUseContent) + case toolResult(ToolResultContent) + + public static func text(_ text: String) -> ContentBlock { + .text(text, annotations: nil, _meta: nil) + } + + public static func image(data: String, mimeType: String) -> ContentBlock { + .image(data: data, mimeType: mimeType, annotations: nil, _meta: nil) + } + + public static func audio(data: String, mimeType: String) -> ContentBlock { + .audio(data: data, mimeType: mimeType, annotations: nil, _meta: nil) + } + + /// Whether this is a basic content type (text, image, or audio). + public var isBasicContent: Bool { + switch self { + case .text, .image, .audio: return true + case .toolUse, .toolResult: return false + } + } } } - /// Context inclusion options for sampling requests + public typealias ModelPreferences = MCP.ModelPreferences + 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 + public typealias StopReason = MCP.StopReason +} + +// MARK: - Tool Use Content + +public struct ToolUseContent: Hashable, Codable, Sendable { + public let type: String + public var name: String + public var id: String + public var input: [String: Value] + public var _meta: [String: Value]? + + public init(name: String, id: String, input: [String: Value], _meta: [String: Value]? = nil) { + self.type = "tool_use" + self.name = name + self.id = id + self.input = input + self._meta = _meta + } +} + +public struct ToolResultContent: Hashable, Codable, Sendable { + public let type: String + public var toolUseId: String + public var content: [Tool.Content] + public var structuredContent: Value? + public var isError: Bool? + public var _meta: [String: Value]? + + public init( + toolUseId: String, + content: [Tool.Content] = [], + structuredContent: Value? = nil, + isError: Bool? = nil, + _meta: [String: Value]? = nil + ) { + self.type = "tool_result" + self.toolUseId = toolUseId + self.content = content + self.structuredContent = structuredContent + self.isError = isError + self._meta = _meta } } -// MARK: - Codable +// MARK: - ContentBlock Codable -extension Sampling.Message.Content: Codable { +extension Sampling.Message.ContentBlock: Codable { private enum CodingKeys: String, CodingKey { - case type, text, data, mimeType + case type, text, data, mimeType, annotations, _meta } public init(from decoder: Decoder) throws { @@ -123,74 +244,239 @@ extension Sampling.Message.Content: Codable { switch type { case "text": - let text = try container.decode(String.self, forKey: .text) - self = .text(text) + self = .text( + try container.decode(String.self, forKey: .text), + annotations: try container.decodeIfPresent(Annotations.self, forKey: .annotations), + _meta: try container.decodeIfPresent([String: Value].self, forKey: ._meta) + ) 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) + self = .image( + data: try container.decode(String.self, forKey: .data), + mimeType: try container.decode(String.self, forKey: .mimeType), + annotations: try container.decodeIfPresent(Annotations.self, forKey: .annotations), + _meta: try container.decodeIfPresent([String: Value].self, forKey: ._meta) + ) + case "audio": + self = .audio( + data: try container.decode(String.self, forKey: .data), + mimeType: try container.decode(String.self, forKey: .mimeType), + annotations: try container.decodeIfPresent(Annotations.self, forKey: .annotations), + _meta: try container.decodeIfPresent([String: Value].self, forKey: ._meta) + ) + case "tool_use": + self = .toolUse(try ToolUseContent(from: decoder)) + case "tool_result": + self = .toolResult(try ToolResultContent(from: decoder)) default: throw DecodingError.dataCorruptedError( forKey: .type, in: container, - debugDescription: "Unknown sampling message content type") + debugDescription: "Unknown content type: \(type)") } } public func encode(to encoder: Encoder) throws { - var container = encoder.container(keyedBy: CodingKeys.self) - switch self { - case .text(let text): + case .text(let text, let annotations, let meta): + var container = encoder.container(keyedBy: CodingKeys.self) try container.encode("text", forKey: .type) try container.encode(text, forKey: .text) - case .image(let data, let mimeType): + try container.encodeIfPresent(annotations, forKey: .annotations) + try container.encodeIfPresent(meta, forKey: ._meta) + case .image(let data, let mimeType, let annotations, let meta): + var container = encoder.container(keyedBy: CodingKeys.self) try container.encode("image", forKey: .type) try container.encode(data, forKey: .data) try container.encode(mimeType, forKey: .mimeType) + try container.encodeIfPresent(annotations, forKey: .annotations) + try container.encodeIfPresent(meta, forKey: ._meta) + case .audio(let data, let mimeType, let annotations, let meta): + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode("audio", forKey: .type) + try container.encode(data, forKey: .data) + try container.encode(mimeType, forKey: .mimeType) + try container.encodeIfPresent(annotations, forKey: .annotations) + try container.encodeIfPresent(meta, forKey: ._meta) + case .toolUse(let toolUse): + try toolUse.encode(to: encoder) + case .toolResult(let toolResult): + try toolResult.encode(to: encoder) } } } -// MARK: - ExpressibleByStringLiteral - -extension Sampling.Message.Content: ExpressibleByStringLiteral { +extension Sampling.Message.ContentBlock: ExpressibleByStringLiteral { public init(stringLiteral value: String) { - self = .text(value) + self = .text(value, annotations: nil, _meta: nil) } } -// MARK: - ExpressibleByStringInterpolation - -extension Sampling.Message.Content: ExpressibleByStringInterpolation { +extension Sampling.Message.ContentBlock: ExpressibleByStringInterpolation { public init(stringInterpolation: DefaultStringInterpolation) { - self = .text(String(stringInterpolation: stringInterpolation)) + self = .text(String(stringInterpolation: stringInterpolation), annotations: nil, _meta: nil) + } +} + +extension Sampling.Message { + /// Type alias for backwards compatibility. + @available(*, deprecated, renamed: "ContentBlock") + public typealias Content = ContentBlock +} + +// MARK: - Sampling Request Parameters (Shared Base) + +/// Common parameters for sampling requests. +/// +/// This struct contains all the shared fields between tool and non-tool sampling requests. +public struct SamplingParameters: Hashable, Codable, Sendable { + public let messages: [Sampling.Message] + public let modelPreferences: Sampling.ModelPreferences? + public let systemPrompt: String? + public let includeContext: Sampling.ContextInclusion? + public let temperature: Double? + public let maxTokens: Int + public let stopSequences: [String]? + public let metadata: [String: Value]? + public var _meta: RequestMeta? + public var task: TaskMetadata? + + 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, + _meta: RequestMeta? = nil, + task: TaskMetadata? = 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 + self._meta = _meta + self.task = task } } -// MARK: - +// MARK: - CreateSamplingMessage (without tools) -/// To request sampling from a client, servers send a `sampling/createMessage` request. -/// - SeeAlso: https://modelcontextprotocol.io/docs/concepts/sampling#how-sampling-works +/// Request sampling from a client without tool support. +/// +/// The result will be a single content block (text, image, or audio). +/// For tool-enabled sampling, use `CreateSamplingMessageWithTools` instead. public enum CreateSamplingMessage: Method { public static let name = "sampling/createMessage" + /// Alias for shared parameters (no tools). + public typealias Parameters = SamplingParameters + + /// Result for a sampling request without tools. + /// Content is a single basic block (text, image, or audio). + public struct Result: ResultWithExtraFields { + public typealias ResultCodingKeys = CodingKeys + + public let model: String + public let stopReason: StopReason? + public let role: Role + /// Single content block (text, image, or audio - no tool use). + public let content: Sampling.Message.ContentBlock + public var _meta: [String: Value]? + public var extraFields: [String: Value]? + + public init( + model: String, + stopReason: StopReason? = nil, + role: Role, + content: Sampling.Message.ContentBlock, + _meta: [String: Value]? = nil, + extraFields: [String: Value]? = nil + ) { + self.model = model + self.stopReason = stopReason + self.role = role + self.content = content + self._meta = _meta + self.extraFields = extraFields + } + + public enum CodingKeys: String, CodingKey, CaseIterable { + case model, stopReason, role, content, _meta + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + model = try container.decode(String.self, forKey: .model) + stopReason = try container.decodeIfPresent(StopReason.self, forKey: .stopReason) + role = try container.decode(Role.self, forKey: .role) + _meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + + // MCP spec allows content to be either a single block or an array. + // Try to decode as array first, then fall back to single block. + if var arrayContainer = try? container.nestedUnkeyedContainer(forKey: .content) { + var blocks: [Sampling.Message.ContentBlock] = [] + while !arrayContainer.isAtEnd { + blocks.append(try arrayContainer.decode(Sampling.Message.ContentBlock.self)) + } + guard let firstBlock = blocks.first else { + throw DecodingError.dataCorruptedError( + forKey: .content, in: container, + debugDescription: "Content array is empty") + } + content = firstBlock + } else { + content = try container.decode(Sampling.Message.ContentBlock.self, forKey: .content) + } + + extraFields = try Self.decodeExtraFields(from: decoder) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode(model, forKey: .model) + try container.encodeIfPresent(stopReason, forKey: .stopReason) + try container.encode(role, forKey: .role) + try container.encode(content, forKey: .content) + try container.encodeIfPresent(_meta, forKey: ._meta) + try encodeExtraFields(to: encoder) + } + } +} + +// MARK: - CreateSamplingMessageWithTools + +/// Request sampling from a client with tool support. +/// +/// The result may contain tool use content, and content can be an array for parallel tool calls. +/// Requires `ClientCapabilities.sampling.tools` to be declared. +public enum CreateSamplingMessageWithTools: Method { + public static let name = "sampling/createMessage" + + /// Parameters for a sampling request with tools. 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]? + /// Base sampling parameters. + public let base: SamplingParameters + /// Tools that the model may use during generation. + public let tools: [Tool] + /// Controls how the model uses tools. + public let toolChoice: ToolChoice? + + // Convenience accessors + public var messages: [Sampling.Message] { base.messages } + public var modelPreferences: Sampling.ModelPreferences? { base.modelPreferences } + public var systemPrompt: String? { base.systemPrompt } + public var includeContext: Sampling.ContextInclusion? { base.includeContext } + public var temperature: Double? { base.temperature } + public var maxTokens: Int { base.maxTokens } + public var stopSequences: [String]? { base.stopSequences } + public var metadata: [String: Value]? { base.metadata } + public var _meta: RequestMeta? { base._meta } + public var task: TaskMetadata? { base.task } public init( messages: [Sampling.Message], @@ -200,39 +486,296 @@ public enum CreateSamplingMessage: Method { temperature: Double? = nil, maxTokens: Int, stopSequences: [String]? = nil, - metadata: [String: Value]? = nil + metadata: [String: Value]? = nil, + tools: [Tool], + toolChoice: ToolChoice? = nil, + _meta: RequestMeta? = nil, + task: TaskMetadata? = 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 + self.base = SamplingParameters( + messages: messages, + modelPreferences: modelPreferences, + systemPrompt: systemPrompt, + includeContext: includeContext, + temperature: temperature, + maxTokens: maxTokens, + stopSequences: stopSequences, + metadata: metadata, + _meta: _meta, + task: task + ) + self.tools = tools + self.toolChoice = toolChoice + } + + // Custom coding to flatten the structure + private enum CodingKeys: String, CodingKey { + case messages, modelPreferences, systemPrompt, includeContext + case temperature, maxTokens, stopSequences, metadata + case tools, toolChoice, _meta, task + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + base = SamplingParameters( + messages: try container.decode([Sampling.Message].self, forKey: .messages), + modelPreferences: try container.decodeIfPresent(Sampling.ModelPreferences.self, forKey: .modelPreferences), + systemPrompt: try container.decodeIfPresent(String.self, forKey: .systemPrompt), + includeContext: try container.decodeIfPresent(Sampling.ContextInclusion.self, forKey: .includeContext), + temperature: try container.decodeIfPresent(Double.self, forKey: .temperature), + maxTokens: try container.decode(Int.self, forKey: .maxTokens), + stopSequences: try container.decodeIfPresent([String].self, forKey: .stopSequences), + metadata: try container.decodeIfPresent([String: Value].self, forKey: .metadata), + _meta: try container.decodeIfPresent(RequestMeta.self, forKey: ._meta), + task: try container.decodeIfPresent(TaskMetadata.self, forKey: .task) + ) + tools = try container.decode([Tool].self, forKey: .tools) + toolChoice = try container.decodeIfPresent(ToolChoice.self, forKey: .toolChoice) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode(base.messages, forKey: .messages) + try container.encodeIfPresent(base.modelPreferences, forKey: .modelPreferences) + try container.encodeIfPresent(base.systemPrompt, forKey: .systemPrompt) + try container.encodeIfPresent(base.includeContext, forKey: .includeContext) + try container.encodeIfPresent(base.temperature, forKey: .temperature) + try container.encode(base.maxTokens, forKey: .maxTokens) + try container.encodeIfPresent(base.stopSequences, forKey: .stopSequences) + try container.encodeIfPresent(base.metadata, forKey: .metadata) + try container.encode(tools, forKey: .tools) + try container.encodeIfPresent(toolChoice, forKey: .toolChoice) + try container.encodeIfPresent(base._meta, forKey: ._meta) + try container.encodeIfPresent(base.task, forKey: .task) } } - public struct Result: Hashable, Codable, Sendable { - /// Name of the model used + /// Result for a sampling request with tools. + public struct Result: ResultWithExtraFields { + public typealias ResultCodingKeys = CodingKeys + 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 let stopReason: StopReason? + public let role: Role + public let content: [Sampling.Message.ContentBlock] + public var _meta: [String: Value]? + public var extraFields: [String: Value]? public init( model: String, - stopReason: Sampling.StopReason? = nil, - role: Sampling.Message.Role, - content: Sampling.Message.Content + stopReason: StopReason? = nil, + role: Role, + content: Sampling.Message.ContentBlock, + _meta: [String: Value]? = nil, + extraFields: [String: Value]? = nil + ) { + self.model = model + self.stopReason = stopReason + self.role = role + self.content = [content] + self._meta = _meta + self.extraFields = extraFields + } + + public init( + model: String, + stopReason: StopReason? = nil, + role: Role, + content: [Sampling.Message.ContentBlock], + _meta: [String: Value]? = nil, + extraFields: [String: Value]? = nil ) { self.model = model self.stopReason = stopReason self.role = role self.content = content + self._meta = _meta + self.extraFields = extraFields + } + + public enum CodingKeys: String, CodingKey, CaseIterable { + case model, stopReason, role, content, _meta + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + model = try container.decode(String.self, forKey: .model) + stopReason = try container.decodeIfPresent(StopReason.self, forKey: .stopReason) + role = try container.decode(Role.self, forKey: .role) + _meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + + if var arrayContainer = try? container.nestedUnkeyedContainer(forKey: .content) { + var blocks: [Sampling.Message.ContentBlock] = [] + while !arrayContainer.isAtEnd { + blocks.append(try arrayContainer.decode(Sampling.Message.ContentBlock.self)) + } + content = blocks + } else { + content = [try container.decode(Sampling.Message.ContentBlock.self, forKey: .content)] + } + + extraFields = try Self.decodeExtraFields(from: decoder) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode(model, forKey: .model) + try container.encodeIfPresent(stopReason, forKey: .stopReason) + try container.encode(role, forKey: .role) + try container.encodeIfPresent(_meta, forKey: ._meta) + + if content.count == 1, let block = content.first { + try container.encode(block, forKey: .content) + } else { + try container.encode(content, forKey: .content) + } + + try encodeExtraFields(to: encoder) + } + } +} + +// MARK: - Client-Side Sampling Handler + +/// Parameters for client-side handling of sampling requests. +/// +/// This extends the base parameters with optional tools for client-side flexibility. +public struct ClientSamplingParameters: Hashable, Codable, Sendable { + public let base: SamplingParameters + public let tools: [Tool]? + public let toolChoice: ToolChoice? + + // Convenience accessors + public var messages: [Sampling.Message] { base.messages } + public var modelPreferences: Sampling.ModelPreferences? { base.modelPreferences } + public var systemPrompt: String? { base.systemPrompt } + public var includeContext: Sampling.ContextInclusion? { base.includeContext } + public var temperature: Double? { base.temperature } + public var maxTokens: Int { base.maxTokens } + public var stopSequences: [String]? { base.stopSequences } + public var metadata: [String: Value]? { base.metadata } + public var _meta: RequestMeta? { base._meta } + public var task: TaskMetadata? { base.task } + + /// Whether this request includes tool support. + public var hasTools: Bool { + tools != nil && !(tools?.isEmpty ?? true) + } + + 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, + tools: [Tool]? = nil, + toolChoice: ToolChoice? = nil, + _meta: RequestMeta? = nil, + task: TaskMetadata? = nil + ) { + self.base = SamplingParameters( + messages: messages, + modelPreferences: modelPreferences, + systemPrompt: systemPrompt, + includeContext: includeContext, + temperature: temperature, + maxTokens: maxTokens, + stopSequences: stopSequences, + metadata: metadata, + _meta: _meta, + task: task + ) + self.tools = tools + self.toolChoice = toolChoice + } + + private enum CodingKeys: String, CodingKey { + case messages, modelPreferences, systemPrompt, includeContext + case temperature, maxTokens, stopSequences, metadata + case tools, toolChoice, _meta, task + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + base = SamplingParameters( + messages: try container.decode([Sampling.Message].self, forKey: .messages), + modelPreferences: try container.decodeIfPresent(Sampling.ModelPreferences.self, forKey: .modelPreferences), + systemPrompt: try container.decodeIfPresent(String.self, forKey: .systemPrompt), + includeContext: try container.decodeIfPresent(Sampling.ContextInclusion.self, forKey: .includeContext), + temperature: try container.decodeIfPresent(Double.self, forKey: .temperature), + maxTokens: try container.decode(Int.self, forKey: .maxTokens), + stopSequences: try container.decodeIfPresent([String].self, forKey: .stopSequences), + metadata: try container.decodeIfPresent([String: Value].self, forKey: .metadata), + _meta: try container.decodeIfPresent(RequestMeta.self, forKey: ._meta), + task: try container.decodeIfPresent(TaskMetadata.self, forKey: .task) + ) + tools = try container.decodeIfPresent([Tool].self, forKey: .tools) + toolChoice = try container.decodeIfPresent(ToolChoice.self, forKey: .toolChoice) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode(base.messages, forKey: .messages) + try container.encodeIfPresent(base.modelPreferences, forKey: .modelPreferences) + try container.encodeIfPresent(base.systemPrompt, forKey: .systemPrompt) + try container.encodeIfPresent(base.includeContext, forKey: .includeContext) + try container.encodeIfPresent(base.temperature, forKey: .temperature) + try container.encode(base.maxTokens, forKey: .maxTokens) + try container.encodeIfPresent(base.stopSequences, forKey: .stopSequences) + try container.encodeIfPresent(base.metadata, forKey: .metadata) + try container.encodeIfPresent(tools, forKey: .tools) + try container.encodeIfPresent(toolChoice, forKey: .toolChoice) + try container.encodeIfPresent(base._meta, forKey: ._meta) + try container.encodeIfPresent(base.task, forKey: .task) + } +} + +/// Method type for client-side handling of sampling requests. +public enum ClientSamplingRequest: Method { + public static let name = "sampling/createMessage" + public typealias Parameters = ClientSamplingParameters + /// Reuse the tools-capable result type. + public typealias Result = CreateSamplingMessageWithTools.Result +} + +// MARK: - Message Validation + +extension Sampling.Message { + /// Validates the structure of tool_use/tool_result messages. + public static func validateToolUseResultMessages(_ messages: [Sampling.Message]) throws { + guard !messages.isEmpty else { return } + + let lastContent = messages[messages.count - 1].content + let hasToolResults = lastContent.contains { if case .toolResult = $0 { return true }; return false } + + let previousContent: [ContentBlock]? = messages.count >= 2 ? messages[messages.count - 2].content : nil + let hasPreviousToolUse = previousContent?.contains { if case .toolUse = $0 { return true }; return false } ?? false + + if hasToolResults { + let hasNonToolResult = lastContent.contains { if case .toolResult = $0 { return false }; return true } + if hasNonToolResult { + throw MCPError.invalidParams("The last message must contain only tool_result content if any is present") + } + + guard previousContent != nil else { + throw MCPError.invalidParams("tool_result requires a previous message containing tool_use") + } + + if !hasPreviousToolUse { + throw MCPError.invalidParams("tool_result blocks do not match any tool_use in the previous message") + } + } + + if hasPreviousToolUse, let previousContent { + let toolUseIds = Set(previousContent.compactMap { if case .toolUse(let c) = $0 { return c.id }; return nil }) + let toolResultIds = Set(lastContent.compactMap { if case .toolResult(let c) = $0 { return c.toolUseId }; return nil }) + + if toolUseIds != toolResultIds { + throw MCPError.invalidParams("IDs of tool_result blocks and tool_use blocks from previous message do not match") + } } } } diff --git a/Sources/MCP/Server/Completions.swift b/Sources/MCP/Server/Completions.swift new file mode 100644 index 00000000..eb0defec --- /dev/null +++ b/Sources/MCP/Server/Completions.swift @@ -0,0 +1,338 @@ +import Foundation + +/// Autocomplete functionality allows servers to provide argument completion +/// suggestions for prompts and resource templates. +/// +/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/server/utilities/completion/ + +/// A reference to a prompt for completion requests. +/// +/// Used in `completion/complete` requests to identify which prompt's arguments +/// should be autocompleted. +/// +/// - SeeAlso: ``CompletionReference`` +public struct PromptReference: Hashable, Codable, Sendable { + /// The type discriminator, always "ref/prompt". + public let type: String + /// The name of the prompt to get completions for. + public let name: String + /// A human-readable title for the prompt, intended for UI display. + /// If not provided, the `name` should be used for display. + public let title: String? + + public init(name: String, title: String? = nil) { + self.type = "ref/prompt" + self.name = name + self.title = title + } + + private enum CodingKeys: String, CodingKey { + case type, name, title + } +} + +/// A reference to a resource template for completion requests. +/// +/// Used in `completion/complete` requests to identify which resource template's +/// URI parameters should be autocompleted. +/// +/// - SeeAlso: ``CompletionReference`` +public struct ResourceTemplateReference: Hashable, Codable, Sendable { + /// The type discriminator, always "ref/resource". + public let type: String + /// The URI or URI template of the resource to get completions for. + public let uri: String + + public init(uri: String) { + self.type = "ref/resource" + self.uri = uri + } + + private enum CodingKeys: String, CodingKey { + case type, uri + } +} + +/// A reference type identifying what to provide completions for. +/// +/// Completion requests can provide suggestions for either: +/// - Prompt arguments (using ``PromptReference``) +/// - Resource template URI parameters (using ``ResourceTemplateReference``) +/// +/// ## Example +/// +/// ```swift +/// // Request completions for a prompt argument +/// let promptRef = CompletionReference.prompt(PromptReference(name: "greet")) +/// +/// // Request completions for a resource template parameter +/// let resourceRef = CompletionReference.resource(ResourceTemplateReference(uri: "file:///{path}")) +/// ``` +public enum CompletionReference: Hashable, Sendable { + /// Reference to a prompt for argument completion. + case prompt(PromptReference) + /// Reference to a resource template for URI parameter completion. + case resource(ResourceTemplateReference) +} + +extension CompletionReference: Codable { + private enum CodingKeys: String, CodingKey { + case type + } + + 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 "ref/prompt": + let ref = try PromptReference(from: decoder) + self = .prompt(ref) + case "ref/resource": + let ref = try ResourceTemplateReference(from: decoder) + self = .resource(ref) + default: + throw DecodingError.dataCorruptedError( + forKey: .type, in: container, + debugDescription: "Unknown reference type: \(type)") + } + } + + public func encode(to encoder: Encoder) throws { + switch self { + case .prompt(let ref): + try ref.encode(to: encoder) + case .resource(let ref): + try ref.encode(to: encoder) + } + } +} + +// MARK: - Completion Request + +/// The argument being completed in a completion request. +/// +/// This identifies which argument/parameter the user is currently typing +/// and provides the partial value for matching suggestions. +public struct CompletionArgument: Hashable, Codable, Sendable { + /// The name of the argument or URI template parameter being completed. + public let name: String + /// The current partial value to use for completion matching. + /// Servers should return suggestions that start with or contain this value. + public let value: String + + public init(name: String, value: String) { + self.name = name + self.value = value + } +} + +/// Additional context for completion requests. +/// +/// Provides previously-resolved argument values that can be used to filter +/// or customize completion suggestions. For example, when completing a file path, +/// the previously-selected directory could be used to show only files in that directory. +public struct CompletionContext: Hashable, Codable, Sendable { + /// Previously-resolved argument values in a URI template or prompt. + /// Keys are argument names, values are their resolved values. + public let arguments: [String: String]? + + public init(arguments: [String: String]? = nil) { + self.arguments = arguments + } +} + +// MARK: - Completion Result + +/// Completion suggestions returned by the server. +/// +/// Contains an array of suggested values for the argument being completed, +/// along with pagination information if there are more results available. +/// +/// ## Example +/// +/// ```swift +/// // Return filtered suggestions +/// let suggestions = CompletionSuggestions( +/// values: ["Alice", "Bob", "Charlie"], +/// total: 3, +/// hasMore: false +/// ) +/// +/// // Return partial results with more available +/// let partialSuggestions = CompletionSuggestions( +/// values: Array(allValues.prefix(100)), +/// total: allValues.count, +/// hasMore: allValues.count > 100 +/// ) +/// +/// // Or use the convenience initializer which handles truncation automatically +/// let autoTruncated = CompletionSuggestions(from: allValues) +/// ``` +public struct CompletionSuggestions: Hashable, Codable, Sendable { + /// The maximum number of values allowed per the MCP specification. + public static let maxValues = 100 + + /// An empty completion result, for use when no suggestions are available. + /// + /// Equivalent to `CompletionSuggestions(values: [], hasMore: false)`. + public static let empty = CompletionSuggestions(values: [], hasMore: false) + + /// An array of completion values. Must not exceed 100 items per the MCP spec. + public let values: [String] + /// The total number of completion options available. + /// This may exceed the number of values in the response if results are truncated. + public let total: Int? + /// Indicates whether there are additional completion options beyond + /// those provided in the current response, even if the exact total is unknown. + public let hasMore: Bool? + + /// Creates a completion suggestions result. + /// + /// - Parameters: + /// - values: The completion values. If more than 100 values are provided, + /// only the first 100 will be used per the MCP specification. + /// - total: The total number of completion options available. + /// - hasMore: Whether there are additional options beyond those provided. + /// + /// - Note: This initializer does not automatically set `total` or `hasMore` based on + /// the values array. Use ``init(from:)`` for automatic handling of these fields. + public init(values: [String], total: Int? = nil, hasMore: Bool? = nil) { + // Enforce the 100-item limit per MCP specification + self.values = Array(values.prefix(Self.maxValues)) + self.total = total + self.hasMore = hasMore + } + + /// Creates a completion suggestions result from an array of values, + /// automatically handling pagination fields. + /// + /// This convenience initializer: + /// - Truncates values to the maximum of 100 allowed by the MCP specification + /// - Sets `total` to the original count of all values + /// - Sets `hasMore` to indicate whether values were truncated + /// + /// ## Example + /// + /// ```swift + /// let allLanguages = ["python", "javascript", "typescript", "java", "go", "rust"] + /// let filtered = allLanguages.filter { $0.hasPrefix(partialValue) } + /// return Complete.Result(completion: CompletionSuggestions(from: filtered)) + /// ``` + /// + /// - Parameter allValues: All available completion values. If more than 100 values + /// are provided, only the first 100 will be returned. + public init(from allValues: [String]) { + let truncated = Array(allValues.prefix(Self.maxValues)) + self.values = truncated + self.total = allValues.count + self.hasMore = allValues.count > Self.maxValues + } +} + +// MARK: - Method + +/// A request from the client to the server, to ask for completion options. +/// +/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/server/utilities/completion/ +public enum Complete: Method { + public static let name = "completion/complete" + + public struct Parameters: Hashable, Codable, Sendable { + /// The reference to the prompt or resource template. + public let ref: CompletionReference + /// The argument information. + public let argument: CompletionArgument + /// Additional, optional context for completions. + public let context: CompletionContext? + /// Request metadata including progress token. + public var _meta: RequestMeta? + + public init( + ref: CompletionReference, + argument: CompletionArgument, + context: CompletionContext? = nil, + _meta: RequestMeta? = nil + ) { + self.ref = ref + self.argument = argument + self.context = context + self._meta = _meta + } + } + + public struct Result: ResultWithExtraFields { + public typealias ResultCodingKeys = CodingKeys + + /// An empty completion result, for use when no suggestions are available. + /// + /// Equivalent to `Complete.Result(completion: .empty)`. + public static let empty = Result(completion: .empty) + + /// The completion options. + public let completion: CompletionSuggestions + /// Reserved for clients and servers to attach additional metadata. + public var _meta: [String: Value]? + /// Additional fields not defined in the schema (for forward compatibility). + public var extraFields: [String: Value]? + + /// Creates a completion result. + /// + /// - Parameters: + /// - completion: The completion suggestions. + /// - _meta: Optional metadata. + /// - extraFields: Additional fields for forward compatibility. + public init( + completion: CompletionSuggestions, + _meta: [String: Value]? = nil, + extraFields: [String: Value]? = nil + ) { + self.completion = completion + self._meta = _meta + self.extraFields = extraFields + } + + /// Creates a completion result from an array of values, + /// automatically handling pagination. + /// + /// This convenience initializer: + /// - Truncates values to the maximum of 100 allowed by the MCP specification + /// - Sets `total` to the original count of all values + /// - Sets `hasMore` to indicate whether values were truncated + /// + /// ## Example + /// + /// ```swift + /// server.withRequestHandler(Complete.self) { params, _ in + /// let allLanguages = ["python", "javascript", "typescript", "java", "go", "rust"] + /// let filtered = allLanguages.filter { $0.hasPrefix(params.argument.value) } + /// return Complete.Result(from: filtered) + /// } + /// ``` + /// + /// - Parameter allValues: All available completion values. + public init(from allValues: [String]) { + self.completion = CompletionSuggestions(from: allValues) + self._meta = nil + self.extraFields = nil + } + + public enum CodingKeys: String, CodingKey, CaseIterable { + case completion, _meta + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + completion = try container.decode(CompletionSuggestions.self, forKey: .completion) + _meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + extraFields = try Self.decodeExtraFields(from: decoder) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode(completion, forKey: .completion) + try container.encodeIfPresent(_meta, forKey: ._meta) + try encodeExtraFields(to: encoder) + } + } +} diff --git a/Sources/MCP/Server/Experimental/ExperimentalServerFeatures.swift b/Sources/MCP/Server/Experimental/ExperimentalServerFeatures.swift new file mode 100644 index 00000000..d1b92081 --- /dev/null +++ b/Sources/MCP/Server/Experimental/ExperimentalServerFeatures.swift @@ -0,0 +1,217 @@ +import Foundation + +/// Experimental APIs for MCP servers. +/// +/// Access via `server.experimental.tasks`: +/// ```swift +/// // Enable task support with in-memory storage +/// await server.experimental.tasks.enable() +/// +/// // Or with custom configuration +/// let taskSupport = TaskSupport.inMemory() +/// await server.experimental.tasks.enable(taskSupport) +/// ``` +/// +/// - Warning: These APIs are experimental and may change without notice. +public struct ExperimentalServerFeatures: Sendable { + private let server: Server + + init(server: Server) { + self.server = server + } + + /// Task-related experimental APIs. + public var tasks: ExperimentalServerTasks { + ExperimentalServerTasks(server: server) + } +} + +/// Experimental task APIs for MCP servers. +/// +/// - Warning: These APIs are experimental and may change without notice. +public struct ExperimentalServerTasks: Sendable { + private let server: Server + + init(server: Server) { + self.server = server + } + + /// Enable task support with default in-memory storage. + /// + /// This is a convenience method that enables task support using + /// an in-memory task store and message queue. Suitable for + /// development, testing, and single-process servers. + /// + /// For production distributed systems, use `enable(_:)` with + /// custom `TaskSupport` configuration. + /// + /// This method: + /// 1. Sets the tasks capability with full support (list, cancel, task-augmented tools/call) + /// 2. Registers default handlers for `tasks/get`, `tasks/list`, `tasks/cancel`, and `tasks/result` + /// + /// - Returns: Self for chaining + @discardableResult + public func enable() async -> Server { + await server.enableTaskSupport(.inMemory()) + } + + /// Enable task support with custom configuration. + /// + /// Use this method when you need custom task storage (e.g., database + /// or distributed cache) or custom message queue implementations. + /// + /// This method: + /// 1. Sets the tasks capability with full support (list, cancel, task-augmented tools/call) + /// 2. Registers the TaskResultHandler as a response router for mid-task elicitation/sampling + /// 3. Registers default handlers for `tasks/get`, `tasks/list`, `tasks/cancel`, and `tasks/result` + /// + /// - Parameter taskSupport: The task support configuration + /// - Returns: Self for chaining + @discardableResult + public func enable(_ taskSupport: TaskSupport) async -> Server { + await server.enableTaskSupport(taskSupport) + } + + // MARK: - Client Task Polling (Server → Client) + + /// Get a task from the client. + /// + /// This sends a `tasks/get` request to the client to retrieve task status. + /// Used when the server has initiated a task-augmented request (like elicitAsTask) + /// and needs to check the client's task status. + /// + /// - Parameter taskId: The client-side task identifier + /// - Returns: The task status from the client + /// - Throws: MCPError if the client doesn't support tasks or the task is not found + public func getClientTask(_ taskId: String) async throws -> GetTask.Result { + try await server.getClientTask(taskId: taskId) + } + + /// Get the result payload of a client task. + /// + /// This sends a `tasks/result` request to the client to retrieve the task result. + /// For non-terminal tasks, this will block until the task completes. + /// + /// - Parameter taskId: The client-side task identifier + /// - Returns: The task result payload + /// - Throws: MCPError if the client doesn't support tasks or the task is not found + public func getClientTaskResult(_ taskId: String) async throws -> GetTaskPayload.Result { + try await server.getClientTaskResult(taskId: taskId) + } + + /// Get the result payload of a client task, decoded as a specific type. + /// + /// This is a convenience method that retrieves the task result and decodes it + /// as the expected result type. + /// + /// ## Example + /// + /// ```swift + /// // Send a task-augmented elicitation to the client + /// let createResult = try await server.experimental.tasks.elicitAsTask(...) + /// + /// // Get the result decoded as ElicitResult + /// let result: ElicitResult = try await server.experimental.tasks.getClientTaskResult( + /// createResult.task.taskId, + /// as: ElicitResult.self + /// ) + /// ``` + /// + /// - Parameters: + /// - taskId: The client-side task identifier + /// - type: The type to decode the result as + /// - Returns: The task result decoded as the specified type + /// - Throws: MCPError or DecodingError if the result cannot be decoded + public func getClientTaskResult( + _ taskId: String, + as type: T.Type + ) async throws -> T { + try await server.getClientTaskResultAs(taskId: taskId, type: type) + } + + /// Poll a client task until it reaches a terminal state. + /// + /// This method repeatedly polls the client for task status until the task + /// reaches a terminal state (completed, failed, or cancelled). + /// + /// The polling interval is determined by the `pollInterval` returned by the client, + /// defaulting to 500ms if not specified. + /// + /// ## Example + /// + /// ```swift + /// // Start a task-augmented request + /// let createResult = try await server.experimental.tasks.elicitAsTask(...) + /// + /// // Poll until complete + /// for try await status in server.experimental.tasks.pollClientTask(createResult.task.taskId) { + /// print("Task status: \(status)") + /// } + /// + /// // Get the final result + /// let result = try await server.experimental.tasks.getClientTaskResult( + /// createResult.task.taskId, + /// as: ElicitResult.self + /// ) + /// ``` + /// + /// - Parameter taskId: The client-side task identifier + /// - Returns: An async stream of task statuses, ending when terminal + public func pollClientTask(_ taskId: String) -> AsyncThrowingStream { + AsyncThrowingStream { continuation in + Task { + do { + while true { + let result = try await server.getClientTask(taskId: taskId) + continuation.yield(result.status) + + if result.status.isTerminal { + continuation.finish() + return + } + + // Wait for poll interval (default 500ms) + let intervalMs = result.pollInterval ?? 500 + try await Task.sleep(for: .milliseconds(intervalMs)) + } + } catch { + continuation.finish(throwing: error) + } + } + } + } + + /// Poll a client task until terminal, then return the final result. + /// + /// This is a convenience method that polls until the task completes and then + /// retrieves and decodes the result. + /// + /// ## Example + /// + /// ```swift + /// // Send a task-augmented elicitation and wait for result + /// let createResult = try await server.experimental.tasks.elicitAsTask(...) + /// let elicitResult: ElicitResult = try await server.experimental.tasks.pollClientTaskResult( + /// createResult.task.taskId, + /// as: ElicitResult.self + /// ) + /// ``` + /// + /// - Parameters: + /// - taskId: The client-side task identifier + /// - type: The type to decode the result as + /// - Returns: The task result decoded as the specified type + /// - Throws: MCPError or DecodingError if the result cannot be decoded + public func pollClientTaskResult( + _ taskId: String, + as type: T.Type + ) async throws -> T { + // Poll until terminal + for try await _ in pollClientTask(taskId) { + // Just consume the stream until terminal + } + + // Get the final result + return try await getClientTaskResult(taskId, as: type) + } +} diff --git a/Sources/MCP/Server/Experimental/Tasks/ServerTaskContext.swift b/Sources/MCP/Server/Experimental/Tasks/ServerTaskContext.swift new file mode 100644 index 00000000..cac8bd27 --- /dev/null +++ b/Sources/MCP/Server/Experimental/Tasks/ServerTaskContext.swift @@ -0,0 +1,936 @@ +import Foundation + +// MARK: - Server Task Context + +/// Context for task handlers to interact with the task lifecycle. +/// +/// This context is passed to task handlers and provides: +/// - Task status updates +/// - Task completion/failure +/// - Cancellation checking +/// - Mid-task elicitation and sampling +/// - Access to the underlying task and store +/// +/// - Important: This is an experimental API that may change without notice. +/// +/// ## Example +/// +/// ```swift +/// async func work(context: ServerTaskContext) async throws -> CallTool.Result { +/// try await context.updateStatus("Starting work...") +/// +/// // Check for cancellation periodically +/// guard !context.isCancelled else { +/// throw CancellationError() +/// } +/// +/// // Request user input mid-task +/// let result = try await context.elicit( +/// message: "Please confirm the operation", +/// requestedSchema: ElicitationSchema(properties: [ +/// "confirm": .boolean(BooleanSchema(title: "Confirm")) +/// ]) +/// ) +/// +/// if result.action != .accept { +/// throw CancellationError() +/// } +/// +/// return CallTool.Result(content: [.text("Done!")]) +/// } +/// ``` +public final class ServerTaskContext: @unchecked Sendable { + /// The task this context is for. + public private(set) var task: MCPTask + + /// The task store for persistence. + private let store: any TaskStore + + /// The message queue for side-channel communication. + private let queue: any TaskMessageQueue + + /// Client capabilities for checking support. + private let clientCapabilities: Client.Capabilities? + + /// Server reference for task-augmented requests (elicitAsTask, createMessageAsTask). + private let server: Server? + + /// Counter for generating request IDs. + private var requestIdCounter: Int = 0 + + /// Whether cancellation has been requested. + private var _isCancelled = false + + /// Check if cancellation has been requested. + public var isCancelled: Bool { _isCancelled } + + /// The task ID. + public var taskId: String { task.taskId } + + /// Create a server task context. + /// + /// - Parameters: + /// - task: The task to manage + /// - store: The task store for persistence + /// - queue: The message queue for side-channel communication + /// - clientCapabilities: Client capabilities for checking support + /// - server: Optional server reference for task-augmented requests + public init( + task: MCPTask, + store: any TaskStore, + queue: any TaskMessageQueue, + clientCapabilities: Client.Capabilities? = nil, + server: Server? = nil + ) { + self.task = task + self.store = store + self.queue = queue + self.clientCapabilities = clientCapabilities + self.server = server + } + + /// Generate a unique request ID for queued requests. + private func nextRequestId() -> RequestId { + requestIdCounter += 1 + return .string("task-\(taskId)-req-\(requestIdCounter)") + } + + /// Request cancellation of the task. + /// + /// This sets the `isCancelled` flag but doesn't immediately stop execution. + /// Task handlers should check this flag periodically and exit gracefully. + public func requestCancellation() { + _isCancelled = true + } + + /// Update the task status with a message. + /// + /// This updates the task to `.working` status with the provided message. + /// Use this to report progress during long-running operations. + /// + /// - Parameters: + /// - message: A human-readable status message + /// - notify: Whether to send a `TaskStatusNotification` to the client (default: true) + /// - Throws: Error if the task cannot be updated + public func updateStatus(_ message: String, notify: Bool = true) async throws { + let updatedTask = try await store.updateTask( + taskId: taskId, + status: .working, + statusMessage: message + ) + task = updatedTask + if notify { + await sendStatusNotification() + } + } + + /// Mark the task as requiring input. + /// + /// This updates the task to `.inputRequired` status, signaling that + /// the task is waiting for user input (e.g., via elicitation). + /// + /// - Parameters: + /// - message: Optional message describing what input is needed + /// - notify: Whether to send a `TaskStatusNotification` to the client (default: true) + /// - Throws: Error if the task cannot be updated + public func setInputRequired(_ message: String? = nil, notify: Bool = true) async throws { + let updatedTask = try await store.updateTask( + taskId: taskId, + status: .inputRequired, + statusMessage: message + ) + task = updatedTask + if notify { + await sendStatusNotification() + } + } + + /// Complete the task successfully with a result. + /// + /// This stores the result and transitions the task to `.completed` status. + /// + /// - Parameters: + /// - result: The result value to store + /// - notify: Whether to send a `TaskStatusNotification` to the client (default: true) + /// - Throws: Error if the task cannot be completed + public func complete(result: Value, notify: Bool = true) async throws { + try await store.storeResult(taskId: taskId, result: result) + let updatedTask = try await store.updateTask( + taskId: taskId, + status: .completed, + statusMessage: nil + ) + task = updatedTask + if notify { + await sendStatusNotification() + } + } + + /// Complete the task successfully with a CallTool.Result. + /// + /// This is a convenience method that encodes the result and stores it. + /// + /// - Parameters: + /// - toolResult: The tool result + /// - notify: Whether to send a `TaskStatusNotification` to the client (default: true) + /// - Throws: Error if encoding fails or the task cannot be completed + public func complete(toolResult: CallTool.Result, notify: Bool = true) async throws { + let encoder = JSONEncoder() + let data = try encoder.encode(toolResult) + let decoder = JSONDecoder() + let value = try decoder.decode(Value.self, from: data) + try await complete(result: value, notify: notify) + } + + /// Fail the task with an error message. + /// + /// This transitions the task to `.failed` status with the error message. + /// + /// - Parameters: + /// - error: A human-readable error message + /// - notify: Whether to send a `TaskStatusNotification` to the client (default: true) + /// - Throws: Error if the task cannot be updated + public func fail(error: String, notify: Bool = true) async throws { + let updatedTask = try await store.updateTask( + taskId: taskId, + status: .failed, + statusMessage: error + ) + task = updatedTask + if notify { + await sendStatusNotification() + } + } + + /// Fail the task with an Error. + /// + /// - Parameters: + /// - error: The error that caused the failure + /// - notify: Whether to send a `TaskStatusNotification` to the client (default: true) + /// - Throws: Error if the task cannot be updated + public func fail(error: any Error, notify: Bool = true) async throws { + try await fail(error: error.localizedDescription, notify: notify) + } + + /// Send a task status notification to the client. + /// + /// This sends a `notifications/tasks/status` notification with the current task state. + /// Per the MCP spec, this is sent when a task's status changes to keep the client informed. + private func sendStatusNotification() async { + guard let server else { return } + do { + try await server.notify(TaskStatusNotification.message(.init(task: task))) + } catch { + // Notification failures shouldn't break task execution + // The client will still get status updates via polling + } + } + + // MARK: - Mid-Task Interactive Requests + + // MARK: - Mid-Task Interactive Requests: Form Elicitation + + /// Request user input via form elicitation mid-task. + /// + /// This queues an elicitation request for delivery via `tasks/result` and waits + /// for the client's response. The task status is automatically transitioned to + /// `inputRequired` while waiting and restored to `working` when the response arrives. + /// + /// - Important: This is an experimental API that may change without notice. + /// + /// ## Example + /// + /// ```swift + /// let result = try await context.elicit( + /// message: "Please confirm the deletion", + /// requestedSchema: ElicitationSchema(properties: [ + /// "confirm": .boolean(BooleanSchema(title: "Confirm deletion")) + /// ]) + /// ) + /// + /// if result.action == .accept, let content = result.content { + /// let confirmed = content["confirm"] + /// // Process the response + /// } + /// ``` + /// + /// - Parameters: + /// - message: The message to present to the user + /// - requestedSchema: The schema defining the form fields + /// - Returns: The elicitation result from the client + /// - Throws: MCPError if the client doesn't support elicitation or if the request fails + public func elicit( + message: String, + requestedSchema: ElicitationSchema + ) async throws -> ElicitResult { + // Check client supports elicitation + guard clientCapabilities?.elicitation?.form != nil else { + throw MCPError.invalidRequest("Client does not support form elicitation") + } + + // Update task status to input_required + try await setInputRequired("Waiting for user input") + + // Build the elicitation request with related task metadata + let requestId = nextRequestId() + let relatedTaskMeta: [String: Value] = [ + relatedTaskMetaKey: .object(["taskId": .string(taskId)]) + ] + + let params = ElicitRequestFormParams( + mode: "form", + message: message, + requestedSchema: requestedSchema, + _meta: RequestMeta(additionalFields: relatedTaskMeta) + ) + + // Build JSON-RPC request + let request = Elicit.request(id: requestId, .form(params)) + let encoder = JSONEncoder() + let requestData = try encoder.encode(request) + + // Create resolver to wait for response + let resolver = Resolver() + + // Queue the request with resolver + let queuedRequest = QueuedRequestWithResolver( + message: .request(requestData, timestamp: Date()), + resolver: resolver, + originalRequestId: requestId + ) + + do { + try await queue.enqueueWithResolver(taskId: taskId, request: queuedRequest, maxSize: nil) + + // Signal that a message is available + await queue.notifyMessageAvailable(taskId: taskId) + // Also signal the store to wake up any waiters + await store.notifyUpdate(taskId: taskId) + + // Wait for response + let responseValue = try await resolver.wait() + + // Restore status to working + try await updateStatus("Continuing after user input") + + // Decode the response + let decoder = JSONDecoder() + let responseData = try encoder.encode(responseValue) + return try decoder.decode(ElicitResult.self, from: responseData) + } catch { + // Restore status to working even on error + try? await updateStatus("Continuing after error") + throw error + } + } + + // MARK: - Mid-Task Interactive Requests: URL Elicitation + + /// Request user interaction via URL-mode elicitation mid-task. + /// + /// This queues a URL elicitation request for delivery via `tasks/result` and waits + /// for the client's response. URL mode is used for out-of-band flows like OAuth + /// or credential collection, where the user needs to navigate to an external URL. + /// + /// The task status is automatically transitioned to `inputRequired` while waiting + /// and restored to `working` when the response arrives. + /// + /// - Important: This is an experimental API that may change without notice. + /// + /// ## Example + /// + /// ```swift + /// let result = try await context.elicitUrl( + /// message: "Please authorize access to your account", + /// url: "https://example.com/oauth?state=abc123", + /// elicitationId: "oauth-flow-123" + /// ) + /// + /// switch result.action { + /// case .accept: + /// // User completed the flow + /// case .decline: + /// // User declined + /// case .cancel: + /// // User cancelled + /// } + /// ``` + /// + /// - Parameters: + /// - message: Human-readable explanation of why the interaction is needed + /// - url: The URL the user should navigate to + /// - elicitationId: Unique identifier for tracking this elicitation + /// - Returns: The elicitation result from the client + /// - Throws: MCPError if the client doesn't support URL elicitation or if the request fails + public func elicitUrl( + message: String, + url: String, + elicitationId: String + ) async throws -> ElicitResult { + // Check client supports URL elicitation + guard clientCapabilities?.elicitation?.url != nil else { + throw MCPError.invalidRequest("Client does not support URL elicitation") + } + + // Update task status to input_required + try await setInputRequired("Waiting for external user interaction") + + // Build the URL elicitation request with related task metadata + let requestId = nextRequestId() + let relatedTaskMeta: [String: Value] = [ + relatedTaskMetaKey: .object(["taskId": .string(taskId)]) + ] + + let params = ElicitRequestURLParams( + message: message, + elicitationId: elicitationId, + url: url, + _meta: RequestMeta(additionalFields: relatedTaskMeta) + ) + + // Build JSON-RPC request + let request = Elicit.request(id: requestId, .url(params)) + let encoder = JSONEncoder() + let requestData = try encoder.encode(request) + + // Create resolver to wait for response + let resolver = Resolver() + + // Queue the request with resolver + let queuedRequest = QueuedRequestWithResolver( + message: .request(requestData, timestamp: Date()), + resolver: resolver, + originalRequestId: requestId + ) + + do { + try await queue.enqueueWithResolver(taskId: taskId, request: queuedRequest, maxSize: nil) + + // Signal that a message is available + await queue.notifyMessageAvailable(taskId: taskId) + // Also signal the store to wake up any waiters + await store.notifyUpdate(taskId: taskId) + + // Wait for response + let responseValue = try await resolver.wait() + + // Restore status to working + try await updateStatus("Continuing after external interaction") + + // Decode the response + let decoder = JSONDecoder() + let responseData = try encoder.encode(responseValue) + return try decoder.decode(ElicitResult.self, from: responseData) + } catch { + // Restore status to working even on error + try? await updateStatus("Continuing after error") + throw error + } + } + + // MARK: - Mid-Task Interactive Requests: Sampling + + /// Request LLM sampling mid-task. + /// + /// This queues a sampling request for delivery via `tasks/result` and waits + /// for the client's response. The task status is automatically transitioned to + /// `inputRequired` while waiting and restored to `working` when the response arrives. + /// + /// - Important: This is an experimental API that may change without notice. + /// + /// ## Example + /// + /// ```swift + /// let result = try await context.createMessage( + /// messages: [ + /// .user(.text("What is the capital of France?")) + /// ], + /// maxTokens: 100 + /// ) + /// + /// // Process the LLM response + /// for block in result.content { + /// if case .text(let text, _, _) = block { + /// print(text) + /// } + /// } + /// ``` + /// + /// - Parameters: + /// - messages: The conversation history + /// - maxTokens: Maximum tokens to generate + /// - modelPreferences: Optional model selection preferences + /// - systemPrompt: Optional system prompt + /// - includeContext: What MCP context to include + /// - temperature: Controls randomness (0.0 to 1.0) + /// - stopSequences: Array of sequences that stop generation + /// - metadata: Additional provider-specific parameters + /// - Returns: The sampling result from the client + /// - Throws: MCPError if the client doesn't support sampling or if the request fails + public func createMessage( + messages: [Sampling.Message], + maxTokens: Int, + modelPreferences: ModelPreferences? = nil, + systemPrompt: String? = nil, + includeContext: Sampling.ContextInclusion? = nil, + temperature: Double? = nil, + stopSequences: [String]? = nil, + metadata: [String: Value]? = nil + ) async throws -> CreateSamplingMessage.Result { + // Check client supports sampling + guard clientCapabilities?.sampling != nil else { + throw MCPError.invalidRequest("Client does not support sampling capability") + } + + // Update task status to input_required + try await setInputRequired("Waiting for LLM response") + + // Build the sampling request with related task metadata + let requestId = nextRequestId() + let relatedTaskMeta: [String: Value] = [ + relatedTaskMetaKey: .object(["taskId": .string(taskId)]) + ] + + let params = CreateSamplingMessage.Parameters( + messages: messages, + modelPreferences: modelPreferences, + systemPrompt: systemPrompt, + includeContext: includeContext, + temperature: temperature, + maxTokens: maxTokens, + stopSequences: stopSequences, + metadata: metadata, + _meta: RequestMeta(additionalFields: relatedTaskMeta) + ) + + // Build JSON-RPC request + let request = CreateSamplingMessage.request(id: requestId, params) + let encoder = JSONEncoder() + let requestData = try encoder.encode(request) + + // Create resolver to wait for response + let resolver = Resolver() + + // Queue the request with resolver + let queuedRequest = QueuedRequestWithResolver( + message: .request(requestData, timestamp: Date()), + resolver: resolver, + originalRequestId: requestId + ) + + do { + try await queue.enqueueWithResolver(taskId: taskId, request: queuedRequest, maxSize: nil) + + // Signal that a message is available + await queue.notifyMessageAvailable(taskId: taskId) + // Also signal the store to wake up any waiters + await store.notifyUpdate(taskId: taskId) + + // Wait for response + let responseValue = try await resolver.wait() + + // Restore status to working + try await updateStatus("Continuing after LLM response") + + // Decode the response + let decoder = JSONDecoder() + let responseData = try encoder.encode(responseValue) + return try decoder.decode(CreateSamplingMessage.Result.self, from: responseData) + } catch { + // Restore status to working even on error + try? await updateStatus("Continuing after error") + throw error + } + } + + // MARK: - Task-Augmented Elicitation (Server → Client Task) + + /// Request user input via task-augmented form elicitation. + /// + /// Unlike regular `elicit()`, this method creates a task on the CLIENT side, + /// allowing the client to handle the elicitation asynchronously. This is useful + /// when the client needs to perform complex operations during elicitation + /// (e.g., OAuth flows that require external callbacks). + /// + /// The method: + /// 1. Sends an elicitation request with a `task` field to the client + /// 2. Client returns a `CreateTaskResult` immediately + /// 3. Polls the client's task until it reaches a terminal state + /// 4. Retrieves and returns the final `ElicitResult` + /// + /// - Important: This is an experimental API that may change without notice. + /// + /// ## Example + /// + /// ```swift + /// // Task-augmented elicitation (client handles as a task) + /// let result = try await context.elicitAsTask( + /// message: "Please authorize access to your account", + /// requestedSchema: ElicitationSchema(properties: [ + /// "authorized": .boolean(BooleanSchema(title: "Authorized")) + /// ]) + /// ) + /// + /// if result.action == .accept { + /// // User completed authorization + /// } + /// ``` + /// + /// - Parameters: + /// - message: The message to present to the user + /// - requestedSchema: The schema defining the form fields + /// - ttl: Optional time-to-live for the client-side task + /// - Returns: The elicitation result from the client + /// - Throws: MCPError if the client doesn't support task-augmented elicitation or if the request fails + public func elicitAsTask( + message: String, + requestedSchema: ElicitationSchema, + ttl: Int? = nil + ) async throws -> ElicitResult { + // Check client supports task-augmented elicitation + guard hasTaskAugmentedElicitation(clientCapabilities) else { + throw MCPError.invalidRequest("Client does not support task-augmented elicitation") + } + + guard clientCapabilities?.elicitation?.form != nil else { + throw MCPError.invalidRequest("Client does not support form elicitation") + } + + // Need server reference to poll client tasks after CreateTaskResult + guard let server else { + throw MCPError.internalError("Server reference required for task-augmented requests") + } + + // Update task status to input_required + try await setInputRequired("Waiting for client task completion") + + // Build the elicitation request with task field and related task metadata + let requestId = nextRequestId() + let relatedTaskMeta: [String: Value] = [ + relatedTaskMetaKey: .object(["taskId": .string(taskId)]) + ] + + let params = ElicitRequestFormParams( + message: message, + requestedSchema: requestedSchema, + _meta: RequestMeta(additionalFields: relatedTaskMeta), + task: TaskMetadata(ttl: ttl) + ) + + // Build JSON-RPC request + let request = Elicit.request(id: requestId, .form(params)) + let encoder = JSONEncoder() + let requestData = try encoder.encode(request) + + // Create resolver to wait for CreateTaskResult response + let resolver = Resolver() + + // Queue the request with resolver (like regular elicit) + let queuedRequest = QueuedRequestWithResolver( + message: .request(requestData, timestamp: Date()), + resolver: resolver, + originalRequestId: requestId + ) + + do { + try await queue.enqueueWithResolver(taskId: taskId, request: queuedRequest, maxSize: nil) + + // Signal that a message is available + await queue.notifyMessageAvailable(taskId: taskId) + await store.notifyUpdate(taskId: taskId) + + // Wait for CreateTaskResult response (delivered when client polls tasks/result) + let responseValue = try await resolver.wait() + + // Decode as CreateTaskResult + let decoder = JSONDecoder() + let responseData = try encoder.encode(responseValue) + let createResult = try decoder.decode(CreateTaskResult.self, from: responseData) + let clientTaskId = createResult.task.taskId + + // NOW poll the client's task DIRECTLY (not through queue) + try await pollClientTaskUntilTerminal(server: server, taskId: clientTaskId) + + // Get the final result from client DIRECTLY + let result: ElicitResult = try await server.getClientTaskResultAs( + taskId: clientTaskId, + type: ElicitResult.self + ) + + // Restore status to working + try await updateStatus("Continuing after client task completion") + + return result + } catch { + // Restore status to working even on error + try? await updateStatus("Continuing after error") + throw error + } + } + + /// Request user input via task-augmented URL elicitation. + /// + /// Similar to `elicitAsTask(message:requestedSchema:)` but for URL-mode elicitation. + /// This creates a task on the CLIENT side for handling out-of-band flows like OAuth. + /// + /// - Important: This is an experimental API that may change without notice. + /// + /// - Parameters: + /// - message: Human-readable explanation of why the interaction is needed + /// - url: The URL the user should navigate to + /// - elicitationId: Unique identifier for tracking this elicitation + /// - ttl: Optional time-to-live for the client-side task + /// - Returns: The elicitation result from the client + /// - Throws: MCPError if the client doesn't support task-augmented URL elicitation + public func elicitUrlAsTask( + message: String, + url: String, + elicitationId: String, + ttl: Int? = nil + ) async throws -> ElicitResult { + // Check client supports task-augmented elicitation + guard hasTaskAugmentedElicitation(clientCapabilities) else { + throw MCPError.invalidRequest("Client does not support task-augmented elicitation") + } + + guard clientCapabilities?.elicitation?.url != nil else { + throw MCPError.invalidRequest("Client does not support URL elicitation") + } + + // Need server reference to poll client tasks after CreateTaskResult + guard let server else { + throw MCPError.internalError("Server reference required for task-augmented requests") + } + + // Update task status to input_required + try await setInputRequired("Waiting for client task completion") + + // Build the URL elicitation request with task field and related task metadata + let requestId = nextRequestId() + let relatedTaskMeta: [String: Value] = [ + relatedTaskMetaKey: .object(["taskId": .string(taskId)]) + ] + + let params = ElicitRequestURLParams( + message: message, + elicitationId: elicitationId, + url: url, + _meta: RequestMeta(additionalFields: relatedTaskMeta), + task: TaskMetadata(ttl: ttl) + ) + + // Build JSON-RPC request + let request = Elicit.request(id: requestId, .url(params)) + let encoder = JSONEncoder() + let requestData = try encoder.encode(request) + + // Create resolver to wait for CreateTaskResult response + let resolver = Resolver() + + // Queue the request with resolver (like regular elicitUrl) + let queuedRequest = QueuedRequestWithResolver( + message: .request(requestData, timestamp: Date()), + resolver: resolver, + originalRequestId: requestId + ) + + do { + try await queue.enqueueWithResolver(taskId: taskId, request: queuedRequest, maxSize: nil) + + // Signal that a message is available + await queue.notifyMessageAvailable(taskId: taskId) + await store.notifyUpdate(taskId: taskId) + + // Wait for CreateTaskResult response (delivered when client polls tasks/result) + let responseValue = try await resolver.wait() + + // Decode as CreateTaskResult + let decoder = JSONDecoder() + let responseData = try encoder.encode(responseValue) + let createResult = try decoder.decode(CreateTaskResult.self, from: responseData) + let clientTaskId = createResult.task.taskId + + // NOW poll the client's task DIRECTLY (not through queue) + try await pollClientTaskUntilTerminal(server: server, taskId: clientTaskId) + + // Get the final result from client DIRECTLY + let result: ElicitResult = try await server.getClientTaskResultAs( + taskId: clientTaskId, + type: ElicitResult.self + ) + + // Restore status to working + try await updateStatus("Continuing after client task completion") + + return result + } catch { + // Restore status to working even on error + try? await updateStatus("Continuing after error") + throw error + } + } + + // MARK: - Task-Augmented Sampling (Server → Client Task) + + /// Request LLM sampling via a task-augmented request. + /// + /// Unlike regular `createMessage()`, this method creates a task on the CLIENT side, + /// allowing the client to handle the sampling request asynchronously. This is useful + /// for long-running LLM operations or when the client needs to perform additional + /// processing during sampling. + /// + /// The method: + /// 1. Sends a sampling request with a `task` field to the client + /// 2. Client returns a `CreateTaskResult` immediately + /// 3. Polls the client's task until it reaches a terminal state + /// 4. Retrieves and returns the final `CreateSamplingMessage.Result` + /// + /// - Important: This is an experimental API that may change without notice. + /// + /// ## Example + /// + /// ```swift + /// // Task-augmented sampling (client handles as a task) + /// let result = try await context.createMessageAsTask( + /// messages: [ + /// .user(.text("Analyze this large document...")) + /// ], + /// maxTokens: 4000 + /// ) + /// + /// for block in result.content { + /// if case .text(let text, _, _) = block { + /// print(text) + /// } + /// } + /// ``` + /// + /// - Parameters: + /// - messages: The conversation history + /// - maxTokens: Maximum tokens to generate + /// - modelPreferences: Optional model selection preferences + /// - systemPrompt: Optional system prompt + /// - includeContext: What MCP context to include + /// - temperature: Controls randomness (0.0 to 1.0) + /// - stopSequences: Array of sequences that stop generation + /// - metadata: Additional provider-specific parameters + /// - ttl: Optional time-to-live for the client-side task + /// - Returns: The sampling result from the client + /// - Throws: MCPError if the client doesn't support task-augmented sampling + public func createMessageAsTask( + messages: [Sampling.Message], + maxTokens: Int, + modelPreferences: ModelPreferences? = nil, + systemPrompt: String? = nil, + includeContext: Sampling.ContextInclusion? = nil, + temperature: Double? = nil, + stopSequences: [String]? = nil, + metadata: [String: Value]? = nil, + ttl: Int? = nil + ) async throws -> CreateSamplingMessage.Result { + // Check client supports task-augmented sampling + guard hasTaskAugmentedSampling(clientCapabilities) else { + throw MCPError.invalidRequest("Client does not support task-augmented sampling") + } + + guard clientCapabilities?.sampling != nil else { + throw MCPError.invalidRequest("Client does not support sampling capability") + } + + // Need server reference to poll client tasks after CreateTaskResult + guard let server else { + throw MCPError.internalError("Server reference required for task-augmented requests") + } + + // Update task status to input_required + try await setInputRequired("Waiting for client LLM task completion") + + // Build the sampling request with task field and related task metadata + let requestId = nextRequestId() + let relatedTaskMeta: [String: Value] = [ + relatedTaskMetaKey: .object(["taskId": .string(taskId)]) + ] + + let params = CreateSamplingMessage.Parameters( + messages: messages, + modelPreferences: modelPreferences, + systemPrompt: systemPrompt, + includeContext: includeContext, + temperature: temperature, + maxTokens: maxTokens, + stopSequences: stopSequences, + metadata: metadata, + _meta: RequestMeta(additionalFields: relatedTaskMeta), + task: TaskMetadata(ttl: ttl) + ) + + // Build JSON-RPC request + let request = CreateSamplingMessage.request(id: requestId, params) + let encoder = JSONEncoder() + let requestData = try encoder.encode(request) + + // Create resolver to wait for CreateTaskResult response + let resolver = Resolver() + + // Queue the request with resolver (like regular createMessage) + let queuedRequest = QueuedRequestWithResolver( + message: .request(requestData, timestamp: Date()), + resolver: resolver, + originalRequestId: requestId + ) + + do { + try await queue.enqueueWithResolver(taskId: taskId, request: queuedRequest, maxSize: nil) + + // Signal that a message is available + await queue.notifyMessageAvailable(taskId: taskId) + await store.notifyUpdate(taskId: taskId) + + // Wait for CreateTaskResult response (delivered when client polls tasks/result) + let responseValue = try await resolver.wait() + + // Decode as CreateTaskResult + let decoder = JSONDecoder() + let responseData = try encoder.encode(responseValue) + let createResult = try decoder.decode(CreateTaskResult.self, from: responseData) + let clientTaskId = createResult.task.taskId + + // NOW poll the client's task DIRECTLY (not through queue) + try await pollClientTaskUntilTerminal(server: server, taskId: clientTaskId) + + // Get the final result from client DIRECTLY + let result: CreateSamplingMessage.Result = try await server.getClientTaskResultAs( + taskId: clientTaskId, + type: CreateSamplingMessage.Result.self + ) + + // Restore status to working + try await updateStatus("Continuing after client LLM task completion") + + return result + } catch { + // Restore status to working even on error + try? await updateStatus("Continuing after error") + throw error + } + } + + /// Poll a client task until it reaches a terminal state. + /// + /// - Parameters: + /// - server: The server to use for polling + /// - taskId: The client-side task identifier + private func pollClientTaskUntilTerminal(server: Server, taskId: String) async throws { + while true { + let result = try await server.getClientTask(taskId: taskId) + + if result.status.isTerminal { + return + } + + // Wait for poll interval (default 500ms) + let intervalMs = result.pollInterval ?? 500 + try await Task.sleep(for: .milliseconds(intervalMs)) + } + } +} diff --git a/Sources/MCP/Server/Experimental/Tasks/TaskContext.swift b/Sources/MCP/Server/Experimental/Tasks/TaskContext.swift new file mode 100644 index 00000000..dff9913c --- /dev/null +++ b/Sources/MCP/Server/Experimental/Tasks/TaskContext.swift @@ -0,0 +1,276 @@ +import Foundation + +// MARK: - Pure Task Context + +/// A pure task context without server dependencies. +/// +/// This context provides basic task management capabilities that work without +/// a server session, making it suitable for distributed workers or background +/// processing. For server-integrated task handling with elicitation and sampling, +/// use `ServerTaskContext` instead. +/// +/// Unlike `ServerTaskContext`, this context: +/// - Does not require client capabilities +/// - Does not support mid-task elicitation or sampling +/// - Can be used in distributed/worker processes with just a TaskStore +/// +/// - Important: This is an experimental API that may change without notice. +/// +/// ## Example (Distributed Worker) +/// +/// ```swift +/// func workerProcess(taskId: String) async { +/// let store = RedisTaskStore(url: redisUrl) +/// let context = try await TaskContext.load(taskId: taskId, from: store) +/// +/// do { +/// await context.updateStatus("Processing...") +/// let result = try await doWork() +/// try await context.complete(result: result) +/// } catch { +/// try await context.fail(error: error) +/// } +/// } +/// ``` +public actor TaskContext { + /// The task this context is for. + public private(set) var task: MCPTask + + /// The task store for persistence. + private let store: any TaskStore + + /// Whether cancellation has been requested. + private var _isCancelled = false + + /// Check if cancellation has been requested. + public var isCancelled: Bool { _isCancelled } + + /// The task ID. + public var taskId: String { task.taskId } + + /// Create a task context. + /// + /// - Parameters: + /// - task: The task to manage + /// - store: The task store for persistence + public init(task: MCPTask, store: any TaskStore) { + self.task = task + self.store = store + } + + /// Load a task context from the store. + /// + /// This is the recommended way to create a context in distributed workers. + /// + /// - Parameters: + /// - taskId: The task identifier + /// - store: The task store + /// - Returns: A TaskContext for the loaded task + /// - Throws: Error if the task is not found + public static func load(taskId: String, from store: any TaskStore) async throws -> TaskContext { + guard let task = await store.getTask(taskId: taskId) else { + throw MCPError.invalidParams("Task not found: \(taskId)") + } + return TaskContext(task: task, store: store) + } + + /// Request cancellation of the task. + /// + /// This sets the `isCancelled` flag but doesn't immediately stop execution. + /// Task handlers should check this flag periodically and exit gracefully. + public func requestCancellation() { + _isCancelled = true + } + + /// Update the task status with a message. + /// + /// This updates the task to `.working` status with the provided message. + /// Use this to report progress during long-running operations. + /// + /// - Parameter message: A human-readable status message + /// - Throws: Error if the task cannot be updated + public func updateStatus(_ message: String) async throws { + let updatedTask = try await store.updateTask( + taskId: taskId, + status: .working, + statusMessage: message + ) + task = updatedTask + } + + /// Mark the task as requiring input. + /// + /// This updates the task to `.inputRequired` status, signaling that + /// the task is waiting for user input. + /// + /// - Note: For mid-task elicitation, use `ServerTaskContext` instead. + /// + /// - Parameter message: Optional message describing what input is needed + /// - Throws: Error if the task cannot be updated + public func setInputRequired(_ message: String? = nil) async throws { + let updatedTask = try await store.updateTask( + taskId: taskId, + status: .inputRequired, + statusMessage: message + ) + task = updatedTask + } + + /// Complete the task successfully with a result. + /// + /// This stores the result and transitions the task to `.completed` status. + /// + /// - Parameter result: The result value to store + /// - Throws: Error if the task cannot be completed + public func complete(result: Value) async throws { + try await store.storeResult(taskId: taskId, result: result) + let updatedTask = try await store.updateTask( + taskId: taskId, + status: .completed, + statusMessage: nil + ) + task = updatedTask + } + + /// Complete the task successfully with a CallTool.Result. + /// + /// This is a convenience method that encodes the result and stores it. + /// + /// - Parameter result: The tool result + /// - Throws: Error if encoding fails or the task cannot be completed + public func complete(toolResult: CallTool.Result) async throws { + let encoder = JSONEncoder() + let data = try encoder.encode(toolResult) + let decoder = JSONDecoder() + let value = try decoder.decode(Value.self, from: data) + try await complete(result: value) + } + + /// Fail the task with an error message. + /// + /// This transitions the task to `.failed` status with the error message. + /// + /// - Parameter error: A human-readable error message + /// - Throws: Error if the task cannot be updated + public func fail(error: String) async throws { + let updatedTask = try await store.updateTask( + taskId: taskId, + status: .failed, + statusMessage: error + ) + task = updatedTask + } + + /// Fail the task with an Error. + /// + /// - Parameter error: The error that caused the failure + /// - Throws: Error if the task cannot be updated + public func fail(error: any Error) async throws { + try await fail(error: error.localizedDescription) + } + + /// Cancel the task. + /// + /// This transitions the task to `.cancelled` status. + /// + /// - Parameter message: Optional message describing why the task was cancelled + /// - Throws: Error if the task cannot be updated + public func cancel(message: String? = nil) async throws { + _isCancelled = true + let updatedTask = try await store.updateTask( + taskId: taskId, + status: .cancelled, + statusMessage: message ?? "Cancelled" + ) + task = updatedTask + } +} + +// MARK: - Task Execution Helper + +/// Execute work within a task context, automatically handling failures. +/// +/// This is similar to Python SDK's `task_execution` context manager. +/// If an unhandled exception occurs, the task is automatically marked as failed +/// and the error is suppressed (since the failure is captured in task state). +/// +/// This is useful for distributed workers that don't have a server session. +/// +/// - Important: This is an experimental API that may change without notice. +/// +/// ## Example (Distributed Worker) +/// +/// ```swift +/// let store = RedisTaskStore(url: redisUrl) +/// try await withTaskExecution(taskId: taskId, store: store) { context in +/// await context.updateStatus("Working...") +/// let result = try await doWork() +/// try await context.complete(result: result) +/// } +/// // If doWork() throws, task is automatically marked as failed +/// ``` +/// +/// - Parameters: +/// - taskId: The task identifier to execute +/// - store: The task store (must be accessible by the worker) +/// - work: The async work function that receives the task context +/// - Throws: Error only if the task cannot be loaded (not for work failures) +public func withTaskExecution( + taskId: String, + store: any TaskStore, + work: @escaping @Sendable (TaskContext) async throws -> Void +) async throws { + let context = try await TaskContext.load(taskId: taskId, from: store) + + do { + try await work(context) + } catch is CancellationError { + // Task was cancelled externally + if !isTerminalStatus(await context.task.status) { + try? await context.cancel(message: "Cancelled") + } + } catch { + // Auto-fail the task if an exception occurs and task isn't already terminal + if !isTerminalStatus(await context.task.status) { + try? await context.fail(error: error) + } + // Don't re-raise - the failure is recorded in task state + } +} + +// MARK: - Task Helper Functions + +/// Generate a unique task ID. +/// +/// This is a helper for TaskStore implementations. +/// +/// - Returns: A unique task identifier +public func generateTaskId() -> String { + UUID().uuidString.lowercased().replacingOccurrences(of: "-", with: "") +} + +/// Create a Task object with initial state. +/// +/// This is a helper for TaskStore implementations. +/// +/// - Parameters: +/// - metadata: Task metadata (TTL, etc.) +/// - taskId: Optional task ID (generated if nil) +/// - pollInterval: Suggested polling interval in milliseconds (default: 500) +/// - Returns: A new Task in "working" status +public func createTaskState( + metadata: TaskMetadata, + taskId: String? = nil, + pollInterval: Int = 500 +) -> MCPTask { + let id = taskId ?? generateTaskId() + let now = ISO8601DateFormatter().string(from: Date()) + return MCPTask( + taskId: id, + status: .working, + ttl: metadata.ttl, + createdAt: now, + lastUpdatedAt: now, + pollInterval: pollInterval + ) +} diff --git a/Sources/MCP/Server/Experimental/Tasks/TaskMessageQueue.swift b/Sources/MCP/Server/Experimental/Tasks/TaskMessageQueue.swift new file mode 100644 index 00000000..3531df59 --- /dev/null +++ b/Sources/MCP/Server/Experimental/Tasks/TaskMessageQueue.swift @@ -0,0 +1,416 @@ +import Foundation + +// MARK: - Response Router + +/// Protocol for routing responses back to waiting task handlers. +/// +/// When a task handler calls `elicit()` or `createMessage()`, it queues a request +/// and waits for a response. This protocol allows the response to be routed back +/// to the waiting handler's resolver when it arrives. +/// +/// Implementations should check if they have a pending resolver for the given +/// request ID and, if so, deliver the response to that resolver. +/// +/// ## Example +/// +/// ```swift +/// // Register the router with the session +/// session.addResponseRouter(taskResultHandler) +/// +/// // When a response arrives: +/// for router in responseRouters { +/// if router.routeResponse(requestId: responseId, response: responseData) { +/// // Response was handled by this router +/// return +/// } +/// } +/// // Fall through to normal response handling +/// ``` +/// +/// - Important: This is an experimental API that may change without notice. +public protocol ResponseRouter: Sendable { + /// Route a response back to a waiting resolver. + /// + /// - Parameters: + /// - requestId: The request ID of the original request + /// - response: The response data + /// - Returns: True if the response was routed (resolver found), false otherwise + func routeResponse(requestId: RequestId, response: Value) async -> Bool + + /// Route an error back to a waiting resolver. + /// + /// - Parameters: + /// - requestId: The request ID of the original request + /// - error: The error + /// - Returns: True if the error was routed (resolver found), false otherwise + func routeError(requestId: RequestId, error: any Error) async -> Bool +} + +// MARK: - Resolver + +/// A resolver for passing results between async contexts. +/// +/// This is used to route responses back to waiting task handlers. +/// When a task-augmented handler calls `elicit()` or `createMessage()`, +/// it creates a resolver and waits on it. When the response arrives +/// via `tasks/result`, the resolver is used to deliver the response. +/// +/// ## Example +/// +/// ```swift +/// // In task handler +/// let resolver = Resolver() +/// await queue.enqueueWithResolver(taskId: taskId, message: request, resolver: resolver) +/// let result = try await resolver.wait() +/// +/// // In tasks/result handler +/// // ... route response back via resolver +/// await resolver.setResult(response) +/// ``` +public actor Resolver { + private var result: Result? + private var continuation: CheckedContinuation? + + public init() {} + + /// Set the result value and wake up waiters. + public func setResult(_ value: T) { + if result != nil { + // Already completed, ignore + return + } + result = .success(value) + continuation?.resume(returning: value) + continuation = nil + } + + /// Set an exception and wake up waiters. + public func setError(_ error: any Error) { + if result != nil { + // Already completed, ignore + return + } + result = .failure(error) + continuation?.resume(throwing: error) + continuation = nil + } + + /// Wait for the result and return it, or throw the exception. + public func wait() async throws -> T { + // Check if already resolved + if let result { + switch result { + case .success(let value): + return value + case .failure(let error): + throw error + } + } + + // Wait for result + return try await withCheckedThrowingContinuation { cont in + // Check again (race with setResult) + if let result { + switch result { + case .success(let value): + cont.resume(returning: value) + case .failure(let error): + cont.resume(throwing: error) + } + } else { + continuation = cont + } + } + } + + /// Return true if the resolver has been completed. + public var isDone: Bool { + result != nil + } +} + +// MARK: - Queued Message + +/// Represents a message queued for side-channel delivery via tasks/result. +/// +/// This is used during task execution to queue requests (like elicitation or sampling) +/// that need to be delivered to the client when it polls for task results. +public enum QueuedMessage: Sendable { + /// A JSON-RPC request to be sent to the client + case request(Data, timestamp: Date) + /// A JSON-RPC notification to be sent to the client + case notification(Data, timestamp: Date) + /// A JSON-RPC response + case response(Data, timestamp: Date) + /// A JSON-RPC error response + case error(Data, timestamp: Date) + + /// The timestamp when this message was queued. + public var timestamp: Date { + switch self { + case .request(_, let ts), .notification(_, let ts), + .response(_, let ts), .error(_, let ts): + return ts + } + } + + /// The message data. + public var data: Data { + switch self { + case .request(let d, _), .notification(let d, _), + .response(let d, _), .error(let d, _): + return d + } + } +} + +/// A queued message with an associated resolver for response routing. +/// +/// When a request is queued that expects a response (like elicitation), +/// this struct pairs the request with a resolver that will receive the response. +public struct QueuedRequestWithResolver: Sendable { + /// The queued message. + public let message: QueuedMessage + /// The resolver to receive the response. + public let resolver: Resolver + /// The original request ID used for routing the response back. + public let originalRequestId: RequestId + + public init(message: QueuedMessage, resolver: Resolver, originalRequestId: RequestId) { + self.message = message + self.resolver = resolver + self.originalRequestId = originalRequestId + } +} + +/// Protocol for managing per-task FIFO message queues. +/// +/// This allows pluggable queue implementations (in-memory, Redis, other distributed queues, etc.). +/// Each method accepts taskId to enable a single queue instance to manage messages for multiple tasks. +/// +/// All methods are async to support external storage implementations. +/// +/// - Important: This is an experimental API that may change without notice. +public protocol TaskMessageQueue: Sendable { + /// Adds a message to the end of the queue for a specific task. + /// + /// - Parameters: + /// - taskId: The task identifier + /// - message: The message to enqueue + /// - maxSize: Optional maximum queue size. If specified and queue is full, throws an error. + /// - Throws: Error if maxSize is specified and would be exceeded + func enqueue(taskId: String, message: QueuedMessage, maxSize: Int?) async throws + + /// Adds a request message with a resolver for response routing. + /// + /// This is used for requests that expect a response (like elicitation or sampling). + /// The resolver will be used to deliver the response when it arrives. + /// + /// - Parameters: + /// - taskId: The task identifier + /// - request: The request message with resolver + /// - maxSize: Optional maximum queue size + /// - Throws: Error if maxSize is specified and would be exceeded + func enqueueWithResolver(taskId: String, request: QueuedRequestWithResolver, maxSize: Int?) async throws + + /// Removes and returns the first message from the queue for a specific task. + /// + /// - Parameter taskId: The task identifier + /// - Returns: The first message, or nil if the queue is empty + func dequeue(taskId: String) async -> QueuedMessage? + + /// Removes and returns the first message with its resolver from the queue. + /// + /// - Parameter taskId: The task identifier + /// - Returns: The first message with resolver, or nil if empty + func dequeueWithResolver(taskId: String) async -> QueuedRequestWithResolver? + + /// Removes and returns all messages from the queue for a specific task. + /// + /// Used when tasks are cancelled or failed to clean up pending messages. + /// + /// - Parameter taskId: The task identifier + /// - Returns: Array of all messages that were in the queue + func dequeueAll(taskId: String) async -> [QueuedMessage] + + /// Check if the queue for a task is empty. + /// + /// - Parameter taskId: The task identifier + /// - Returns: True if the queue is empty or doesn't exist + func isEmpty(taskId: String) async -> Bool + + /// Wait for a message to become available for the specified task. + /// + /// This method blocks until a message is enqueued for the task. + /// Used by `tasks/result` to implement long-polling behavior. + /// + /// - Parameter taskId: The task identifier + /// - Throws: Error if waiting is interrupted + func waitForMessage(taskId: String) async throws + + /// Notify waiters that a message is available for the specified task. + /// + /// This should be called after enqueueing a message. + /// + /// - Parameter taskId: The task identifier + func notifyMessageAvailable(taskId: String) async + + /// Get the resolver for a pending request by its request ID. + /// + /// This is used to route responses back to the waiting handler. + /// + /// - Parameter requestId: The request ID to look up + /// - Returns: The resolver if found, or nil + func getResolver(forRequestId requestId: RequestId) async -> Resolver? + + /// Remove and return the resolver for a pending request. + /// + /// - Parameter requestId: The request ID to look up + /// - Returns: The resolver if found, or nil + func removeResolver(forRequestId requestId: RequestId) async -> Resolver? +} + +/// An in-memory implementation of ``TaskMessageQueue`` for demonstration purposes. +/// +/// This implementation stores messages in memory, organized by task ID. +/// Messages are stored in FIFO queues per task. +/// +/// - Important: This is not suitable for production use in distributed systems. +/// For production, consider implementing TaskMessageQueue with Redis or other distributed queues. +public actor InMemoryTaskMessageQueue: TaskMessageQueue { + /// Internal storage for a queued item that may have a resolver. + private struct QueuedItem { + let message: QueuedMessage + let resolver: Resolver? + let originalRequestId: RequestId? + } + + /// Dictionary of message queues keyed by task ID. + private var queues: [String: [QueuedItem]] = [:] + + /// Pending request resolvers keyed by request ID for response routing. + private var pendingResolvers: [RequestId: Resolver] = [:] + + /// Waiters for message availability, keyed by task ID. + private var messageWaiters: [String: [CheckedContinuation]] = [:] + + /// Create an in-memory task message queue. + public init() {} + + public func enqueue(taskId: String, message: QueuedMessage, maxSize: Int?) async throws { + var queue = queues[taskId, default: []] + + if let maxSize, queue.count >= maxSize { + throw MCPError.internalError("Task message queue overflow: queue size (\(queue.count)) exceeds maximum (\(maxSize))") + } + + queue.append(QueuedItem(message: message, resolver: nil, originalRequestId: nil)) + queues[taskId] = queue + + // Notify waiters that a message is available + await notifyMessageAvailable(taskId: taskId) + } + + public func enqueueWithResolver(taskId: String, request: QueuedRequestWithResolver, maxSize: Int?) async throws { + var queue = queues[taskId, default: []] + + if let maxSize, queue.count >= maxSize { + throw MCPError.internalError("Task message queue overflow: queue size (\(queue.count)) exceeds maximum (\(maxSize))") + } + + let item = QueuedItem( + message: request.message, + resolver: request.resolver, + originalRequestId: request.originalRequestId + ) + queue.append(item) + queues[taskId] = queue + + // Store the resolver for response routing + pendingResolvers[request.originalRequestId] = request.resolver + + // Notify waiters that a message is available + await notifyMessageAvailable(taskId: taskId) + } + + public func dequeue(taskId: String) async -> QueuedMessage? { + guard var queue = queues[taskId], !queue.isEmpty else { + return nil + } + + let item = queue.removeFirst() + queues[taskId] = queue + return item.message + } + + public func dequeueWithResolver(taskId: String) async -> QueuedRequestWithResolver? { + guard var queue = queues[taskId], !queue.isEmpty else { + return nil + } + + let item = queue.removeFirst() + queues[taskId] = queue + + guard let resolver = item.resolver, let originalRequestId = item.originalRequestId else { + // Re-queue the message if it doesn't have a resolver + // and try the next one + queues[taskId, default: []].insert(QueuedItem(message: item.message, resolver: nil, originalRequestId: nil), at: 0) + return nil + } + + return QueuedRequestWithResolver( + message: item.message, + resolver: resolver, + originalRequestId: originalRequestId + ) + } + + public func dequeueAll(taskId: String) async -> [QueuedMessage] { + let items = queues.removeValue(forKey: taskId) ?? [] + return items.map(\.message) + } + + public func isEmpty(taskId: String) async -> Bool { + queues[taskId]?.isEmpty ?? true + } + + public func waitForMessage(taskId: String) async throws { + try await withCheckedThrowingContinuation { (continuation: CheckedContinuation) in + messageWaiters[taskId, default: []].append(continuation) + } + } + + public func notifyMessageAvailable(taskId: String) async { + guard let waiters = messageWaiters.removeValue(forKey: taskId), !waiters.isEmpty else { + return + } + for continuation in waiters { + continuation.resume() + } + } + + public func getResolver(forRequestId requestId: RequestId) async -> Resolver? { + pendingResolvers[requestId] + } + + public func removeResolver(forRequestId requestId: RequestId) async -> Resolver? { + pendingResolvers.removeValue(forKey: requestId) + } + + /// Clear all queues (useful for testing or graceful shutdown). + public func cleanUp() async { + queues.removeAll() + // Cancel all message waiters + for (_, continuations) in messageWaiters { + for continuation in continuations { + continuation.resume(throwing: CancellationError()) + } + } + messageWaiters.removeAll() + // Error out all pending resolvers + for (_, resolver) in pendingResolvers { + await resolver.setError(CancellationError()) + } + pendingResolvers.removeAll() + } +} diff --git a/Sources/MCP/Server/Experimental/Tasks/TaskResultHandler.swift b/Sources/MCP/Server/Experimental/Tasks/TaskResultHandler.swift new file mode 100644 index 00000000..9f728929 --- /dev/null +++ b/Sources/MCP/Server/Experimental/Tasks/TaskResultHandler.swift @@ -0,0 +1,136 @@ +import Foundation + +// MARK: - Task Result Handler + +/// Handler for `tasks/result` that implements the message queue pattern. +/// +/// This handler: +/// 1. Dequeues pending messages (elicitations, sampling) for the task +/// 2. Sends them to the client via the response stream +/// 3. Waits for responses and resolves them back to callers +/// 4. Blocks until task reaches terminal state +/// 5. Returns the final result +/// +/// The handler also implements `ResponseRouter` to route incoming responses +/// back to waiting task handlers. +/// +/// - Important: This is an experimental API that may change without notice. +public final class TaskResultHandler: Sendable, ResponseRouter { + private let store: any TaskStore + private let queue: any TaskMessageQueue + + /// Create a task result handler. + /// + /// - Parameters: + /// - store: The task store for reading task state + /// - queue: The message queue for pending messages + public init(store: any TaskStore, queue: any TaskMessageQueue) { + self.store = store + self.queue = queue + } + + /// Handle a `tasks/result` request. + /// + /// This implements the dequeue-send-wait loop: + /// 1. Dequeue all pending messages + /// 2. Send each via the provided send function with relatedRequestId + /// 3. If task not terminal, wait for status change or new messages + /// 4. Loop until task is terminal + /// 5. Return final result with related-task metadata + /// + /// - Parameters: + /// - taskId: The task to get results for + /// - sendMessage: Closure to send queued messages to the client + /// - Returns: The task result with related-task metadata + /// - Throws: MCPError if task not found or processing fails + public func handle( + taskId: String, + sendMessage: @Sendable (Data) async throws -> Void + ) async throws -> GetTaskPayload.Result { + while true { + // Check task exists + guard let task = await store.getTask(taskId: taskId) else { + throw MCPError.invalidParams("Task not found: \(taskId)") + } + + // Deliver all queued messages + try await deliverQueuedMessages(taskId: taskId, sendMessage: sendMessage) + + // If task is terminal, return result + if isTerminalStatus(task.status) { + let result = await store.getResult(taskId: taskId) + let relatedTaskMeta: [String: Value] = [ + relatedTaskMetaKey: .object(["taskId": .string(taskId)]) + ] + + // Flatten result fields into extraFields + return GetTaskPayload.Result( + fromResultValue: result, + _meta: relatedTaskMeta + ) + } + + // Wait for task update (status change or new messages) + try await waitForTaskUpdate(taskId: taskId) + } + } + + /// Deliver all queued messages for a task. + /// + /// - Parameters: + /// - taskId: The task identifier + /// - sendMessage: Closure to send messages to the client + private func deliverQueuedMessages( + taskId: String, + sendMessage: @Sendable (Data) async throws -> Void + ) async throws { + while let message = await queue.dequeue(taskId: taskId) { + // Send the message to the client + try await sendMessage(message.data) + } + } + + /// Wait for a task update (status change or new message). + /// + /// - Parameter taskId: The task identifier + private func waitForTaskUpdate(taskId: String) async throws { + // We need to wait for either: + // 1. Task status to change (via store.waitForUpdate) + // 2. A new message to be queued (via queue.waitForMessage) + // + // For simplicity, we'll use the store's wait which is signaled + // both on status changes and when messages are queued. + try await store.waitForUpdate(taskId: taskId) + } + + /// Route a response back to a waiting resolver. + /// + /// This is called when a response arrives for a queued request + /// (e.g., elicitation or sampling response). + /// + /// - Parameters: + /// - requestId: The request ID of the original request + /// - response: The response value + /// - Returns: True if the response was routed, false if no resolver found + public func routeResponse(requestId: RequestId, response: Value) async -> Bool { + guard let resolver = await queue.removeResolver(forRequestId: requestId) else { + return false + } + await resolver.setResult(response) + return true + } + + /// Route an error back to a waiting resolver. + /// + /// - Parameters: + /// - requestId: The request ID of the original request + /// - error: The error + /// - Returns: True if the error was routed, false if no resolver found + public func routeError(requestId: RequestId, error: any Error) async -> Bool { + guard let resolver = await queue.removeResolver(forRequestId: requestId) else { + return false + } + await resolver.setError(error) + return true + } +} diff --git a/Sources/MCP/Server/Experimental/Tasks/TaskStore.swift b/Sources/MCP/Server/Experimental/Tasks/TaskStore.swift new file mode 100644 index 00000000..bc17ab85 --- /dev/null +++ b/Sources/MCP/Server/Experimental/Tasks/TaskStore.swift @@ -0,0 +1,333 @@ +import Foundation + +/// Protocol for storing and retrieving task state and results. +/// +/// This abstraction allows pluggable task storage implementations +/// (in-memory, database, distributed cache, etc.). +/// +/// All methods are async to support various backends. +/// +/// - Important: This is an experimental API that may change without notice. +public protocol TaskStore: Sendable { + /// Create a new task with the given metadata. + /// + /// - Parameters: + /// - metadata: Task metadata (TTL, etc.) + /// - taskId: Optional task ID. If nil, implementation should generate one. + /// - Returns: The created Task with status `working` + /// - Throws: Error if taskId already exists + func createTask(metadata: TaskMetadata, taskId: String?) async throws -> MCPTask + + /// Get a task by ID. + /// + /// - Parameter taskId: The task identifier + /// - Returns: The Task, or nil if not found + func getTask(taskId: String) async -> MCPTask? + + /// Update a task's status and/or message. + /// + /// - Parameters: + /// - taskId: The task identifier + /// - status: New status (if changing) + /// - statusMessage: New status message (if changing) + /// - Returns: The updated Task + /// - Throws: Error if task not found or if attempting to transition from a terminal status + func updateTask(taskId: String, status: TaskStatus?, statusMessage: String?) async throws -> MCPTask + + /// Store the result for a task. + /// + /// - Parameters: + /// - taskId: The task identifier + /// - result: The result to store + /// - Throws: Error if task not found + func storeResult(taskId: String, result: Value) async throws + + /// Get the stored result for a task. + /// + /// - Parameter taskId: The task identifier + /// - Returns: The stored result, or nil if not available + func getResult(taskId: String) async -> Value? + + /// List tasks with pagination. + /// + /// - Parameter cursor: Optional cursor for pagination + /// - Returns: Tuple of (tasks, nextCursor). nextCursor is nil if no more pages. + func listTasks(cursor: String?) async -> (tasks: [MCPTask], nextCursor: String?) + + /// Delete a task. + /// + /// - Parameter taskId: The task identifier + /// - Returns: True if deleted, false if not found + func deleteTask(taskId: String) async -> Bool + + /// Wait for an update to the specified task. + /// + /// This method blocks until the task's status changes or a message becomes available. + /// Used by `tasks/result` to implement long-polling behavior. + /// + /// - Parameter taskId: The task identifier + /// - Throws: Error if waiting is interrupted + func waitForUpdate(taskId: String) async throws + + /// Notify waiters that a task has been updated. + /// + /// This should be called after updating a task's status or queueing a message. + /// + /// - Parameter taskId: The task identifier + func notifyUpdate(taskId: String) async +} + +/// Checks if a task status represents a terminal state. +/// +/// Terminal states are those where the task has finished and will not change. +/// +/// - Parameter status: The task status to check +/// - Returns: True if the status is terminal (completed, failed, or cancelled) +public func isTerminalStatus(_ status: TaskStatus) -> Bool { + switch status { + case .completed, .failed, .cancelled: + return true + case .working, .inputRequired: + return false + } +} + +/// An in-memory implementation of ``TaskStore`` for demonstration and testing purposes. +/// +/// This implementation stores all tasks in memory and provides lazy cleanup +/// based on the TTL duration specified in the task metadata. +/// +/// - Important: This is not suitable for production use as all data is lost on restart. +/// For production, consider implementing TaskStore with a database or distributed cache. +public actor InMemoryTaskStore: TaskStore { + /// Internal storage for a task and its result. + private struct StoredTask { + var task: MCPTask + var result: Value? + /// Time when this task should be removed (nil = never) + var expiresAt: Date? + } + + /// Dictionary of stored tasks keyed by task ID. + private var tasks: [String: StoredTask] = [:] + + /// Page size for listing tasks. + private let pageSize: Int + + /// A waiter entry with unique ID for cancellation tracking. + private struct Waiter { + let id: UUID + let continuation: CheckedContinuation + } + + /// Waiters for task updates, keyed by task ID. + /// Each waiter has a unique ID so it can be individually cancelled. + private var waiters: [String: [Waiter]] = [:] + + /// Create an in-memory task store. + /// + /// - Parameter pageSize: The number of tasks to return per page in `listTasks`. Defaults to 10. + public init(pageSize: Int = 10) { + self.pageSize = pageSize + } + + /// Calculate expiry date from TTL in milliseconds. + private func calculateExpiry(ttl: Int?) -> Date? { + guard let ttl else { return nil } + return Date().addingTimeInterval(Double(ttl) / 1000.0) + } + + /// Check if a stored task has expired. + private func isExpired(_ stored: StoredTask) -> Bool { + guard let expiresAt = stored.expiresAt else { return false } + return Date() >= expiresAt + } + + /// Remove all expired tasks (called lazily during access operations). + private func cleanUpExpired() { + let expiredIds = tasks.filter { isExpired($0.value) }.map(\.key) + for id in expiredIds { + tasks.removeValue(forKey: id) + } + } + + /// Generate a unique task ID using UUID. + private func generateTaskId() -> String { + UUID().uuidString.lowercased().replacingOccurrences(of: "-", with: "") + } + + /// Create an ISO 8601 timestamp for the current time. + private func currentTimestamp() -> String { + ISO8601DateFormatter().string(from: Date()) + } + + public func createTask(metadata: TaskMetadata, taskId: String?) async throws -> MCPTask { + cleanUpExpired() + + let id = taskId ?? generateTaskId() + + guard tasks[id] == nil else { + throw MCPError.invalidRequest("Task with ID \(id) already exists") + } + + let now = currentTimestamp() + let task = MCPTask( + taskId: id, + status: .working, + ttl: metadata.ttl, + createdAt: now, + lastUpdatedAt: now, + pollInterval: 1000 // Default 1 second poll interval + ) + + tasks[id] = StoredTask( + task: task, + result: nil, + expiresAt: calculateExpiry(ttl: metadata.ttl) + ) + + return task + } + + public func getTask(taskId: String) async -> MCPTask? { + cleanUpExpired() + return tasks[taskId]?.task + } + + public func updateTask(taskId: String, status: TaskStatus?, statusMessage: String?) async throws -> MCPTask { + guard var stored = tasks[taskId] else { + throw MCPError.invalidParams("Task with ID \(taskId) not found") + } + + // Per spec: Terminal states MUST NOT transition to any other status + if let newStatus = status, newStatus != stored.task.status, isTerminalStatus(stored.task.status) { + throw MCPError.invalidRequest("Cannot transition from terminal status '\(stored.task.status.rawValue)'") + } + + if let newStatus = status { + stored.task.status = newStatus + } + + if let message = statusMessage { + stored.task.statusMessage = message + } + + stored.task.lastUpdatedAt = currentTimestamp() + + // If task is now terminal and has TTL, reset expiry timer + if let newStatus = status, isTerminalStatus(newStatus), let ttl = stored.task.ttl { + stored.expiresAt = calculateExpiry(ttl: ttl) + } + + tasks[taskId] = stored + + // Notify waiters that the task has been updated + await notifyUpdate(taskId: taskId) + + return stored.task + } + + public func storeResult(taskId: String, result: Value) async throws { + guard var stored = tasks[taskId] else { + throw MCPError.invalidParams("Task with ID \(taskId) not found") + } + + stored.result = result + tasks[taskId] = stored + + // Notify waiters that the task has been updated + await notifyUpdate(taskId: taskId) + } + + public func getResult(taskId: String) async -> Value? { + tasks[taskId]?.result + } + + public func listTasks(cursor: String?) async -> (tasks: [MCPTask], nextCursor: String?) { + cleanUpExpired() + + let allTaskIds = Array(tasks.keys).sorted() + + var startIndex = 0 + if let cursor { + if let index = allTaskIds.firstIndex(of: cursor) { + startIndex = index + 1 + } + } + + let pageTaskIds = Array(allTaskIds.dropFirst(startIndex).prefix(pageSize)) + let pageTasks = pageTaskIds.compactMap { tasks[$0]?.task } + + let nextCursor: String? = if startIndex + pageSize < allTaskIds.count, let lastId = pageTaskIds.last { + lastId + } else { + nil + } + + return (tasks: pageTasks, nextCursor: nextCursor) + } + + public func deleteTask(taskId: String) async -> Bool { + tasks.removeValue(forKey: taskId) != nil + } + + /// Clear all tasks (useful for testing or graceful shutdown). + public func cleanUp() { + tasks.removeAll() + // Cancel all waiters + for (_, taskWaiters) in waiters { + for waiter in taskWaiters { + waiter.continuation.resume(throwing: CancellationError()) + } + } + waiters.removeAll() + } + + /// Get all tasks (useful for debugging). + public func getAllTasks() -> [MCPTask] { + cleanUpExpired() + return tasks.values.map(\.task) + } + + public func waitForUpdate(taskId: String) async throws { + let waiterId = UUID() + + try await withTaskCancellationHandler { + try await withCheckedThrowingContinuation { (continuation: CheckedContinuation) in + waiters[taskId, default: []].append(Waiter(id: waiterId, continuation: continuation)) + } + } onCancel: { + // Schedule cancellation on the actor + // Note: This runs synchronously when the Task is cancelled + Task { [weak self] in + await self?.cancelWaiter(taskId: taskId, waiterId: waiterId) + } + } + } + + /// Cancel a specific waiter by ID. + /// Called when the waiting Task is cancelled. + private func cancelWaiter(taskId: String, waiterId: UUID) { + guard var taskWaiters = waiters[taskId] else { return } + + if let index = taskWaiters.firstIndex(where: { $0.id == waiterId }) { + let waiter = taskWaiters.remove(at: index) + waiter.continuation.resume(throwing: CancellationError()) + + if taskWaiters.isEmpty { + waiters.removeValue(forKey: taskId) + } else { + waiters[taskId] = taskWaiters + } + } + } + + public func notifyUpdate(taskId: String) async { + guard let taskWaiters = waiters.removeValue(forKey: taskId), !taskWaiters.isEmpty else { + return + } + for waiter in taskWaiters { + waiter.continuation.resume() + } + } +} diff --git a/Sources/MCP/Server/Experimental/Tasks/TaskSupport.swift b/Sources/MCP/Server/Experimental/Tasks/TaskSupport.swift new file mode 100644 index 00000000..93fff668 --- /dev/null +++ b/Sources/MCP/Server/Experimental/Tasks/TaskSupport.swift @@ -0,0 +1,287 @@ +import Foundation + +// MARK: - Task Mode Validation + +/// Validates that a request is compatible with a tool's task execution mode. +/// +/// Per MCP spec: +/// - `required`: Clients MUST invoke as task. Server returns error if not. +/// - `forbidden` (or nil): Clients MUST NOT invoke as task. Server returns error if they do. +/// - `optional`: Either is acceptable. +/// +/// - Parameters: +/// - isTaskRequest: Whether the request includes task metadata +/// - taskSupport: The tool's task support setting (nil defaults to .forbidden) +/// - Throws: MCPError if the request is incompatible with the tool's task mode +public func validateTaskMode( + isTaskRequest: Bool, + taskSupport: Tool.Execution.TaskSupport? +) throws { + let mode = taskSupport ?? .forbidden + + switch mode { + case .required: + if !isTaskRequest { + throw MCPError.methodNotFound("This tool requires task-augmented invocation") + } + case .forbidden: + if isTaskRequest { + throw MCPError.methodNotFound("This tool does not support task-augmented invocation") + } + case .optional: + // Both task and non-task requests are acceptable + break + } +} + +/// Validates that a request is compatible with a tool's configuration. +/// +/// - Parameters: +/// - isTaskRequest: Whether the request includes task metadata +/// - tool: The tool being invoked +/// - Throws: MCPError if the request is incompatible with the tool's configuration +public func validateTaskMode(isTaskRequest: Bool, for tool: Tool) throws { + let taskSupport = tool.execution?.taskSupport + try validateTaskMode(isTaskRequest: isTaskRequest, taskSupport: taskSupport) +} + +/// Check if a client can invoke a tool with the given task mode. +/// +/// - Parameters: +/// - clientSupportsTask: Whether the client supports task-augmented requests +/// - taskSupport: The tool's task support setting +/// - Returns: True if the client can use this tool +public func canUseToolWithTaskMode( + clientSupportsTask: Bool, + taskSupport: Tool.Execution.TaskSupport? +) -> Bool { + let mode = taskSupport ?? .forbidden + switch mode { + case .required: + return clientSupportsTask + case .forbidden, .optional: + return true + } +} + +// MARK: - Task Support Configuration + +/// Configuration for experimental task support on the server. +/// +/// TaskSupport encapsulates the task store and message queue infrastructure +/// needed for task-augmented requests. When enabled on a server, it provides +/// default handlers for task operations. +/// +/// - Important: This is an experimental API that may change without notice. +/// +/// ## Example +/// +/// ```swift +/// let server = Server(name: "MyServer", version: "1.0") +/// +/// // Enable task support with in-memory storage +/// let taskSupport = TaskSupport.inMemory() +/// server.enableTaskSupport(taskSupport) +/// ``` +public final class TaskSupport: Sendable { + /// The task store for persisting task state. + public let store: any TaskStore + + /// The message queue for side-channel communication during task execution. + public let queue: any TaskMessageQueue + + /// The result handler for processing tasks/result requests. + public let resultHandler: TaskResultHandler + + /// Create task support with custom store and queue. + /// + /// - Parameters: + /// - store: The task store implementation + /// - queue: The message queue implementation + public init(store: any TaskStore, queue: any TaskMessageQueue) { + self.store = store + self.queue = queue + self.resultHandler = TaskResultHandler(store: store, queue: queue) + } + + /// Create in-memory task support. + /// + /// Suitable for development, testing, and single-process servers. + /// For distributed systems, provide custom store and queue implementations. + /// + /// - Returns: TaskSupport configured with in-memory store and queue + public static func inMemory() -> TaskSupport { + TaskSupport( + store: InMemoryTaskStore(), + queue: InMemoryTaskMessageQueue() + ) + } + + /// Run a work function as a background task. + /// + /// This is the recommended way to handle task-augmented tool calls. It: + /// 1. Creates a task in the store + /// 2. Spawns the work function in a background task + /// 3. Returns `CreateTaskResult` immediately + /// + /// The work function receives a `ServerTaskContext` with: + /// - `updateStatus()` for progress updates + /// - `complete(result:)` / `fail(error:)` for finishing the task + /// - `isCancelled` to check for cancellation + /// + /// - Important: This is an experimental API that may change without notice. + /// + /// ## Example + /// + /// ```swift + /// // In your tool handler: + /// server.withRequestHandler(CallTool.self) { params, context in + /// guard let taskMetadata = params.task else { + /// // Handle non-task request + /// return CallTool.Result(content: [.text("Done")]) + /// } + /// + /// // Run as a task + /// let createTaskResult = try await taskSupport.runTask( + /// metadata: taskMetadata, + /// modelImmediateResponse: "Starting to process..." + /// ) { taskContext in + /// try await taskContext.updateStatus("Working...") + /// // Do work... + /// return CallTool.Result(content: [.text("Done!")]) + /// } + /// + /// // Return CreateTaskResult as the response + /// // Note: This requires the response type to be flexible + /// } + /// ``` + /// + /// - Parameters: + /// - metadata: The task metadata from the request + /// - taskId: Optional specific task ID (generated if nil) + /// - modelImmediateResponse: Optional immediate feedback for the model + /// - clientCapabilities: Optional client capabilities for mid-task elicitation/sampling + /// - server: Optional server reference for task-augmented requests (elicitAsTask, createMessageAsTask) + /// - work: The async work function that receives the task context + /// - Returns: CreateTaskResult to return to the client + /// - Throws: Error if task creation fails + public func runTask( + metadata: TaskMetadata, + taskId: String? = nil, + modelImmediateResponse: String? = nil, + clientCapabilities: Client.Capabilities? = nil, + server: Server? = nil, + work: @escaping @Sendable (ServerTaskContext) async throws -> CallTool.Result + ) async throws -> CreateTaskResult { + // Create the task + let task = try await store.createTask(metadata: metadata, taskId: taskId) + + // Create the context + let context = ServerTaskContext( + task: task, + store: store, + queue: queue, + clientCapabilities: clientCapabilities, + server: server + ) + + // Spawn the work in a background task + Task.detached { [store] in + do { + let result = try await work(context) + // If the task isn't already in a terminal state, complete it + if !isTerminalStatus(context.task.status) { + try await context.complete(toolResult: result) + } + } catch is CancellationError { + // If cancelled, update status + if !isTerminalStatus(context.task.status) { + _ = try? await store.updateTask( + taskId: context.taskId, + status: .cancelled, + statusMessage: "Cancelled" + ) + } + } catch { + // If the task isn't already in a terminal state, fail it + if !isTerminalStatus(context.task.status) { + try? await context.fail(error: error) + } + } + } + + // Return immediately + return CreateTaskResult(task: task, modelImmediateResponse: modelImmediateResponse) + } +} + +// MARK: - Server Extension + +extension Server { + // Note: This method is internal. Access via server.experimental.enableTasks() + @discardableResult + func enableTaskSupport(_ taskSupport: TaskSupport) -> Self { + // Set the tasks capability with full support + capabilities.tasks = .full() + + // Register the result handler as a response router + // This routes responses back to waiting task handlers (elicit/createMessage) + addResponseRouter(taskSupport.resultHandler) + + // Register default task handlers + registerDefaultTaskHandlers(taskSupport) + + return self + } + + /// Register default handlers for task operations. + private func registerDefaultTaskHandlers(_ taskSupport: TaskSupport) { + // tasks/get - Get task status + withRequestHandler(GetTask.self) { params, _ in + guard let task = await taskSupport.store.getTask(taskId: params.taskId) else { + throw MCPError.invalidParams("Task not found: \(params.taskId)") + } + return GetTask.Result(task: task) + } + + // tasks/list - List all tasks + withRequestHandler(ListTasks.self) { params, _ in + let (tasks, nextCursor) = await taskSupport.store.listTasks(cursor: params.cursor) + return ListTasks.Result(tasks: tasks, nextCursor: nextCursor) + } + + // tasks/cancel - Cancel a running task + withRequestHandler(CancelTask.self) { params, _ in + guard let task = await taskSupport.store.getTask(taskId: params.taskId) else { + throw MCPError.invalidParams("Task not found: \(params.taskId)") + } + + // Can't cancel a task that's already in a terminal state + // Per spec: return -32602 (Invalid params) for terminal status tasks + if isTerminalStatus(task.status) { + throw MCPError.invalidParams("Cannot cancel task in terminal status: \(task.status.rawValue)") + } + + // Update task status to cancelled + let updatedTask = try await taskSupport.store.updateTask( + taskId: params.taskId, + status: .cancelled, + statusMessage: "Cancelled by client request" + ) + + // Clean up any queued messages for this task + _ = await taskSupport.queue.dequeueAll(taskId: params.taskId) + + return CancelTask.Result(task: updatedTask) + } + + // tasks/result - Get task result (with blocking until terminal) + // Uses TaskResultHandler to deliver queued messages (elicitation/sampling) + withRequestHandler(GetTaskPayload.self) { params, context in + try await taskSupport.resultHandler.handle( + taskId: params.taskId, + sendMessage: context.sendData + ) + } + } +} diff --git a/Sources/MCP/Server/Experimental/Tasks/Tasks.swift b/Sources/MCP/Server/Experimental/Tasks/Tasks.swift new file mode 100644 index 00000000..58e944ec --- /dev/null +++ b/Sources/MCP/Server/Experimental/Tasks/Tasks.swift @@ -0,0 +1,1013 @@ +import Foundation + +/// Tasks provide a way to track the progress of long-running operations. +/// This is an experimental feature in MCP protocol version 2025-11-25. +/// +/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2025-11-25/ + +// MARK: - Task Metadata Keys + +/// Metadata key for associating messages with a related task. +/// +/// This constant is used in the `_meta` field of requests, responses, and notifications +/// to indicate they are related to a specific task. +/// +/// ## Example +/// +/// ```swift +/// let meta: [String: Value] = [ +/// relatedTaskMetaKey: .object(["taskId": .string(taskId)]) +/// ] +/// ``` +public let relatedTaskMetaKey = "io.modelcontextprotocol/related-task" + +/// Metadata key for providing an immediate response to the model while a task continues. +/// +/// When a task is created, the server can include this in `_meta` to provide +/// immediate feedback to the model while the actual work continues in the background. +/// +/// ## Example +/// +/// ```swift +/// let meta: [String: Value] = [ +/// modelImmediateResponseKey: .string("Starting to process your request...") +/// ] +/// ``` +public let modelImmediateResponseKey = "io.modelcontextprotocol/model-immediate-response" + +// MARK: - Task Status + +/// The status of a task. +public enum TaskStatus: String, Hashable, Codable, Sendable { + /// Task is actively being worked on + case working + /// Task requires user input to continue + case inputRequired = "input_required" + /// Task completed successfully + case completed + /// Task failed + case failed + /// Task was cancelled + case cancelled + + /// Whether this status represents a terminal state. + /// + /// Terminal states are: completed, failed, cancelled. + /// Once a task reaches a terminal state, no further status updates will occur. + public var isTerminal: Bool { + switch self { + case .completed, .failed, .cancelled: + return true + case .working, .inputRequired: + return false + } + } +} + +// MARK: - Task + +/// Represents a running or completed task. +/// +/// Note: This type represents the `Task` schema which does not include `_meta`. +/// When combined with `Result` or `NotificationParams` (via allOf), the `_meta` +/// field comes from those base types. +public struct MCPTask: Hashable, Sendable { + /// Unique identifier for the task + public var taskId: String + /// Current status of the task + public var status: TaskStatus + /// Time in milliseconds to keep task results available after completion. + /// If nil, the task has unlimited lifetime until manually cleaned up. + /// Note: Per the MCP spec, this field is always present in the JSON (encoded as null when nil). + public var ttl: Int? + /// ISO 8601 timestamp when the task was created + public var createdAt: String + /// ISO 8601 timestamp when the task was last updated + public var lastUpdatedAt: String + /// Suggested polling interval in milliseconds for clients + public var pollInterval: Int? + /// Optional diagnostic message for failed tasks or other status information + public var statusMessage: String? + + public init( + taskId: String, + status: TaskStatus, + ttl: Int? = nil, + createdAt: String, + lastUpdatedAt: String, + pollInterval: Int? = nil, + statusMessage: String? = nil + ) { + self.taskId = taskId + self.status = status + self.ttl = ttl + self.createdAt = createdAt + self.lastUpdatedAt = lastUpdatedAt + self.pollInterval = pollInterval + self.statusMessage = statusMessage + } +} + +extension MCPTask: Codable { + enum CodingKeys: String, CodingKey { + case taskId, status, ttl, createdAt, lastUpdatedAt, pollInterval, statusMessage + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + taskId = try container.decode(String.self, forKey: .taskId) + status = try container.decode(TaskStatus.self, forKey: .status) + // ttl is required in the spec but can be null for unlimited + ttl = try container.decode(Int?.self, forKey: .ttl) + createdAt = try container.decode(String.self, forKey: .createdAt) + lastUpdatedAt = try container.decode(String.self, forKey: .lastUpdatedAt) + pollInterval = try container.decodeIfPresent(Int.self, forKey: .pollInterval) + statusMessage = try container.decodeIfPresent(String.self, forKey: .statusMessage) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode(taskId, forKey: .taskId) + try container.encode(status, forKey: .status) + // ttl is required in the spec - always encode it (as null when nil) + try container.encode(ttl, forKey: .ttl) + try container.encode(createdAt, forKey: .createdAt) + try container.encode(lastUpdatedAt, forKey: .lastUpdatedAt) + try container.encodeIfPresent(pollInterval, forKey: .pollInterval) + try container.encodeIfPresent(statusMessage, forKey: .statusMessage) + } +} + +/// Metadata for task creation, passed via `_meta.task` in request parameters. +/// +/// When a client sends a request that may become a long-running task, it can include +/// this metadata to configure task behavior. +/// +/// ## Example +/// +/// ```swift +/// // Include task metadata in a tool call +/// let params = CallTool.Parameters( +/// name: "long_running_operation", +/// arguments: ["input": .string("data")], +/// _meta: ["task": .object(["ttl": .int(3600000)])] // Keep for 1 hour +/// ) +/// ``` +public struct TaskMetadata: Hashable, Codable, Sendable { + /// Time-to-live in milliseconds for task results after completion. + /// If nil, the task has unlimited lifetime until manually cleaned up. + public var ttl: Int? + + public init(ttl: Int? = nil) { + self.ttl = ttl + } +} + +/// Metadata indicating an operation is related to an existing task. +/// +/// When a server sends notifications or requests during task execution, +/// it can include this metadata in `_meta.task` to associate them with +/// the originating task. +/// +/// ## Example +/// +/// ```swift +/// // Send progress notification for a task +/// let notification = ProgressNotification.Parameters( +/// progressToken: token, +/// progress: 50, +/// total: 100, +/// _meta: ["task": .object(["taskId": .string(taskId)])] +/// ) +/// ``` +public struct RelatedTaskMetadata: Hashable, Codable, Sendable { + /// The ID of the task this operation is related to. + public var taskId: String + + public init(taskId: String) { + self.taskId = taskId + } +} + +// MARK: - Create Task Result + +/// Result returned when a task-augmented request creates a task. +/// +/// When a client sends a request with a `task` field in the parameters, +/// the server returns this result instead of the normal method result. +/// The client can then poll for the actual result using `tasks/result`. +/// +/// ## Example +/// +/// ```swift +/// // Client sends task-augmented tool call +/// let params = CallTool.Parameters( +/// name: "long_running_tool", +/// arguments: ["input": .string("data")], +/// task: TaskMetadata(ttl: 60000) +/// ) +/// +/// // Server returns CreateTaskResult instead of CallTool.Result +/// let createTaskResult = CreateTaskResult( +/// task: MCPTask( +/// taskId: "abc123", +/// status: .working, +/// ttl: 60000, +/// createdAt: ISO8601DateFormatter().string(from: Date()), +/// lastUpdatedAt: ISO8601DateFormatter().string(from: Date()), +/// pollInterval: 1000 +/// ) +/// ) +/// ``` +public struct CreateTaskResult: ResultWithExtraFields { + public typealias ResultCodingKeys = CodingKeys + + /// The created task. + public var task: MCPTask + /// Reserved for clients and servers to attach additional metadata. + /// May include `io.modelcontextprotocol/model-immediate-response` for feedback. + public var _meta: [String: Value]? + /// Additional fields not defined in the schema (for forward compatibility). + public var extraFields: [String: Value]? + + public init( + task: MCPTask, + _meta: [String: Value]? = nil, + extraFields: [String: Value]? = nil + ) { + self.task = task + self._meta = _meta + self.extraFields = extraFields + } + + /// Convenience initializer with optional model immediate response. + /// + /// - Parameters: + /// - task: The created task + /// - modelImmediateResponse: Optional immediate feedback for the model + public init(task: MCPTask, modelImmediateResponse: String?) { + self.task = task + if let response = modelImmediateResponse { + self._meta = [modelImmediateResponseKey: .string(response)] + } else { + self._meta = nil + } + self.extraFields = nil + } + + public enum CodingKeys: String, CodingKey, CaseIterable { + case task, _meta + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + task = try container.decode(MCPTask.self, forKey: .task) + _meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + extraFields = try Self.decodeExtraFields(from: decoder) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode(task, forKey: .task) + try container.encodeIfPresent(_meta, forKey: ._meta) + try encodeExtraFields(to: encoder) + } +} + +// MARK: - Get Task + +/// Request to get information about a specific task. +public enum GetTask: Method { + public static let name = "tasks/get" + + public struct Parameters: Hashable, Codable, Sendable { + /// The ID of the task to retrieve + public var taskId: String + /// Request metadata including progress token. + public var _meta: RequestMeta? + + public init(taskId: String, _meta: RequestMeta? = nil) { + self.taskId = taskId + self._meta = _meta + } + } + + /// The response to a tasks/get request. + /// + /// This type flattens `Result` and `Task` fields per the spec's `allOf[Result, Task]`. + public struct Result: ResultWithExtraFields { + public typealias ResultCodingKeys = CodingKeys + + // Task fields (flattened from MCPTask) + /// Unique identifier for the task + public var taskId: String + /// Current status of the task + public var status: TaskStatus + /// Time in milliseconds to keep task results available after completion. + public var ttl: Int? + /// ISO 8601 timestamp when the task was created + public var createdAt: String + /// ISO 8601 timestamp when the task was last updated + public var lastUpdatedAt: String + /// Suggested polling interval in milliseconds for clients + public var pollInterval: Int? + /// Optional diagnostic message for failed tasks or other status information + public var statusMessage: String? + + // Result fields + /// Reserved for clients and servers to attach additional metadata + public var _meta: [String: Value]? + /// Additional fields not defined in the schema (for forward compatibility). + public var extraFields: [String: Value]? + + public init( + taskId: String, + status: TaskStatus, + ttl: Int? = nil, + createdAt: String, + lastUpdatedAt: String, + pollInterval: Int? = nil, + statusMessage: String? = nil, + _meta: [String: Value]? = nil, + extraFields: [String: Value]? = nil + ) { + self.taskId = taskId + self.status = status + self.ttl = ttl + self.createdAt = createdAt + self.lastUpdatedAt = lastUpdatedAt + self.pollInterval = pollInterval + self.statusMessage = statusMessage + self._meta = _meta + self.extraFields = extraFields + } + + /// Convenience initializer from MCPTask + public init(task: MCPTask, _meta: [String: Value]? = nil, extraFields: [String: Value]? = nil) { + self.taskId = task.taskId + self.status = task.status + self.ttl = task.ttl + self.createdAt = task.createdAt + self.lastUpdatedAt = task.lastUpdatedAt + self.pollInterval = task.pollInterval + self.statusMessage = task.statusMessage + self._meta = _meta + self.extraFields = extraFields + } + + public enum CodingKeys: String, CodingKey, CaseIterable { + case taskId, status, ttl, createdAt, lastUpdatedAt, pollInterval, statusMessage, _meta + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + taskId = try container.decode(String.self, forKey: .taskId) + status = try container.decode(TaskStatus.self, forKey: .status) + ttl = try container.decodeIfPresent(Int.self, forKey: .ttl) + createdAt = try container.decode(String.self, forKey: .createdAt) + lastUpdatedAt = try container.decode(String.self, forKey: .lastUpdatedAt) + pollInterval = try container.decodeIfPresent(Int.self, forKey: .pollInterval) + statusMessage = try container.decodeIfPresent(String.self, forKey: .statusMessage) + _meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + extraFields = try Self.decodeExtraFields(from: decoder) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode(taskId, forKey: .taskId) + try container.encode(status, forKey: .status) + // ttl is required in the spec - always encode it (as null when nil) + try container.encode(ttl, forKey: .ttl) + try container.encode(createdAt, forKey: .createdAt) + try container.encode(lastUpdatedAt, forKey: .lastUpdatedAt) + try container.encodeIfPresent(pollInterval, forKey: .pollInterval) + try container.encodeIfPresent(statusMessage, forKey: .statusMessage) + try container.encodeIfPresent(_meta, forKey: ._meta) + try encodeExtraFields(to: encoder) + } + } +} + +// MARK: - Get Task Payload + +/// Request to get the result payload of a completed task. +/// +/// This method retrieves the actual result data for a completed task. +/// The result type depends on the original request that created the task +/// (e.g., a tool call result for a task created from `tools/call`). +/// +/// - Note: This should only be called for tasks with status `.completed`. +/// For failed or cancelled tasks, check `MCPTask.statusMessage` instead. +public enum GetTaskPayload: Method { + public static let name = "tasks/result" + + public struct Parameters: Hashable, Codable, Sendable { + /// The ID of the task to get results for + public var taskId: String + /// Request metadata including progress token. + public var _meta: RequestMeta? + + public init(taskId: String, _meta: RequestMeta? = nil) { + self.taskId = taskId + self._meta = _meta + } + } + + /// The result type for tasks/result. + /// + /// Per the MCP spec, this is a "loose" Result type where the actual result fields + /// are flattened directly into the response (via `extraFields`), not wrapped + /// in a separate field. For example, a tools/call task would have `content` and + /// `isError` as top-level fields in the response. + /// + /// ## Example response for a completed tools/call task: + /// ```json + /// { + /// "_meta": {"io.modelcontextprotocol/related-task": {"taskId": "..."}}, + /// "content": [{"type": "text", "text": "Result"}], + /// "isError": false + /// } + /// ``` + public struct Result: ResultWithExtraFields { + public typealias ResultCodingKeys = CodingKeys + + /// Reserved for clients and servers to attach additional metadata. + /// Typically includes `io.modelcontextprotocol/related-task` with the task ID. + public var _meta: [String: Value]? + + /// The actual result payload fields from the original request's result type. + /// For a tools/call task, this would contain `content`, `isError`, etc. + /// These fields are encoded/decoded as top-level fields in the JSON. + public var extraFields: [String: Value]? + + public init(_meta: [String: Value]? = nil, extraFields: [String: Value]? = nil) { + self._meta = _meta + self.extraFields = extraFields + } + + /// Convenience initializer from a Value representing the original result. + /// + /// This extracts the fields from the Value and stores them in extraFields. + /// - Parameters: + /// - resultValue: The result as a Value (typically from task storage) + /// - _meta: Optional metadata + public init(fromResultValue resultValue: Value?, _meta: [String: Value]? = nil) { + self._meta = _meta + if case .object(let fields) = resultValue { + self.extraFields = fields + } else { + self.extraFields = nil + } + } + + public enum CodingKeys: String, CodingKey, CaseIterable { + case _meta + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + _meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + extraFields = try Self.decodeExtraFields(from: decoder) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encodeIfPresent(_meta, forKey: ._meta) + try encodeExtraFields(to: encoder) + } + } +} + +// MARK: - List Tasks + +/// Request to list all tasks. +public enum ListTasks: Method { + public static let name = "tasks/list" + + public struct Parameters: NotRequired, Hashable, Codable, Sendable { + /// Pagination cursor + public var cursor: String? + /// Request metadata including progress token. + public var _meta: RequestMeta? + + public init() { + self.cursor = nil + self._meta = nil + } + + public init(cursor: String? = nil, _meta: RequestMeta? = nil) { + self.cursor = cursor + self._meta = _meta + } + } + + public struct Result: ResultWithExtraFields { + public typealias ResultCodingKeys = CodingKeys + + /// List of tasks + public var tasks: [MCPTask] + /// Next pagination cursor + public var nextCursor: String? + /// Reserved for clients and servers to attach additional metadata + public var _meta: [String: Value]? + /// Additional fields not defined in the schema (for forward compatibility). + public var extraFields: [String: Value]? + + public init( + tasks: [MCPTask], + nextCursor: String? = nil, + _meta: [String: Value]? = nil, + extraFields: [String: Value]? = nil + ) { + self.tasks = tasks + self.nextCursor = nextCursor + self._meta = _meta + self.extraFields = extraFields + } + + public enum CodingKeys: String, CodingKey, CaseIterable { + case tasks, nextCursor, _meta + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + tasks = try container.decode([MCPTask].self, forKey: .tasks) + nextCursor = try container.decodeIfPresent(String.self, forKey: .nextCursor) + _meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + extraFields = try Self.decodeExtraFields(from: decoder) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode(tasks, forKey: .tasks) + try container.encodeIfPresent(nextCursor, forKey: .nextCursor) + try container.encodeIfPresent(_meta, forKey: ._meta) + try encodeExtraFields(to: encoder) + } + } +} + +// MARK: - Cancel Task + +/// Request to cancel a running task. +public enum CancelTask: Method { + public static let name = "tasks/cancel" + + public struct Parameters: Hashable, Codable, Sendable { + /// The ID of the task to cancel + public var taskId: String + /// Request metadata including progress token. + public var _meta: RequestMeta? + + public init(taskId: String, _meta: RequestMeta? = nil) { + self.taskId = taskId + self._meta = _meta + } + } + + /// The response to a tasks/cancel request. + /// + /// This type flattens `Result` and `Task` fields per the spec's `allOf[Result, Task]`. + public struct Result: ResultWithExtraFields { + public typealias ResultCodingKeys = CodingKeys + + // Task fields (flattened from MCPTask) + /// Unique identifier for the task + public var taskId: String + /// Current status of the task + public var status: TaskStatus + /// Time in milliseconds to keep task results available after completion. + public var ttl: Int? + /// ISO 8601 timestamp when the task was created + public var createdAt: String + /// ISO 8601 timestamp when the task was last updated + public var lastUpdatedAt: String + /// Suggested polling interval in milliseconds for clients + public var pollInterval: Int? + /// Optional diagnostic message for failed tasks or other status information + public var statusMessage: String? + + // Result fields + /// Reserved for clients and servers to attach additional metadata + public var _meta: [String: Value]? + /// Additional fields not defined in the schema (for forward compatibility). + public var extraFields: [String: Value]? + + public init( + taskId: String, + status: TaskStatus, + ttl: Int? = nil, + createdAt: String, + lastUpdatedAt: String, + pollInterval: Int? = nil, + statusMessage: String? = nil, + _meta: [String: Value]? = nil, + extraFields: [String: Value]? = nil + ) { + self.taskId = taskId + self.status = status + self.ttl = ttl + self.createdAt = createdAt + self.lastUpdatedAt = lastUpdatedAt + self.pollInterval = pollInterval + self.statusMessage = statusMessage + self._meta = _meta + self.extraFields = extraFields + } + + /// Convenience initializer from MCPTask + public init(task: MCPTask, _meta: [String: Value]? = nil, extraFields: [String: Value]? = nil) { + self.taskId = task.taskId + self.status = task.status + self.ttl = task.ttl + self.createdAt = task.createdAt + self.lastUpdatedAt = task.lastUpdatedAt + self.pollInterval = task.pollInterval + self.statusMessage = task.statusMessage + self._meta = _meta + self.extraFields = extraFields + } + + public enum CodingKeys: String, CodingKey, CaseIterable { + case taskId, status, ttl, createdAt, lastUpdatedAt, pollInterval, statusMessage, _meta + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + taskId = try container.decode(String.self, forKey: .taskId) + status = try container.decode(TaskStatus.self, forKey: .status) + ttl = try container.decodeIfPresent(Int.self, forKey: .ttl) + createdAt = try container.decode(String.self, forKey: .createdAt) + lastUpdatedAt = try container.decode(String.self, forKey: .lastUpdatedAt) + pollInterval = try container.decodeIfPresent(Int.self, forKey: .pollInterval) + statusMessage = try container.decodeIfPresent(String.self, forKey: .statusMessage) + _meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + extraFields = try Self.decodeExtraFields(from: decoder) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode(taskId, forKey: .taskId) + try container.encode(status, forKey: .status) + // ttl is required in the spec - always encode it (as null when nil) + try container.encode(ttl, forKey: .ttl) + try container.encode(createdAt, forKey: .createdAt) + try container.encode(lastUpdatedAt, forKey: .lastUpdatedAt) + try container.encodeIfPresent(pollInterval, forKey: .pollInterval) + try container.encodeIfPresent(statusMessage, forKey: .statusMessage) + try container.encodeIfPresent(_meta, forKey: ._meta) + try encodeExtraFields(to: encoder) + } + } +} + +// MARK: - Task Status Notification + +/// Notification sent when a task's status changes. +public struct TaskStatusNotification: Notification { + public static let name = "notifications/tasks/status" + + /// Parameters for task status notification. + /// + /// This type flattens `NotificationParams` and `Task` fields per the spec's + /// `allOf[NotificationParams, Task]`. + public struct Parameters: Hashable, Sendable { + // Task fields (flattened from MCPTask) + /// Unique identifier for the task + public var taskId: String + /// Current status of the task + public var status: TaskStatus + /// Time in milliseconds to keep task results available after completion. + public var ttl: Int? + /// ISO 8601 timestamp when the task was created + public var createdAt: String + /// ISO 8601 timestamp when the task was last updated + public var lastUpdatedAt: String + /// Suggested polling interval in milliseconds for clients + public var pollInterval: Int? + /// Optional diagnostic message for failed tasks or other status information + public var statusMessage: String? + + // NotificationParams fields + /// Reserved for additional metadata. + public var _meta: [String: Value]? + + public init( + taskId: String, + status: TaskStatus, + ttl: Int? = nil, + createdAt: String, + lastUpdatedAt: String, + pollInterval: Int? = nil, + statusMessage: String? = nil, + _meta: [String: Value]? = nil + ) { + self.taskId = taskId + self.status = status + self.ttl = ttl + self.createdAt = createdAt + self.lastUpdatedAt = lastUpdatedAt + self.pollInterval = pollInterval + self.statusMessage = statusMessage + self._meta = _meta + } + + /// Convenience initializer from MCPTask + public init(task: MCPTask, _meta: [String: Value]? = nil) { + self.taskId = task.taskId + self.status = task.status + self.ttl = task.ttl + self.createdAt = task.createdAt + self.lastUpdatedAt = task.lastUpdatedAt + self.pollInterval = task.pollInterval + self.statusMessage = task.statusMessage + self._meta = _meta + } + + enum CodingKeys: String, CodingKey { + case taskId, status, ttl, createdAt, lastUpdatedAt, pollInterval, statusMessage, _meta + } + } +} + +extension TaskStatusNotification.Parameters: Codable { + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + taskId = try container.decode(String.self, forKey: .taskId) + status = try container.decode(TaskStatus.self, forKey: .status) + ttl = try container.decodeIfPresent(Int.self, forKey: .ttl) + createdAt = try container.decode(String.self, forKey: .createdAt) + lastUpdatedAt = try container.decode(String.self, forKey: .lastUpdatedAt) + pollInterval = try container.decodeIfPresent(Int.self, forKey: .pollInterval) + statusMessage = try container.decodeIfPresent(String.self, forKey: .statusMessage) + _meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode(taskId, forKey: .taskId) + try container.encode(status, forKey: .status) + // ttl is required in the spec - always encode it (as null when nil) + try container.encode(ttl, forKey: .ttl) + try container.encode(createdAt, forKey: .createdAt) + try container.encode(lastUpdatedAt, forKey: .lastUpdatedAt) + try container.encodeIfPresent(pollInterval, forKey: .pollInterval) + try container.encodeIfPresent(statusMessage, forKey: .statusMessage) + try container.encodeIfPresent(_meta, forKey: ._meta) + } +} + +// MARK: - Server Capabilities + +extension Server.Capabilities { + /// Tasks capabilities for servers. + /// + /// Servers advertise these capabilities during initialization to indicate + /// what task-related features they support. + /// + /// ## Example + /// + /// ```swift + /// let capabilities = Server.Capabilities( + /// tasks: .init( + /// list: .init(), + /// cancel: .init(), + /// requests: .init(tools: .init(call: .init())) + /// ) + /// ) + /// ``` + public struct Tasks: Hashable, Codable, Sendable { + /// Capability marker for list operations. + public struct List: Hashable, Codable, Sendable { + public init() {} + } + + /// Capability marker for cancel operations. + public struct Cancel: Hashable, Codable, Sendable { + public init() {} + } + + /// Task-augmented request capabilities. + public struct Requests: Hashable, Codable, Sendable { + /// Tools request capabilities. + public struct Tools: Hashable, Codable, Sendable { + /// Capability marker for task-augmented tools/call. + public struct Call: Hashable, Codable, Sendable { + public init() {} + } + + /// Whether task-augmented tools/call is supported. + public var call: Call? + + public init(call: Call? = nil) { + self.call = call + } + } + + /// Whether task-augmented tools requests are supported. + public var tools: Tools? + + public init(tools: Tools? = nil) { + self.tools = tools + } + } + + /// Whether the server supports tasks/list. + public var list: List? + /// Whether the server supports tasks/cancel. + public var cancel: Cancel? + /// Task-augmented request capabilities. + public var requests: Requests? + + public init( + list: List? = nil, + cancel: Cancel? = nil, + requests: Requests? = nil + ) { + self.list = list + self.cancel = cancel + self.requests = requests + } + + /// Convenience initializer for full task support. + /// + /// Creates a capability declaration with list, cancel, and task-augmented tools/call. + public static func full() -> Tasks { + Tasks( + list: List(), + cancel: Cancel(), + requests: Requests(tools: .init(call: .init())) + ) + } + } +} + +// MARK: - Client Capabilities + +extension Client.Capabilities { + /// Tasks capabilities for clients. + /// + /// Clients advertise these capabilities during initialization to indicate + /// what task-related features they support. This is for bidirectional task + /// support where servers can initiate tasks on clients. + /// + /// ## Example + /// + /// ```swift + /// let capabilities = Client.Capabilities( + /// tasks: .init( + /// list: .init(), + /// cancel: .init(), + /// requests: .init( + /// sampling: .init(createMessage: .init()), + /// elicitation: .init(create: .init()) + /// ) + /// ) + /// ) + /// ``` + public struct Tasks: Hashable, Codable, Sendable { + /// Capability marker for list operations. + public struct List: Hashable, Codable, Sendable { + public init() {} + } + + /// Capability marker for cancel operations. + public struct Cancel: Hashable, Codable, Sendable { + public init() {} + } + + /// Task-augmented request capabilities for client. + public struct Requests: Hashable, Codable, Sendable { + /// Sampling request capabilities. + public struct Sampling: Hashable, Codable, Sendable { + /// Capability marker for task-augmented sampling/createMessage. + public struct CreateMessage: Hashable, Codable, Sendable { + public init() {} + } + + /// Whether task-augmented sampling/createMessage is supported. + public var createMessage: CreateMessage? + + public init(createMessage: CreateMessage? = nil) { + self.createMessage = createMessage + } + } + + /// Elicitation request capabilities. + public struct Elicitation: Hashable, Codable, Sendable { + /// Capability marker for task-augmented elicitation/create. + public struct Create: Hashable, Codable, Sendable { + public init() {} + } + + /// Whether task-augmented elicitation/create is supported. + public var create: Create? + + public init(create: Create? = nil) { + self.create = create + } + } + + /// Whether task-augmented sampling requests are supported. + public var sampling: Sampling? + /// Whether task-augmented elicitation requests are supported. + public var elicitation: Elicitation? + + public init( + sampling: Sampling? = nil, + elicitation: Elicitation? = nil + ) { + self.sampling = sampling + self.elicitation = elicitation + } + } + + /// Whether the client supports tasks/list. + public var list: List? + /// Whether the client supports tasks/cancel. + public var cancel: Cancel? + /// Task-augmented request capabilities. + public var requests: Requests? + + public init( + list: List? = nil, + cancel: Cancel? = nil, + requests: Requests? = nil + ) { + self.list = list + self.cancel = cancel + self.requests = requests + } + + /// Convenience initializer for full task support. + /// + /// Creates a capability declaration with list, cancel, and all task-augmented requests. + public static func full() -> Tasks { + Tasks( + list: List(), + cancel: Cancel(), + requests: Requests( + sampling: .init(createMessage: .init()), + elicitation: .init(create: .init()) + ) + ) + } + } +} + +// MARK: - Capability Checking Helpers + +/// Check if server capabilities include task-augmented tools/call support. +/// +/// - Parameter caps: The server capabilities +/// - Returns: True if task-augmented tools/call is supported +public func hasTaskAugmentedToolsCall(_ caps: Server.Capabilities?) -> Bool { + caps?.tasks?.requests?.tools?.call != nil +} + +/// Check if client capabilities include task-augmented elicitation support. +/// +/// - Parameter caps: The client capabilities +/// - Returns: True if task-augmented elicitation/create is supported +public func hasTaskAugmentedElicitation(_ caps: Client.Capabilities?) -> Bool { + caps?.tasks?.requests?.elicitation?.create != nil +} + +/// Check if client capabilities include task-augmented sampling support. +/// +/// - Parameter caps: The client capabilities +/// - Returns: True if task-augmented sampling/createMessage is supported +public func hasTaskAugmentedSampling(_ caps: Client.Capabilities?) -> Bool { + caps?.tasks?.requests?.sampling?.createMessage != nil +} + +/// Require task-augmented elicitation support from client. +/// +/// - Parameter caps: The client capabilities +/// - Throws: MCPError if client doesn't support task-augmented elicitation +public func requireTaskAugmentedElicitation(_ caps: Client.Capabilities?) throws { + if !hasTaskAugmentedElicitation(caps) { + throw MCPError.invalidRequest("Client does not support task-augmented elicitation") + } +} + +/// Require task-augmented sampling support from client. +/// +/// - Parameter caps: The client capabilities +/// - Throws: MCPError if client doesn't support task-augmented sampling +public func requireTaskAugmentedSampling(_ caps: Client.Capabilities?) throws { + if !hasTaskAugmentedSampling(caps) { + throw MCPError.invalidRequest("Client does not support task-augmented sampling") + } +} + +/// Require task-augmented tools/call support from server. +/// +/// - Parameter caps: The server capabilities +/// - Throws: MCPError if server doesn't support task-augmented tools/call +public func requireTaskAugmentedToolsCall(_ caps: Server.Capabilities?) throws { + if !hasTaskAugmentedToolsCall(caps) { + throw MCPError.invalidRequest("Server does not support task-augmented tools/call") + } +} + diff --git a/Sources/MCP/Server/Logging.swift b/Sources/MCP/Server/Logging.swift new file mode 100644 index 00000000..f9b00cc6 --- /dev/null +++ b/Sources/MCP/Server/Logging.swift @@ -0,0 +1,100 @@ +/// Server logging capabilities. +/// +/// Servers can send log messages to clients, and clients can control +/// the minimum log level they wish to receive. +/// +/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2025-11-25/utilities/logging/ + +/// Log severity levels following RFC 5424 syslog conventions. +/// +/// Levels are ordered by increasing severity: +/// debug < info < notice < warning < error < critical < alert < emergency +public enum LoggingLevel: String, Hashable, Codable, Sendable, CaseIterable { + case debug + case info + case notice + case warning + case error + case critical + case alert + case emergency + + /// The severity index of this log level (0 = debug, 7 = emergency). + public var severity: Int { + switch self { + case .debug: return 0 + case .info: return 1 + case .notice: return 2 + case .warning: return 3 + case .error: return 4 + case .critical: return 5 + case .alert: return 6 + case .emergency: return 7 + } + } + + /// Returns true if this level is at least as severe as the given level. + public func isAtLeast(_ level: LoggingLevel) -> Bool { + self.severity >= level.severity + } +} + +/// Request from client to set the minimum log level for messages. +/// +/// After receiving this request, servers should only send log messages +/// at the specified level or higher (more severe). +/// +/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2025-11-25/utilities/logging/ +public enum SetLoggingLevel: Method { + public static let name: String = "logging/setLevel" + + public struct Parameters: Hashable, Codable, Sendable { + /// The minimum log level to receive. + public let level: LoggingLevel + /// Request metadata including progress token. + public var _meta: RequestMeta? + + public init(level: LoggingLevel, _meta: RequestMeta? = nil) { + self.level = level + self._meta = _meta + } + } + + public typealias Result = Empty +} + +/// Notification sent by servers to deliver log messages to clients. +/// +/// Servers should respect the log level set by the client via `SetLoggingLevel`, +/// only sending messages at or above that severity. +/// +/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2025-11-25/utilities/logging/ +public struct LogMessageNotification: Notification { + public static let name: String = "notifications/message" + + public struct Parameters: Hashable, Codable, Sendable { + /// The severity level of this log message. + public let level: LoggingLevel + + /// An optional name identifying the logger source. + public let logger: String? + + /// The log message data. Can be any JSON-serializable value. + public let data: Value + + /// Reserved for additional metadata. + public var _meta: [String: Value]? + + public init( + level: LoggingLevel, + logger: String? = nil, + data: Value, + _meta: [String: Value]? = nil + ) { + self.level = level + self.logger = logger + self.data = data + self._meta = _meta + } + } +} diff --git a/Sources/MCP/Server/Prompts.swift b/Sources/MCP/Server/Prompts.swift index c194b28a..c1b87946 100644 --- a/Sources/MCP/Server/Prompts.swift +++ b/Sources/MCP/Server/Prompts.swift @@ -7,32 +7,58 @@ import Foundation /// Clients can discover available prompts, retrieve their contents, /// and provide arguments to customize them. /// -/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2024-11-05/server/prompts/ +/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/server/prompts/ public struct Prompt: Hashable, Codable, Sendable { - /// The prompt name + /// The prompt name (intended for programmatic or logical use) public let name: String + /// A human-readable title for the prompt, intended for UI display. + /// If not provided, the `name` should be used for display. + public let title: String? /// The prompt description public let description: String? /// The prompt arguments public let arguments: [Argument]? - - public init(name: String, description: String? = nil, arguments: [Argument]? = nil) { + /// Reserved for clients and servers to attach additional metadata. + public var _meta: [String: Value]? + /// Optional icons representing this prompt. + public var icons: [Icon]? + + public init( + name: String, + title: String? = nil, + description: String? = nil, + arguments: [Argument]? = nil, + _meta: [String: Value]? = nil, + icons: [Icon]? = nil + ) { self.name = name + self.title = title self.description = description self.arguments = arguments + self._meta = _meta + self.icons = icons } /// An argument for a prompt public struct Argument: Hashable, Codable, Sendable { - /// The argument name + /// The argument name (intended for programmatic or logical use) public let name: String + /// A human-readable title for the argument, intended for UI display. + /// If not provided, the `name` should be used for display. + public let title: String? /// The argument description public let description: String? /// Whether the argument is required public let required: Bool? - public init(name: String, description: String? = nil, required: Bool? = nil) { + public init( + name: String, + title: String? = nil, + description: String? = nil, + required: Bool? = nil + ) { self.name = name + self.title = title self.description = description self.required = required } @@ -40,13 +66,9 @@ public struct Prompt: Hashable, Codable, Sendable { /// A message in a prompt 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 - } + // TODO: Deprecate in a future version + /// Backwards compatibility alias for top-level `Role`. + public typealias Role = MCP.Role /// The message role public let role: Role @@ -78,16 +100,50 @@ public struct Prompt: Hashable, Codable, Sendable { return Message(_role: .assistant, _content: content) } - /// Content types for messages + // TODO: Consider consolidating with Tool.Content into a shared ContentBlock type + // in a future breaking change release. The spec uses a single ContentBlock type. + /// Content types for messages. + /// + /// Matches the MCP spec (2025-11-25) ContentBlock union: + /// - TextContent, ImageContent, AudioContent, ResourceLink, EmbeddedResource public enum Content: Hashable, Sendable { /// Text content - case text(text: String) + case text(text: String, annotations: Annotations?, _meta: [String: Value]?) /// Image content - case image(data: String, mimeType: String) + case image(data: String, mimeType: String, annotations: Annotations?, _meta: [String: Value]?) /// Audio content - case audio(data: String, mimeType: String) - /// Embedded resource content - case resource(uri: String, mimeType: String, text: String?, blob: String?) + case audio(data: String, mimeType: String, annotations: Annotations?, _meta: [String: Value]?) + /// Embedded resource content (includes actual content) + case resource(resource: Resource.Content, annotations: Annotations?, _meta: [String: Value]?) + /// Resource link (reference to a resource that can be read) + case resourceLink(ResourceLink) + + // MARK: - Convenience initializers (backwards compatibility) + + /// Creates text content + public static func text(_ text: String) -> Content { + .text(text: text, annotations: nil, _meta: nil) + } + + /// Creates image content + public static func image(data: String, mimeType: String) -> Content { + .image(data: data, mimeType: mimeType, annotations: nil, _meta: nil) + } + + /// Creates audio content + public static func audio(data: String, mimeType: String) -> Content { + .audio(data: data, mimeType: mimeType, annotations: nil, _meta: nil) + } + + /// Creates embedded resource content with text + public static func resource(uri: String, mimeType: String? = nil, text: String) -> Content { + .resource(resource: .text(text, uri: uri, mimeType: mimeType), annotations: nil, _meta: nil) + } + + /// Creates embedded resource content with binary data + public static func resource(uri: String, mimeType: String? = nil, blob: Data) -> Content { + .resource(resource: .binary(blob, uri: uri, mimeType: mimeType), annotations: nil, _meta: nil) + } } } @@ -95,25 +151,31 @@ public struct Prompt: Hashable, Codable, Sendable { public struct Reference: Hashable, Codable, Sendable { /// The prompt reference name public let name: String + /// A human-readable title for the prompt, intended for UI display. + /// If not provided, the `name` should be used for display. + public let title: String? - public init(name: String) { + public init(name: String, title: String? = nil) { self.name = name + self.title = title } private enum CodingKeys: String, CodingKey { - case type, name + case type, name, title } public func encode(to encoder: Encoder) throws { var container = encoder.container(keyedBy: CodingKeys.self) try container.encode("ref/prompt", forKey: .type) try container.encode(name, forKey: .name) + try container.encodeIfPresent(title, forKey: .title) } public init(from decoder: Decoder) throws { let container = try decoder.container(keyedBy: CodingKeys.self) _ = try container.decode(String.self, forKey: .type) name = try container.decode(String.self, forKey: .name) + title = try container.decodeIfPresent(String.self, forKey: .title) } } } @@ -122,30 +184,37 @@ public struct Prompt: Hashable, Codable, Sendable { extension Prompt.Message.Content: Codable { private enum CodingKeys: String, CodingKey { - case type, text, data, mimeType, uri, blob + case type, text, data, mimeType, resource, annotations, _meta } public func encode(to encoder: Encoder) throws { var container = encoder.container(keyedBy: CodingKeys.self) switch self { - case .text(let text): + case .text(let text, let annotations, let meta): try container.encode("text", forKey: .type) try container.encode(text, forKey: .text) - case .image(let data, let mimeType): + try container.encodeIfPresent(annotations, forKey: .annotations) + try container.encodeIfPresent(meta, forKey: ._meta) + case .image(let data, let mimeType, let annotations, let meta): try container.encode("image", forKey: .type) try container.encode(data, forKey: .data) try container.encode(mimeType, forKey: .mimeType) - case .audio(let data, let mimeType): + try container.encodeIfPresent(annotations, forKey: .annotations) + try container.encodeIfPresent(meta, forKey: ._meta) + case .audio(let data, let mimeType, let annotations, let meta): try container.encode("audio", forKey: .type) try container.encode(data, forKey: .data) try container.encode(mimeType, forKey: .mimeType) - case .resource(let uri, let mimeType, let text, let blob): + try container.encodeIfPresent(annotations, forKey: .annotations) + try container.encodeIfPresent(meta, forKey: ._meta) + case .resource(let resourceContent, let annotations, let meta): try container.encode("resource", forKey: .type) - try container.encode(uri, forKey: .uri) - try container.encode(mimeType, forKey: .mimeType) - try container.encodeIfPresent(text, forKey: .text) - try container.encodeIfPresent(blob, forKey: .blob) + try container.encode(resourceContent, forKey: .resource) + try container.encodeIfPresent(annotations, forKey: .annotations) + try container.encodeIfPresent(meta, forKey: ._meta) + case .resourceLink(let link): + try link.encode(to: encoder) } } @@ -156,21 +225,29 @@ extension Prompt.Message.Content: Codable { switch type { case "text": let text = try container.decode(String.self, forKey: .text) - self = .text(text: text) + let annotations = try container.decodeIfPresent(Annotations.self, forKey: .annotations) + let meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + self = .text(text: text, annotations: annotations, _meta: meta) 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) + let annotations = try container.decodeIfPresent(Annotations.self, forKey: .annotations) + let meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + self = .image(data: data, mimeType: mimeType, annotations: annotations, _meta: meta) case "audio": let data = try container.decode(String.self, forKey: .data) let mimeType = try container.decode(String.self, forKey: .mimeType) - self = .audio(data: data, mimeType: mimeType) + let annotations = try container.decodeIfPresent(Annotations.self, forKey: .annotations) + let meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + self = .audio(data: data, mimeType: mimeType, annotations: annotations, _meta: meta) case "resource": - let uri = try container.decode(String.self, forKey: .uri) - let mimeType = try container.decode(String.self, forKey: .mimeType) - let text = try container.decodeIfPresent(String.self, forKey: .text) - let blob = try container.decodeIfPresent(String.self, forKey: .blob) - self = .resource(uri: uri, mimeType: mimeType, text: text, blob: blob) + let resourceContent = try container.decode(Resource.Content.self, forKey: .resource) + let annotations = try container.decodeIfPresent(Annotations.self, forKey: .annotations) + let meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + self = .resource(resource: resourceContent, annotations: annotations, _meta: meta) + case "resource_link": + let link = try ResourceLink(from: decoder) + self = .resourceLink(link) default: throw DecodingError.dataCorruptedError( forKey: .type, @@ -184,7 +261,7 @@ extension Prompt.Message.Content: Codable { extension Prompt.Message.Content: ExpressibleByStringLiteral { public init(stringLiteral value: String) { - self = .text(text: value) + self = .text(text: value, annotations: nil, _meta: nil) } } @@ -192,7 +269,7 @@ extension Prompt.Message.Content: ExpressibleByStringLiteral { extension Prompt.Message.Content: ExpressibleByStringInterpolation { public init(stringInterpolation: DefaultStringInterpolation) { - self = .text(text: String(stringInterpolation: stringInterpolation)) + self = .text(text: String(stringInterpolation: stringInterpolation), annotations: nil, _meta: nil) } } @@ -205,23 +282,60 @@ public enum ListPrompts: Method { public struct Parameters: NotRequired, Hashable, Codable, Sendable { public let cursor: String? + /// Request metadata including progress token. + public var _meta: RequestMeta? public init() { self.cursor = nil + self._meta = nil } - public init(cursor: String) { + public init(cursor: String? = nil, _meta: RequestMeta? = nil) { self.cursor = cursor + self._meta = _meta } } - public struct Result: Hashable, Codable, Sendable { + public struct Result: ResultWithExtraFields { + public typealias ResultCodingKeys = CodingKeys + public let prompts: [Prompt] public let nextCursor: String? - - public init(prompts: [Prompt], nextCursor: String? = nil) { + /// Reserved for clients and servers to attach additional metadata. + public var _meta: [String: Value]? + /// Additional fields not defined in the schema (for forward compatibility). + public var extraFields: [String: Value]? + + public init( + prompts: [Prompt], + nextCursor: String? = nil, + _meta: [String: Value]? = nil, + extraFields: [String: Value]? = nil + ) { self.prompts = prompts self.nextCursor = nextCursor + self._meta = _meta + self.extraFields = extraFields + } + + public enum CodingKeys: String, CodingKey, CaseIterable { + case prompts, nextCursor, _meta + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + prompts = try container.decode([Prompt].self, forKey: .prompts) + nextCursor = try container.decodeIfPresent(String.self, forKey: .nextCursor) + _meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + extraFields = try Self.decodeExtraFields(from: decoder) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode(prompts, forKey: .prompts) + try container.encodeIfPresent(nextCursor, forKey: .nextCursor) + try container.encodeIfPresent(_meta, forKey: ._meta) + try encodeExtraFields(to: encoder) } } } @@ -234,21 +348,59 @@ public enum GetPrompt: Method { public struct Parameters: Hashable, Codable, Sendable { public let name: String - public let arguments: [String: Value]? + /// Arguments to use for templating the prompt. + /// Per the MCP spec, argument values must be strings. + public let arguments: [String: String]? + /// Request metadata including progress token. + public var _meta: RequestMeta? - public init(name: String, arguments: [String: Value]? = nil) { + public init(name: String, arguments: [String: String]? = nil, _meta: RequestMeta? = nil) { self.name = name self.arguments = arguments + self._meta = _meta } } - public struct Result: Hashable, Codable, Sendable { + public struct Result: ResultWithExtraFields { + public typealias ResultCodingKeys = CodingKeys + public let description: String? public let messages: [Prompt.Message] - - public init(description: String?, messages: [Prompt.Message]) { + /// Reserved for clients and servers to attach additional metadata. + public var _meta: [String: Value]? + /// Additional fields not defined in the schema (for forward compatibility). + public var extraFields: [String: Value]? + + public init( + description: String?, + messages: [Prompt.Message], + _meta: [String: Value]? = nil, + extraFields: [String: Value]? = nil + ) { self.description = description self.messages = messages + self._meta = _meta + self.extraFields = extraFields + } + + public enum CodingKeys: String, CodingKey, CaseIterable { + case description, messages, _meta + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + description = try container.decodeIfPresent(String.self, forKey: .description) + messages = try container.decode([Prompt.Message].self, forKey: .messages) + _meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + extraFields = try Self.decodeExtraFields(from: decoder) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encodeIfPresent(description, forKey: .description) + try container.encode(messages, forKey: .messages) + try container.encodeIfPresent(_meta, forKey: ._meta) + try encodeExtraFields(to: encoder) } } } @@ -257,4 +409,6 @@ public enum GetPrompt: Method { /// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2024-11-05/server/prompts/#list-changed-notification public struct PromptListChangedNotification: Notification { public static let name: String = "notifications/prompts/list_changed" + + public typealias Parameters = NotificationParams } diff --git a/Sources/MCP/Server/Resources.swift b/Sources/MCP/Server/Resources.swift index 12f67335..47992d39 100644 --- a/Sources/MCP/Server/Resources.swift +++ b/Sources/MCP/Server/Resources.swift @@ -6,35 +6,52 @@ import Foundation /// such as files, database schemas, or application-specific information. /// Each resource is uniquely identified by a URI. /// -/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2024-11-05/server/resources/ +/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/server/resources/ public struct Resource: Hashable, Codable, Sendable { - /// The resource name + /// The resource name (intended for programmatic or logical use) public var name: String + /// A human-readable title for the resource, intended for UI display. + /// If not provided, the `name` should be used for display. + public var title: String? /// The resource URI public var uri: String /// The resource description public var description: String? /// The resource MIME type public var mimeType: String? - /// The resource metadata - public var metadata: [String: String]? + /// The size of the raw resource content, in bytes, if known. + public var size: Int? + /// Optional annotations for the client. + public var annotations: Annotations? + /// Reserved for clients and servers to attach additional metadata. + public var _meta: [String: Value]? + /// Optional icons representing this resource. + public var icons: [Icon]? public init( name: String, + title: String? = nil, uri: String, description: String? = nil, mimeType: String? = nil, - metadata: [String: String]? = nil + size: Int? = nil, + annotations: Annotations? = nil, + _meta: [String: Value]? = nil, + icons: [Icon]? = nil ) { self.name = name + self.title = title self.uri = uri self.description = description self.mimeType = mimeType - self.metadata = metadata + self.size = size + self.annotations = annotations + self._meta = _meta + self.icons = icons } /// Content of a resource. - public struct Content: Hashable, Codable, Sendable { + public struct Contents: Hashable, Codable, Sendable { /// The resource URI public let uri: String /// The resource MIME type @@ -43,6 +60,8 @@ public struct Resource: Hashable, Codable, Sendable { public let text: String? /// The resource binary content public let blob: String? + /// Reserved for clients and servers to attach additional metadata. + public var _meta: [String: Value]? public static func text(_ content: String, uri: String, mimeType: String? = nil) -> Self { .init(uri: uri, mimeType: mimeType, text: content) @@ -57,6 +76,7 @@ public struct Resource: Hashable, Codable, Sendable { self.mimeType = mimeType self.text = text self.blob = nil + self._meta = nil } private init(uri: String, mimeType: String? = nil, blob: String) { @@ -64,34 +84,184 @@ public struct Resource: Hashable, Codable, Sendable { self.mimeType = mimeType self.text = nil self.blob = blob + self._meta = nil } } - /// A resource template. + // TODO: Deprecate in a future version + /// Backwards compatibility alias for `Contents`. + public typealias Content = Contents + + /// A resource template that can generate multiple resources via URI pattern matching. + /// + /// Resource templates use [RFC 6570 URI Templates](https://datatracker.ietf.org/doc/html/rfc6570) + /// to define patterns for dynamic resource URIs. Clients can use these templates to construct + /// resource URIs by substituting template variables. + /// + /// ## Example + /// + /// ```swift + /// // Define a template for user profiles + /// let template = Resource.Template( + /// uriTemplate: "users://{userId}/profile", + /// name: "user_profile", + /// title: "User Profile", + /// description: "Profile information for a specific user", + /// mimeType: "application/json" + /// ) + /// + /// // Register with a server + /// server.registerResources { + /// listTemplates: { _ in [template] }, + /// read: { uri in + /// // Parse userId from URI and return profile data + /// let userId = parseUserId(from: uri) + /// return [.text(getProfile(userId), uri: uri)] + /// } + /// } + /// ``` + /// + /// - SeeAlso: https://spec.modelcontextprotocol.io/specification/server/resources/#resource-templates public struct Template: Hashable, Codable, Sendable { - /// The URI template pattern + /// The URI template pattern (RFC 6570 format, e.g., "file:///{path}"). public var uriTemplate: String - /// The template name + /// The template name (intended for programmatic or logical use). public var name: String - /// The template description + /// A human-readable title for the template, intended for UI display. + /// If not provided, the `name` should be used for display. + public var title: String? + /// A description of what resources this template provides. public var description: String? - /// The resource MIME type + /// The MIME type of resources generated from this template. public var mimeType: String? + /// Optional annotations for the client. + public var annotations: Annotations? + /// Reserved for clients and servers to attach additional metadata. + public var _meta: [String: Value]? + /// Optional icons representing this resource template. + public var icons: [Icon]? public init( uriTemplate: String, name: String, + title: String? = nil, description: String? = nil, - mimeType: String? = nil + mimeType: String? = nil, + annotations: Annotations? = nil, + _meta: [String: Value]? = nil, + icons: [Icon]? = nil ) { self.uriTemplate = uriTemplate self.name = name + self.title = title self.description = description self.mimeType = mimeType + self.annotations = annotations + self._meta = _meta + self.icons = icons } } } +/// A resource link returned in tool results, referencing a resource that can be read. +/// +/// Resource links differ from embedded resources in that they don't include +/// the actual content - they're references to resources that can be read later. +/// +/// Note: Resource links returned by tools are not guaranteed to appear +/// in the results of `resources/list` requests. +/// +/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/server/tools/#resource-links +public struct ResourceLink: Hashable, Codable, Sendable { + /// The resource name (intended for programmatic or logical use) + public var name: String + /// A human-readable title for the resource, intended for UI display. + public var title: String? + /// The resource URI + public var uri: String + /// The resource description + public var description: String? + /// The resource MIME type + public var mimeType: String? + /// The size of the raw resource content, in bytes, if known. + public var size: Int? + /// Optional annotations for the client. + public var annotations: Annotations? + /// Optional icons representing this resource. + public var icons: [Icon]? + /// Reserved for clients and servers to attach additional metadata. + public var _meta: [String: Value]? + + public init( + name: String, + title: String? = nil, + uri: String, + description: String? = nil, + mimeType: String? = nil, + size: Int? = nil, + annotations: Annotations? = nil, + icons: [Icon]? = nil, + _meta: [String: Value]? = nil + ) { + self.name = name + self.title = title + self.uri = uri + self.description = description + self.mimeType = mimeType + self.size = size + self.annotations = annotations + self.icons = icons + self._meta = _meta + } + + private enum CodingKeys: String, CodingKey { + case type + case name + case title + case uri + case description + case mimeType + case size + case annotations + case icons + case _meta + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + // Verify type is "resource_link" + let type = try container.decodeIfPresent(String.self, forKey: .type) + if let type, type != "resource_link" { + throw DecodingError.dataCorruptedError( + forKey: .type, in: container, + debugDescription: "Expected type 'resource_link', got '\(type)'") + } + name = try container.decode(String.self, forKey: .name) + title = try container.decodeIfPresent(String.self, forKey: .title) + uri = try container.decode(String.self, forKey: .uri) + description = try container.decodeIfPresent(String.self, forKey: .description) + mimeType = try container.decodeIfPresent(String.self, forKey: .mimeType) + size = try container.decodeIfPresent(Int.self, forKey: .size) + annotations = try container.decodeIfPresent(Annotations.self, forKey: .annotations) + icons = try container.decodeIfPresent([Icon].self, forKey: .icons) + _meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode("resource_link", forKey: .type) + try container.encode(name, forKey: .name) + try container.encodeIfPresent(title, forKey: .title) + try container.encode(uri, forKey: .uri) + try container.encodeIfPresent(description, forKey: .description) + try container.encodeIfPresent(mimeType, forKey: .mimeType) + try container.encodeIfPresent(size, forKey: .size) + try container.encodeIfPresent(annotations, forKey: .annotations) + try container.encodeIfPresent(icons, forKey: .icons) + try container.encodeIfPresent(_meta, forKey: ._meta) + } +} + // MARK: - /// To discover available resources, clients send a `resources/list` request. @@ -101,23 +271,60 @@ public enum ListResources: Method { public struct Parameters: NotRequired, Hashable, Codable, Sendable { public let cursor: String? + /// Request metadata including progress token. + public var _meta: RequestMeta? public init() { self.cursor = nil + self._meta = nil } - - public init(cursor: String) { + + public init(cursor: String? = nil, _meta: RequestMeta? = nil) { self.cursor = cursor + self._meta = _meta } } - public struct Result: Hashable, Codable, Sendable { + public struct Result: ResultWithExtraFields { + public typealias ResultCodingKeys = CodingKeys + public let resources: [Resource] public let nextCursor: String? + /// Reserved for clients and servers to attach additional metadata. + public var _meta: [String: Value]? + /// Additional fields not defined in the schema (for forward compatibility). + public var extraFields: [String: Value]? - public init(resources: [Resource], nextCursor: String? = nil) { + public init( + resources: [Resource], + nextCursor: String? = nil, + _meta: [String: Value]? = nil, + extraFields: [String: Value]? = nil + ) { self.resources = resources self.nextCursor = nextCursor + self._meta = _meta + self.extraFields = extraFields + } + + public enum CodingKeys: String, CodingKey, CaseIterable { + case resources, nextCursor, _meta + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + resources = try container.decode([Resource].self, forKey: .resources) + nextCursor = try container.decodeIfPresent(String.self, forKey: .nextCursor) + _meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + extraFields = try Self.decodeExtraFields(from: decoder) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode(resources, forKey: .resources) + try container.encodeIfPresent(nextCursor, forKey: .nextCursor) + try container.encodeIfPresent(_meta, forKey: ._meta) + try encodeExtraFields(to: encoder) } } } @@ -129,17 +336,50 @@ public enum ReadResource: Method { public struct Parameters: Hashable, Codable, Sendable { public let uri: String + /// Request metadata including progress token. + public var _meta: RequestMeta? - public init(uri: String) { + public init(uri: String, _meta: RequestMeta? = nil) { self.uri = uri + self._meta = _meta } } - public struct Result: Hashable, Codable, Sendable { + public struct Result: ResultWithExtraFields { + public typealias ResultCodingKeys = CodingKeys + public let contents: [Resource.Content] + /// Reserved for clients and servers to attach additional metadata. + public var _meta: [String: Value]? + /// Additional fields not defined in the schema (for forward compatibility). + public var extraFields: [String: Value]? - public init(contents: [Resource.Content]) { + public init( + contents: [Resource.Content], + _meta: [String: Value]? = nil, + extraFields: [String: Value]? = nil + ) { self.contents = contents + self._meta = _meta + self.extraFields = extraFields + } + + public enum CodingKeys: String, CodingKey, CaseIterable { + case contents, _meta + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + contents = try container.decode([Resource.Content].self, forKey: .contents) + _meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + extraFields = try Self.decodeExtraFields(from: decoder) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode(contents, forKey: .contents) + try container.encodeIfPresent(_meta, forKey: ._meta) + try encodeExtraFields(to: encoder) } } } @@ -151,28 +391,62 @@ public enum ListResourceTemplates: Method { public struct Parameters: NotRequired, Hashable, Codable, Sendable { public let cursor: String? + /// Request metadata including progress token. + public var _meta: RequestMeta? public init() { self.cursor = nil + self._meta = nil } - - public init(cursor: String) { + + public init(cursor: String? = nil, _meta: RequestMeta? = nil) { self.cursor = cursor + self._meta = _meta } } - public struct Result: Hashable, Codable, Sendable { + public struct Result: ResultWithExtraFields { + public typealias ResultCodingKeys = CodingKeys + public let templates: [Resource.Template] public let nextCursor: String? + /// Reserved for clients and servers to attach additional metadata. + public var _meta: [String: Value]? + /// Additional fields not defined in the schema (for forward compatibility). + public var extraFields: [String: Value]? - public init(templates: [Resource.Template], nextCursor: String? = nil) { + public init( + templates: [Resource.Template], + nextCursor: String? = nil, + _meta: [String: Value]? = nil, + extraFields: [String: Value]? = nil + ) { self.templates = templates self.nextCursor = nextCursor + self._meta = _meta + self.extraFields = extraFields } - private enum CodingKeys: String, CodingKey { + public enum CodingKeys: String, CodingKey, CaseIterable { case templates = "resourceTemplates" case nextCursor + case _meta + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + templates = try container.decode([Resource.Template].self, forKey: .templates) + nextCursor = try container.decodeIfPresent(String.self, forKey: .nextCursor) + _meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + extraFields = try Self.decodeExtraFields(from: decoder) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode(templates, forKey: .templates) + try container.encodeIfPresent(nextCursor, forKey: .nextCursor) + try container.encodeIfPresent(_meta, forKey: ._meta) + try encodeExtraFields(to: encoder) } } } @@ -182,16 +456,42 @@ public enum ListResourceTemplates: Method { public struct ResourceListChangedNotification: Notification { public static let name: String = "notifications/resources/list_changed" - public typealias Parameters = Empty + public typealias Parameters = NotificationParams } /// Clients can subscribe to specific resources and receive notifications when they change. -/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2024-11-05/server/resources/#subscriptions +/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2025-11-25/server/resources/#subscriptions public enum ResourceSubscribe: Method { public static let name: String = "resources/subscribe" public struct Parameters: Hashable, Codable, Sendable { public let uri: String + /// Request metadata including progress token. + public var _meta: RequestMeta? + + public init(uri: String, _meta: RequestMeta? = nil) { + self.uri = uri + self._meta = _meta + } + } + + public typealias Result = Empty +} + +/// Clients can unsubscribe from resources to stop receiving update notifications. +/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2025-11-25/server/resources/#subscriptions +public enum ResourceUnsubscribe: Method { + public static let name: String = "resources/unsubscribe" + + public struct Parameters: Hashable, Codable, Sendable { + public let uri: String + /// Request metadata including progress token. + public var _meta: RequestMeta? + + public init(uri: String, _meta: RequestMeta? = nil) { + self.uri = uri + self._meta = _meta + } } public typealias Result = Empty @@ -204,9 +504,12 @@ public struct ResourceUpdatedNotification: Notification { public struct Parameters: Hashable, Codable, Sendable { public let uri: String + /// Reserved for additional metadata. + public var _meta: [String: Value]? - public init(uri: String) { + public init(uri: String, _meta: [String: Value]? = nil) { self.uri = uri + self._meta = _meta } } } diff --git a/Sources/MCP/Server/Server.swift b/Sources/MCP/Server/Server.swift index f2c975af..3e3e01f0 100644 --- a/Sources/MCP/Server/Server.swift +++ b/Sources/MCP/Server/Server.swift @@ -32,10 +32,30 @@ public actor Server { public let name: String /// The server version public let version: String + /// A human-readable title for the server, intended for UI display. + /// If not provided, the `name` should be used for display. + public let title: String? + /// An optional human-readable description of what this implementation does. + public let description: String? + /// Optional icons representing this implementation. + public let icons: [Icon]? + /// An optional URL of the website for this implementation. + public let websiteUrl: String? - public init(name: String, version: String) { + public init( + name: String, + version: String, + title: String? = nil, + description: String? = nil, + icons: [Icon]? = nil, + websiteUrl: String? = nil + ) { self.name = name self.version = version + self.title = title + self.description = description + self.icons = icons + self.websiteUrl = websiteUrl } } @@ -82,8 +102,8 @@ public actor Server { public init() {} } - /// Sampling capabilities - public struct Sampling: Hashable, Codable, Sendable { + /// Completions capabilities + public struct Completions: Hashable, Codable, Sendable { public init() {} } @@ -93,23 +113,319 @@ public actor Server { public var prompts: Prompts? /// Resources capabilities public var resources: Resources? - /// Sampling capabilities - public var sampling: Sampling? /// Tools capabilities public var tools: Tools? + /// Completions capabilities + public var completions: Completions? + /// Tasks capabilities (experimental) + public var tasks: Tasks? + /// Experimental, non-standard capabilities that the server supports. + public var experimental: [String: [String: Value]]? public init( logging: Logging? = nil, prompts: Prompts? = nil, resources: Resources? = nil, - sampling: Sampling? = nil, - tools: Tools? = nil + tools: Tools? = nil, + completions: Completions? = nil, + tasks: Tasks? = nil, + experimental: [String: [String: Value]]? = nil ) { self.logging = logging self.prompts = prompts self.resources = resources - self.sampling = sampling self.tools = tools + self.completions = completions + self.tasks = tasks + self.experimental = experimental + } + } + + /// Context provided to request handlers for sending notifications during execution. + /// + /// When a request handler needs to send notifications (e.g., progress updates during + /// a long-running tool), it should use this context to ensure the notification is + /// routed to the correct client, even if other clients have connected in the meantime. + /// + /// Example: + /// ```swift + /// server.withRequestHandler(CallTool.self) { params, context in + /// // Send progress notification using convenience method + /// try await context.sendProgress( + /// token: progressToken, + /// progress: 50.0, + /// total: 100.0, + /// message: "Processing..." + /// ) + /// // ... do work ... + /// return result + /// } + /// ``` + public struct RequestHandlerContext: Sendable { + /// Send a notification without parameters to the client that initiated this request. + /// + /// The notification will be routed to the correct client even if other clients + /// have connected since the request was received. + /// + /// - Parameter notification: The notification to send (for notifications without parameters) + public let sendNotification: @Sendable (any Notification) async throws -> Void + + /// Send a notification message with parameters to the client that initiated this request. + /// + /// Use this method to send notifications that have parameters, such as `ProgressNotification` + /// or `LogMessageNotification`. + /// + /// Example: + /// ```swift + /// try await context.sendMessage(ProgressNotification.message(.init( + /// progressToken: token, + /// progress: 50.0, + /// total: 100.0, + /// message: "Halfway done" + /// ))) + /// ``` + /// + /// - Parameter message: The notification message to send + public let sendMessage: @Sendable (any NotificationMessageProtocol) async throws -> Void + + /// Send raw data to the client that initiated this request. + /// + /// This is used internally for sending queued task messages (such as elicitation + /// or sampling requests that were queued during task execution). + /// + /// - Important: This is an internal API primarily used by the task system. + /// + /// - Parameter data: The raw JSON data to send + public let sendData: @Sendable (Data) async throws -> Void + + /// The session identifier for the client that initiated this request. + /// + /// For HTTP transports with multiple concurrent clients, each client session + /// has a unique identifier. This can be used for per-session features like + /// independent log levels. + /// + /// For simple transports (stdio, single-connection), this is `nil`. + public let sessionId: String? + + /// Check if a log message at the given level should be sent. + /// + /// This respects the minimum log level set by the client via `logging/setLevel`. + /// Messages below the threshold will be silently dropped. + let shouldSendLogMessage: @Sendable (LoggingLevel) async -> Bool + + // MARK: - Convenience Methods + + /// Send a progress notification to the client. + /// + /// Use this to report progress on long-running operations. + /// + /// - Parameters: + /// - token: The progress token from the request's `_meta.progressToken` + /// - progress: The current progress value (should increase monotonically) + /// - total: The total progress value, if known + /// - message: An optional human-readable message describing current progress + public func sendProgress( + token: ProgressToken, + progress: Double, + total: Double? = nil, + message: String? = nil + ) async throws { + try await sendMessage(ProgressNotification.message(.init( + progressToken: token, + progress: progress, + total: total, + message: message + ))) + } + + /// Send a log message notification to the client. + /// + /// The message will only be sent if its level is at or above the minimum + /// log level set by the client via `logging/setLevel`. Messages below the + /// threshold are silently dropped. + /// + /// - Parameters: + /// - level: The severity level of the log message + /// - logger: An optional name for the logger producing the message + /// - data: The log message data (can be a string or structured data) + public func sendLogMessage( + level: LoggingLevel, + logger: String? = nil, + data: Value + ) async throws { + // Check if this message should be sent based on the current log level + guard await shouldSendLogMessage(level) else { return } + + try await sendMessage(LogMessageNotification.message(.init( + level: level, + logger: logger, + data: data + ))) + } + + /// Send a resource list changed notification to the client. + /// + /// Call this when the list of available resources has changed. + public func sendResourceListChanged() async throws { + try await sendNotification(ResourceListChangedNotification()) + } + + /// Send a resource updated notification to the client. + /// + /// Call this when a specific resource's content has been updated. + /// + /// - Parameter uri: The URI of the resource that was updated + public func sendResourceUpdated(uri: String) async throws { + try await sendMessage(ResourceUpdatedNotification.message(.init(uri: uri))) + } + + /// Send a tool list changed notification to the client. + /// + /// Call this when the list of available tools has changed. + public func sendToolListChanged() async throws { + try await sendNotification(ToolListChangedNotification()) + } + + /// Send a prompt list changed notification to the client. + /// + /// Call this when the list of available prompts has changed. + public func sendPromptListChanged() async throws { + try await sendNotification(PromptListChangedNotification()) + } + + /// Send a cancellation notification to the client. + /// + /// - Parameters: + /// - requestId: The ID of the request being cancelled (optional in protocol 2025-11-25+) + /// - reason: An optional reason for the cancellation + public func sendCancelled(requestId: RequestId? = nil, reason: String? = nil) async throws { + try await sendMessage(CancelledNotification.message(.init( + requestId: requestId, + reason: reason + ))) + } + + /// Send an elicitation complete notification to the client. + /// + /// This notifies the client that an out-of-band (URL mode) elicitation + /// request has been completed. + /// + /// - Parameter elicitationId: The ID of the elicitation that completed. + public func sendElicitationComplete(elicitationId: String) async throws { + try await sendMessage(ElicitationCompleteNotification.message(.init( + elicitationId: elicitationId + ))) + } + + /// Send a task status notification to the client. + /// + /// This notifies the client of a change in task status. + /// + /// - Parameter task: The task to send the status notification for. + public func sendTaskStatus(task: MCPTask) async throws { + try await sendMessage(TaskStatusNotification.message(.init(task: task))) + } + + // MARK: - Cancellation Checking + + /// Whether the request has been cancelled. + /// + /// Check this property periodically during long-running operations + /// to respond to cancellation requests from the client. + /// + /// This returns `true` when: + /// - The client sends a `CancelledNotification` for this request + /// - The server is shutting down + /// + /// When cancelled, the handler should clean up resources and return + /// or throw an error. Per MCP spec, responses are not sent for cancelled requests. + /// + /// ## Example + /// + /// ```swift + /// server.withRequestHandler(CallTool.self) { params, context in + /// for item in largeDataset { + /// // Check cancellation periodically + /// guard !context.isCancelled else { + /// throw CancellationError() + /// } + /// try await process(item) + /// } + /// return CallTool.Result(content: [.text("Done")]) + /// } + /// ``` + public var isCancelled: Bool { + Task.isCancelled + } + + /// Check if the request has been cancelled and throw if so. + /// + /// Call this method periodically during long-running operations. + /// If the request has been cancelled, this throws `CancellationError`. + /// + /// This is equivalent to checking `isCancelled` and throwing manually, + /// but provides a more idiomatic Swift concurrency pattern. + /// + /// ## Example + /// + /// ```swift + /// server.withRequestHandler(CallTool.self) { params, context in + /// for item in largeDataset { + /// try context.checkCancellation() // Throws if cancelled + /// try await process(item) + /// } + /// return CallTool.Result(content: [.text("Done")]) + /// } + /// ``` + /// + /// - Throws: `CancellationError` if the request has been cancelled. + public func checkCancellation() throws { + try Task.checkCancellation() + } + } + + /// A type-erased pending request for server→client requests (bidirectional communication). + private struct AnyServerPendingRequest { + private let _yield: (Result) -> Void + private let _finish: () -> Void + + init( + continuation: AsyncThrowingStream.Continuation + ) { + _yield = { result in + switch result { + case .success(let value): + if let typedValue = value as? T { + continuation.yield(typedValue) + continuation.finish() + } else if let value = value as? Value, + let data = try? JSONEncoder().encode(value), + let decoded = try? JSONDecoder().decode(T.self, from: data) + { + continuation.yield(decoded) + continuation.finish() + } else { + continuation.finish(throwing: MCPError.internalError("Type mismatch in response")) + } + case .failure(let error): + continuation.finish(throwing: error) + } + } + _finish = { + continuation.finish() + } + } + + func resume(returning value: Any) { + _yield(.success(value)) + } + + func resume(throwing error: Swift.Error) { + _yield(.failure(error)) + } + + func finish() { + _finish() } } @@ -139,13 +455,36 @@ public actor Server { public var capabilities: Capabilities /// The server configuration public var configuration: Configuration - + + /// Experimental APIs for tasks and other features. + /// + /// Access experimental features via this property: + /// ```swift + /// // Enable task support with in-memory storage + /// await server.experimental.tasks.enable() + /// + /// // Or with custom configuration + /// let taskSupport = TaskSupport.inMemory() + /// await server.experimental.tasks.enable(taskSupport) + /// ``` + /// + /// - Warning: These APIs are experimental and may change without notice. + public var experimental: ExperimentalServerFeatures { + ExperimentalServerFeatures(server: self) + } /// Request handlers private var methodHandlers: [String: RequestHandlerBox] = [:] /// Notification handlers private var notificationHandlers: [String: [NotificationHandlerBox]] = [:] + /// Pending requests sent from server to client (for bidirectional communication) + private var pendingRequests: [RequestId: AnyServerPendingRequest] = [:] + /// Counter for generating unique request IDs + private var nextRequestId = 0 + /// Response routers for intercepting responses before normal handling + private var responseRouters: [any ResponseRouter] = [] + /// Whether the server is initialized private var isInitialized = false /// The client information @@ -155,9 +494,21 @@ public actor Server { /// The protocol version private var protocolVersion: String? /// The list of subscriptions - private var subscriptions: [String: Set] = [:] + private var subscriptions: [String: Set] = [:] /// The task for the message handling loop private var task: Task? + /// Per-session minimum log levels set by clients. + /// + /// For HTTP transports with multiple concurrent clients, each session can + /// independently set its own log level. The key is the session ID (`nil` for + /// transports without session support like stdio). + /// + /// Log messages below a session's level will be filtered out for that session. + private var loggingLevels: [String?: LoggingLevel] = [:] + + /// In-flight request handler Tasks, tracked by request ID. + /// Used for protocol-level cancellation when CancelledNotification is received. + private var inFlightHandlerTasks: [RequestId: Task] = [:] public init( name: String, @@ -194,14 +545,52 @@ public actor Server { for try await data in stream { if Task.isCancelled { break } // Check cancellation inside loop - var requestID: ID? + var requestID: RequestId? do { - // Attempt to decode as batch first, then as individual request or notification + // Attempt to decode as batch first, then as individual request, response, or notification let decoder = JSONDecoder() if let batch = try? decoder.decode(Server.Batch.self, from: data) { - try await handleBatch(batch) + // Spawn batch handler in a separate task for the same reason + // as individual requests - to support nested server-to-client + // requests within batch item handlers. + Task { [weak self] in + guard let self else { return } + do { + try await self.handleBatch(batch) + } catch { + await self.logger?.error( + "Error handling batch", + metadata: ["error": "\(error)"] + ) + } + } + } else if let response = try? decoder.decode(AnyResponse.self, from: data) { + // Handle response from client (for server→client requests) + await handleClientResponse(response) } else if let request = try? decoder.decode(AnyRequest.self, from: data) { - _ = try await handleRequest(request, sendResponse: true) + // Spawn request handler in a separate task to avoid blocking + // the message loop. This allows nested server-to-client requests + // (like elicitation or sampling) to work correctly - the handler + // can await a response while the message loop continues processing + // incoming messages including that response. + let requestId = request.id + let handlerTask = Task { [weak self] in + guard let self else { return } + defer { + Task { await self.removeInFlightRequest(requestId) } + } + do { + _ = try await self.handleRequest(request, sendResponse: true) + } catch { + // handleRequest already sends error responses, so this + // only catches errors from send() itself + await self.logger?.error( + "Error sending response", + metadata: ["error": "\(error)", "requestId": "\(request.id)"] + ) + } + } + trackInFlightRequest(requestId, task: handlerTask) } else if let message = try? decoder.decode(AnyMessage.self, from: data) { try await handleMessage(message) } else { @@ -218,11 +607,9 @@ public actor Server { } throw MCPError.parseError("Invalid message format") } - } catch let error where MCPError.isResourceTemporarilyUnavailable(error) { - // Resource temporarily unavailable, retry after a short delay - try? await Task.sleep(for: .milliseconds(10)) - continue } catch { + // Note: EAGAIN handling is not needed here - the transport layer + // handles it internally. Message handling code won't throw EAGAIN. await logger?.error( "Error processing message", metadata: ["error": "\(error)"]) let response = AnyMethod.response( @@ -243,6 +630,16 @@ public actor Server { /// Stop the server public func stop() async { + // Cancel all in-flight request handlers + for (requestId, handlerTask) in inFlightHandlerTasks { + handlerTask.cancel() + await logger?.debug( + "Cancelled in-flight request during shutdown", + metadata: ["id": "\(requestId)"] + ) + } + inFlightHandlerTasks.removeAll() + task?.cancel() task = nil if let connection = connection { @@ -257,19 +654,66 @@ public actor Server { // MARK: - Registration - /// Register a method handler + /// Register a method handler with access to request context. + /// + /// The context provides capabilities like sending notifications during request + /// processing, with correct routing to the requesting client. + /// + /// - Parameters: + /// - type: The method type to handle + /// - handler: The handler function receiving parameters and context + /// - Returns: Self for chaining @discardableResult - public func withMethodHandler( + public func withRequestHandler( _ type: M.Type, - handler: @escaping @Sendable (M.Parameters) async throws -> M.Result + handler: @escaping @Sendable (M.Parameters, RequestHandlerContext) async throws -> M.Result ) -> Self { - methodHandlers[M.name] = TypedRequestHandler { (request: Request) -> Response in - let result = try await handler(request.params) + methodHandlers[M.name] = TypedRequestHandler { (request: Request, context: RequestHandlerContext) -> Response in + let result = try await handler(request.params, context) return Response(id: request.id, result: result) } return self } + /// Register a method handler without context. + /// + /// - Parameters: + /// - type: The method type to handle + /// - handler: The handler function receiving only parameters + /// - Returns: Self for chaining + @available(*, deprecated, message: "Use withRequestHandler(_:handler:) with RequestHandlerContext for correct notification routing") + @discardableResult + public func withRequestHandler( + _ type: M.Type, + handler: @escaping @Sendable (M.Parameters) async throws -> M.Result + ) -> Self { + withRequestHandler(type) { params, _ in + try await handler(params) + } + } + + // MARK: - Deprecated Method Handler Registration + + /// Register a request handler for a method (deprecated, use withRequestHandler instead) + @available(*, deprecated, renamed: "withRequestHandler") + @discardableResult + public func withMethodHandler( + _ type: M.Type, + handler: @escaping @Sendable (M.Parameters, RequestHandlerContext) async throws -> M.Result + ) -> Self { + withRequestHandler(type, handler: handler) + } + + /// Register a request handler for a method (deprecated, use withRequestHandler instead) + @available(*, deprecated, renamed: "withRequestHandler") + @discardableResult + public func withMethodHandler( + _ type: M.Type, + handler: @escaping @Sendable (M.Parameters) async throws -> M.Result + ) -> Self { + withRequestHandler(type, handler: handler) + } + /// Register a notification handler @discardableResult public func onNotification( @@ -280,6 +724,22 @@ public actor Server { return self } + /// Register a response router to intercept responses before normal handling. + /// + /// Response routers are checked in order before falling back to the default + /// pending request handling. This is used by TaskResultHandler to route + /// responses for queued task requests back to their resolvers. + /// + /// - Important: This is an experimental API that may change without notice. + /// + /// - Parameter router: The response router to add + /// - Returns: Self for chaining + @discardableResult + public func addResponseRouter(_ router: any ResponseRouter) -> Self { + responseRouters.append(router) + return self + } + // MARK: - Sending /// Send a response to a request @@ -308,67 +768,39 @@ public actor Server { try await connection.send(notificationData) } - // MARK: - Sampling - - /// Request sampling from the connected client + /// Send a log message notification to connected clients. /// - /// Sampling allows servers to request LLM completions through the client, - /// enabling sophisticated agentic behaviors while maintaining human-in-the-loop control. + /// This method can be called outside of request handlers to send log messages + /// asynchronously. The message will only be sent if: + /// - The server has declared the `logging` capability + /// - The message's level is at or above the minimum level set by the session /// - /// 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 + /// If the logging capability is not declared, this method silently returns without + /// sending (matching TypeScript SDK behavior). /// /// - 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") - } + /// - level: The severity level of the log message + /// - logger: An optional name for the logger producing the message + /// - data: The log message data (can be a string or structured data) + /// - sessionId: Optional session ID for per-session log level filtering. + /// If `nil`, the log level for the nil-session (default) is used. + public func sendLogMessage( + level: LoggingLevel, + logger: String? = nil, + data: Value, + sessionId: String? = nil + ) async throws { + // Check if logging capability is declared (matching TypeScript SDK behavior) + guard capabilities.logging != nil else { return } - // 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 - ) - ) + // Check if this message should be sent based on the session's log level + guard shouldSendLogMessage(at: level, forSession: sessionId) else { return } - // 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") + try await notify(LogMessageNotification.message(.init( + level: level, + logger: logger, + data: data + ))) } /// A JSON-RPC batch containing multiple requests and/or notifications @@ -389,13 +821,23 @@ public actor Server { /// Process a batch of requests and/or notifications private func handleBatch(_ batch: Batch) async throws { + // Capture the connection at batch start. + // This ensures all batch responses go to the correct client. + let capturedConnection = self.connection + await logger?.trace("Processing batch request", metadata: ["size": "\(batch.items.count)"]) if batch.items.isEmpty { // Empty batch is invalid according to JSON-RPC spec let error = MCPError.invalidRequest("Batch array must not be empty") let response = AnyMethod.response(id: .random, error: error) - try await send(response) + // Use captured connection for error response + if let connection = capturedConnection { + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + let responseData = try encoder.encode(response) + try await connection.send(responseData) + } return } @@ -425,22 +867,87 @@ public actor Server { } } - // Send collected responses if any + // Send collected responses if any (using captured connection) if !responses.isEmpty { + guard let connection = capturedConnection else { + await logger?.warning("Cannot send batch response - connection was nil at batch start") + return + } + let encoder = JSONEncoder() encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] let responseData = try encoder.encode(responses) - guard let connection = connection else { - throw MCPError.internalError("Server connection not initialized") - } - try await connection.send(responseData) } } // MARK: - Request and Message Handling + /// Internal context for routing responses to the correct transport. + /// + /// When handling requests, we capture the current connection at request time. + /// This ensures that when the handler completes (which may be async), the response + /// is sent to the correct client even if `self.connection` has changed in the meantime. + /// + /// This pattern is critical for HTTP transports where multiple clients can connect + /// and the server's `connection` reference gets reassigned. + private struct RequestContext { + /// The transport connection captured at request time + let capturedConnection: (any Transport)? + /// The ID of the request being handled + let requestId: RequestId + /// The session ID from the transport, if available. + /// + /// For HTTP transports with multiple concurrent clients, this identifies + /// the specific session. Used for per-session features like log levels. + let sessionId: String? + } + + /// Wrapper for encoding type-erased notifications as JSON-RPC messages. + private struct NotificationWrapper: Encodable { + let jsonrpc = "2.0" + let method: String + let params: Value + + init(notification: any Notification) { + self.method = type(of: notification).name + + // Encode the notification's params to Value + // Since Notification is Codable, we encode it and extract the params field + let encoder = JSONEncoder() + let decoder = JSONDecoder() + if let data = try? encoder.encode(notification), + let dict = try? decoder.decode([String: Value].self, from: data), + let params = dict["params"] { + self.params = params + } else { + self.params = .object([:]) + } + } + } + + /// Send a response using the captured request context. + /// + /// This ensures responses are routed to the correct client by: + /// 1. Using the connection that was active when the request was received + /// 2. Passing the request ID so multiplexing transports can route correctly + private func send(_ response: Response, using context: RequestContext) async throws { + guard let connection = context.capturedConnection else { + await logger?.warning( + "Cannot send response - connection was nil at request time", + metadata: ["requestId": "\(context.requestId)"] + ) + return + } + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + + let responseData = try encoder.encode(response) + try await connection.send(responseData, relatedRequestId: context.requestId) + } + /// Handle a request and either send the response immediately or return it /// /// - Parameters: @@ -450,6 +957,16 @@ public actor Server { private func handleRequest(_ request: Request, sendResponse: Bool = true) async throws -> Response? { + // Capture the connection and session ID at request time. + // This ensures responses go to the correct client even if self.connection + // changes while the handler is executing (e.g., another client connects). + let capturedConnection = self.connection + let context = RequestContext( + capturedConnection: capturedConnection, + requestId: request.id, + sessionId: await capturedConnection?.sessionId + ) + // Check if this is a pre-processed error request (empty method) if request.method.isEmpty && !sendResponse { // This is a placeholder for an invalid request that couldn't be parsed in batch mode @@ -483,29 +1000,92 @@ public actor Server { let response = AnyMethod.response(id: request.id, error: error) if sendResponse { - try await send(response) + try await send(response, using: context) return nil } return response } + // Create the public handler context with sendNotification capability + let handlerContext = RequestHandlerContext( + sendNotification: { [context] notification in + guard let connection = context.capturedConnection else { + throw MCPError.internalError("Cannot send notification - connection was nil at request time") + } + + // Wrap the notification in a JSON-RPC message structure + let wrapper = NotificationWrapper(notification: notification) + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + + let notificationData = try encoder.encode(wrapper) + try await connection.send(notificationData, relatedRequestId: context.requestId) + }, + sendMessage: { [context] message in + guard let connection = context.capturedConnection else { + throw MCPError.internalError("Cannot send notification - connection was nil at request time") + } + + // Message already encodes to JSON-RPC format with method and params + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + + let messageData = try encoder.encode(message) + try await connection.send(messageData, relatedRequestId: context.requestId) + }, + sendData: { [context] data in + guard let connection = context.capturedConnection else { + throw MCPError.internalError("Cannot send data - connection was nil at request time") + } + + // Send raw data (used for queued task messages) + try await connection.send(data, relatedRequestId: context.requestId) + }, + sessionId: context.sessionId, + shouldSendLogMessage: { [weak self, context] level in + guard let self else { return true } + return await self.shouldSendLogMessage(at: level, forSession: context.sessionId) + } + ) + do { // Handle request and get response - let response = try await handler(request) + let response: Response = try await handler(request, context: handlerContext) + + // Check cancellation before sending response (per MCP spec: + // "Receivers of a cancellation notification SHOULD... Not send a response + // for the cancelled request") + if Task.isCancelled { + await logger?.debug( + "Request cancelled, suppressing response", + metadata: ["id": "\(request.id)"] + ) + return nil + } if sendResponse { - try await send(response) + try await send(response, using: context) return nil } return response } catch { + // Also check cancellation on error path - don't send error response if cancelled + if Task.isCancelled { + await logger?.debug( + "Request cancelled during error handling, suppressing response", + metadata: ["id": "\(request.id)"] + ) + return nil + } + let mcpError = error as? MCPError ?? MCPError.internalError(error.localizedDescription) - let response = AnyMethod.response(id: request.id, error: mcpError) + let response: Response = AnyMethod.response(id: request.id, error: mcpError) if sendResponse { - try await send(response) + try await send(response, using: context) return nil } @@ -543,17 +1123,408 @@ public actor Server { } } + /// Handle a response from the client (for server→client requests). + private func handleClientResponse(_ response: Response) async { + await logger?.trace( + "Processing client response", + metadata: ["id": "\(response.id)"]) + + // Check response routers first (e.g., for task-related responses) + for router in responseRouters { + switch response.result { + case .success(let value): + if await router.routeResponse(requestId: response.id, response: value) { + await logger?.trace( + "Response routed via router", + metadata: ["id": "\(response.id)"]) + return + } + case .failure(let error): + if await router.routeError(requestId: response.id, error: error) { + await logger?.trace( + "Error routed via router", + metadata: ["id": "\(response.id)"]) + return + } + } + } + + // Fall back to normal pending request handling + if let pendingRequest = pendingRequests.removeValue(forKey: response.id) { + switch response.result { + case .success(let value): + pendingRequest.resume(returning: value) + case .failure(let error): + pendingRequest.resume(throwing: error) + } + } else { + await logger?.warning( + "Received response for unknown request", + metadata: ["id": "\(response.id)"]) + } + } + + // MARK: - Server→Client Requests (Bidirectional Communication) + + /// Send a request to the client and wait for a response. + /// + /// This enables bidirectional communication where the server can request + /// information from the client (e.g., roots, sampling, elicitation). + /// + /// - Parameter request: The request to send + /// - Returns: The result from the client + public func sendRequest(_ request: Request) async throws -> M.Result { + guard let connection = connection else { + throw MCPError.internalError("Server connection not initialized") + } + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + let requestData = try encoder.encode(request) + + // Create stream for receiving the response + let (stream, continuation) = AsyncThrowingStream.makeStream() + + // Clean up pending request if cancelled + let requestId = request.id + continuation.onTermination = { @Sendable [weak self] _ in + Task { await self?.cleanupPendingRequest(id: requestId) } + } + + // Register the pending request + pendingRequests[request.id] = AnyServerPendingRequest(continuation: continuation) + + // Send the request + do { + try await connection.send(requestData) + } catch { + pendingRequests.removeValue(forKey: request.id) + continuation.finish(throwing: error) + throw error + } + + // Wait for response + for try await result in stream { + return result + } + + throw MCPError.internalError("No response received from client") + } + + private func cleanupPendingRequest(id: RequestId) { + pendingRequests.removeValue(forKey: id) + } + + // MARK: - In-Flight Request Tracking (Protocol-Level Cancellation) + + /// Track an in-flight request handler Task. + private func trackInFlightRequest(_ requestId: RequestId, task: Task) { + inFlightHandlerTasks[requestId] = task + } + + /// Remove an in-flight request handler Task. + private func removeInFlightRequest(_ requestId: RequestId) { + inFlightHandlerTasks.removeValue(forKey: requestId) + } + + /// Cancel an in-flight request handler Task. + /// + /// Called when a CancelledNotification is received for a specific requestId. + /// Per MCP spec, if the request is unknown or already completed, this is a no-op. + private func cancelInFlightRequest(_ requestId: RequestId, reason: String?) async { + if let task = inFlightHandlerTasks[requestId] { + task.cancel() + await logger?.debug( + "Cancelled in-flight request", + metadata: [ + "id": "\(requestId)", + "reason": "\(reason ?? "none")", + ] + ) + } + // Per spec: MAY ignore if request is unknown - no error needed + } + + /// Generate a unique request ID for server→client requests. + private func generateRequestId() -> RequestId { + let id = nextRequestId + nextRequestId += 1 + return .number(id) + } + + /// Request the list of roots from the client. + /// + /// Roots represent filesystem directories that the client has access to. + /// Servers can use this to understand the scope of files they can work with. + /// + /// - Throws: MCPError if the client doesn't support roots or if the request fails. + /// - Returns: The list of roots from the client. + public func listRoots() async throws -> [Root] { + // Check that client supports roots + guard clientCapabilities?.roots != nil else { + throw MCPError.invalidRequest("Client does not support roots capability") + } + + let request: Request = ListRoots.request(id: generateRequestId()) + let result = try await sendRequest(request) + return result.roots + } + + /// Request a sampling completion from the client (without tools). + /// + /// This enables servers to request LLM completions through the client, + /// allowing sophisticated agentic behaviors while maintaining security. + /// + /// The result will be a single content block (text, image, or audio). + /// For tool-enabled sampling, use `createMessageWithTools(_:)` instead. + /// + /// - Parameter params: The sampling parameters including messages, model preferences, etc. + /// - Throws: MCPError if the client doesn't support sampling or if the request fails. + /// - Returns: The sampling result from the client containing a single content block. + public func createMessage(_ params: CreateSamplingMessage.Parameters) async throws -> CreateSamplingMessage.Result { + // Check that client supports sampling + guard clientCapabilities?.sampling != nil else { + throw MCPError.invalidRequest("Client does not support sampling capability") + } + + let request: Request = CreateSamplingMessage.request(id: generateRequestId(), params) + return try await sendRequest(request) + } + + /// Request a sampling completion from the client with tool support. + /// + /// This enables servers to request LLM completions that may involve tool use. + /// The result may contain tool use content, and content can be an array for parallel tool calls. + /// + /// - Parameter params: The sampling parameters including messages, tools, and model preferences. + /// - Throws: MCPError if the client doesn't support sampling or tool capabilities. + /// - Returns: The sampling result from the client, which may include tool use content. + public func createMessageWithTools(_ params: CreateSamplingMessageWithTools.Parameters) async throws -> CreateSamplingMessageWithTools.Result { + // Check that client supports sampling + guard clientCapabilities?.sampling != nil else { + throw MCPError.invalidRequest("Client does not support sampling capability") + } + + // Check tools capability + guard clientCapabilities?.sampling?.tools != nil else { + throw MCPError.invalidRequest("Client does not support sampling tools capability") + } + + // Validate tool_use/tool_result message structure per MCP specification + try Sampling.Message.validateToolUseResultMessages(params.messages) + + let request: Request = CreateSamplingMessageWithTools.request(id: generateRequestId(), params) + return try await sendRequest(request) + } + + /// Request user input via elicitation from the client. + /// + /// Elicitation allows servers to request structured input from users through + /// the client, either via forms or external URLs (e.g., OAuth flows). + /// + /// - Parameter params: The elicitation parameters. + /// - Throws: MCPError if the client doesn't support elicitation or if the request fails. + /// - Returns: The elicitation result from the client. + public func elicit(_ params: Elicit.Parameters) async throws -> Elicit.Result { + // Check that client supports elicitation + guard clientCapabilities?.elicitation != nil else { + throw MCPError.invalidRequest("Client does not support elicitation capability") + } + + // Check mode-specific capabilities + switch params { + case .form: + guard clientCapabilities?.elicitation?.form != nil else { + throw MCPError.invalidRequest("Client does not support form elicitation") + } + case .url: + guard clientCapabilities?.elicitation?.url != nil else { + throw MCPError.invalidRequest("Client does not support URL elicitation") + } + } + + let request: Request = Elicit.request(id: generateRequestId(), params) + let result = try await sendRequest(request) + + // TODO: Add elicitation response validation against the requestedSchema. + // TypeScript SDK uses JSON Schema validators (AJV, CfWorker) to validate + // elicitation responses against the requestedSchema. Python SDK uses Pydantic. + // The ideal solution is to use the same JSON Schema validator for both + // elicitation and tool validation, for spec compliance and consistency. + + return result + } + private func checkInitialized() throws { guard isInitialized else { throw MCPError.invalidRequest("Server is not initialized") } } + // MARK: - Client Task Polling (Server → Client) + + /// Get a task from the client. + /// + /// Internal method used by experimental server task features. + func getClientTask(taskId: String) async throws -> GetTask.Result { + guard clientCapabilities?.tasks != nil else { + throw MCPError.invalidRequest("Client does not support tasks capability") + } + + let request = GetTask.request(.init(taskId: taskId)) + return try await sendRequest(request) + } + + /// Get the result payload of a client task. + /// + /// Internal method used by experimental server task features. + func getClientTaskResult(taskId: String) async throws -> GetTaskPayload.Result { + guard clientCapabilities?.tasks != nil else { + throw MCPError.invalidRequest("Client does not support tasks capability") + } + + let request = GetTaskPayload.request(.init(taskId: taskId)) + return try await sendRequest(request) + } + + /// Get the task result decoded as a specific type. + /// + /// Internal method used by experimental server task features. + func getClientTaskResultAs(taskId: String, type: T.Type) async throws -> T { + let result = try await getClientTaskResult(taskId: taskId) + + // The result's extraFields contain the actual result payload + guard let extraFields = result.extraFields else { + throw MCPError.invalidParams("Task result has no payload") + } + + // Convert extraFields to the target type + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let jsonData = try encoder.encode(extraFields) + return try decoder.decode(T.self, from: jsonData) + } + + // MARK: - Task-Augmented Requests (Server → Client) + + /// Send a task-augmented elicitation request to the client. + /// + /// The client returns a `CreateTaskResult` instead of an `ElicitResult`. + /// Use client task polling to get the final result. + /// + /// Internal method used by experimental server task features. + func sendElicitAsTask(_ params: Elicit.Parameters) async throws -> CreateTaskResult { + // Check that client supports task-augmented elicitation + try requireTaskAugmentedElicitation(clientCapabilities) + + // Check mode-specific capabilities + switch params { + case .form: + guard clientCapabilities?.elicitation?.form != nil else { + throw MCPError.invalidRequest("Client does not support form elicitation") + } + case .url: + guard clientCapabilities?.elicitation?.url != nil else { + throw MCPError.invalidRequest("Client does not support URL elicitation") + } + } + + guard let connection else { + throw MCPError.internalError("Server connection not initialized") + } + + // Build the request + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + let request: Request = Elicit.request(id: generateRequestId(), params) + let requestData = try encoder.encode(request) + + // Create stream for receiving the response + let (stream, continuation) = AsyncThrowingStream.makeStream() + + let requestId = request.id + continuation.onTermination = { @Sendable [weak self] _ in + Task { await self?.cleanupPendingRequest(id: requestId) } + } + + // Register the pending request + pendingRequests[requestId] = AnyServerPendingRequest(continuation: continuation) + + // Send the request + do { + try await connection.send(requestData) + } catch { + pendingRequests.removeValue(forKey: requestId) + continuation.finish(throwing: error) + throw error + } + + // Wait for single result + for try await result in stream { + return result + } + + throw MCPError.internalError("No response received") + } + + /// Send a task-augmented sampling request to the client. + /// + /// The client returns a `CreateTaskResult` instead of a `CreateSamplingMessage.Result`. + /// Use client task polling to get the final result. + /// + /// Internal method used by experimental server task features. + func sendCreateMessageAsTask(_ params: CreateSamplingMessage.Parameters) async throws -> CreateTaskResult { + // Check that client supports task-augmented sampling + try requireTaskAugmentedSampling(clientCapabilities) + + guard clientCapabilities?.sampling != nil else { + throw MCPError.invalidRequest("Client does not support sampling capability") + } + + guard let connection else { + throw MCPError.internalError("Server connection not initialized") + } + + // Build the request + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + let request = CreateSamplingMessage.request(id: generateRequestId(), params) + let requestData = try encoder.encode(request) + + // Create stream for receiving the response + let (stream, continuation) = AsyncThrowingStream.makeStream() + + let requestId = request.id + continuation.onTermination = { @Sendable [weak self] _ in + Task { await self?.cleanupPendingRequest(id: requestId) } + } + + // Register the pending request + pendingRequests[requestId] = AnyServerPendingRequest(continuation: continuation) + + // Send the request + do { + try await connection.send(requestData) + } catch { + pendingRequests.removeValue(forKey: requestId) + continuation.finish(throwing: error) + throw error + } + + // Wait for single result + for try await result in stream { + return result + } + + throw MCPError.internalError("No response received") + } + private func registerDefaultHandlers( initializeHook: (@Sendable (Client.Info, Client.Capabilities) async throws -> Void)? ) { // Initialize - withMethodHandler(Initialize.self) { [weak self] params in + withRequestHandler(Initialize.self) { [weak self] params, _ in guard let self = self else { throw MCPError.internalError("Server was deallocated") } @@ -588,7 +1559,64 @@ public actor Server { } // Ping - withMethodHandler(Ping.self) { _ in return Empty() } + withRequestHandler(Ping.self) { _, _ in return Empty() } + + // CancelledNotification: Handle cancellation of in-flight requests + onNotification(CancelledNotification.self) { [weak self] message in + guard let self else { return } + guard let requestId = message.params.requestId else { + // Per protocol 2025-11-25+, requestId is optional. + // If not provided, we cannot cancel a specific request. + return + } + await self.cancelInFlightRequest(requestId, reason: message.params.reason) + } + + // Logging: Set minimum log level (only if logging capability is enabled) + if capabilities.logging != nil { + withRequestHandler(SetLoggingLevel.self) { [weak self] params, context in + guard let self else { + throw MCPError.internalError("Server was deallocated") + } + await self.setLoggingLevel(params.level, forSession: context.sessionId) + return Empty() + } + } + } + + /// Set the minimum log level for messages sent to a specific session. + /// + /// After this is set, only log messages at this level or higher (more severe) + /// will be sent to clients in this session via `sendLogMessage`. + /// + /// - Parameters: + /// - level: The minimum log level to send. + /// - sessionId: The session identifier, or `nil` for transports without sessions. + private func setLoggingLevel(_ level: LoggingLevel, forSession sessionId: String?) { + loggingLevels[sessionId] = level + } + + /// Check if a log message at the given level should be sent to a specific session. + /// + /// Returns `false` if: + /// - The logging capability is not declared, OR + /// - The message level is below the minimum level set by the client for this session + /// + /// - Parameters: + /// - level: The level of the log message to check. + /// - sessionId: The session identifier, or `nil` for transports without sessions. + /// - Returns: `true` if the message should be sent, `false` if it should be filtered out. + func shouldSendLogMessage(at level: LoggingLevel, forSession sessionId: String?) -> Bool { + // Check if logging capability is declared (matching TypeScript SDK behavior) + guard capabilities.logging != nil else { return false } + + guard let sessionLevel = loggingLevels[sessionId] else { + // If no level is set for this session, send all messages (per MCP spec: + // "If no logging/setLevel request has been sent from the client, the server + // MAY decide which messages to send automatically") + return true + } + return level.isAtLeast(sessionLevel) } private func setInitialState( diff --git a/Sources/MCP/Server/SessionManager.swift b/Sources/MCP/Server/SessionManager.swift new file mode 100644 index 00000000..62d84159 --- /dev/null +++ b/Sources/MCP/Server/SessionManager.swift @@ -0,0 +1,221 @@ +import Foundation + +/// Optional helper actor for managing HTTP sessions. +/// +/// `SessionManager` provides a thread-safe storage layer for managing multiple concurrent +/// HTTP sessions. It's a Swift equivalent of the TypeScript SDK's session dictionary pattern: +/// +/// ```typescript +/// // TypeScript SDK pattern: +/// const transports: { [sessionId: string]: HTTPServerTransport } = {}; +/// ``` +/// +/// In Swift, we use an actor for thread-safe access: +/// ```swift +/// let sessionManager = SessionManager() +/// await sessionManager.store(transport, forSessionId: sessionId) +/// ``` +/// +/// This actor is **optional** - applications can implement their own session management +/// logic directly using dictionaries or other data structures if preferred. +/// +/// ## Usage Pattern +/// +/// - One `Server` instance can be shared across all sessions +/// - Each session has its own `HTTPServerTransport` +/// - The application routes requests to the correct transport by session ID +/// +/// ```swift +/// // Create session manager +/// let sessionManager = SessionManager() +/// +/// // In your HTTP handler: +/// func handleMCPRequest(_ request: HTTPRequest) async -> HTTPResponse { +/// let sessionId = request.headers[HTTPHeader.sessionId] +/// let isInitializeRequest = request.body.contains("\"method\":\"initialize\"") +/// +/// // Get or create transport +/// let transport: HTTPServerTransport +/// if let sessionId, let existing = await sessionManager.transport(forSessionId: sessionId) { +/// transport = existing +/// } else if isInitializeRequest { +/// // Create new transport with session callbacks +/// transport = HTTPServerTransport( +/// options: .init( +/// sessionIdGenerator: { UUID().uuidString }, +/// onSessionInitialized: { id in +/// await sessionManager.store(transport, forSessionId: id) +/// }, +/// onSessionClosed: { id in +/// await sessionManager.remove(id) +/// } +/// ) +/// ) +/// try await server.start(transport: transport) +/// } else { +/// // No session, no initialization - reject +/// return HTTPResponse(statusCode: 400, ...) +/// } +/// +/// return await transport.handleRequest(request) +/// } +/// ``` +public actor SessionManager { + /// Storage for transports by session ID + private var transports: [String: HTTPServerTransport] = [:] + + /// Last activity time for each session (for cleanup) + private var lastActivity: [String: Date] = [:] + + /// Maximum number of concurrent sessions (nil = unlimited) + public var maxSessions: Int? + + /// Creates a new SessionManager. + /// + /// - Parameter maxSessions: Maximum concurrent sessions allowed (nil for unlimited) + public init(maxSessions: Int? = nil) { + self.maxSessions = maxSessions + } + + /// Gets an existing transport for the given session ID. + /// + /// - Parameter sessionId: The session ID to look up + /// - Returns: The transport for this session, or nil if not found + public func transport(forSessionId sessionId: String) -> HTTPServerTransport? { + if let transport = transports[sessionId] { + lastActivity[sessionId] = Date() + return transport + } + return nil + } + + /// Stores a transport for a session ID. + /// + /// - Parameters: + /// - transport: The transport to store + /// - sessionId: The session ID + public func store(_ transport: HTTPServerTransport, forSessionId sessionId: String) { + transports[sessionId] = transport + lastActivity[sessionId] = Date() + } + + /// Removes a transport from the session manager. + /// + /// - Parameter sessionId: The session ID to remove + public func remove(_ sessionId: String) { + transports.removeValue(forKey: sessionId) + lastActivity.removeValue(forKey: sessionId) + } + + /// Checks if capacity allows adding a new session. + /// + /// - Returns: true if a new session can be added, false if at capacity + public func canAddSession() -> Bool { + guard let max = maxSessions else { return true } + return transports.count < max + } + + /// Removes all sessions that have been inactive for longer than the specified duration. + /// + /// Call this periodically to clean up stale sessions. + /// + /// - Parameter timeout: Sessions inactive for longer than this duration will be removed + /// - Returns: The number of sessions removed + @discardableResult + public func cleanUpStaleSessions(olderThan timeout: Duration) async -> Int { + let cutoff = Date().addingTimeInterval(-timeout.timeInterval) + var removed = 0 + + for (sessionId, activity) in lastActivity where activity < cutoff { + if let transport = transports[sessionId] { + await transport.close() + } + transports.removeValue(forKey: sessionId) + lastActivity.removeValue(forKey: sessionId) + removed += 1 + } + + return removed + } + + /// The number of active sessions. + public var activeSessionCount: Int { + transports.count + } + + /// All active session IDs. + public var activeSessionIds: [String] { + Array(transports.keys) + } + + /// Closes all sessions and clears the session manager. + public func closeAll() async { + for (_, transport) in transports { + await transport.close() + } + transports.removeAll() + lastActivity.removeAll() + } +} + +/// Errors that can occur during HTTP session management. +/// +/// These errors are used with ``SessionManager`` to handle common session-related +/// failure scenarios when routing HTTP requests to the appropriate transport. +/// +/// ## Example Usage +/// +/// ```swift +/// func handleMCPRequest(_ request: HTTPRequest) async throws -> HTTPResponse { +/// let sessionId = request.headers[HTTPHeader.sessionId] +/// let isInitialize = request.body?.contains("initialize") ?? false +/// +/// guard let sessionId else { +/// if isInitialize { +/// // Create new session +/// } else { +/// throw SessionError.missingSessionId +/// } +/// } +/// +/// guard let transport = await sessionManager.transport(forSessionId: sessionId) else { +/// throw SessionError.sessionNotFound(sessionId) +/// } +/// +/// return await transport.handleRequest(request) +/// } +/// ``` +public enum SessionError: Error, CustomStringConvertible { + /// The requested session was not found. + /// This typically means the session expired or was never created. + case sessionNotFound(String) + + /// A session ID was required but not provided. + /// This occurs when a non-initialization request lacks the `Mcp-Session-Id` header. + case missingSessionId + + /// The maximum number of concurrent sessions has been reached. + /// The server should reject new connections until existing sessions are closed. + case capacityReached(Int) + + public var description: String { + switch self { + case .sessionNotFound(let sessionId): + return "Session not found: \(sessionId)" + case .missingSessionId: + return "Session ID required for non-initialization requests" + case .capacityReached(let max): + return "Maximum session capacity reached (\(max))" + } + } +} + +// MARK: - Duration Extension + +extension Duration { + /// Converts the duration to a TimeInterval (seconds). + var timeInterval: TimeInterval { + let (seconds, attoseconds) = self.components + return TimeInterval(seconds) + TimeInterval(attoseconds) / 1e18 + } +} diff --git a/Sources/MCP/Server/ToolNameValidation.swift b/Sources/MCP/Server/ToolNameValidation.swift new file mode 100644 index 00000000..27dec55b --- /dev/null +++ b/Sources/MCP/Server/ToolNameValidation.swift @@ -0,0 +1,111 @@ +import Foundation + +/// Result of tool name validation. +public struct ToolNameValidationResult: Sendable { + /// Whether the tool name is valid. + public let isValid: Bool + /// Warnings about the tool name (may be present even if valid). + public let warnings: [String] + + public init(isValid: Bool, warnings: [String]) { + self.isValid = isValid + self.warnings = warnings + } +} + +/// Validates tool names according to MCP specification. +/// +/// Tool names must: +/// - Be 1-128 characters long +/// - Contain only alphanumeric characters, underscores, dashes, and dots +/// +/// - SeeAlso: https://github.com/modelcontextprotocol/modelcontextprotocol/issues/986 +public func validateToolName(_ name: String) -> ToolNameValidationResult { + var warnings: [String] = [] + + // Check length + if name.isEmpty { + return ToolNameValidationResult(isValid: false, warnings: ["Tool name cannot be empty"]) + } + + if name.count > 128 { + return ToolNameValidationResult( + isValid: false, + warnings: ["Tool name exceeds maximum length of 128 characters (current: \(name.count))"] + ) + } + + // Check for specific problematic patterns (warnings, not validation failures) + if name.contains(" ") { + warnings.append("Tool name contains spaces, which may cause parsing issues") + } + + if name.contains(",") { + warnings.append("Tool name contains commas, which may cause parsing issues") + } + + // Check for potentially confusing patterns + if name.hasPrefix("-") || name.hasSuffix("-") { + warnings.append( + "Tool name starts or ends with a dash, which may cause parsing issues in some contexts" + ) + } + + if name.hasPrefix(".") || name.hasSuffix(".") { + warnings.append( + "Tool name starts or ends with a dot, which may cause parsing issues in some contexts" + ) + } + + // Check for invalid characters + let validCharacterSet = CharacterSet( + charactersIn: "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789._-" + ) + let nameCharacterSet = CharacterSet(charactersIn: name) + + if !validCharacterSet.isSuperset(of: nameCharacterSet) { + // Find invalid characters + let invalidChars = name.filter { char in + let charString = String(char) + let charSet = CharacterSet(charactersIn: charString) + return !validCharacterSet.isSuperset(of: charSet) + } + let uniqueInvalidChars = Set(invalidChars).map { "\"\($0)\"" }.joined(separator: ", ") + + warnings.append("Tool name contains invalid characters: \(uniqueInvalidChars)") + warnings.append( + "Allowed characters are: A-Z, a-z, 0-9, underscore (_), dash (-), and dot (.)" + ) + + return ToolNameValidationResult(isValid: false, warnings: warnings) + } + + return ToolNameValidationResult(isValid: true, warnings: warnings) +} + +/// Validates a tool name and logs any warnings. +/// +/// - Parameters: +/// - name: The tool name to validate +/// - logger: Optional logger for warnings +/// - Returns: Whether the tool name is valid +@discardableResult +public func validateAndWarnToolName(_ name: String) -> Bool { + let result = validateToolName(name) + + if !result.warnings.isEmpty { + print("Tool name validation warning for \"\(name)\":") + for warning in result.warnings { + print(" - \(warning)") + } + if result.isValid { + print("Tool registration will proceed, but this may cause compatibility issues.") + print("Consider updating the tool name to conform to the MCP tool naming standard.") + print( + "See SEP: Specify Format for Tool Names (https://github.com/modelcontextprotocol/modelcontextprotocol/issues/986) for more details." + ) + } + } + + return result.isValid +} diff --git a/Sources/MCP/Server/Tools.swift b/Sources/MCP/Server/Tools.swift index fd10d934..0ded3e4c 100644 --- a/Sources/MCP/Server/Tools.swift +++ b/Sources/MCP/Server/Tools.swift @@ -7,14 +7,50 @@ import Foundation /// Each tool is uniquely identified by a name and includes metadata /// describing its schema. /// -/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2024-11-05/server/tools/ +/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/server/tools/ public struct Tool: Hashable, Codable, Sendable { - /// The tool name + /// The tool name (intended for programmatic or logical use) public let name: String + /// A human-readable title for the tool, intended for UI display. + /// If not provided, the `annotations.title` or `name` should be used for display. + public let title: String? /// The tool description public let description: String? /// The tool input schema public let inputSchema: Value + /// An optional JSON Schema object defining the structure of the tool's output + /// returned in the `structuredContent` field of a `CallTool.Result`. + public let outputSchema: Value? + + /// Reserved for clients and servers to attach additional metadata. + public var _meta: [String: Value]? + + /// Optional icons representing this tool. + public var icons: [Icon]? + + /// Execution-related properties for a tool. + public struct Execution: Hashable, Codable, Sendable { + /// The tool's preference for task-augmented execution. + public enum TaskSupport: String, Hashable, Codable, Sendable { + /// Clients MUST invoke the tool as a task + case required + /// Clients MAY invoke the tool as a task or normal request + case optional + /// Clients MUST NOT attempt to invoke the tool as a task (default) + case forbidden + } + + /// Indicates the tool's preference for task-augmented execution. + /// If not present, defaults to "forbidden". + public var taskSupport: TaskSupport? + + public init(taskSupport: TaskSupport? = nil) { + self.taskSupport = taskSupport + } + } + + /// Execution-related properties for the tool. + public var execution: Execution? /// Annotations that provide display-facing and operational information for a Tool. /// @@ -85,37 +121,77 @@ public struct Tool: Hashable, Codable, Sendable { /// Initialize a tool with a name, description, input schema, and annotations public init( name: String, - description: String?, + title: String? = nil, + description: String? = nil, inputSchema: Value, + outputSchema: Value? = nil, + _meta: [String: Value]? = nil, + icons: [Icon]? = nil, + execution: Execution? = nil, annotations: Annotations = nil ) { self.name = name + self.title = title self.description = description self.inputSchema = inputSchema + self.outputSchema = outputSchema + self._meta = _meta + self.icons = icons + self.execution = execution self.annotations = annotations } - /// Content types that can be returned by a tool + // TODO: Consider consolidating with Prompt.Message.Content into a shared ContentBlock type + // in a future breaking change release. The spec uses a single ContentBlock type. + /// Content types that can be returned by a tool. + /// + /// Matches the MCP spec (2025-11-25) ContentBlock union: + /// - TextContent, ImageContent, AudioContent, ResourceLink, EmbeddedResource public enum Content: Hashable, Codable, Sendable { + /// Type alias for content-level annotations (with audience, priority, lastModified). + /// Not to be confused with `Tool.Annotations` which are tool-specific hints. + public typealias ContentAnnotations = MCP.Annotations + /// Text content - case text(String) + case text(String, annotations: ContentAnnotations?, _meta: [String: Value]?) /// Image content - case image(data: String, mimeType: String, metadata: [String: String]?) + case image(data: String, mimeType: String, annotations: ContentAnnotations?, _meta: [String: Value]?) /// Audio content - case audio(data: String, mimeType: String) - /// Embedded resource content - case resource(uri: String, mimeType: String, text: String?) + case audio(data: String, mimeType: String, annotations: ContentAnnotations?, _meta: [String: Value]?) + /// Embedded resource content (includes actual content) + case resource(resource: Resource.Content, annotations: ContentAnnotations?, _meta: [String: Value]?) + /// Resource link (reference to a resource that can be read) + case resourceLink(ResourceLink) + + // MARK: - Convenience initializers (backwards compatibility) + + /// Creates text content + public static func text(_ text: String) -> Content { + .text(text, annotations: nil, _meta: nil) + } + + /// Creates image content + public static func image(data: String, mimeType: String) -> Content { + .image(data: data, mimeType: mimeType, annotations: nil, _meta: nil) + } + + /// Creates audio content + public static func audio(data: String, mimeType: String) -> Content { + .audio(data: data, mimeType: mimeType, annotations: nil, _meta: nil) + } + + /// Creates embedded resource content with text + public static func resource(uri: String, mimeType: String? = nil, text: String) -> Content { + .resource(resource: .text(text, uri: uri, mimeType: mimeType), annotations: nil, _meta: nil) + } + + /// Creates embedded resource content with binary data + public static func resource(uri: String, mimeType: String? = nil, blob: Data) -> Content { + .resource(resource: .binary(blob, uri: uri, mimeType: mimeType), annotations: nil, _meta: nil) + } private enum CodingKeys: String, CodingKey { - case type - case text - case image - case resource - case audio - case uri - case mimeType - case data - case metadata + case type, text, data, mimeType, resource, annotations, _meta } public init(from decoder: Decoder) throws { @@ -125,22 +201,29 @@ public struct Tool: Hashable, Codable, Sendable { switch type { case "text": let text = try container.decode(String.self, forKey: .text) - self = .text(text) + let annotations = try container.decodeIfPresent(ContentAnnotations.self, forKey: .annotations) + let meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + self = .text(text, annotations: annotations, _meta: meta) case "image": let data = try container.decode(String.self, forKey: .data) let mimeType = try container.decode(String.self, forKey: .mimeType) - let metadata = try container.decodeIfPresent( - [String: String].self, forKey: .metadata) - self = .image(data: data, mimeType: mimeType, metadata: metadata) + let annotations = try container.decodeIfPresent(ContentAnnotations.self, forKey: .annotations) + let meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + self = .image(data: data, mimeType: mimeType, annotations: annotations, _meta: meta) case "audio": let data = try container.decode(String.self, forKey: .data) let mimeType = try container.decode(String.self, forKey: .mimeType) - self = .audio(data: data, mimeType: mimeType) + let annotations = try container.decodeIfPresent(ContentAnnotations.self, forKey: .annotations) + let meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + self = .audio(data: data, mimeType: mimeType, annotations: annotations, _meta: meta) case "resource": - let uri = try container.decode(String.self, forKey: .uri) - let mimeType = try container.decode(String.self, forKey: .mimeType) - let text = try container.decodeIfPresent(String.self, forKey: .text) - self = .resource(uri: uri, mimeType: mimeType, text: text) + let resourceContent = try container.decode(Resource.Content.self, forKey: .resource) + let annotations = try container.decodeIfPresent(ContentAnnotations.self, forKey: .annotations) + let meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + self = .resource(resource: resourceContent, annotations: annotations, _meta: meta) + case "resource_link": + let link = try ResourceLink(from: decoder) + self = .resourceLink(link) default: throw DecodingError.dataCorruptedError( forKey: .type, in: container, debugDescription: "Unknown tool content type") @@ -151,39 +234,56 @@ public struct Tool: Hashable, Codable, Sendable { var container = encoder.container(keyedBy: CodingKeys.self) switch self { - case .text(let text): + case .text(let text, let annotations, let meta): try container.encode("text", forKey: .type) try container.encode(text, forKey: .text) - case .image(let data, let mimeType, let metadata): + try container.encodeIfPresent(annotations, forKey: .annotations) + try container.encodeIfPresent(meta, forKey: ._meta) + case .image(let data, let mimeType, let annotations, let meta): try container.encode("image", forKey: .type) try container.encode(data, forKey: .data) try container.encode(mimeType, forKey: .mimeType) - try container.encodeIfPresent(metadata, forKey: .metadata) - case .audio(let data, let mimeType): + try container.encodeIfPresent(annotations, forKey: .annotations) + try container.encodeIfPresent(meta, forKey: ._meta) + case .audio(let data, let mimeType, let annotations, let meta): try container.encode("audio", forKey: .type) try container.encode(data, forKey: .data) try container.encode(mimeType, forKey: .mimeType) - case .resource(let uri, let mimeType, let text): + try container.encodeIfPresent(annotations, forKey: .annotations) + try container.encodeIfPresent(meta, forKey: ._meta) + case .resource(let resourceContent, let annotations, let meta): try container.encode("resource", forKey: .type) - try container.encode(uri, forKey: .uri) - try container.encode(mimeType, forKey: .mimeType) - try container.encodeIfPresent(text, forKey: .text) + try container.encode(resourceContent, forKey: .resource) + try container.encodeIfPresent(annotations, forKey: .annotations) + try container.encodeIfPresent(meta, forKey: ._meta) + case .resourceLink(let link): + try link.encode(to: encoder) } } } private enum CodingKeys: String, CodingKey { case name + case title case description case inputSchema + case outputSchema + case _meta + case icons + case execution case annotations } public init(from decoder: Decoder) throws { let container = try decoder.container(keyedBy: CodingKeys.self) name = try container.decode(String.self, forKey: .name) + title = try container.decodeIfPresent(String.self, forKey: .title) description = try container.decodeIfPresent(String.self, forKey: .description) inputSchema = try container.decode(Value.self, forKey: .inputSchema) + outputSchema = try container.decodeIfPresent(Value.self, forKey: .outputSchema) + _meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + icons = try container.decodeIfPresent([Icon].self, forKey: .icons) + execution = try container.decodeIfPresent(Tool.Execution.self, forKey: .execution) annotations = try container.decodeIfPresent(Tool.Annotations.self, forKey: .annotations) ?? .init() } @@ -191,8 +291,13 @@ public struct Tool: Hashable, Codable, Sendable { public func encode(to encoder: Encoder) throws { var container = encoder.container(keyedBy: CodingKeys.self) try container.encode(name, forKey: .name) - try container.encode(description, forKey: .description) + try container.encodeIfPresent(title, forKey: .title) + try container.encodeIfPresent(description, forKey: .description) try container.encode(inputSchema, forKey: .inputSchema) + try container.encodeIfPresent(outputSchema, forKey: .outputSchema) + try container.encodeIfPresent(_meta, forKey: ._meta) + try container.encodeIfPresent(icons, forKey: .icons) + try container.encodeIfPresent(execution, forKey: .execution) if !annotations.isEmpty { try container.encode(annotations, forKey: .annotations) } @@ -208,23 +313,60 @@ public enum ListTools: Method { public struct Parameters: NotRequired, Hashable, Codable, Sendable { public let cursor: String? + /// Request metadata including progress token. + public var _meta: RequestMeta? public init() { self.cursor = nil + self._meta = nil } - public init(cursor: String) { + public init(cursor: String? = nil, _meta: RequestMeta? = nil) { self.cursor = cursor + self._meta = _meta } } - public struct Result: Hashable, Codable, Sendable { + public struct Result: ResultWithExtraFields { + public typealias ResultCodingKeys = CodingKeys + public let tools: [Tool] public let nextCursor: String? + /// Reserved for clients and servers to attach additional metadata. + public var _meta: [String: Value]? + /// Additional fields not defined in the schema (for forward compatibility). + public var extraFields: [String: Value]? - public init(tools: [Tool], nextCursor: String? = nil) { + public init( + tools: [Tool], + nextCursor: String? = nil, + _meta: [String: Value]? = nil, + extraFields: [String: Value]? = nil + ) { self.tools = tools self.nextCursor = nextCursor + self._meta = _meta + self.extraFields = extraFields + } + + public enum CodingKeys: String, CodingKey, CaseIterable { + case tools, nextCursor, _meta + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + tools = try container.decode([Tool].self, forKey: .tools) + nextCursor = try container.decodeIfPresent(String.self, forKey: .nextCursor) + _meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + extraFields = try Self.decodeExtraFields(from: decoder) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode(tools, forKey: .tools) + try container.encodeIfPresent(nextCursor, forKey: .nextCursor) + try container.encodeIfPresent(_meta, forKey: ._meta) + try encodeExtraFields(to: encoder) } } } @@ -236,21 +378,82 @@ public enum CallTool: Method { public struct Parameters: Hashable, Codable, Sendable { public let name: String + // TODO: Add server-side input validation against the tool's inputSchema. + // TypeScript and Python SDKs validate arguments against the tool's inputSchema + // before calling the handler. This requires a tool cache to look up the schema. public let arguments: [String: Value]? + /// Task metadata for task-augmented requests. + /// When present, the request becomes task-augmented and returns a `CreateTaskResult` + /// instead of a normal result. + public var task: TaskMetadata? + /// Request metadata including progress token. + public var _meta: RequestMeta? - public init(name: String, arguments: [String: Value]? = nil) { + public init( + name: String, + arguments: [String: Value]? = nil, + task: TaskMetadata? = nil, + _meta: RequestMeta? = nil + ) { self.name = name self.arguments = arguments + self.task = task + self._meta = _meta } } - public struct Result: Hashable, Codable, Sendable { + public struct Result: ResultWithExtraFields { + public typealias ResultCodingKeys = CodingKeys + + /// A list of content objects that represent the unstructured result of the tool call. public let content: [Tool.Content] + /// An optional JSON object that represents the structured result of the tool call. + /// If the tool defined an `outputSchema`, this should conform to that schema. + // TODO: Add server-side output validation against the tool's outputSchema. + // TypeScript and Python SDKs validate structuredContent against outputSchema + // after the handler returns. This requires a tool cache to look up the schema. + public let structuredContent: Value? + /// Whether the tool call ended in an error. public let isError: Bool? + /// Reserved for clients and servers to attach additional metadata. + public var _meta: [String: Value]? + /// Additional fields not defined in the schema (for forward compatibility). + public var extraFields: [String: Value]? - public init(content: [Tool.Content], isError: Bool? = nil) { + public init( + content: [Tool.Content], + structuredContent: Value? = nil, + isError: Bool? = nil, + _meta: [String: Value]? = nil, + extraFields: [String: Value]? = nil + ) { self.content = content + self.structuredContent = structuredContent self.isError = isError + self._meta = _meta + self.extraFields = extraFields + } + + public enum CodingKeys: String, CodingKey, CaseIterable { + case content, structuredContent, isError, _meta + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + content = try container.decode([Tool.Content].self, forKey: .content) + structuredContent = try container.decodeIfPresent(Value.self, forKey: .structuredContent) + isError = try container.decodeIfPresent(Bool.self, forKey: .isError) + _meta = try container.decodeIfPresent([String: Value].self, forKey: ._meta) + extraFields = try Self.decodeExtraFields(from: decoder) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode(content, forKey: .content) + try container.encodeIfPresent(structuredContent, forKey: .structuredContent) + try container.encodeIfPresent(isError, forKey: .isError) + try container.encodeIfPresent(_meta, forKey: ._meta) + try encodeExtraFields(to: encoder) } } } @@ -259,4 +462,6 @@ public enum CallTool: Method { /// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2024-11-05/server/tools/#list-changed-notification public struct ToolListChangedNotification: Notification { public static let name: String = "notifications/tools/list_changed" + + public typealias Parameters = NotificationParams } diff --git a/Tests/MCPTests/AdditionalServerTests.swift b/Tests/MCPTests/AdditionalServerTests.swift new file mode 100644 index 00000000..326255de --- /dev/null +++ b/Tests/MCPTests/AdditionalServerTests.swift @@ -0,0 +1,677 @@ +import Foundation +import Testing + +@testable import MCP + +/// Tests for additional server functionality. +/// +/// These tests follow the TypeScript SDK patterns from: +/// - `packages/server/test/server/streamableHttp.test.ts` +/// +/// Additional TypeScript tests that are covered elsewhere: +/// - DNS rebinding protection tests - see HTTPServerTransportTests +/// - JSON response mode tests - see HTTPServerTransportTests +/// - Pre-parsed body tests - Swift SDK handles body parsing differently +/// +/// TypeScript tests not applicable in Swift: +/// +/// 1. `should support sync onsessioninitialized callback (backwards compatibility)` +/// Rationale: Swift callbacks are always async with signature `@Sendable (String) async -> Void`. +/// The language doesn't support sync/async callback overloading like TypeScript does. +/// +/// 2. `should propagate errors from async onsessioninitialized callback` +/// 3. `should propagate errors from async onsessionclosed callback` +/// Rationale: Swift callback signature is `@Sendable (String) async -> Void` (non-throwing). +/// TypeScript callbacks can throw because JavaScript functions can always throw. +/// This is a language difference - Swift callbacks cannot throw by design. +/// If error handling is needed, Swift users should handle errors inside the callback itself. +@Suite("Additional Server Tests") +struct AdditionalServerTests { + + // MARK: - Test Helpers + + /// Creates a configured MCP Server with tools for testing + func createTestServer() -> Server { + let server = Server( + name: "test-server", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + return server + } + + /// Sets up tool handlers on the server + func setUpToolHandlers(_ server: Server) async { + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool( + name: "greet", + description: "A simple greeting tool", + inputSchema: [ + "type": "object", + "properties": ["name": ["type": "string"]] + ] + ) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, _ in + switch request.name { + case "greet": + let name = request.arguments?["name"]?.stringValue ?? "World" + return CallTool.Result(content: [.text("Hello, \(name)!")]) + default: + return CallTool.Result(content: [.text("Unknown tool")], isError: true) + } + } + } + + // MARK: - 5.1 Response routing with concurrent requests + + @Test("Response messages are sent to the connection that sent the request") + func responseRoutingWithConcurrentRequests() async throws { + let server = createTestServer() + await setUpToolHandlers(server) + + let sessionId = UUID().uuidString + + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { sessionId }) + ) + try await server.start(transport: transport) + + // Initialize + let initRequest = TestPayloads.initializeRequest(id: "init") + _ = await transport.handleRequest(TestPayloads.postRequest(body: initRequest)) + + // Prepare two different requests + let listToolsRequest = TestPayloads.listToolsRequest(id: "req-1") + let callToolRequest = TestPayloads.callToolRequest( + id: "req-2", + name: "greet", + arguments: ["name": "Connection2"] + ) + + // Send requests concurrently + async let response1Task = transport.handleRequest(TestPayloads.postRequest(body: listToolsRequest, sessionId: sessionId)) + async let response2Task = transport.handleRequest(TestPayloads.postRequest(body: callToolRequest, sessionId: sessionId)) + + let response1 = await response1Task + let response2 = await response2Task + + #expect(response1.statusCode == 200) + #expect(response2.statusCode == 200) + + // Verify response 1 contains tools/list result with correct ID + if let body1 = response1.body { + let text1 = String(data: body1, encoding: .utf8) ?? "" + #expect(text1.contains("\"id\":\"req-1\"") || text1.contains("\"id\": \"req-1\""), "Response 1 should have ID req-1") + #expect(text1.contains("tools") || text1.contains("greet"), "Response 1 should contain tools list") + } + + // Verify response 2 contains tools/call result with correct ID + if let body2 = response2.body { + let text2 = String(data: body2, encoding: .utf8) ?? "" + #expect(text2.contains("\"id\":\"req-2\"") || text2.contains("\"id\": \"req-2\""), "Response 2 should have ID req-2") + #expect(text2.contains("Hello, Connection2"), "Response 2 should contain greeting result") + } + } + + @Test("Multiple sequential requests maintain isolation") + func multipleSequentialRequestsMaintainIsolation() async throws { + let server = createTestServer() + await setUpToolHandlers(server) + + let sessionId = UUID().uuidString + + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { sessionId }) + ) + try await server.start(transport: transport) + + // Initialize + let initRequest = TestPayloads.initializeRequest(id: "init") + _ = await transport.handleRequest(TestPayloads.postRequest(body: initRequest)) + + // Send 3 sequential requests with different greet names + let names = ["Alice", "Bob", "Charlie"] + var successCount = 0 + + for (index, name) in names.enumerated() { + let requestId = "req-\(index)" + let request = """ + {"jsonrpc":"2.0","method":"tools/call","id":"\(requestId)","params":{"name":"greet","arguments":{"name":"\(name)"}}} + """ + let response = await transport.handleRequest(TestPayloads.postRequest(body: request, sessionId: sessionId)) + + #expect(response.statusCode == 200, "Request for \(name) should succeed") + + // Response can be in body or stream depending on implementation + var responseText: String? = nil + + if let body = response.body, let text = String(data: body, encoding: .utf8) { + responseText = text + } else if let stream = response.stream { + // Try to read from stream with timeout + var data = Data() + let deadline = Date().addingTimeInterval(1.0) + for try await chunk in stream { + data.append(chunk) + if Date() > deadline { break } + if data.count > 0 { break } // Got some data + } + responseText = String(data: data, encoding: .utf8) + } + + if let text = responseText { + #expect(text.contains("Hello, \(name)!") || text.contains(requestId), "Response should contain greeting or request ID for \(name)") + successCount += 1 + } + } + + #expect(successCount == 3, "Should have 3 successful requests") + } + + // MARK: - 5.2 Error data in parse error response + + @Test("Include error data in parse error response for invalid JSON") + func includeErrorDataInParseErrorResponse() async throws { + let server = createTestServer() + await setUpToolHandlers(server) + + let sessionId = UUID().uuidString + + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { sessionId }) + ) + try await server.start(transport: transport) + + // Initialize + let initRequest = TestPayloads.initializeRequest(id: "init") + _ = await transport.handleRequest(TestPayloads.postRequest(body: initRequest)) + + // Send invalid JSON + let invalidJSONRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + HTTPHeader.sessionId: sessionId, + HTTPHeader.protocolVersion: Version.v2024_11_05, + ], + body: "{ invalid json }".data(using: .utf8) + ) + let response = await transport.handleRequest(invalidJSONRequest) + + #expect(response.statusCode == 400, "Should return 400 for parse error") + + // Verify error response contains proper error data + if let body = response.body, let text = String(data: body, encoding: .utf8) { + #expect(text.contains("error"), "Response should contain error field") + #expect(text.contains("\(ErrorCode.parseError)") || text.contains("Parse error"), "Should contain parse error code or message") + } + } + + @Test("Include error data for invalid JSON-RPC messages") + func includeErrorDataForInvalidJSONRPCMessages() async throws { + let server = createTestServer() + await setUpToolHandlers(server) + + let sessionId = UUID().uuidString + + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { sessionId }) + ) + try await server.start(transport: transport) + + // Initialize + let initRequest = TestPayloads.initializeRequest(id: "init") + _ = await transport.handleRequest(TestPayloads.postRequest(body: initRequest)) + + // Send invalid JSON-RPC (missing jsonrpc field) + // Note: Swift SDK may process this and return error in body with 200 status + let invalidJSONRPC = """ + {"method":"tools/list","id":"test"} + """ + let response = await transport.handleRequest(TestPayloads.postRequest(body: invalidJSONRPC, sessionId: sessionId)) + + // The Swift SDK may return 200 with error in body, or 400 + // Either is acceptable as long as the error is communicated + #expect(response.statusCode == 200 || response.statusCode == 400, "Should handle invalid JSON-RPC") + + // Verify error response contains proper error structure if there's a body + if let body = response.body, let text = String(data: body, encoding: .utf8) { + #expect(text.contains("jsonrpc") || text.contains("error") || text.contains("result"), "Response should be JSON-RPC formatted") + } + } + + @Test("Reject requests to uninitialized server") + func rejectRequestsToUninitializedServer() async throws { + let server = createTestServer() + await setUpToolHandlers(server) + + let sessionId = UUID().uuidString + + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { sessionId }) + ) + try await server.start(transport: transport) + + // Don't initialize - send request directly with a session ID + let listToolsRequest = """ + {"jsonrpc":"2.0","method":"tools/list","id":"test"} + """ + let response = await transport.handleRequest(TestPayloads.postRequest(body: listToolsRequest, sessionId: "any-session-id")) + + // Should reject because session doesn't exist + #expect(response.statusCode == 400 || response.statusCode == 404, "Should reject request to uninitialized session") + } + + // MARK: - Additional Edge Cases + + @Test("Empty batch request returns appropriate response") + func emptyBatchRequestReturnsAppropriateResponse() async throws { + let server = createTestServer() + await setUpToolHandlers(server) + + let sessionId = UUID().uuidString + + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { sessionId }) + ) + try await server.start(transport: transport) + + // Initialize + let initRequest = TestPayloads.initializeRequest(id: "init") + _ = await transport.handleRequest(TestPayloads.postRequest(body: initRequest)) + + // Send empty batch + let emptyBatch = "[]" + let response = await transport.handleRequest(TestPayloads.postRequest(body: emptyBatch, sessionId: sessionId)) + + // Empty batch returns 202 - same as TypeScript SDK + // Empty batch has no requests, so it's treated like a notification-only batch + // Per JSON-RPC spec, notification-only batches return no response (202 Accepted) + #expect(response.statusCode == 202, "Empty batch should return 202 (no requests to respond to)") + } + + @Test("Notifications in batch don't generate individual responses") + func notificationsInBatchDontGenerateResponses() async throws { + let server = createTestServer() + await setUpToolHandlers(server) + + let sessionId = UUID().uuidString + + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { sessionId }) + ) + try await server.start(transport: transport) + + // Initialize + let initRequest = TestPayloads.initializeRequest(id: "init") + _ = await transport.handleRequest(TestPayloads.postRequest(body: initRequest)) + + // Send batch with notification (no id) and request + let batchWithNotification = """ + [ + {"jsonrpc":"2.0","method":"notifications/initialized"}, + {"jsonrpc":"2.0","method":"tools/list","id":"req-1"} + ] + """ + let response = await transport.handleRequest(TestPayloads.postRequest(body: batchWithNotification, sessionId: sessionId)) + + #expect(response.statusCode == 200, "Batch should succeed") + + // Should only have response for the request, not the notification + if let body = response.body, let text = String(data: body, encoding: .utf8) { + // The response should contain req-1's result + #expect(text.contains("req-1") || text.contains("tools"), "Should contain response for request") + } + } + + @Test("Batch requests work for protocol versions before 2025-06-18") + func batchRequestsWorkForOlderProtocolVersions() async throws { + let server = createTestServer() + await setUpToolHandlers(server) + + let sessionId = UUID().uuidString + + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { sessionId }) + ) + try await server.start(transport: transport) + + // Initialize with protocol version 2024-11-05 (before batch removal) + let initRequest = TestPayloads.initializeRequest(id: "init") + _ = await transport.handleRequest(TestPayloads.postRequest(body: initRequest)) + + // Send batch of requests (both have IDs, so both need responses) + // Batching is supported in protocol versions < 2025-06-18 + let batchRequests = """ + [ + {"jsonrpc":"2.0","method":"tools/list","id":"req-1"}, + {"jsonrpc":"2.0","method":"tools/call","id":"req-2","params":{"name":"greet","arguments":{"name":"BatchUser"}}} + ] + """ + let response = await transport.handleRequest(TestPayloads.postRequest(body: batchRequests, sessionId: sessionId)) + + #expect(response.statusCode == 200, "Batch requests should succeed for protocol version 2024-11-05") + + // Check the response body contains both results + if let body = response.body, let text = String(data: body, encoding: .utf8) { + // Verify both request IDs are in the response + #expect(text.contains("req-1") || text.contains("tools"), + "Should contain response for req-1 (tools/list)") + #expect(text.contains("req-2") || text.contains("Hello, BatchUser"), + "Should contain response for req-2 (tools/call)") + } else { + // If no body, at least verify we got a response + #expect(response.stream != nil, "Should have either body or stream") + } + } + + @Test("Batch requests rejected for protocol versions >= 2025-06-18") + func batchRequestsRejectedForNewerProtocolVersions() async throws { + let server = createTestServer() + await setUpToolHandlers(server) + + let sessionId = UUID().uuidString + + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { sessionId }) + ) + try await server.start(transport: transport) + + // Initialize with protocol version 2025-06-18 (batch removal version) + let initRequest = TestPayloads.initializeRequest(id: "init", protocolVersion: Version.v2025_06_18) + _ = await transport.handleRequest(TestPayloads.postRequest(body: initRequest, protocolVersion: Version.v2025_06_18)) + + // Send batch of requests - should be rejected per spec + let batchRequests = """ + [ + {"jsonrpc":"2.0","method":"tools/list","id":"req-1"}, + {"jsonrpc":"2.0","method":"tools/call","id":"req-2","params":{"name":"greet","arguments":{"name":"BatchUser"}}} + ] + """ + let response = await transport.handleRequest(TestPayloads.postRequest(body: batchRequests, sessionId: sessionId, protocolVersion: Version.v2025_06_18)) + + // Batch requests should be rejected with 400 for protocol version >= 2025-06-18 + #expect(response.statusCode == 400, "Batch requests should be rejected for protocol version 2025-06-18") + + // Verify error message + if let body = response.body, let text = String(data: body, encoding: .utf8) { + #expect(text.contains("\(ErrorCode.invalidRequest)") || text.contains("not supported") || text.contains("Invalid"), + "Should return Invalid Request error for batch in newer protocol") + } + } + + @Test("Keep stream open after sending server notifications") + func keepStreamOpenAfterNotifications() async throws { + let server = createTestServer() + await setUpToolHandlers(server) + + let sessionId = UUID().uuidString + + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { sessionId }) + ) + try await server.start(transport: transport) + + // Initialize + let initRequest = TestPayloads.initializeRequest(id: "init") + _ = await transport.handleRequest(TestPayloads.postRequest(body: initRequest)) + + // Open standalone SSE stream + let getRequest = HTTPRequest( + method: "GET", + headers: [ + HTTPHeader.accept: "text/event-stream", + HTTPHeader.sessionId: sessionId, + HTTPHeader.protocolVersion: Version.v2024_11_05, + ] + ) + let response = await transport.handleRequest(getRequest) + + #expect(response.statusCode == 200, "GET should succeed") + #expect(response.stream != nil, "Should return a stream") + + // Send a notification through the transport + let notification = """ + {"jsonrpc":"2.0","method":"notifications/message","params":{"level":"info","data":"Test notification"}} + """ + try await transport.send(notification.data(using: .utf8)!) + + // Stream should still be open (we just verify the transport is still functional) + // We can't easily verify the stream is still open, but we can verify transport works + let listToolsRequest = """ + {"jsonrpc":"2.0","method":"tools/list","id":"test"} + """ + let listResponse = await transport.handleRequest(TestPayloads.postRequest(body: listToolsRequest, sessionId: sessionId)) + #expect(listResponse.statusCode == 200, "Transport should still be functional after sending notifications") + } + + // MARK: - Session Callback Tests + + @Test("Async onSessionInitialized callback is called") + func asyncOnSessionInitializedCallbackIsCalled() async throws { + actor CallbackTracker { + var events: [String] = [] + func add(_ event: String) { events.append(event) } + func getEvents() -> [String] { events } + } + + let tracker = CallbackTracker() + let sessionId = UUID().uuidString + + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { sessionId }, + onSessionInitialized: { id in + await tracker.add("initialized:\(id)") + } + ) + ) + + let server = createTestServer() + try await server.start(transport: transport) + + // Initialize to trigger the callback + let initRequest = TestPayloads.initializeRequest(id: "init") + let response = await transport.handleRequest(TestPayloads.postRequest(body: initRequest)) + + #expect(response.statusCode == 200) + + // Give time for async callback to complete + try await Task.sleep(for: .milliseconds(50)) + + let events = await tracker.getEvents() + #expect(events.contains("initialized:\(sessionId)"), "onSessionInitialized should be called with session ID") + } + + @Test("Async onSessionClosed callback is called on DELETE") + func asyncOnSessionClosedCallbackIsCalledOnDelete() async throws { + actor CallbackTracker { + var events: [String] = [] + func add(_ event: String) { events.append(event) } + func getEvents() -> [String] { events } + } + + let tracker = CallbackTracker() + let sessionId = UUID().uuidString + + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { sessionId }, + onSessionClosed: { id in + await tracker.add("closed:\(id)") + } + ) + ) + + let server = createTestServer() + try await server.start(transport: transport) + + // Initialize first + let initRequest = TestPayloads.initializeRequest(id: "init") + _ = await transport.handleRequest(TestPayloads.postRequest(body: initRequest)) + + // DELETE the session + let deleteRequest = HTTPRequest( + method: "DELETE", + headers: [ + HTTPHeader.sessionId: sessionId, + HTTPHeader.protocolVersion: Version.v2024_11_05, + ] + ) + let deleteResponse = await transport.handleRequest(deleteRequest) + + #expect(deleteResponse.statusCode == 200) + + // Give time for async callback to complete + try await Task.sleep(for: .milliseconds(50)) + + let events = await tracker.getEvents() + #expect(events.contains("closed:\(sessionId)"), "onSessionClosed should be called with session ID") + } + + @Test("Both async callbacks work together") + func bothAsyncCallbacksWorkTogether() async throws { + actor CallbackTracker { + var events: [String] = [] + func add(_ event: String) { events.append(event) } + func getEvents() -> [String] { events } + } + + let tracker = CallbackTracker() + let sessionId = UUID().uuidString + + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { sessionId }, + onSessionInitialized: { id in + await tracker.add("initialized:\(id)") + }, + onSessionClosed: { id in + await tracker.add("closed:\(id)") + } + ) + ) + + let server = createTestServer() + try await server.start(transport: transport) + + // Initialize to trigger first callback + let initRequest = TestPayloads.initializeRequest(id: "init") + let initResponse = await transport.handleRequest(TestPayloads.postRequest(body: initRequest)) + #expect(initResponse.statusCode == 200) + + // Give time for async callback + try await Task.sleep(for: .milliseconds(50)) + + var events = await tracker.getEvents() + #expect(events.contains("initialized:\(sessionId)"), "onSessionInitialized should be called") + + // DELETE to trigger second callback + let deleteRequest = HTTPRequest( + method: "DELETE", + headers: [ + HTTPHeader.sessionId: sessionId, + HTTPHeader.protocolVersion: Version.v2024_11_05, + ] + ) + let deleteResponse = await transport.handleRequest(deleteRequest) + #expect(deleteResponse.statusCode == 200) + + // Give time for async callback + try await Task.sleep(for: .milliseconds(50)) + + events = await tracker.getEvents() + #expect(events.contains("closed:\(sessionId)"), "onSessionClosed should be called") + #expect(events.count == 2, "Should have exactly 2 events") + } + + @Test("onSessionClosed called with correct session ID for multiple sessions") + func onSessionClosedCalledWithCorrectSessionIdForMultipleSessions() async throws { + actor CallbackTracker { + var closedSessions: [String] = [] + func add(_ sessionId: String) { closedSessions.append(sessionId) } + func getSessions() -> [String] { closedSessions } + } + + let tracker = CallbackTracker() + + // Create first transport with unique session + let sessionId1 = "session-1-\(UUID().uuidString)" + let transport1 = HTTPServerTransport( + options: .init( + sessionIdGenerator: { sessionId1 }, + onSessionClosed: { id in + await tracker.add(id) + } + ) + ) + + let server1 = createTestServer() + try await server1.start(transport: transport1) + + // Create second transport with unique session + let sessionId2 = "session-2-\(UUID().uuidString)" + let transport2 = HTTPServerTransport( + options: .init( + sessionIdGenerator: { sessionId2 }, + onSessionClosed: { id in + await tracker.add(id) + } + ) + ) + + let server2 = createTestServer() + try await server2.start(transport: transport2) + + // Initialize both transports + let initRequest = TestPayloads.initializeRequest(id: "init") + + let initResponse1 = await transport1.handleRequest(TestPayloads.postRequest(body: initRequest)) + #expect(initResponse1.statusCode == 200) + #expect(initResponse1.headers[HTTPHeader.sessionId] == sessionId1) + + let initResponse2 = await transport2.handleRequest(TestPayloads.postRequest(body: initRequest)) + #expect(initResponse2.statusCode == 200) + #expect(initResponse2.headers[HTTPHeader.sessionId] == sessionId2) + + // DELETE first session + let deleteRequest1 = HTTPRequest( + method: "DELETE", + headers: [ + HTTPHeader.sessionId: sessionId1, + HTTPHeader.protocolVersion: Version.v2024_11_05, + ] + ) + let deleteResponse1 = await transport1.handleRequest(deleteRequest1) + #expect(deleteResponse1.statusCode == 200) + + try await Task.sleep(for: .milliseconds(50)) + + var closedSessions = await tracker.getSessions() + #expect(closedSessions.count == 1) + #expect(closedSessions.contains(sessionId1)) + + // DELETE second session + let deleteRequest2 = HTTPRequest( + method: "DELETE", + headers: [ + HTTPHeader.sessionId: sessionId2, + HTTPHeader.protocolVersion: Version.v2024_11_05, + ] + ) + let deleteResponse2 = await transport2.handleRequest(deleteRequest2) + #expect(deleteResponse2.statusCode == 200) + + try await Task.sleep(for: .milliseconds(50)) + + closedSessions = await tracker.getSessions() + #expect(closedSessions.count == 2) + #expect(closedSessions.contains(sessionId1)) + #expect(closedSessions.contains(sessionId2)) + } +} diff --git a/Tests/MCPTests/CancellationTests.swift b/Tests/MCPTests/CancellationTests.swift new file mode 100644 index 00000000..c4f0c1a1 --- /dev/null +++ b/Tests/MCPTests/CancellationTests.swift @@ -0,0 +1,1848 @@ +import Foundation +import Logging +import Testing + +#if canImport(System) + import System +#else + @preconcurrency import SystemPackage +#endif + +@testable import MCP + +// MARK: - Cancellation Tests + +/// Tests for request cancellation functionality in MCP. +/// +/// Cancellation allows either client or server to signal that an ongoing +/// operation should be terminated. This is done via the `notifications/cancelled` +/// notification which can optionally include a request ID and reason. +/// +/// Reference: MCP Specification 2025-11-25 (cancellation support) +/// Based on: +/// - Python SDK: tests/server/test_cancel_handling.py +/// - TypeScript SDK: packages/core/test/shared/protocol.test.ts +@Suite("Cancellation Tests") +struct CancellationTests { + + // MARK: - CancelledNotification Encoding/Decoding Tests + + @Suite("CancelledNotification encoding/decoding") + struct CancelledNotificationEncodingTests { + + @Test("Encodes with requestId and reason") + func encodesWithRequestIdAndReason() throws { + let params = CancelledNotification.Parameters( + requestId: .string("req-123"), + reason: "User cancelled the operation" + ) + let notification = CancelledNotification.message(params) + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + let data = try encoder.encode(notification) + let json = try JSONDecoder().decode([String: Value].self, from: data) + + #expect(json["jsonrpc"] == "2.0") + #expect(json["method"] == "notifications/cancelled") + + let notificationParams = json["params"]?.objectValue + #expect(notificationParams?["requestId"]?.stringValue == "req-123") + #expect(notificationParams?["reason"]?.stringValue == "User cancelled the operation") + } + + @Test("Encodes with integer requestId") + func encodesWithIntegerRequestId() throws { + let params = CancelledNotification.Parameters( + requestId: .number(42), + reason: "Timeout" + ) + let notification = CancelledNotification.message(params) + + let encoder = JSONEncoder() + let data = try encoder.encode(notification) + let json = try JSONDecoder().decode([String: Value].self, from: data) + + let notificationParams = json["params"]?.objectValue + #expect(notificationParams?["requestId"]?.intValue == 42) + #expect(notificationParams?["reason"]?.stringValue == "Timeout") + } + + @Test("Encodes with only requestId") + func encodesWithOnlyRequestId() throws { + let params = CancelledNotification.Parameters(requestId: .string("req-456")) + let notification = CancelledNotification.message(params) + + let encoder = JSONEncoder() + let data = try encoder.encode(notification) + let json = try JSONDecoder().decode([String: Value].self, from: data) + + let notificationParams = json["params"]?.objectValue + #expect(notificationParams?["requestId"]?.stringValue == "req-456") + #expect(notificationParams?["reason"] == nil) + } + + @Test("Encodes with only reason (protocol 2025-11-25+)") + func encodesWithOnlyReason() throws { + // In protocol 2025-11-25+, requestId is optional + let params = CancelledNotification.Parameters(reason: "General cancellation") + let notification = CancelledNotification.message(params) + + let encoder = JSONEncoder() + let data = try encoder.encode(notification) + let json = try JSONDecoder().decode([String: Value].self, from: data) + + let notificationParams = json["params"]?.objectValue + #expect(notificationParams?["requestId"] == nil) + #expect(notificationParams?["reason"]?.stringValue == "General cancellation") + } + + @Test("Encodes with no parameters (protocol 2025-11-25+)") + func encodesWithNoParameters() throws { + // In protocol 2025-11-25+, both requestId and reason are optional + let params = CancelledNotification.Parameters() + let notification = CancelledNotification.message(params) + + let encoder = JSONEncoder() + let data = try encoder.encode(notification) + let json = try JSONDecoder().decode([String: Value].self, from: data) + + #expect(json["jsonrpc"] == "2.0") + #expect(json["method"] == "notifications/cancelled") + // Params should be present but may be empty + #expect(json["params"] != nil) + } + + @Test("Decodes with requestId and reason") + func decodesWithRequestIdAndReason() throws { + let jsonString = """ + { + "jsonrpc": "2.0", + "method": "notifications/cancelled", + "params": { + "requestId": "req-789", + "reason": "Operation timed out" + } + } + """ + let data = jsonString.data(using: .utf8)! + let decoded = try JSONDecoder().decode( + Message.self, from: data) + + #expect(decoded.method == "notifications/cancelled") + #expect(decoded.params.requestId == .string("req-789")) + #expect(decoded.params.reason == "Operation timed out") + } + + @Test("Decodes with integer requestId") + func decodesWithIntegerRequestId() throws { + let jsonString = """ + { + "jsonrpc": "2.0", + "method": "notifications/cancelled", + "params": { + "requestId": 123, + "reason": "Client disconnected" + } + } + """ + let data = jsonString.data(using: .utf8)! + let decoded = try JSONDecoder().decode( + Message.self, from: data) + + #expect(decoded.params.requestId == .number(123)) + #expect(decoded.params.reason == "Client disconnected") + } + + @Test("Decodes with empty params") + func decodesWithEmptyParams() throws { + let jsonString = """ + { + "jsonrpc": "2.0", + "method": "notifications/cancelled", + "params": {} + } + """ + let data = jsonString.data(using: .utf8)! + let decoded = try JSONDecoder().decode( + Message.self, from: data) + + #expect(decoded.method == "notifications/cancelled") + #expect(decoded.params.requestId == nil) + #expect(decoded.params.reason == nil) + } + + @Test("Round-trip encoding/decoding preserves all fields") + func roundTripPreservesAllFields() throws { + let original = CancelledNotification.Parameters( + requestId: .string("round-trip-test"), + reason: "Testing round-trip encoding", + _meta: ["key": .string("value")] + ) + let notification = CancelledNotification.message(original) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(notification) + let decoded = try decoder.decode(Message.self, from: data) + + #expect(decoded.params.requestId == original.requestId) + #expect(decoded.params.reason == original.reason) + #expect(decoded.params._meta?["key"]?.stringValue == "value") + } + + @Test("Notification name is correct") + func notificationNameIsCorrect() { + #expect(CancelledNotification.name == "notifications/cancelled") + } + } + + // MARK: - JSON Format Compatibility Tests + + @Suite("CancelledNotification JSON format") + struct CancelledNotificationJSONFormatTests { + + @Test("Matches TypeScript SDK format with all fields") + func matchesTypeScriptFormatAllFields() throws { + // TypeScript SDK format for cancelled notification + let params = CancelledNotification.Parameters( + requestId: .string("test-request"), + reason: "User cancelled" + ) + let notification = CancelledNotification.message(params) + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + let data = try encoder.encode(notification) + let jsonString = String(data: data, encoding: .utf8)! + + // Verify JSON structure matches expected format + #expect(jsonString.contains("\"jsonrpc\":\"2.0\"")) + #expect(jsonString.contains("\"method\":\"notifications/cancelled\"")) + #expect(jsonString.contains("\"params\"")) + #expect(jsonString.contains("\"requestId\":\"test-request\"")) + #expect(jsonString.contains("\"reason\":\"User cancelled\"")) + } + + @Test("Matches Python SDK format") + func matchesPythonFormat() throws { + // Python SDK test uses CancelledNotificationParams with requestId and reason + let params = CancelledNotification.Parameters( + requestId: .string("first-request-id"), + reason: "Testing server recovery" + ) + let notification = CancelledNotification.message(params) + + let encoder = JSONEncoder() + let data = try encoder.encode(notification) + let json = try JSONDecoder().decode([String: Value].self, from: data) + + // Verify structure matches Python SDK expectations + #expect(json["jsonrpc"] == "2.0") + #expect(json["method"] == "notifications/cancelled") + + let notificationParams = json["params"]?.objectValue + #expect(notificationParams != nil) + #expect(notificationParams?["requestId"] == .string("first-request-id")) + #expect(notificationParams?["reason"] == .string("Testing server recovery")) + } + } + + // MARK: - Integration Tests + + @Suite("Cancellation integration") + struct CancellationIntegrationTests { + + /// Actor to track received cancellation notifications + private actor CancellationTracker { + private var cancellations: [CancelledNotification.Parameters] = [] + + func add(_ params: CancelledNotification.Parameters) { + cancellations.append(params) + } + + var count: Int { cancellations.count } + var all: [CancelledNotification.Parameters] { cancellations } + } + + /// Test that client can send a CancelledNotification to the server. + /// + /// This tests the basic notification flow from client to server. + @Test(.timeLimit(.minutes(1))) + func clientSendsCancelledNotificationToServer() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.cancellation") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let cancellationTracker = CancellationTracker() + + // Set up server with cancellation notification handler + let server = Server( + name: "CancellationTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.onNotification(CancelledNotification.self) { message in + await cancellationTracker.add(message.params) + } + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "test_tool", inputSchema: ["type": "object"]) + ]) + } + + let client = Client(name: "CancellationTestClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Send a cancellation notification + let cancelParams = CancelledNotification.Parameters( + requestId: .string("req-to-cancel"), + reason: "User requested cancellation" + ) + try await client.notify(CancelledNotification.message(cancelParams)) + + // Give time for notification to be processed + try await Task.sleep(for: .milliseconds(100)) + + // Verify server received the cancellation notification + let count = await cancellationTracker.count + #expect(count == 1, "Server should receive exactly one cancellation notification") + + let cancellations = await cancellationTracker.all + if let first = cancellations.first { + #expect(first.requestId == .string("req-to-cancel")) + #expect(first.reason == "User requested cancellation") + } + } + + /// Test that server remains functional after receiving a cancellation notification. + /// + /// This is based on Python SDK's test_server_remains_functional_after_cancel test. + /// The key insight is that cancellation notifications should not break the server's + /// ability to handle subsequent requests. + @Test(.timeLimit(.minutes(1))) + func serverRemainsFunctionalAfterCancel() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.cancellation.recovery") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let callCounter = CallCounter() + let cancellationTracker = CancellationTracker() + + // Set up server + let server = Server( + name: "RecoveryTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.onNotification(CancelledNotification.self) { message in + await cancellationTracker.add(message.params) + } + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool( + name: "test_tool", + description: "A tool for testing cancellation recovery", + inputSchema: ["type": "object"] + ) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, _ in + guard request.name == "test_tool" else { + return CallTool.Result(content: [.text("Unknown tool")], isError: true) + } + let count = await callCounter.increment() + return CallTool.Result(content: [.text("Call number: \(count)")]) + } + + let client = Client(name: "RecoveryTestClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // First tool call - should succeed + let result1 = try await client.send( + CallTool.request(.init(name: "test_tool", arguments: [:])) + ) + if case .text(let text, _, _) = result1.content.first { + #expect(text == "Call number: 1") + } else { + Issue.record("Expected text content for first call") + } + + // Send cancellation notification (simulating a cancelled request) + let cancelParams = CancelledNotification.Parameters( + requestId: .string("some-cancelled-request"), + reason: "Testing server recovery" + ) + try await client.notify(CancelledNotification.message(cancelParams)) + + // Give time for notification to be processed + try await Task.sleep(for: .milliseconds(50)) + + // Verify cancellation was received + let cancellationCount = await cancellationTracker.count + #expect(cancellationCount == 1, "Server should have received the cancellation") + + // Second tool call - should also succeed (server recovered) + let result2 = try await client.send( + CallTool.request(.init(name: "test_tool", arguments: [:])) + ) + if case .text(let text, _, _) = result2.content.first { + #expect(text == "Call number: 2") + } else { + Issue.record("Expected text content for second call") + } + + // Verify call count + let finalCount = await callCounter.count + #expect(finalCount == 2, "Both tool calls should have been processed") + } + + /// Test that server can send a cancellation notification to the client. + /// + /// The server may need to cancel a pending request (e.g., if it takes too long + /// or if the server is shutting down). + @Test(.timeLimit(.minutes(1))) + func serverSendsCancelledNotificationToClient() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.cancellation.server-to-client") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let clientCancellationReceived = ClientCancellationTracker() + + let server = Server( + name: "ServerCancelTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "trigger_cancel", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, context in + guard request.name == "trigger_cancel" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + // Server sends a cancellation notification for some other request + // This simulates the server cancelling a client's pending request + try await context.sendMessage(CancelledNotification.message(.init( + requestId: .string("client-pending-request"), + reason: "Server is cancelling this request" + ))) + + return CallTool.Result(content: [.text("Cancel notification sent")]) + } + + let client = Client(name: "ServerCancelTestClient", version: "1.0") + + // Register client handler for cancellation notifications + await client.onNotification(CancelledNotification.self) { message in + await clientCancellationReceived.add(message.params) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Call the tool that triggers a cancellation notification + let result = try await client.send( + CallTool.request(.init(name: "trigger_cancel", arguments: [:])) + ) + + // Verify tool completed + if case .text(let text, _, _) = result.content.first { + #expect(text == "Cancel notification sent") + } + + // Give time for notification to arrive + try await Task.sleep(for: .milliseconds(100)) + + // Verify client received the cancellation + let count = await clientCancellationReceived.count + #expect(count == 1, "Client should receive the cancellation notification") + + let cancellations = await clientCancellationReceived.all + if let first = cancellations.first { + #expect(first.requestId == .string("client-pending-request")) + #expect(first.reason == "Server is cancelling this request") + } + } + + /// Test multiple cancellation notifications can be processed. + @Test(.timeLimit(.minutes(1))) + func multipleCancellationNotifications() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.cancellation.multiple") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let cancellationTracker = CancellationTracker() + + let server = Server( + name: "MultipleCancelServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.onNotification(CancelledNotification.self) { message in + await cancellationTracker.add(message.params) + } + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: []) + } + + let client = Client(name: "MultipleCancelClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Send multiple cancellation notifications + for i in 1...5 { + let cancelParams = CancelledNotification.Parameters( + requestId: .string("req-\(i)"), + reason: "Cancellation \(i)" + ) + try await client.notify(CancelledNotification.message(cancelParams)) + } + + // Give time for all notifications to be processed + try await Task.sleep(for: .milliseconds(200)) + + // Verify all cancellations were received + let count = await cancellationTracker.count + #expect(count == 5, "Server should receive all 5 cancellation notifications") + + let cancellations = await cancellationTracker.all + for i in 1...5 { + let expected = CancelledNotification.Parameters( + requestId: .string("req-\(i)"), + reason: "Cancellation \(i)" + ) + #expect( + cancellations.contains { $0.requestId == expected.requestId }, + "Should contain cancellation for req-\(i)" + ) + } + } + + /// Test cancellation notification with no requestId (protocol 2025-11-25+). + /// + /// In newer protocol versions, the requestId is optional. + @Test(.timeLimit(.minutes(1))) + func cancellationWithoutRequestId() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.cancellation.no-request-id") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let cancellationTracker = CancellationTracker() + + let server = Server( + name: "NoRequestIdCancelServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.onNotification(CancelledNotification.self) { message in + await cancellationTracker.add(message.params) + } + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: []) + } + + let client = Client(name: "NoRequestIdCancelClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Send cancellation without requestId (general cancellation) + let cancelParams = CancelledNotification.Parameters( + reason: "General operation cancellation" + ) + try await client.notify(CancelledNotification.message(cancelParams)) + + // Give time for notification to be processed + try await Task.sleep(for: .milliseconds(100)) + + // Verify cancellation was received + let count = await cancellationTracker.count + #expect(count == 1) + + let cancellations = await cancellationTracker.all + if let first = cancellations.first { + #expect(first.requestId == nil, "requestId should be nil") + #expect(first.reason == "General operation cancellation") + } + } + } + + // MARK: - Server Context Cancellation Tests + + @Suite("Server.Context cancellation") + struct ServerContextCancellationTests { + + /// Test that Server.Context.sendCancelled works correctly. + /// + /// The Server.Context has a sendCancelled convenience method for sending + /// cancellation notifications. + @Test(.timeLimit(.minutes(1))) + func serverContextSendCancelled() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.cancellation.context") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let clientCancellationReceived = ClientCancellationTracker() + + let server = Server( + name: "ContextCancelServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "cancel_via_context", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, context in + guard request.name == "cancel_via_context" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + // Use the context's sendMessage to send a CancelledNotification + try await context.sendMessage(CancelledNotification.message(.init( + requestId: .string("ctx-cancel-request"), + reason: "Cancelled via server context" + ))) + + return CallTool.Result(content: [.text("Cancellation sent via context")]) + } + + let client = Client(name: "ContextCancelClient", version: "1.0") + + await client.onNotification(CancelledNotification.self) { message in + await clientCancellationReceived.add(message.params) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Call the tool + let result = try await client.send( + CallTool.request(.init(name: "cancel_via_context", arguments: [:])) + ) + + if case .text(let text, _, _) = result.content.first { + #expect(text == "Cancellation sent via context") + } + + // Give time for notification + try await Task.sleep(for: .milliseconds(100)) + + // Verify client received cancellation + let count = await clientCancellationReceived.count + #expect(count == 1) + + let cancellations = await clientCancellationReceived.all + if let first = cancellations.first { + #expect(first.requestId == .string("ctx-cancel-request")) + #expect(first.reason == "Cancelled via server context") + } + } + } + + // MARK: - Client Task Cancellation Tests + + @Suite("Client Swift Task cancellation") + struct ClientTaskCancellationTests { + + /// Actor to track received cancellation notifications + private actor CancellationTracker { + private var cancellations: [CancelledNotification.Parameters] = [] + + func add(_ params: CancelledNotification.Parameters) { + cancellations.append(params) + } + + var count: Int { cancellations.count } + var all: [CancelledNotification.Parameters] { cancellations } + } + + /// Test that cancelling a Swift Task that's waiting for a response properly cleans up. + /// + /// This mirrors the TypeScript SDK's AbortController behavior - when the client + /// cancels the Task waiting for a response, the pending request is cleaned up + /// and an appropriate error is thrown. + @Test(.timeLimit(.minutes(1))) + func clientTaskCancellationCleansUpPendingRequest() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.cancellation.task-cancel") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let toolCallStarted = ToolCallStartedTracker() + + let server = Server( + name: "TaskCancelServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "slow_tool", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, _ in + guard request.name == "slow_tool" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + // Signal that the tool call has started + await toolCallStarted.markStarted() + + // Simulate a slow operation - this should be interrupted by task cancellation + do { + try await Task.sleep(for: .seconds(10)) + return CallTool.Result(content: [.text("Completed")]) + } catch { + // Task was cancelled - this is expected + return CallTool.Result(content: [.text("Cancelled")], isError: true) + } + } + + let client = Client(name: "TaskCancelClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Start a tool call in a separate Task that we can cancel + let callTask = Task { + try await client.send( + CallTool.request(.init(name: "slow_tool", arguments: [:])) + ) + } + + // Wait for the tool call to start + try await toolCallStarted.waitForStart() + + // Cancel the Task + callTask.cancel() + + // Verify the task throws an error (CancellationError, connection closed, or no response) + do { + _ = try await callTask.value + Issue.record("Expected task to throw an error when cancelled") + } catch is CancellationError { + // Expected - Swift Task cancellation + } catch let error as MCPError { + // Also expected - connection closed or no response received + // When a Task is cancelled, the pending request stream is terminated + // which can result in "No response received" or "connectionClosed" + let errorDescription = String(describing: error) + #expect( + error == .connectionClosed || + errorDescription.contains("cancel") || + errorDescription.contains("No response received"), + "Error should be related to cancellation or no response: \(error)" + ) + } + + // Verify client is still functional after cancellation + // List tools should still work + let tools = try await client.send(ListTools.request(.init())) + #expect(tools.tools.count == 1) + #expect(tools.tools.first?.name == "slow_tool") + } + + /// Test that multiple concurrent requests can be individually cancelled. + @Test(.timeLimit(.minutes(1))) + func multipleConcurrentRequestsCancellation() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.cancellation.concurrent") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let server = Server( + name: "ConcurrentCancelServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "variable_tool", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, _ in + guard request.name == "variable_tool" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + // Delay based on the "delay" argument + let delay = request.arguments?["delay"]?.doubleValue ?? 1.0 + try? await Task.sleep(for: .seconds(delay)) + return CallTool.Result(content: [.text("Done after \(delay)s")]) + } + + let client = Client(name: "ConcurrentCancelClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Start two concurrent requests + let fastTask = Task { + try await client.send( + CallTool.request(.init(name: "variable_tool", arguments: ["delay": .double(0.1)])) + ) + } + + let slowTask = Task { + try await client.send( + CallTool.request(.init(name: "variable_tool", arguments: ["delay": .double(10.0)])) + ) + } + + // Wait a bit for both requests to start + try await Task.sleep(for: .milliseconds(50)) + + // Cancel only the slow task + slowTask.cancel() + + // Fast task should complete successfully + let fastResult = try await fastTask.value + if case .text(let text, _, _) = fastResult.content.first { + #expect(text.contains("0.1")) + } + + // Slow task should be cancelled + do { + _ = try await slowTask.value + Issue.record("Slow task should have been cancelled") + } catch { + // Expected - task was cancelled + } + } + + /// Test that when a client Task is cancelled, the client sends a CancelledNotification to the server. + /// + /// This is per MCP spec: "When a party wants to cancel an in-progress request, + /// it sends a `notifications/cancelled` notification" + /// + /// This mirrors the TypeScript SDK's behavior where AbortSignal abort triggers + /// sending notifications/cancelled. + @Test(.timeLimit(.minutes(1))) + func clientTaskCancellationSendsCancelledNotification() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.cancellation.client-sends-notification") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let cancellationReceived = CancellationTracker() + let handlerStarted = ToolCallStartedTracker() + + let server = Server( + name: "ClientCancellationNotificationServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + // Track cancellation notifications received by the server + await server.onNotification(CancelledNotification.self) { message in + await cancellationReceived.add(message.params) + } + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "slow_tool", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, _ in + guard request.name == "slow_tool" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + await handlerStarted.markStarted() + + // Slow operation + try? await Task.sleep(for: .seconds(10)) + return CallTool.Result(content: [.text("Completed")]) + } + + let client = Client(name: "ClientCancellationNotificationClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Create a request with a known ID + let knownRequestId = RequestId.string("client-will-cancel-this") + let request = Request( + id: knownRequestId, + method: CallTool.name, + params: CallTool.Parameters(name: "slow_tool", arguments: [:]) + ) + + // Start the request in a Task we can cancel + let callTask = Task { + try await client.send(request) + } + + // Wait for handler to start + try await handlerStarted.waitForStart() + + // Cancel the client Task - this should trigger sending CancelledNotification + callTask.cancel() + + // Give time for cancellation notification to be sent and processed + try await Task.sleep(for: .milliseconds(200)) + + // Verify server received the cancellation notification + let count = await cancellationReceived.count + #expect(count >= 1, "Server should receive a cancellation notification when client Task is cancelled") + + let cancellations = await cancellationReceived.all + let hasCancellationForRequest = cancellations.contains { $0.requestId == knownRequestId } + #expect(hasCancellationForRequest, "Server should receive cancellation for the specific request ID") + } + } + + // MARK: - Protocol-Level Cancellation Tests + + @Suite("Protocol-level cancellation (CancelledNotification aborts in-flight handlers)") + struct ProtocolLevelCancellationTests { + + /// Test that when a CancelledNotification is received, the in-flight request handler + /// is cancelled and no response is sent. + /// + /// This mirrors the Python SDK's `test_server_remains_functional_after_cancel` test. + @Test(.timeLimit(.minutes(1))) + func serverCancelsInFlightRequestOnNotification() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.cancellation.protocol-level") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let handlerStarted = ToolCallStartedTracker() + let handlerCompleted = HandlerCompletionTracker() + + let server = Server( + name: "ProtocolCancelServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "slow_tool", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, _ in + guard request.name == "slow_tool" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + // Signal that the handler has started + await handlerStarted.markStarted() + + // This is a slow operation that should be cancelled + do { + try await Task.sleep(for: .seconds(10)) + await handlerCompleted.markCompleted() + return CallTool.Result(content: [.text("Completed")]) + } catch is CancellationError { + // Expected - handler was cancelled + await handlerCompleted.markCancelled() + throw CancellationError() + } + } + + let client = Client(name: "ProtocolCancelClient", version: "1.0") + + try await server.start(transport: serverTransport) + let initResult = try await client.connect(transport: clientTransport) + #expect(initResult.serverInfo.name == "ProtocolCancelServer") + + // Create a request with a known ID so we can cancel it + let knownRequestId = RequestId.string("test-request-to-cancel") + let toolCallRequest = Request( + id: knownRequestId, + method: CallTool.name, + params: CallTool.Parameters(name: "slow_tool", arguments: [:]) + ) + + // Start a slow tool call in a separate Task + let callTask = Task { + try await client.send(toolCallRequest) + } + + // Wait for the handler to start + try await handlerStarted.waitForStart() + + // Send cancellation notification from client to server with the known request ID + try await client.notify(CancelledNotification.message(.init( + requestId: knownRequestId, + reason: "Test cancellation" + ))) + + // Give time for cancellation to propagate + try await Task.sleep(for: .milliseconds(100)) + + // The call task should eventually fail or hang waiting for response + // Cancel it from the client side as well to clean up + callTask.cancel() + + // Verify the handler was cancelled (not completed normally) + let wasCompleted = await handlerCompleted.wasCompleted + let wasCancelled = await handlerCompleted.wasCancelled + #expect(!wasCompleted, "Handler should not have completed normally") + #expect(wasCancelled, "Handler should have been cancelled") + + // Verify server is still functional after cancellation + let tools = try await client.send(ListTools.request(.init())) + #expect(tools.tools.count == 1) + #expect(tools.tools.first?.name == "slow_tool") + } + + /// Test that response is suppressed when Task.isCancelled is true + /// even if the handler completes normally. + @Test(.timeLimit(.minutes(1))) + func serverSuppressesResponseWhenCancelled() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.cancellation.suppress-response") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let handlerStarted = ToolCallStartedTracker() + + let server = Server( + name: "SuppressResponseServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "quick_tool", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, _ in + guard request.name == "quick_tool" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + await handlerStarted.markStarted() + + // Small delay to allow cancellation to arrive + try? await Task.sleep(for: .milliseconds(50)) + + // Handler completes normally, but response should be suppressed + // if Task.isCancelled is true + return CallTool.Result(content: [.text("Should be suppressed")]) + } + + let client = Client(name: "SuppressResponseClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Create a request with a known ID so we can cancel it + let knownRequestId = RequestId.string("suppress-response-request") + let toolCallRequest = Request( + id: knownRequestId, + method: CallTool.name, + params: CallTool.Parameters(name: "quick_tool", arguments: [:]) + ) + + // Start a tool call + let callTask = Task { + try await client.send(toolCallRequest) + } + + // Wait for handler to start + try await handlerStarted.waitForStart() + + // Send cancellation immediately + try await client.notify(CancelledNotification.message(.init( + requestId: knownRequestId, + reason: "Suppress response test" + ))) + + // Cancel the client task as well since no response will come + try await Task.sleep(for: .milliseconds(100)) + callTask.cancel() + + // The test passes if we get here - the server didn't crash + // and handled the cancellation gracefully + } + + /// Test that server shutdown cancels all in-flight handlers. + @Test(.timeLimit(.minutes(1))) + func serverShutdownCancelsInFlightHandlers() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.cancellation.shutdown") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let handlerStarted = ToolCallStartedTracker() + let handlerCompleted = HandlerCompletionTracker() + + let server = Server( + name: "ShutdownCancelServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "very_slow_tool", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, _ in + guard request.name == "very_slow_tool" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + await handlerStarted.markStarted() + + do { + try await Task.sleep(for: .seconds(60)) + await handlerCompleted.markCompleted() + return CallTool.Result(content: [.text("Completed")]) + } catch is CancellationError { + await handlerCompleted.markCancelled() + throw CancellationError() + } + } + + let client = Client(name: "ShutdownCancelClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Start a very slow tool call + let callTask = Task { + try await client.send( + CallTool.request(.init(name: "very_slow_tool", arguments: [:])) + ) + } + + // Wait for handler to start + try await handlerStarted.waitForStart() + + // Stop the server while the handler is running + await server.stop() + + // Give time for cancellation to propagate + try await Task.sleep(for: .milliseconds(100)) + + // Verify the handler was cancelled + let wasCancelled = await handlerCompleted.wasCancelled + #expect(wasCancelled, "Handler should have been cancelled on shutdown") + + // Clean up the client task + callTask.cancel() + } + } +} + +// MARK: - Request Timeout Tests + +@Suite("Request timeout") +struct RequestTimeoutTests { + + /// Test that request timeout triggers cancellation and throws the correct error. + @Test(.timeLimit(.minutes(1))) + func requestTimeoutTriggersError() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.timeout") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let server = Server( + name: "TimeoutServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "slow_tool", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, _ in + guard request.name == "slow_tool" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + // Simulate a slow operation that takes longer than the timeout + try? await Task.sleep(for: .seconds(10)) + return CallTool.Result(content: [.text("Completed")]) + } + + let client = Client(name: "TimeoutClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Send request with a short timeout + do { + _ = try await client.send( + CallTool.request(.init(name: "slow_tool", arguments: [:])), + options: .init(timeout: .milliseconds(100)) + ) + Issue.record("Expected request to timeout") + } catch let error as MCPError { + // Verify we get a requestTimeout error + if case .requestTimeout(let timeout, let message) = error { + #expect(timeout == .milliseconds(100)) + #expect(message?.contains("timed out") == true) + } else { + Issue.record("Expected MCPError.requestTimeout, got: \(error)") + } + } + + // Verify client is still functional after timeout + let tools = try await client.send(ListTools.request(.init())) + #expect(tools.tools.count == 1) + + await client.disconnect() + await server.stop() + } + + /// Test that request without timeout waits indefinitely (until completed). + @Test(.timeLimit(.minutes(1))) + func requestWithoutTimeoutWaitsForResponse() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.no-timeout") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let server = Server( + name: "NoTimeoutServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "fast_tool", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, _ in + guard request.name == "fast_tool" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + // Small delay, but should complete fine without timeout + try? await Task.sleep(for: .milliseconds(50)) + return CallTool.Result(content: [.text("Completed")]) + } + + let client = Client(name: "NoTimeoutClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Send request without timeout - should complete normally + let result = try await client.send( + CallTool.request(.init(name: "fast_tool", arguments: [:])), + options: nil + ) + + if case .text(let text, _, _) = result.content.first { + #expect(text == "Completed") + } else { + Issue.record("Expected text content") + } + + await client.disconnect() + await server.stop() + } + + /// Test that timeout sends CancelledNotification to server. + @Test(.timeLimit(.minutes(1))) + func timeoutSendsCancelledNotification() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.timeout-cancellation") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let cancellationReceived = CancellationReceivedTracker() + + let server = Server( + name: "CancellationTrackingServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + // Track cancellation notifications + await server.onNotification(CancelledNotification.self) { message in + await cancellationReceived.add(message.params) + } + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "slow_tool", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, _ in + guard request.name == "slow_tool" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + // Simulate a slow operation + try? await Task.sleep(for: .seconds(10)) + return CallTool.Result(content: [.text("Completed")]) + } + + let client = Client(name: "CancellationClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Send request with timeout + do { + _ = try await client.send( + CallTool.request(.init(name: "slow_tool", arguments: [:])), + options: .init(timeout: .milliseconds(100)) + ) + } catch { + // Expected timeout error + } + + // Give time for cancellation notification to be sent and processed + try await Task.sleep(for: .milliseconds(100)) + + // Verify cancellation notification was received + let count = await cancellationReceived.count + #expect(count >= 1, "Server should have received at least one cancellation notification") + + if count > 0 { + let cancellations = await cancellationReceived.all + let lastCancellation = cancellations.last! + #expect(lastCancellation.reason?.contains("timed out") == true) + } + + await client.disconnect() + await server.stop() + } + +} + +// MARK: - Helper Types + +/// Actor to track cancellation notifications received by the server +private actor CancellationReceivedTracker { + private var cancellations: [CancelledNotification.Parameters] = [] + + func add(_ params: CancelledNotification.Parameters) { + cancellations.append(params) + } + + var count: Int { cancellations.count } + var all: [CancelledNotification.Parameters] { cancellations } +} + +/// Actor to track when a tool call has started +private actor ToolCallStartedTracker { + private var started = false + + func markStarted() { + started = true + } + + func waitForStart() async throws { + while !started { + try await Task.sleep(for: .milliseconds(10)) + } + } +} + +/// Actor to track whether a handler completed or was cancelled +private actor HandlerCompletionTracker { + private var _completed = false + private var _cancelled = false + + func markCompleted() { + _completed = true + } + + func markCancelled() { + _cancelled = true + } + + var wasCompleted: Bool { _completed } + var wasCancelled: Bool { _cancelled } +} + +/// Actor to track tool call counts +private actor CallCounter { + private var _count = 0 + + func increment() -> Int { + _count += 1 + return _count + } + + var count: Int { _count } +} + +/// Actor to track cancellation notifications received by the client +private actor ClientCancellationTracker { + private var cancellations: [CancelledNotification.Parameters] = [] + + func add(_ params: CancelledNotification.Parameters) { + cancellations.append(params) + } + + var count: Int { cancellations.count } + var all: [CancelledNotification.Parameters] { cancellations } +} + +// MARK: - Client.cancelRequest() Tests + +/// Tests for the public `Client.cancelRequest(_:reason:)` API. +/// +/// This API allows explicit cancellation of in-flight requests by ID, +/// similar to TypeScript SDK's AbortController pattern. +@Suite("Client.cancelRequest() API") +struct ClientCancelRequestAPITests { + + /// Test that cancelRequest properly cancels an in-flight request and sends CancelledNotification. + /// + /// Based on Python SDK's test pattern where cancellation is sent for a specific request ID. + @Test(.timeLimit(.minutes(1))) + func cancelRequestSendsCancelledNotificationAndThrowsError() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.cancel-request-api") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let cancellationReceived = CancellationReceivedTracker() + let handlerStarted = ToolCallStartedTracker() + + let server = Server( + name: "CancelRequestAPIServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + // Track cancellation notifications + await server.onNotification(CancelledNotification.self) { message in + await cancellationReceived.add(message.params) + } + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "slow_tool", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, _ in + guard request.name == "slow_tool" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + await handlerStarted.markStarted() + + // Slow operation - should be interrupted + try? await Task.sleep(for: .seconds(10)) + return CallTool.Result(content: [.text("Completed")]) + } + + let client = Client(name: "CancelRequestAPIClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Create a request with a known ID + let knownRequestId = RequestId.string("api-cancel-test-\(UUID().uuidString)") + let request = Request( + id: knownRequestId, + method: CallTool.name, + params: CallTool.Parameters(name: "slow_tool", arguments: [:]) + ) + + // Start the request in a separate Task + let requestTask = Task { + try await client.send(request) + } + + // Wait for handler to start + try await handlerStarted.waitForStart() + + // Use the cancelRequest API to cancel the request + await client.cancelRequest(knownRequestId, reason: "Cancelled via cancelRequest API") + + // Give time for cancellation to propagate + try await Task.sleep(for: .milliseconds(100)) + + // Verify the request throws MCPError.requestCancelled + do { + _ = try await requestTask.value + Issue.record("Expected request to throw MCPError.requestCancelled") + } catch let error as MCPError { + // Should be requestCancelled error + if case .requestCancelled(let reason) = error { + #expect(reason == "Cancelled via cancelRequest API") + } else { + // May also be connectionClosed or similar if timing is different + // This is acceptable behavior + } + } catch is CancellationError { + // Also acceptable - Swift Task cancellation propagated + } + + // Verify server received the cancellation notification + let count = await cancellationReceived.count + #expect(count >= 1, "Server should receive at least one CancelledNotification") + + let cancellations = await cancellationReceived.all + let matchingCancellation = cancellations.first { $0.requestId == knownRequestId } + #expect(matchingCancellation != nil, "Should have received cancellation for the specific request ID") + #expect(matchingCancellation?.reason == "Cancelled via cancelRequest API") + + // Verify client is still functional after cancellation + let tools = try await client.send(ListTools.request(.init())) + #expect(tools.tools.count == 1) + } + + /// Test that cancelRequest for unknown request ID is a no-op (per MCP spec). + /// + /// Per MCP spec: "The receiver MUST NOT assume that the request will be cancelled; + /// it MAY still complete normally." + @Test(.timeLimit(.minutes(1))) + func cancelRequestForUnknownIdIsNoOp() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.cancel-unknown-request") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let cancellationReceived = CancellationReceivedTracker() + + let server = Server( + name: "CancelUnknownServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.onNotification(CancelledNotification.self) { message in + await cancellationReceived.add(message.params) + } + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: []) + } + + let client = Client(name: "CancelUnknownClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Cancel a request that doesn't exist + let unknownId = RequestId.string("non-existent-request") + await client.cancelRequest(unknownId, reason: "This request doesn't exist") + + // Give time for notification to be processed + try await Task.sleep(for: .milliseconds(100)) + + // The cancellation notification should still be sent (best effort) + // but the client should not crash + let count = await cancellationReceived.count + #expect(count >= 1, "Cancellation notification should still be sent") + + // Client should still be functional + let tools = try await client.send(ListTools.request(.init())) + #expect(tools.tools.isEmpty) + } + + /// Test that cancelRequest can be called without a reason. + @Test(.timeLimit(.minutes(1))) + func cancelRequestWithoutReason() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.cancel-no-reason") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let cancellationReceived = CancellationReceivedTracker() + + let server = Server( + name: "CancelNoReasonServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.onNotification(CancelledNotification.self) { message in + await cancellationReceived.add(message.params) + } + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: []) + } + + let client = Client(name: "CancelNoReasonClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Cancel with no reason + let requestId = RequestId.string("some-request") + await client.cancelRequest(requestId) // No reason provided + + try await Task.sleep(for: .milliseconds(100)) + + // Verify cancellation was sent + let count = await cancellationReceived.count + #expect(count >= 1) + + let cancellations = await cancellationReceived.all + if let first = cancellations.first { + #expect(first.requestId == requestId) + #expect(first.reason == nil) + } + } +} diff --git a/Tests/MCPTests/CapabilitiesTests.swift b/Tests/MCPTests/CapabilitiesTests.swift new file mode 100644 index 00000000..3f473fa9 --- /dev/null +++ b/Tests/MCPTests/CapabilitiesTests.swift @@ -0,0 +1,919 @@ +import Foundation +import Testing + +@testable import MCP + +// MARK: - Client Capabilities Encoding Tests + +@Suite("Client Capabilities Encoding Tests") +struct ClientCapabilitiesEncodingTests { + + @Test("Empty client capabilities encodes correctly") + func testEmptyClientCapabilities() throws { + let capabilities = Client.Capabilities() + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + + let data = try encoder.encode(capabilities) + let json = String(data: data, encoding: .utf8)! + + // Empty capabilities should encode to empty object + #expect(json == "{}") + + // Verify roundtrip + let decoder = JSONDecoder() + let decoded = try decoder.decode(Client.Capabilities.self, from: data) + #expect(decoded.sampling == nil) + #expect(decoded.elicitation == nil) + #expect(decoded.roots == nil) + #expect(decoded.experimental == nil) + #expect(decoded.tasks == nil) + } + + @Test("Client capabilities with roots encodes correctly") + func testClientCapabilitiesWithRoots() throws { + let capabilities = Client.Capabilities( + roots: .init(listChanged: true) + ) + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + + let data = try encoder.encode(capabilities) + let json = String(data: data, encoding: .utf8)! + + #expect(json.contains("\"roots\"")) + #expect(json.contains("\"listChanged\":true")) + + let decoder = JSONDecoder() + let decoded = try decoder.decode(Client.Capabilities.self, from: data) + #expect(decoded.roots?.listChanged == true) + } + + @Test("Client capabilities with experimental encodes correctly") + func testClientCapabilitiesWithExperimental() throws { + let capabilities = Client.Capabilities( + experimental: [ + "feature": [ + "enabled": .bool(true), + "count": .int(42) + ] + ] + ) + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + + let data = try encoder.encode(capabilities) + let json = String(data: data, encoding: .utf8)! + + #expect(json.contains("\"experimental\"")) + #expect(json.contains("\"feature\"")) + #expect(json.contains("\"enabled\":true")) + + let decoder = JSONDecoder() + let decoded = try decoder.decode(Client.Capabilities.self, from: data) + #expect(decoded.experimental?["feature"]?["enabled"] == .bool(true)) + #expect(decoded.experimental?["feature"]?["count"] == .int(42)) + } + + @Test("Client capabilities all fields roundtrip") + func testClientCapabilitiesAllFieldsRoundtrip() throws { + let capabilities = Client.Capabilities( + sampling: .init(context: .init(), tools: .init()), + elicitation: .init(form: .init(applyDefaults: true), url: .init()), + experimental: ["test": ["value": .string("data")]], + roots: .init(listChanged: true) + ) + + 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?.context != nil) + #expect(decoded.sampling?.tools != nil) + #expect(decoded.elicitation?.form?.applyDefaults == true) + #expect(decoded.elicitation?.url != nil) + #expect(decoded.experimental?["test"]?["value"] == .string("data")) + #expect(decoded.roots?.listChanged == true) + } +} + +// MARK: - Server Capabilities Encoding Tests + +@Suite("Server Capabilities Encoding Tests") +struct ServerCapabilitiesEncodingTests { + + @Test("Empty server capabilities encodes correctly") + func testEmptyServerCapabilities() throws { + let capabilities = Server.Capabilities() + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + + let data = try encoder.encode(capabilities) + let json = String(data: data, encoding: .utf8)! + + // Empty capabilities should encode to empty object + #expect(json == "{}") + + let decoder = JSONDecoder() + let decoded = try decoder.decode(Server.Capabilities.self, from: data) + #expect(decoded.logging == nil) + #expect(decoded.prompts == nil) + #expect(decoded.resources == nil) + #expect(decoded.tools == nil) + #expect(decoded.completions == nil) + #expect(decoded.experimental == nil) + } + + @Test("Server capabilities with logging encodes correctly") + func testServerCapabilitiesWithLogging() throws { + let capabilities = Server.Capabilities( + logging: .init() + ) + + let encoder = JSONEncoder() + let data = try encoder.encode(capabilities) + let json = String(data: data, encoding: .utf8)! + + #expect(json.contains("\"logging\"")) + + let decoder = JSONDecoder() + let decoded = try decoder.decode(Server.Capabilities.self, from: data) + #expect(decoded.logging != nil) + } + + @Test("Server capabilities with prompts listChanged true") + func testServerCapabilitiesWithPromptsListChangedTrue() throws { + let capabilities = Server.Capabilities( + prompts: .init(listChanged: true) + ) + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + + let data = try encoder.encode(capabilities) + let json = String(data: data, encoding: .utf8)! + + #expect(json.contains("\"prompts\"")) + #expect(json.contains("\"listChanged\":true")) + + let decoder = JSONDecoder() + let decoded = try decoder.decode(Server.Capabilities.self, from: data) + #expect(decoded.prompts?.listChanged == true) + } + + @Test("Server capabilities with prompts listChanged false") + func testServerCapabilitiesWithPromptsListChangedFalse() throws { + let capabilities = Server.Capabilities( + prompts: .init(listChanged: false) + ) + + let encoder = JSONEncoder() + let data = try encoder.encode(capabilities) + + let decoder = JSONDecoder() + let decoded = try decoder.decode(Server.Capabilities.self, from: data) + #expect(decoded.prompts?.listChanged == false) + } + + @Test("Server capabilities with resources encodes correctly") + func testServerCapabilitiesWithResources() throws { + let capabilities = Server.Capabilities( + resources: .init(subscribe: true, listChanged: true) + ) + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + + let data = try encoder.encode(capabilities) + let json = String(data: data, encoding: .utf8)! + + #expect(json.contains("\"resources\"")) + #expect(json.contains("\"subscribe\":true")) + #expect(json.contains("\"listChanged\":true")) + + let decoder = JSONDecoder() + let decoded = try decoder.decode(Server.Capabilities.self, from: data) + #expect(decoded.resources?.subscribe == true) + #expect(decoded.resources?.listChanged == true) + } + + @Test("Server capabilities with tools listChanged") + func testServerCapabilitiesWithToolsListChanged() throws { + let capabilities = Server.Capabilities( + tools: .init(listChanged: true) + ) + + let encoder = JSONEncoder() + let data = try encoder.encode(capabilities) + + let decoder = JSONDecoder() + let decoded = try decoder.decode(Server.Capabilities.self, from: data) + #expect(decoded.tools?.listChanged == true) + } + + @Test("Server capabilities with completions encodes correctly") + func testServerCapabilitiesWithCompletions() throws { + let capabilities = Server.Capabilities( + completions: .init() + ) + + let encoder = JSONEncoder() + let data = try encoder.encode(capabilities) + let json = String(data: data, encoding: .utf8)! + + #expect(json.contains("\"completions\"")) + + let decoder = JSONDecoder() + let decoded = try decoder.decode(Server.Capabilities.self, from: data) + #expect(decoded.completions != nil) + } + + @Test("Server capabilities with experimental encodes correctly") + func testServerCapabilitiesWithExperimental() throws { + let capabilities = Server.Capabilities( + experimental: [ + "customFeature": [ + "supported": .bool(true) + ] + ] + ) + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + + let data = try encoder.encode(capabilities) + let json = String(data: data, encoding: .utf8)! + + #expect(json.contains("\"experimental\"")) + #expect(json.contains("\"customFeature\"")) + + let decoder = JSONDecoder() + let decoded = try decoder.decode(Server.Capabilities.self, from: data) + #expect(decoded.experimental?["customFeature"]?["supported"] == .bool(true)) + } + + @Test("Server capabilities all fields roundtrip") + func testServerCapabilitiesAllFieldsRoundtrip() throws { + let capabilities = Server.Capabilities( + logging: .init(), + prompts: .init(listChanged: true), + resources: .init(subscribe: true, listChanged: true), + tools: .init(listChanged: false), + completions: .init(), + experimental: ["test": ["enabled": .bool(true)]] + ) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(capabilities) + let decoded = try decoder.decode(Server.Capabilities.self, from: data) + + #expect(decoded.logging != nil) + #expect(decoded.prompts?.listChanged == true) + #expect(decoded.resources?.subscribe == true) + #expect(decoded.resources?.listChanged == true) + #expect(decoded.tools?.listChanged == false) + #expect(decoded.completions != nil) + #expect(decoded.experimental?["test"]?["enabled"] == .bool(true)) + } +} + +// MARK: - Initialize Request Encoding Tests + +@Suite("Initialize Request Encoding Tests") +struct InitializeRequestEncodingTests { + + @Test("Initialize parameters encodes with capabilities") + func testInitializeParametersEncoding() throws { + let params = Initialize.Parameters( + protocolVersion: Version.latest, + capabilities: Client.Capabilities( + sampling: .init(tools: .init()), + roots: .init(listChanged: true) + ), + clientInfo: Client.Info(name: "TestClient", version: "1.0.0") + ) + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + + let data = try encoder.encode(params) + let json = String(data: data, encoding: .utf8)! + + #expect(json.contains("\"protocolVersion\":\"\(Version.latest)\"")) + #expect(json.contains("\"clientInfo\"")) + #expect(json.contains("\"name\":\"TestClient\"")) + #expect(json.contains("\"capabilities\"")) + #expect(json.contains("\"sampling\"")) + #expect(json.contains("\"roots\"")) + } + + @Test("Initialize parameters decodes correctly") + func testInitializeParametersDecoding() throws { + let json = """ + { + "protocolVersion": "2025-11-25", + "capabilities": { + "sampling": {"tools": {}}, + "roots": {"listChanged": true} + }, + "clientInfo": { + "name": "TestClient", + "version": "1.0.0" + } + } + """ + + let decoder = JSONDecoder() + let params = try decoder.decode(Initialize.Parameters.self, from: json.data(using: .utf8)!) + + #expect(params.protocolVersion == Version.v2025_11_25) + #expect(params.capabilities.sampling?.tools != nil) + #expect(params.capabilities.roots?.listChanged == true) + #expect(params.clientInfo.name == "TestClient") + #expect(params.clientInfo.version == "1.0.0") + } + + @Test("Initialize parameters defaults when fields missing") + func testInitializeParametersDefaults() throws { + let json = "{}" + + let decoder = JSONDecoder() + let params = try decoder.decode(Initialize.Parameters.self, from: json.data(using: .utf8)!) + + // Should use defaults + #expect(params.protocolVersion == Version.latest) + #expect(params.clientInfo.name == "unknown") + #expect(params.clientInfo.version == "0.0.0") + } + + @Test("Initialize result encodes with server capabilities") + func testInitializeResultEncoding() throws { + let result = Initialize.Result( + protocolVersion: Version.latest, + capabilities: Server.Capabilities( + logging: .init(), + prompts: .init(listChanged: true), + resources: .init(subscribe: true, listChanged: true), + tools: .init(listChanged: false) + ), + serverInfo: Server.Info(name: "TestServer", version: "2.0.0"), + instructions: "Server instructions." + ) + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + + let data = try encoder.encode(result) + let json = String(data: data, encoding: .utf8)! + + #expect(json.contains("\"protocolVersion\":\"\(Version.latest)\"")) + #expect(json.contains("\"serverInfo\"")) + #expect(json.contains("\"name\":\"TestServer\"")) + #expect(json.contains("\"instructions\":\"Server instructions.\"")) + #expect(json.contains("\"capabilities\"")) + #expect(json.contains("\"logging\"")) + #expect(json.contains("\"prompts\"")) + #expect(json.contains("\"resources\"")) + #expect(json.contains("\"tools\"")) + } + + @Test("Initialize result decodes correctly") + func testInitializeResultDecoding() throws { + let json = """ + { + "protocolVersion": "2025-11-25", + "capabilities": { + "logging": {}, + "prompts": {"listChanged": true}, + "resources": {"subscribe": true, "listChanged": true}, + "tools": {"listChanged": false} + }, + "serverInfo": { + "name": "TestServer", + "version": "2.0.0" + }, + "instructions": "Server instructions." + } + """ + + let decoder = JSONDecoder() + let result = try decoder.decode(Initialize.Result.self, from: json.data(using: .utf8)!) + + #expect(result.protocolVersion == Version.v2025_11_25) + #expect(result.capabilities.logging != nil) + #expect(result.capabilities.prompts?.listChanged == true) + #expect(result.capabilities.resources?.subscribe == true) + #expect(result.capabilities.resources?.listChanged == true) + #expect(result.capabilities.tools?.listChanged == false) + #expect(result.serverInfo.name == "TestServer") + #expect(result.serverInfo.version == "2.0.0") + #expect(result.instructions == "Server instructions.") + } + + @Test("Initialize result roundtrip") + func testInitializeResultRoundtrip() throws { + let original = Initialize.Result( + protocolVersion: Version.latest, + capabilities: Server.Capabilities( + logging: .init(), + prompts: .init(listChanged: true), + resources: .init(subscribe: true, listChanged: true), + tools: .init(listChanged: false), + completions: .init() + ), + serverInfo: Server.Info( + name: "TestServer", + version: "2.0.0", + title: "Test Server Title", + description: "A test server" + ), + instructions: "Follow these instructions." + ) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(original) + let decoded = try decoder.decode(Initialize.Result.self, from: data) + + #expect(decoded.protocolVersion == original.protocolVersion) + #expect(decoded.capabilities.logging != nil) + #expect(decoded.capabilities.prompts?.listChanged == true) + #expect(decoded.capabilities.resources?.subscribe == true) + #expect(decoded.capabilities.tools?.listChanged == false) + #expect(decoded.capabilities.completions != nil) + #expect(decoded.serverInfo.name == original.serverInfo.name) + #expect(decoded.serverInfo.version == original.serverInfo.version) + #expect(decoded.serverInfo.title == original.serverInfo.title) + #expect(decoded.serverInfo.description == original.serverInfo.description) + #expect(decoded.instructions == original.instructions) + } +} + +// MARK: - Capability Negotiation Integration Tests + +@Suite("Capability Negotiation Integration Tests") +struct CapabilityNegotiationTests { + + @Test("Client sends capabilities to server during initialization") + func testClientSendsCapabilitiesToServer() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + // Set up server with specific capabilities + let server = Server( + name: "CapabilityTestServer", + version: "1.0.0", + capabilities: .init( + logging: .init(), + prompts: .init(listChanged: true), + tools: .init() + ) + ) + + // Register a tools handler + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: []) + } + + try await server.start(transport: serverTransport) + + // Set up client with specific capabilities + let client = Client( + name: "CapabilityTestClient", + version: "1.0.0" + ) + + // Set capabilities before connecting + await client.setCapabilities(.init( + sampling: .init(tools: .init()), + roots: .init(listChanged: true) + )) + + // Connect and verify + try await client.connect(transport: clientTransport) + + // Verify the server is running correctly + let tools = try await client.listTools() + // Just verify the connection works - server has no tools registered + #expect(tools.tools.isEmpty) + + await client.disconnect() + await server.stop() + } + + @Test("Server responds with its capabilities") + func testServerRespondsWithCapabilities() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + // Set up server with all capabilities + let server = Server( + name: "FullCapabilityServer", + version: "1.0.0", + capabilities: .init( + logging: .init(), + prompts: .init(listChanged: true), + resources: .init(subscribe: true, listChanged: true), + tools: .init(listChanged: true), + completions: .init() + ) + ) + + // Register handlers for capabilities that require them + await server.withRequestHandler(ListPrompts.self) { _, _ in + ListPrompts.Result(prompts: []) + } + await server.withRequestHandler(ListResources.self) { _, _ in + ListResources.Result(resources: []) + } + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: []) + } + + try await server.start(transport: serverTransport) + + // Set up client + let client = Client( + name: "TestClient", + version: "1.0.0" + ) + + try await client.connect(transport: clientTransport) + + // Verify we can use the capabilities + let prompts = try await client.listPrompts() + #expect(prompts.prompts.isEmpty) + + let resources = try await client.listResources() + #expect(resources.resources.isEmpty) + + let tools = try await client.listTools() + #expect(tools.tools.isEmpty) + + await client.disconnect() + await server.stop() + } + + @Test("Client in strict mode fails on missing capability") + func testStrictModeFailsOnMissingCapability() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + // Server without completions capability + let server = Server( + name: "LimitedServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: []) + } + + try await server.start(transport: serverTransport) + + // Client in strict mode + let client = Client( + name: "StrictClient", + version: "1.0.0", + configuration: .strict + ) + + try await client.connect(transport: clientTransport) + + // Attempting to use completions should fail + do { + let _ = try await client.complete( + ref: .prompt(PromptReference(name: "test")), + argument: CompletionArgument(name: "arg", value: "val") + ) + #expect(Bool(false), "Should have thrown error") + } catch { + // Expected to fail + #expect(error is MCPError) + } + + await client.disconnect() + await server.stop() + } + + @Test("getServerCapabilities returns nil before connect") + func testGetServerCapabilitiesReturnsNilBeforeConnect() async throws { + let client = Client( + name: "TestClient", + version: "1.0.0" + ) + + // Before connecting, server capabilities should be nil + let capabilities = await client.getServerCapabilities() + #expect(capabilities == nil) + } + + @Test("getServerCapabilities returns capabilities after connect") + func testGetServerCapabilitiesReturnsCapabilitiesAfterConnect() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + // Set up server with specific capabilities + let server = Server( + name: "CapabilityServer", + version: "1.0.0", + capabilities: .init( + logging: .init(), + prompts: .init(listChanged: true), + resources: .init(subscribe: true, listChanged: false), + tools: .init(listChanged: true) + ) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: []) + } + + try await server.start(transport: serverTransport) + + let client = Client( + name: "TestClient", + version: "1.0.0" + ) + + // Before connecting + let beforeCapabilities = await client.getServerCapabilities() + #expect(beforeCapabilities == nil) + + // Connect + try await client.connect(transport: clientTransport) + + // After connecting, should have server capabilities + let afterCapabilities = await client.getServerCapabilities() + #expect(afterCapabilities != nil) + #expect(afterCapabilities?.logging != nil) + #expect(afterCapabilities?.prompts?.listChanged == true) + #expect(afterCapabilities?.resources?.subscribe == true) + #expect(afterCapabilities?.resources?.listChanged == false) + #expect(afterCapabilities?.tools?.listChanged == true) + + await client.disconnect() + await server.stop() + } +} + +// MARK: - JSON Format Compatibility Tests + +@Suite("Capability JSON Format Compatibility Tests") +struct CapabilityJSONCompatibilityTests { + + @Test("Client capabilities matches TypeScript format") + func testClientCapabilitiesMatchesTypeScriptFormat() throws { + // TypeScript format: { "sampling": {}, "roots": { "listChanged": true } } + let typeScriptJSON = """ + {"sampling":{},"roots":{"listChanged":true}} + """ + + let decoder = JSONDecoder() + let capabilities = try decoder.decode( + Client.Capabilities.self, from: typeScriptJSON.data(using: .utf8)!) + + #expect(capabilities.sampling != nil) + #expect(capabilities.roots?.listChanged == true) + + // Encode and verify format + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + let data = try encoder.encode(capabilities) + let json = String(data: data, encoding: .utf8)! + + // Should match the TypeScript format + #expect(json.contains("\"sampling\":{}")) + #expect(json.contains("\"roots\":{\"listChanged\":true}")) + } + + @Test("Server capabilities matches TypeScript format") + func testServerCapabilitiesMatchesTypeScriptFormat() throws { + // TypeScript format from protocol.test.ts + let typeScriptJSON = """ + {"logging":{},"prompts":{"listChanged":true},"resources":{"subscribe":true},"tools":{"listChanged":false}} + """ + + let decoder = JSONDecoder() + let capabilities = try decoder.decode( + Server.Capabilities.self, from: typeScriptJSON.data(using: .utf8)!) + + #expect(capabilities.logging != nil) + #expect(capabilities.prompts?.listChanged == true) + #expect(capabilities.resources?.subscribe == true) + #expect(capabilities.tools?.listChanged == false) + } + + @Test("Client elicitation capability with form matches TypeScript format") + func testClientElicitationFormMatchesTypeScriptFormat() throws { + // TypeScript format: { "elicitation": { "form": {} } } + let typeScriptJSON = """ + {"elicitation":{"form":{}}} + """ + + let decoder = JSONDecoder() + let capabilities = try decoder.decode( + Client.Capabilities.self, from: typeScriptJSON.data(using: .utf8)!) + + #expect(capabilities.elicitation?.form != nil) + #expect(capabilities.elicitation?.url == nil) + } + + @Test("Client elicitation capability with form applyDefaults matches TypeScript format") + func testClientElicitationFormApplyDefaultsMatchesTypeScriptFormat() throws { + // TypeScript format: { "elicitation": { "form": { "applyDefaults": true } } } + let typeScriptJSON = """ + {"elicitation":{"form":{"applyDefaults":true}}} + """ + + let decoder = JSONDecoder() + let capabilities = try decoder.decode( + Client.Capabilities.self, from: typeScriptJSON.data(using: .utf8)!) + + #expect(capabilities.elicitation?.form?.applyDefaults == true) + } + + @Test("Client elicitation capability with url matches TypeScript format") + func testClientElicitationURLMatchesTypeScriptFormat() throws { + // TypeScript format: { "elicitation": { "url": {} } } + let typeScriptJSON = """ + {"elicitation":{"url":{}}} + """ + + let decoder = JSONDecoder() + let capabilities = try decoder.decode( + Client.Capabilities.self, from: typeScriptJSON.data(using: .utf8)!) + + #expect(capabilities.elicitation?.form == nil) + #expect(capabilities.elicitation?.url != nil) + } + + @Test("Client elicitation capability with both form and url matches TypeScript format") + func testClientElicitationBothMatchesTypeScriptFormat() throws { + // TypeScript format: { "elicitation": { "form": {}, "url": {} } } + let typeScriptJSON = """ + {"elicitation":{"form":{},"url":{}}} + """ + + let decoder = JSONDecoder() + let capabilities = try decoder.decode( + Client.Capabilities.self, from: typeScriptJSON.data(using: .utf8)!) + + #expect(capabilities.elicitation?.form != nil) + #expect(capabilities.elicitation?.url != nil) + } + + @Test("Initialize request matches Python format") + func testInitializeRequestMatchesPythonFormat() throws { + // Python format from test_session.py + let pythonJSON = """ + { + "protocolVersion": "2025-11-25", + "capabilities": { + "sampling": {} + }, + "clientInfo": { + "name": "mcp-client", + "version": "0.1.0" + } + } + """ + + let decoder = JSONDecoder() + let params = try decoder.decode(Initialize.Parameters.self, from: pythonJSON.data(using: .utf8)!) + + #expect(params.protocolVersion == Version.v2025_11_25) + #expect(params.capabilities.sampling != nil) + #expect(params.clientInfo.name == "mcp-client") + #expect(params.clientInfo.version == "0.1.0") + } + + @Test("Initialize result matches Python format") + func testInitializeResultMatchesPythonFormat() throws { + // Python format from test_session.py + let pythonJSON = """ + { + "protocolVersion": "2025-11-25", + "capabilities": { + "logging": {}, + "prompts": {"listChanged": true}, + "resources": {"subscribe": true, "listChanged": true}, + "tools": {"listChanged": false} + }, + "serverInfo": { + "name": "mock-server", + "version": "0.1.0" + }, + "instructions": "The server instructions." + } + """ + + let decoder = JSONDecoder() + let result = try decoder.decode(Initialize.Result.self, from: pythonJSON.data(using: .utf8)!) + + #expect(result.protocolVersion == Version.v2025_11_25) + #expect(result.capabilities.logging != nil) + #expect(result.capabilities.prompts?.listChanged == true) + #expect(result.capabilities.resources?.subscribe == true) + #expect(result.capabilities.resources?.listChanged == true) + #expect(result.capabilities.tools?.listChanged == false) + #expect(result.serverInfo.name == "mock-server") + #expect(result.serverInfo.version == "0.1.0") + #expect(result.instructions == "The server instructions.") + } +} + +// MARK: - Sampling Capability Tests (additional coverage) + +@Suite("Sampling Capability Encoding Tests") +struct SamplingCapabilityEncodingTests { + + @Test("Client sampling with no sub-capabilities encodes correctly") + func testClientSamplingBasic() throws { + let capabilities = Client.Capabilities( + sampling: .init() + ) + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + + let data = try encoder.encode(capabilities) + let json = String(data: data, encoding: .utf8)! + + // Should have empty sampling object + #expect(json == "{\"sampling\":{}}") + + let decoder = JSONDecoder() + let decoded = try decoder.decode(Client.Capabilities.self, from: data) + #expect(decoded.sampling != nil) + #expect(decoded.sampling?.tools == nil) + #expect(decoded.sampling?.context == nil) + } +} + +// MARK: - Tasks Capability Tests + +@Suite("Tasks Capability Encoding Tests") +struct TasksCapabilityEncodingTests { + + @Test("Server tasks capability encodes correctly") + func testServerTasksCapability() throws { + let capabilities = Server.Capabilities( + tasks: .init( + list: .init(), + cancel: .init(), + requests: .init(tools: .init(call: .init())) + ) + ) + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + + let data = try encoder.encode(capabilities) + let json = String(data: data, encoding: .utf8)! + + #expect(json.contains("\"tasks\"")) + #expect(json.contains("\"list\"")) + #expect(json.contains("\"cancel\"")) + #expect(json.contains("\"requests\"")) + + let decoder = JSONDecoder() + let decoded = try decoder.decode(Server.Capabilities.self, from: data) + #expect(decoded.tasks != nil) + #expect(decoded.tasks?.list != nil) + #expect(decoded.tasks?.cancel != nil) + #expect(decoded.tasks?.requests?.tools?.call != nil) + } + + @Test("Client tasks capability encodes correctly") + func testClientTasksCapability() throws { + let capabilities = Client.Capabilities( + tasks: .init( + list: .init(), + cancel: .init() + ) + ) + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + + let data = try encoder.encode(capabilities) + let json = String(data: data, encoding: .utf8)! + + #expect(json.contains("\"tasks\"")) + #expect(json.contains("\"list\"")) + #expect(json.contains("\"cancel\"")) + + let decoder = JSONDecoder() + let decoded = try decoder.decode(Client.Capabilities.self, from: data) + #expect(decoded.tasks != nil) + #expect(decoded.tasks?.list != nil) + #expect(decoded.tasks?.cancel != nil) + } +} diff --git a/Tests/MCPTests/ClientReconnectionTests.swift b/Tests/MCPTests/ClientReconnectionTests.swift new file mode 100644 index 00000000..47328926 --- /dev/null +++ b/Tests/MCPTests/ClientReconnectionTests.swift @@ -0,0 +1,238 @@ +import Foundation +import Testing + +@testable import MCP + +#if canImport(FoundationNetworking) + import FoundationNetworking +#endif + +/// Tests for HTTPClientTransport reconnection behavior. +/// +/// These tests verify the client transport's ability to handle disconnections +/// and reconnect with proper Last-Event-ID headers. +/// +/// TypeScript tests not yet implemented (require streaming mock infrastructure): +/// +/// The following tests from `packages/client/test/client/streamableHttp.test.ts` require +/// the ability to mock streaming HTTP responses with controlled failures. The TypeScript SDK +/// uses vi.fn() mocks that can return ReadableStream objects that error mid-stream. +/// Swift's MockURLProtocol doesn't easily support this pattern. +/// +/// Possible solutions for future implementation: +/// 1. Create a custom URLProtocol that supports async stream injection +/// 2. Use a real local HTTP server in tests (like TypeScript does with node's http.Server) +/// 3. Test at a lower level by mocking the EventSource directly +/// +/// Tests pending implementation: +/// - `should reconnect a GET-initiated notification stream that fails` +/// - `should NOT reconnect a POST-initiated stream that fails` +/// - `should reconnect a POST-initiated stream after receiving a priming event` +/// - `should NOT reconnect a POST stream when response was received` +/// - `should not attempt reconnection after close() is called` +/// - `should use server-provided retry value for reconnection delay` +/// - `should reconnect on graceful stream close` +/// - `should not schedule any reconnection attempts when maxRetries is 0` +@Suite("Client Reconnection Tests") +struct ClientReconnectionTests { + + // MARK: - Reconnection Options Tests + + @Test("Default reconnection options") + func defaultReconnectionOptions() async throws { + let transport = HTTPClientTransport( + endpoint: URL(string: "http://localhost:8080/mcp")!, + streaming: false + ) + + let options = transport.reconnectionOptions + #expect(options.initialReconnectionDelay == 1.0) + #expect(options.maxReconnectionDelay == 30.0) + #expect(options.reconnectionDelayGrowFactor == 1.5) + #expect(options.maxRetries == 2) + } + + @Test("Custom reconnection options") + func customReconnectionOptions() async throws { + let customOptions = HTTPReconnectionOptions( + initialReconnectionDelay: 0.5, + maxReconnectionDelay: 10.0, + reconnectionDelayGrowFactor: 2.0, + maxRetries: 5 + ) + + let transport = HTTPClientTransport( + endpoint: URL(string: "http://localhost:8080/mcp")!, + streaming: false, + reconnectionOptions: customOptions + ) + + let options = transport.reconnectionOptions + #expect(options.initialReconnectionDelay == 0.5) + #expect(options.maxReconnectionDelay == 10.0) + #expect(options.reconnectionDelayGrowFactor == 2.0) + #expect(options.maxRetries == 5) + } + + // MARK: - Last Event ID Tracking Tests + + @Test("Last received event ID is initially nil") + func lastEventIdInitiallyNil() async throws { + let transport = HTTPClientTransport( + endpoint: URL(string: "http://localhost:8080/mcp")!, + streaming: false + ) + + let lastEventId = await transport.lastReceivedEventId + #expect(lastEventId == nil) + } + + // MARK: - Resumption Token Callback Tests + + @Test("Resumption token callback can be set") + func resumptionTokenCallbackCanBeSet() async throws { + let transport = HTTPClientTransport( + endpoint: URL(string: "http://localhost:8080/mcp")!, + streaming: false + ) + + actor TokenCollector { + var tokens: [String] = [] + func add(_ token: String) { tokens.append(token) } + func get() -> [String] { tokens } + } + + let collector = TokenCollector() + + await transport.setOnResumptionToken { token in + Task { + await collector.add(token) + } + } + + // The callback is set but won't be triggered without an actual SSE stream + // This test just verifies the API works + let tokens = await collector.get() + #expect(tokens.isEmpty) + } + + // MARK: - Reconnection Options Struct Tests + + @Test("HTTPReconnectionOptions default values") + func reconnectionOptionsDefaultValues() { + let options = HTTPReconnectionOptions() + #expect(options.initialReconnectionDelay == 1.0) + #expect(options.maxReconnectionDelay == 30.0) + #expect(options.reconnectionDelayGrowFactor == 1.5) + #expect(options.maxRetries == 2) + } + + @Test("HTTPReconnectionOptions static default") + func reconnectionOptionsStaticDefault() { + let options = HTTPReconnectionOptions.default + #expect(options.initialReconnectionDelay == 1.0) + #expect(options.maxReconnectionDelay == 30.0) + #expect(options.reconnectionDelayGrowFactor == 1.5) + #expect(options.maxRetries == 2) + } + + @Test("HTTPReconnectionOptions custom initialization") + func reconnectionOptionsCustomInit() { + let options = HTTPReconnectionOptions( + initialReconnectionDelay: 2.0, + maxReconnectionDelay: 60.0, + reconnectionDelayGrowFactor: 3.0, + maxRetries: 10 + ) + #expect(options.initialReconnectionDelay == 2.0) + #expect(options.maxReconnectionDelay == 60.0) + #expect(options.reconnectionDelayGrowFactor == 3.0) + #expect(options.maxRetries == 10) + } + + // MARK: - Exponential Backoff Logic Tests + + @Test("Exponential backoff calculation") + func exponentialBackoffCalculation() { + // Test the math of exponential backoff + let options = HTTPReconnectionOptions( + initialReconnectionDelay: 1.0, + maxReconnectionDelay: 30.0, + reconnectionDelayGrowFactor: 2.0, + maxRetries: 10 + ) + + // delay = initialDelay * growFactor^attempt + // Attempt 0: 1.0 * 2^0 = 1.0 + // Attempt 1: 1.0 * 2^1 = 2.0 + // Attempt 2: 1.0 * 2^2 = 4.0 + // Attempt 3: 1.0 * 2^3 = 8.0 + // Attempt 4: 1.0 * 2^4 = 16.0 + // Attempt 5: 1.0 * 2^5 = 32.0 -> capped at 30.0 + + let delays = (0...5).map { attempt -> TimeInterval in + let delay = options.initialReconnectionDelay * pow(options.reconnectionDelayGrowFactor, Double(attempt)) + return min(delay, options.maxReconnectionDelay) + } + + #expect(delays[0] == 1.0) + #expect(delays[1] == 2.0) + #expect(delays[2] == 4.0) + #expect(delays[3] == 8.0) + #expect(delays[4] == 16.0) + #expect(delays[5] == 30.0) // Capped at max + } + + // MARK: - Transport State Tests + + @Test("Transport tracks session ID") + func transportTracksSessionId() async throws { + let transport = HTTPClientTransport( + endpoint: URL(string: "http://localhost:8080/mcp")!, + streaming: false + ) + + // Initially nil + let sessionId = await transport.sessionID + #expect(sessionId == nil) + } + + @Test("Transport tracks protocol version") + func transportTracksProtocolVersion() async throws { + let transport = HTTPClientTransport( + endpoint: URL(string: "http://localhost:8080/mcp")!, + streaming: false + ) + + // Initially nil + let version = await transport.protocolVersion + #expect(version == nil) + } + + // MARK: - Resume Stream API Tests + + #if !os(Linux) + @Test("Resume stream method exists and is callable") + func resumeStreamMethodExists() async throws { + let transport = HTTPClientTransport( + endpoint: URL(string: "http://localhost:8080/mcp")!, + streaming: true + ) + + // The method exists and is callable (though it won't do anything useful + // without a real connection). This test just verifies the API exists. + // When not connected, it should return early without throwing. + try await transport.resumeStream(from: "test-event-id") + // If we get here, the method exists and is callable + } + #endif +} + +// MARK: - HTTPClientTransport Extension for Testing + +extension HTTPClientTransport { + /// Sets the onResumptionToken callback + func setOnResumptionToken(_ callback: @escaping (String) -> Void) async { + self.onResumptionToken = callback + } +} diff --git a/Tests/MCPTests/ClientTests.swift b/Tests/MCPTests/ClientTests.swift index 80096cc1..5fa39701 100644 --- a/Tests/MCPTests/ClientTests.swift +++ b/Tests/MCPTests/ClientTests.swift @@ -790,4 +790,158 @@ struct ClientTests { await client.disconnect() } + + @Test( + "Unexpected transport closure with pending requests", + .timeLimit(.minutes(1)) + ) + func testUnexpectedTransportClosureWithPendingRequests() async throws { + // Based on: mcp-python-sdk/tests/client/test_stdio.py::test_stdio_client_bad_path + let transport = MockTransport() + let client = Client(name: "TestClient", version: "1.0") + + // Set up a task to handle the initialize response + let initTask = Task { + try await Task.sleep(for: .milliseconds(10)) + if let lastMessage = await transport.sentMessages.last, + let data = lastMessage.data(using: .utf8), + let request = try? JSONDecoder().decode(Request.self, from: data) + { + let response = Initialize.response( + id: request.id, + result: .init( + protocolVersion: Version.latest, + capabilities: .init(), + serverInfo: .init(name: "TestServer", version: "1.0"), + instructions: nil + ) + ) + try await transport.queue(response: response) + } + } + + try await client.connect(transport: transport) + try await Task.sleep(for: .milliseconds(10)) + initTask.cancel() + + // Start a ping request in a separate task - we intentionally don't queue + // a response, so this request will remain pending + let pingTask = Task { + try await client.ping() + } + + // Give it time to send the request and register as pending + try await Task.sleep(for: .milliseconds(20)) + + // Verify the ping request was sent (Initialize + Initialized notification + Ping) + #expect(await transport.sentMessages.count >= 3) + + // Simulate unexpected transport closure (e.g., server process exits) + // by disconnecting the transport directly without calling client.disconnect() + await transport.disconnect() + + // Wait for the receive loop to detect the closed transport and clean up + try await Task.sleep(for: .milliseconds(50)) + + // The pending request should receive a connectionClosed error + do { + _ = try await pingTask.value + #expect(Bool(false), "Expected request to fail with connectionClosed error") + } catch let error as MCPError { + #expect(error.code == ErrorCode.connectionClosed, "Expected CONNECTION_CLOSED error code") + let errorMessage = error.errorDescription ?? "" + #expect(errorMessage.contains("Connection closed")) + } catch { + #expect(Bool(false), "Expected MCPError, got \(error)") + } + + // Clean up + await client.disconnect() + } + + @Test("Client rejects unsupported server protocol version") + func testClientRejectsUnsupportedProtocolVersion() async throws { + let transport = MockTransport() + let client = Client(name: "TestClient", version: "1.0") + + // Set up a task to handle the initialize response with an unsupported version + let initTask = Task { + try await Task.sleep(for: .milliseconds(10)) + if let lastMessage = await transport.sentMessages.last, + let data = lastMessage.data(using: .utf8), + let request = try? JSONDecoder().decode(Request.self, from: data) + { + // Respond with an unsupported protocol version + let response = Initialize.response( + id: request.id, + result: .init( + protocolVersion: "2099-01-01", // Future unsupported version + capabilities: .init(), + serverInfo: .init(name: "TestServer", version: "1.0"), + instructions: nil + ) + ) + try await transport.queue(response: response) + } + } + + defer { initTask.cancel() } + + // Connect should fail with an error about unsupported protocol version + do { + try await client.connect(transport: transport) + #expect(Bool(false), "Expected connection to fail due to unsupported protocol version") + } catch let error as MCPError { + // Should be an invalidRequest error about unsupported version + if case .invalidRequest(let message) = error { + #expect(message?.contains("unsupported protocol version") == true) + #expect(message?.contains("2099-01-01") == true) + } else { + #expect(Bool(false), "Expected invalidRequest error, got \(error)") + } + } catch { + #expect(Bool(false), "Expected MCPError, got \(error)") + } + + // Client should have disconnected + #expect(await transport.isConnected == false) + } + + @Test("Client accepts supported server protocol version") + func testClientAcceptsSupportedProtocolVersion() async throws { + let transport = MockTransport() + let client = Client(name: "TestClient", version: "1.0") + + // Test with an older but supported version + let olderSupportedVersion = Version.v2024_11_05 + + // Set up a task to handle the initialize response + let initTask = Task { + try await Task.sleep(for: .milliseconds(10)) + if let lastMessage = await transport.sentMessages.last, + let data = lastMessage.data(using: .utf8), + let request = try? JSONDecoder().decode(Request.self, from: data) + { + let response = Initialize.response( + id: request.id, + result: .init( + protocolVersion: olderSupportedVersion, + capabilities: .init(), + serverInfo: .init(name: "TestServer", version: "1.0"), + instructions: nil + ) + ) + try await transport.queue(response: response) + } + } + + defer { initTask.cancel() } + + // Connect should succeed + let result = try await client.connect(transport: transport) + #expect(result.protocolVersion == olderSupportedVersion) + #expect(await transport.isConnected == true) + + await client.disconnect() + } } diff --git a/Tests/MCPTests/CompletionTests.swift b/Tests/MCPTests/CompletionTests.swift new file mode 100644 index 00000000..8dfaa0fa --- /dev/null +++ b/Tests/MCPTests/CompletionTests.swift @@ -0,0 +1,1035 @@ +import Foundation +import Testing + +@testable import MCP + +/// Tests for completion (autocomplete) functionality. +/// +/// These tests follow the patterns from: +/// - Python SDK: `tests/server/test_completion_with_context.py` +/// - TypeScript SDK: `packages/core/test/types.test.ts` (CompleteRequest tests) +/// - TypeScript SDK: `test/integration/test/server/mcp.test.ts` (completion integration tests) +@Suite("Completion Tests") +struct CompletionTests { + + // MARK: - Type Encoding/Decoding Tests + + @Test("PromptReference encoding and decoding") + func testPromptReferenceEncodingDecoding() throws { + let reference = PromptReference(name: "greeting") + + #expect(reference.type == "ref/prompt") + #expect(reference.name == "greeting") + #expect(reference.title == nil) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(reference) + let decoded = try decoder.decode(PromptReference.self, from: data) + + #expect(decoded.type == "ref/prompt") + #expect(decoded.name == "greeting") + #expect(decoded.title == nil) + } + + @Test("PromptReference with title encoding and decoding") + func testPromptReferenceWithTitleEncodingDecoding() throws { + let reference = PromptReference(name: "greeting", title: "Send Greeting") + + #expect(reference.type == "ref/prompt") + #expect(reference.name == "greeting") + #expect(reference.title == "Send Greeting") + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(reference) + let decoded = try decoder.decode(PromptReference.self, from: data) + + #expect(decoded.type == "ref/prompt") + #expect(decoded.name == "greeting") + #expect(decoded.title == "Send Greeting") + + // Verify JSON structure includes title + let jsonObject = try JSONSerialization.jsonObject(with: data) as! [String: Any] + #expect(jsonObject["type"] as? String == "ref/prompt") + #expect(jsonObject["name"] as? String == "greeting") + #expect(jsonObject["title"] as? String == "Send Greeting") + } + + @Test("ResourceTemplateReference encoding and decoding") + func testResourceTemplateReferenceEncodingDecoding() throws { + let reference = ResourceTemplateReference(uri: "github://repos/{owner}/{repo}") + + #expect(reference.type == "ref/resource") + #expect(reference.uri == "github://repos/{owner}/{repo}") + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(reference) + let decoded = try decoder.decode(ResourceTemplateReference.self, from: data) + + #expect(decoded.type == "ref/resource") + #expect(decoded.uri == "github://repos/{owner}/{repo}") + } + + @Test("CompletionReference prompt case encoding and decoding") + func testCompletionReferencePromptCase() throws { + let reference = CompletionReference.prompt(PromptReference(name: "test-prompt")) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(reference) + let decoded = try decoder.decode(CompletionReference.self, from: data) + + if case .prompt(let promptRef) = decoded { + #expect(promptRef.name == "test-prompt") + #expect(promptRef.type == "ref/prompt") + } else { + Issue.record("Expected prompt reference") + } + } + + @Test("CompletionReference resource case encoding and decoding") + func testCompletionReferenceResourceCase() throws { + let reference = CompletionReference.resource( + ResourceTemplateReference(uri: "file:///{path}") + ) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(reference) + let decoded = try decoder.decode(CompletionReference.self, from: data) + + if case .resource(let resourceRef) = decoded { + #expect(resourceRef.uri == "file:///{path}") + #expect(resourceRef.type == "ref/resource") + } else { + Issue.record("Expected resource reference") + } + } + + @Test("CompletionReference decoding unknown type throws error") + func testCompletionReferenceUnknownTypeThrows() throws { + let json = """ + {"type":"ref/unknown","name":"test"} + """ + let data = json.data(using: .utf8)! + let decoder = JSONDecoder() + + #expect(throws: DecodingError.self) { + _ = try decoder.decode(CompletionReference.self, from: data) + } + } + + @Test("CompletionArgument encoding and decoding") + func testCompletionArgumentEncodingDecoding() throws { + let argument = CompletionArgument(name: "language", value: "py") + + #expect(argument.name == "language") + #expect(argument.value == "py") + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(argument) + let decoded = try decoder.decode(CompletionArgument.self, from: data) + + #expect(decoded.name == "language") + #expect(decoded.value == "py") + } + + @Test("CompletionContext encoding and decoding with arguments") + func testCompletionContextWithArguments() throws { + let context = CompletionContext(arguments: [ + "owner": "modelcontextprotocol", + "database": "users_db", + ]) + + #expect(context.arguments?["owner"] == "modelcontextprotocol") + #expect(context.arguments?["database"] == "users_db") + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(context) + let decoded = try decoder.decode(CompletionContext.self, from: data) + + #expect(decoded.arguments?["owner"] == "modelcontextprotocol") + #expect(decoded.arguments?["database"] == "users_db") + } + + @Test("CompletionContext encoding and decoding without arguments") + func testCompletionContextWithoutArguments() throws { + let context = CompletionContext(arguments: nil) + + #expect(context.arguments == nil) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(context) + let decoded = try decoder.decode(CompletionContext.self, from: data) + + #expect(decoded.arguments == nil) + } + + @Test("CompletionContext with empty arguments") + func testCompletionContextWithEmptyArguments() throws { + let context = CompletionContext(arguments: [:]) + + #expect(context.arguments?.isEmpty == true) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(context) + let decoded = try decoder.decode(CompletionContext.self, from: data) + + #expect(decoded.arguments?.isEmpty == true) + } + + @Test("CompletionSuggestions encoding and decoding") + func testCompletionSuggestionsEncodingDecoding() throws { + let suggestions = CompletionSuggestions( + values: ["python", "javascript", "typescript"], + total: 10, + hasMore: true + ) + + #expect(suggestions.values == ["python", "javascript", "typescript"]) + #expect(suggestions.total == 10) + #expect(suggestions.hasMore == true) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(suggestions) + let decoded = try decoder.decode(CompletionSuggestions.self, from: data) + + #expect(decoded.values == ["python", "javascript", "typescript"]) + #expect(decoded.total == 10) + #expect(decoded.hasMore == true) + } + + @Test("CompletionSuggestions with minimal fields") + func testCompletionSuggestionsMinimal() throws { + let suggestions = CompletionSuggestions(values: ["a", "b"]) + + #expect(suggestions.values == ["a", "b"]) + #expect(suggestions.total == nil) + #expect(suggestions.hasMore == nil) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(suggestions) + let decoded = try decoder.decode(CompletionSuggestions.self, from: data) + + #expect(decoded.values == ["a", "b"]) + #expect(decoded.total == nil) + #expect(decoded.hasMore == nil) + } + + // MARK: - Spec Compliance Tests (100-item limit) + + @Test("CompletionSuggestions maxValues constant is 100") + func testMaxValuesConstant() { + #expect(CompletionSuggestions.maxValues == 100) + } + + @Test("CompletionSuggestions.empty returns empty result") + func testEmptySuggestions() { + let empty = CompletionSuggestions.empty + + #expect(empty.values.isEmpty) + #expect(empty.total == nil) + #expect(empty.hasMore == false) + } + + @Test("CompletionSuggestions init truncates values over 100") + func testInitTruncatesOverMaxValues() { + // Create 150 values + let allValues = (1...150).map { "value\($0)" } + let suggestions = CompletionSuggestions(values: allValues, total: 150, hasMore: true) + + // Should be truncated to 100 + #expect(suggestions.values.count == 100) + #expect(suggestions.values.first == "value1") + #expect(suggestions.values.last == "value100") + // User-specified total and hasMore are preserved + #expect(suggestions.total == 150) + #expect(suggestions.hasMore == true) + } + + @Test("CompletionSuggestions init(from:) with few values") + func testInitFromFewValues() { + let values = ["python", "javascript", "typescript"] + let suggestions = CompletionSuggestions(from: values) + + #expect(suggestions.values == values) + #expect(suggestions.total == 3) + #expect(suggestions.hasMore == false) + } + + @Test("CompletionSuggestions init(from:) with exactly 100 values") + func testInitFromExactly100Values() { + let values = (1...100).map { "value\($0)" } + let suggestions = CompletionSuggestions(from: values) + + #expect(suggestions.values.count == 100) + #expect(suggestions.total == 100) + #expect(suggestions.hasMore == false) + } + + @Test("CompletionSuggestions init(from:) with over 100 values") + func testInitFromOver100Values() { + let values = (1...250).map { "value\($0)" } + let suggestions = CompletionSuggestions(from: values) + + #expect(suggestions.values.count == 100) + #expect(suggestions.values.first == "value1") + #expect(suggestions.values.last == "value100") + #expect(suggestions.total == 250) + #expect(suggestions.hasMore == true) + } + + @Test("CompletionSuggestions init(from:) with empty array") + func testInitFromEmptyArray() { + let suggestions = CompletionSuggestions(from: []) + + #expect(suggestions.values.isEmpty) + #expect(suggestions.total == 0) + #expect(suggestions.hasMore == false) + } + + @Test("Complete.Result.empty returns empty result") + func testCompleteResultEmpty() { + let empty = Complete.Result.empty + + #expect(empty.completion.values.isEmpty) + #expect(empty.completion.hasMore == false) + #expect(empty._meta == nil) + #expect(empty.extraFields == nil) + } + + @Test("Complete.Result init(from:) convenience initializer") + func testCompleteResultInitFrom() { + let values = ["alice", "bob", "charlie"] + let result = Complete.Result(from: values) + + #expect(result.completion.values == values) + #expect(result.completion.total == 3) + #expect(result.completion.hasMore == false) + #expect(result._meta == nil) + #expect(result.extraFields == nil) + } + + @Test("Complete.Result init(from:) with over 100 values") + func testCompleteResultInitFromOver100() { + let values = (1...200).map { "item\($0)" } + let result = Complete.Result(from: values) + + #expect(result.completion.values.count == 100) + #expect(result.completion.total == 200) + #expect(result.completion.hasMore == true) + } + + // MARK: - Complete Request/Result Tests + + @Test("Complete.Parameters encoding without context") + func testCompleteParametersWithoutContext() throws { + let params = Complete.Parameters( + ref: .prompt(PromptReference(name: "greeting")), + argument: CompletionArgument(name: "name", value: "A") + ) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(params) + let decoded = try decoder.decode(Complete.Parameters.self, from: data) + + if case .prompt(let promptRef) = decoded.ref { + #expect(promptRef.name == "greeting") + #expect(promptRef.type == "ref/prompt") + } else { + Issue.record("Expected prompt reference") + } + #expect(decoded.argument.name == "name") + #expect(decoded.argument.value == "A") + #expect(decoded.context == nil) + } + + @Test("Complete.Parameters encoding with context") + func testCompleteParametersWithContext() throws { + let params = Complete.Parameters( + ref: .resource(ResourceTemplateReference(uri: "github://repos/{owner}/{repo}")), + argument: CompletionArgument(name: "repo", value: "t"), + context: CompletionContext(arguments: ["{owner}": "microsoft"]) + ) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(params) + let decoded = try decoder.decode(Complete.Parameters.self, from: data) + + if case .resource(let resourceRef) = decoded.ref { + #expect(resourceRef.uri == "github://repos/{owner}/{repo}") + } else { + Issue.record("Expected resource reference") + } + #expect(decoded.argument.name == "repo") + #expect(decoded.argument.value == "t") + #expect(decoded.context?.arguments?["{owner}"] == "microsoft") + } + + @Test("Complete.Parameters with multiple resolved variables") + func testCompleteParametersWithMultipleResolvedVariables() throws { + let params = Complete.Parameters( + ref: .resource(ResourceTemplateReference(uri: "api://v1/{tenant}/{resource}/{id}")), + argument: CompletionArgument(name: "id", value: "123"), + context: CompletionContext(arguments: [ + "{tenant}": "acme-corp", + "{resource}": "users", + ]) + ) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(params) + let decoded = try decoder.decode(Complete.Parameters.self, from: data) + + #expect(decoded.context?.arguments?["{tenant}"] == "acme-corp") + #expect(decoded.context?.arguments?["{resource}"] == "users") + } + + @Test("Complete.Result encoding and decoding") + func testCompleteResultEncodingDecoding() throws { + let result = Complete.Result( + completion: CompletionSuggestions( + values: ["typescript-sdk", "python-sdk", "specification"], + total: 3, + hasMore: false + ) + ) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(result) + let decoded = try decoder.decode(Complete.Result.self, from: data) + + #expect(decoded.completion.values == ["typescript-sdk", "python-sdk", "specification"]) + #expect(decoded.completion.total == 3) + #expect(decoded.completion.hasMore == false) + } + + @Test("Complete request JSON-RPC format") + func testCompleteRequestJsonRpcFormat() throws { + let request = Complete.request(.init( + ref: .prompt(PromptReference(name: "review_code")), + argument: CompletionArgument(name: "language", value: "py") + )) + + #expect(request.method == Complete.name) + #expect(Complete.name == "completion/complete") + + // Verify it can roundtrip through JSON + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(request) + let decoded = try decoder.decode(Request.self, from: data) + + #expect(decoded.method == "completion/complete") + if case .prompt(let promptRef) = decoded.params.ref { + #expect(promptRef.name == "review_code") + } else { + Issue.record("Expected prompt reference") + } + } + + // MARK: - Server Handler Integration Tests + + /// Actor to safely track received parameters across async closures + private actor ReceivedParams { + var ref: CompletionReference? + var argument: CompletionArgument? + var context: CompletionContext? + var contextWasNil = false + + func set(ref: CompletionReference, argument: CompletionArgument, context: CompletionContext?) { + self.ref = ref + self.argument = argument + self.context = context + self.contextWasNil = context == nil + } + + func getRef() -> CompletionReference? { ref } + func getArgument() -> CompletionArgument? { argument } + func getContext() -> CompletionContext? { context } + func wasContextNil() -> Bool { contextWasNil } + } + + @Test("Completion handler receives context correctly") + func testCompletionHandlerReceivesContext() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + // Track what the handler receives + let received = ReceivedParams() + + let server = Server( + name: "test-server", + version: "1.0.0", + capabilities: .init(completions: .init()) + ) + + await server.withRequestHandler(Complete.self) { [received] params, _ in + await received.set(ref: params.ref, argument: params.argument, context: params.context) + return Complete.Result( + completion: CompletionSuggestions( + values: ["test-completion"], + total: 1, + hasMore: false + ) + ) + } + + try await server.start(transport: serverTransport) + + let client = Client(name: "test-client", version: "1.0.0") + _ = try await client.connect(transport: clientTransport) + + // Request completion with context + let result = try await client.complete( + ref: .resource(ResourceTemplateReference(uri: "test://resource/{param}")), + argument: CompletionArgument(name: "param", value: "test"), + context: CompletionContext(arguments: ["previous": "value"]) + ) + + // Verify handler received the context + let receivedContext = await received.getContext() + #expect(receivedContext != nil) + #expect(receivedContext?.arguments?["previous"] == "value") + #expect(result.values == ["test-completion"]) + + // Verify the ref and argument were received correctly + let receivedRef = await received.getRef() + if case .resource(let resourceRef) = receivedRef { + #expect(resourceRef.uri == "test://resource/{param}") + } else { + Issue.record("Expected resource reference") + } + let receivedArgument = await received.getArgument() + #expect(receivedArgument?.name == "param") + #expect(receivedArgument?.value == "test") + + await client.disconnect() + await server.stop() + } + + @Test("Completion works without context (backward compatibility)") + func testCompletionBackwardCompatibility() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let received = ReceivedParams() + + let server = Server( + name: "test-server", + version: "1.0.0", + capabilities: .init(completions: .init()) + ) + + await server.withRequestHandler(Complete.self) { [received] params, _ in + await received.set(ref: params.ref, argument: params.argument, context: params.context) + return Complete.Result( + completion: CompletionSuggestions( + values: ["no-context-completion"], + total: 1, + hasMore: false + ) + ) + } + + try await server.start(transport: serverTransport) + + let client = Client(name: "test-client", version: "1.0.0") + _ = try await client.connect(transport: clientTransport) + + // Request completion without context + let result = try await client.complete( + ref: .prompt(PromptReference(name: "test-prompt")), + argument: CompletionArgument(name: "arg", value: "val") + ) + + #expect(await received.wasContextNil()) + #expect(result.values == ["no-context-completion"]) + + await client.disconnect() + await server.stop() + } + + @Test("Dependent completion scenario (database/table)") + func testDependentCompletionScenario() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "test-server", + version: "1.0.0", + capabilities: .init(completions: .init()) + ) + + // Handler that returns different completions based on context + await server.withRequestHandler(Complete.self) { params, _ in + if case .resource(let resourceRef) = params.ref { + if resourceRef.uri == "db://{database}/{table}" { + if params.argument.name == "database" { + // Complete database names + return Complete.Result( + completion: CompletionSuggestions( + values: ["users_db", "products_db", "analytics_db"], + total: 3, + hasMore: false + ) + ) + } else if params.argument.name == "table" { + // Complete table names based on selected database + let db = params.context?.arguments?["database"] + let tables: [String] + switch db { + case "users_db": + tables = ["users", "sessions", "permissions"] + case "products_db": + tables = ["products", "categories", "inventory"] + default: + tables = [] + } + return Complete.Result( + completion: CompletionSuggestions( + values: tables, + total: tables.count, + hasMore: false + ) + ) + } + } + } + return Complete.Result( + completion: CompletionSuggestions(values: [], total: 0, hasMore: false) + ) + } + + try await server.start(transport: serverTransport) + + let client = Client(name: "test-client", version: "1.0.0") + _ = try await client.connect(transport: clientTransport) + + // First, complete database + let dbResult = try await client.complete( + ref: .resource(ResourceTemplateReference(uri: "db://{database}/{table}")), + argument: CompletionArgument(name: "database", value: "") + ) + #expect(dbResult.values.contains("users_db")) + #expect(dbResult.values.contains("products_db")) + + // Then complete table with database context + let tableResult = try await client.complete( + ref: .resource(ResourceTemplateReference(uri: "db://{database}/{table}")), + argument: CompletionArgument(name: "table", value: ""), + context: CompletionContext(arguments: ["database": "users_db"]) + ) + #expect(tableResult.values == ["users", "sessions", "permissions"]) + + // Different database gives different tables + let tableResult2 = try await client.complete( + ref: .resource(ResourceTemplateReference(uri: "db://{database}/{table}")), + argument: CompletionArgument(name: "table", value: ""), + context: CompletionContext(arguments: ["database": "products_db"]) + ) + #expect(tableResult2.values == ["products", "categories", "inventory"]) + + await client.disconnect() + await server.stop() + } + + @Test("Completion error handling when context is required") + func testCompletionErrorOnMissingContext() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "test-server", + version: "1.0.0", + capabilities: .init(completions: .init()) + ) + + await server.withRequestHandler(Complete.self) { params, _ in + if case .resource(let resourceRef) = params.ref { + if resourceRef.uri == "db://{database}/{table}" && params.argument.name == "table" { + // Check if database context is provided + guard let arguments = params.context?.arguments, + arguments["database"] != nil + else { + throw MCPError.invalidParams( + "Please select a database first to see available tables") + } + // Return completions if context is provided + return Complete.Result( + completion: CompletionSuggestions( + values: ["users", "orders", "products"], + total: 3, + hasMore: false + ) + ) + } + } + return Complete.Result( + completion: CompletionSuggestions(values: [], total: 0, hasMore: false) + ) + } + + try await server.start(transport: serverTransport) + + let client = Client(name: "test-client", version: "1.0.0") + _ = try await client.connect(transport: clientTransport) + + // Try to complete table without database context - should fail + do { + _ = try await client.complete( + ref: .resource(ResourceTemplateReference(uri: "db://{database}/{table}")), + argument: CompletionArgument(name: "table", value: "") + ) + Issue.record("Expected error for missing context") + } catch { + let errorMessage = String(describing: error) + #expect(errorMessage.contains("database") || errorMessage.contains("select")) + } + + // Now complete with proper context - should work + let result = try await client.complete( + ref: .resource(ResourceTemplateReference(uri: "db://{database}/{table}")), + argument: CompletionArgument(name: "table", value: ""), + context: CompletionContext(arguments: ["database": "test_db"]) + ) + #expect(result.values == ["users", "orders", "products"]) + + await client.disconnect() + await server.stop() + } + + @Test("Prompt completion with filtered results") + func testPromptCompletionWithFilteredResults() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "test-server", + version: "1.0.0", + capabilities: .init(completions: .init()) + ) + + await server.withRequestHandler(Complete.self) { params, _ in + if case .prompt(let promptRef) = params.ref { + if promptRef.name == "review_code" && params.argument.name == "language" { + let allLanguages = ["python", "javascript", "typescript", "java", "go", "rust"] + let filtered = allLanguages.filter { $0.hasPrefix(params.argument.value) } + return Complete.Result( + completion: CompletionSuggestions( + values: filtered, + total: filtered.count, + hasMore: false + ) + ) + } + } + return Complete.Result( + completion: CompletionSuggestions(values: [], total: 0, hasMore: false) + ) + } + + try await server.start(transport: serverTransport) + + let client = Client(name: "test-client", version: "1.0.0") + _ = try await client.connect(transport: clientTransport) + + // Request completion with "py" prefix + let result = try await client.complete( + ref: .prompt(PromptReference(name: "review_code")), + argument: CompletionArgument(name: "language", value: "py") + ) + + #expect(result.values == ["python"]) + + // Request completion with "j" prefix + let result2 = try await client.complete( + ref: .prompt(PromptReference(name: "review_code")), + argument: CompletionArgument(name: "language", value: "j") + ) + + #expect(result2.values.contains("javascript")) + #expect(result2.values.contains("java")) + #expect(result2.values.allSatisfy { $0.hasPrefix("j") }) + + await client.disconnect() + await server.stop() + } + + // MARK: - Capability Tests + + @Test("Server advertises completions capability") + func testServerAdvertisesCompletionsCapability() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "test-server", + version: "1.0.0", + capabilities: .init(completions: .init()) + ) + + await server.withRequestHandler(Complete.self) { _, _ in + Complete.Result(completion: CompletionSuggestions(values: [])) + } + + try await server.start(transport: serverTransport) + + let client = Client(name: "test-client", version: "1.0.0") + let initResult = try await client.connect(transport: clientTransport) + + #expect(initResult.capabilities.completions != nil) + + await client.disconnect() + await server.stop() + } + + @Test("Server without completions capability does not advertise it") + func testServerWithoutCompletionsCapability() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + // Server without completions capability + let server = Server( + name: "test-server", + version: "1.0.0", + capabilities: .init(tools: .init()) // Only tools, no completions + ) + + try await server.start(transport: serverTransport) + + let client = Client(name: "test-client", version: "1.0.0") + let initResult = try await client.connect(transport: clientTransport) + + #expect(initResult.capabilities.completions == nil) + + await client.disconnect() + await server.stop() + } + + @Test("Client in strict mode rejects completion when server lacks capability") + func testStrictClientRejectsCompletionWithoutCapability() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + // Server without completions capability + let server = Server( + name: "test-server", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + try await server.start(transport: serverTransport) + + let client = Client( + name: "test-client", + version: "1.0.0", + configuration: .init(strict: true) + ) + _ = try await client.connect(transport: clientTransport) + + // Attempt to complete should throw + do { + _ = try await client.complete( + ref: .prompt(PromptReference(name: "test")), + argument: CompletionArgument(name: "arg", value: "") + ) + Issue.record("Expected error when server lacks completions capability") + } catch let error as MCPError { + if case .methodNotFound(let message) = error { + let msg = message ?? "" + #expect(msg.contains("Completions") || msg.contains("not supported")) + } else { + Issue.record("Expected methodNotFound error, got: \(error)") + } + } + + await client.disconnect() + await server.stop() + } + + // MARK: - Resource Template Completion Tests + + @Test("Resource template completion with context") + func testResourceTemplateCompletionWithContext() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "test-server", + version: "1.0.0", + capabilities: .init(completions: .init()) + ) + + // Simulate GitHub repos completion based on owner + await server.withRequestHandler(Complete.self) { params, _ in + if case .resource(let resourceRef) = params.ref, + resourceRef.uri == "github://repos/{owner}/{repo}" + { + let owner = params.context?.arguments?["owner"] + let repos: [String] + switch owner { + case "modelcontextprotocol": + repos = ["python-sdk", "typescript-sdk", "specification"] + case "microsoft": + repos = ["vscode", "typescript", "playwright"] + case "facebook": + repos = ["react", "react-native", "jest"] + default: + repos = ["repo1", "repo2", "repo3"] + } + return Complete.Result( + completion: CompletionSuggestions( + values: repos, + total: repos.count, + hasMore: false + ) + ) + } + return Complete.Result( + completion: CompletionSuggestions(values: [], total: 0, hasMore: false) + ) + } + + try await server.start(transport: serverTransport) + + let client = Client(name: "test-client", version: "1.0.0") + _ = try await client.connect(transport: clientTransport) + + // Test with modelcontextprotocol owner + let result1 = try await client.complete( + ref: .resource(ResourceTemplateReference(uri: "github://repos/{owner}/{repo}")), + argument: CompletionArgument(name: "repo", value: ""), + context: CompletionContext(arguments: ["owner": "modelcontextprotocol"]) + ) + #expect(result1.values.contains("python-sdk")) + #expect(result1.values.contains("typescript-sdk")) + #expect(result1.values.contains("specification")) + #expect(result1.total == 3) + + // Test with microsoft owner + let result2 = try await client.complete( + ref: .resource(ResourceTemplateReference(uri: "github://repos/{owner}/{repo}")), + argument: CompletionArgument(name: "repo", value: ""), + context: CompletionContext(arguments: ["owner": "microsoft"]) + ) + #expect(result2.values.contains("vscode")) + #expect(result2.values.contains("typescript")) + #expect(result2.values.contains("playwright")) + + // Test with no context + let result3 = try await client.complete( + ref: .resource(ResourceTemplateReference(uri: "github://repos/{owner}/{repo}")), + argument: CompletionArgument(name: "repo", value: "") + ) + #expect(result3.values == ["repo1", "repo2", "repo3"]) + + await client.disconnect() + await server.stop() + } + + // MARK: - Prompt Completion with Context Tests + + @Test("Prompt argument completion with resolved context") + func testPromptArgumentCompletionWithResolvedContext() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "test-server", + version: "1.0.0", + capabilities: .init(completions: .init()) + ) + + // Simulate team member completion based on department + await server.withRequestHandler(Complete.self) { params, _ in + if case .prompt(let promptRef) = params.ref, + promptRef.name == "team-greeting", + params.argument.name == "name" + { + let department = params.context?.arguments?["department"] + let value = params.argument.value + let names: [String] + switch department { + case "engineering": + names = ["Alice", "Bob", "Charlie"] + case "sales": + names = ["David", "Eve", "Frank"] + case "marketing": + names = ["Grace", "Henry", "Ivy"] + default: + names = ["Unknown1", "Unknown2"] + } + let filtered = names.filter { $0.lowercased().hasPrefix(value.lowercased()) } + return Complete.Result( + completion: CompletionSuggestions( + values: filtered, + total: filtered.count, + hasMore: false + ) + ) + } + return Complete.Result( + completion: CompletionSuggestions(values: [], total: 0, hasMore: false) + ) + } + + try await server.start(transport: serverTransport) + + let client = Client(name: "test-client", version: "1.0.0") + _ = try await client.connect(transport: clientTransport) + + // Test with engineering department + let result1 = try await client.complete( + ref: .prompt(PromptReference(name: "team-greeting")), + argument: CompletionArgument(name: "name", value: "A"), + context: CompletionContext(arguments: ["department": "engineering"]) + ) + #expect(result1.values == ["Alice"]) + + // Test with sales department + let result2 = try await client.complete( + ref: .prompt(PromptReference(name: "team-greeting")), + argument: CompletionArgument(name: "name", value: "D"), + context: CompletionContext(arguments: ["department": "sales"]) + ) + #expect(result2.values == ["David"]) + + // Test with marketing department + let result3 = try await client.complete( + ref: .prompt(PromptReference(name: "team-greeting")), + argument: CompletionArgument(name: "name", value: "G"), + context: CompletionContext(arguments: ["department": "marketing"]) + ) + #expect(result3.values == ["Grace"]) + + // Test with no context + let result4 = try await client.complete( + ref: .prompt(PromptReference(name: "team-greeting")), + argument: CompletionArgument(name: "name", value: "U") + ) + #expect(result4.values.contains("Unknown1")) + #expect(result4.values.contains("Unknown2")) + + await client.disconnect() + await server.stop() + } +} diff --git a/Tests/MCPTests/ElicitationTests.swift b/Tests/MCPTests/ElicitationTests.swift new file mode 100644 index 00000000..36caeae6 --- /dev/null +++ b/Tests/MCPTests/ElicitationTests.swift @@ -0,0 +1,2460 @@ +import Foundation +import Testing + +@testable import MCP + +// MARK: - Test Helpers + +/// Extension to provide convenient value accessors for testing +extension ElicitValue { + var stringValue: String? { + if case .string(let value) = self { return value } + return nil + } + + var intValue: Int? { + if case .int(let value) = self { return value } + return nil + } + + var doubleValue: Double? { + if case .double(let value) = self { return value } + return nil + } + + var boolValue: Bool? { + if case .bool(let value) = self { return value } + return nil + } + + var stringsValue: [String]? { + if case .strings(let value) = self { return value } + return nil + } +} + +// MARK: - Schema Encoding/Decoding Tests + +@Suite("Schema Encoding/Decoding Tests") +struct SchemaEncodingTests { + + @Test("StringSchema encodes and decodes correctly") + func testStringSchemaRoundtrip() throws { + let schema = StringSchema( + title: "Name", + description: "The user's name", + minLength: 1, + maxLength: 100, + format: .email, + defaultValue: "user@example.com" + ) + + let encoder = JSONEncoder() + encoder.outputFormatting = .sortedKeys + let data = try encoder.encode(schema) + let decoded = try JSONDecoder().decode(StringSchema.self, from: data) + + #expect(decoded.title == "Name") + #expect(decoded.description == "The user's name") + #expect(decoded.minLength == 1) + #expect(decoded.maxLength == 100) + #expect(decoded.format == .email) + #expect(decoded.defaultValue == "user@example.com") + } + + @Test("StringSchema with format encodes correctly") + func testStringSchemaFormats() throws { + let formats: [StringSchemaFormat] = [.email, .uri, .date, .dateTime] + + for format in formats { + let schema = StringSchema(format: format) + let data = try JSONEncoder().encode(schema) + let decoded = try JSONDecoder().decode(StringSchema.self, from: data) + #expect(decoded.format == format) + } + } + + @Test("StringSchema with pattern encodes correctly") + func testStringSchemaPattern() throws { + let schema = StringSchema( + title: "ZIP Code", + pattern: "^[0-9]{5}$" + ) + + let encoder = JSONEncoder() + encoder.outputFormatting = .sortedKeys + let data = try encoder.encode(schema) + let json = String(data: data, encoding: .utf8)! + + #expect(json.contains("\"pattern\":\"^[0-9]{5}$\"")) + + let decoded = try JSONDecoder().decode(StringSchema.self, from: data) + #expect(decoded.pattern == "^[0-9]{5}$") + #expect(decoded.title == "ZIP Code") + } + + @Test("NumberSchema encodes and decodes correctly") + func testNumberSchemaRoundtrip() throws { + let schema = NumberSchema( + isInteger: true, + title: "Age", + description: "User age", + minimum: 0, + maximum: 150, + defaultValue: 25 + ) + + let data = try JSONEncoder().encode(schema) + let decoded = try JSONDecoder().decode(NumberSchema.self, from: data) + + #expect(decoded.type == "integer") // isInteger is encoded as type + #expect(decoded.title == "Age") + #expect(decoded.minimum == 0) + #expect(decoded.maximum == 150) + #expect(decoded.defaultValue == 25) + } + + @Test("BooleanSchema encodes and decodes correctly") + func testBooleanSchemaRoundtrip() throws { + let schema = BooleanSchema( + title: "Subscribe", + description: "Subscribe to newsletter", + defaultValue: true + ) + + let data = try JSONEncoder().encode(schema) + let decoded = try JSONDecoder().decode(BooleanSchema.self, from: data) + + #expect(decoded.title == "Subscribe") + #expect(decoded.defaultValue == true) + } + + @Test("TitledEnumSchema encodes and decodes correctly") + func testTitledEnumSchemaRoundtrip() throws { + let schema = TitledEnumSchema( + title: "Color", + description: "Pick a color", + oneOf: [ + TitledEnumOption(const: "red", title: "Red"), + TitledEnumOption(const: "green", title: "Green"), + TitledEnumOption(const: "blue", title: "Blue"), + ], + defaultValue: "red" + ) + + let data = try JSONEncoder().encode(schema) + let decoded = try JSONDecoder().decode(TitledEnumSchema.self, from: data) + + #expect(decoded.title == "Color") + #expect(decoded.oneOf.count == 3) + #expect(decoded.oneOf[0].const == "red") + #expect(decoded.oneOf[0].title == "Red") + #expect(decoded.defaultValue == "red") + } + + @Test("UntitledEnumSchema encodes and decodes correctly") + func testUntitledEnumSchemaRoundtrip() throws { + let schema = UntitledEnumSchema( + title: "Size", + enumValues: ["small", "medium", "large"], + defaultValue: "medium" + ) + + let data = try JSONEncoder().encode(schema) + let decoded = try JSONDecoder().decode(UntitledEnumSchema.self, from: data) + + #expect(decoded.title == "Size") + #expect(decoded.enumValues == ["small", "medium", "large"]) + #expect(decoded.defaultValue == "medium") + } + + @Test("TitledMultiSelectEnumSchema encodes and decodes correctly") + func testTitledMultiSelectEnumSchemaRoundtrip() throws { + let schema = TitledMultiSelectEnumSchema( + title: "Interests", + description: "Select your interests", + options: [ + TitledEnumOption(const: "tech", title: "Technology"), + TitledEnumOption(const: "sports", title: "Sports"), + TitledEnumOption(const: "music", title: "Music"), + ], + defaultValue: ["tech"] + ) + + let data = try JSONEncoder().encode(schema) + let decoded = try JSONDecoder().decode(TitledMultiSelectEnumSchema.self, from: data) + + #expect(decoded.title == "Interests") + #expect(decoded.items.anyOf.count == 3) // options are in items.anyOf + #expect(decoded.defaultValue == ["tech"]) + } + + @Test("UntitledMultiSelectEnumSchema encodes and decodes correctly") + func testUntitledMultiSelectEnumSchemaRoundtrip() throws { + let schema = UntitledMultiSelectEnumSchema( + title: "Tags", + enumValues: ["tag1", "tag2", "tag3"], + defaultValue: ["tag1", "tag2"] + ) + + let data = try JSONEncoder().encode(schema) + let decoded = try JSONDecoder().decode(UntitledMultiSelectEnumSchema.self, from: data) + + #expect(decoded.title == "Tags") + #expect(decoded.items.enumValues == ["tag1", "tag2", "tag3"]) // enumValues are in items + #expect(decoded.defaultValue == ["tag1", "tag2"]) + } + + @Test("ElicitationSchema with mixed property types encodes correctly") + func testElicitationSchemaWithMixedTypes() throws { + let schema = ElicitationSchema( + properties: [ + "name": .string(StringSchema(title: "Name")), + "age": .number(NumberSchema(isInteger: true, title: "Age")), + "subscribe": .boolean(BooleanSchema(title: "Subscribe")), + "color": .titledEnum(TitledEnumSchema( + title: "Color", + oneOf: [TitledEnumOption(const: "red", title: "Red")] + )), + ], + required: ["name", "age"] + ) + + let data = try JSONEncoder().encode(schema) + let decoded = try JSONDecoder().decode(ElicitationSchema.self, from: data) + + #expect(decoded.properties.count == 4) + #expect(decoded.required == ["name", "age"]) + } +} + +// MARK: - ElicitRequestParams Tests + +@Suite("ElicitRequestParams Tests") +struct ElicitRequestParamsTests { + + @Test("Form mode params encode and decode correctly") + func testFormModeParams() throws { + let params = ElicitRequestParams.form(ElicitRequestFormParams( + message: "Please fill out this form", + requestedSchema: ElicitationSchema( + properties: [ + "name": .string(StringSchema(title: "Name")), + ], + required: ["name"] + ) + )) + + let data = try JSONEncoder().encode(params) + let decoded = try JSONDecoder().decode(ElicitRequestParams.self, from: data) + + if case .form(let formParams) = decoded { + #expect(formParams.message == "Please fill out this form") + #expect(formParams.requestedSchema.properties.count == 1) + } else { + Issue.record("Expected form params") + } + } + + @Test("URL mode params encode and decode correctly") + func testURLModeParams() throws { + let params = ElicitRequestParams.url(ElicitRequestURLParams( + message: "Please authorize access", + elicitationId: "auth-123", + url: "https://example.com/oauth/authorize" + )) + + let data = try JSONEncoder().encode(params) + let decoded = try JSONDecoder().decode(ElicitRequestParams.self, from: data) + + if case .url(let urlParams) = decoded { + #expect(urlParams.message == "Please authorize access") + #expect(urlParams.elicitationId == "auth-123") + #expect(urlParams.url == "https://example.com/oauth/authorize") + } else { + Issue.record("Expected URL params") + } + } +} + +// MARK: - ElicitResult Tests + +@Suite("ElicitResult Tests") +struct ElicitResultTests { + + @Test("ElicitResult with accept action encodes correctly") + func testElicitResultAccept() throws { + let result = ElicitResult( + action: .accept, + content: [ + "name": .string("Alice"), + "age": .int(30), + ] + ) + + let data = try JSONEncoder().encode(result) + let decoded = try JSONDecoder().decode(ElicitResult.self, from: data) + + #expect(decoded.action == .accept) + #expect(decoded.content?["name"]?.stringValue == "Alice") + #expect(decoded.content?["age"]?.intValue == 30) + } + + @Test("ElicitResult with decline action encodes correctly") + func testElicitResultDecline() throws { + let result = ElicitResult(action: .decline) + + let data = try JSONEncoder().encode(result) + let decoded = try JSONDecoder().decode(ElicitResult.self, from: data) + + #expect(decoded.action == .decline) + #expect(decoded.content == nil) + } + + @Test("ElicitResult with cancel action encodes correctly") + func testElicitResultCancel() throws { + let result = ElicitResult(action: .cancel) + + let data = try JSONEncoder().encode(result) + let decoded = try JSONDecoder().decode(ElicitResult.self, from: data) + + #expect(decoded.action == .cancel) + } +} + +// MARK: - ElicitValue Tests + +@Suite("ElicitValue Tests") +struct ElicitValueTests { + + @Test("String value encodes and decodes correctly") + func testStringValue() throws { + let value = ElicitValue.string("hello") + let data = try JSONEncoder().encode(value) + let decoded = try JSONDecoder().decode(ElicitValue.self, from: data) + #expect(decoded.stringValue == "hello") + } + + @Test("Int value encodes and decodes correctly") + func testIntValue() throws { + let value = ElicitValue.int(42) + let data = try JSONEncoder().encode(value) + let decoded = try JSONDecoder().decode(ElicitValue.self, from: data) + #expect(decoded.intValue == 42) + } + + @Test("Double value encodes and decodes correctly") + func testDoubleValue() throws { + let value = ElicitValue.double(3.14) + let data = try JSONEncoder().encode(value) + let decoded = try JSONDecoder().decode(ElicitValue.self, from: data) + #expect(decoded.doubleValue == 3.14) + } + + @Test("Bool value encodes and decodes correctly") + func testBoolValue() throws { + let value = ElicitValue.bool(true) + let data = try JSONEncoder().encode(value) + let decoded = try JSONDecoder().decode(ElicitValue.self, from: data) + #expect(decoded.boolValue == true) + } + + @Test("Strings array value encodes and decodes correctly") + func testStringsValue() throws { + let value = ElicitValue.strings(["a", "b", "c"]) + let data = try JSONEncoder().encode(value) + let decoded = try JSONDecoder().decode(ElicitValue.self, from: data) + #expect(decoded.stringsValue == ["a", "b", "c"]) + } +} + +// MARK: - ElicitAction Tests + +@Suite("ElicitAction Tests") +struct ElicitActionTests { + + @Test("All action types encode and decode correctly") + func testActionTypes() throws { + let actions: [ElicitAction] = [.accept, .decline, .cancel] + + for action in actions { + let data = try JSONEncoder().encode(action) + let decoded = try JSONDecoder().decode(ElicitAction.self, from: data) + #expect(decoded == action) + } + } +} + +// MARK: - Capability Tests + +@Suite("Elicitation Capability Tests") +struct ElicitationCapabilityTests { + + @Test("Form capability encodes correctly") + func testFormCapability() throws { + let capability = Client.Capabilities.Elicitation( + form: Client.Capabilities.Elicitation.Form(applyDefaults: true) + ) + + let data = try JSONEncoder().encode(capability) + let decoded = try JSONDecoder().decode(Client.Capabilities.Elicitation.self, from: data) + + #expect(decoded.form?.applyDefaults == true) + } + + @Test("URL capability encodes correctly") + func testURLCapability() throws { + let capability = Client.Capabilities.Elicitation( + url: Client.Capabilities.Elicitation.URL() + ) + + let data = try JSONEncoder().encode(capability) + let decoded = try JSONDecoder().decode(Client.Capabilities.Elicitation.self, from: data) + + #expect(decoded.url != nil) + } + + @Test("Combined form and URL capability encodes correctly") + func testCombinedCapability() throws { + let capability = Client.Capabilities.Elicitation( + form: Client.Capabilities.Elicitation.Form(), + url: Client.Capabilities.Elicitation.URL() + ) + + let data = try JSONEncoder().encode(capability) + let decoded = try JSONDecoder().decode(Client.Capabilities.Elicitation.self, from: data) + + #expect(decoded.form != nil) + #expect(decoded.url != nil) + } +} + +// MARK: - ElicitationRequiredErrorData Tests + +@Suite("ElicitationRequiredErrorData Tests") +struct ElicitationRequiredErrorDataTests { + + @Test("Error data encodes correctly") + func testErrorDataEncoding() throws { + let errorData = ElicitationRequiredErrorData( + elicitations: [ + ElicitRequestURLParams( + message: "Authorize", + elicitationId: "auth-1", + url: "https://example.com/auth" + ) + ] + ) + + let data = try JSONEncoder().encode(errorData) + let decoded = try JSONDecoder().decode(ElicitationRequiredErrorData.self, from: data) + + #expect(decoded.elicitations.count == 1) + #expect(decoded.elicitations[0].elicitationId == "auth-1") + #expect(decoded.elicitations[0].url == "https://example.com/auth") + } +} + +// MARK: - Server-Client Integration Tests + +@Suite("Elicitation Integration Tests") +struct ElicitationIntegrationTests { + + // MARK: - Form Mode Flow Tests + + @Test("Server can elicit form input from client - accept with content") + func testFormElicitationAccept() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ElicitTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "askName", description: "Ask for name", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] _, _ in + let result = try await server.elicit(ElicitRequestParams.form(ElicitRequestFormParams( + message: "Please enter your name", + requestedSchema: ElicitationSchema( + properties: ["name": .string(StringSchema(title: "Name"))], + required: ["name"] + ) + ))) + + if result.action == .accept, let name = result.content?["name"]?.stringValue { + return CallTool.Result(content: [.text("Hello, \(name)!")]) + } else { + return CallTool.Result(content: [.text("No name provided")], isError: true) + } + } + + let client = Client(name: "ElicitTestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(form: Client.Capabilities.Elicitation.Form()) + )) + + await client.withElicitationHandler { params, _ in + guard case .form(let formParams) = params else { + return ElicitResult(action: .decline) + } + #expect(formParams.message == "Please enter your name") + return ElicitResult(action: .accept, content: ["name": .string("Alice")]) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "askName", arguments: [:]) + + #expect(result.isError == nil) + if case .text(let text, _, _) = result.content[0] { + #expect(text == "Hello, Alice!") + } + + await client.disconnect() + } + + @Test("Server can elicit form input from client - user declines") + func testFormElicitationDecline() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ElicitTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "confirm", description: "Confirm action", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] _, _ in + let result = try await server.elicit(ElicitRequestParams.form(ElicitRequestFormParams( + message: "Confirm action?", + requestedSchema: ElicitationSchema( + properties: ["confirm": .boolean(BooleanSchema(title: "Confirm"))] + ) + ))) + + return switch result.action { + case .accept: CallTool.Result(content: [.text("Confirmed")]) + case .decline: CallTool.Result(content: [.text("Declined")]) + case .cancel: CallTool.Result(content: [.text("Cancelled")]) + } + } + + let client = Client(name: "ElicitTestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(form: Client.Capabilities.Elicitation.Form()) + )) + + await client.withElicitationHandler { _, _ in + return ElicitResult(action: .decline) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "confirm", arguments: [:]) + + if case .text(let text, _, _) = result.content[0] { + #expect(text == "Declined") + } + + await client.disconnect() + } + + @Test("Server can elicit form input from client - user cancels") + func testFormElicitationCancel() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ElicitTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "getData", description: "Get data", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] _, _ in + let result = try await server.elicit(ElicitRequestParams.form(ElicitRequestFormParams( + message: "Enter data", + requestedSchema: ElicitationSchema( + properties: ["data": .string(StringSchema())] + ) + ))) + + return switch result.action { + case .accept: CallTool.Result(content: [.text("Got data")]) + case .decline: CallTool.Result(content: [.text("Declined")]) + case .cancel: CallTool.Result(content: [.text("Cancelled")]) + } + } + + let client = Client(name: "ElicitTestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(form: Client.Capabilities.Elicitation.Form()) + )) + + await client.withElicitationHandler { _, _ in + return ElicitResult(action: .cancel) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "getData", arguments: [:]) + + if case .text(let text, _, _) = result.content[0] { + #expect(text == "Cancelled") + } + + await client.disconnect() + } + + @Test("Form elicitation with multiple field types") + func testFormElicitationMultipleFieldTypes() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ElicitTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "survey", description: "Survey", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] _, _ in + let result = try await server.elicit(ElicitRequestParams.form(ElicitRequestFormParams( + message: "Complete the survey", + requestedSchema: ElicitationSchema( + properties: [ + "name": .string(StringSchema(title: "Name")), + "age": .number(NumberSchema(isInteger: true, title: "Age")), + "score": .number(NumberSchema(isInteger: false, title: "Score")), + "subscribe": .boolean(BooleanSchema(title: "Subscribe")), + ], + required: ["name"] + ) + ))) + + guard result.action == .accept, let content = result.content else { + return CallTool.Result(content: [.text("No response")]) + } + + let name = content["name"]?.stringValue ?? "unknown" + let age = content["age"]?.intValue ?? 0 + let score = content["score"]?.doubleValue ?? 0.0 + let subscribe = content["subscribe"]?.boolValue ?? false + + return CallTool.Result(content: [.text( + "Name: \(name), Age: \(age), Score: \(score), Subscribe: \(subscribe)" + )]) + } + + let client = Client(name: "ElicitTestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(form: Client.Capabilities.Elicitation.Form()) + )) + + await client.withElicitationHandler { _, _ in + return ElicitResult( + action: .accept, + content: [ + "name": .string("Bob"), + "age": .int(30), + "score": .double(95.5), + "subscribe": .bool(true), + ] + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "survey", arguments: [:]) + + if case .text(let text, _, _) = result.content[0] { + #expect(text.contains("Name: Bob")) + #expect(text.contains("Age: 30")) + #expect(text.contains("Score: 95.5")) + #expect(text.contains("Subscribe: true")) + } + + await client.disconnect() + } + + // MARK: - URL Mode Flow Tests + + @Test("Server can elicit URL authorization from client - accept") + func testURLElicitationAccept() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ElicitTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "authorize", description: "Authorize", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] _, _ in + let result = try await server.elicit(ElicitRequestParams.url(ElicitRequestURLParams( + message: "Please authorize access", + elicitationId: "auth-123", + url: "https://example.com/oauth" + ))) + + return switch result.action { + case .accept: CallTool.Result(content: [.text("Authorized")]) + case .decline: CallTool.Result(content: [.text("Declined")]) + case .cancel: CallTool.Result(content: [.text("Cancelled")]) + } + } + + let client = Client(name: "ElicitTestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(url: Client.Capabilities.Elicitation.URL()) + )) + + await client.withElicitationHandler { params, _ in + guard case .url(let urlParams) = params else { + return ElicitResult(action: .decline) + } + #expect(urlParams.message == "Please authorize access") + #expect(urlParams.elicitationId == "auth-123") + #expect(urlParams.url == "https://example.com/oauth") + return ElicitResult(action: .accept) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "authorize", arguments: [:]) + + if case .text(let text, _, _) = result.content[0] { + #expect(text == "Authorized") + } + + await client.disconnect() + } + + // MARK: - Capability Checking Tests + + @Test("Server rejects form elicitation when client only supports URL mode") + func testFormElicitationRejectedWhenClientOnlySupportsURL() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ElicitTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "formTool", description: "Tool requiring form", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] _, _ in + do { + _ = try await server.elicit(ElicitRequestParams.form(ElicitRequestFormParams( + message: "This should fail", + requestedSchema: ElicitationSchema(properties: [:]) + ))) + return CallTool.Result(content: [.text("Should not reach here")]) + } catch let error as MCPError { + return CallTool.Result(content: [.text("Error: \(error)")], isError: true) + } + } + + // Client only supports URL mode, not form + let client = Client(name: "ElicitTestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(url: Client.Capabilities.Elicitation.URL()) + )) + + await client.withElicitationHandler { _, _ in + Issue.record("Handler should not be called") + return ElicitResult(action: .decline) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "formTool", arguments: [:]) + + #expect(result.isError == true) + + await client.disconnect() + } + + @Test("Server rejects URL elicitation when client only supports form mode") + func testURLElicitationRejectedWhenClientOnlySupportsForm() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ElicitTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "urlTool", description: "Tool requiring URL", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] _, _ in + do { + _ = try await server.elicit(ElicitRequestParams.url(ElicitRequestURLParams( + message: "This should fail", + elicitationId: "test", + url: "https://example.com/auth" + ))) + return CallTool.Result(content: [.text("Should not reach here")]) + } catch let error as MCPError { + return CallTool.Result(content: [.text("Error: \(error)")], isError: true) + } + } + + // Client only supports form mode, not URL + let client = Client(name: "ElicitTestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(form: Client.Capabilities.Elicitation.Form()) + )) + + await client.withElicitationHandler { _, _ in + Issue.record("Handler should not be called") + return ElicitResult(action: .decline) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "urlTool", arguments: [:]) + + #expect(result.isError == true) + + await client.disconnect() + } + + @Test("Server rejects elicitation when client has no elicitation capability") + func testElicitationRejectedWhenNoCapability() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ElicitTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "elicitTool", description: "Tool that elicits", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] _, _ in + do { + _ = try await server.elicit(ElicitRequestParams.form(ElicitRequestFormParams( + message: "This should fail", + requestedSchema: ElicitationSchema(properties: [:]) + ))) + return CallTool.Result(content: [.text("Should not reach here")]) + } catch let error as MCPError { + return CallTool.Result(content: [.text("Error: \(error)")], isError: true) + } + } + + // Client has no elicitation capability + let client = Client(name: "ElicitTestClient", version: "1.0.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "elicitTool", arguments: [:]) + + #expect(result.isError == true) + + await client.disconnect() + } + + // MARK: - Multiple Elicitations Tests + + @Test("Server can perform multiple sequential elicitations") + func testMultipleSequentialElicitations() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ElicitTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "wizard", description: "Multi-step wizard", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] _, _ in + // Step 1: Get name + let step1 = try await server.elicit(ElicitRequestParams.form(ElicitRequestFormParams( + message: "Step 1: Enter name", + requestedSchema: ElicitationSchema( + properties: ["name": .string(StringSchema())], + required: ["name"] + ) + ))) + + guard step1.action == .accept, let name = step1.content?["name"]?.stringValue else { + return CallTool.Result(content: [.text("Cancelled at step 1")]) + } + + // Step 2: Get age + let step2 = try await server.elicit(ElicitRequestParams.form(ElicitRequestFormParams( + message: "Step 2: Enter age", + requestedSchema: ElicitationSchema( + properties: ["age": .number(NumberSchema(isInteger: true))], + required: ["age"] + ) + ))) + + guard step2.action == .accept, let age = step2.content?["age"]?.intValue else { + return CallTool.Result(content: [.text("Cancelled at step 2")]) + } + + return CallTool.Result(content: [.text("Completed: \(name), \(age)")]) + } + + let client = Client(name: "ElicitTestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(form: Client.Capabilities.Elicitation.Form()) + )) + + actor Counter { + var count = 0 + func increment() { count += 1 } + func getCount() -> Int { count } + } + let counter = Counter() + + await client.withElicitationHandler { params, _ in + await counter.increment() + guard case .form(let formParams) = params else { + return ElicitResult(action: .decline) + } + + if formParams.message.contains("Step 1") { + return ElicitResult(action: .accept, content: ["name": .string("Charlie")]) + } else if formParams.message.contains("Step 2") { + return ElicitResult(action: .accept, content: ["age": .int(25)]) + } + + return ElicitResult(action: .decline) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "wizard", arguments: [:]) + + #expect(await counter.getCount() == 2) + if case .text(let text, _, _) = result.content[0] { + #expect(text == "Completed: Charlie, 25") + } + + await client.disconnect() + } + + // MARK: - Enum Schema Tests + + @Test("Form elicitation with titled enum") + func testFormElicitationWithTitledEnum() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ElicitTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "pickColor", description: "Pick a color", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] _, _ in + let result = try await server.elicit(ElicitRequestParams.form(ElicitRequestFormParams( + message: "Choose your color", + requestedSchema: ElicitationSchema( + properties: [ + "color": .titledEnum(TitledEnumSchema( + title: "Color", + oneOf: [ + TitledEnumOption(const: "#FF0000", title: "Red"), + TitledEnumOption(const: "#00FF00", title: "Green"), + TitledEnumOption(const: "#0000FF", title: "Blue"), + ] + )) + ], + required: ["color"] + ) + ))) + + guard result.action == .accept, let color = result.content?["color"]?.stringValue else { + return CallTool.Result(content: [.text("No color")]) + } + + return CallTool.Result(content: [.text("Color: \(color)")]) + } + + let client = Client(name: "ElicitTestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(form: Client.Capabilities.Elicitation.Form()) + )) + + await client.withElicitationHandler { params, _ in + guard case .form(let formParams) = params else { + return ElicitResult(action: .decline) + } + + // Verify enum options are present + if case .titledEnum(let enumSchema) = formParams.requestedSchema.properties["color"] { + #expect(enumSchema.oneOf.count == 3) + #expect(enumSchema.oneOf[0].title == "Red") + } + + return ElicitResult(action: .accept, content: ["color": .string("#00FF00")]) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "pickColor", arguments: [:]) + + if case .text(let text, _, _) = result.content[0] { + #expect(text == "Color: #00FF00") + } + + await client.disconnect() + } + + @Test("Form elicitation with multi-select enum") + func testFormElicitationWithMultiSelect() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ElicitTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "selectTags", description: "Select tags", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] _, _ in + let result = try await server.elicit(ElicitRequestParams.form(ElicitRequestFormParams( + message: "Select your interests", + requestedSchema: ElicitationSchema( + properties: [ + "interests": .titledMultiSelect(TitledMultiSelectEnumSchema( + title: "Interests", + options: [ + TitledEnumOption(const: "tech", title: "Technology"), + TitledEnumOption(const: "sports", title: "Sports"), + TitledEnumOption(const: "music", title: "Music"), + ] + )) + ] + ) + ))) + + guard result.action == .accept, + let interests = result.content?["interests"]?.stringsValue else { + return CallTool.Result(content: [.text("No interests")]) + } + + return CallTool.Result(content: [.text("Interests: \(interests.joined(separator: ", "))")]) + } + + let client = Client(name: "ElicitTestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(form: Client.Capabilities.Elicitation.Form()) + )) + + await client.withElicitationHandler { _, _ in + return ElicitResult(action: .accept, content: ["interests": .strings(["tech", "music"])]) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "selectTags", arguments: [:]) + + if case .text(let text, _, _) = result.content[0] { + #expect(text == "Interests: tech, music") + } + + await client.disconnect() + } + + // MARK: - URL Mode Decline/Cancel Tests + + @Test("Server can elicit URL authorization from client - decline") + func testURLElicitationDecline() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ElicitTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "authorize", description: "Authorize", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] _, _ in + let result = try await server.elicit(ElicitRequestParams.url(ElicitRequestURLParams( + message: "Please authorize access", + elicitationId: "auth-decline-123", + url: "https://example.com/oauth" + ))) + + return switch result.action { + case .accept: CallTool.Result(content: [.text("Authorized")]) + case .decline: CallTool.Result(content: [.text("User declined authorization")]) + case .cancel: CallTool.Result(content: [.text("Cancelled")]) + } + } + + let client = Client(name: "ElicitTestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(url: Client.Capabilities.Elicitation.URL()) + )) + + await client.withElicitationHandler { params, _ in + guard case .url(let urlParams) = params else { + return ElicitResult(action: .cancel) + } + #expect(urlParams.elicitationId == "auth-decline-123") + return ElicitResult(action: .decline) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "authorize", arguments: [:]) + + if case .text(let text, _, _) = result.content[0] { + #expect(text == "User declined authorization") + } + + await client.disconnect() + } + + @Test("Server can elicit URL authorization from client - cancel") + func testURLElicitationCancel() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ElicitTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "authorize", description: "Authorize", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] _, _ in + let result = try await server.elicit(ElicitRequestParams.url(ElicitRequestURLParams( + message: "Please authorize access", + elicitationId: "auth-cancel-456", + url: "https://example.com/oauth" + ))) + + return switch result.action { + case .accept: CallTool.Result(content: [.text("Authorized")]) + case .decline: CallTool.Result(content: [.text("Declined")]) + case .cancel: CallTool.Result(content: [.text("User cancelled authorization")]) + } + } + + let client = Client(name: "ElicitTestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(url: Client.Capabilities.Elicitation.URL()) + )) + + await client.withElicitationHandler { params, _ in + guard case .url = params else { + return ElicitResult(action: .decline) + } + return ElicitResult(action: .cancel) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "authorize", arguments: [:]) + + if case .text(let text, _, _) = result.content[0] { + #expect(text == "User cancelled authorization") + } + + await client.disconnect() + } + + @Test("URL mode elicitation response should not include content") + func testURLElicitationNoContent() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ElicitTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "checkContent", description: "Check content", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] _, _ in + let result = try await server.elicit(ElicitRequestParams.url(ElicitRequestURLParams( + message: "Complete authorization", + elicitationId: "content-check-789", + url: "https://example.com/auth" + ))) + + // URL mode responses should not have content + let hasContent = result.content != nil + return CallTool.Result(content: [.text("Action: \(result.action.rawValue), HasContent: \(hasContent)")]) + } + + let client = Client(name: "ElicitTestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(url: Client.Capabilities.Elicitation.URL()) + )) + + await client.withElicitationHandler { params, _ in + guard case .url = params else { + return ElicitResult(action: .decline) + } + // URL mode should return accept without content + return ElicitResult(action: .accept) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "checkContent", arguments: [:]) + + if case .text(let text, _, _) = result.content[0] { + #expect(text == "Action: accept, HasContent: false") + } + + await client.disconnect() + } + + // MARK: - Legacy Enum Format Tests + + @Test("Form elicitation with legacy enumNames format") + func testFormElicitationWithLegacyEnumNames() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ElicitTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "pickColor", description: "Pick a color", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] _, _ in + let result = try await server.elicit(ElicitRequestParams.form(ElicitRequestFormParams( + message: "Choose your color", + requestedSchema: ElicitationSchema( + properties: [ + "color": .legacyTitledEnum(LegacyTitledEnumSchema( + title: "Color", + description: "Choose your favorite color", + enumValues: ["#FF0000", "#00FF00", "#0000FF"], + enumNames: ["Red", "Green", "Blue"], + defaultValue: "#00FF00" + )) + ], + required: ["color"] + ) + ))) + + guard result.action == .accept, let color = result.content?["color"]?.stringValue else { + return CallTool.Result(content: [.text("No color selected")]) + } + + return CallTool.Result(content: [.text("Selected color: \(color)")]) + } + + let client = Client(name: "ElicitTestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(form: Client.Capabilities.Elicitation.Form()) + )) + + await client.withElicitationHandler { params, _ in + guard case .form(let formParams) = params else { + return ElicitResult(action: .decline) + } + + // Verify legacy enum schema is decoded correctly + if case .legacyTitledEnum(let enumSchema) = formParams.requestedSchema.properties["color"] { + #expect(enumSchema.enumValues == ["#FF0000", "#00FF00", "#0000FF"]) + #expect(enumSchema.enumNames == ["Red", "Green", "Blue"]) + #expect(enumSchema.defaultValue == "#00FF00") + } + + // Return the const value, not the display name + return ElicitResult(action: .accept, content: ["color": .string("#FF0000")]) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "pickColor", arguments: [:]) + + if case .text(let text, _, _) = result.content[0] { + #expect(text == "Selected color: #FF0000") + } + + await client.disconnect() + } + + // MARK: - Optional Fields Tests + + @Test("Form elicitation with optional fields - all provided") + func testFormElicitationOptionalFieldsAllProvided() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ElicitTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "userInfo", description: "Get user info", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] _, _ in + let result = try await server.elicit(ElicitRequestParams.form(ElicitRequestFormParams( + message: "Please provide your information", + requestedSchema: ElicitationSchema( + properties: [ + "name": .string(StringSchema(title: "Name", description: "Your name (required)")), + "nickname": .string(StringSchema(title: "Nickname", description: "Optional nickname")), + "age": .number(NumberSchema(isInteger: true, title: "Age", description: "Optional age")), + "subscribe": .boolean(BooleanSchema(title: "Subscribe", description: "Optional subscription")), + ], + required: ["name"] // Only name is required + ) + ))) + + guard result.action == .accept, let content = result.content else { + return CallTool.Result(content: [.text("No response")]) + } + + var parts: [String] = [] + if let name = content["name"]?.stringValue { + parts.append("Name: \(name)") + } + if let nickname = content["nickname"]?.stringValue { + parts.append("Nickname: \(nickname)") + } + if let age = content["age"]?.intValue { + parts.append("Age: \(age)") + } + if let subscribe = content["subscribe"]?.boolValue { + parts.append("Subscribe: \(subscribe)") + } + + return CallTool.Result(content: [.text(parts.joined(separator: ", "))]) + } + + let client = Client(name: "ElicitTestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(form: Client.Capabilities.Elicitation.Form()) + )) + + await client.withElicitationHandler { _, _ in + return ElicitResult( + action: .accept, + content: [ + "name": .string("John Doe"), + "nickname": .string("Johnny"), + "age": .int(30), + "subscribe": .bool(true), + ] + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "userInfo", arguments: [:]) + + if case .text(let text, _, _) = result.content[0] { + #expect(text.contains("Name: John Doe")) + #expect(text.contains("Nickname: Johnny")) + #expect(text.contains("Age: 30")) + #expect(text.contains("Subscribe: true")) + } + + await client.disconnect() + } + + @Test("Form elicitation with optional fields - only required provided") + func testFormElicitationOptionalFieldsOnlyRequired() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ElicitTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "userInfo", description: "Get user info", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] _, _ in + let result = try await server.elicit(ElicitRequestParams.form(ElicitRequestFormParams( + message: "Please provide your information", + requestedSchema: ElicitationSchema( + properties: [ + "name": .string(StringSchema(title: "Name")), + "nickname": .string(StringSchema(title: "Nickname")), + "email": .string(StringSchema(title: "Email", format: .email)), + ], + required: ["name"] + ) + ))) + + guard result.action == .accept, let content = result.content else { + return CallTool.Result(content: [.text("No response")]) + } + + let name = content["name"]?.stringValue ?? "unknown" + let hasNickname = content["nickname"] != nil + let hasEmail = content["email"] != nil + + return CallTool.Result(content: [.text("Name: \(name), HasNickname: \(hasNickname), HasEmail: \(hasEmail)")]) + } + + let client = Client(name: "ElicitTestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(form: Client.Capabilities.Elicitation.Form()) + )) + + await client.withElicitationHandler { _, _ in + // Only provide the required field + return ElicitResult(action: .accept, content: ["name": .string("Jane Smith")]) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "userInfo", arguments: [:]) + + if case .text(let text, _, _) = result.content[0] { + #expect(text == "Name: Jane Smith, HasNickname: false, HasEmail: false") + } + + await client.disconnect() + } + + // MARK: - Default Values Tests + + @Test("Schema default values are included in encoded JSON") + func testSchemaDefaultValuesEncoding() throws { + let schema = ElicitationSchema( + properties: [ + "name": .string(StringSchema(title: "Name", defaultValue: "Guest")), + "age": .number(NumberSchema(isInteger: true, title: "Age", defaultValue: 18)), + "subscribe": .boolean(BooleanSchema(title: "Subscribe", defaultValue: true)), + "color": .untitledEnum(UntitledEnumSchema( + title: "Color", + enumValues: ["red", "green", "blue"], + defaultValue: "green" + )), + "interests": .untitledMultiSelect(UntitledMultiSelectEnumSchema( + title: "Interests", + enumValues: ["tech", "sports", "music"], + defaultValue: ["tech", "music"] + )), + ], + required: ["name"] + ) + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + let data = try encoder.encode(schema) + let json = String(data: data, encoding: .utf8)! + + // Verify defaults are present in the encoded JSON + #expect(json.contains("\"default\":\"Guest\"")) + #expect(json.contains("\"default\":18")) + #expect(json.contains("\"default\":true")) + #expect(json.contains("\"default\":\"green\"")) + #expect(json.contains("\"default\":[\"tech\",\"music\"]")) + } + + @Test("Form elicitation with default values in schema") + func testFormElicitationWithDefaultValues() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ElicitTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "preferences", description: "Get preferences", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] _, _ in + let result = try await server.elicit(ElicitRequestParams.form(ElicitRequestFormParams( + message: "Set your preferences", + requestedSchema: ElicitationSchema( + properties: [ + "email": .string(StringSchema(title: "Email", format: .email)), + "nickname": .string(StringSchema(title: "Nickname", defaultValue: "Guest")), + "volume": .number(NumberSchema(isInteger: true, title: "Volume", minimum: 0, maximum: 100, defaultValue: 50)), + "darkMode": .boolean(BooleanSchema(title: "Dark Mode", defaultValue: false)), + ], + required: ["email"] + ) + ))) + + guard result.action == .accept, let content = result.content else { + return CallTool.Result(content: [.text("No response")]) + } + + let email = content["email"]?.stringValue ?? "none" + let nickname = content["nickname"]?.stringValue ?? "Guest" // Use default if not provided + let volume = content["volume"]?.intValue ?? 50 // Use default if not provided + let darkMode = content["darkMode"]?.boolValue ?? false // Use default if not provided + + return CallTool.Result(content: [.text("Email: \(email), Nickname: \(nickname), Volume: \(volume), DarkMode: \(darkMode)")]) + } + + let client = Client(name: "ElicitTestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(form: Client.Capabilities.Elicitation.Form()) + )) + + await client.withElicitationHandler { params, _ in + // Verify the schema contains default values + guard case .form(let formParams) = params else { + return ElicitResult(action: .decline) + } + + if case .string(let nicknameSchema) = formParams.requestedSchema.properties["nickname"] { + #expect(nicknameSchema.defaultValue == "Guest") + } + if case .number(let volumeSchema) = formParams.requestedSchema.properties["volume"] { + #expect(volumeSchema.defaultValue == 50) + } + if case .boolean(let darkModeSchema) = formParams.requestedSchema.properties["darkMode"] { + #expect(darkModeSchema.defaultValue == false) + } + + // Client provides only email, using defaults for others + return ElicitResult(action: .accept, content: ["email": .string("test@example.com")]) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "preferences", arguments: [:]) + + if case .text(let text, _, _) = result.content[0] { + #expect(text == "Email: test@example.com, Nickname: Guest, Volume: 50, DarkMode: false") + } + + await client.disconnect() + } + + // MARK: - Complex Schema Tests (matching TypeScript SDK) + + @Test("Form elicitation with complex object - multiple fields like TypeScript SDK") + func testFormElicitationComplexObject() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ElicitTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "userProfile", description: "Get user profile", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] _, _ in + let result = try await server.elicit(ElicitRequestParams.form(ElicitRequestFormParams( + message: "Please provide your information", + requestedSchema: ElicitationSchema( + properties: [ + "name": .string(StringSchema(title: "Name", minLength: 1)), + "email": .string(StringSchema(title: "Email", format: .email)), + "age": .number(NumberSchema(isInteger: true, title: "Age", minimum: 0, maximum: 150)), + "street": .string(StringSchema(title: "Street")), + "city": .string(StringSchema(title: "City")), + "zipCode": .string(StringSchema(title: "Zip Code")), + "newsletter": .boolean(BooleanSchema(title: "Newsletter")), + "notifications": .boolean(BooleanSchema(title: "Notifications")), + ], + required: ["name", "email", "age", "street", "city", "zipCode"] + ) + ))) + + guard result.action == .accept, let content = result.content else { + return CallTool.Result(content: [.text("No response")]) + } + + let name = content["name"]?.stringValue ?? "" + let email = content["email"]?.stringValue ?? "" + let age = content["age"]?.intValue ?? 0 + let city = content["city"]?.stringValue ?? "" + let zipCode = content["zipCode"]?.stringValue ?? "" + let newsletter = content["newsletter"]?.boolValue ?? false + + return CallTool.Result(content: [.text( + "\(name), \(email), \(age), \(city), \(zipCode), newsletter=\(newsletter)" + )]) + } + + let client = Client(name: "ElicitTestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(form: Client.Capabilities.Elicitation.Form()) + )) + + await client.withElicitationHandler { _, _ in + return ElicitResult( + action: .accept, + content: [ + "name": .string("Jane Smith"), + "email": .string("jane@example.com"), + "age": .int(28), + "street": .string("123 Main St"), + "city": .string("San Francisco"), + "zipCode": .string("94105"), + "newsletter": .bool(true), + "notifications": .bool(false), + ] + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "userProfile", arguments: [:]) + + if case .text(let text, _, _) = result.content[0] { + #expect(text == "Jane Smith, jane@example.com, 28, San Francisco, 94105, newsletter=true") + } + + await client.disconnect() + } +} + +// MARK: - Additional Schema Encoding Tests + +@Suite("LegacyTitledEnumSchema Tests") +struct LegacyTitledEnumSchemaTests { + + @Test("LegacyTitledEnumSchema encodes and decodes correctly") + func testLegacyTitledEnumSchemaRoundtrip() throws { + let schema = LegacyTitledEnumSchema( + title: "Priority", + description: "Select priority level", + enumValues: ["low", "medium", "high"], + enumNames: ["Low Priority", "Medium Priority", "High Priority"], + defaultValue: "medium" + ) + + let encoder = JSONEncoder() + encoder.outputFormatting = .sortedKeys + let data = try encoder.encode(schema) + let decoded = try JSONDecoder().decode(LegacyTitledEnumSchema.self, from: data) + + #expect(decoded.title == "Priority") + #expect(decoded.description == "Select priority level") + #expect(decoded.enumValues == ["low", "medium", "high"]) + #expect(decoded.enumNames == ["Low Priority", "Medium Priority", "High Priority"]) + #expect(decoded.defaultValue == "medium") + } + + @Test("LegacyTitledEnumSchema decodes via PrimitiveSchemaDefinition") + func testLegacyTitledEnumSchemaDecodingViaPrimitive() throws { + let json = """ + { + "type": "string", + "title": "Status", + "enum": ["active", "inactive", "pending"], + "enumNames": ["Active", "Inactive", "Pending"], + "default": "pending" + } + """ + + let data = json.data(using: .utf8)! + let decoded = try JSONDecoder().decode(PrimitiveSchemaDefinition.self, from: data) + + guard case .legacyTitledEnum(let schema) = decoded else { + Issue.record("Expected legacyTitledEnum") + return + } + + #expect(schema.title == "Status") + #expect(schema.enumValues == ["active", "inactive", "pending"]) + #expect(schema.enumNames == ["Active", "Inactive", "Pending"]) + #expect(schema.defaultValue == "pending") + } +} + +// MARK: - ElicitationCompleteNotification Tests + +@Suite("ElicitationCompleteNotification Tests") +struct ElicitationCompleteNotificationTests { + + @Test("ElicitationCompleteNotification name is correct") + func testNotificationName() { + #expect(ElicitationCompleteNotification.name == "notifications/elicitation/complete") + } + + @Test("ElicitationCompleteNotification parameters encode correctly") + func testNotificationParametersEncoding() throws { + let params = ElicitationCompleteNotification.Parameters( + elicitationId: "test-elicitation-123" + ) + + let data = try JSONEncoder().encode(params) + let decoded = try JSONDecoder().decode(ElicitationCompleteNotification.Parameters.self, from: data) + + #expect(decoded.elicitationId == "test-elicitation-123") + } + + @Test("ElicitationCompleteNotification message encoding") + func testNotificationMessageEncoding() throws { + let notification = ElicitationCompleteNotification.message( + ElicitationCompleteNotification.Parameters(elicitationId: "complete-456") + ) + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + let data = try encoder.encode(notification) + let json = String(data: data, encoding: .utf8)! + + #expect(json.contains("\"method\":\"notifications/elicitation/complete\"")) + #expect(json.contains("\"elicitationId\":\"complete-456\"")) + #expect(json.contains("\"jsonrpc\":\"2.0\"")) + } +} + +// MARK: - Untitled Enum Schema Tests + +@Suite("UntitledEnumSchema Tests") +struct UntitledEnumSchemaTests { + + @Test("UntitledEnumSchema with minItems/maxItems constraints") + func testUntitledMultiSelectWithConstraints() throws { + let schema = UntitledMultiSelectEnumSchema( + title: "Tags", + description: "Select 1-3 tags", + minItems: 1, + maxItems: 3, + enumValues: ["tag1", "tag2", "tag3", "tag4", "tag5"], + defaultValue: ["tag1"] + ) + + let encoder = JSONEncoder() + encoder.outputFormatting = .sortedKeys + let data = try encoder.encode(schema) + let decoded = try JSONDecoder().decode(UntitledMultiSelectEnumSchema.self, from: data) + + #expect(decoded.title == "Tags") + #expect(decoded.minItems == 1) + #expect(decoded.maxItems == 3) + #expect(decoded.items.enumValues == ["tag1", "tag2", "tag3", "tag4", "tag5"]) + #expect(decoded.defaultValue == ["tag1"]) + } +} + +// MARK: - Additional Capability Tests + +@Suite("Elicitation Capability applyDefaults Tests") +struct ElicitationCapabilityApplyDefaultsTests { + + @Test("Form capability with applyDefaults encodes correctly") + func testFormCapabilityWithApplyDefaults() throws { + let capability = Client.Capabilities.Elicitation( + form: Client.Capabilities.Elicitation.Form(applyDefaults: true) + ) + + let data = try JSONEncoder().encode(capability) + let json = String(data: data, encoding: .utf8)! + + #expect(json.contains("applyDefaults")) + #expect(json.contains("true")) + + let decoded = try JSONDecoder().decode(Client.Capabilities.Elicitation.self, from: data) + #expect(decoded.form?.applyDefaults == true) + } + + @Test("Form capability with applyDefaults false encodes correctly") + func testFormCapabilityWithApplyDefaultsFalse() throws { + let capability = Client.Capabilities.Elicitation( + form: Client.Capabilities.Elicitation.Form(applyDefaults: false) + ) + + let data = try JSONEncoder().encode(capability) + let decoded = try JSONDecoder().decode(Client.Capabilities.Elicitation.self, from: data) + + #expect(decoded.form?.applyDefaults == false) + } +} + +// MARK: - URLElicitationRequiredError Tests + +@Suite("URLElicitationRequiredError Tests") +struct URLElicitationRequiredErrorTests { + + @Test("MCPError.urlElicitationRequired creates error with correct code") + func testErrorCode() { + let error = MCPError.urlElicitationRequired( + elicitations: [ + ElicitRequestURLParams( + message: "Please authorize", + elicitationId: "auth-123", + url: "https://example.com/oauth" + ) + ] + ) + + #expect(error.code == ErrorCode.urlElicitationRequired) + } + + @Test("MCPError.urlElicitationRequired default message") + func testDefaultMessage() { + let singleError = MCPError.urlElicitationRequired( + elicitations: [ + ElicitRequestURLParams( + message: "Authorize", + elicitationId: "auth-1", + url: "https://example.com/auth" + ) + ] + ) + + #expect(singleError.errorDescription == "URL elicitation required") + + let multipleError = MCPError.urlElicitationRequired( + elicitations: [ + ElicitRequestURLParams(message: "Auth 1", elicitationId: "a1", url: "https://example.com/1"), + ElicitRequestURLParams(message: "Auth 2", elicitationId: "a2", url: "https://example.com/2"), + ] + ) + + #expect(multipleError.errorDescription == "URL elicitations required") + } + + @Test("MCPError.urlElicitationRequired custom message") + func testCustomMessage() { + let error = MCPError.urlElicitationRequired( + elicitations: [ + ElicitRequestURLParams( + message: "Authorize", + elicitationId: "auth-1", + url: "https://example.com/auth" + ) + ], + message: "Custom authorization required" + ) + + #expect(error.errorDescription == "Custom authorization required") + } + + @Test("MCPError.urlElicitationRequired elicitations accessor") + func testElicitationsAccessor() { + let elicitations = [ + ElicitRequestURLParams(message: "Auth 1", elicitationId: "a1", url: "https://example.com/1"), + ElicitRequestURLParams(message: "Auth 2", elicitationId: "a2", url: "https://example.com/2"), + ] + + let error = MCPError.urlElicitationRequired(elicitations: elicitations) + + #expect(error.elicitations?.count == 2) + #expect(error.elicitations?[0].elicitationId == "a1") + #expect(error.elicitations?[1].elicitationId == "a2") + } + + @Test("MCPError.urlElicitationRequired encodes correctly to JSON") + func testEncoding() throws { + let error = MCPError.urlElicitationRequired( + elicitations: [ + ElicitRequestURLParams( + message: "Please authorize", + elicitationId: "auth-123", + url: "https://example.com/oauth" + ) + ], + message: "Authorization required" + ) + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + let data = try encoder.encode(error) + let json = String(data: data, encoding: .utf8)! + + #expect(json.contains("\"code\":-32042")) + #expect(json.contains("\"message\":\"Authorization required\"")) + #expect(json.contains("\"elicitations\"")) + #expect(json.contains("\"elicitationId\":\"auth-123\"")) + #expect(json.contains("\"url\":\"https://example.com/oauth\"")) + } + + @Test("MCPError.urlElicitationRequired decodes correctly from JSON") + func testDecoding() throws { + let json = """ + { + "code": -32042, + "message": "Authorization required", + "data": { + "elicitations": [ + { + "mode": "url", + "message": "Please authorize", + "elicitationId": "auth-456", + "url": "https://example.com/authorize" + } + ] + } + } + """ + + let data = json.data(using: .utf8)! + let error = try JSONDecoder().decode(MCPError.self, from: data) + + #expect(error.code == ErrorCode.urlElicitationRequired) + #expect(error.elicitations?.count == 1) + #expect(error.elicitations?[0].elicitationId == "auth-456") + #expect(error.elicitations?[0].url == "https://example.com/authorize") + } + + @Test("MCPError.urlElicitationRequired roundtrip encoding") + func testRoundtrip() throws { + let original = MCPError.urlElicitationRequired( + elicitations: [ + ElicitRequestURLParams( + message: "Authorize OAuth", + elicitationId: "oauth-789", + url: "https://provider.com/oauth/authorize" + ) + ], + message: "OAuth authorization needed" + ) + + let data = try JSONEncoder().encode(original) + let decoded = try JSONDecoder().decode(MCPError.self, from: data) + + #expect(decoded.code == original.code) + #expect(decoded.errorDescription == original.errorDescription) + #expect(decoded.elicitations?.count == original.elicitations?.count) + #expect(decoded.elicitations?[0].elicitationId == original.elicitations?[0].elicitationId) + } + + @Test("MCPError.fromError reconstructs urlElicitationRequired") + func testFromError() { + let data: Value = .object([ + "elicitations": .array([ + .object([ + "mode": .string("url"), + "message": .string("Authorize access"), + "elicitationId": .string("from-error-123"), + "url": .string("https://example.com/auth"), + ]) + ]) + ]) + + let error = MCPError.fromError( + code: ErrorCode.urlElicitationRequired, + message: "Elicitation required", + data: data + ) + + #expect(error.code == ErrorCode.urlElicitationRequired) + #expect(error.elicitations?.count == 1) + #expect(error.elicitations?[0].elicitationId == "from-error-123") + } + + @Test("MCPError.fromError falls back to serverError for invalid data") + func testFromErrorFallback() { + let error = MCPError.fromError( + code: ErrorCode.urlElicitationRequired, + message: "Elicitation required", + data: nil + ) + + // Should fall back to serverError when data is missing + if case .serverError(let code, let message) = error { + #expect(code == ErrorCode.urlElicitationRequired) + #expect(message == "Elicitation required") + } else { + Issue.record("Expected serverError fallback") + } + } + + @Test("Tool handler can throw URLElicitationRequiredError") + func testThrowFromToolHandler() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ElicitTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "requiresAuth", description: "Requires auth", inputSchema: [:]) + ]) + } + + // Tool handler throws URLElicitationRequiredError + await server.withRequestHandler(CallTool.self) { _, _ in + throw MCPError.urlElicitationRequired( + elicitations: [ + ElicitRequestURLParams( + message: "Please authorize access to your files", + elicitationId: "file-access-auth", + url: "https://files.example.com/oauth" + ) + ], + message: "Authorization required to access files" + ) + } + + let client = Client(name: "ElicitTestClient", version: "1.0.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + do { + _ = try await client.callTool(name: "requiresAuth", arguments: [:]) + Issue.record("Expected error to be thrown") + } catch let error as MCPError { + #expect(error.code == ErrorCode.urlElicitationRequired) + #expect(error.elicitations?.count == 1) + #expect(error.elicitations?[0].elicitationId == "file-access-auth") + #expect(error.elicitations?[0].url == "https://files.example.com/oauth") + } + + await client.disconnect() + } + + @Test("Client receives URLElicitationRequiredError with multiple elicitations") + func testMultipleElicitations() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ElicitTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "multiAuth", description: "Multiple auth", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { _, _ in + throw MCPError.urlElicitationRequired( + elicitations: [ + ElicitRequestURLParams( + message: "Authorize Google Drive", + elicitationId: "google-drive", + url: "https://accounts.google.com/oauth" + ), + ElicitRequestURLParams( + message: "Authorize Dropbox", + elicitationId: "dropbox", + url: "https://www.dropbox.com/oauth" + ), + ElicitRequestURLParams( + message: "Authorize OneDrive", + elicitationId: "onedrive", + url: "https://login.microsoftonline.com/oauth" + ), + ], + message: "Multiple cloud storage authorizations required" + ) + } + + let client = Client(name: "ElicitTestClient", version: "1.0.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + do { + _ = try await client.callTool(name: "multiAuth", arguments: [:]) + Issue.record("Expected error to be thrown") + } catch let error as MCPError { + #expect(error.code == ErrorCode.urlElicitationRequired) + #expect(error.elicitations?.count == 3) + #expect(error.elicitations?[0].elicitationId == "google-drive") + #expect(error.elicitations?[1].elicitationId == "dropbox") + #expect(error.elicitations?[2].elicitationId == "onedrive") + } + + await client.disconnect() + } +} + +// MARK: - ElicitationComplete Notification Integration Tests + +/// Actor to track notification receipt in tests +private actor NotificationState { + var received = false + var elicitationId: String? + var count = 0 + + func markReceived(elicitationId: String? = nil) { + received = true + self.elicitationId = elicitationId + count += 1 + } +} + +@Suite("ElicitationComplete Notification Integration Tests") +struct ElicitationCompleteNotificationIntegrationTests { + + @Test("Server can send elicitation complete notification") + func testServerSendsElicitationCompleteNotification() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ElicitTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + let notificationState = NotificationState() + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "completeAuth", description: "Complete auth", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { _, context in + // Simulate async operation completion (e.g., user finished OAuth in browser) + let elicitationId = "complete-test-123" + + // Send the completion notification + try await context.sendMessage(ElicitationCompleteNotification.message(.init( + elicitationId: elicitationId + ))) + + return CallTool.Result(content: [.text("Elicitation completed")]) + } + + let client = Client(name: "ElicitTestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(url: Client.Capabilities.Elicitation.URL()) + )) + + // Set up notification handler + await client.onNotification(ElicitationCompleteNotification.self) { [notificationState] message in + await notificationState.markReceived(elicitationId: message.params.elicitationId) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "completeAuth", arguments: [:]) + + // Verify tool result + if case .text(let text, _, _) = result.content[0] { + #expect(text == "Elicitation completed") + } + + // Give time for notification to be processed + try await Task.sleep(for: .milliseconds(50)) + + // Verify notification was received + let received = await notificationState.received + let receivedId = await notificationState.elicitationId + #expect(received == true) + #expect(receivedId == "complete-test-123") + + await client.disconnect() + } + + @Test("Server sends elicitation complete after URL mode elicitation") + func testElicitationCompleteAfterURLElicitation() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ElicitTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + let notificationState = NotificationState() + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "authorize", description: "Authorize", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] _, context in + // First, send a URL elicitation + let elicitationId = "oauth-flow-456" + let result = try await server.elicit(ElicitRequestParams.url(ElicitRequestURLParams( + message: "Complete OAuth", + elicitationId: elicitationId, + url: "https://example.com/oauth" + ))) + + // After client responds, send completion notification + if result.action == .accept { + try await context.sendMessage(ElicitationCompleteNotification.message(.init( + elicitationId: elicitationId + ))) + } + + return CallTool.Result(content: [.text("Authorization complete")]) + } + + let client = Client(name: "ElicitTestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(url: Client.Capabilities.Elicitation.URL()) + )) + + await client.withElicitationHandler { params, _ in + guard case .url(let urlParams) = params else { + return ElicitResult(action: .decline) + } + #expect(urlParams.elicitationId == "oauth-flow-456") + return ElicitResult(action: .accept) + } + + await client.onNotification(ElicitationCompleteNotification.self) { [notificationState] _ in + await notificationState.markReceived() + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "authorize", arguments: [:]) + + if case .text(let text, _, _) = result.content[0] { + #expect(text == "Authorization complete") + } + + // Give time for notification to be processed + try await Task.sleep(for: .milliseconds(50)) + + let count = await notificationState.count + #expect(count == 1) + + await client.disconnect() + } +} + +// MARK: - Task-Augmented Elicitation Tests + +/// Actor to track task-augmented elicitation state +private actor TaskAugmentedState { + var received = false + var taskId: String? + + func markReceived(taskId: String) { + received = true + self.taskId = taskId + } +} + +@Suite("Task-Augmented Elicitation Tests") +struct TaskAugmentedElicitationTests { + + /// Helper to create a simple MCPTask for testing + private static func makeTask(taskId: String, status: TaskStatus) -> MCPTask { + let now = ISO8601DateFormatter().string(from: Date()) + return MCPTask( + taskId: taskId, + status: status, + createdAt: now, + lastUpdatedAt: now + ) + } + + @Test("Client can register task-augmented elicitation handler") + func testTaskAugmentedElicitationHandler() async throws { + let elicitationState = TaskAugmentedState() + + let client = Client(name: "ElicitTestClient", version: "1.0.0") + + // Set up task-augmented elicitation handler + var taskHandlers = ExperimentalClientTaskHandlers() + taskHandlers.taskAugmentedElicitation = { [elicitationState] _, _ in + let taskId = UUID().uuidString + await elicitationState.markReceived(taskId: taskId) + + // Create a task to handle the elicitation + return CreateTaskResult(task: Self.makeTask(taskId: taskId, status: .completed)) + } + + await client.enableTaskHandlers(ClientTaskSupport.inMemory(handlers: taskHandlers)) + + // Verify the capability was built correctly + let caps = await client.capabilities + #expect(caps.tasks?.requests?.elicitation?.create != nil) + } + + @Test("TaskAugmentedElicitationHandler type alias exists") + func testTaskAugmentedElicitationHandlerType() { + // Verify the type alias compiles correctly + let handler: ExperimentalClientTaskHandlers.TaskAugmentedElicitationHandler = { _, _ in + let taskId = UUID().uuidString + return CreateTaskResult(task: Self.makeTask(taskId: taskId, status: .working)) + } + + // Just verify it compiles - the type system enforces correctness + _ = handler + } + + @Test("ExperimentalClientTaskHandlers builds correct capability for elicitation") + func testBuildCapabilityWithElicitation() { + var handlers = ExperimentalClientTaskHandlers() + + // With elicitation handler, capability should include requests.elicitation + handlers.taskAugmentedElicitation = { _, _ in + CreateTaskResult(task: Self.makeTask(taskId: UUID().uuidString, status: .completed)) + } + + let capability = handlers.buildCapability() + #expect(capability != nil) + #expect(capability?.requests?.elicitation?.create != nil) + } + + @Test("ExperimentalClientTaskHandlers builds capability with both sampling and elicitation") + func testBuildCapabilityWithBoth() { + var handlers = ExperimentalClientTaskHandlers() + + handlers.taskAugmentedSampling = { _, _ in + CreateTaskResult(task: Self.makeTask(taskId: UUID().uuidString, status: .completed)) + } + + handlers.taskAugmentedElicitation = { _, _ in + CreateTaskResult(task: Self.makeTask(taskId: UUID().uuidString, status: .completed)) + } + + let capability = handlers.buildCapability() + #expect(capability != nil) + #expect(capability?.requests?.sampling?.createMessage != nil) + #expect(capability?.requests?.elicitation?.create != nil) + } + + @Test("hasTaskAugmentedElicitation returns correct value") + func testHasTaskAugmentedElicitation() { + // Without capability + #expect(hasTaskAugmentedElicitation(nil) == false) + + // With empty capabilities + let emptyCaps = Client.Capabilities() + #expect(hasTaskAugmentedElicitation(emptyCaps) == false) + + // With tasks but no requests + let withTasks = Client.Capabilities(tasks: .init()) + #expect(hasTaskAugmentedElicitation(withTasks) == false) + + // With requests but no elicitation + let withRequests = Client.Capabilities(tasks: .init(requests: .init())) + #expect(hasTaskAugmentedElicitation(withRequests) == false) + + // With elicitation.create + let withElicitation = Client.Capabilities( + tasks: .init(requests: .init(elicitation: .init(create: .init()))) + ) + #expect(hasTaskAugmentedElicitation(withElicitation) == true) + } + + @Test("requireTaskAugmentedElicitation throws when not supported") + func testRequireTaskAugmentedElicitation() { + // Should throw when capability is nil + #expect(throws: MCPError.self) { + try requireTaskAugmentedElicitation(nil) + } + + // Should throw when tasks.requests.elicitation is nil + let withoutElicitation = Client.Capabilities(tasks: .init()) + #expect(throws: MCPError.self) { + try requireTaskAugmentedElicitation(withoutElicitation) + } + + // Should not throw when supported + let withElicitation = Client.Capabilities( + tasks: .init(requests: .init(elicitation: .init(create: .init()))) + ) + #expect(throws: Never.self) { + try requireTaskAugmentedElicitation(withElicitation) + } + } +} diff --git a/Tests/MCPTests/ErrorTests.swift b/Tests/MCPTests/ErrorTests.swift new file mode 100644 index 00000000..4cf86764 --- /dev/null +++ b/Tests/MCPTests/ErrorTests.swift @@ -0,0 +1,417 @@ +import Foundation +import Testing + +@testable import MCP + +@Suite("MCPError Roundtrip Tests") +struct MCPErrorRoundtripTests { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + // MARK: - Standard JSON-RPC Errors + + @Test("parseError roundtrip with nil detail") + func testParseErrorNilRoundtrip() throws { + let original = MCPError.parseError(nil) + let decoded = try roundtrip(original) + #expect(decoded == original) + #expect(decoded.code == ErrorCode.parseError) + #expect(decoded.message == "Invalid JSON") + } + + @Test("parseError roundtrip with custom detail") + func testParseErrorDetailRoundtrip() throws { + let original = MCPError.parseError("Unexpected token at position 5") + let decoded = try roundtrip(original) + #expect(decoded == original) + #expect(decoded.message == "Unexpected token at position 5") + } + + @Test("invalidRequest roundtrip with nil detail") + func testInvalidRequestNilRoundtrip() throws { + let original = MCPError.invalidRequest(nil) + let decoded = try roundtrip(original) + #expect(decoded == original) + #expect(decoded.code == ErrorCode.invalidRequest) + } + + @Test("invalidRequest roundtrip with custom detail") + func testInvalidRequestDetailRoundtrip() throws { + let original = MCPError.invalidRequest("Missing id field") + let decoded = try roundtrip(original) + #expect(decoded == original) + #expect(decoded.message == "Missing id field") + } + + @Test("methodNotFound roundtrip with nil detail") + func testMethodNotFoundNilRoundtrip() throws { + let original = MCPError.methodNotFound(nil) + let decoded = try roundtrip(original) + #expect(decoded == original) + #expect(decoded.code == ErrorCode.methodNotFound) + } + + @Test("methodNotFound roundtrip with custom detail") + func testMethodNotFoundDetailRoundtrip() throws { + let original = MCPError.methodNotFound("tools/call") + let decoded = try roundtrip(original) + #expect(decoded == original) + #expect(decoded.message == "tools/call") + } + + @Test("invalidParams roundtrip with nil detail") + func testInvalidParamsNilRoundtrip() throws { + let original = MCPError.invalidParams(nil) + let decoded = try roundtrip(original) + #expect(decoded == original) + #expect(decoded.code == ErrorCode.invalidParams) + } + + @Test("invalidParams roundtrip with custom detail") + func testInvalidParamsDetailRoundtrip() throws { + let original = MCPError.invalidParams("name is required") + let decoded = try roundtrip(original) + #expect(decoded == original) + #expect(decoded.message == "name is required") + } + + @Test("internalError roundtrip with nil detail") + func testInternalErrorNilRoundtrip() throws { + let original = MCPError.internalError(nil) + let decoded = try roundtrip(original) + #expect(decoded == original) + #expect(decoded.code == ErrorCode.internalError) + } + + @Test("internalError roundtrip with custom detail") + func testInternalErrorDetailRoundtrip() throws { + let original = MCPError.internalError("Database connection failed") + let decoded = try roundtrip(original) + #expect(decoded == original) + #expect(decoded.message == "Database connection failed") + } + + // MARK: - MCP-Specific Errors + + @Test("resourceNotFound roundtrip with nil URI") + func testResourceNotFoundNilRoundtrip() throws { + let original = MCPError.resourceNotFound(uri: nil) + let decoded = try roundtrip(original) + #expect(decoded == original) + #expect(decoded.code == ErrorCode.resourceNotFound) + #expect(decoded.data == nil) + } + + @Test("resourceNotFound roundtrip with URI") + func testResourceNotFoundUriRoundtrip() throws { + let original = MCPError.resourceNotFound(uri: "file:///path/to/file.txt") + let decoded = try roundtrip(original) + #expect(decoded == original) + #expect(decoded.message == "Resource not found: file:///path/to/file.txt") + #expect(decoded.data == .object(["uri": .string("file:///path/to/file.txt")])) + } + + @Test("urlElicitationRequired roundtrip") + func testUrlElicitationRequiredRoundtrip() throws { + let elicitations = [ + ElicitRequestURLParams( + message: "Please authorize", + elicitationId: "auth-123", + url: "https://example.com/oauth" + ) + ] + let original = MCPError.urlElicitationRequired( + message: "Authorization required", + elicitations: elicitations + ) + let decoded = try roundtrip(original) + #expect(decoded == original) + #expect(decoded.code == ErrorCode.urlElicitationRequired) + #expect(decoded.elicitations?.count == 1) + #expect(decoded.elicitations?[0].elicitationId == "auth-123") + } + + // MARK: - Server Errors + + @Test("serverError roundtrip") + func testServerErrorRoundtrip() throws { + let original = MCPError.serverError(code: -32050, message: "Custom server error") + let decoded = try roundtrip(original) + #expect(decoded == original) + #expect(decoded.code == -32050) + #expect(decoded.message == "Custom server error") + } + + @Test("serverErrorWithData roundtrip") + func testServerErrorWithDataRoundtrip() throws { + let data: Value = .object(["detail": .string("Extra info"), "count": .int(42)]) + let original = MCPError.serverErrorWithData(code: -32051, message: "Error with data", data: data) + let decoded = try roundtrip(original) + #expect(decoded == original) + #expect(decoded.code == -32051) + #expect(decoded.data == data) + } + + // MARK: - SDK-Specific Errors + + @Test("connectionClosed roundtrip") + func testConnectionClosedRoundtrip() throws { + let original = MCPError.connectionClosed + let decoded = try roundtrip(original) + #expect(decoded == original) + #expect(decoded.code == ErrorCode.connectionClosed) + } + + @Test("requestTimeout roundtrip") + func testRequestTimeoutRoundtrip() throws { + let original = MCPError.requestTimeout(timeout: .seconds(30), message: nil) + let decoded = try roundtrip(original) + #expect(decoded.code == ErrorCode.requestTimeout) + // Note: Duration precision may not be exact due to ms conversion + } + + @Test("requestTimeout roundtrip with message") + func testRequestTimeoutMessageRoundtrip() throws { + let original = MCPError.requestTimeout(timeout: .seconds(60), message: "Server unresponsive") + let decoded = try roundtrip(original) + #expect(decoded.code == ErrorCode.requestTimeout) + #expect(decoded.message == "Server unresponsive") + } + + @Test("requestCancelled roundtrip") + func testRequestCancelledRoundtrip() throws { + let original = MCPError.requestCancelled(reason: nil) + let decoded = try roundtrip(original) + #expect(decoded.code == ErrorCode.requestCancelled) + #expect(decoded.message == "Request cancelled") + } + + @Test("requestCancelled roundtrip with reason") + func testRequestCancelledWithReasonRoundtrip() throws { + let original = MCPError.requestCancelled(reason: "User cancelled the operation") + let decoded = try roundtrip(original) + #expect(decoded.code == ErrorCode.requestCancelled) + #expect(decoded.message == "User cancelled the operation") + if case .requestCancelled(let reason) = decoded { + #expect(reason == "User cancelled the operation") + } else { + Issue.record("Expected requestCancelled") + } + } + + // MARK: - Helpers + + private func roundtrip(_ error: MCPError) throws -> MCPError { + let data = try encoder.encode(error) + return try decoder.decode(MCPError.self, from: data) + } +} + +@Suite("MCPError Message and Data Properties Tests") +struct MCPErrorPropertyTests { + @Test("message property returns raw message for wire format") + func testMessageProperty() { + #expect(MCPError.parseError(nil).message == "Invalid JSON") + #expect(MCPError.parseError("custom").message == "custom") + #expect(MCPError.invalidRequest(nil).message == "Invalid Request") + #expect(MCPError.methodNotFound(nil).message == "Method not found") + #expect(MCPError.invalidParams(nil).message == "Invalid params") + #expect(MCPError.internalError(nil).message == "Internal error") + #expect(MCPError.resourceNotFound(uri: nil).message == "Resource not found") + #expect(MCPError.resourceNotFound(uri: "test://uri").message == "Resource not found: test://uri") + #expect(MCPError.connectionClosed.message == "Connection closed") + #expect(MCPError.serverError(code: -32050, message: "Custom").message == "Custom") + } + + @Test("data property returns correct payload for wire format") + func testDataProperty() { + // Standard errors have no data + #expect(MCPError.parseError(nil).data == nil) + #expect(MCPError.invalidRequest("detail").data == nil) + + // Resource not found with URI includes URI in data + #expect(MCPError.resourceNotFound(uri: nil).data == nil) + let resourceData = MCPError.resourceNotFound(uri: "test://uri").data + #expect(resourceData == .object(["uri": .string("test://uri")])) + + // Server error with data includes the data + let customData: Value = .object(["key": .string("value")]) + #expect(MCPError.serverErrorWithData(code: -32050, message: "msg", data: customData).data == customData) + #expect(MCPError.serverError(code: -32050, message: "msg").data == nil) + } + + @Test("code property returns correct error codes") + func testCodeProperty() { + #expect(MCPError.parseError(nil).code == ErrorCode.parseError) + #expect(MCPError.invalidRequest(nil).code == ErrorCode.invalidRequest) + #expect(MCPError.methodNotFound(nil).code == ErrorCode.methodNotFound) + #expect(MCPError.invalidParams(nil).code == ErrorCode.invalidParams) + #expect(MCPError.internalError(nil).code == ErrorCode.internalError) + #expect(MCPError.resourceNotFound(uri: nil).code == ErrorCode.resourceNotFound) + #expect(MCPError.urlElicitationRequired(message: "", elicitations: []).code == ErrorCode.urlElicitationRequired) + #expect(MCPError.connectionClosed.code == ErrorCode.connectionClosed) + #expect(MCPError.requestTimeout(timeout: .seconds(1), message: nil).code == ErrorCode.requestTimeout) + #expect(MCPError.transportError(NSError(domain: "", code: 0)).code == ErrorCode.transportError) + #expect(MCPError.requestCancelled(reason: nil).code == ErrorCode.requestCancelled) + // Custom server error codes (no constants defined - these are arbitrary test values) + #expect(MCPError.serverError(code: -32050, message: "").code == -32050) + #expect(MCPError.serverErrorWithData(code: -32051, message: "", data: .null).code == -32051) + } +} + +@Suite("MCPError fromError Factory Tests") +struct MCPErrorFromErrorTests { + @Test("fromError reconstructs standard errors") + func testFromErrorStandardErrors() { + let parseError = MCPError.fromError(code: ErrorCode.parseError, message: "Invalid JSON") + #expect(parseError == .parseError(nil)) + + let parseErrorCustom = MCPError.fromError(code: ErrorCode.parseError, message: "Custom parse error") + #expect(parseErrorCustom == .parseError("Custom parse error")) + + let invalidRequest = MCPError.fromError(code: ErrorCode.invalidRequest, message: "Invalid Request") + #expect(invalidRequest == .invalidRequest(nil)) + + let methodNotFound = MCPError.fromError(code: ErrorCode.methodNotFound, message: "Method not found") + #expect(methodNotFound == .methodNotFound(nil)) + } + + @Test("fromError reconstructs resourceNotFound with URI from data") + func testFromErrorResourceNotFound() { + let withoutData = MCPError.fromError(code: ErrorCode.resourceNotFound, message: "Resource not found") + if case .resourceNotFound(let uri) = withoutData { + #expect(uri == nil) + } else { + Issue.record("Expected resourceNotFound") + } + + let data: Value = .object(["uri": .string("file:///test.txt")]) + let withData = MCPError.fromError(code: ErrorCode.resourceNotFound, message: "Resource not found", data: data) + if case .resourceNotFound(let uri) = withData { + #expect(uri == "file:///test.txt") + } else { + Issue.record("Expected resourceNotFound with uri") + } + } + + @Test("fromError reconstructs urlElicitationRequired from data") + func testFromErrorUrlElicitation() { + let data: Value = .object([ + "elicitations": .array([ + .object([ + "mode": .string("url"), + "elicitationId": .string("test-123"), + "url": .string("https://example.com"), + "message": .string("Authorize") + ]) + ]) + ]) + let error = MCPError.fromError(code: ErrorCode.urlElicitationRequired, message: "Elicitation required", data: data) + #expect(error.code == ErrorCode.urlElicitationRequired) + #expect(error.elicitations?.count == 1) + #expect(error.elicitations?[0].elicitationId == "test-123") + } + + @Test("fromError falls back to serverError for unknown codes") + func testFromErrorUnknownCode() { + let error = MCPError.fromError(code: -32099, message: "Unknown error") + if case .serverError(let code, let message) = error { + #expect(code == -32099) + #expect(message == "Unknown error") + } else { + Issue.record("Expected serverError") + } + } + + @Test("fromError preserves data for unknown codes") + func testFromErrorUnknownCodeWithData() { + let data: Value = .object(["extra": .string("info")]) + let error = MCPError.fromError(code: -32099, message: "Error with data", data: data) + if case .serverErrorWithData(let code, let message, let errorData) = error { + #expect(code == -32099) + #expect(message == "Error with data") + #expect(errorData == data) + } else { + Issue.record("Expected serverErrorWithData") + } + } + + @Test("fromError reconstructs requestCancelled") + func testFromErrorRequestCancelled() { + // Without reason + let withoutReason = MCPError.fromError(code: ErrorCode.requestCancelled, message: "Request cancelled") + if case .requestCancelled(let reason) = withoutReason { + #expect(reason == nil) + } else { + Issue.record("Expected requestCancelled") + } + + // With reason in message + let withMessage = MCPError.fromError(code: ErrorCode.requestCancelled, message: "User cancelled") + if case .requestCancelled(let reason) = withMessage { + #expect(reason == "User cancelled") + } else { + Issue.record("Expected requestCancelled") + } + + // With reason in data + let data: Value = .object(["reason": .string("Operation aborted by user")]) + let withData = MCPError.fromError(code: ErrorCode.requestCancelled, message: "Request cancelled", data: data) + if case .requestCancelled(let reason) = withData { + #expect(reason == "Operation aborted by user") + } else { + Issue.record("Expected requestCancelled") + } + } +} + +@Suite("MCPError Wire Format Tests") +struct MCPErrorWireFormatTests { + let encoder = JSONEncoder() + + @Test("Standard errors encode without data field") + func testStandardErrorsNoData() throws { + let error = MCPError.parseError(nil) + let data = try encoder.encode(error) + let json = try JSONSerialization.jsonObject(with: data) as! [String: Any] + + #expect(json["code"] as? Int == ErrorCode.parseError) + #expect(json["message"] as? String == "Invalid JSON") + #expect(json["data"] == nil) + } + + @Test("resourceNotFound encodes URI in data field") + func testResourceNotFoundWireFormat() throws { + let error = MCPError.resourceNotFound(uri: "file:///test.txt") + let data = try encoder.encode(error) + let json = try JSONSerialization.jsonObject(with: data) as! [String: Any] + + #expect(json["code"] as? Int == ErrorCode.resourceNotFound) + #expect(json["message"] as? String == "Resource not found: file:///test.txt") + let dataField = json["data"] as? [String: Any] + #expect(dataField?["uri"] as? String == "file:///test.txt") + } + + @Test("urlElicitationRequired encodes elicitations in data field") + func testUrlElicitationWireFormat() throws { + let error = MCPError.urlElicitationRequired( + elicitations: [ + ElicitRequestURLParams( + message: "Auth", + elicitationId: "e1", + url: "https://auth.example.com" + ) + ], + message: "Authorization needed" + ) + let data = try encoder.encode(error) + let json = try JSONSerialization.jsonObject(with: data) as! [String: Any] + + #expect(json["code"] as? Int == ErrorCode.urlElicitationRequired) + #expect(json["message"] as? String == "Authorization needed") + let dataField = json["data"] as? [String: Any] + let elicitations = dataField?["elicitations"] as? [[String: Any]] + #expect(elicitations?.count == 1) + #expect(elicitations?[0]["elicitationId"] as? String == "e1") + } +} diff --git a/Tests/MCPTests/FullRoundtripTests.swift b/Tests/MCPTests/FullRoundtripTests.swift new file mode 100644 index 00000000..d86d596d --- /dev/null +++ b/Tests/MCPTests/FullRoundtripTests.swift @@ -0,0 +1,342 @@ +import Foundation +import Testing + +@testable import MCP + +/// Tests for full client-server roundtrip flows through the HTTP transport layer +/// with a real MCP Server instance. +/// +/// These tests follow the TypeScript SDK patterns from: +/// - `test/integration/test/stateManagementStreamableHttp.test.ts` +@Suite("Full Roundtrip Tests") +struct FullRoundtripTests { + + // MARK: - Test Helpers + + /// Creates a configured MCP Server with tools for testing + func createTestServer() -> Server { + let server = Server( + name: "test-server", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + return server + } + + /// Sets up tool handlers on the server + func setUpToolHandlers(_ server: Server) async { + // Register tool list handler + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool( + name: "greet", + description: "A simple greeting tool", + inputSchema: [ + "type": "object", + "properties": [ + "name": ["type": "string", "description": "Name to greet"] + ] + ] + ), + Tool( + name: "add", + description: "Adds two numbers", + inputSchema: [ + "type": "object", + "properties": [ + "a": ["type": "number", "description": "First number"], + "b": ["type": "number", "description": "Second number"] + ], + "required": ["a", "b"] + ] + ), + ]) + } + + // Register tool call handler + await server.withRequestHandler(CallTool.self) { request, _ in + switch request.name { + case "greet": + let name = request.arguments?["name"]?.stringValue ?? "World" + return CallTool.Result(content: [.text("Hello, \(name)!")]) + + case "add": + let a = request.arguments?["a"]?.doubleValue ?? 0 + let b = request.arguments?["b"]?.doubleValue ?? 0 + return CallTool.Result(content: [.text("Result: \(a + b)")]) + + default: + return CallTool.Result(content: [.text("Unknown tool: \(request.name)")], isError: true) + } + } + } + + + // MARK: - 2.1 Multiple client connections (stateless mode) + + @Test("Multiple client connections in stateless mode") + func multipleClientConnectionsStateless() async throws { + let server = createTestServer() + await setUpToolHandlers(server) + + // Create transport in stateless mode (no sessionIdGenerator) + let transport = HTTPServerTransport() + try await server.start(transport: transport) + + // Client 1 initializes + let init1 = TestPayloads.initializeRequest(id: "c1-init", clientName: "client1") + let response1 = await transport.handleRequest(TestPayloads.postRequest(body: init1)) + #expect(response1.statusCode == 200) + #expect(response1.headers[HTTPHeader.sessionId] == nil, "Stateless mode should not return session ID") + + // Client 1 lists tools + let listTools1 = TestPayloads.listToolsRequest(id: "c1-list") + let toolsResponse1 = await transport.handleRequest(TestPayloads.postRequest(body: listTools1)) + #expect(toolsResponse1.statusCode == 200) + + if let body = toolsResponse1.body, let text = String(data: body, encoding: .utf8) { + #expect(text.contains("greet"), "Should list greet tool") + } + + // Client 2 initializes (separate connection) + let init2 = TestPayloads.initializeRequest(id: "c2-init", clientName: "client2") + let response2 = await transport.handleRequest(TestPayloads.postRequest(body: init2)) + #expect(response2.statusCode == 200) + + // Client 2 calls a tool + let callTool2 = TestPayloads.callToolRequest(id: "c2-call", name: "greet", arguments: ["name": "Client2"]) + let toolResponse2 = await transport.handleRequest(TestPayloads.postRequest(body: callTool2)) + #expect(toolResponse2.statusCode == 200) + + if let body = toolResponse2.body, let text = String(data: body, encoding: .utf8) { + #expect(text.contains("Hello, Client2!"), "Should return greeting for Client2") + } + } + + // MARK: - 2.2 Operate with session management (stateful mode) + + @Test("Operate with session management in stateful mode") + func operateWithSessionManagement() async throws { + let server = createTestServer() + await setUpToolHandlers(server) + + let sessionId = UUID().uuidString + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { sessionId }) + ) + try await server.start(transport: transport) + + // Initialize and get session ID + let initRequest = TestPayloads.initializeRequest(id: "init", clientName: "test-client") + let initResponse = await transport.handleRequest(TestPayloads.postRequest(body: initRequest)) + #expect(initResponse.statusCode == 200) + #expect(initResponse.headers[HTTPHeader.sessionId] == sessionId, "Should return session ID") + + // Make subsequent request with session ID - should succeed + let listTools = TestPayloads.listToolsRequest(id: "list") + let toolsResponse = await transport.handleRequest(TestPayloads.postRequest(body: listTools, sessionId: sessionId)) + #expect(toolsResponse.statusCode == 200) + + if let body = toolsResponse.body, let text = String(data: body, encoding: .utf8) { + #expect(text.contains("greet"), "Should list greet tool") + #expect(text.contains("add"), "Should list add tool") + } + + // Make request without session ID - should fail + let noSessionRequest = await transport.handleRequest(TestPayloads.postRequest(body: listTools)) + #expect(noSessionRequest.statusCode == 400, "Should reject request without session ID in stateful mode") + } + + // MARK: - 2.3 Full tool call roundtrip + + @Test("Full tool call roundtrip with real server") + func fullToolCallRoundtrip() async throws { + let server = createTestServer() + await setUpToolHandlers(server) + + let sessionId = UUID().uuidString + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { sessionId }) + ) + try await server.start(transport: transport) + + // Step 1: Initialize + let initRequest = TestPayloads.initializeRequest() + let initResponse = await transport.handleRequest(TestPayloads.postRequest(body: initRequest)) + #expect(initResponse.statusCode == 200) + + if let body = initResponse.body, let text = String(data: body, encoding: .utf8) { + #expect(text.contains("protocolVersion"), "Init response should include protocolVersion") + #expect(text.contains("capabilities"), "Init response should include capabilities") + } + + // Step 2: List tools + let listToolsRequest = """ + {"jsonrpc":"2.0","method":"tools/list","id":"2"} + """ + let listResponse = await transport.handleRequest(TestPayloads.postRequest(body: listToolsRequest, sessionId: sessionId)) + #expect(listResponse.statusCode == 200) + + if let body = listResponse.body, let text = String(data: body, encoding: .utf8) { + #expect(text.contains("greet"), "Should include greet tool") + #expect(text.contains("add"), "Should include add tool") + } + + // Step 3: Call greet tool + let greetRequest = """ + {"jsonrpc":"2.0","method":"tools/call","id":"3","params":{"name":"greet","arguments":{"name":"MCP Swift"}}} + """ + let greetResponse = await transport.handleRequest(TestPayloads.postRequest(body: greetRequest, sessionId: sessionId)) + #expect(greetResponse.statusCode == 200) + + if let body = greetResponse.body, let text = String(data: body, encoding: .utf8) { + #expect(text.contains("Hello, MCP Swift!"), "Should return correct greeting") + } + + // Step 4: Call add tool + let addRequest = """ + {"jsonrpc":"2.0","method":"tools/call","id":"4","params":{"name":"add","arguments":{"a":5,"b":3}}} + """ + let addResponse = await transport.handleRequest(TestPayloads.postRequest(body: addRequest, sessionId: sessionId)) + #expect(addResponse.statusCode == 200) + + if let body = addResponse.body, let text = String(data: body, encoding: .utf8) { + #expect(text.contains("Result: 8"), "Should return correct sum") + } + } + + // MARK: - 2.4 Protocol version negotiation + + @Test("Protocol version negotiation stores correct version") + func protocolVersionNegotiation() async throws { + let server = createTestServer() + + let sessionId = UUID().uuidString + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { sessionId }) + ) + try await server.start(transport: transport) + + // Initialize with specific protocol version + let initRequest = TestPayloads.initializeRequest() + let initResponse = await transport.handleRequest(TestPayloads.postRequest(body: initRequest)) + #expect(initResponse.statusCode == 200) + + // Verify response includes the negotiated version + if let body = initResponse.body, let text = String(data: body, encoding: .utf8) { + #expect(text.contains("protocolVersion"), "Response should include protocol version") + #expect(text.contains(Version.v2024_11_05) || text.contains("2025"), "Response should include a valid version") + } + + // Verify subsequent requests work with the protocol version header + let pingRequest = """ + {"jsonrpc":"2.0","method":"ping","id":"2"} + """ + let headers = [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + HTTPHeader.sessionId: sessionId, + HTTPHeader.protocolVersion: Version.v2024_11_05, + ] + let pingResponse = await transport.handleRequest(HTTPRequest( + method: "POST", + headers: headers, + body: pingRequest.data(using: .utf8) + )) + #expect(pingResponse.statusCode == 200) + } + + @Test("Reject mismatched protocol version") + func rejectMismatchedProtocolVersion() async throws { + let server = createTestServer() + + let sessionId = UUID().uuidString + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { sessionId }) + ) + try await server.start(transport: transport) + + // Initialize first + let initRequest = TestPayloads.initializeRequest() + _ = await transport.handleRequest(TestPayloads.postRequest(body: initRequest)) + + // Send request with different protocol version + let pingRequest = """ + {"jsonrpc":"2.0","method":"ping","id":"2"} + """ + let headers = [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + HTTPHeader.sessionId: sessionId, + HTTPHeader.protocolVersion: "9999-99-99", // Invalid version + ] + let response = await transport.handleRequest(HTTPRequest( + method: "POST", + headers: headers, + body: pingRequest.data(using: .utf8) + )) + + // Should reject the mismatched version + #expect(response.statusCode == 400, "Should reject mismatched protocol version") + } + + // MARK: - Additional Integration Tests + + @Test("Unknown tool returns error") + func unknownToolReturnsError() async throws { + let server = createTestServer() + await setUpToolHandlers(server) + + let sessionId = UUID().uuidString + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { sessionId }) + ) + try await server.start(transport: transport) + + // Initialize + let initRequest = TestPayloads.initializeRequest() + _ = await transport.handleRequest(TestPayloads.postRequest(body: initRequest)) + + // Call unknown tool + let unknownToolRequest = """ + {"jsonrpc":"2.0","method":"tools/call","id":"2","params":{"name":"nonexistent","arguments":{}}} + """ + let response = await transport.handleRequest(TestPayloads.postRequest(body: unknownToolRequest, sessionId: sessionId)) + #expect(response.statusCode == 200) // JSON-RPC error is 200 with error in body + + if let body = response.body, let text = String(data: body, encoding: .utf8) { + #expect(text.contains("Unknown tool") || text.contains("isError"), "Should indicate error for unknown tool") + } + } + + @Test("Batch requests work correctly") + func batchRequestsWorkCorrectly() async throws { + let server = createTestServer() + await setUpToolHandlers(server) + + let sessionId = UUID().uuidString + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { sessionId }) + ) + try await server.start(transport: transport) + + // Initialize + let initRequest = TestPayloads.initializeRequest() + _ = await transport.handleRequest(TestPayloads.postRequest(body: initRequest)) + + // Send batch request with multiple tool calls + let batchRequest = """ + [ + {"jsonrpc":"2.0","method":"tools/call","id":"b1","params":{"name":"greet","arguments":{"name":"Alice"}}}, + {"jsonrpc":"2.0","method":"tools/call","id":"b2","params":{"name":"add","arguments":{"a":10,"b":20}}} + ] + """ + let response = await transport.handleRequest(TestPayloads.postRequest(body: batchRequest, sessionId: sessionId)) + #expect(response.statusCode == 200) + + if let body = response.body, let text = String(data: body, encoding: .utf8) { + #expect(text.contains("Hello, Alice!"), "Should include first tool result") + #expect(text.contains("Result: 30"), "Should include second tool result") + } + } +} diff --git a/Tests/MCPTests/HTTPClientTransportTests.swift b/Tests/MCPTests/HTTPClientTransportTests.swift index cf2a25d8..6ff4547b 100644 --- a/Tests/MCPTests/HTTPClientTransportTests.swift +++ b/Tests/MCPTests/HTTPClientTransportTests.swift @@ -178,7 +178,7 @@ import Testing let response = HTTPURLResponse( url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", - headerFields: ["Content-Type": "application/json"])! + headerFields: [HTTPHeader.contentType: "application/json"])! return (response, responseData) } @@ -209,12 +209,12 @@ import Testing await MockURLProtocol.requestHandlerStorage.setHandler { [testEndpoint] (request: URLRequest) in - #expect(request.value(forHTTPHeaderField: "Mcp-Session-Id") == nil) + #expect(request.value(forHTTPHeaderField: HTTPHeader.sessionId) == nil) let response = HTTPURLResponse( url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", headerFields: [ - "Content-Type": "application/json", - "Mcp-Session-Id": newSessionID, + HTTPHeader.contentType: "application/json", + HTTPHeader.sessionId: newSessionID, ])! return (response, Data()) } @@ -247,12 +247,12 @@ import Testing await MockURLProtocol.requestHandlerStorage.setHandler { [testEndpoint] (request: URLRequest) in #expect(request.readBody() == firstMessageData) - #expect(request.value(forHTTPHeaderField: "Mcp-Session-Id") == nil) + #expect(request.value(forHTTPHeaderField: HTTPHeader.sessionId) == nil) let response = HTTPURLResponse( url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", headerFields: [ - "Content-Type": "application/json", - "Mcp-Session-Id": initialSessionID, + HTTPHeader.contentType: "application/json", + HTTPHeader.sessionId: initialSessionID, ])! return (response, Data()) } @@ -262,11 +262,11 @@ import Testing await MockURLProtocol.requestHandlerStorage.setHandler { [testEndpoint] (request: URLRequest) in #expect(request.readBody() == secondMessageData) - #expect(request.value(forHTTPHeaderField: "Mcp-Session-Id") == initialSessionID) + #expect(request.value(forHTTPHeaderField: HTTPHeader.sessionId) == initialSessionID) let response = HTTPURLResponse( url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", - headerFields: ["Content-Type": "application/json"])! + headerFields: [HTTPHeader.contentType: "application/json"])! return (response, Data()) } try await transport.send(secondMessageData) @@ -367,8 +367,8 @@ import Testing let response = HTTPURLResponse( url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", headerFields: [ - "Content-Type": "application/json", - "Mcp-Session-Id": initialSessionID, + HTTPHeader.contentType: "application/json", + HTTPHeader.sessionId: initialSessionID, ])! return (response, Data()) } @@ -387,7 +387,7 @@ import Testing // Set up the second handler for the 404 response await MockURLProtocol.requestHandlerStorage.setHandler { [testEndpoint, initialSessionID] (request: URLRequest) in - #expect(request.value(forHTTPHeaderField: "Mcp-Session-Id") == initialSessionID) + #expect(request.value(forHTTPHeaderField: HTTPHeader.sessionId) == initialSessionID) let response = HTTPURLResponse( url: testEndpoint, statusCode: 404, httpVersion: "HTTP/1.1", headerFields: nil)! return (response, Data("Not Found".utf8)) @@ -409,6 +409,624 @@ import Testing } } + // MARK: - Additional HTTP Error Codes + // These tests verify handling of additional HTTP status codes per the MCP spec + + @Test("HTTP 400 Bad Request Error", .httpClientTransportSetup) + func testHTTPBadRequestError() async throws { + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let messageData = #"{"jsonrpc":"2.0","method":"test","id":1}"#.data(using: .utf8)! + + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 400, httpVersion: "HTTP/1.1", headerFields: nil)! + return (response, Data("Bad Request".utf8)) + } + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: false, + logger: nil + ) + try await transport.connect() + + do { + try await transport.send(messageData) + Issue.record("Expected send to throw an error for 400") + } catch let error as MCPError { + guard case .internalError(let message) = error else { + Issue.record("Expected MCPError.internalError, got \(error)") + throw error + } + #expect(message?.contains("Bad request") ?? false) + } catch { + Issue.record("Expected MCPError, got \(error)") + throw error + } + } + + @Test("HTTP 401 Unauthorized Error", .httpClientTransportSetup) + func testHTTPUnauthorizedError() async throws { + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let messageData = #"{"jsonrpc":"2.0","method":"test","id":1}"#.data(using: .utf8)! + + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 401, httpVersion: "HTTP/1.1", headerFields: nil)! + return (response, Data("Unauthorized".utf8)) + } + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: false, + logger: nil + ) + try await transport.connect() + + do { + try await transport.send(messageData) + Issue.record("Expected send to throw an error for 401") + } catch let error as MCPError { + guard case .internalError(let message) = error else { + Issue.record("Expected MCPError.internalError, got \(error)") + throw error + } + #expect(message?.contains("Authentication required") ?? false) + } catch { + Issue.record("Expected MCPError, got \(error)") + throw error + } + } + + @Test("HTTP 403 Forbidden Error", .httpClientTransportSetup) + func testHTTPForbiddenError() async throws { + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let messageData = #"{"jsonrpc":"2.0","method":"test","id":1}"#.data(using: .utf8)! + + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 403, httpVersion: "HTTP/1.1", headerFields: nil)! + return (response, Data("Forbidden".utf8)) + } + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: false, + logger: nil + ) + try await transport.connect() + + do { + try await transport.send(messageData) + Issue.record("Expected send to throw an error for 403") + } catch let error as MCPError { + guard case .internalError(let message) = error else { + Issue.record("Expected MCPError.internalError, got \(error)") + throw error + } + #expect(message?.contains("Access forbidden") ?? false) + } catch { + Issue.record("Expected MCPError, got \(error)") + throw error + } + } + + @Test("HTTP 405 Method Not Allowed Error", .httpClientTransportSetup) + func testHTTPMethodNotAllowedError() async throws { + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let messageData = #"{"jsonrpc":"2.0","method":"test","id":1}"#.data(using: .utf8)! + + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 405, httpVersion: "HTTP/1.1", headerFields: nil)! + return (response, Data("Method Not Allowed".utf8)) + } + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: false, + logger: nil + ) + try await transport.connect() + + do { + try await transport.send(messageData) + Issue.record("Expected send to throw an error for 405") + } catch let error as MCPError { + guard case .internalError(let message) = error else { + Issue.record("Expected MCPError.internalError, got \(error)") + throw error + } + #expect(message?.contains("Method not allowed") ?? false) + } catch { + Issue.record("Expected MCPError, got \(error)") + throw error + } + } + + @Test("HTTP 408 Request Timeout Error", .httpClientTransportSetup) + func testHTTPRequestTimeoutError() async throws { + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let messageData = #"{"jsonrpc":"2.0","method":"test","id":1}"#.data(using: .utf8)! + + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 408, httpVersion: "HTTP/1.1", headerFields: nil)! + return (response, Data("Request Timeout".utf8)) + } + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: false, + logger: nil + ) + try await transport.connect() + + do { + try await transport.send(messageData) + Issue.record("Expected send to throw an error for 408") + } catch let error as MCPError { + guard case .internalError(let message) = error else { + Issue.record("Expected MCPError.internalError, got \(error)") + throw error + } + #expect(message?.contains("Request timeout") ?? false) + } catch { + Issue.record("Expected MCPError, got \(error)") + throw error + } + } + + @Test("HTTP 429 Too Many Requests Error", .httpClientTransportSetup) + func testHTTPTooManyRequestsError() async throws { + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let messageData = #"{"jsonrpc":"2.0","method":"test","id":1}"#.data(using: .utf8)! + + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 429, httpVersion: "HTTP/1.1", headerFields: nil)! + return (response, Data("Too Many Requests".utf8)) + } + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: false, + logger: nil + ) + try await transport.connect() + + do { + try await transport.send(messageData) + Issue.record("Expected send to throw an error for 429") + } catch let error as MCPError { + guard case .internalError(let message) = error else { + Issue.record("Expected MCPError.internalError, got \(error)") + throw error + } + #expect(message?.contains("Too many requests") ?? false) + } catch { + Issue.record("Expected MCPError, got \(error)") + throw error + } + } + + @Test("HTTP 202 Accepted with no content", .httpClientTransportSetup) + func testHTTP202AcceptedNoContent() async throws { + // TypeScript SDK tests: 'should send JSON-RPC messages via POST' with status 202 + // This verifies that 202 Accepted responses (no content body) are handled correctly + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let messageData = #"{"jsonrpc":"2.0","method":"notifications/initialized"}"#.data( + using: .utf8)! + + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + // Server accepts the notification with 202 and no response body + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 202, httpVersion: "HTTP/1.1", + headerFields: [HTTPHeader.contentType: "application/json"])! + return (response, Data()) // Empty body for 202 Accepted + } + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: false, + logger: nil + ) + try await transport.connect() + + // Should not throw - 202 is a valid success response + try await transport.send(messageData) + } + + @Test("Unexpected content-type throws error for requests", .httpClientTransportSetup) + func testUnexpectedContentTypeThrowsErrorForRequest() async throws { + // Per MCP spec: requests MUST receive application/json or text/event-stream + // This aligns with TypeScript/Python SDKs which validate content-type for requests + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + // Request has both "method" and "id" - content-type validation applies + let messageData = #"{"jsonrpc":"2.0","method":"test","id":1}"#.data(using: .utf8)! + + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + // Server returns unexpected content-type (text/plain instead of application/json) + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [HTTPHeader.contentType: "text/plain"])! + return (response, Data("unexpected plain text response".utf8)) + } + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: false, + logger: nil + ) + try await transport.connect() + + do { + try await transport.send(messageData) + Issue.record("Expected send to throw an error for unexpected content-type") + } catch let error as MCPError { + guard case .internalError(let message) = error else { + Issue.record("Expected MCPError.internalError, got \(error)") + throw error + } + #expect(message?.contains("Unexpected content type") ?? false) + } catch { + Issue.record("Expected MCPError, got \(error)") + throw error + } + } + + @Test("Unexpected content-type ignored for notifications", .httpClientTransportSetup) + func testUnexpectedContentTypeIgnoredForNotification() async throws { + // Per MCP spec: notifications expect 202 Accepted with no body + // Content-type validation does not apply to notifications + // This aligns with TypeScript/Python SDKs behavior + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + // Notification has "method" but NO "id" - content-type validation does not apply + let messageData = #"{"jsonrpc":"2.0","method":"notifications/initialized"}"#.data( + using: .utf8)! + + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + // Server returns unexpected content-type with body (unusual but allowed for notifications) + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [HTTPHeader.contentType: "text/plain"])! + return (response, Data("some unexpected response".utf8)) + } + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: false, + logger: nil + ) + try await transport.connect() + + // Should not throw - notifications don't require content-type validation + try await transport.send(messageData) + } + + @Test("Empty response with unexpected content-type does not throw", .httpClientTransportSetup) + func testEmptyResponseUnexpectedContentTypeNoError() async throws { + // Even for requests, empty responses with unexpected content-type are acceptable + // (e.g., server returns 200 OK with empty body instead of proper 202/204) + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + // Request has both "method" and "id" + let messageData = #"{"jsonrpc":"2.0","method":"test","id":1}"#.data(using: .utf8)! + + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + // Server returns unexpected content-type but empty body + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [HTTPHeader.contentType: "text/plain"])! + return (response, Data()) // Empty body + } + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: false, + logger: nil + ) + try await transport.connect() + + // Should not throw - empty body is acceptable even with unexpected content-type + try await transport.send(messageData) + } + + // MARK: - Protocol Version Header Tests + + @Test("Protocol version header sent after initialization", .httpClientTransportSetup) + func testProtocolVersionHeaderSentAfterInit() async throws { + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: false, + logger: nil + ) + try await transport.connect() + + let firstMessageData = #"{"jsonrpc":"2.0","method":"initialize","id":1}"#.data(using: .utf8)! + let secondMessageData = #"{"jsonrpc":"2.0","method":"ping","id":2}"#.data(using: .utf8)! + let protocolVersion = Version.v2024_11_05 + + // First request - no protocol version header expected + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + // Before initialization, no protocol version header should be sent + #expect(request.value(forHTTPHeaderField: HTTPHeader.protocolVersion) == nil) + + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [HTTPHeader.contentType: "application/json"])! + return (response, Data()) + } + try await transport.send(firstMessageData) + + // Set the protocol version (simulating what Client does after init) + await transport.setProtocolVersion(protocolVersion) + + // Second request - protocol version header should be present + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint, protocolVersion] (request: URLRequest) in + #expect(request.value(forHTTPHeaderField: HTTPHeader.protocolVersion) == protocolVersion) + + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [HTTPHeader.contentType: "application/json"])! + return (response, Data()) + } + try await transport.send(secondMessageData) + } + + // MARK: - Session Termination Tests + + @Test("Terminate session sends DELETE request", .httpClientTransportSetup) + func testTerminateSessionSendsDelete() async throws { + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: false, + logger: nil + ) + try await transport.connect() + + let sessionID = "session-to-terminate-123" + let initMessageData = #"{"jsonrpc":"2.0","method":"initialize","id":1}"#.data(using: .utf8)! + + // First, establish a session + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint, sessionID] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [ + HTTPHeader.contentType: "application/json", + HTTPHeader.sessionId: sessionID, + ])! + return (response, Data()) + } + try await transport.send(initMessageData) + #expect(await transport.sessionID == sessionID) + + // Now set up handler for DELETE request + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint, sessionID] (request: URLRequest) in + #expect(request.httpMethod == "DELETE") + #expect(request.value(forHTTPHeaderField: HTTPHeader.sessionId) == sessionID) + + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 204, httpVersion: "HTTP/1.1", headerFields: nil)! + return (response, Data()) + } + + try await transport.terminateSession() + + // Session ID should be cleared + #expect(await transport.sessionID == nil) + } + + @Test("Terminate session handles 405 gracefully", .httpClientTransportSetup) + func testTerminateSessionHandles405() async throws { + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: false, + logger: nil + ) + try await transport.connect() + + let sessionID = "session-405-test" + let initMessageData = #"{"jsonrpc":"2.0","method":"initialize","id":1}"#.data(using: .utf8)! + + // First, establish a session + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint, sessionID] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [ + HTTPHeader.contentType: "application/json", + HTTPHeader.sessionId: sessionID, + ])! + return (response, Data()) + } + try await transport.send(initMessageData) + + // Server returns 405 - doesn't support session termination + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 405, httpVersion: "HTTP/1.1", headerFields: nil)! + return (response, Data()) + } + + // Should not throw - 405 is handled gracefully per spec + try await transport.terminateSession() + + // Session ID is NOT cleared when server returns 405 + // (server doesn't support termination, session may still be valid) + #expect(await transport.sessionID == sessionID) + } + + @Test("Terminate session handles 404 (session expired)", .httpClientTransportSetup) + func testTerminateSessionHandles404() async throws { + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: false, + logger: nil + ) + try await transport.connect() + + let sessionID = "session-404-test" + let initMessageData = #"{"jsonrpc":"2.0","method":"initialize","id":1}"#.data(using: .utf8)! + + // First, establish a session + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint, sessionID] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [ + HTTPHeader.contentType: "application/json", + HTTPHeader.sessionId: sessionID, + ])! + return (response, Data()) + } + try await transport.send(initMessageData) + + // Server returns 404 - session already expired + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 404, httpVersion: "HTTP/1.1", headerFields: nil)! + return (response, Data()) + } + + // Should not throw - 404 means session already gone + try await transport.terminateSession() + + // Session ID should be cleared + #expect(await transport.sessionID == nil) + } + + @Test("Terminate session with no session ID does nothing", .httpClientTransportSetup) + func testTerminateSessionNoSessionId() async throws { + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: false, + logger: nil + ) + try await transport.connect() + + // No session established - should return early without making request + #expect(await transport.sessionID == nil) + + // This should not throw and should not make any HTTP request + try await transport.terminateSession() + + #expect(await transport.sessionID == nil) + } + + @Test("Terminate session includes protocol version header", .httpClientTransportSetup) + func testTerminateSessionIncludesProtocolVersion() async throws { + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: false, + logger: nil + ) + try await transport.connect() + + let sessionID = "session-protocol-version-test" + let protocolVersion = Version.v2024_11_05 + let initMessageData = #"{"jsonrpc":"2.0","method":"initialize","id":1}"#.data(using: .utf8)! + + // First, establish a session + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint, sessionID] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [ + HTTPHeader.contentType: "application/json", + HTTPHeader.sessionId: sessionID, + ])! + return (response, Data()) + } + try await transport.send(initMessageData) + + // Set protocol version + await transport.setProtocolVersion(protocolVersion) + + // Verify DELETE includes protocol version + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint, sessionID, protocolVersion] (request: URLRequest) in + #expect(request.httpMethod == "DELETE") + #expect(request.value(forHTTPHeaderField: HTTPHeader.sessionId) == sessionID) + #expect(request.value(forHTTPHeaderField: HTTPHeader.protocolVersion) == protocolVersion) + + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 204, httpVersion: "HTTP/1.1", headerFields: nil)! + return (response, Data()) + } + + try await transport.terminateSession() + } + // Skip SSE tests on platforms that don't support streaming #if !canImport(FoundationNetworking) @Test("Receive Server-Sent Event (SSE)", .httpClientTransportSetup) @@ -424,56 +1042,837 @@ import Testing logger: nil ) - let eventString = "id: event1\ndata: {\"key\":\"value\"}\n\n" - let sseEventData = eventString.data(using: .utf8)! - - // First, set up a handler for the initial POST that will provide a session ID + let eventString = "id: event1\ndata: {\"key\":\"value\"}\n\n" + let sseEventData = eventString.data(using: .utf8)! + + // First, set up a handler for the initial POST that will provide a session ID + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [ + HTTPHeader.contentType: "text/plain", + HTTPHeader.sessionId: "test-session-123", + ])! + return (response, Data()) + } + + // Connect and send a dummy message to get the session ID + try await transport.connect() + try await transport.send(Data()) + + // Now set up the handler for the SSE GET request + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint, sseEventData] (request: URLRequest) in // sseEventData is now empty Data() + #expect(request.url == testEndpoint) + #expect(request.httpMethod == "GET") + #expect(request.value(forHTTPHeaderField: "Accept") == "text/event-stream") + #expect( + request.value(forHTTPHeaderField: HTTPHeader.sessionId) == "test-session-123") + + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [HTTPHeader.contentType: "text/event-stream"])! + + return (response, sseEventData) // Will return empty Data for SSE + } + + try await Task.sleep(for: .milliseconds(100)) + + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + + let expectedData = #"{"key":"value"}"#.data(using: .utf8)! + let receivedData = try await iterator.next() + + #expect(receivedData == expectedData) + + await transport.disconnect() + } + + @Test("Receive Server-Sent Event (SSE) (CR-NL)", .httpClientTransportSetup) + func testReceiveSSE_CRNL() async throws { + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: true, + sseInitializationTimeout: 1, + logger: nil + ) + + let eventString = "id: event1\r\ndata: {\"key\":\"value\"}\r\n\n" + let sseEventData = eventString.data(using: .utf8)! + + // First, set up a handler for the initial POST that will provide a session ID + // Use text/plain to prevent its (empty) body from being yielded to messageStream + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [ + HTTPHeader.contentType: "text/plain", + HTTPHeader.sessionId: "test-session-123", + ])! + return (response, Data()) + } + + // Connect and send a dummy message to get the session ID + try await transport.connect() + try await transport.send(Data()) + + // Now set up the handler for the SSE GET request + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint, sseEventData] (request: URLRequest) in + #expect(request.url == testEndpoint) + #expect(request.httpMethod == "GET") + #expect(request.value(forHTTPHeaderField: "Accept") == "text/event-stream") + #expect( + request.value(forHTTPHeaderField: HTTPHeader.sessionId) == "test-session-123") + + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [HTTPHeader.contentType: "text/event-stream"])! + + return (response, sseEventData) + } + + try await Task.sleep(for: .milliseconds(100)) + + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + + let expectedData = #"{"key":"value"}"#.data(using: .utf8)! + let receivedData = try await iterator.next() + + #expect(receivedData == expectedData) + + await transport.disconnect() + } + + @Test( + "Client with HTTP Transport complete flow", .httpClientTransportSetup, + .timeLimit(.minutes(1))) + func testClientFlow() async throws { + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: false, + logger: nil + ) + + let client = Client(name: "TestClient", version: "1.0.0") + + // Use an actor to track request sequence + actor RequestTracker { + enum RequestType { + case initialize + case callTool + } + + private(set) var lastRequest: RequestType? + + func setRequest(_ type: RequestType) { + lastRequest = type + } + + func getLastRequest() -> RequestType? { + return lastRequest + } + } + + let tracker = RequestTracker() + + // Setup mock responses + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint, tracker] (request: URLRequest) in + switch request.httpMethod { + case "GET": + #expect( + request.allHTTPHeaderFields?["Accept"]?.contains("text/event-stream") + == true) + case "POST": + #expect( + request.allHTTPHeaderFields?["Accept"]?.contains("application/json") + == true + ) + default: + Issue.record( + "Unsupported HTTP method \(String(describing: request.httpMethod))") + } + + #expect(request.url == testEndpoint) + + let bodyData = request.readBody() + + guard let bodyData = bodyData, + let json = try JSONSerialization.jsonObject(with: bodyData) + as? [String: Any], + let method = json["method"] as? String + else { + throw NSError( + domain: "MockURLProtocolError", code: 0, + userInfo: [ + NSLocalizedDescriptionKey: + "Invalid JSON-RPC message \(#file):\(#line)" + ]) + } + + if method == "initialize" { + await tracker.setRequest(.initialize) + + let requestID = json["id"] as! String + let result = Initialize.Result( + protocolVersion: Version.latest, + capabilities: .init(tools: .init()), + serverInfo: .init(name: "Mock Server", version: "0.0.1"), + instructions: nil + ) + let response = Initialize.response(id: .string(requestID), result: result) + let responseData = try JSONEncoder().encode(response) + + let httpResponse = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [HTTPHeader.contentType: "application/json"])! + return (httpResponse, responseData) + } else if method == "tools/call" { + // Verify initialize was called first + if let lastRequest = await tracker.getLastRequest(), + lastRequest != .initialize + { + #expect(Bool(false), "Initialize should be called before callTool") + } + + await tracker.setRequest(.callTool) + + let params = json["params"] as? [String: Any] + let toolName = params?["name"] as? String + #expect(toolName == "calculator") + + let requestID = json["id"] as! String + let result = CallTool.Result(content: [.text("42")]) + let response = CallTool.response(id: .string(requestID), result: result) + let responseData = try JSONEncoder().encode(response) + + let httpResponse = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [HTTPHeader.contentType: "application/json"])! + return (httpResponse, responseData) + } else if method == "notifications/initialized" { + // Ignore initialized notifications + let httpResponse = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [HTTPHeader.contentType: "application/json"])! + return (httpResponse, Data()) + } else { + throw NSError( + domain: "MockURLProtocolError", code: 0, + userInfo: [ + NSLocalizedDescriptionKey: + "Unexpected request method: \(method) \(#file):\(#line)" + ]) + } + } + + // Step 1: Initialize client + let initResult = try await client.connect(transport: transport) + #expect(initResult.protocolVersion == Version.latest) + #expect(initResult.capabilities.tools != nil) + + // Step 2: Call a tool + let toolResult = try await client.callTool(name: "calculator") + #expect(toolResult.content.count == 1) + if case .text(let text, _, _) = toolResult.content[0] { + #expect(text == "42") + } else { + #expect(Bool(false), "Expected text content") + } + + // Step 3: Verify request sequence + #expect(await tracker.getLastRequest() == .callTool) + + // Step 4: Disconnect + await client.disconnect() + } + + @Test("Request modifier functionality", .httpClientTransportSetup) + func testRequestModifier() async throws { + let testEndpoint = URL(string: "https://api.example.com/mcp")! + let testToken = "test-bearer-token-12345" + + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint, testToken] (request: URLRequest) in + // Verify the Authorization header was added by the requestModifier + #expect( + request.value(forHTTPHeaderField: "Authorization") == "Bearer \(testToken)") + + // Return a successful response + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [HTTPHeader.contentType: "application/json"])! + return (response, Data()) + } + + // Create transport with requestModifier that adds Authorization header + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: false, + requestModifier: { request in + var modifiedRequest = request + modifiedRequest.addValue( + "Bearer \(testToken)", forHTTPHeaderField: "Authorization") + return modifiedRequest + }, + logger: nil + ) + + try await transport.connect() + + let messageData = #"{"jsonrpc":"2.0","method":"test","id":5}"#.data(using: .utf8)! + + try await transport.send(messageData) + await transport.disconnect() + } + + // MARK: - Reconnection and Resumption Tests + // These tests verify the reconnection logic aligns with TypeScript/Python SDKs + + @Test("Custom reconnection options are respected", .httpClientTransportSetup) + func testCustomReconnectionOptions() async throws { + // TypeScript SDK test: 'should support custom reconnection options' + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let customOptions = HTTPReconnectionOptions( + initialReconnectionDelay: 0.5, + maxReconnectionDelay: 10.0, + reconnectionDelayGrowFactor: 2.0, + maxRetries: 5 + ) + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: false, + reconnectionOptions: customOptions, + logger: nil + ) + + // Verify options were set correctly + let options = transport.reconnectionOptions + #expect(options.initialReconnectionDelay == 0.5) + #expect(options.maxReconnectionDelay == 10.0) + #expect(options.reconnectionDelayGrowFactor == 2.0) + #expect(options.maxRetries == 5) + } + + @Test("Exponential backoff options configuration", .httpClientTransportSetup) + func testExponentialBackoffConfiguration() async throws { + // TypeScript SDK test: 'should have exponential backoff with configurable maxRetries' + // This test verifies that exponential backoff options can be configured + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let customOptions = HTTPReconnectionOptions( + initialReconnectionDelay: 0.1, // 100ms + maxReconnectionDelay: 5.0, // 5000ms + reconnectionDelayGrowFactor: 2.0, + maxRetries: 3 + ) + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: false, + reconnectionOptions: customOptions, + logger: nil + ) + + // Verify exponential backoff options are set correctly + let options = transport.reconnectionOptions + #expect(options.initialReconnectionDelay == 0.1) + #expect(options.maxReconnectionDelay == 5.0) + #expect(options.reconnectionDelayGrowFactor == 2.0) + #expect(options.maxRetries == 3) + + // The actual exponential backoff delay calculation is: + // delay = initialReconnectionDelay * pow(reconnectionDelayGrowFactor, attempt) + // Capped at maxReconnectionDelay + // This is tested indirectly through reconnection behavior + } + + @Test("Resumption token callback is invoked", .httpClientTransportSetup) + func testResumptionTokenCallback() async throws { + // TypeScript SDK test: related to 'onresumptiontoken' callback + // Python SDK: 'on_resumption_token_update' callback + // This test verifies the callback works by checking lastReceivedEventId + // which is set from the same event processing + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: true, + sseInitializationTimeout: 1, + logger: nil + ) + + // Set up handler for initial POST to get session ID + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [ + HTTPHeader.contentType: "text/plain", + HTTPHeader.sessionId: "test-session-resumption", + ])! + return (response, Data()) + } + + try await transport.connect() + try await transport.send(Data()) + + // Set up SSE response with event ID (priming event) + let sseWithEventId = "id: event-123\ndata: {\"test\":\"data\"}\n\n" + let sseData = sseWithEventId.data(using: .utf8)! + + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint, sseData] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [HTTPHeader.contentType: "text/event-stream"])! + return (response, sseData) + } + + try await Task.sleep(for: .milliseconds(200)) + + // Verify the event ID was captured (same mechanism as callback) + // The onResumptionToken callback and lastReceivedEventId are both set + // when an event with ID is received + let lastEventId = await transport.lastReceivedEventId + #expect(lastEventId == "event-123") + + await transport.disconnect() + } + + @Test("Last event ID is stored for resumption", .httpClientTransportSetup) + func testLastEventIdStoredForResumption() async throws { + // TypeScript SDK test: 'should pass lastEventId when reconnecting' + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: true, + sseInitializationTimeout: 1, + logger: nil + ) + + // Set up handler for initial POST + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [ + HTTPHeader.contentType: "text/plain", + HTTPHeader.sessionId: "test-session-last-event", + ])! + return (response, Data()) + } + + try await transport.connect() + try await transport.send(Data()) + + // Set up SSE response with event ID + let sseWithEventId = "id: last-event-456\ndata: {}\n\n" + let sseData = sseWithEventId.data(using: .utf8)! + + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint, sseData] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [HTTPHeader.contentType: "text/event-stream"])! + return (response, sseData) + } + + try await Task.sleep(for: .milliseconds(200)) + + // Verify last event ID is stored (via public API) + let lastEventId = await transport.lastReceivedEventId + #expect(lastEventId == "last-event-456") + + await transport.disconnect() + } + + @Test("SSE priming event with empty data does not throw", .httpClientTransportSetup) + func testPrimingEventEmptyDataNoError() async throws { + // TypeScript SDK test: 'should not throw JSON parse error on priming events with empty data' + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: true, + sseInitializationTimeout: 1, + logger: nil + ) + + // Set up handler for initial POST + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [ + HTTPHeader.contentType: "text/plain", + HTTPHeader.sessionId: "test-session-priming", + ])! + return (response, Data()) + } + + try await transport.connect() + try await transport.send(Data()) + + // Priming event: has ID but empty data (this is valid per MCP spec) + // Followed by a real message + let sseWithPriming = "id: priming-123\ndata: \n\nid: msg-456\ndata: {\"result\":\"ok\"}\n\n" + let sseData = sseWithPriming.data(using: .utf8)! + + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint, sseData] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [HTTPHeader.contentType: "text/event-stream"])! + return (response, sseData) + } + + try await Task.sleep(for: .milliseconds(200)) + + // Should not have thrown - priming events with empty data are valid + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + let receivedData = try await iterator.next() + + // Should only receive the actual message, not the priming event + let expectedData = #"{"result":"ok"}"#.data(using: .utf8)! + #expect(receivedData == expectedData) + + await transport.disconnect() + } + + @Test("Server retry directive does not cause errors", .httpClientTransportSetup) + func testServerRetryDirectiveHandled() async throws { + // TypeScript SDK test: 'should use server-provided retry value for reconnection delay' + // Python SDK: 'test_streamable_http_client_respects_retry_interval' + // This test verifies that SSE retry directives are handled without error + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: true, + sseInitializationTimeout: 1, + logger: nil + ) + + // Set up handler for initial POST + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [ + HTTPHeader.contentType: "text/plain", + HTTPHeader.sessionId: "test-session-retry", + ])! + return (response, Data()) + } + + try await transport.connect() + try await transport.send(Data()) + + // SSE response with retry directive (3000ms = 3 seconds) + // The transport should parse this without error + let sseWithRetry = "retry: 3000\nid: evt-1\ndata: {\"result\":\"ok\"}\n\n" + let sseData = sseWithRetry.data(using: .utf8)! + + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint, sseData] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [HTTPHeader.contentType: "text/event-stream"])! + return (response, sseData) + } + + try await Task.sleep(for: .milliseconds(200)) + + // Verify the event was processed successfully (no error thrown) + // The server retry value is stored internally for reconnection logic + let lastEventId = await transport.lastReceivedEventId + #expect(lastEventId == "evt-1") + + await transport.disconnect() + } + + @Test("Default reconnection options use exponential backoff", .httpClientTransportSetup) + func testDefaultReconnectionOptions() async throws { + // TypeScript SDK test: 'should fall back to exponential backoff when no server retry value' + // This test verifies that default reconnection options are properly configured + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + // Use default options + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: false, + logger: nil + ) + + // Verify default options are set correctly + let options = transport.reconnectionOptions + #expect(options.initialReconnectionDelay == 1.0) // Default: 1 second + #expect(options.maxReconnectionDelay == 30.0) // Default: 30 seconds + #expect(options.reconnectionDelayGrowFactor == 1.5) // Default: 1.5x growth + #expect(options.maxRetries == 2) // Default: 2 retries + + // Test that HTTPReconnectionOptions.default has the same values + let defaultOptions = HTTPReconnectionOptions.default + #expect(defaultOptions.initialReconnectionDelay == 1.0) + #expect(defaultOptions.maxReconnectionDelay == 30.0) + #expect(defaultOptions.reconnectionDelayGrowFactor == 1.5) + #expect(defaultOptions.maxRetries == 2) + } + + @Test("SSE notifications do not stop reconnection", .httpClientTransportSetup) + func testSSENotificationsDoNotStopReconnection() async throws { + // This test verifies that server notifications via SSE don't incorrectly + // mark receivedResponse=true (which would stop reconnection). + // Per MCP spec and TypeScript/Python SDKs, only actual JSON-RPC responses + // should stop reconnection. Server requests and notifications should not. + // + // Bug fix: Previously any non-empty SSE data would set receivedResponse=true. + // Now only JSON-RPC responses (with id + result/error, no method) do. + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: true, + sseInitializationTimeout: 1, + logger: nil + ) + + // Set up handler for initial POST + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [ + HTTPHeader.contentType: "text/plain", + HTTPHeader.sessionId: "test-session-notifications", + ])! + return (response, Data()) + } + + try await transport.connect() + try await transport.send(Data()) + + // SSE stream with: + // 1. A notification (has method, no id) - should NOT stop reconnection + // 2. A server request (has method AND id) - should NOT stop reconnection + // 3. A response (has id + result, no method) - SHOULD stop reconnection + // Note: SSE format requires no leading spaces on field lines + let sseWithMixedMessages = "id: evt-1\ndata: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/progress\",\"params\":{\"progress\":50}}\n\nid: evt-2\ndata: {\"jsonrpc\":\"2.0\",\"method\":\"sampling/createMessage\",\"id\":\"server-req-1\",\"params\":{}}\n\nid: evt-3\ndata: {\"jsonrpc\":\"2.0\",\"id\":1,\"result\":{\"status\":\"ok\"}}\n\n" + let sseData = sseWithMixedMessages.data(using: .utf8)! + + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint, sseData] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [HTTPHeader.contentType: "text/event-stream"])! + return (response, sseData) + } + + try await Task.sleep(for: .milliseconds(200)) + + // Verify all three messages were received + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + + // First: notification + let msg1 = try await iterator.next() + #expect(msg1 != nil) + let msg1String = String(data: msg1!, encoding: .utf8)! + #expect(msg1String.contains("notifications/progress")) + + // Second: server request + let msg2 = try await iterator.next() + #expect(msg2 != nil) + let msg2String = String(data: msg2!, encoding: .utf8)! + #expect(msg2String.contains("sampling/createMessage")) + + // Third: response + let msg3 = try await iterator.next() + #expect(msg3 != nil) + let msg3String = String(data: msg3!, encoding: .utf8)! + #expect(msg3String.contains("\"result\"")) + + // The lastReceivedEventId should be evt-3 (last event with ID) + let lastEventId = await transport.lastReceivedEventId + #expect(lastEventId == "evt-3") + + await transport.disconnect() + } + + @Test("SSE error response stops reconnection", .httpClientTransportSetup) + func testSSEErrorResponseStopsReconnection() async throws { + // Per JSON-RPC 2.0, error responses also count as responses and should + // stop reconnection. An error response has id + error fields, no method. + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: true, + sseInitializationTimeout: 1, + logger: nil + ) + + // Set up handler for initial POST + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [ + HTTPHeader.contentType: "text/plain", + HTTPHeader.sessionId: "test-session-error", + ])! + return (response, Data()) + } + + try await transport.connect() + try await transport.send(Data()) + + // SSE stream with an error response + let sseWithError = """ + id: evt-1 + data: {"jsonrpc":"2.0","id":1,"error":{"code":-32600,"message":"Invalid request"}} + + """ + let sseData = sseWithError.data(using: .utf8)! + + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint, sseData] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [HTTPHeader.contentType: "text/event-stream"])! + return (response, sseData) + } + + try await Task.sleep(for: .milliseconds(200)) + + // Verify error response was received + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + + let msg = try await iterator.next() + #expect(msg != nil) + let msgString = String(data: msg!, encoding: .utf8)! + #expect(msgString.contains("\"error\"")) + #expect(msgString.contains("\(ErrorCode.invalidRequest)")) + + await transport.disconnect() + } + + @Test("Response ID remapping with string ID", .httpClientTransportSetup) + func testResponseIdRemappingStringId() async throws { + // This test verifies that response IDs are remapped to the original + // request ID during stream resumption, aligning with TypeScript and + // Python SDK behavior. This is a defensive feature for edge cases + // where servers might send responses with different IDs during replay. + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: true, + sseInitializationTimeout: 1, + logger: nil + ) + + // Set up handler for initial connection await MockURLProtocol.requestHandlerStorage.setHandler { [testEndpoint] (request: URLRequest) in let response = HTTPURLResponse( url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", headerFields: [ - "Content-Type": "text/plain", - "Mcp-Session-Id": "test-session-123", + HTTPHeader.contentType: "text/plain", + HTTPHeader.sessionId: "test-session-remap", ])! return (response, Data()) } - // Connect and send a dummy message to get the session ID try await transport.connect() - try await transport.send(Data()) - // Now set up the handler for the SSE GET request - await MockURLProtocol.requestHandlerStorage.setHandler { - [testEndpoint, sseEventData] (request: URLRequest) in // sseEventData is now empty Data() - #expect(request.url == testEndpoint) - #expect(request.httpMethod == "GET") - #expect(request.value(forHTTPHeaderField: "Accept") == "text/event-stream") - #expect( - request.value(forHTTPHeaderField: "Mcp-Session-Id") == "test-session-123") + // SSE stream with a response that has a DIFFERENT ID than the original request + // The server sends id: "server-generated-id" but our original request had id: "original-req-42" + let sseWithDifferentId = """ + id: evt-1 + data: {"jsonrpc":"2.0","id":"server-generated-id","result":{"status":"ok"}} + + """ + let sseData = sseWithDifferentId.data(using: .utf8)! + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint, sseData] (request: URLRequest) in + // Verify Last-Event-ID header is sent + #expect(request.value(forHTTPHeaderField: HTTPHeader.lastEventId) == "last-evt-123") let response = HTTPURLResponse( url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", - headerFields: ["Content-Type": "text/event-stream"])! - - return (response, sseEventData) // Will return empty Data for SSE + headerFields: [HTTPHeader.contentType: "text/event-stream"])! + return (response, sseData) } - try await Task.sleep(for: .milliseconds(100)) + // Resume with original request ID + let originalRequestId: RequestId = "original-req-42" + try await transport.resumeStream(from: "last-evt-123", forRequestId: originalRequestId) + try await Task.sleep(for: .milliseconds(200)) + + // Verify the response was received with REMAPPED ID let stream = await transport.receive() var iterator = stream.makeAsyncIterator() - let expectedData = #"{"key":"value"}"#.data(using: .utf8)! - let receivedData = try await iterator.next() + let msg = try await iterator.next() + #expect(msg != nil) + let msgString = String(data: msg!, encoding: .utf8)! - #expect(receivedData == expectedData) + // The ID should be remapped to "original-req-42" + #expect(msgString.contains("\"id\":\"original-req-42\"")) + #expect(!msgString.contains("server-generated-id")) + #expect(msgString.contains("\"result\"")) await transport.disconnect() } - @Test("Receive Server-Sent Event (SSE) (CR-NL)", .httpClientTransportSetup) - func testReceiveSSE_CRNL() async throws { + @Test("Response ID remapping with numeric ID", .httpClientTransportSetup) + func testResponseIdRemappingNumericId() async throws { + // Test ID remapping with numeric IDs (JSON-RPC allows both string and number IDs) let configuration = URLSessionConfiguration.ephemeral configuration.protocolClasses = [MockURLProtocol.self] @@ -485,243 +1884,267 @@ import Testing logger: nil ) - let eventString = "id: event1\r\ndata: {\"key\":\"value\"}\r\n\n" - let sseEventData = eventString.data(using: .utf8)! - - // First, set up a handler for the initial POST that will provide a session ID - // Use text/plain to prevent its (empty) body from being yielded to messageStream + // Set up handler for initial connection await MockURLProtocol.requestHandlerStorage.setHandler { [testEndpoint] (request: URLRequest) in let response = HTTPURLResponse( url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", headerFields: [ - "Content-Type": "text/plain", - "Mcp-Session-Id": "test-session-123", + HTTPHeader.contentType: "text/plain", + HTTPHeader.sessionId: "test-session-remap-num", ])! return (response, Data()) } - // Connect and send a dummy message to get the session ID try await transport.connect() - try await transport.send(Data()) - // Now set up the handler for the SSE GET request - await MockURLProtocol.requestHandlerStorage.setHandler { - [testEndpoint, sseEventData] (request: URLRequest) in - #expect(request.url == testEndpoint) - #expect(request.httpMethod == "GET") - #expect(request.value(forHTTPHeaderField: "Accept") == "text/event-stream") - #expect( - request.value(forHTTPHeaderField: "Mcp-Session-Id") == "test-session-123") + // SSE stream with a response that has id: 999 (different from original) + let sseWithDifferentId = """ + id: evt-1 + data: {"jsonrpc":"2.0","id":999,"result":{"value":42}} + """ + let sseData = sseWithDifferentId.data(using: .utf8)! + + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint, sseData] (request: URLRequest) in let response = HTTPURLResponse( url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", - headerFields: ["Content-Type": "text/event-stream"])! - - return (response, sseEventData) + headerFields: [HTTPHeader.contentType: "text/event-stream"])! + return (response, sseData) } - try await Task.sleep(for: .milliseconds(100)) + // Resume with original numeric request ID + let originalRequestId: RequestId = 42 + try await transport.resumeStream(from: "last-evt-456", forRequestId: originalRequestId) + + try await Task.sleep(for: .milliseconds(200)) + // Verify the response was received with REMAPPED numeric ID let stream = await transport.receive() var iterator = stream.makeAsyncIterator() - let expectedData = #"{"key":"value"}"#.data(using: .utf8)! - let receivedData = try await iterator.next() + let msg = try await iterator.next() + #expect(msg != nil) + let msgString = String(data: msg!, encoding: .utf8)! - #expect(receivedData == expectedData) + // The ID should be remapped to 42 (numeric) + #expect(msgString.contains("\"id\":42")) + #expect(!msgString.contains("999")) + #expect(msgString.contains("\"result\"")) await transport.disconnect() } - @Test( - "Client with HTTP Transport complete flow", .httpClientTransportSetup, - .timeLimit(.minutes(1))) - func testClientFlow() async throws { + @Test("No ID remapping without originalRequestId", .httpClientTransportSetup) + func testNoRemappingWithoutOriginalRequestId() async throws { + // When originalRequestId is nil (default), IDs should NOT be remapped let configuration = URLSessionConfiguration.ephemeral configuration.protocolClasses = [MockURLProtocol.self] let transport = HTTPClientTransport( endpoint: testEndpoint, configuration: configuration, - streaming: false, + streaming: true, + sseInitializationTimeout: 1, logger: nil ) - let client = Client(name: "TestClient", version: "1.0.0") + // Set up handler for initial connection + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [ + HTTPHeader.contentType: "text/plain", + HTTPHeader.sessionId: "test-session-no-remap", + ])! + return (response, Data()) + } - // Use an actor to track request sequence - actor RequestTracker { - enum RequestType { - case initialize - case callTool - } + try await transport.connect() - private(set) var lastRequest: RequestType? + // SSE stream with a response + let sseResponse = """ + id: evt-1 + data: {"jsonrpc":"2.0","id":"original-id","result":{"status":"ok"}} - func setRequest(_ type: RequestType) { - lastRequest = type - } + """ + let sseData = sseResponse.data(using: .utf8)! - func getLastRequest() -> RequestType? { - return lastRequest - } + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint, sseData] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [HTTPHeader.contentType: "text/event-stream"])! + return (response, sseData) } - let tracker = RequestTracker() + // Resume WITHOUT providing originalRequestId (default nil) + try await transport.resumeStream(from: "last-evt-789") - // Setup mock responses - await MockURLProtocol.requestHandlerStorage.setHandler { - [testEndpoint, tracker] (request: URLRequest) in - switch request.httpMethod { - case "GET": - #expect( - request.allHTTPHeaderFields?["Accept"]?.contains("text/event-stream") - == true) - case "POST": - #expect( - request.allHTTPHeaderFields?["Accept"]?.contains("application/json") - == true - ) - default: - Issue.record( - "Unsupported HTTP method \(String(describing: request.httpMethod))") - } + try await Task.sleep(for: .milliseconds(200)) - #expect(request.url == testEndpoint) + // Verify the response was received with ORIGINAL ID (no remapping) + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() - let bodyData = request.readBody() + let msg = try await iterator.next() + #expect(msg != nil) + let msgString = String(data: msg!, encoding: .utf8)! - guard let bodyData = bodyData, - let json = try JSONSerialization.jsonObject(with: bodyData) - as? [String: Any], - let method = json["method"] as? String - else { - throw NSError( - domain: "MockURLProtocolError", code: 0, - userInfo: [ - NSLocalizedDescriptionKey: - "Invalid JSON-RPC message \(#file):\(#line)" - ]) - } + // The ID should remain as "original-id" + #expect(msgString.contains("\"id\":\"original-id\"")) - if method == "initialize" { - await tracker.setRequest(.initialize) + await transport.disconnect() + } - let requestID = json["id"] as! String - let result = Initialize.Result( - protocolVersion: Version.latest, - capabilities: .init(tools: .init()), - serverInfo: .init(name: "Mock Server", version: "0.0.1"), - instructions: nil - ) - let response = Initialize.response(id: .string(requestID), result: result) - let responseData = try JSONEncoder().encode(response) + @Test("Error response ID remapping", .httpClientTransportSetup) + func testErrorResponseIdRemapping() async throws { + // Per JSON-RPC 2.0, error responses are also responses and should have + // their IDs remapped. This aligns with Python SDK behavior (which handles + // both JSONRPCResponse and JSONRPCError), and is more complete than + // TypeScript which only handles success responses. + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] - let httpResponse = HTTPURLResponse( - url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", - headerFields: ["Content-Type": "application/json"])! - return (httpResponse, responseData) - } else if method == "tools/call" { - // Verify initialize was called first - if let lastRequest = await tracker.getLastRequest(), - lastRequest != .initialize - { - #expect(Bool(false), "Initialize should be called before callTool") - } + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: true, + sseInitializationTimeout: 1, + logger: nil + ) - await tracker.setRequest(.callTool) + // Set up handler for initial connection + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [ + HTTPHeader.contentType: "text/plain", + HTTPHeader.sessionId: "test-session-error-remap", + ])! + return (response, Data()) + } - let params = json["params"] as? [String: Any] - let toolName = params?["name"] as? String - #expect(toolName == "calculator") + try await transport.connect() - let requestID = json["id"] as! String - let result = CallTool.Result(content: [.text("42")]) - let response = CallTool.response(id: .string(requestID), result: result) - let responseData = try JSONEncoder().encode(response) + // SSE stream with an ERROR response that has a different ID + let sseWithError = """ + id: evt-1 + data: {"jsonrpc":"2.0","id":"server-error-id","error":{"code":-32600,"message":"Invalid request"}} - let httpResponse = HTTPURLResponse( - url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", - headerFields: ["Content-Type": "application/json"])! - return (httpResponse, responseData) - } else if method == "notifications/initialized" { - // Ignore initialized notifications - let httpResponse = HTTPURLResponse( - url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", - headerFields: ["Content-Type": "application/json"])! - return (httpResponse, Data()) - } else { - throw NSError( - domain: "MockURLProtocolError", code: 0, - userInfo: [ - NSLocalizedDescriptionKey: - "Unexpected request method: \(method) \(#file):\(#line)" - ]) - } + """ + let sseData = sseWithError.data(using: .utf8)! + + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint, sseData] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [HTTPHeader.contentType: "text/event-stream"])! + return (response, sseData) } - // Step 1: Initialize client - let initResult = try await client.connect(transport: transport) - #expect(initResult.protocolVersion == Version.latest) - #expect(initResult.capabilities.tools != nil) + // Resume with original request ID - error response ID should be remapped + let originalRequestId: RequestId = "my-failed-request" + try await transport.resumeStream(from: "last-evt", forRequestId: originalRequestId) - // Step 2: Call a tool - let toolResult = try await client.callTool(name: "calculator") - #expect(toolResult.content.count == 1) - if case let .text(text) = toolResult.content[0] { - #expect(text == "42") - } else { - #expect(Bool(false), "Expected text content") - } + try await Task.sleep(for: .milliseconds(200)) - // Step 3: Verify request sequence - #expect(await tracker.getLastRequest() == .callTool) + // Verify error response was received with REMAPPED ID + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() - // Step 4: Disconnect - await client.disconnect() - } + let msg = try await iterator.next() + #expect(msg != nil) + let msgString = String(data: msg!, encoding: .utf8)! - @Test("Request modifier functionality", .httpClientTransportSetup) - func testRequestModifier() async throws { - let testEndpoint = URL(string: "https://api.example.com/mcp")! - let testToken = "test-bearer-token-12345" + // The ID should be remapped to "my-failed-request" + #expect(msgString.contains("\"id\":\"my-failed-request\"")) + #expect(!msgString.contains("server-error-id")) + #expect(msgString.contains("\"error\"")) + #expect(msgString.contains("\(ErrorCode.invalidRequest)")) + + await transport.disconnect() + } + @Test("ID remapping only affects responses, not requests/notifications", .httpClientTransportSetup) + func testIdRemappingOnlyAffectsResponses() async throws { + // ID remapping should only apply to responses, not to server requests or notifications let configuration = URLSessionConfiguration.ephemeral configuration.protocolClasses = [MockURLProtocol.self] - await MockURLProtocol.requestHandlerStorage.setHandler { - [testEndpoint, testToken] (request: URLRequest) in - // Verify the Authorization header was added by the requestModifier - #expect( - request.value(forHTTPHeaderField: "Authorization") == "Bearer \(testToken)") - - // Return a successful response - let response = HTTPURLResponse( - url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", - headerFields: ["Content-Type": "application/json"])! - return (response, Data()) - } - - // Create transport with requestModifier that adds Authorization header let transport = HTTPClientTransport( endpoint: testEndpoint, configuration: configuration, - streaming: false, - requestModifier: { request in - var modifiedRequest = request - modifiedRequest.addValue( - "Bearer \(testToken)", forHTTPHeaderField: "Authorization") - return modifiedRequest - }, + streaming: true, + sseInitializationTimeout: 1, logger: nil ) + // Set up handler for initial connection + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [ + HTTPHeader.contentType: "text/plain", + HTTPHeader.sessionId: "test-session-selective", + ])! + return (response, Data()) + } + try await transport.connect() - let messageData = #"{"jsonrpc":"2.0","method":"test","id":5}"#.data(using: .utf8)! + // SSE stream with: + // 1. A server request (has method AND id) - should NOT be remapped + // 2. A notification (has method, no id) - should NOT be remapped + // 3. A response (has id + result, no method) - SHOULD be remapped + // Note: SSE format requires no leading spaces on field lines + let sseWithMixed = "id: evt-1\ndata: {\"jsonrpc\":\"2.0\",\"method\":\"sampling/createMessage\",\"id\":\"server-req-1\",\"params\":{}}\n\nid: evt-2\ndata: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/progress\",\"params\":{\"progress\":50}}\n\nid: evt-3\ndata: {\"jsonrpc\":\"2.0\",\"id\":\"server-resp-id\",\"result\":{\"status\":\"ok\"}}\n\n" + let sseData = sseWithMixed.data(using: .utf8)! + + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint, sseData] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [HTTPHeader.contentType: "text/event-stream"])! + return (response, sseData) + } + + // Resume with original request ID + let originalRequestId: RequestId = "my-original-request" + try await transport.resumeStream(from: "last-evt", forRequestId: originalRequestId) + + try await Task.sleep(for: .milliseconds(200)) + + // Verify all messages were received + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + + // First: server request - ID should NOT be remapped + let msg1 = try await iterator.next() + #expect(msg1 != nil) + let msg1String = String(data: msg1!, encoding: .utf8)! + #expect(msg1String.contains("\"id\":\"server-req-1\"")) // Original ID preserved + #expect(msg1String.contains("sampling/createMessage")) + + // Second: notification - no ID field, should pass through unchanged + let msg2 = try await iterator.next() + #expect(msg2 != nil) + let msg2String = String(data: msg2!, encoding: .utf8)! + #expect(msg2String.contains("notifications/progress")) + #expect(!msg2String.contains("my-original-request")) + + // Third: response - ID SHOULD be remapped + let msg3 = try await iterator.next() + #expect(msg3 != nil) + let msg3String = String(data: msg3!, encoding: .utf8)! + #expect(msg3String.contains("\"id\":\"my-original-request\"")) // Remapped ID + #expect(!msg3String.contains("server-resp-id")) // Original ID replaced + #expect(msg3String.contains("\"result\"")) - try await transport.send(messageData) await transport.disconnect() } #endif // !canImport(FoundationNetworking) diff --git a/Tests/MCPTests/HTTPIntegrationTests.swift b/Tests/MCPTests/HTTPIntegrationTests.swift new file mode 100644 index 00000000..786cd099 --- /dev/null +++ b/Tests/MCPTests/HTTPIntegrationTests.swift @@ -0,0 +1,536 @@ +import Foundation +import Testing + +@testable import MCP + +/// Integration tests for HTTP transport following TypeScript SDK patterns. +/// +/// These tests verify: +/// - Multi-client scenarios (10+ concurrent clients) +/// - Session lifecycle (create, use, delete) +/// - Stateful vs stateless mode +/// - Response routing +@Suite("HTTP Integration Tests") +struct HTTPIntegrationTests { + + // MARK: - Test Message Templates (matching TypeScript SDK) + + static let initializeMessage = TestPayloads.initializeRequest(id: "init-1", clientName: "test-client") + + static let toolsListMessage = TestPayloads.listToolsRequest(id: "tools-1") + + // MARK: - Helper Functions + + // MARK: - Initialization Tests (matching TypeScript SDK) + + @Test("Initialize server and generate session ID") + func initializeServerAndGenerateSessionId() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { UUID().uuidString }) + ) + try await transport.connect() + + let request = TestPayloads.postRequest(body: Self.initializeMessage) + let response = await transport.handleRequest(request) + + #expect(response.statusCode == 200) + #expect(response.headers[HTTPHeader.contentType] == "text/event-stream") + #expect(response.headers[HTTPHeader.sessionId] != nil) + } + + @Test("Reject second initialization request") + func rejectSecondInitializationRequest() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { UUID().uuidString }) + ) + try await transport.connect() + + // First initialize + let request1 = TestPayloads.postRequest(body: Self.initializeMessage) + let response1 = await transport.handleRequest(request1) + #expect(response1.statusCode == 200) + + let sessionId = response1.headers[HTTPHeader.sessionId]! + + // Second initialize - should fail + let secondInitMessage = TestPayloads.initializeRequest(id: "init-2", clientName: "test-client") + let request2 = TestPayloads.postRequest(body: secondInitMessage, sessionId: sessionId) + let response2 = await transport.handleRequest(request2) + + #expect(response2.statusCode == 400) + } + + @Test("Reject batch initialize request") + func rejectBatchInitializeRequest() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { UUID().uuidString }) + ) + try await transport.connect() + + let batchInitMessages = TestPayloads.batchRequest([ + TestPayloads.initializeRequest(id: "init-1", clientName: "test-client-1"), + TestPayloads.initializeRequest(id: "init-2", clientName: "test-client-2"), + ]) + let request = TestPayloads.postRequest(body: batchInitMessages) + let response = await transport.handleRequest(request) + + #expect(response.statusCode == 400) + } + + // MARK: - Session Validation Tests (matching TypeScript SDK) + + @Test("Reject requests without valid session ID") + func rejectRequestsWithoutValidSessionId() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { UUID().uuidString }) + ) + try await transport.connect() + + // Initialize first + let initRequest = TestPayloads.postRequest(body: Self.initializeMessage) + _ = await transport.handleRequest(initRequest) + + // Try without session ID + let request = TestPayloads.postRequest(body: Self.toolsListMessage) + let response = await transport.handleRequest(request) + + #expect(response.statusCode == 400) + } + + @Test("Reject invalid session ID") + func rejectInvalidSessionId() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { UUID().uuidString }) + ) + try await transport.connect() + + // Initialize first + let initRequest = TestPayloads.postRequest(body: Self.initializeMessage) + _ = await transport.handleRequest(initRequest) + + // Try with invalid session ID + let request = TestPayloads.postRequest(body: Self.toolsListMessage, sessionId: "invalid-session-id") + let response = await transport.handleRequest(request) + + #expect(response.statusCode == 404) + } + + // MARK: - SSE Stream Tests (matching TypeScript SDK) + + @Test("Reject second SSE stream for same session") + func rejectSecondSSEStreamForSameSession() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { "test-session" }) + ) + try await transport.connect() + + // Initialize + let initRequest = TestPayloads.postRequest(body: Self.initializeMessage) + let initResponse = await transport.handleRequest(initRequest) + #expect(initResponse.statusCode == 200) + + // First GET - should succeed + let getRequest1 = HTTPRequest( + method: "GET", + headers: [ + HTTPHeader.accept: "text/event-stream", + HTTPHeader.sessionId: "test-session", + HTTPHeader.protocolVersion: Version.v2024_11_05, + ] + ) + let response1 = await transport.handleRequest(getRequest1) + #expect(response1.statusCode == 200) + + // Second GET - should fail (only one stream allowed) + let getRequest2 = HTTPRequest( + method: "GET", + headers: [ + HTTPHeader.accept: "text/event-stream", + HTTPHeader.sessionId: "test-session", + HTTPHeader.protocolVersion: Version.v2024_11_05, + ] + ) + let response2 = await transport.handleRequest(getRequest2) + #expect(response2.statusCode == 409) + } + + @Test("Reject GET requests without Accept header") + func rejectGETRequestsWithoutAcceptHeader() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { "test-session" }) + ) + try await transport.connect() + + // Initialize + let initRequest = TestPayloads.postRequest(body: Self.initializeMessage) + _ = await transport.handleRequest(initRequest) + + // GET without proper Accept header + let getRequest = HTTPRequest( + method: "GET", + headers: [ + HTTPHeader.accept: "application/json", // Wrong Accept header + HTTPHeader.sessionId: "test-session", + HTTPHeader.protocolVersion: Version.v2024_11_05, + ] + ) + let response = await transport.handleRequest(getRequest) + #expect(response.statusCode == 406) + } + + // MARK: - Content Type Validation (matching TypeScript SDK) + + @Test("Reject POST requests without proper Accept header") + func rejectPOSTRequestsWithoutProperAcceptHeader() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { "test-session" }) + ) + try await transport.connect() + + // Initialize + let initRequest = TestPayloads.postRequest(body: Self.initializeMessage) + _ = await transport.handleRequest(initRequest) + + // POST without text/event-stream in Accept + let request = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json", // Missing text/event-stream + HTTPHeader.contentType: "application/json", + HTTPHeader.sessionId: "test-session", + HTTPHeader.protocolVersion: Version.v2024_11_05, + ], + body: Self.toolsListMessage.data(using: .utf8) + ) + let response = await transport.handleRequest(request) + #expect(response.statusCode == 406) + } + + @Test("Reject unsupported Content-Type") + func rejectUnsupportedContentType() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { "test-session" }) + ) + try await transport.connect() + + // Initialize + let initRequest = TestPayloads.postRequest(body: Self.initializeMessage) + _ = await transport.handleRequest(initRequest) + + // POST with wrong Content-Type + let request = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "text/plain", // Wrong Content-Type + HTTPHeader.sessionId: "test-session", + HTTPHeader.protocolVersion: Version.v2024_11_05, + ], + body: "This is plain text".data(using: .utf8) + ) + let response = await transport.handleRequest(request) + #expect(response.statusCode == 415) + } + + // MARK: - Notification Handling (matching TypeScript SDK) + + @Test("Handle JSON-RPC batch notification messages with 202 response") + func handleJSONRPCBatchNotificationMessagesWith202Response() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { "test-session" }) + ) + try await transport.connect() + + // Initialize + let initRequest = TestPayloads.postRequest(body: Self.initializeMessage) + _ = await transport.handleRequest(initRequest) + + // Send batch of notifications (no IDs) + let batchNotifications = """ + [{"jsonrpc":"2.0","method":"someNotification1","params":{}},{"jsonrpc":"2.0","method":"someNotification2","params":{}}] + """ + let request = TestPayloads.postRequest(body: batchNotifications, sessionId: "test-session") + let response = await transport.handleRequest(request) + + #expect(response.statusCode == 202) + } + + // MARK: - JSON Parsing (matching TypeScript SDK) + + @Test("Handle invalid JSON data properly") + func handleInvalidJSONDataProperly() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { "test-session" }) + ) + try await transport.connect() + + // Initialize + let initRequest = TestPayloads.postRequest(body: Self.initializeMessage) + _ = await transport.handleRequest(initRequest) + + // Send invalid JSON + let request = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + HTTPHeader.sessionId: "test-session", + HTTPHeader.protocolVersion: Version.v2024_11_05, + ], + body: "This is not valid JSON".data(using: .utf8) + ) + let response = await transport.handleRequest(request) + #expect(response.statusCode == 400) + } + + // MARK: - DELETE Tests (matching TypeScript SDK) + + @Test("Handle DELETE requests and close session properly") + func handleDELETERequestsAndCloseSession() async throws { + actor ClosedState { + var closed = false + func markClosed() { closed = true } + func isClosed() -> Bool { closed } + } + let state = ClosedState() + + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { "test-session" }, + onSessionClosed: { _ in await state.markClosed() } + ) + ) + try await transport.connect() + + // Initialize + let initRequest = TestPayloads.postRequest(body: Self.initializeMessage) + _ = await transport.handleRequest(initRequest) + + // DELETE + let deleteRequest = HTTPRequest( + method: "DELETE", + headers: [ + HTTPHeader.sessionId: "test-session", + HTTPHeader.protocolVersion: Version.v2024_11_05, + ] + ) + let response = await transport.handleRequest(deleteRequest) + + #expect(response.statusCode == 200) + let closed = await state.isClosed() + #expect(closed == true) + } + + @Test("Reject DELETE requests with invalid session ID") + func rejectDELETERequestsWithInvalidSessionId() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { "valid-session" }) + ) + try await transport.connect() + + // Initialize + let initRequest = TestPayloads.postRequest(body: Self.initializeMessage) + _ = await transport.handleRequest(initRequest) + + // DELETE with invalid session ID + let deleteRequest = HTTPRequest( + method: "DELETE", + headers: [ + HTTPHeader.sessionId: "invalid-session-id", + HTTPHeader.protocolVersion: Version.v2024_11_05, + ] + ) + let response = await transport.handleRequest(deleteRequest) + + #expect(response.statusCode == 404) + } + + // MARK: - Protocol Version Tests (matching TypeScript SDK) + + @Test("Accept requests with matching protocol version") + func acceptRequestsWithMatchingProtocolVersion() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { "test-session" }) + ) + try await transport.connect() + + // Initialize + let initRequest = TestPayloads.postRequest(body: Self.initializeMessage) + _ = await transport.handleRequest(initRequest) + + // Request with valid protocol version + let request = TestPayloads.postRequest(body: Self.toolsListMessage, sessionId: "test-session") + let response = await transport.handleRequest(request) + + #expect(response.statusCode == 200) + } + + @Test("Reject unsupported protocol version") + func rejectUnsupportedProtocolVersion() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { "test-session" }) + ) + try await transport.connect() + + // Initialize + let initRequest = TestPayloads.postRequest(body: Self.initializeMessage) + _ = await transport.handleRequest(initRequest) + + // Request with unsupported protocol version + let request = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + HTTPHeader.sessionId: "test-session", + HTTPHeader.protocolVersion: "1999-01-01", // Unsupported version + ], + body: Self.toolsListMessage.data(using: .utf8) + ) + let response = await transport.handleRequest(request) + + #expect(response.statusCode == 400) + } + + // MARK: - Stateless Mode Tests (matching TypeScript SDK) + + @Test("Stateless mode - no session ID validation") + func statelessModeNoSessionIdValidation() async throws { + // Stateless mode - no session ID generator + let transport = HTTPServerTransport() + try await transport.connect() + + // Initialize + let initRequest = TestPayloads.postRequest(body: Self.initializeMessage) + let initResponse = await transport.handleRequest(initRequest) + + #expect(initResponse.statusCode == 200) + // Should NOT have session ID header in stateless mode + #expect(initResponse.headers[HTTPHeader.sessionId] == nil) + + // Request without session ID should work in stateless mode + let toolsRequest = TestPayloads.postRequest(body: Self.toolsListMessage) + let toolsResponse = await transport.handleRequest(toolsRequest) + + #expect(toolsResponse.statusCode == 200) + } + + @Test("Stateless mode accepts requests with various session IDs") + func statelessModeAcceptsRequestsWithVariousSessionIds() async throws { + let transport = HTTPServerTransport() + try await transport.connect() + + // Initialize + let initRequest = TestPayloads.postRequest(body: Self.initializeMessage) + _ = await transport.handleRequest(initRequest) + + // Try with random session ID - should be accepted in stateless mode + let request1 = TestPayloads.postRequest(body: Self.toolsListMessage, sessionId: "random-id-1") + let response1 = await transport.handleRequest(request1) + #expect(response1.statusCode == 200) + + // Try with another random session ID - should also be accepted + let request2 = TestPayloads.postRequest(body: Self.toolsListMessage, sessionId: "different-id-2") + let response2 = await transport.handleRequest(request2) + #expect(response2.statusCode == 200) + } + + @Test("Stateless mode rejects second SSE stream") + func statelessModeRejectsSecondSSEStream() async throws { + // Despite no session ID requirement, only one SSE stream allowed + let transport = HTTPServerTransport() + try await transport.connect() + + // Initialize + let initRequest = TestPayloads.postRequest(body: Self.initializeMessage) + _ = await transport.handleRequest(initRequest) + + // First GET + let getRequest1 = HTTPRequest( + method: "GET", + headers: [HTTPHeader.accept: "text/event-stream"] + ) + let response1 = await transport.handleRequest(getRequest1) + #expect(response1.statusCode == 200) + + // Second GET - should be rejected + let getRequest2 = HTTPRequest( + method: "GET", + headers: [HTTPHeader.accept: "text/event-stream"] + ) + let response2 = await transport.handleRequest(getRequest2) + #expect(response2.statusCode == 409) + } + + // MARK: - Multi-Client Tests + + @Test("Ten concurrent clients") + func tenConcurrentClients() async throws { + // Simulate 10 clients connecting concurrently + // Each client gets its own transport + actor Counter { + var count = 0 + func increment() { count += 1 } + func value() -> Int { count } + } + let successCounter = Counter() + + await withTaskGroup(of: Int.self) { group in + for clientId in 0..<10 { + group.addTask { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { "session-\(clientId)" }) + ) + try? await transport.connect() + + let initMessage = TestPayloads.initializeRequest( + id: "\(clientId)", + clientName: "client-\(clientId)" + ) + let request = TestPayloads.postRequest(body: initMessage) + + let response = await transport.handleRequest(request) + return response.statusCode + } + } + + // Verify all clients connected successfully + for await statusCode in group { + if statusCode == 200 { + await successCounter.increment() + } + } + } + + let count = await successCounter.value() + #expect(count == 10) + } + + // MARK: - Session Callbacks + + @Test("Session initialized callback fires") + func sessionInitializedCallbackFires() async throws { + actor CallbackTracker { + var sessionId: String? + func set(_ id: String) { sessionId = id } + func get() -> String? { sessionId } + } + let tracker = CallbackTracker() + + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { "callback-test-session" }, + onSessionInitialized: { sessionId in + await tracker.set(sessionId) + } + ) + ) + try await transport.connect() + + let initRequest = TestPayloads.postRequest(body: Self.initializeMessage) + _ = await transport.handleRequest(initRequest) + + let trackedSessionId = await tracker.get() + #expect(trackedSessionId == "callback-test-session") + } +} diff --git a/Tests/MCPTests/Helpers/MockTransport.swift b/Tests/MCPTests/Helpers/MockTransport.swift index 5979e9ad..e35b8357 100644 --- a/Tests/MCPTests/Helpers/MockTransport.swift +++ b/Tests/MCPTests/Helpers/MockTransport.swift @@ -121,4 +121,49 @@ actor MockTransport: Transport { sentData.removeAll() dataToReceive.removeAll() } + + /// Queue a raw JSON string for the server to receive + func queueRaw(_ jsonString: String) { + if let data = jsonString.data(using: .utf8) { + queue(data: data) + } + } + + /// Wait until the sent message count reaches the expected value, with timeout. + /// - Parameters: + /// - count: The expected number of sent messages + /// - timeout: Maximum time to wait (default 2 seconds) + /// - Returns: `true` if the count was reached, `false` if timeout occurred + func waitForSentMessageCount( + _ count: Int, + timeout: Duration = .seconds(2) + ) async -> Bool { + let deadline = ContinuousClock.now.advanced(by: timeout) + while ContinuousClock.now < deadline { + if sentData.count >= count { + return true + } + try? await Task.sleep(for: .milliseconds(10)) + } + return sentData.count >= count + } + + /// Wait until a sent message matches the predicate, with timeout. + /// - Parameters: + /// - timeout: Maximum time to wait (default 2 seconds) + /// - predicate: Closure that returns true when the expected message is found + /// - Returns: `true` if a matching message was found, `false` if timeout occurred + func waitForSentMessage( + timeout: Duration = .seconds(2), + matching predicate: @escaping (String) -> Bool + ) async -> Bool { + let deadline = ContinuousClock.now.advanced(by: timeout) + while ContinuousClock.now < deadline { + if sentMessages.contains(where: predicate) { + return true + } + try? await Task.sleep(for: .milliseconds(10)) + } + return sentMessages.contains(where: predicate) + } } diff --git a/Tests/MCPTests/Helpers/TestPayloads.swift b/Tests/MCPTests/Helpers/TestPayloads.swift new file mode 100644 index 00000000..29a22259 --- /dev/null +++ b/Tests/MCPTests/Helpers/TestPayloads.swift @@ -0,0 +1,199 @@ +import Foundation + +@testable import MCP + +/// Common JSON-RPC payloads used in tests. +/// +/// These helpers centralize test payload construction to: +/// - Eliminate duplicate JSON strings across tests +/// - Ensure consistent use of version constants +/// - Make version-specific testing easier +enum TestPayloads { + + // MARK: - Default Values + + /// Default protocol version for tests (initial stable release). + static let defaultVersion = Version.v2024_11_05 + + // MARK: - Initialize + + /// Creates a JSON-RPC initialize request. + static func initializeRequest( + id: String = "1", + protocolVersion: String = defaultVersion, + clientName: String = "test", + clientVersion: String = "1.0" + ) -> String { + """ + {"jsonrpc":"2.0","method":"initialize","id":"\(id)","params":{"protocolVersion":"\(protocolVersion)","capabilities":{},"clientInfo":{"name":"\(clientName)","version":"\(clientVersion)"}}} + """ + } + + /// Creates a JSON-RPC initialize result. + static func initializeResult( + id: String = "1", + protocolVersion: String = defaultVersion, + serverName: String = "test", + serverVersion: String = "1.0" + ) -> String { + """ + {"jsonrpc":"2.0","result":{"protocolVersion":"\(protocolVersion)","capabilities":{},"serverInfo":{"name":"\(serverName)","version":"\(serverVersion)"}},"id":"\(id)"} + """ + } + + // MARK: - Initialized Notification + + /// Creates a JSON-RPC initialized notification. + static func initializedNotification() -> String { + """ + {"jsonrpc":"2.0","method":"notifications/initialized"} + """ + } + + // MARK: - Tools + + /// Creates a JSON-RPC tools/list request. + static func listToolsRequest(id: String = "2") -> String { + """ + {"jsonrpc":"2.0","method":"tools/list","id":"\(id)","params":{}} + """ + } + + /// Creates a JSON-RPC tools/call request. + static func callToolRequest( + id: String = "2", + name: String, + arguments: [String: Any] = [:] + ) -> String { + let argsJSON = arguments.isEmpty ? "{}" : serializeJSON(arguments) + return """ + {"jsonrpc":"2.0","method":"tools/call","id":"\(id)","params":{"name":"\(name)","arguments":\(argsJSON)}} + """ + } + + // MARK: - Resources + + /// Creates a JSON-RPC resources/list request. + static func listResourcesRequest(id: String = "2") -> String { + """ + {"jsonrpc":"2.0","method":"resources/list","id":"\(id)","params":{}} + """ + } + + /// Creates a JSON-RPC resources/read request. + static func readResourceRequest(id: String = "2", uri: String) -> String { + """ + {"jsonrpc":"2.0","method":"resources/read","id":"\(id)","params":{"uri":"\(uri)"}} + """ + } + + // MARK: - Prompts + + /// Creates a JSON-RPC prompts/list request. + static func listPromptsRequest(id: String = "2") -> String { + """ + {"jsonrpc":"2.0","method":"prompts/list","id":"\(id)","params":{}} + """ + } + + /// Creates a JSON-RPC prompts/get request. + static func getPromptRequest(id: String = "2", name: String) -> String { + """ + {"jsonrpc":"2.0","method":"prompts/get","id":"\(id)","params":{"name":"\(name)"}} + """ + } + + // MARK: - Ping + + /// Creates a JSON-RPC ping request. + static func pingRequest(id: String = "ping") -> String { + """ + {"jsonrpc":"2.0","method":"ping","id":"\(id)"} + """ + } + + // MARK: - Batch Requests + + /// Creates a batch of JSON-RPC requests. + static func batchRequest(_ requests: [String]) -> String { + "[\(requests.joined(separator: ","))]" + } + + // MARK: - Helpers + + private static func serializeJSON(_ dict: [String: Any]) -> String { + guard let data = try? JSONSerialization.data(withJSONObject: dict), + let string = String(data: data, encoding: .utf8) + else { + return "{}" + } + return string + } +} + +// MARK: - HTTPRequest Helpers + +extension TestPayloads { + + /// Creates an HTTP POST request for MCP. + static func postRequest( + body: String, + sessionId: String? = nil, + protocolVersion: String = defaultVersion, + lastEventId: String? = nil, + accept: String = "application/json, text/event-stream" + ) -> HTTPRequest { + var headers = [ + HTTPHeader.accept: accept, + HTTPHeader.contentType: "application/json", + HTTPHeader.protocolVersion: protocolVersion, + ] + if let sessionId { + headers[HTTPHeader.sessionId] = sessionId + } + if let lastEventId { + headers[HTTPHeader.lastEventId] = lastEventId + } + return HTTPRequest( + method: "POST", + headers: headers, + body: body.data(using: .utf8) + ) + } + + /// Creates an HTTP GET request for SSE streams. + static func getRequest( + sessionId: String, + protocolVersion: String = defaultVersion, + lastEventId: String? = nil + ) -> HTTPRequest { + var headers = [ + HTTPHeader.accept: "text/event-stream", + HTTPHeader.sessionId: sessionId, + HTTPHeader.protocolVersion: protocolVersion, + ] + if let lastEventId { + headers[HTTPHeader.lastEventId] = lastEventId + } + return HTTPRequest( + method: "GET", + headers: headers, + body: nil + ) + } + + /// Creates an HTTP DELETE request for session termination. + static func deleteRequest( + sessionId: String, + protocolVersion: String = defaultVersion + ) -> HTTPRequest { + HTTPRequest( + method: "DELETE", + headers: [ + HTTPHeader.sessionId: sessionId, + HTTPHeader.protocolVersion: protocolVersion, + ], + body: nil + ) + } +} diff --git a/Tests/MCPTests/IDTests.swift b/Tests/MCPTests/IDTests.swift index 0abf28e6..bf5e1103 100644 --- a/Tests/MCPTests/IDTests.swift +++ b/Tests/MCPTests/IDTests.swift @@ -10,34 +10,34 @@ import struct Foundation.UUID struct IDTests { @Test("String ID initialization and encoding") func testStringID() throws { - let id: ID = "test-id" + let id: RequestId = "test-id" #expect(id.description == "test-id") let encoder = JSONEncoder() let decoder = JSONDecoder() let data = try encoder.encode(id) - let decoded = try decoder.decode(ID.self, from: data) + let decoded = try decoder.decode(RequestId.self, from: data) #expect(decoded == id) } @Test("Number ID initialization and encoding") func testNumberID() throws { - let id: ID = 42 + let id: RequestId = 42 #expect(id.description == "42") let encoder = JSONEncoder() let decoder = JSONDecoder() let data = try encoder.encode(id) - let decoded = try decoder.decode(ID.self, from: data) + let decoded = try decoder.decode(RequestId.self, from: data) #expect(decoded == id) } @Test("Random ID generation") func testRandomID() throws { - let id1 = ID.random - let id2 = ID.random + let id1 = RequestId.random + let id2 = RequestId.random #expect(id1 != id2, "Random IDs should be unique") if case .string(let str) = id1 { diff --git a/Tests/MCPTests/InMemoryEventStoreTests.swift b/Tests/MCPTests/InMemoryEventStoreTests.swift new file mode 100644 index 00000000..1bd79f5f --- /dev/null +++ b/Tests/MCPTests/InMemoryEventStoreTests.swift @@ -0,0 +1,464 @@ +import Foundation +import Testing + +@testable import MCP + +/// Tests for InMemoryEventStore - event storage for resumability support. +@Suite("InMemory Event Store Tests") +struct InMemoryEventStoreTests { + + // MARK: - Basic Operations + + @Test("Initialization creates empty store") + func initialization() async { + let store = InMemoryEventStore() + let count = await store.eventCount + #expect(count == 0) + } + + @Test("Store event") + func storeEvent() async throws { + let store = InMemoryEventStore() + let message = #"{"jsonrpc":"2.0","result":"test","id":"1"}"#.data(using: .utf8)! + + let eventId = try await store.storeEvent(streamId: "stream-1", message: message) + + #expect(!eventId.isEmpty) + #expect(eventId.contains("stream-1")) + + let count = await store.eventCount + #expect(count == 1) + } + + @Test("Store multiple events") + func storeMultipleEvents() async throws { + let store = InMemoryEventStore() + + for i in 0..<5 { + let message = #"{"jsonrpc":"2.0","result":"\#(i)","id":"\#(i)"}"#.data(using: .utf8)! + _ = try await store.storeEvent(streamId: "stream-1", message: message) + } + + let count = await store.eventCount + #expect(count == 5) + } + + @Test("Stream ID for event ID") + func streamIdForEventId() async throws { + let store = InMemoryEventStore() + let message = #"{"jsonrpc":"2.0","result":"test","id":"1"}"#.data(using: .utf8)! + + let eventId = try await store.storeEvent(streamId: "my-stream-id", message: message) + + let streamId = await store.streamIdForEventId(eventId) + #expect(streamId == "my-stream-id") + } + + @Test("Stream ID for unknown event ID returns nil") + func streamIdForUnknownEventId() async { + let store = InMemoryEventStore() + + let streamId = await store.streamIdForEventId("unknown-event-id") + #expect(streamId == nil) + } + + @Test("Stream ID for event ID with underscores") + func streamIdForEventIdWithUnderscores() async throws { + let store = InMemoryEventStore() + let message = Data() + + // Stream ID with underscores + let eventId = try await store.storeEvent(streamId: "stream_with_underscores", message: message) + + let streamId = await store.streamIdForEventId(eventId) + #expect(streamId == "stream_with_underscores") + } + + // MARK: - Event Replay + + @Test("Replay events after") + func replayEventsAfter() async throws { + let store = InMemoryEventStore() + + // Store some events + var eventIds: [String] = [] + for i in 0..<5 { + let message = #"{"jsonrpc":"2.0","result":"\#(i)","id":"\#(i)"}"#.data(using: .utf8)! + let eventId = try await store.storeEvent(streamId: "stream-1", message: message) + eventIds.append(eventId) + } + + // Replay events after the second one + actor MessageCollector { + var messages: [String] = [] + func add(_ msg: String) { messages.append(msg) } + func get() -> [String] { messages } + } + let collector = MessageCollector() + + let streamId = try await store.replayEventsAfter(eventIds[1]) { _, message in + if let json = try? JSONSerialization.jsonObject(with: message) as? [String: Any], + let result = json["result"] as? String + { + await collector.add(result) + } + } + + #expect(streamId == "stream-1") + let replayedMessages = await collector.get() + #expect(replayedMessages == ["2", "3", "4"]) // Events 2, 3, 4 (after event 1) + } + + @Test("Replay events only from same stream") + func replayEventsOnlyFromSameStream() async throws { + let store = InMemoryEventStore() + + // Store events for two different streams + let message1 = #"{"stream":"1","id":"a"}"#.data(using: .utf8)! + let eventId1 = try await store.storeEvent(streamId: "stream-1", message: message1) + + let message2 = #"{"stream":"2","id":"b"}"#.data(using: .utf8)! + _ = try await store.storeEvent(streamId: "stream-2", message: message2) + + let message3 = #"{"stream":"1","id":"c"}"#.data(using: .utf8)! + _ = try await store.storeEvent(streamId: "stream-1", message: message3) + + // Replay from stream-1's first event + actor Counter { + var count = 0 + func increment() { count += 1 } + func value() -> Int { count } + } + let counter = Counter() + + _ = try await store.replayEventsAfter(eventId1) { _, _ in + await counter.increment() + } + + // Should only replay event "c" from stream-1 (not "b" from stream-2) + let replayedCount = await counter.value() + #expect(replayedCount == 1) + } + + @Test("Replay events with unknown event ID throws") + func replayEventsUnknownEventId() async { + let store = InMemoryEventStore() + + await #expect(throws: EventStoreError.self) { + _ = try await store.replayEventsAfter("unknown-event") { _, _ in } + } + } + + // MARK: - Cleanup + + @Test("Clear removes all events") + func clear() async throws { + let store = InMemoryEventStore() + + // Store some events + for _ in 0..<5 { + let message = Data() + _ = try await store.storeEvent(streamId: "stream", message: message) + } + + var count = await store.eventCount + #expect(count == 5) + + await store.clear() + + count = await store.eventCount + #expect(count == 0) + } + + @Test("Remove events for stream") + func removeEventsForStream() async throws { + let store = InMemoryEventStore() + + // Store events for two streams + for _ in 0..<3 { + _ = try await store.storeEvent(streamId: "stream-1", message: Data()) + } + for _ in 0..<2 { + _ = try await store.storeEvent(streamId: "stream-2", message: Data()) + } + + var count = await store.eventCount + #expect(count == 5) + + let removed = await store.removeEvents(forStream: "stream-1") + #expect(removed == 3) + + count = await store.eventCount + #expect(count == 2) + } + + @Test("Cleanup old events") + func cleanUpOldEvents() async throws { + let store = InMemoryEventStore() + + // Store an event + _ = try await store.storeEvent(streamId: "stream", message: Data()) + + var count = await store.eventCount + #expect(count == 1) + + // Clean up with zero age - should remove all + let removed = await store.cleanUp(olderThan: .zero) + #expect(removed == 1) + + count = await store.eventCount + #expect(count == 0) + } + + @Test("Cleanup does not remove recent events") + func cleanUpDoesNotRemoveRecentEvents() async throws { + let store = InMemoryEventStore() + + // Store an event + _ = try await store.storeEvent(streamId: "stream", message: Data()) + + // Clean up with 1 hour age - should not remove recent event + let removed = await store.cleanUp(olderThan: .seconds(3600)) + #expect(removed == 0) + + let count = await store.eventCount + #expect(count == 1) + } + + // MARK: - Concurrency + + @Test("Concurrent store and retrieve") + func concurrentStoreAndRetrieve() async throws { + let store = InMemoryEventStore() + + // Concurrently store events + await withTaskGroup(of: String.self) { group in + for i in 0..<100 { + group.addTask { + let message = Data() + return try! await store.storeEvent(streamId: "stream-\(i % 10)", message: message) + } + } + } + + let count = await store.eventCount + #expect(count == 100) + } + + @Test("Concurrent replay") + func concurrentReplay() async throws { + let store = InMemoryEventStore() + + // Store events for multiple streams + // Note: Use non-empty data because empty data is treated as priming events and skipped during replay + var firstEventIds: [String] = [] + for stream in 0..<5 { + let message = Data("test".utf8) + let eventId = try await store.storeEvent(streamId: "stream-\(stream)", message: message) + firstEventIds.append(eventId) + + // Add more events to each stream + for _ in 0..<10 { + _ = try await store.storeEvent(streamId: "stream-\(stream)", message: message) + } + } + + // Concurrently replay from each stream + actor Counter { + var value = 0 + func increment() { value += 1 } + func reset() -> Int { + let v = value + value = 0 + return v + } + } + + await withTaskGroup(of: Int.self) { group in + for (_, eventId) in firstEventIds.enumerated() { + group.addTask { + let counter = Counter() + _ = try? await store.replayEventsAfter(eventId) { _, _ in + await counter.increment() + } + return await counter.reset() + } + } + + for await count in group { + #expect(count == 10) // Each stream should replay 10 events + } + } + } + + // MARK: - Event ID Format + + @Test("Event ID contains stream ID") + func eventIdContainsStreamId() async throws { + let store = InMemoryEventStore() + let message = Data() + + let eventId = try await store.storeEvent(streamId: "unique-stream-123", message: message) + + #expect(eventId.hasPrefix("unique-stream-123_")) + } + + @Test("Event IDs are unique") + func eventIdsAreUnique() async throws { + let store = InMemoryEventStore() + let message = Data() + + var eventIds = Set() + for _ in 0..<100 { + let eventId = try await store.storeEvent(streamId: "stream", message: message) + eventIds.insert(eventId) + } + + #expect(eventIds.count == 100) // All IDs should be unique + } + + // MARK: - Max Events Per Stream + + @Test("Default maxEventsPerStream is 100") + func defaultMaxEventsPerStream() async { + let store = InMemoryEventStore() + #expect(store.maxEventsPerStream == 100) + } + + @Test("Custom maxEventsPerStream is respected") + func customMaxEventsPerStream() async { + let store = InMemoryEventStore(maxEventsPerStream: 50) + #expect(store.maxEventsPerStream == 50) + } + + @Test("Automatic eviction when max events reached") + func automaticEviction() async throws { + let store = InMemoryEventStore(maxEventsPerStream: 5) + let message = Data("test".utf8) + + // Store 5 events (at capacity) + var eventIds: [String] = [] + for i in 0..<5 { + let msg = #"{"id":"\#(i)"}"#.data(using: .utf8)! + let eventId = try await store.storeEvent(streamId: "stream", message: msg) + eventIds.append(eventId) + } + + var count = await store.eventCount + #expect(count == 5) + + // Store one more - should evict the oldest + let newEventId = try await store.storeEvent(streamId: "stream", message: message) + + count = await store.eventCount + #expect(count == 5) // Still 5 events + + // The oldest event should be evicted + let oldestStreamId = await store.streamIdForEventId(eventIds[0]) + // The event is no longer in the index, so we fall back to parsing + #expect(oldestStreamId == "stream") // Parsing still works + + // But replay should fail for the evicted event + await #expect(throws: EventStoreError.self) { + _ = try await store.replayEventsAfter(eventIds[0]) { _, _ in } + } + + // The new event should be retrievable + let newStreamId = await store.streamIdForEventId(newEventId) + #expect(newStreamId == "stream") + } + + @Test("Eviction is per-stream") + func evictionIsPerStream() async throws { + let store = InMemoryEventStore(maxEventsPerStream: 3) + let message = Data("test".utf8) + + // Fill stream-1 to capacity + for _ in 0..<3 { + _ = try await store.storeEvent(streamId: "stream-1", message: message) + } + + // Fill stream-2 to capacity + for _ in 0..<3 { + _ = try await store.storeEvent(streamId: "stream-2", message: message) + } + + var count = await store.eventCount + #expect(count == 6) // 3 per stream + + // Add to stream-1 - should only evict from stream-1 + _ = try await store.storeEvent(streamId: "stream-1", message: message) + + count = await store.eventCount + #expect(count == 6) // Still 6 total (3 + 3) + + let streamCount = await store.streamCount + #expect(streamCount == 2) + } + + @Test("Replay works correctly after eviction") + func replayAfterEviction() async throws { + let store = InMemoryEventStore(maxEventsPerStream: 5) + + // Store 5 events + var eventIds: [String] = [] + for i in 0..<5 { + let msg = #"{"id":"\#(i)"}"#.data(using: .utf8)! + let eventId = try await store.storeEvent(streamId: "stream", message: msg) + eventIds.append(eventId) + } + + // Store 2 more (evicting the first 2) + for i in 5..<7 { + let msg = #"{"id":"\#(i)"}"#.data(using: .utf8)! + let eventId = try await store.storeEvent(streamId: "stream", message: msg) + eventIds.append(eventId) + } + + // eventIds[2] (id: "2") should still be valid and allow replay of 3, 4, 5, 6 + actor MessageCollector { + var ids: [String] = [] + func add(_ id: String) { ids.append(id) } + func get() -> [String] { ids } + } + let collector = MessageCollector() + + _ = try await store.replayEventsAfter(eventIds[2]) { _, message in + if let json = try? JSONSerialization.jsonObject(with: message) as? [String: Any], + let id = json["id"] as? String + { + await collector.add(id) + } + } + + let replayedIds = await collector.get() + #expect(replayedIds == ["3", "4", "5", "6"]) + } + + @Test("Stream count tracks active streams") + func streamCountTracksActiveStreams() async throws { + let store = InMemoryEventStore() + let message = Data() + + var streamCount = await store.streamCount + #expect(streamCount == 0) + + _ = try await store.storeEvent(streamId: "stream-1", message: message) + streamCount = await store.streamCount + #expect(streamCount == 1) + + _ = try await store.storeEvent(streamId: "stream-2", message: message) + streamCount = await store.streamCount + #expect(streamCount == 2) + + // Adding to existing stream doesn't increase count + _ = try await store.storeEvent(streamId: "stream-1", message: message) + streamCount = await store.streamCount + #expect(streamCount == 2) + + // Removing stream reduces count + _ = await store.removeEvents(forStream: "stream-1") + streamCount = await store.streamCount + #expect(streamCount == 1) + } +} diff --git a/Tests/MCPTests/NotificationTests.swift b/Tests/MCPTests/NotificationTests.swift index 4bdd1873..68df08da 100644 --- a/Tests/MCPTests/NotificationTests.swift +++ b/Tests/MCPTests/NotificationTests.swift @@ -169,4 +169,243 @@ struct NotificationTests { #expect(decoded.method == ResourceUpdatedNotification.name) #expect(decoded.params.objectValue?["uri"]?.stringValue == "test://resource") } + + // MARK: - LogMessageNotification Tests + + @Test("LogMessageNotification encoding with all fields") + func testLogMessageNotificationEncodingAllFields() throws { + let params = LogMessageNotification.Parameters( + level: .info, + logger: "test-logger", + data: .string("Test log message") + ) + let notification = LogMessageNotification.message(params) + + #expect(notification.method == LogMessageNotification.name) + #expect(notification.params.level == .info) + #expect(notification.params.logger == "test-logger") + #expect(notification.params.data == .string("Test log message")) + + let encoder = JSONEncoder() + let data = try encoder.encode(notification) + + // Verify JSON structure + let json = try JSONDecoder().decode([String: Value].self, from: data) + #expect(json["jsonrpc"] == "2.0") + #expect(json["method"] == "notifications/message") + #expect(json["params"]?.objectValue?["level"] == "info") + #expect(json["params"]?.objectValue?["logger"] == "test-logger") + #expect(json["params"]?.objectValue?["data"] == "Test log message") + } + + @Test("LogMessageNotification encoding with minimal fields") + func testLogMessageNotificationEncodingMinimal() throws { + let params = LogMessageNotification.Parameters( + level: .warning, + data: .string("Warning message") + ) + let notification = LogMessageNotification.message(params) + + let encoder = JSONEncoder() + let data = try encoder.encode(notification) + + // Verify JSON structure (logger should be omitted) + let json = try JSONDecoder().decode([String: Value].self, from: data) + #expect(json["method"] == "notifications/message") + #expect(json["params"]?.objectValue?["level"] == "warning") + #expect(json["params"]?.objectValue?["logger"] == nil) + #expect(json["params"]?.objectValue?["data"] == "Warning message") + } + + @Test("LogMessageNotification decoding") + func testLogMessageNotificationDecoding() throws { + let jsonString = """ + {"jsonrpc":"2.0","method":"notifications/message","params":{"level":"error","logger":"app","data":"Error occurred"}} + """ + let data = jsonString.data(using: .utf8)! + + let decoder = JSONDecoder() + let decoded = try decoder.decode(Message.self, from: data) + + #expect(decoded.method == LogMessageNotification.name) + #expect(decoded.params.level == .error) + #expect(decoded.params.logger == "app") + #expect(decoded.params.data == .string("Error occurred")) + } + + @Test("LogMessageNotification with object data") + func testLogMessageNotificationWithObjectData() throws { + let params = LogMessageNotification.Parameters( + level: .debug, + data: .object(["key": .string("value"), "count": .int(42)]) + ) + let notification = LogMessageNotification.message(params) + + let encoder = JSONEncoder() + let data = try encoder.encode(notification) + let decoded = try JSONDecoder().decode(Message.self, from: data) + + #expect(decoded.params.level == .debug) + #expect(decoded.params.data.objectValue?["key"] == .string("value")) + #expect(decoded.params.data.objectValue?["count"] == .int(42)) + } + + @Test("LogMessageNotification all log levels") + func testLogMessageNotificationAllLogLevels() throws { + let levels: [LoggingLevel] = [ + .debug, .info, .notice, .warning, .error, .critical, .alert, .emergency + ] + + for level in levels { + let params = LogMessageNotification.Parameters(level: level, data: .string("test")) + let notification = LogMessageNotification.message(params) + + let encoder = JSONEncoder() + let data = try encoder.encode(notification) + let decoded = try JSONDecoder().decode(Message.self, from: data) + + #expect(decoded.params.level == level, "Log level \(level) should roundtrip correctly") + } + } + + // MARK: - ToolListChangedNotification Tests + + @Test("ToolListChangedNotification encoding") + func testToolListChangedNotificationEncoding() throws { + let notification = ToolListChangedNotification.message() + + #expect(notification.method == ToolListChangedNotification.name) + + let encoder = JSONEncoder() + let data = try encoder.encode(notification) + + // Verify JSON structure + let json = try JSONDecoder().decode([String: Value].self, from: data) + #expect(json["jsonrpc"] == "2.0") + #expect(json["method"] == "notifications/tools/list_changed") + // Empty params may be included as {} per JSON-RPC conventions + if let params = json["params"] { + #expect(params == .object([:]), "Params should be empty object if present") + } + } + + @Test("ToolListChangedNotification decoding") + func testToolListChangedNotificationDecoding() throws { + let jsonString = """ + {"jsonrpc":"2.0","method":"notifications/tools/list_changed"} + """ + let data = jsonString.data(using: .utf8)! + + let decoder = JSONDecoder() + let decoded = try decoder.decode(Message.self, from: data) + + #expect(decoded.method == ToolListChangedNotification.name) + } + + @Test("ToolListChangedNotification decoding with empty params") + func testToolListChangedNotificationDecodingWithEmptyParams() throws { + let jsonString = """ + {"jsonrpc":"2.0","method":"notifications/tools/list_changed","params":{}} + """ + let data = jsonString.data(using: .utf8)! + + let decoder = JSONDecoder() + let decoded = try decoder.decode(Message.self, from: data) + + #expect(decoded.method == ToolListChangedNotification.name) + } + + // MARK: - PromptListChangedNotification Tests + + @Test("PromptListChangedNotification encoding") + func testPromptListChangedNotificationEncoding() throws { + let notification = PromptListChangedNotification.message() + + #expect(notification.method == PromptListChangedNotification.name) + + let encoder = JSONEncoder() + let data = try encoder.encode(notification) + + // Verify JSON structure + let json = try JSONDecoder().decode([String: Value].self, from: data) + #expect(json["jsonrpc"] == "2.0") + #expect(json["method"] == "notifications/prompts/list_changed") + // Empty params may be included as {} per JSON-RPC conventions + if let params = json["params"] { + #expect(params == .object([:]), "Params should be empty object if present") + } + } + + @Test("PromptListChangedNotification decoding") + func testPromptListChangedNotificationDecoding() throws { + let jsonString = """ + {"jsonrpc":"2.0","method":"notifications/prompts/list_changed"} + """ + let data = jsonString.data(using: .utf8)! + + let decoder = JSONDecoder() + let decoded = try decoder.decode(Message.self, from: data) + + #expect(decoded.method == PromptListChangedNotification.name) + } + + @Test("PromptListChangedNotification decoding with empty params") + func testPromptListChangedNotificationDecodingWithEmptyParams() throws { + let jsonString = """ + {"jsonrpc":"2.0","method":"notifications/prompts/list_changed","params":{}} + """ + let data = jsonString.data(using: .utf8)! + + let decoder = JSONDecoder() + let decoded = try decoder.decode(Message.self, from: data) + + #expect(decoded.method == PromptListChangedNotification.name) + } + + // MARK: - ResourceListChangedNotification Tests + + @Test("ResourceListChangedNotification encoding") + func testResourceListChangedNotificationEncoding() throws { + let notification = ResourceListChangedNotification.message() + + #expect(notification.method == ResourceListChangedNotification.name) + + let encoder = JSONEncoder() + let data = try encoder.encode(notification) + + // Verify JSON structure + let json = try JSONDecoder().decode([String: Value].self, from: data) + #expect(json["jsonrpc"] == "2.0") + #expect(json["method"] == "notifications/resources/list_changed") + // Empty params may be included as {} per JSON-RPC conventions + if let params = json["params"] { + #expect(params == .object([:]), "Params should be empty object if present") + } + } + + @Test("ResourceListChangedNotification decoding") + func testResourceListChangedNotificationDecoding() throws { + let jsonString = """ + {"jsonrpc":"2.0","method":"notifications/resources/list_changed"} + """ + let data = jsonString.data(using: .utf8)! + + let decoder = JSONDecoder() + let decoded = try decoder.decode(Message.self, from: data) + + #expect(decoded.method == ResourceListChangedNotification.name) + } + + @Test("ResourceListChangedNotification decoding with empty params") + func testResourceListChangedNotificationDecodingWithEmptyParams() throws { + let jsonString = """ + {"jsonrpc":"2.0","method":"notifications/resources/list_changed","params":{}} + """ + let data = jsonString.data(using: .utf8)! + + let decoder = JSONDecoder() + let decoded = try decoder.decode(Message.self, from: data) + + #expect(decoded.method == ResourceListChangedNotification.name) + } } diff --git a/Tests/MCPTests/PrimingEventsTests.swift b/Tests/MCPTests/PrimingEventsTests.swift new file mode 100644 index 00000000..aeb6da25 --- /dev/null +++ b/Tests/MCPTests/PrimingEventsTests.swift @@ -0,0 +1,371 @@ +import Foundation +import Testing + +@testable import MCP + +/// Tests for SSE priming events on POST streams. +/// +/// Priming events are empty SSE events sent at the beginning of a POST SSE stream +/// to enable resumability. They contain an event ID (for resumption) and optionally +/// a retry field (for reconnection timing). +/// +/// These tests follow the TypeScript SDK patterns from: +/// - `packages/server/test/server/streamableHttp.test.ts` +/// +/// TypeScript tests not yet implemented (require protocol version 2025-11-25): +/// +/// Rationale: The MCP protocol requires priming events only for protocol version >= 2025-11-25. +/// The Swift SDK currently supports ['2024-11-05', '2025-03-26'], while TypeScript supports +/// ['2025-11-25', '2025-06-18', '2025-03-26', '2024-11-05', '2024-10-07']. +/// +/// The priming event code exists in HTTPServerTransport.writePrimingEvent() but +/// is gated by: `guard protocolVersion >= "2025-11-25" else { return }` +/// +/// Once the Swift SDK adds support for 2025-11-25 (see Versioning.swift TODO), implement: +/// - `should send priming event with retry field on POST SSE stream` +/// - `should send priming event without retry field when retryInterval is not configured` +/// - `should close POST SSE stream when extra.closeSSEStream is called` +/// - `should provide closeSSEStream callback in extra when eventStore is configured` +/// - `should NOT provide closeSSEStream callback for old protocol versions` +/// - `should NOT provide closeSSEStream callback when eventStore is NOT configured` +/// - `should provide closeStandaloneSSEStream callback in extra when eventStore is configured` +/// - `should close standalone GET SSE stream when extra.closeStandaloneSSEStream is called` +/// - `should allow client to reconnect after standalone SSE stream is closed` +/// +/// The current tests verify that priming events are NOT sent for the currently supported +/// protocol versions, which is the correct behavior for backwards compatibility. +@Suite("Priming Events Tests") +struct PrimingEventsTests { + + // MARK: - Test Helpers + + /// Helper to read from stream with timeout + func readFromStream( + _ stream: AsyncThrowingStream, + maxChunks: Int = 1, + timeout: Duration = .seconds(2) + ) async throws -> Data { + var receivedData = Data() + + try await withThrowingTaskGroup(of: Data?.self) { group in + group.addTask { + var data = Data() + var count = 0 + for try await chunk in stream { + data.append(chunk) + count += 1 + if count >= maxChunks { + break + } + } + return data + } + + group.addTask { + try await Task.sleep(for: timeout) + return nil + } + + if let result = try await group.next(), let data = result { + receivedData = data + } + group.cancelAll() + } + + return receivedData + } + + /// Creates a configured MCP Server with tools for testing + func createTestServer() -> Server { + let server = Server( + name: "test-server", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + return server + } + + /// Sets up tool handlers on the server + func setUpToolHandlers(_ server: Server) async { + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool( + name: "greet", + description: "A simple greeting tool", + inputSchema: [ + "type": "object", + "properties": ["name": ["type": "string"]] + ] + ) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, _ in + switch request.name { + case "greet": + let name = request.arguments?["name"]?.stringValue ?? "World" + return CallTool.Result(content: [.text("Hello, \(name)!")]) + default: + return CallTool.Result(content: [.text("Unknown tool")], isError: true) + } + } + } + + // MARK: - 4.1 Priming event configuration + // + // Note: The Swift SDK currently requires protocol version >= "2025-11-25" for priming events, + // but the supported versions are "2024-11-05" and "2025-03-26". This means priming events + // won't be sent until a newer protocol version is added. These tests verify the current behavior. + + @Test("Priming event configuration - retryInterval can be configured") + func primingEventRetryIntervalConfig() async throws { + // Test that retryInterval can be set in transport options + let options1 = HTTPServerTransportOptions(retryInterval: 5000) + #expect(options1.retryInterval == 5000) + + let options2 = HTTPServerTransportOptions() + #expect(options2.retryInterval == nil) + } + + @Test("Priming events not sent for current supported protocol versions") + func primingEventsNotSentForCurrentVersions() async throws { + let server = createTestServer() + await setUpToolHandlers(server) + + let eventStore = InMemoryEventStore() + let sessionId = UUID().uuidString + + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { sessionId }, + eventStore: eventStore, + retryInterval: 5000 + ) + ) + try await server.start(transport: transport) + + // Initialize with latest supported version (2025-03-26) + // Note: Priming events require >= 2025-11-25 which is not yet supported + let initRequest = TestPayloads.initializeRequest(protocolVersion: Version.v2025_03_26) + let initResponse = await transport.handleRequest(TestPayloads.postRequest(body: initRequest, protocolVersion: Version.v2025_03_26)) + #expect(initResponse.statusCode == 200) + + // Send a tool call request + let toolCallRequest = """ + {"jsonrpc":"2.0","method":"tools/call","id":"100","params":{"name":"greet","arguments":{"name":"Test"}}} + """ + let response = await transport.handleRequest(TestPayloads.postRequest(body: toolCallRequest, sessionId: sessionId, protocolVersion: Version.v2025_03_26)) + + #expect(response.statusCode == 200) + + if let stream = response.stream { + let data = try await readFromStream(stream, maxChunks: 2) + let text = String(data: data, encoding: .utf8) ?? "" + + // Priming events have empty data - current versions won't have them + #expect(!text.contains("data: \n\n"), "Should NOT have empty priming event for current protocol versions") + #expect(text.contains("Hello, Test!") || text.contains("result"), "Should contain tool result") + } else if let body = response.body { + let text = String(data: body, encoding: .utf8) ?? "" + #expect(text.contains("Hello, Test!") || text.contains("result"), "Should contain tool result") + } + } + + // MARK: - 4.2 Event ID on messages (even without priming events) + + @Test("Event IDs are included in SSE messages when event store is configured") + func eventIdsIncludedInMessages() async throws { + let server = createTestServer() + await setUpToolHandlers(server) + + let eventStore = InMemoryEventStore() + let sessionId = UUID().uuidString + + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { sessionId }, + eventStore: eventStore + ) + ) + try await server.start(transport: transport) + + // Initialize + let initRequest = TestPayloads.initializeRequest() + _ = await transport.handleRequest(TestPayloads.postRequest(body: initRequest)) + + // Send a tool call request + let toolCallRequest = """ + {"jsonrpc":"2.0","method":"tools/call","id":"100","params":{"name":"greet","arguments":{"name":"Test"}}} + """ + let response = await transport.handleRequest(TestPayloads.postRequest(body: toolCallRequest, sessionId: sessionId)) + + #expect(response.statusCode == 200) + + if let stream = response.stream { + let data = try await readFromStream(stream, maxChunks: 2) + let text = String(data: data, encoding: .utf8) ?? "" + + // Even without priming events, messages should have event IDs for resumability + #expect(text.contains("id: "), "SSE messages should include event IDs for resumability") + #expect(text.contains("Hello, Test!") || text.contains("result"), "Should contain tool result") + + // Verify events were stored + let eventCount = await eventStore.eventCount + #expect(eventCount > 0, "Events should be stored in event store") + } + } + + // MARK: - Priming Event Content Tests + + @Test("No priming event for old protocol versions (backwards compatibility)") + func noPrimingEventForOldProtocolVersions() async throws { + let server = createTestServer() + await setUpToolHandlers(server) + + let eventStore = InMemoryEventStore() + let sessionId = UUID().uuidString + + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { sessionId }, + eventStore: eventStore, + retryInterval: 5000 + ) + ) + try await server.start(transport: transport) + + // Initialize with OLD protocol version (< 2025-11-25) + let initRequest = TestPayloads.initializeRequest() + _ = await transport.handleRequest(TestPayloads.postRequest(body: initRequest, protocolVersion: Version.v2024_11_05)) + + // Send a tool call request + let toolCallRequest = """ + {"jsonrpc":"2.0","method":"tools/call","id":"100","params":{"name":"greet","arguments":{"name":"Test"}}} + """ + let response = await transport.handleRequest(TestPayloads.postRequest(body: toolCallRequest, sessionId: sessionId, protocolVersion: Version.v2024_11_05)) + + #expect(response.statusCode == 200) + + if let stream = response.stream { + let data = try await readFromStream(stream, maxChunks: 2) + let text = String(data: data, encoding: .utf8) ?? "" + + // Should NOT have retry field for old protocol versions + // Priming events are not sent for backwards compatibility + #expect(!text.contains("data: \n\n"), "Should NOT have empty priming event data for old protocol versions") + #expect(text.contains("Hello, Test!") || text.contains("result"), "Should contain actual tool result") + } else if let body = response.body { + let text = String(data: body, encoding: .utf8) ?? "" + #expect(text.contains("Hello, Test!") || text.contains("result"), "Should contain tool result") + } + } + + @Test("No priming event when event store is not configured") + func noPrimingEventWithoutEventStore() async throws { + let server = createTestServer() + await setUpToolHandlers(server) + + let sessionId = UUID().uuidString + + // No event store configured + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { sessionId } + // No eventStore + ) + ) + try await server.start(transport: transport) + + // Initialize + let initRequest = TestPayloads.initializeRequest() + _ = await transport.handleRequest(TestPayloads.postRequest(body: initRequest)) + + // Send a tool call request + let toolCallRequest = """ + {"jsonrpc":"2.0","method":"tools/call","id":"100","params":{"name":"greet","arguments":{"name":"Test"}}} + """ + let response = await transport.handleRequest(TestPayloads.postRequest(body: toolCallRequest, sessionId: sessionId)) + + #expect(response.statusCode == 200) + + // Without event store, the response might be JSON directly or SSE without priming + // Either way, the actual response should be there + if let body = response.body { + let text = String(data: body, encoding: .utf8) ?? "" + #expect(text.contains("Hello, Test!") || text.contains("result"), "Should contain tool result") + } else if let stream = response.stream { + let data = try await readFromStream(stream, maxChunks: 2) + let text = String(data: data, encoding: .utf8) ?? "" + // Without event store, there should be no event ID in the stream + // Note: The first message could still have an id field depending on implementation + #expect(text.contains("Hello, Test!") || text.contains("result"), "Should contain tool result") + } + } + + // MARK: - Close SSE Stream Tests + + @Test("Close SSE stream for specific request") + func closeSSEStreamForSpecificRequest() async throws { + let server = createTestServer() + + let eventStore = InMemoryEventStore() + let sessionId = UUID().uuidString + + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { sessionId }, + eventStore: eventStore, + retryInterval: 1000 + ) + ) + + // Set up a tool that takes time, allowing us to close the stream mid-execution + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool( + name: "slow-tool", + description: "A slow tool", + inputSchema: ["type": "object"] + ) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, context in + if request.name == "slow-tool" { + // Simulate slow operation + try? await Task.sleep(for: .milliseconds(500)) + return CallTool.Result(content: [.text("Done")]) + } + return CallTool.Result(content: [.text("Unknown tool")], isError: true) + } + + try await server.start(transport: transport) + + // Initialize + let initRequest = TestPayloads.initializeRequest() + _ = await transport.handleRequest(TestPayloads.postRequest(body: initRequest)) + + // The closeSSEStream method exists and is callable + // We can't fully test the stream closure without complex async coordination + // but we can verify the method exists and doesn't crash + let requestId: RequestId = .string("test-request") + await transport.closeSSEStream(for: requestId) + + // If we get here without crashing, the method works + } + + // MARK: - Retry Interval Configuration Tests + + @Test("Retry interval is configurable in transport options") + func retryIntervalIsConfigurable() async throws { + // Test that different retry interval configurations can be set + let options1 = HTTPServerTransportOptions(retryInterval: 1000) + #expect(options1.retryInterval == 1000) + + let options2 = HTTPServerTransportOptions(retryInterval: 30000) + #expect(options2.retryInterval == 30000) + + let options3 = HTTPServerTransportOptions() + #expect(options3.retryInterval == nil) + } +} diff --git a/Tests/MCPTests/ProgressTests.swift b/Tests/MCPTests/ProgressTests.swift new file mode 100644 index 00000000..ed7f1327 --- /dev/null +++ b/Tests/MCPTests/ProgressTests.swift @@ -0,0 +1,2737 @@ +import Foundation +import Logging +import Testing + +#if canImport(System) + import System +#else + @preconcurrency import SystemPackage +#endif + +@testable import MCP + +@Suite("Progress Tests") +struct ProgressTests { + + // MARK: - ProgressToken Tests + + @Suite("ProgressToken encoding/decoding") + struct ProgressTokenTests { + + @Test("String token encodes as JSON string") + func stringTokenEncoding() throws { + let token: ProgressToken = .string("abc-123") + let encoder = JSONEncoder() + let data = try encoder.encode(token) + let json = String(data: data, encoding: .utf8) + #expect(json == "\"abc-123\"") + } + + @Test("Integer token encodes as JSON number") + func integerTokenEncoding() throws { + let token: ProgressToken = .integer(42) + let encoder = JSONEncoder() + let data = try encoder.encode(token) + let json = String(data: data, encoding: .utf8) + #expect(json == "42") + } + + @Test("String token decodes from JSON string") + func stringTokenDecoding() throws { + let json = "\"my-token\"" + let data = json.data(using: .utf8)! + let decoder = JSONDecoder() + let token = try decoder.decode(ProgressToken.self, from: data) + #expect(token == .string("my-token")) + } + + @Test("Integer token decodes from JSON number") + func integerTokenDecoding() throws { + let json = "123" + let data = json.data(using: .utf8)! + let decoder = JSONDecoder() + let token = try decoder.decode(ProgressToken.self, from: data) + #expect(token == .integer(123)) + } + + @Test("Integer token zero decodes correctly") + func integerTokenZero() throws { + // Edge case from Python SDK test #176 - progress token 0 should work + let json = "0" + let data = json.data(using: .utf8)! + let decoder = JSONDecoder() + let token = try decoder.decode(ProgressToken.self, from: data) + #expect(token == .integer(0)) + } + + @Test("Negative integer token decodes correctly") + func negativeIntegerToken() throws { + let json = "-1" + let data = json.data(using: .utf8)! + let decoder = JSONDecoder() + let token = try decoder.decode(ProgressToken.self, from: data) + #expect(token == .integer(-1)) + } + + @Test("String literal initialization") + func stringLiteralInit() { + let token: ProgressToken = "my-token" + #expect(token == .string("my-token")) + } + + @Test("Integer literal initialization") + func integerLiteralInit() { + let token: ProgressToken = 42 + #expect(token == .integer(42)) + } + + @Test("Round-trip encoding/decoding for string token") + func stringTokenRoundTrip() throws { + let original: ProgressToken = .string("test-token-abc") + let encoder = JSONEncoder() + let decoder = JSONDecoder() + let data = try encoder.encode(original) + let decoded = try decoder.decode(ProgressToken.self, from: data) + #expect(decoded == original) + } + + @Test("Round-trip encoding/decoding for integer token") + func integerTokenRoundTrip() throws { + let original: ProgressToken = .integer(999) + let encoder = JSONEncoder() + let decoder = JSONDecoder() + let data = try encoder.encode(original) + let decoded = try decoder.decode(ProgressToken.self, from: data) + #expect(decoded == original) + } + + @Test("Invalid token type throws error") + func invalidTokenType() throws { + let json = "true" // Boolean is not a valid progress token + let data = json.data(using: .utf8)! + let decoder = JSONDecoder() + #expect(throws: DecodingError.self) { + _ = try decoder.decode(ProgressToken.self, from: data) + } + } + } + + // MARK: - ProgressNotification Tests + + @Suite("ProgressNotification encoding/decoding") + struct ProgressNotificationTests { + + @Test("Notification with string token encodes correctly") + func notificationWithStringToken() throws { + let params = ProgressNotification.Parameters( + progressToken: .string("abc-123"), + progress: 50.0, + total: 100.0, + message: "Halfway done" + ) + let notification = ProgressNotification.message(params) + + let encoder = JSONEncoder() + let data = try encoder.encode(notification) + let json = try JSONDecoder().decode([String: Value].self, from: data) + + #expect(json["jsonrpc"] == "2.0") + #expect(json["method"] == "notifications/progress") + + let notificationParams = json["params"]?.objectValue + #expect(notificationParams?["progressToken"]?.stringValue == "abc-123") + // Use Double(_ value:) which handles both .int and .double cases + #expect(notificationParams?["progress"].flatMap { Double($0) } == 50.0) + #expect(notificationParams?["total"].flatMap { Double($0) } == 100.0) + #expect(notificationParams?["message"]?.stringValue == "Halfway done") + } + + @Test("Notification with integer token encodes correctly") + func notificationWithIntegerToken() throws { + let params = ProgressNotification.Parameters( + progressToken: .integer(42), + progress: 25.0, + total: 100.0, + message: "Quarter done" + ) + let notification = ProgressNotification.message(params) + + let encoder = JSONEncoder() + let data = try encoder.encode(notification) + let json = try JSONDecoder().decode([String: Value].self, from: data) + + let notificationParams = json["params"]?.objectValue + #expect(notificationParams?["progressToken"]?.intValue == 42) + #expect(notificationParams?["progress"].flatMap { Double($0) } == 25.0) + } + + @Test("Notification with zero token encodes correctly") + func notificationWithZeroToken() throws { + // Edge case from Python SDK test #176 + let params = ProgressNotification.Parameters( + progressToken: .integer(0), + progress: 0.0, + total: 10.0, + message: nil + ) + let notification = ProgressNotification.message(params) + + let encoder = JSONEncoder() + let data = try encoder.encode(notification) + let json = try JSONDecoder().decode([String: Value].self, from: data) + + let notificationParams = json["params"]?.objectValue + #expect(notificationParams?["progressToken"]?.intValue == 0) + #expect(notificationParams?["progress"].flatMap { Double($0) } == 0.0) + } + + @Test("Notification decodes from JSON with string token") + func decodeNotificationWithStringToken() throws { + let json = """ + { + "jsonrpc": "2.0", + "method": "notifications/progress", + "params": { + "progressToken": "token-abc", + "progress": 75.0, + "total": 100.0, + "message": "Almost done" + } + } + """ + let data = json.data(using: .utf8)! + let decoder = JSONDecoder() + let notification = try decoder.decode(Message.self, from: data) + + #expect(notification.method == "notifications/progress") + #expect(notification.params.progressToken == .string("token-abc")) + #expect(notification.params.progress == 75.0) + #expect(notification.params.total == 100.0) + #expect(notification.params.message == "Almost done") + } + + @Test("Notification decodes from JSON with integer token") + func decodeNotificationWithIntegerToken() throws { + let json = """ + { + "jsonrpc": "2.0", + "method": "notifications/progress", + "params": { + "progressToken": 123, + "progress": 50.0, + "total": 200.0 + } + } + """ + let data = json.data(using: .utf8)! + let decoder = JSONDecoder() + let notification = try decoder.decode(Message.self, from: data) + + #expect(notification.params.progressToken == .integer(123)) + #expect(notification.params.progress == 50.0) + #expect(notification.params.total == 200.0) + #expect(notification.params.message == nil) + } + + @Test("Notification without optional fields decodes correctly") + func decodeNotificationMinimalFields() throws { + let json = """ + { + "jsonrpc": "2.0", + "method": "notifications/progress", + "params": { + "progressToken": "min-token", + "progress": 10.0 + } + } + """ + let data = json.data(using: .utf8)! + let decoder = JSONDecoder() + let notification = try decoder.decode(Message.self, from: data) + + #expect(notification.params.progressToken == .string("min-token")) + #expect(notification.params.progress == 10.0) + #expect(notification.params.total == nil) + #expect(notification.params.message == nil) + } + + @Test("Notification with _meta field decodes correctly") + func decodeNotificationWithMeta() throws { + let json = """ + { + "jsonrpc": "2.0", + "method": "notifications/progress", + "params": { + "progressToken": "meta-token", + "progress": 50.0, + "_meta": { + "customField": "customValue" + } + } + } + """ + let data = json.data(using: .utf8)! + let decoder = JSONDecoder() + let notification = try decoder.decode(Message.self, from: data) + + #expect(notification.params.progressToken == .string("meta-token")) + #expect(notification.params._meta?["customField"]?.stringValue == "customValue") + } + + @Test("Round-trip encoding/decoding for notification") + func notificationRoundTrip() throws { + let params = ProgressNotification.Parameters( + progressToken: .string("round-trip-token"), + progress: 33.3, + total: 100.0, + message: "Processing..." + ) + let original = ProgressNotification.message(params) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + let data = try encoder.encode(original) + let decoded = try decoder.decode(Message.self, from: data) + + #expect(decoded.method == original.method) + #expect(decoded.params.progressToken == original.params.progressToken) + #expect(decoded.params.progress == original.params.progress) + #expect(decoded.params.total == original.params.total) + #expect(decoded.params.message == original.params.message) + } + } + + // MARK: - RequestMeta Tests + + @Suite("RequestMeta encoding/decoding") + struct RequestMetaTests { + + @Test("RequestMeta with progressToken encodes correctly") + func requestMetaWithProgressToken() throws { + let meta = RequestMeta(progressToken: .string("request-token")) + let encoder = JSONEncoder() + let data = try encoder.encode(meta) + let json = try JSONDecoder().decode([String: Value].self, from: data) + + #expect(json["progressToken"]?.stringValue == "request-token") + } + + @Test("RequestMeta with integer progressToken encodes correctly") + func requestMetaWithIntegerToken() throws { + let meta = RequestMeta(progressToken: .integer(42)) + let encoder = JSONEncoder() + let data = try encoder.encode(meta) + let json = try JSONDecoder().decode([String: Value].self, from: data) + + #expect(json["progressToken"]?.intValue == 42) + } + + @Test("RequestMeta with additional fields encodes correctly") + func requestMetaWithAdditionalFields() throws { + let meta = RequestMeta( + progressToken: .string("token"), + additionalFields: [ + "customField": .string("customValue"), + "numericField": .int(123) + ] + ) + let encoder = JSONEncoder() + let data = try encoder.encode(meta) + let json = try JSONDecoder().decode([String: Value].self, from: data) + + #expect(json["progressToken"]?.stringValue == "token") + #expect(json["customField"]?.stringValue == "customValue") + #expect(json["numericField"]?.intValue == 123) + } + + @Test("RequestMeta decodes from JSON") + func requestMetaDecoding() throws { + let json = """ + { + "progressToken": "decoded-token", + "extraField": "extraValue" + } + """ + let data = json.data(using: .utf8)! + let decoder = JSONDecoder() + let meta = try decoder.decode(RequestMeta.self, from: data) + + #expect(meta.progressToken == .string("decoded-token")) + #expect(meta.additionalFields?["extraField"]?.stringValue == "extraValue") + } + + @Test("RequestMeta decodes integer progressToken from JSON") + func requestMetaIntegerTokenDecoding() throws { + let json = """ + { + "progressToken": 999 + } + """ + let data = json.data(using: .utf8)! + let decoder = JSONDecoder() + let meta = try decoder.decode(RequestMeta.self, from: data) + + #expect(meta.progressToken == .integer(999)) + } + + @Test("RequestMeta without progressToken decodes correctly") + func requestMetaWithoutToken() throws { + let json = """ + { + "customField": "value" + } + """ + let data = json.data(using: .utf8)! + let decoder = JSONDecoder() + let meta = try decoder.decode(RequestMeta.self, from: data) + + #expect(meta.progressToken == nil) + #expect(meta.additionalFields?["customField"]?.stringValue == "value") + } + + @Test("Round-trip encoding/decoding for RequestMeta") + func requestMetaRoundTrip() throws { + let original = RequestMeta( + progressToken: .string("round-trip"), + additionalFields: ["key": .string("value")] + ) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + let data = try encoder.encode(original) + let decoded = try decoder.decode(RequestMeta.self, from: data) + + #expect(decoded.progressToken == original.progressToken) + #expect(decoded.additionalFields?["key"] == original.additionalFields?["key"]) + } + } + + // MARK: - Integration Tests + + /// Integration tests for progress notifications through actual client/server communication. + /// Based on Python SDK's test_progress_notifications.py tests. + @Suite("Progress notification integration") + struct ProgressIntegrationTests { + + /// Test that server can send progress notifications to client during tool execution. + /// Based on TypeScript SDK's "should send progress notifications with message field" test. + /// + /// Flow matches TS/Python pattern: + /// 1. Client sends request WITH progressToken in _meta + /// 2. Server extracts token from request._meta.progressToken + /// 3. Server sends notifications using that token + /// 4. Client receives and correlates by token + @Test(.timeLimit(.minutes(1))) + func serverSendsProgressNotificationsToClient() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.progress") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + // Track received progress updates + let receivedProgress = ProgressUpdateTracker() + + // Set up server with a tool that sends progress notifications + let server = Server( + name: "ProgressTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool( + name: "slow_operation", + description: "A tool that reports progress", + inputSchema: ["type": "object", "properties": ["steps": ["type": "integer"]]] + ) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, context in + guard request.name == "slow_operation" else { + return CallTool.Result(content: [.text("Unknown tool")], isError: true) + } + + // Extract progress token from request _meta (matching TS/Python pattern) + guard let progressToken = request._meta?.progressToken else { + return CallTool.Result(content: [.text("No progress token provided")], isError: true) + } + + let steps = request.arguments?["steps"]?.intValue ?? 3 + + // Send progress notifications for each step (like TS SDK test) + for step in 1...steps { + try await context.sendProgress( + token: progressToken, + progress: Double(step), + total: Double(steps), + message: "Completed step \(step) of \(steps)" + ) + } + + return CallTool.Result(content: [.text("Operation completed with \(steps) steps")]) + } + + let client = Client(name: "ProgressTestClient", version: "1.0") + + // Register progress notification handler + await client.onNotification(ProgressNotification.self) { message in + await receivedProgress.add( + token: message.params.progressToken, + progress: message.params.progress, + total: message.params.total, + message: message.params.message + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Call the tool WITH progressToken in _meta (matching TS/Python pattern) + let result = try await client.send( + CallTool.request(.init( + name: "slow_operation", + arguments: ["steps": .int(3)], + _meta: RequestMeta(progressToken: .string("progress-test-1")) + )) + ) + + // Give time for notifications to be processed + try await Task.sleep(for: .milliseconds(100)) + + // Verify tool result + if case .text(let text, _, _) = result.content.first { + #expect(text == "Operation completed with 3 steps") + } else { + Issue.record("Expected text content") + } + + // Verify progress notifications were received (matching TS SDK assertions) + let updates = await receivedProgress.updates + #expect(updates.count == 3, "Should receive 3 progress notifications") + + if updates.count >= 3 { + // Verify each notification has the correct token from the request + for update in updates { + #expect(update.token == .string("progress-test-1"), "Token should match request") + } + + // Verify progress values match TS SDK test pattern + #expect(updates[0].progress == 1.0) + #expect(updates[0].total == 3.0) + #expect(updates[0].message == "Completed step 1 of 3") + + #expect(updates[1].progress == 2.0) + #expect(updates[1].total == 3.0) + #expect(updates[1].message == "Completed step 2 of 3") + + #expect(updates[2].progress == 3.0) + #expect(updates[2].total == 3.0) + #expect(updates[2].message == "Completed step 3 of 3") + } + } + + /// Test that progress token 0 works correctly in actual communication. + /// Based on Python SDK's test_176_progress_token.py (issue #176 - falsy token value). + /// + /// This tests the edge case where progressToken is 0 (a falsy value in many languages). + /// The token must flow correctly: client → server → notification → client. + @Test(.timeLimit(.minutes(1))) + func progressTokenZeroWorks() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.progress.zero") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let receivedProgress = ProgressUpdateTracker() + + let server = Server( + name: "ZeroTokenServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "zero_token_tool", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, context in + guard request.name == "zero_token_tool" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + // Extract token from request (should be integer 0) + guard let progressToken = request._meta?.progressToken else { + return CallTool.Result(content: [.text("No progress token")], isError: true) + } + + // The key test: token 0 should NOT be treated as "no token" + // This was bug #176 in Python SDK + try await context.sendProgress(token: progressToken, progress: 0.0, total: 10.0) + try await context.sendProgress(token: progressToken, progress: 5.0, total: 10.0) + try await context.sendProgress(token: progressToken, progress: 10.0, total: 10.0) + + return CallTool.Result(content: [.text("Done")]) + } + + let client = Client(name: "ZeroTokenClient", version: "1.0") + + await client.onNotification(ProgressNotification.self) { message in + await receivedProgress.add( + token: message.params.progressToken, + progress: message.params.progress, + total: message.params.total, + message: message.params.message + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Client sends request with progressToken = 0 (the edge case) + _ = try await client.send( + CallTool.request(.init( + name: "zero_token_tool", + arguments: [:], + _meta: RequestMeta(progressToken: .integer(0)) + )) + ) + + try await Task.sleep(for: .milliseconds(100)) + + let updates = await receivedProgress.updates + #expect(updates.count == 3, "Should receive all 3 progress notifications with token 0") + + // Verify token 0 was correctly transmitted through the entire flow + for update in updates { + #expect(update.token == .integer(0), "Token should be integer 0") + } + + if updates.count >= 3 { + #expect(updates[0].progress == 0.0) + #expect(updates[1].progress == 5.0) + #expect(updates[2].progress == 10.0) + } + } + + /// Test that server correctly extracts progressToken from request _meta. + /// This matches the typical flow where client includes progressToken in _meta. + @Test(.timeLimit(.minutes(1))) + func serverExtractsProgressTokenFromRequestMeta() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.progress.meta") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let receivedProgress = ProgressUpdateTracker() + + let server = Server( + name: "MetaExtractServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "meta_test", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, context in + guard request.name == "meta_test" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + // Extract progress token from request _meta (the recommended pattern) + if let token = request._meta?.progressToken { + try await context.sendProgress( + token: token, + progress: 50.0, + total: 100.0, + message: "Using token from _meta" + ) + } + + return CallTool.Result(content: [.text("Done")]) + } + + let client = Client(name: "MetaExtractClient", version: "1.0") + + await client.onNotification(ProgressNotification.self) { message in + await receivedProgress.add( + token: message.params.progressToken, + progress: message.params.progress, + total: message.params.total, + message: message.params.message + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Call tool WITH progressToken in _meta + let result = try await client.send( + CallTool.request(.init( + name: "meta_test", + arguments: [:], + _meta: RequestMeta(progressToken: .string("client-provided-token")) + )) + ) + + try await Task.sleep(for: .milliseconds(100)) + + // Verify the tool result succeeded + #expect(result.content.count == 1) + + // Verify progress notification was received with the client-provided token + let updates = await receivedProgress.updates + #expect(updates.count == 1, "Should receive 1 progress notification") + + if let update = updates.first { + #expect(update.token == .string("client-provided-token")) + #expect(update.message == "Using token from _meta") + } + } + + /// Test with integer progress token in full client-server roundtrip. + /// Ensures integer tokens work end-to-end, not just in serialization. + @Test(.timeLimit(.minutes(1))) + func integerTokenRoundtripIntegration() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.progress.int") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let receivedProgress = ProgressUpdateTracker() + + let server = Server( + name: "IntTokenServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "int_token_test", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, context in + guard request.name == "int_token_test" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + // Extract integer token from request _meta + if let token = request._meta?.progressToken { + try await context.sendProgress( + token: token, + progress: 100.0, + total: 100.0, + message: "Complete" + ) + } + + return CallTool.Result(content: [.text("Done")]) + } + + let client = Client(name: "IntTokenClient", version: "1.0") + + await client.onNotification(ProgressNotification.self) { message in + await receivedProgress.add( + token: message.params.progressToken, + progress: message.params.progress, + total: message.params.total, + message: message.params.message + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Call tool with INTEGER progressToken in _meta + _ = try await client.send( + CallTool.request(.init( + name: "int_token_test", + arguments: [:], + _meta: RequestMeta(progressToken: .integer(12345)) + )) + ) + + try await Task.sleep(for: .milliseconds(100)) + + let updates = await receivedProgress.updates + #expect(updates.count == 1, "Should receive 1 progress notification") + + if let update = updates.first { + // Verify integer token was preserved through the roundtrip + #expect(update.token == .integer(12345), "Integer token should be preserved") + } + } + + /// Test that sendMessage can be used to send notifications with custom parameters. + @Test(.timeLimit(.minutes(1))) + func sendMessageWorksForCustomNotifications() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.progress.sendMessage") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let receivedProgress = ProgressUpdateTracker() + + let server = Server( + name: "SendMessageServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "message_test", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, context in + guard request.name == "message_test" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + // Use sendMessage directly instead of convenience method + try await context.sendMessage(ProgressNotification.message(.init( + progressToken: .string("via-sendMessage"), + progress: 42.0, + total: 100.0, + message: "Sent via sendMessage" + ))) + + return CallTool.Result(content: [.text("Done")]) + } + + let client = Client(name: "SendMessageClient", version: "1.0") + + await client.onNotification(ProgressNotification.self) { message in + await receivedProgress.add( + token: message.params.progressToken, + progress: message.params.progress, + total: message.params.total, + message: message.params.message + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + _ = try await client.callTool(name: "message_test", arguments: [:]) + + try await Task.sleep(for: .milliseconds(100)) + + let updates = await receivedProgress.updates + #expect(updates.count == 1, "Should receive 1 progress notification") + + if let update = updates.first { + #expect(update.token == .string("via-sendMessage")) + #expect(update.progress == 42.0) + #expect(update.message == "Sent via sendMessage") + } + } + + /// Test sendLogMessage convenience method. + @Test(.timeLimit(.minutes(1))) + func sendLogMessageWorks() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.logging") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let receivedLogs = LogTracker() + + let server = Server( + name: "LogTestServer", + version: "1.0.0", + capabilities: .init(logging: .init(), tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "log_test", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, context in + guard request.name == "log_test" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + // Send log messages using convenience method + try await context.sendLogMessage( + level: .info, + logger: "test-logger", + data: .string("Starting operation") + ) + + try await context.sendLogMessage( + level: .warning, + data: .object(["status": .string("in-progress"), "step": .int(1)]) + ) + + return CallTool.Result(content: [.text("Done")]) + } + + let client = Client(name: "LogTestClient", version: "1.0") + + await client.onNotification(LogMessageNotification.self) { message in + await receivedLogs.add( + level: message.params.level, + logger: message.params.logger, + data: message.params.data + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + _ = try await client.callTool(name: "log_test", arguments: [:]) + + try await Task.sleep(for: .milliseconds(100)) + + let logs = await receivedLogs.logs + #expect(logs.count == 2, "Should receive 2 log notifications") + + if logs.count >= 2 { + #expect(logs[0].level == .info) + #expect(logs[0].logger == "test-logger") + #expect(logs[0].data.stringValue == "Starting operation") + + #expect(logs[1].level == .warning) + #expect(logs[1].logger == nil) + #expect(logs[1].data.objectValue?["status"]?.stringValue == "in-progress") + } + } + + /// Test that log level filtering works correctly. + /// + /// When the client sets a minimum log level, the server should only + /// send log messages at that level or higher (more severe). + @Test(.timeLimit(.minutes(1))) + func logLevelFilteringWorks() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.loglevel") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let receivedLogs = LogTracker() + + let server = Server( + name: "LogLevelTestServer", + version: "1.0.0", + capabilities: .init(logging: .init(), tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "log_all_levels", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, context in + guard request.name == "log_all_levels" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + // Send log messages at all levels + try await context.sendLogMessage(level: .debug, data: .string("Debug message")) + try await context.sendLogMessage(level: .info, data: .string("Info message")) + try await context.sendLogMessage(level: .notice, data: .string("Notice message")) + try await context.sendLogMessage(level: .warning, data: .string("Warning message")) + try await context.sendLogMessage(level: .error, data: .string("Error message")) + try await context.sendLogMessage(level: .critical, data: .string("Critical message")) + + return CallTool.Result(content: [.text("Done")]) + } + + let client = Client(name: "LogLevelTestClient", version: "1.0") + + await client.onNotification(LogMessageNotification.self) { message in + await receivedLogs.add( + level: message.params.level, + logger: message.params.logger, + data: message.params.data + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Set minimum log level to warning - only warning and above should be received + try await client.setLoggingLevel(.warning) + + // Give time for the level to be set + try await Task.sleep(for: .milliseconds(50)) + + // Call the tool that sends logs at all levels + _ = try await client.callTool(name: "log_all_levels", arguments: [:]) + + try await Task.sleep(for: .milliseconds(100)) + + let logs = await receivedLogs.logs + // Should only receive warning, error, critical (3 messages) + // debug, info, notice should be filtered out + #expect(logs.count == 3, "Should receive only 3 log notifications (warning and above), got \(logs.count)") + + if logs.count >= 3 { + #expect(logs[0].level == .warning) + #expect(logs[0].data.stringValue == "Warning message") + + #expect(logs[1].level == .error) + #expect(logs[1].data.stringValue == "Error message") + + #expect(logs[2].level == .critical) + #expect(logs[2].data.stringValue == "Critical message") + } + } + + /// Test that setLoggingLevel throws when server doesn't have logging capability. + /// + /// This matches TypeScript SDK behavior where `client.setLoggingLevel('error')` + /// throws "Server does not support logging" when capability is not declared. + @Test(.timeLimit(.minutes(1))) + func setLoggingLevelThrowsWithoutCapability() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.logging.nocap") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + // Server WITHOUT logging capability + let server = Server( + name: "NoLoggingServer", + version: "1.0.0", + capabilities: .init(tools: .init()) // No logging capability + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: []) + } + + let client = Client(name: "LogTestClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Should throw because server doesn't support logging + await #expect(throws: MCPError.self) { + try await client.setLoggingLevel(.warning) + } + } + + /// Test that all 8 RFC 5424 log levels work correctly. + /// + /// The MCP spec uses syslog severity levels: + /// debug < info < notice < warning < error < critical < alert < emergency + @Test(.timeLimit(.minutes(1))) + func allEightLogLevelsWork() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.alllevels") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let receivedLogs = LogTracker() + + let server = Server( + name: "AllLevelsServer", + version: "1.0.0", + capabilities: .init(logging: .init(), tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "log_all_eight", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, context in + guard request.name == "log_all_eight" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + // Send all 8 log levels + try await context.sendLogMessage(level: .debug, data: .string("Debug")) + try await context.sendLogMessage(level: .info, data: .string("Info")) + try await context.sendLogMessage(level: .notice, data: .string("Notice")) + try await context.sendLogMessage(level: .warning, data: .string("Warning")) + try await context.sendLogMessage(level: .error, data: .string("Error")) + try await context.sendLogMessage(level: .critical, data: .string("Critical")) + try await context.sendLogMessage(level: .alert, data: .string("Alert")) + try await context.sendLogMessage(level: .emergency, data: .string("Emergency")) + + return CallTool.Result(content: [.text("Done")]) + } + + let client = Client(name: "AllLevelsClient", version: "1.0") + + await client.onNotification(LogMessageNotification.self) { message in + await receivedLogs.add( + level: message.params.level, + logger: message.params.logger, + data: message.params.data + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Set level to debug to receive all messages + try await client.setLoggingLevel(.debug) + + _ = try await client.callTool(name: "log_all_eight", arguments: [:]) + + try await Task.sleep(for: .milliseconds(100)) + + let logs = await receivedLogs.logs + #expect(logs.count == 8, "Should receive all 8 log levels, got \(logs.count)") + + // Verify each level in order + let expectedLevels: [LoggingLevel] = [ + .debug, .info, .notice, .warning, .error, .critical, .alert, .emergency + ] + let expectedMessages = ["Debug", "Info", "Notice", "Warning", "Error", "Critical", "Alert", "Emergency"] + + for (index, expectedLevel) in expectedLevels.enumerated() { + if index < logs.count { + #expect(logs[index].level == expectedLevel, "Level at index \(index) should be \(expectedLevel)") + #expect(logs[index].data.stringValue == expectedMessages[index]) + } + } + } + + /// Test that when no logging level is set, all messages are sent. + /// + /// Per MCP spec: "If no logging/setLevel request has been sent from the client, + /// the server MAY decide which messages to send automatically." + /// Our implementation sends all messages when no level is set. + @Test(.timeLimit(.minutes(1))) + func defaultLoggingBehaviorSendsAllMessages() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.defaultlog") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let receivedLogs = LogTracker() + + let server = Server( + name: "DefaultLogServer", + version: "1.0.0", + capabilities: .init(logging: .init(), tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "log_without_level_set", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, context in + guard request.name == "log_without_level_set" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + // Send messages at various levels without client setting a level + try await context.sendLogMessage(level: .debug, data: .string("Debug message")) + try await context.sendLogMessage(level: .info, data: .string("Info message")) + try await context.sendLogMessage(level: .error, data: .string("Error message")) + + return CallTool.Result(content: [.text("Done")]) + } + + let client = Client(name: "DefaultLogClient", version: "1.0") + + await client.onNotification(LogMessageNotification.self) { message in + await receivedLogs.add( + level: message.params.level, + logger: message.params.logger, + data: message.params.data + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Do NOT set logging level - test default behavior + _ = try await client.callTool(name: "log_without_level_set", arguments: [:]) + + try await Task.sleep(for: .milliseconds(100)) + + let logs = await receivedLogs.logs + // All 3 messages should be received since no level filter was set + #expect(logs.count == 3, "Should receive all 3 log messages when no level is set, got \(logs.count)") + + if logs.count >= 3 { + #expect(logs[0].level == .debug) + #expect(logs[1].level == .info) + #expect(logs[2].level == .error) + } + } + + /// Test server-level sendLogMessage method (outside request handlers). + /// + /// This matches TypeScript SDK behavior where `server.sendLoggingMessage()` + /// can be called outside of request handlers. + @Test(.timeLimit(.minutes(1))) + func serverLevelSendLogMessageWorks() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.serverlevellog") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let receivedLogs = LogTracker() + + let server = Server( + name: "ServerLevelLogServer", + version: "1.0.0", + capabilities: .init(logging: .init(), tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "trigger", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, _ in + guard request.name == "trigger" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + // Just return, we'll send log via server-level method + return CallTool.Result(content: [.text("Done")]) + } + + let client = Client(name: "ServerLevelLogClient", version: "1.0") + + await client.onNotification(LogMessageNotification.self) { message in + await receivedLogs.add( + level: message.params.level, + logger: message.params.logger, + data: message.params.data + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Use server-level sendLogMessage (not context-level) + try await server.sendLogMessage( + level: .info, + logger: "server-logger", + data: .string("Server-level log message") + ) + + try await Task.sleep(for: .milliseconds(100)) + + let logs = await receivedLogs.logs + #expect(logs.count == 1, "Should receive 1 log message from server-level sendLogMessage") + + if let log = logs.first { + #expect(log.level == .info) + #expect(log.logger == "server-logger") + #expect(log.data.stringValue == "Server-level log message") + } + } + + /// Test that sendLogMessage from context silently drops messages when logging + /// capability is not declared. + /// + /// This matches TypeScript SDK behavior where logging messages are silently + /// dropped when the server doesn't have the logging capability. + @Test(.timeLimit(.minutes(1))) + func contextSendLogMessageSilentlyDropsWithoutCapability() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.nocaplog") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let receivedLogs = LogTracker() + + // Server WITHOUT logging capability + let server = Server( + name: "NoLogCapServer", + version: "1.0.0", + capabilities: .init(tools: .init()) // No logging capability! + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "try_log", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, context in + guard request.name == "try_log" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + // This should be silently dropped since logging capability is not declared + try await context.sendLogMessage( + level: .info, + data: .string("This should be dropped") + ) + + return CallTool.Result(content: [.text("Done")]) + } + + let client = Client(name: "NoLogCapClient", version: "1.0") + + await client.onNotification(LogMessageNotification.self) { message in + await receivedLogs.add( + level: message.params.level, + logger: message.params.logger, + data: message.params.data + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + _ = try await client.callTool(name: "try_log", arguments: [:]) + + try await Task.sleep(for: .milliseconds(100)) + + let logs = await receivedLogs.logs + // No logs should be received because logging capability is not declared + #expect(logs.count == 0, "Should receive 0 logs when logging capability is not declared, got \(logs.count)") + } + + /// Test sendToolListChanged convenience method. + @Test(.timeLimit(.minutes(1))) + func sendToolListChangedWorks() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.toolchange") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let notificationReceived = NotificationTracker() + + let server = Server( + name: "ToolChangeServer", + version: "1.0.0", + capabilities: .init(tools: .init(listChanged: true)) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "notify_test", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, context in + guard request.name == "notify_test" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + // Send tool list changed notification + try await context.sendToolListChanged() + + return CallTool.Result(content: [.text("Done")]) + } + + let client = Client(name: "ToolChangeClient", version: "1.0") + + await client.onNotification(ToolListChangedNotification.self) { _ in + await notificationReceived.recordToolListChanged() + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + _ = try await client.callTool(name: "notify_test", arguments: [:]) + + try await Task.sleep(for: .milliseconds(100)) + + let count = await notificationReceived.toolListChangedCount + #expect(count == 1, "Should receive 1 tool list changed notification") + } + + /// Test sendResourceListChanged convenience method. + /// + /// This tests that the server can notify the client when the list of + /// available resources has changed. + @Test(.timeLimit(.minutes(1))) + func sendResourceListChangedWorks() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.resourcelistchange") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let notificationReceived = NotificationTracker() + + let server = Server( + name: "ResourceListChangeServer", + version: "1.0.0", + capabilities: .init(resources: .init(listChanged: true), tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "notify_test", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, context in + guard request.name == "notify_test" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + // Send resource list changed notification + try await context.sendResourceListChanged() + + return CallTool.Result(content: [.text("Done")]) + } + + let client = Client(name: "ResourceListChangeClient", version: "1.0") + + await client.onNotification(ResourceListChangedNotification.self) { _ in + await notificationReceived.recordResourceListChanged() + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + _ = try await client.callTool(name: "notify_test", arguments: [:]) + + try await Task.sleep(for: .milliseconds(100)) + + let count = await notificationReceived.resourceListChangedCount + #expect(count == 1, "Should receive 1 resource list changed notification") + } + + /// Test sendPromptListChanged convenience method. + /// + /// This tests that the server can notify the client when the list of + /// available prompts has changed. + @Test(.timeLimit(.minutes(1))) + func sendPromptListChangedWorks() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.promptlistchange") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let notificationReceived = NotificationTracker() + + let server = Server( + name: "PromptListChangeServer", + version: "1.0.0", + capabilities: .init(prompts: .init(listChanged: true), tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "notify_test", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, context in + guard request.name == "notify_test" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + // Send prompt list changed notification + try await context.sendPromptListChanged() + + return CallTool.Result(content: [.text("Done")]) + } + + let client = Client(name: "PromptListChangeClient", version: "1.0") + + await client.onNotification(PromptListChangedNotification.self) { _ in + await notificationReceived.recordPromptListChanged() + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + _ = try await client.callTool(name: "notify_test", arguments: [:]) + + try await Task.sleep(for: .milliseconds(100)) + + let count = await notificationReceived.promptListChangedCount + #expect(count == 1, "Should receive 1 prompt list changed notification") + } + + /// Test sendResourceUpdated convenience method. + @Test(.timeLimit(.minutes(1))) + func sendResourceUpdatedWorks() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.resourceupdate") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let notificationReceived = NotificationTracker() + + let server = Server( + name: "ResourceUpdateServer", + version: "1.0.0", + capabilities: .init(resources: .init(subscribe: true), tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "update_resource", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, context in + guard request.name == "update_resource" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + // Send resource updated notification + try await context.sendResourceUpdated(uri: "file:///path/to/resource.txt") + + return CallTool.Result(content: [.text("Done")]) + } + + let client = Client(name: "ResourceUpdateClient", version: "1.0") + + await client.onNotification(ResourceUpdatedNotification.self) { message in + await notificationReceived.recordResourceUpdated(uri: message.params.uri) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + _ = try await client.callTool(name: "update_resource", arguments: [:]) + + try await Task.sleep(for: .milliseconds(100)) + + let uris = await notificationReceived.resourceUpdatedURIs + #expect(uris.count == 1, "Should receive 1 resource updated notification") + #expect(uris.first == "file:///path/to/resource.txt") + } + + /// Test that client can send progress notifications to server (bidirectional progress). + /// Based on Python SDK's test_bidirectional_progress_notifications. + /// + /// This tests the reverse direction from serverSendsProgressNotificationsToClient: + /// - Client sends progress notifications using notify() + /// - Server receives them via onNotification(ProgressNotification.self) + @Test(.timeLimit(.minutes(1))) + func clientSendsProgressNotificationsToServer() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.progress.bidirectional") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + // Track progress updates received by server + let serverReceivedProgress = ProgressUpdateTracker() + + // Progress token that client will use + let clientProgressToken: ProgressToken = "client-progress-token-123" + + let server = Server( + name: "BidirectionalProgressServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + // Server registers handler to receive progress notifications from client + await server.onNotification(ProgressNotification.self) { message in + await serverReceivedProgress.add( + token: message.params.progressToken, + progress: message.params.progress, + total: message.params.total, + message: message.params.message + ) + } + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "simple_tool", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, _ in + guard request.name == "simple_tool" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + return CallTool.Result(content: [.text("Done")]) + } + + let client = Client(name: "BidirectionalProgressClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Client sends progress notifications to server (like Python test) + try await client.notify(ProgressNotification.message(.init( + progressToken: clientProgressToken, + progress: 0.33, + total: 1.0, + message: "Client progress 33%" + ))) + + try await client.notify(ProgressNotification.message(.init( + progressToken: clientProgressToken, + progress: 0.66, + total: 1.0, + message: "Client progress 66%" + ))) + + try await client.notify(ProgressNotification.message(.init( + progressToken: clientProgressToken, + progress: 1.0, + total: 1.0, + message: "Client progress 100%" + ))) + + // Give time for notifications to be processed + try await Task.sleep(for: .milliseconds(100)) + + // Verify server received progress updates from client + let updates = await serverReceivedProgress.updates + #expect(updates.count == 3, "Server should receive 3 progress notifications from client") + + if updates.count >= 3 { + // Verify first update + #expect(updates[0].token == clientProgressToken) + #expect(updates[0].progress == 0.33) + #expect(updates[0].message == "Client progress 33%") + + // Verify last update + #expect(updates[2].progress == 1.0) + #expect(updates[2].message == "Client progress 100%") + } + } + + /// Test bidirectional progress: both client→server and server→client in same session. + /// Based on Python SDK's test_bidirectional_progress_notifications which tests both + /// directions simultaneously. + @Test(.timeLimit(.minutes(1))) + func bidirectionalProgressNotifications() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.progress.fullbidirectional") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + // Track progress updates received by both sides + let serverReceivedProgress = ProgressUpdateTracker() + let clientReceivedProgress = ProgressUpdateTracker() + + // Tokens + let serverProgressToken: ProgressToken = "server-token-abc" + let clientProgressToken: ProgressToken = "client-token-xyz" + + let server = Server( + name: "FullBidirectionalServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + // Server registers handler to receive progress notifications from client + await server.onNotification(ProgressNotification.self) { message in + await serverReceivedProgress.add( + token: message.params.progressToken, + progress: message.params.progress, + total: message.params.total, + message: message.params.message + ) + } + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "progress_tool", inputSchema: ["type": "object"]) + ]) + } + + // Tool that sends progress back to client + await server.withRequestHandler(CallTool.self) { request, context in + guard request.name == "progress_tool" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + // Server sends progress notifications to client + try await context.sendProgress( + token: serverProgressToken, + progress: 0.5, + total: 1.0, + message: "Server progress 50%" + ) + + try await context.sendProgress( + token: serverProgressToken, + progress: 1.0, + total: 1.0, + message: "Server progress 100%" + ) + + return CallTool.Result(content: [.text("Done")]) + } + + let client = Client(name: "FullBidirectionalClient", version: "1.0") + + // Client registers handler to receive progress notifications from server + await client.onNotification(ProgressNotification.self) { message in + await clientReceivedProgress.add( + token: message.params.progressToken, + progress: message.params.progress, + total: message.params.total, + message: message.params.message + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Client sends progress notifications to server + try await client.notify(ProgressNotification.message(.init( + progressToken: clientProgressToken, + progress: 0.25, + total: 1.0, + message: "Client progress 25%" + ))) + + try await client.notify(ProgressNotification.message(.init( + progressToken: clientProgressToken, + progress: 0.75, + total: 1.0, + message: "Client progress 75%" + ))) + + // Call tool to trigger server→client progress + _ = try await client.callTool(name: "progress_tool", arguments: [:]) + + // Give time for all notifications to be processed + try await Task.sleep(for: .milliseconds(100)) + + // Verify client received progress updates from server + let clientUpdates = await clientReceivedProgress.updates + #expect(clientUpdates.count == 2, "Client should receive 2 progress notifications from server") + + if clientUpdates.count >= 2 { + #expect(clientUpdates[0].token == serverProgressToken) + #expect(clientUpdates[0].progress == 0.5) + #expect(clientUpdates[1].progress == 1.0) + } + + // Verify server received progress updates from client + let serverUpdates = await serverReceivedProgress.updates + #expect(serverUpdates.count == 2, "Server should receive 2 progress notifications from client") + + if serverUpdates.count >= 2 { + #expect(serverUpdates[0].token == clientProgressToken) + #expect(serverUpdates[0].progress == 0.25) + #expect(serverUpdates[1].progress == 0.75) + } + } + + /// Test that exceptions in progress notification handlers are logged but don't crash the session. + /// Based on Python SDK's test_progress_callback_exception_logging. + /// + /// This ensures that if a progress handler throws, the error is handled gracefully + /// and subsequent operations continue to work. + @Test(.timeLimit(.minutes(1))) + func progressNotificationHandlerExceptionDoesNotCrashSession() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.progress.exception") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + // Track that handler was called + let handlerCallTracker = HandlerCallTracker() + + // Custom error for testing + struct ProgressHandlerError: Error { + let message: String + } + + let server = Server( + name: "ProgressExceptionServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "progress_tool", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, context in + guard request.name == "progress_tool" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + // Server sends progress notification + guard let progressToken = request._meta?.progressToken else { + return CallTool.Result(content: [.text("No progress token")], isError: true) + } + + try await context.sendProgress( + token: progressToken, + progress: 50.0, + total: 100.0, + message: "Halfway" + ) + + return CallTool.Result(content: [.text("progress_result")]) + } + + let client = Client(name: "ProgressExceptionClient", version: "1.0") + + // Register a handler that throws an exception + await client.onNotification(ProgressNotification.self) { _ in + await handlerCallTracker.recordCall() + throw ProgressHandlerError(message: "Progress callback failed!") + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Call tool with progress token - the progress handler will throw + let result = try await client.send( + CallTool.request(.init( + name: "progress_tool", + arguments: [:], + _meta: RequestMeta(progressToken: .string("exception-test-token")) + )) + ) + + // Give time for notification to be processed + try await Task.sleep(for: .milliseconds(100)) + + // Verify the request completed successfully despite the callback failure + #expect(result.content.count == 1) + if case .text(let text, _, _) = result.content.first { + #expect(text == "progress_result") + } else { + Issue.record("Expected text content") + } + + // Verify the progress handler was called (even though it threw) + let callCount = await handlerCallTracker.callCount + #expect(callCount == 1, "Progress handler should have been called") + + // Session should still be functional - verify by making another request + let pingResult = try await client.send(Ping.request()) + // Ping returns empty result - just verify it doesn't throw + _ = pingResult + } + + /// Test that progress notifications with integer token 0 work in bidirectional flow. + /// Edge case: token 0 is falsy in many languages but should be treated as valid. + @Test(.timeLimit(.minutes(1))) + func clientSendsProgressWithZeroToken() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.progress.client.zero") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let serverReceivedProgress = ProgressUpdateTracker() + + let server = Server( + name: "ZeroTokenClientServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.onNotification(ProgressNotification.self) { message in + await serverReceivedProgress.add( + token: message.params.progressToken, + progress: message.params.progress, + total: message.params.total, + message: message.params.message + ) + } + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: []) + } + + let client = Client(name: "ZeroTokenClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Client sends progress with token 0 (edge case) + try await client.notify(ProgressNotification.message(.init( + progressToken: .integer(0), + progress: 50.0, + total: 100.0, + message: "Progress with zero token" + ))) + + try await Task.sleep(for: .milliseconds(100)) + + let updates = await serverReceivedProgress.updates + #expect(updates.count == 1, "Server should receive progress with token 0") + + if let update = updates.first { + #expect(update.token == .integer(0), "Token should be integer 0") + #expect(update.progress == 50.0) + } + } + } + + // MARK: - ProgressTracker Actor Tests + + @Suite("ProgressTracker actor (server-side cumulative progress)") + struct ProgressTrackerTests { + + /// Test that ProgressTracker accumulates progress correctly with advance(by:). + @Test(.timeLimit(.minutes(1))) + func progressTrackerAdvanceAccumulatesProgress() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.tracker.advance") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let receivedProgress = ProgressUpdateTracker() + + let server = Server( + name: "TrackerAdvanceServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "tracker_advance_test", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, context in + guard request.name == "tracker_advance_test" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + guard let token = request._meta?.progressToken else { + return CallTool.Result(content: [.text("No token")], isError: true) + } + + // Use ProgressTracker to accumulate progress + let tracker = ProgressTracker(token: token, total: 100, context: context) + + try await tracker.advance(by: 25, message: "Step 1") + try await tracker.advance(by: 25, message: "Step 2") + try await tracker.advance(by: 50, message: "Step 3") + + return CallTool.Result(content: [.text("Done")]) + } + + let client = Client(name: "TrackerAdvanceClient", version: "1.0") + + await client.onNotification(ProgressNotification.self) { message in + await receivedProgress.add( + token: message.params.progressToken, + progress: message.params.progress, + total: message.params.total, + message: message.params.message + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + _ = try await client.send( + CallTool.request(.init( + name: "tracker_advance_test", + arguments: [:], + _meta: RequestMeta(progressToken: .string("tracker-test")) + )) + ) + + try await Task.sleep(for: .milliseconds(100)) + + let updates = await receivedProgress.updates + #expect(updates.count == 3, "Should receive 3 progress notifications") + + if updates.count >= 3 { + // Verify cumulative progress values + #expect(updates[0].progress == 25.0, "First advance should be 25") + #expect(updates[0].message == "Step 1") + + #expect(updates[1].progress == 50.0, "Second advance should be 50 (25+25)") + #expect(updates[1].message == "Step 2") + + #expect(updates[2].progress == 100.0, "Third advance should be 100 (50+50)") + #expect(updates[2].message == "Step 3") + } + } + + /// Test that ProgressTracker.set(to:) sets absolute progress. + @Test(.timeLimit(.minutes(1))) + func progressTrackerSetToAbsoluteValue() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.tracker.set") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let receivedProgress = ProgressUpdateTracker() + + let server = Server( + name: "TrackerSetServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "tracker_set_test", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, context in + guard request.name == "tracker_set_test" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + guard let token = request._meta?.progressToken else { + return CallTool.Result(content: [.text("No token")], isError: true) + } + + let tracker = ProgressTracker(token: token, total: 100, context: context) + + // Use set(to:) for absolute values + try await tracker.set(to: 10, message: "10%") + try await tracker.set(to: 50, message: "50%") + try await tracker.set(to: 100, message: "100%") + + return CallTool.Result(content: [.text("Done")]) + } + + let client = Client(name: "TrackerSetClient", version: "1.0") + + await client.onNotification(ProgressNotification.self) { message in + await receivedProgress.add( + token: message.params.progressToken, + progress: message.params.progress, + total: message.params.total, + message: message.params.message + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + _ = try await client.send( + CallTool.request(.init( + name: "tracker_set_test", + arguments: [:], + _meta: RequestMeta(progressToken: .string("set-test")) + )) + ) + + try await Task.sleep(for: .milliseconds(100)) + + let updates = await receivedProgress.updates + #expect(updates.count == 3, "Should receive 3 progress notifications") + + if updates.count >= 3 { + #expect(updates[0].progress == 10.0) + #expect(updates[1].progress == 50.0) + #expect(updates[2].progress == 100.0) + } + } + + /// Test that ProgressTracker.update(message:) sends notification without changing progress. + @Test(.timeLimit(.minutes(1))) + func progressTrackerUpdateMessageOnly() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.tracker.update") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let receivedProgress = ProgressUpdateTracker() + + let server = Server( + name: "TrackerUpdateServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "tracker_update_test", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, context in + guard request.name == "tracker_update_test" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + guard let token = request._meta?.progressToken else { + return CallTool.Result(content: [.text("No token")], isError: true) + } + + let tracker = ProgressTracker(token: token, total: 100, context: context) + + // First advance to set initial progress + try await tracker.advance(by: 50, message: "Halfway") + // Use update(message:) to send message without changing value + try await tracker.update(message: "Still at 50%, processing...") + try await tracker.update(message: "Almost done with phase 1") + // Advance again to complete + try await tracker.advance(by: 50, message: "Complete") + + return CallTool.Result(content: [.text("Done")]) + } + + let client = Client(name: "TrackerUpdateClient", version: "1.0") + + await client.onNotification(ProgressNotification.self) { message in + await receivedProgress.add( + token: message.params.progressToken, + progress: message.params.progress, + total: message.params.total, + message: message.params.message + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + _ = try await client.send( + CallTool.request(.init( + name: "tracker_update_test", + arguments: [:], + _meta: RequestMeta(progressToken: .string("update-test")) + )) + ) + + try await Task.sleep(for: .milliseconds(100)) + + let updates = await receivedProgress.updates + // Should receive 4 notifications: advance(50), update(), update(), advance(100) + #expect(updates.count == 4, "Should receive 4 progress notifications") + + if updates.count >= 4 { + // First advance: progress = 50 + #expect(updates[0].progress == 50.0) + #expect(updates[0].message == "Halfway") + // First update: progress should still be 50 + #expect(updates[1].progress == 50.0) + #expect(updates[1].message == "Still at 50%, processing...") + // Second update: progress should still be 50 + #expect(updates[2].progress == 50.0) + #expect(updates[2].message == "Almost done with phase 1") + // Final advance: progress = 100 + #expect(updates[3].progress == 100.0) + #expect(updates[3].message == "Complete") + } + } + } + + // MARK: - Client onProgress Callback Tests + + @Suite("Client send with onProgress callback") + struct ClientOnProgressTests { + + /// Test that send(_:onProgress:) receives progress updates via callback. + @Test(.timeLimit(.minutes(1))) + func clientReceivesProgressViaCallback() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.client.callback") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let receivedProgress = ProgressUpdateTracker() + + let server = Server( + name: "CallbackTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "callback_test", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, context in + guard request.name == "callback_test" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + // Server extracts the auto-generated progress token from _meta + guard let token = request._meta?.progressToken else { + return CallTool.Result(content: [.text("No progress token - callback should have set one")], isError: true) + } + + // Send progress notifications + try await context.sendProgress(token: token, progress: 1, total: 3, message: "Step 1") + try await context.sendProgress(token: token, progress: 2, total: 3, message: "Step 2") + try await context.sendProgress(token: token, progress: 3, total: 3, message: "Step 3") + + return CallTool.Result(content: [.text("Done")]) + } + + let client = Client(name: "CallbackTestClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Use send with onProgress callback - no need to manually set progressToken + let result = try await client.send( + CallTool.request(.init(name: "callback_test", arguments: [:])), + onProgress: { progress in + Task { + // Record the progress (use a dummy token since we just care about values) + await receivedProgress.add( + token: .string("callback"), + progress: progress.value, + total: progress.total, + message: progress.message + ) + } + } + ) + + try await Task.sleep(for: .milliseconds(100)) + + // Verify result + if case .text(let text, _, _) = result.content.first { + #expect(text == "Done") + } + + // Verify progress callback was invoked + let updates = await receivedProgress.updates + #expect(updates.count == 3, "Should receive 3 progress updates via callback") + + if updates.count >= 3 { + #expect(updates[0].progress == 1.0) + #expect(updates[0].message == "Step 1") + #expect(updates[2].progress == 3.0) + #expect(updates[2].message == "Step 3") + } + } + + /// Test that send(_:onProgress:) automatically injects progressToken into _meta. + @Test(.timeLimit(.minutes(1))) + func clientAutoInjectsProgressToken() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.client.autoinject") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + // Track whether server received a progress token + let tokenTracker = TokenTracker() + + let server = Server( + name: "AutoInjectServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "token_check", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, context in + guard request.name == "token_check" else { + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + + // Record whether we received a progress token + if let token = request._meta?.progressToken { + await tokenTracker.record(token) + } + + return CallTool.Result(content: [.text("Done")]) + } + + let client = Client(name: "AutoInjectClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Send request WITHOUT manually setting _meta.progressToken + // The send(_:onProgress:) should auto-inject it + _ = try await client.send( + CallTool.request(.init(name: "token_check", arguments: [:])), + onProgress: { _ in } + ) + + // Verify server received a progress token + let receivedToken = await tokenTracker.token + #expect(receivedToken != nil, "Server should have received an auto-injected progress token") + } + } + + // MARK: - Task-Augmented Progress Tests + + @Suite("Task-augmented progress (progress continues after CreateTaskResult)") + struct TaskAugmentedProgressTests { + + /// Test that TaskStatus.isTerminal correctly identifies terminal statuses. + @Test("TaskStatus.isTerminal identifies terminal states") + func taskStatusIsTerminal() { + #expect(TaskStatus.working.isTerminal == false) + #expect(TaskStatus.inputRequired.isTerminal == false) + #expect(TaskStatus.completed.isTerminal == true) + #expect(TaskStatus.failed.isTerminal == true) + #expect(TaskStatus.cancelled.isTerminal == true) + } + + /// Test that cleanupTaskProgressHandler removes progress callback. + @Test("cleanupTaskProgressHandler removes progress callback for task") + func cleanUpTaskProgressHandler() async throws { + let client = Client(name: "CleanupTest", version: "1.0") + + // Manually test the cleanup method by verifying it doesn't crash + // when called with non-existent task ID (should be no-op) + await client.cleanUpTaskProgressHandler(taskId: "non-existent-task") + // If we get here without crashing, the test passes + } + + /// Test that checkForTaskResponse correctly identifies task responses. + /// This tests the internal logic by verifying the Value structure parsing. + @Test("Task response detection parses task.taskId from response") + func taskResponseDetection() throws { + // CreateTaskResult structure: { "task": { "taskId": "...", "status": "...", ... } } + let taskResponse: [String: Value] = [ + "task": .object([ + "taskId": .string("test-task-123"), + "status": .string("working"), + "ttl": .null, + "createdAt": .string("2024-01-01T00:00:00Z"), + "lastUpdatedAt": .string("2024-01-01T00:00:00Z") + ]) + ] + + // Verify the structure can be parsed + guard let taskValue = taskResponse["task"], + case .object(let taskObject) = taskValue, + let taskIdValue = taskObject["taskId"], + case .string(let taskId) = taskIdValue else { + Issue.record("Failed to parse task response structure") + return + } + + #expect(taskId == "test-task-123") + } + + /// Test that non-task responses are correctly identified. + @Test("Non-task response detection returns nil taskId") + func nonTaskResponseDetection() throws { + // Regular CallTool.Result structure (no task field) + let regularResponse: [String: Value] = [ + "content": .array([ + .object([ + "type": .string("text"), + "text": .string("Hello") + ]) + ]) + ] + + // Verify the task field is not present + let taskValue = regularResponse["task"] + #expect(taskValue == nil, "Regular response should not have task field") + } + + /// Test TaskStatusNotification.Parameters decoding. + @Test("TaskStatusNotification.Parameters decodes correctly") + func taskStatusNotificationDecoding() throws { + let json = """ + { + "taskId": "task-abc", + "status": "completed", + "ttl": null, + "createdAt": "2024-01-01T00:00:00Z", + "lastUpdatedAt": "2024-01-01T00:00:01Z" + } + """ + + let decoder = JSONDecoder() + let params = try decoder.decode( + TaskStatusNotification.Parameters.self, + from: json.data(using: .utf8)! + ) + + #expect(params.taskId == "task-abc") + #expect(params.status == .completed) + #expect(params.status.isTerminal == true) + } + + /// Test that terminal task status notification triggers cleanup. + @Test("Terminal task status triggers progress cleanup") + func terminalTaskStatusTriggersCleanup() throws { + // Test that the isTerminal check works as expected for cleanup logic + let completedStatus = TaskStatus.completed + let workingStatus = TaskStatus.working + + #expect(completedStatus.isTerminal == true, "Completed should trigger cleanup") + #expect(workingStatus.isTerminal == false, "Working should not trigger cleanup") + } + } +} + +// MARK: - Test Helpers + +/// Tracker for progress tokens received by server. +private actor TokenTracker { + private(set) var token: ProgressToken? + + func record(_ token: ProgressToken) { + self.token = token + } +} + +/// Thread-safe tracker for handler calls. +private actor HandlerCallTracker { + private(set) var callCount = 0 + + func recordCall() { + callCount += 1 + } +} + +/// Thread-safe tracker for received progress updates. +private actor ProgressUpdateTracker { + struct Update { + let token: ProgressToken + let progress: Double + let total: Double? + let message: String? + } + + private(set) var updates: [Update] = [] + + func add(token: ProgressToken, progress: Double, total: Double?, message: String?) { + updates.append(Update(token: token, progress: progress, total: total, message: message)) + } +} + +/// Thread-safe tracker for received log messages. +private actor LogTracker { + struct Log { + let level: LoggingLevel + let logger: String? + let data: Value + } + + private(set) var logs: [Log] = [] + + func add(level: LoggingLevel, logger: String?, data: Value) { + logs.append(Log(level: level, logger: logger, data: data)) + } +} + +/// Thread-safe tracker for various notification types. +private actor NotificationTracker { + private(set) var toolListChangedCount = 0 + private(set) var resourceListChangedCount = 0 + private(set) var promptListChangedCount = 0 + private(set) var resourceUpdatedURIs: [String] = [] + + func recordToolListChanged() { + toolListChangedCount += 1 + } + + func recordResourceListChanged() { + resourceListChangedCount += 1 + } + + func recordPromptListChanged() { + promptListChangedCount += 1 + } + + func recordResourceUpdated(uri: String) { + resourceUpdatedURIs.append(uri) + } +} diff --git a/Tests/MCPTests/PromptTests.swift b/Tests/MCPTests/PromptTests.swift index 20e5c279..3c2e9607 100644 --- a/Tests/MCPTests/PromptTests.swift +++ b/Tests/MCPTests/PromptTests.swift @@ -38,7 +38,7 @@ struct PromptTests { let decoded = try decoder.decode(Prompt.Message.self, from: data) #expect(decoded.role == .user) - if case .text(let text) = decoded.content { + if case .text(let text, _, _) = decoded.content { #expect(text == "Hello, world!") } else { #expect(Bool(false), "Expected text content") @@ -51,10 +51,10 @@ struct PromptTests { let decoder = JSONDecoder() // Test text content - let textContent = Prompt.Message.Content.text(text: "Test text") + let textContent = Prompt.Message.Content.text("Test text") let textData = try encoder.encode(textContent) let decodedText = try decoder.decode(Prompt.Message.Content.self, from: textData) - if case .text(let text) = decodedText { + if case .text(let text, _, _) = decodedText { #expect(text == "Test text") } else { #expect(Bool(false), "Expected text content") @@ -65,7 +65,7 @@ struct PromptTests { data: "base64audiodata", mimeType: "audio/wav") let audioData = try encoder.encode(audioContent) let decodedAudio = try decoder.decode(Prompt.Message.Content.self, from: audioData) - if case .audio(let data, let mimeType) = decodedAudio { + if case .audio(let data, let mimeType, _, _) = decodedAudio { #expect(data == "base64audiodata") #expect(mimeType == "audio/wav") } else { @@ -76,7 +76,7 @@ struct PromptTests { let imageContent = Prompt.Message.Content.image(data: "base64data", mimeType: "image/png") let imageData = try encoder.encode(imageContent) let decodedImage = try decoder.decode(Prompt.Message.Content.self, from: imageData) - if case .image(let data, let mimeType) = decodedImage { + if case .image(let data, let mimeType, _, _) = decodedImage { #expect(data == "base64data") #expect(mimeType == "image/png") } else { @@ -87,16 +87,14 @@ struct PromptTests { let resourceContent = Prompt.Message.Content.resource( uri: "file://test.txt", mimeType: "text/plain", - text: "Sample text", - blob: "blob_data" + text: "Sample text" ) let resourceData = try encoder.encode(resourceContent) let decodedResource = try decoder.decode(Prompt.Message.Content.self, from: resourceData) - if case .resource(let uri, let mimeType, let text, let blob) = decodedResource { - #expect(uri == "file://test.txt") - #expect(mimeType == "text/plain") - #expect(text == "Sample text") - #expect(blob == "blob_data") + if case .resource(let resourceData, _, _) = decodedResource { + #expect(resourceData.uri == "file://test.txt") + #expect(resourceData.mimeType == "text/plain") + #expect(resourceData.text == "Sample text") } else { #expect(Bool(false), "Expected resource content") } @@ -106,6 +104,7 @@ struct PromptTests { func testPromptReference() throws { let reference = Prompt.Reference(name: "test_prompt") #expect(reference.name == "test_prompt") + #expect(reference.title == nil) let encoder = JSONEncoder() let decoder = JSONDecoder() @@ -114,19 +113,43 @@ struct PromptTests { let decoded = try decoder.decode(Prompt.Reference.self, from: data) #expect(decoded.name == "test_prompt") + #expect(decoded.title == nil) + } + + @Test("Prompt Reference with title validation") + func testPromptReferenceWithTitle() throws { + let reference = Prompt.Reference(name: "test_prompt", title: "Test Prompt") + #expect(reference.name == "test_prompt") + #expect(reference.title == "Test Prompt") + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(reference) + let decoded = try decoder.decode(Prompt.Reference.self, from: data) + + #expect(decoded.name == "test_prompt") + #expect(decoded.title == "Test Prompt") + + // Verify JSON structure includes title + let jsonObject = try JSONSerialization.jsonObject(with: data) as! [String: Any] + #expect(jsonObject["type"] as? String == "ref/prompt") + #expect(jsonObject["name"] as? String == "test_prompt") + #expect(jsonObject["title"] as? String == "Test Prompt") } @Test("GetPrompt parameters validation") func testGetPromptParameters() throws { - let arguments: [String: Value] = [ - "param1": .string("value1"), - "param2": .int(42), + // Per MCP spec, prompt arguments must be string values + let arguments: [String: String] = [ + "param1": "value1", + "param2": "42", ] let params = GetPrompt.Parameters(name: "test_prompt", arguments: arguments) #expect(params.name == "test_prompt") - #expect(params.arguments?["param1"] == .string("value1")) - #expect(params.arguments?["param2"] == .int(42)) + #expect(params.arguments?["param1"] == "value1") + #expect(params.arguments?["param2"] == "42") } @Test("GetPrompt result validation") @@ -206,7 +229,7 @@ struct PromptTests { // Test user message factory method let userMessage: Prompt.Message = .user("Hello, world!") #expect(userMessage.role == .user) - if case .text(let text) = userMessage.content { + if case .text(let text, _, _) = userMessage.content { #expect(text == "Hello, world!") } else { #expect(Bool(false), "Expected text content") @@ -215,7 +238,7 @@ struct PromptTests { // Test assistant message factory method let assistantMessage: Prompt.Message = .assistant("Hi there!") #expect(assistantMessage.role == .assistant) - if case .text(let text) = assistantMessage.content { + if case .text(let text, _, _) = assistantMessage.content { #expect(text == "Hi there!") } else { #expect(Bool(false), "Expected text content") @@ -224,7 +247,7 @@ struct PromptTests { // Test with image content let imageMessage: Prompt.Message = .user(.image(data: "base64data", mimeType: "image/png")) #expect(imageMessage.role == .user) - if case .image(let data, let mimeType) = imageMessage.content { + if case .image(let data, let mimeType, _, _) = imageMessage.content { #expect(data == "base64data") #expect(mimeType == "image/png") } else { @@ -235,7 +258,7 @@ struct PromptTests { let audioMessage: Prompt.Message = .assistant( .audio(data: "base64audio", mimeType: "audio/wav")) #expect(audioMessage.role == .assistant) - if case .audio(let data, let mimeType) = audioMessage.content { + if case .audio(let data, let mimeType, _, _) = audioMessage.content { #expect(data == "base64audio") #expect(mimeType == "audio/wav") } else { @@ -244,14 +267,12 @@ struct PromptTests { // Test with resource content let resourceMessage: Prompt.Message = .user( - .resource( - uri: "file://test.txt", mimeType: "text/plain", text: "Sample text", blob: nil)) + .resource(uri: "file://test.txt", mimeType: "text/plain", text: "Sample text")) #expect(resourceMessage.role == .user) - if case .resource(let uri, let mimeType, let text, let blob) = resourceMessage.content { - #expect(uri == "file://test.txt") - #expect(mimeType == "text/plain") - #expect(text == "Sample text") - #expect(blob == nil) + if case .resource(let resourceContent, _, _) = resourceMessage.content { + #expect(resourceContent.uri == "file://test.txt") + #expect(resourceContent.mimeType == "text/plain") + #expect(resourceContent.text == "Sample text") } else { #expect(Bool(false), "Expected resource content") } @@ -262,7 +283,7 @@ struct PromptTests { // Test string literal assignment let content: Prompt.Message.Content = "Hello from string literal" - if case .text(let text) = content { + if case .text(let text, _, _) = content { #expect(text == "Hello from string literal") } else { #expect(Bool(false), "Expected text content") @@ -270,7 +291,7 @@ struct PromptTests { // Test in message creation let message: Prompt.Message = .user("Direct string literal") - if case .text(let text) = message.content { + if case .text(let text, _, _) = message.content { #expect(text == "Direct string literal") } else { #expect(Bool(false), "Expected text content") @@ -299,7 +320,7 @@ struct PromptTests { let content: Prompt.Message.Content = "Hello \(userName), welcome to your \(position) interview at \(company)" - if case .text(let text) = content { + if case .text(let text, _, _) = content { #expect(text == "Hello Alice, welcome to your Software Engineer interview at TechCorp") } else { #expect(Bool(false), "Expected text content") @@ -308,7 +329,7 @@ struct PromptTests { // Test in message creation with interpolation let message: Prompt.Message = .user( "Hi \(userName), I'm excited about the \(position) role at \(company)") - if case .text(let text) = message.content { + if case .text(let text, _, _) = message.content { #expect(text == "Hi Alice, I'm excited about the Software Engineer role at TechCorp") } else { #expect(Bool(false), "Expected text content") @@ -321,7 +342,7 @@ struct PromptTests { "I see you have \(experience) years of experience with \(skills.joined(separator: ", ")). That's impressive!" ) - if case .text(let text) = interviewMessage.content { + if case .text(let text, _, _) = interviewMessage.content { #expect( text == "I see you have 5 years of experience with Swift, Python, JavaScript. That's impressive!" @@ -342,7 +363,7 @@ struct PromptTests { let userMessage: Prompt.Message = .user( "Hello, I'm \(candidateName) and I'm interviewing for the \(position) position") #expect(userMessage.role == .user) - if case .text(let text) = userMessage.content { + if case .text(let text, _, _) = userMessage.content { #expect(text == "Hello, I'm Bob and I'm interviewing for the Data Scientist position") } else { #expect(Bool(false), "Expected text content") @@ -353,7 +374,7 @@ struct PromptTests { "Welcome \(candidateName)! Tell me about your \(experience) years of experience in data science" ) #expect(assistantMessage.role == .assistant) - if case .text(let text) = assistantMessage.content { + if case .text(let text, _, _) = assistantMessage.content { #expect(text == "Welcome Bob! Tell me about your 3 years of experience in data science") } else { #expect(Bool(false), "Expected text content") @@ -372,7 +393,7 @@ struct PromptTests { #expect(conversation.count == 4) // Verify interpolated content - if case .text(let text) = conversation[2].content { + if case .text(let text, _, _) = conversation[2].content { #expect(text == "I have 3 years of experience in the field") } else { #expect(Bool(false), "Expected text content") @@ -443,3 +464,672 @@ struct PromptTests { #expect(decoded[3].role == .assistant) } } + +// MARK: - Prompt Pagination Tests + +@Suite("Prompt Pagination Tests") +struct PromptPaginationTests { + + @Test("ListPrompts cursor parameter encodes correctly") + func cursorParameterEncoding() throws { + let testCursor = "test-cursor-123" + let params = ListPrompts.Parameters(cursor: testCursor) + + let encoder = JSONEncoder() + let data = try encoder.encode(params) + let jsonString = String(data: data, encoding: .utf8)! + + #expect(jsonString.contains("\"cursor\":\"test-cursor-123\"")) + } + + @Test("ListPrompts result with nextCursor encodes correctly") + func resultWithNextCursor() throws { + let prompts = [ + Prompt(name: "prompt1", description: "First prompt"), + Prompt(name: "prompt2", description: "Second prompt"), + ] + let result = ListPrompts.Result(prompts: prompts, nextCursor: "next-page-token") + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(result) + let decoded = try decoder.decode(ListPrompts.Result.self, from: data) + + #expect(decoded.prompts.count == 2) + #expect(decoded.nextCursor == "next-page-token") + } + + @Test("ListPrompts result without nextCursor indicates end of pagination") + func resultWithoutNextCursor() throws { + let prompts = [ + Prompt(name: "final_prompt", description: "Final prompt") + ] + let result = ListPrompts.Result(prompts: prompts, nextCursor: nil) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(result) + let decoded = try decoder.decode(ListPrompts.Result.self, from: data) + + #expect(decoded.nextCursor == nil) + + // Verify null cursor is not included in JSON + let jsonString = String(data: data, encoding: .utf8)! + #expect(!jsonString.contains("nextCursor")) + } + + @Test("ListPrompts request with cursor decodes correctly") + func requestWithCursorDecoding() throws { + let jsonString = """ + {"jsonrpc":"2.0","id":"page-2","method":"prompts/list","params":{"cursor":"page-1-token"}} + """ + let jsonData = jsonString.data(using: .utf8)! + + let decoded = try JSONDecoder().decode(Request.self, from: jsonData) + + #expect(decoded.id == "page-2") + #expect(decoded.params.cursor == "page-1-token") + } + + @Test("Simulated multi-page prompt listing") + func simulatedMultiPagePromptListing() throws { + // Simulate a server that returns 20 prompts across multiple pages + let allPrompts = (0..<20).map { i in + Prompt(name: "prompt_\(i)", description: "Prompt number \(i)") + } + + let pageSize = 7 + var collectedPrompts: [Prompt] = [] + var currentCursor: String? = nil + + // Simulate pagination + for pageIndex in 0..<3 { + let startIndex = pageIndex * pageSize + let endIndex = min(startIndex + pageSize, allPrompts.count) + let pagePrompts = Array(allPrompts[startIndex.. ListPrompts.Result { + let startIndex: Int + let nextCursor: String? + + switch cursor { + case nil: + startIndex = 0 + nextCursor = "page_2" + case "page_2": + startIndex = 3 + nextCursor = "page_3" + case "page_3": + startIndex = 6 + nextCursor = nil + default: + return ListPrompts.Result(prompts: []) + } + + let endIndex = min(startIndex + pageSize, allPrompts.count) + let pagePrompts = Array(allPrompts[startIndex..(id: id, method: CallTool.name, params: params) diff --git a/Tests/MCPTests/ResourceSubscriptionTests.swift b/Tests/MCPTests/ResourceSubscriptionTests.swift new file mode 100644 index 00000000..88f94263 --- /dev/null +++ b/Tests/MCPTests/ResourceSubscriptionTests.swift @@ -0,0 +1,446 @@ +import Foundation +import Logging +import Testing + +#if canImport(System) + import System +#else + @preconcurrency import SystemPackage +#endif + +@testable import MCP + +@Suite("Resource Subscription Tests") +struct ResourceSubscriptionTests { + + // MARK: - End-to-End Subscribe Tests + + @Test("Client can subscribe to a resource when server supports subscriptions") + @available(macOS 14.0, *) + func subscribeToResourceWorks() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.subscribe") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let subscriptionTracker = SubscriptionTracker() + + let server = Server( + name: "SubscriptionServer", + version: "1.0.0", + capabilities: .init(resources: .init(subscribe: true)) + ) + + // Handle resource subscription requests + await server.withRequestHandler(ResourceSubscribe.self) { request, _ in + await subscriptionTracker.recordSubscribe(uri: request.uri) + return Empty() + } + + // Handle list resources (required for resources capability) + await server.withRequestHandler(ListResources.self) { _, _ in + ListResources.Result(resources: [ + Resource(name: "test", uri: "file:///test.txt") + ]) + } + + let client = Client(name: "SubscriptionClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Subscribe to a resource + try await client.subscribeToResource(uri: "file:///test.txt") + + // Verify subscription was received by server + let subscriptions = await subscriptionTracker.subscribedURIs + #expect(subscriptions.count == 1) + #expect(subscriptions.first == "file:///test.txt") + } + + @Test("Client can unsubscribe from a resource when server supports subscriptions") + @available(macOS 14.0, *) + func unsubscribeFromResourceWorks() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.unsubscribe") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let subscriptionTracker = SubscriptionTracker() + + let server = Server( + name: "UnsubscriptionServer", + version: "1.0.0", + capabilities: .init(resources: .init(subscribe: true)) + ) + + // Handle resource unsubscription requests + await server.withRequestHandler(ResourceUnsubscribe.self) { request, _ in + await subscriptionTracker.recordUnsubscribe(uri: request.uri) + return Empty() + } + + // Handle list resources (required for resources capability) + await server.withRequestHandler(ListResources.self) { _, _ in + ListResources.Result(resources: [ + Resource(name: "test", uri: "file:///test.txt") + ]) + } + + let client = Client(name: "UnsubscriptionClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Unsubscribe from a resource + try await client.unsubscribeFromResource(uri: "file:///test.txt") + + // Verify unsubscription was received by server + let unsubscriptions = await subscriptionTracker.unsubscribedURIs + #expect(unsubscriptions.count == 1) + #expect(unsubscriptions.first == "file:///test.txt") + } + + @Test("Subscribe and unsubscribe cycle works correctly") + @available(macOS 14.0, *) + func subscribeUnsubscribeCycleWorks() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.subscription.cycle") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let subscriptionTracker = SubscriptionTracker() + + let server = Server( + name: "SubscriptionCycleServer", + version: "1.0.0", + capabilities: .init(resources: .init(subscribe: true)) + ) + + await server.withRequestHandler(ResourceSubscribe.self) { request, _ in + await subscriptionTracker.recordSubscribe(uri: request.uri) + return Empty() + } + + await server.withRequestHandler(ResourceUnsubscribe.self) { request, _ in + await subscriptionTracker.recordUnsubscribe(uri: request.uri) + return Empty() + } + + await server.withRequestHandler(ListResources.self) { _, _ in + ListResources.Result(resources: [ + Resource(name: "doc1", uri: "file:///doc1.txt"), + Resource(name: "doc2", uri: "file:///doc2.txt"), + ]) + } + + let client = Client(name: "CycleClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Subscribe to multiple resources + try await client.subscribeToResource(uri: "file:///doc1.txt") + try await client.subscribeToResource(uri: "file:///doc2.txt") + + // Unsubscribe from one + try await client.unsubscribeFromResource(uri: "file:///doc1.txt") + + // Verify state + let subscriptions = await subscriptionTracker.subscribedURIs + let unsubscriptions = await subscriptionTracker.unsubscribedURIs + + #expect(subscriptions.count == 2) + #expect(subscriptions.contains("file:///doc1.txt")) + #expect(subscriptions.contains("file:///doc2.txt")) + #expect(unsubscriptions.count == 1) + #expect(unsubscriptions.first == "file:///doc1.txt") + } + + // MARK: - Capability Validation Tests + + @Test("Subscribe throws when server does not support subscriptions") + @available(macOS 14.0, *) + func subscribeThrowsWithoutCapability() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.subscribe.nocap") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + // Server with resources capability but WITHOUT subscribe + let server = Server( + name: "NoSubscribeServer", + version: "1.0.0", + capabilities: .init(resources: .init(subscribe: false)) + ) + + await server.withRequestHandler(ListResources.self) { _, _ in + ListResources.Result(resources: [ + Resource(name: "test", uri: "file:///test.txt") + ]) + } + + // Client configured as strict (default) to enforce capability checks + let client = Client( + name: "StrictClient", + version: "1.0", + configuration: .init(strict: true) + ) + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Attempt to subscribe should throw + await #expect(throws: MCPError.self) { + try await client.subscribeToResource(uri: "file:///test.txt") + } + } + + @Test("Unsubscribe throws when server does not support subscriptions") + @available(macOS 14.0, *) + func unsubscribeThrowsWithoutCapability() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.unsubscribe.nocap") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + // Server with resources capability but WITHOUT subscribe + let server = Server( + name: "NoUnsubscribeServer", + version: "1.0.0", + capabilities: .init(resources: .init(subscribe: false)) + ) + + await server.withRequestHandler(ListResources.self) { _, _ in + ListResources.Result(resources: [ + Resource(name: "test", uri: "file:///test.txt") + ]) + } + + // Client configured as strict to enforce capability checks + let client = Client( + name: "StrictClient", + version: "1.0", + configuration: .init(strict: true) + ) + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Attempt to unsubscribe should throw + await #expect(throws: MCPError.self) { + try await client.unsubscribeFromResource(uri: "file:///test.txt") + } + } + + @Test("Subscribe throws when server has no resources capability at all") + @available(macOS 14.0, *) + func subscribeThrowsWithNoResourcesCapability() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.subscribe.nores") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + // Server with NO resources capability + let server = Server( + name: "NoResourcesServer", + version: "1.0.0", + capabilities: .init(tools: .init()) // Only tools, no resources + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: []) + } + + // Client configured as strict to enforce capability checks + let client = Client( + name: "StrictClient", + version: "1.0", + configuration: .init(strict: true) + ) + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Attempt to subscribe should throw because resources capability is nil + await #expect(throws: MCPError.self) { + try await client.subscribeToResource(uri: "file:///test.txt") + } + } + + // MARK: - Notification Flow Tests + + @Test("Server can send resource updated notification after subscription") + @available(macOS 14.0, *) + func resourceUpdatedNotificationAfterSubscribe() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.subscribe.notify") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let notificationTracker = NotificationTracker() + + let server = Server( + name: "NotifyingServer", + version: "1.0.0", + capabilities: .init(resources: .init(subscribe: true), tools: .init()) + ) + + await server.withRequestHandler(ResourceSubscribe.self) { _, _ in + Empty() + } + + await server.withRequestHandler(ListResources.self) { _, _ in + ListResources.Result(resources: [ + Resource(name: "watched", uri: "file:///watched.txt") + ]) + } + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "trigger_update", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, context in + guard request.name == "trigger_update" else { + return CallTool.Result(content: [.text("Unknown tool")], isError: true) + } + // Send resource updated notification + try await context.sendResourceUpdated(uri: "file:///watched.txt") + return CallTool.Result(content: [.text("Update sent")]) + } + + let client = Client(name: "WatchingClient", version: "1.0") + + await client.onNotification(ResourceUpdatedNotification.self) { message in + await notificationTracker.recordResourceUpdated(uri: message.params.uri) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Subscribe first + try await client.subscribeToResource(uri: "file:///watched.txt") + + // Trigger an update via tool call + _ = try await client.callTool(name: "trigger_update", arguments: [:]) + + // Wait for notification to be processed + try await Task.sleep(for: .milliseconds(100)) + + // Verify notification was received + let updatedURIs = await notificationTracker.resourceUpdatedURIs + #expect(updatedURIs.count == 1) + #expect(updatedURIs.first == "file:///watched.txt") + } +} + +// MARK: - Test Helpers + +/// Thread-safe tracker for subscription requests +private actor SubscriptionTracker { + private(set) var subscribedURIs: [String] = [] + private(set) var unsubscribedURIs: [String] = [] + + func recordSubscribe(uri: String) { + subscribedURIs.append(uri) + } + + func recordUnsubscribe(uri: String) { + unsubscribedURIs.append(uri) + } +} + +/// Thread-safe tracker for notifications +private actor NotificationTracker { + private(set) var resourceUpdatedURIs: [String] = [] + + func recordResourceUpdated(uri: String) { + resourceUpdatedURIs.append(uri) + } +} diff --git a/Tests/MCPTests/ResourceTests.swift b/Tests/MCPTests/ResourceTests.swift index 54036327..b75fdbdb 100644 --- a/Tests/MCPTests/ResourceTests.swift +++ b/Tests/MCPTests/ResourceTests.swift @@ -12,14 +12,14 @@ struct ResourceTests { uri: "file://test.txt", description: "A test resource", mimeType: "text/plain", - metadata: ["key": "value"] + _meta: ["key": "value"] ) #expect(resource.name == "test_resource") #expect(resource.uri == "file://test.txt") #expect(resource.description == "A test resource") #expect(resource.mimeType == "text/plain") - #expect(resource.metadata?["key"] == "value") + #expect(resource._meta?["key"] == "value") } @Test("Resource encoding and decoding") @@ -29,7 +29,7 @@ struct ResourceTests { uri: "file://test.txt", description: "Test resource description", mimeType: "text/plain", - metadata: ["key1": "value1", "key2": "value2"] + _meta: ["key1": "value1", "key2": "value2"] ) let encoder = JSONEncoder() @@ -42,7 +42,7 @@ struct ResourceTests { #expect(decoded.uri == resource.uri) #expect(decoded.description == resource.description) #expect(decoded.mimeType == resource.mimeType) - #expect(decoded.metadata == resource.metadata) + #expect(decoded._meta == resource._meta) } @Test("Resource.Content text initialization and encoding") @@ -152,6 +152,14 @@ struct ResourceTests { func testResourceSubscribeParameters() throws { let params = ResourceSubscribe.Parameters(uri: "file://test.txt") #expect(params.uri == "file://test.txt") + #expect(ResourceSubscribe.name == "resources/subscribe") + } + + @Test("ResourceUnsubscribe parameters validation") + func testResourceUnsubscribeParameters() throws { + let params = ResourceUnsubscribe.Parameters(uri: "file://test.txt") + #expect(params.uri == "file://test.txt") + #expect(ResourceUnsubscribe.name == "resources/unsubscribe") } @Test("ResourceUpdatedNotification parameters validation") @@ -165,4 +173,794 @@ struct ResourceTests { func testResourceListChangedNotification() throws { #expect(ResourceListChangedNotification.name == "notifications/resources/list_changed") } + + // MARK: - MIME Type Parameter Tests (RFC 2045) + + /// Tests for MIME types with parameters as specified in RFC 2045. + /// Based on: Python SDK `tests/issues/test_1754_mime_type_parameters.py` + @Test("Resource with MIME type parameters (RFC 2045)") + func testMimeTypeWithParameters() throws { + // MIME types with parameters should be accepted per RFC 2045 + let resource = Resource( + name: "widget", + uri: "ui://widget", + mimeType: "text/html;profile=mcp-app" + ) + + #expect(resource.mimeType == "text/html;profile=mcp-app") + + // Verify encoding/decoding preserves the full MIME type + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(resource) + let decoded = try decoder.decode(Resource.self, from: data) + + #expect(decoded.mimeType == "text/html;profile=mcp-app") + } + + @Test("Resource with MIME type parameters and space after semicolon") + func testMimeTypeWithParametersAndSpace() throws { + let resource = Resource( + name: "data", + uri: "data://json", + mimeType: "application/json; charset=utf-8" + ) + + #expect(resource.mimeType == "application/json; charset=utf-8") + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(resource) + let decoded = try decoder.decode(Resource.self, from: data) + + #expect(decoded.mimeType == "application/json; charset=utf-8") + } + + @Test("Resource with multiple MIME type parameters") + func testMimeTypeWithMultipleParameters() throws { + let resource = Resource( + name: "multi", + uri: "data://multi", + mimeType: "text/plain; charset=utf-8; format=fixed" + ) + + #expect(resource.mimeType == "text/plain; charset=utf-8; format=fixed") + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(resource) + let decoded = try decoder.decode(Resource.self, from: data) + + #expect(decoded.mimeType == "text/plain; charset=utf-8; format=fixed") + } + + @Test("Resource.Content preserves MIME type with parameters") + func testResourceContentPreservesMimeTypeWithParameters() throws { + let content = Resource.Content.text( + "Hello MCP-UI", + uri: "ui://my-widget", + mimeType: "text/html;profile=mcp-app" + ) + + #expect(content.mimeType == "text/html;profile=mcp-app") + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(content) + let decoded = try decoder.decode(Resource.Content.self, from: data) + + #expect(decoded.mimeType == "text/html;profile=mcp-app") + } + + // MARK: - Resource Template Tests + + /// Tests for Resource.Template encoding and decoding. + /// Based on: Python SDK `tests/issues/test_129_resource_templates.py` + @Test("Resource.Template initialization and encoding") + func testResourceTemplateInitialization() throws { + let template = Resource.Template( + uriTemplate: "greeting://{name}", + name: "greeting", + title: "Greeting Resource", + description: "Get a personalized greeting", + mimeType: "text/plain" + ) + + #expect(template.uriTemplate == "greeting://{name}") + #expect(template.name == "greeting") + #expect(template.title == "Greeting Resource") + #expect(template.description == "Get a personalized greeting") + #expect(template.mimeType == "text/plain") + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(template) + let decoded = try decoder.decode(Resource.Template.self, from: data) + + #expect(decoded.uriTemplate == "greeting://{name}") + #expect(decoded.name == "greeting") + #expect(decoded.title == "Greeting Resource") + #expect(decoded.description == "Get a personalized greeting") + #expect(decoded.mimeType == "text/plain") + } + + @Test("Resource.Template with multiple URI parameters") + func testResourceTemplateMultipleParams() throws { + let template = Resource.Template( + uriTemplate: "users://{user_id}/posts/{post_id}", + name: "user_post", + description: "User post resource" + ) + + #expect(template.uriTemplate == "users://{user_id}/posts/{post_id}") + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(template) + let decoded = try decoder.decode(Resource.Template.self, from: data) + + #expect(decoded.uriTemplate == "users://{user_id}/posts/{post_id}") + } + + @Test("Resource.Template with annotations and metadata") + func testResourceTemplateWithAnnotations() throws { + let template = Resource.Template( + uriTemplate: "file:///{path}", + name: "file", + annotations: Annotations(audience: [.user], priority: 0.8), + _meta: ["custom": "value"] + ) + + #expect(template.uriTemplate == "file:///{path}") + #expect(template.annotations?.audience == [.user]) + #expect(template.annotations?.priority == 0.8) + #expect(template._meta?["custom"] == "value") + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(template) + let decoded = try decoder.decode(Resource.Template.self, from: data) + + #expect(decoded.annotations?.audience == [.user]) + #expect(decoded.annotations?.priority == 0.8) + #expect(decoded._meta?["custom"] == "value") + } + + @Test("ListResourceTemplates parameters validation") + func testListResourceTemplatesParameters() throws { + let params = ListResourceTemplates.Parameters(cursor: "template_page_2") + #expect(params.cursor == "template_page_2") + + let emptyParams = ListResourceTemplates.Parameters() + #expect(emptyParams.cursor == nil) + } + + @Test("ListResourceTemplates result encoding/decoding") + func testListResourceTemplatesResult() throws { + let templates = [ + Resource.Template( + uriTemplate: "greeting://{name}", + name: "greeting", + description: "Get a personalized greeting" + ), + Resource.Template( + uriTemplate: "users://{user_id}/profile", + name: "user_profile", + description: "User profile resource" + ), + ] + + let result = ListResourceTemplates.Result( + templates: templates, + nextCursor: "next_template_page" + ) + + #expect(result.templates.count == 2) + #expect(result.nextCursor == "next_template_page") + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(result) + let decoded = try decoder.decode(ListResourceTemplates.Result.self, from: data) + + #expect(decoded.templates.count == 2) + #expect(decoded.templates[0].uriTemplate == "greeting://{name}") + #expect(decoded.templates[1].uriTemplate == "users://{user_id}/profile") + #expect(decoded.nextCursor == "next_template_page") + } + + @Test("ListResourceTemplates result uses 'resourceTemplates' key in JSON") + func testListResourceTemplatesUsesCorrectJsonKey() throws { + let result = ListResourceTemplates.Result( + templates: [ + Resource.Template(uriTemplate: "test://{id}", name: "test") + ] + ) + + let encoder = JSONEncoder() + let data = try encoder.encode(result) + let jsonString = String(data: data, encoding: .utf8)! + + // Verify the JSON uses "resourceTemplates" as the key (per MCP spec) + #expect(jsonString.contains("resourceTemplates")) + #expect(!jsonString.contains("\"templates\"")) + } + + // MARK: - ResourceLink Tests + + @Test("ResourceLink initialization and encoding") + func testResourceLinkInitialization() throws { + let link = ResourceLink( + name: "example_file", + title: "Example File", + uri: "file:///example.txt", + description: "An example resource link", + mimeType: "text/plain", + size: 1024 + ) + + #expect(link.name == "example_file") + #expect(link.title == "Example File") + #expect(link.uri == "file:///example.txt") + #expect(link.description == "An example resource link") + #expect(link.mimeType == "text/plain") + #expect(link.size == 1024) + } + + @Test("ResourceLink encodes with 'resource_link' type") + func testResourceLinkEncodesWithType() throws { + let link = ResourceLink( + name: "test", + uri: "file:///test.txt" + ) + + let encoder = JSONEncoder() + let data = try encoder.encode(link) + let jsonString = String(data: data, encoding: .utf8)! + + #expect(jsonString.contains("\"type\":\"resource_link\"")) + } + + @Test("ResourceLink decoding validates type field") + func testResourceLinkDecodingValidatesType() throws { + // Valid resource_link type + let validJson = """ + {"type":"resource_link","name":"test","uri":"file:///test.txt"} + """ + let validData = validJson.data(using: .utf8)! + let decoder = JSONDecoder() + + let decoded = try decoder.decode(ResourceLink.self, from: validData) + #expect(decoded.name == "test") + #expect(decoded.uri == "file:///test.txt") + + // Invalid type should throw + let invalidJson = """ + {"type":"wrong_type","name":"test","uri":"file:///test.txt"} + """ + let invalidData = invalidJson.data(using: .utf8)! + + #expect(throws: DecodingError.self) { + _ = try decoder.decode(ResourceLink.self, from: invalidData) + } + } + + @Test("ResourceLink decodes without type field (backward compatibility)") + func testResourceLinkDecodesWithoutType() throws { + // Type field is optional for backward compatibility + let json = """ + {"name":"test","uri":"file:///test.txt"} + """ + let data = json.data(using: .utf8)! + let decoder = JSONDecoder() + + let decoded = try decoder.decode(ResourceLink.self, from: data) + #expect(decoded.name == "test") + #expect(decoded.uri == "file:///test.txt") + } + + @Test("ResourceLink encoding/decoding roundtrip with all optional fields") + func testResourceLinkRoundtripAllFields() throws { + let link = ResourceLink( + name: "complete_resource", + title: "Complete Resource Title", + uri: "file:///complete.txt", + description: "A complete resource link", + mimeType: "application/json; charset=utf-8", + size: 4096, + annotations: Annotations(audience: [.user, .assistant], priority: 0.75), + _meta: ["version": "2.0", "author": "test"] + ) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(link) + let decoded = try decoder.decode(ResourceLink.self, from: data) + + #expect(decoded.name == link.name) + #expect(decoded.title == link.title) + #expect(decoded.uri == link.uri) + #expect(decoded.description == link.description) + #expect(decoded.mimeType == link.mimeType) + #expect(decoded.size == link.size) + #expect(decoded.annotations?.audience == link.annotations?.audience) + #expect(decoded.annotations?.priority == link.annotations?.priority) + #expect(decoded._meta?["version"] == "2.0") + #expect(decoded._meta?["author"] == "test") + } + + @Test("ResourceLink decoding fails when required 'name' field is missing") + func testResourceLinkDecodingFailsMissingName() throws { + // Missing required 'name' field + let json = """ + {"type":"resource_link","uri":"file:///test.txt"} + """ + let data = json.data(using: .utf8)! + let decoder = JSONDecoder() + + #expect(throws: DecodingError.self) { + _ = try decoder.decode(ResourceLink.self, from: data) + } + } + + @Test("ResourceLink decoding fails when required 'uri' field is missing") + func testResourceLinkDecodingFailsMissingUri() throws { + // Missing required 'uri' field + let json = """ + {"type":"resource_link","name":"test"} + """ + let data = json.data(using: .utf8)! + let decoder = JSONDecoder() + + #expect(throws: DecodingError.self) { + _ = try decoder.decode(ResourceLink.self, from: data) + } + } + + // MARK: - Pagination Encoding Tests + + @Test("ListResources pagination cursor encoding roundtrip") + func testPaginationCursorEncodingRoundtrip() throws { + // Test that cursor values are properly encoded and decoded + let cursor = "page_2_cursor_abc123" + let params = ListResources.Parameters(cursor: cursor) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(params) + let decoded = try decoder.decode(ListResources.Parameters.self, from: data) + + #expect(decoded.cursor == cursor) + } + + @Test("ListResources result nextCursor encoding roundtrip") + func testPaginationNextCursorEncodingRoundtrip() throws { + let resources = [ + Resource(name: "resource1", uri: "file://test1.txt"), + Resource(name: "resource2", uri: "file://test2.txt"), + ] + + let result = ListResources.Result( + resources: resources, + nextCursor: "next_page_cursor_xyz789" + ) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(result) + let decoded = try decoder.decode(ListResources.Result.self, from: data) + + #expect(decoded.resources.count == 2) + #expect(decoded.nextCursor == "next_page_cursor_xyz789") + } + + @Test("Pagination result with no more pages") + func testPaginationResultNoMorePages() throws { + let resources = [ + Resource(name: "last_resource", uri: "file://last.txt") + ] + + let result = ListResources.Result(resources: resources, nextCursor: nil) + + #expect(result.nextCursor == nil) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(result) + let decoded = try decoder.decode(ListResources.Result.self, from: data) + + #expect(decoded.nextCursor == nil) + } + + // MARK: - Resource with All Fields Tests + + @Test("Resource with all optional fields") + func testResourceWithAllFields() throws { + let resource = Resource( + name: "complete_resource", + title: "Complete Resource Title", + uri: "file:///complete.txt", + description: "A resource with all fields populated", + mimeType: "application/json; charset=utf-8", + size: 2048, + annotations: Annotations(audience: [.user, .assistant], priority: 0.9), + _meta: ["version": "1.0", "author": "test"] + ) + + #expect(resource.name == "complete_resource") + #expect(resource.title == "Complete Resource Title") + #expect(resource.uri == "file:///complete.txt") + #expect(resource.description == "A resource with all fields populated") + #expect(resource.mimeType == "application/json; charset=utf-8") + #expect(resource.size == 2048) + #expect(resource.annotations?.audience == [.user, .assistant]) + #expect(resource.annotations?.priority == 0.9) + #expect(resource._meta?["version"] == "1.0") + #expect(resource._meta?["author"] == "test") + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(resource) + let decoded = try decoder.decode(Resource.self, from: data) + + #expect(decoded.name == resource.name) + #expect(decoded.title == resource.title) + #expect(decoded.uri == resource.uri) + #expect(decoded.description == resource.description) + #expect(decoded.mimeType == resource.mimeType) + #expect(decoded.size == resource.size) + #expect(decoded.annotations?.audience == resource.annotations?.audience) + #expect(decoded.annotations?.priority == resource.annotations?.priority) + } + + // MARK: - Integration Tests + + @Test("Paginated resource listing") + func testPaginatedResourceListing() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + // Track pagination state + let paginationState = PaginationState() + + let server = Server( + name: "PaginatedResourceServer", + version: "1.0.0", + capabilities: .init(resources: .init()) + ) + + // Handler returns paginated resources + await server.withRequestHandler(ListResources.self) { [paginationState] params, _ in + let cursor = params.cursor + return await paginationState.getPage(cursor: cursor) + } + + try await server.start(transport: serverTransport) + + let client = Client(name: "PaginationClient", version: "1.0.0") + _ = try await client.connect(transport: clientTransport) + + // First page (no cursor) + let page1 = try await client.listResources() + #expect(page1.resources.count == 3) + #expect(page1.resources[0].name == "resource_1") + #expect(page1.resources[1].name == "resource_2") + #expect(page1.resources[2].name == "resource_3") + #expect(page1.nextCursor == "page_2") + + // Second page (with cursor) + let page2 = try await client.listResources(cursor: page1.nextCursor) + #expect(page2.resources.count == 3) + #expect(page2.resources[0].name == "resource_4") + #expect(page2.resources[1].name == "resource_5") + #expect(page2.resources[2].name == "resource_6") + #expect(page2.nextCursor == "page_3") + + // Third page (last page) + let page3 = try await client.listResources(cursor: page2.nextCursor) + #expect(page3.resources.count == 2) + #expect(page3.resources[0].name == "resource_7") + #expect(page3.resources[1].name == "resource_8") + #expect(page3.nextCursor == nil) + + await client.disconnect() + await server.stop() + } + + @Test("Resource template listing and reading") + func testResourceTemplateListingAndReading() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "TemplateServer", + version: "1.0.0", + capabilities: .init(resources: .init()) + ) + + // Handler for listing resource templates + await server.withRequestHandler(ListResourceTemplates.self) { _, _ in + ListResourceTemplates.Result(templates: [ + Resource.Template( + uriTemplate: "greeting://{name}", + name: "greeting", + description: "Get a personalized greeting" + ), + Resource.Template( + uriTemplate: "users://{user_id}/profile", + name: "user_profile", + description: "User profile resource" + ), + ]) + } + + // Handler for reading resources (handles templated URIs) + await server.withRequestHandler(ReadResource.self) { params, _ in + let uri = params.uri + + if uri.hasPrefix("greeting://") { + let name = String(uri.dropFirst("greeting://".count)) + return ReadResource.Result(contents: [ + .text("Hello, \(name)!", uri: uri, mimeType: "text/plain") + ]) + } else if uri.hasPrefix("users://") && uri.hasSuffix("/profile") { + let userId = uri + .replacingOccurrences(of: "users://", with: "") + .replacingOccurrences(of: "/profile", with: "") + return ReadResource.Result(contents: [ + .text("Profile for user \(userId)", uri: uri, mimeType: "text/plain") + ]) + } + + throw MCPError.invalidParams("Unknown resource: \(uri)") + } + + // Handler for listing resources (required for resources capability) + await server.withRequestHandler(ListResources.self) { _, _ in + ListResources.Result(resources: []) + } + + try await server.start(transport: serverTransport) + + let client = Client(name: "TemplateClient", version: "1.0.0") + _ = try await client.connect(transport: clientTransport) + + // List templates + let templates = try await client.listResourceTemplates() + #expect(templates.templates.count == 2) + #expect(templates.templates[0].uriTemplate == "greeting://{name}") + #expect(templates.templates[1].uriTemplate == "users://{user_id}/profile") + + // Read a resource using a templated URI + let greetingContent = try await client.readResource(uri: "greeting://World") + #expect(greetingContent.count == 1) + #expect(greetingContent[0].text == "Hello, World!") + #expect(greetingContent[0].mimeType == "text/plain") + + // Read another templated resource + let profileContent = try await client.readResource(uri: "users://123/profile") + #expect(profileContent.count == 1) + #expect(profileContent[0].text == "Profile for user 123") + + await client.disconnect() + await server.stop() + } + + @Test("Resource with MIME type parameters in integration") + func testResourceWithMimeTypeParametersIntegration() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "MimeTypeServer", + version: "1.0.0", + capabilities: .init(resources: .init()) + ) + + await server.withRequestHandler(ListResources.self) { _, _ in + ListResources.Result(resources: [ + Resource( + name: "widget", + uri: "ui://widget", + mimeType: "text/html;profile=mcp-app" + ), + Resource( + name: "data", + uri: "data://json", + mimeType: "application/json; charset=utf-8" + ), + ]) + } + + await server.withRequestHandler(ReadResource.self) { params, _ in + if params.uri == "ui://widget" { + return ReadResource.Result(contents: [ + .text( + "Hello MCP-UI", + uri: params.uri, + mimeType: "text/html;profile=mcp-app" + ) + ]) + } + throw MCPError.invalidParams("Unknown resource") + } + + try await server.start(transport: serverTransport) + + let client = Client(name: "MimeTypeClient", version: "1.0.0") + _ = try await client.connect(transport: clientTransport) + + // List resources and verify MIME types are preserved + let resources = try await client.listResources() + #expect(resources.resources.count == 2) + #expect(resources.resources[0].mimeType == "text/html;profile=mcp-app") + #expect(resources.resources[1].mimeType == "application/json; charset=utf-8") + + // Read resource and verify MIME type is preserved + let content = try await client.readResource(uri: "ui://widget") + #expect(content.count == 1) + #expect(content[0].mimeType == "text/html;profile=mcp-app") + + await client.disconnect() + await server.stop() + } + + @Test("Paginated resource template listing") + func testPaginatedResourceTemplateListing() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "PaginatedTemplateServer", + version: "1.0.0", + capabilities: .init(resources: .init()) + ) + + await server.withRequestHandler(ListResourceTemplates.self) { params, _ in + if params.cursor == nil { + return ListResourceTemplates.Result( + templates: [ + Resource.Template(uriTemplate: "template://{id1}", name: "template1"), + Resource.Template(uriTemplate: "template://{id2}", name: "template2"), + ], + nextCursor: "page_2" + ) + } else if params.cursor == "page_2" { + return ListResourceTemplates.Result( + templates: [ + Resource.Template(uriTemplate: "template://{id3}", name: "template3") + ], + nextCursor: nil + ) + } + return ListResourceTemplates.Result(templates: []) + } + + await server.withRequestHandler(ListResources.self) { _, _ in + ListResources.Result(resources: []) + } + + try await server.start(transport: serverTransport) + + let client = Client(name: "PaginatedTemplateClient", version: "1.0.0") + _ = try await client.connect(transport: clientTransport) + + // First page + let page1 = try await client.listResourceTemplates() + #expect(page1.templates.count == 2) + #expect(page1.nextCursor == "page_2") + + // Second page + let page2 = try await client.listResourceTemplates(cursor: page1.nextCursor) + #expect(page2.templates.count == 1) + #expect(page2.nextCursor == nil) + + await client.disconnect() + await server.stop() + } + + @Test("Empty resource listing result") + func testEmptyResourceListingResult() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "EmptyResourceServer", + version: "1.0.0", + capabilities: .init(resources: .init()) + ) + + // Handler returns empty resource list + await server.withRequestHandler(ListResources.self) { _, _ in + ListResources.Result(resources: []) + } + + try await server.start(transport: serverTransport) + + let client = Client(name: "EmptyResourceClient", version: "1.0.0") + _ = try await client.connect(transport: clientTransport) + + // Verify empty list is handled correctly + let result = try await client.listResources() + #expect(result.resources.isEmpty) + #expect(result.nextCursor == nil) + + await client.disconnect() + await server.stop() + } + + @Test("Empty resource template listing result") + func testEmptyResourceTemplateListingResult() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "EmptyTemplateServer", + version: "1.0.0", + capabilities: .init(resources: .init()) + ) + + // Handler returns empty template list + await server.withRequestHandler(ListResourceTemplates.self) { _, _ in + ListResourceTemplates.Result(templates: []) + } + + await server.withRequestHandler(ListResources.self) { _, _ in + ListResources.Result(resources: []) + } + + try await server.start(transport: serverTransport) + + let client = Client(name: "EmptyTemplateClient", version: "1.0.0") + _ = try await client.connect(transport: clientTransport) + + // Verify empty list is handled correctly + let result = try await client.listResourceTemplates() + #expect(result.templates.isEmpty) + #expect(result.nextCursor == nil) + + await client.disconnect() + await server.stop() + } +} + +// MARK: - Test Helpers + +/// Actor to track pagination state for tests +private actor PaginationState { + private let allResources: [Resource] = (1...8).map { i in + Resource(name: "resource_\(i)", uri: "file:///resource_\(i).txt") + } + private let pageSize = 3 + + func getPage(cursor: String?) -> ListResources.Result { + let startIndex: Int + let nextCursor: String? + + switch cursor { + case nil: + startIndex = 0 + nextCursor = "page_2" + case "page_2": + startIndex = 3 + nextCursor = "page_3" + case "page_3": + startIndex = 6 + nextCursor = nil + default: + return ListResources.Result(resources: []) + } + + let endIndex = min(startIndex + pageSize, allResources.count) + let pageResources = Array(allResources[startIndex..(id: id, result: result) @@ -42,7 +42,7 @@ struct ResponseTests { @Test("Error response initialization and encoding") func testErrorResponse() throws { - let id: ID = "test-id" + let id: RequestId = "test-id" let error = MCPError.parseError(nil) let response = Response(id: id, error: error) @@ -53,10 +53,10 @@ struct ResponseTests { let decoded = try decoder.decode(Response.self, from: data) if case .failure(let decodedError) = decoded.result { - #expect(decodedError.code == -32700) - #expect( - decodedError.localizedDescription - == "Parse error: Invalid JSON: Parse error: Invalid JSON") + #expect(decodedError.code == ErrorCode.parseError) + // Roundtrip preserves the error: parseError(nil) encodes as "Invalid JSON", + // which decodes back to parseError(nil) since it matches the default message + #expect(decodedError.localizedDescription == "Parse error: Invalid JSON") } else { #expect(Bool(false), "Expected error result") } @@ -64,7 +64,7 @@ struct ResponseTests { @Test("Error response with detail") func testErrorResponseWithDetail() throws { - let id: ID = "test-id" + let id: RequestId = "test-id" let error = MCPError.parseError("Invalid syntax") let response = Response(id: id, error: error) @@ -75,7 +75,7 @@ struct ResponseTests { let decoded = try decoder.decode(Response.self, from: data) if case .failure(let decodedError) = decoded.result { - #expect(decodedError.code == -32700) + #expect(decodedError.code == ErrorCode.parseError) #expect( decodedError.localizedDescription == "Parse error: Invalid JSON: Invalid syntax") diff --git a/Tests/MCPTests/ResumabilityTests.swift b/Tests/MCPTests/ResumabilityTests.swift new file mode 100644 index 00000000..73d341ad --- /dev/null +++ b/Tests/MCPTests/ResumabilityTests.swift @@ -0,0 +1,684 @@ +import Foundation +import Testing + +@testable import MCP + +/// Tests for resumability support - verifying that clients can reconnect and resume +/// receiving events after disconnection using InMemoryEventStore integration with +/// HTTPServerTransport. +/// +/// These tests follow the TypeScript SDK patterns from: +/// - `packages/server/test/server/streamableHttp.test.ts` +/// +/// Note: These tests use protocol version 2024-11-05. The TypeScript SDK uses 2025-11-25. +/// TODO: Update tests when Swift SDK adds support for protocol version 2025-11-25. +/// +/// TypeScript test not yet implemented: +/// +/// `should resume long-running notifications with lastEventId` (taskResumability.test.ts:180) +/// +/// Rationale: This is a full end-to-end integration test that requires: +/// 1. A Client instance connected to a Server via HTTPClientTransport +/// 2. A tool handler that sends multiple progress notifications over time +/// 3. Ability to disconnect the client mid-stream and reconnect with lastEventId +/// 4. The `onresumptiontoken` callback on the client side +/// +/// The Swift SDK has all the building blocks, but testing this requires either: +/// - A real HTTP server running in tests (like TypeScript's node http.Server) +/// - Full client-server integration test infrastructure +/// +/// The server-side resumability is tested here; client-side reconnection with +/// resumption token is tested in HTTPClientTransportTests and ClientReconnectionTests. +@Suite("Resumability Tests") +struct ResumabilityTests { + + // MARK: - Test Helpers + + static let initializeMessage = TestPayloads.initializeRequest(id: "init-1", clientName: "test-client") + + /// Parses SSE data to extract event ID and data content + func parseSSEEvents(_ data: Data) -> [(id: String?, data: String)] { + guard let text = String(data: data, encoding: .utf8) else { return [] } + + var events: [(id: String?, data: String)] = [] + var currentId: String? + var currentData: [String] = [] + + for line in text.components(separatedBy: "\n") { + if line.hasPrefix("id: ") || line.hasPrefix("id:") { + let idValue = line.hasPrefix("id: ") ? String(line.dropFirst(4)) : String(line.dropFirst(3)) + currentId = idValue.trimmingCharacters(in: .whitespaces) + } else if line.hasPrefix("data: ") || line.hasPrefix("data:") { + let dataValue = line.hasPrefix("data: ") ? String(line.dropFirst(6)) : String(line.dropFirst(5)) + currentData.append(dataValue) + } else if line.isEmpty && !currentData.isEmpty { + // End of event + events.append((id: currentId, data: currentData.joined(separator: "\n"))) + currentId = nil + currentData = [] + } + } + + // Handle case where last event doesn't end with empty line + if !currentData.isEmpty { + events.append((id: currentId, data: currentData.joined(separator: "\n"))) + } + + return events + } + + /// Helper to read from stream with timeout + func readFromStream( + _ stream: AsyncThrowingStream, + maxChunks: Int = 1, + timeout: Duration = .seconds(2) + ) async throws -> Data { + var receivedData = Data() + + try await withThrowingTaskGroup(of: Data?.self) { group in + // Task to read from stream + group.addTask { + var data = Data() + var count = 0 + for try await chunk in stream { + data.append(chunk) + count += 1 + if count >= maxChunks { + break + } + } + return data + } + + // Timeout task + group.addTask { + try await Task.sleep(for: timeout) + return nil // Return nil on timeout + } + + // Wait for first to complete + if let result = try await group.next(), let data = result { + receivedData = data + } + group.cancelAll() + } + + return receivedData + } + + // MARK: - 1.1 Store and include event IDs in server SSE messages + + @Test("Store and include event IDs in server SSE messages") + func storeAndIncludeEventIds() async throws { + let eventStore = InMemoryEventStore() + let sessionId = "test-session-\(UUID().uuidString)" + + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { sessionId }, + eventStore: eventStore + ) + ) + try await transport.connect() + + // Initialize the transport + let initRequest = TestPayloads.postRequest(body: Self.initializeMessage) + let initResponse = await transport.handleRequest(initRequest) + #expect(initResponse.statusCode == 200) + #expect(initResponse.headers[HTTPHeader.sessionId] == sessionId) + + // Open a standalone SSE stream (GET request) + let getRequest = TestPayloads.getRequest(sessionId: sessionId) + let getResponse = await transport.handleRequest(getRequest) + + #expect(getResponse.statusCode == 200) + #expect(getResponse.headers[HTTPHeader.contentType] == "text/event-stream") + #expect(getResponse.stream != nil) + + guard let stream = getResponse.stream else { + Issue.record("Expected stream in response") + return + } + + // Send a notification through the transport (in a concurrent task) + // and read from the stream simultaneously + let notification = """ + {"jsonrpc":"2.0","method":"notifications/message","params":{"level":"info","data":"Test notification with event ID"}} + """ + + // Start reading task first + let readTask = Task { + try await readFromStream(stream, maxChunks: 1, timeout: .seconds(2)) + } + + // Give a small delay then send notification + try await Task.sleep(for: .milliseconds(50)) + try await transport.send(notification.data(using: .utf8)!) + + // Wait for read to complete + let receivedData = try await readTask.value + + // Parse the SSE events + let events = parseSSEEvents(receivedData) + + // Verify we got at least one event with an ID + #expect(!events.isEmpty, "Should have received at least one event") + + if let firstEvent = events.first { + #expect(firstEvent.id != nil, "Event should have an ID") + #expect(firstEvent.data.contains("notifications/message"), "Event should contain the notification") + + // Verify the event was stored + let eventCount = await eventStore.eventCount + #expect(eventCount > 0, "Event should be stored in event store") + } + } + + // MARK: - 1.2 Store and replay MCP server tool notifications + + @Test("Store and replay MCP server notifications") + func storeAndReplayServerNotifications() async throws { + let eventStore = InMemoryEventStore() + let sessionId = "test-session-\(UUID().uuidString)" + + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { sessionId }, + eventStore: eventStore + ) + ) + try await transport.connect() + + // Initialize + let initRequest = TestPayloads.postRequest(body: Self.initializeMessage) + _ = await transport.handleRequest(initRequest) + + // Open first SSE stream + let getRequest1 = TestPayloads.getRequest(sessionId: sessionId) + let getResponse1 = await transport.handleRequest(getRequest1) + #expect(getResponse1.statusCode == 200) + + guard let stream1 = getResponse1.stream else { + Issue.record("Expected stream in response") + return + } + + // Start reading and send first notification + let readTask1 = Task { + try await readFromStream(stream1, maxChunks: 1, timeout: .seconds(2)) + } + + try await Task.sleep(for: .milliseconds(50)) + + let notification1 = """ + {"jsonrpc":"2.0","method":"notifications/message","params":{"level":"info","data":"First notification"}} + """ + try await transport.send(notification1.data(using: .utf8)!) + + let receivedData1 = try await readTask1.value + let events1 = parseSSEEvents(receivedData1) + #expect(!events1.isEmpty, "Should have received first notification") + + guard let firstEventId = events1.first?.id else { + Issue.record("First event should have an ID") + return + } + + // Close the first stream (simulating disconnect) + await transport.closeStandaloneSSEStream() + + // Send second notification while "disconnected" + let notification2 = """ + {"jsonrpc":"2.0","method":"notifications/message","params":{"level":"info","data":"Second notification"}} + """ + try await transport.send(notification2.data(using: .utf8)!) + + // Reconnect with Last-Event-ID to get missed messages + let getRequest2 = TestPayloads.getRequest(sessionId: sessionId, lastEventId: firstEventId) + let getResponse2 = await transport.handleRequest(getRequest2) + + #expect(getResponse2.statusCode == 200) + + guard let stream2 = getResponse2.stream else { + Issue.record("Expected stream in reconnection response") + return + } + + // Read replayed notifications + let receivedData2 = try await readFromStream(stream2, maxChunks: 1, timeout: .seconds(2)) + let events2 = parseSSEEvents(receivedData2) + + // Verify we received the second notification that was sent after our stored eventId + let hasSecondNotification = events2.contains { event in + event.data.contains("Second notification") + } + #expect(hasSecondNotification, "Should have received the second notification on reconnect") + } + + // MARK: - 1.3 Store and replay multiple notifications + + @Test("Store and replay multiple notifications sent while client is disconnected") + func storeAndReplayMultipleNotifications() async throws { + let eventStore = InMemoryEventStore() + let sessionId = "test-session-\(UUID().uuidString)" + + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { sessionId }, + eventStore: eventStore + ) + ) + try await transport.connect() + + // Initialize + let initRequest = TestPayloads.postRequest(body: Self.initializeMessage) + _ = await transport.handleRequest(initRequest) + + // Open first SSE stream + let getRequest1 = TestPayloads.getRequest(sessionId: sessionId) + let getResponse1 = await transport.handleRequest(getRequest1) + #expect(getResponse1.statusCode == 200) + + guard let stream1 = getResponse1.stream else { + Issue.record("Expected stream in response") + return + } + + // Start reading and send initial notification + let readTask1 = Task { + try await readFromStream(stream1, maxChunks: 1, timeout: .seconds(2)) + } + + try await Task.sleep(for: .milliseconds(50)) + + let initialNotification = """ + {"jsonrpc":"2.0","method":"notifications/message","params":{"level":"info","data":"Initial notification"}} + """ + try await transport.send(initialNotification.data(using: .utf8)!) + + let receivedData1 = try await readTask1.value + let events1 = parseSSEEvents(receivedData1) + + guard let lastEventId = events1.first?.id else { + Issue.record("Initial event should have an ID") + return + } + + // Close the SSE stream (simulate disconnect) + await transport.closeStandaloneSSEStream() + + // Send MULTIPLE notifications while the client is disconnected + for i in 1...3 { + let notification = """ + {"jsonrpc":"2.0","method":"notifications/message","params":{"level":"info","data":"Missed notification \(i)"}} + """ + try await transport.send(notification.data(using: .utf8)!) + } + + // Reconnect with the Last-Event-ID to get all missed messages + let getRequest2 = TestPayloads.getRequest(sessionId: sessionId, lastEventId: lastEventId) + let getResponse2 = await transport.handleRequest(getRequest2) + + #expect(getResponse2.statusCode == 200) + + guard let stream2 = getResponse2.stream else { + Issue.record("Expected stream in reconnection response") + return + } + + // Read replayed notifications (expect 3 chunks) + let receivedData2 = try await readFromStream(stream2, maxChunks: 3, timeout: .seconds(3)) + let allText = String(data: receivedData2, encoding: .utf8) ?? "" + + // Verify we received ALL notifications that were sent while disconnected + #expect(allText.contains("Missed notification 1"), "Should have received missed notification 1") + #expect(allText.contains("Missed notification 2"), "Should have received missed notification 2") + #expect(allText.contains("Missed notification 3"), "Should have received missed notification 3") + } + + // MARK: - Event Store Integration + + @Test("Event store receives events with correct stream ID") + func eventStoreReceivesEventsWithCorrectStreamId() async throws { + let eventStore = InMemoryEventStore() + let sessionId = "test-session-\(UUID().uuidString)" + + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { sessionId }, + eventStore: eventStore + ) + ) + try await transport.connect() + + // Initialize + let initRequest = TestPayloads.postRequest(body: Self.initializeMessage) + _ = await transport.handleRequest(initRequest) + + // Open SSE stream + let getRequest = TestPayloads.getRequest(sessionId: sessionId) + let getResponse = await transport.handleRequest(getRequest) + #expect(getResponse.statusCode == 200) + + guard let stream = getResponse.stream else { + Issue.record("Expected stream in response") + return + } + + // Start reading task + let readTask = Task { + try await readFromStream(stream, maxChunks: 1, timeout: .seconds(2)) + } + + try await Task.sleep(for: .milliseconds(50)) + + // Send a notification + let notification = """ + {"jsonrpc":"2.0","method":"test/notification","params":{}} + """ + try await transport.send(notification.data(using: .utf8)!) + + // Wait for read to complete + _ = try await readTask.value + + // Verify events were stored + let eventCount = await eventStore.eventCount + #expect(eventCount >= 1, "At least one event should be stored") + } + + @Test("Replay returns correct stream ID") + func replayReturnsCorrectStreamId() async throws { + let eventStore = InMemoryEventStore() + + // Store some test events directly + let message1 = """ + {"jsonrpc":"2.0","method":"test","params":{"msg":"first"}} + """.data(using: .utf8)! + let eventId1 = try await eventStore.storeEvent(streamId: "stream-A", message: message1) + + let message2 = """ + {"jsonrpc":"2.0","method":"test","params":{"msg":"second"}} + """.data(using: .utf8)! + _ = try await eventStore.storeEvent(streamId: "stream-A", message: message2) + + // Replay events after the first one + actor MessageCollector { + var messages: [String] = [] + func add(_ msg: String) { messages.append(msg) } + func get() -> [String] { messages } + } + let collector = MessageCollector() + + let streamId = try await eventStore.replayEventsAfter(eventId1) { _, message in + if let text = String(data: message, encoding: .utf8) { + await collector.add(text) + } + } + + #expect(streamId == "stream-A") + let messages = await collector.get() + #expect(messages.count == 1) + #expect(messages.first?.contains("second") == true) + } + + // MARK: - Edge Cases + + @Test("Replay with unknown event ID returns error") + func replayWithUnknownEventIdReturnsError() async throws { + let eventStore = InMemoryEventStore() + let sessionId = "test-session-\(UUID().uuidString)" + + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { sessionId }, + eventStore: eventStore + ) + ) + try await transport.connect() + + // Initialize + let initRequest = TestPayloads.postRequest(body: Self.initializeMessage) + _ = await transport.handleRequest(initRequest) + + // Try to reconnect with an unknown event ID + let getRequest = TestPayloads.getRequest(sessionId: sessionId, lastEventId: "unknown-event-id") + let getResponse = await transport.handleRequest(getRequest) + + // Should return 400 Bad Request for unknown event ID + #expect(getResponse.statusCode == 400, "Should return 400 for unknown event ID") + } + + @Test("Transport without event store does not include event IDs") + func transportWithoutEventStoreDoesNotIncludeEventIds() async throws { + let sessionId = "test-session-\(UUID().uuidString)" + + // Transport without event store + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { sessionId }) + ) + try await transport.connect() + + // Initialize + let initRequest = TestPayloads.postRequest(body: Self.initializeMessage) + _ = await transport.handleRequest(initRequest) + + // Open SSE stream + let getRequest = TestPayloads.getRequest(sessionId: sessionId) + let getResponse = await transport.handleRequest(getRequest) + #expect(getResponse.statusCode == 200) + + guard let stream = getResponse.stream else { + Issue.record("Expected stream in response") + return + } + + // Start reading task + let readTask = Task { + try await readFromStream(stream, maxChunks: 1, timeout: .seconds(2)) + } + + try await Task.sleep(for: .milliseconds(50)) + + // Send a notification + let notification = """ + {"jsonrpc":"2.0","method":"test/notification","params":{}} + """ + try await transport.send(notification.data(using: .utf8)!) + + let receivedData = try await readTask.value + let events = parseSSEEvents(receivedData) + + // Events should NOT have IDs when there's no event store + if let firstEvent = events.first { + #expect(firstEvent.id == nil, "Events should not have IDs without event store") + } + } + + // MARK: - Priming Events After Replay + + /// Creates a GET request with protocol version >= 2025-11-25 to enable priming events + + @Test("Replay sends a new priming event after replayed events") + func replaySendsNewPrimingEventAfterReplayedEvents() async throws { + let eventStore = InMemoryEventStore() + let sessionId = "test-session-\(UUID().uuidString)" + + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { sessionId }, + eventStore: eventStore, + retryInterval: 5000 // Include retry to make priming event more visible + ) + ) + try await transport.connect() + + // Initialize with modern protocol version + let initMessage = TestPayloads.initializeRequest(id: "init-1", protocolVersion: Version.v2025_11_25, clientName: "test-client") + let initRequest = TestPayloads.postRequest(body: initMessage) + let initResponse = await transport.handleRequest(initRequest) + #expect(initResponse.statusCode == 200) + + // Open first SSE stream with modern protocol + let getRequest1 = TestPayloads.getRequest(sessionId: sessionId, protocolVersion: Version.v2025_11_25) + let getResponse1 = await transport.handleRequest(getRequest1) + #expect(getResponse1.statusCode == 200) + + guard let stream1 = getResponse1.stream else { + Issue.record("Expected stream in response") + return + } + + // Read the priming event + let readTask1 = Task { + try await readFromStream(stream1, maxChunks: 2, timeout: .seconds(2)) + } + + try await Task.sleep(for: .milliseconds(50)) + + // Send a notification + let notification = """ + {"jsonrpc":"2.0","method":"notifications/message","params":{"level":"info","data":"Test notification"}} + """ + try await transport.send(notification.data(using: .utf8)!) + + let receivedData1 = try await readTask1.value + let events1 = parseSSEEvents(receivedData1) + + // Should have priming event (empty data) followed by notification + #expect(events1.count >= 2, "Should have at least priming event + notification") + + // Find the priming event (empty data) and notification event + let primingEvent = events1.first { $0.data.isEmpty || $0.data.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty } + let notificationEvent = events1.first { $0.data.contains("notifications/message") } + + guard let firstPrimingEventId = primingEvent?.id else { + Issue.record("Priming event should have an ID") + return + } + + guard let notificationEventId = notificationEvent?.id else { + Issue.record("Notification event should have an ID") + return + } + + // Close the stream + await transport.closeStandaloneSSEStream() + + // Send another notification while disconnected + let notification2 = """ + {"jsonrpc":"2.0","method":"notifications/message","params":{"level":"info","data":"Second notification"}} + """ + try await transport.send(notification2.data(using: .utf8)!) + + // Reconnect with Last-Event-ID pointing to the notification + // This should replay the second notification AND send a new priming event + let getRequest2 = TestPayloads.getRequest(sessionId: sessionId, protocolVersion: Version.v2025_11_25, lastEventId: notificationEventId) + let getResponse2 = await transport.handleRequest(getRequest2) + + #expect(getResponse2.statusCode == 200) + + guard let stream2 = getResponse2.stream else { + Issue.record("Expected stream in reconnection response") + return + } + + // Read replayed events + new priming event + let receivedData2 = try await readFromStream(stream2, maxChunks: 2, timeout: .seconds(3)) + let events2 = parseSSEEvents(receivedData2) + + // Should have replayed notification AND a NEW priming event + let hasSecondNotification = events2.contains { $0.data.contains("Second notification") } + #expect(hasSecondNotification, "Should have replayed the second notification") + + // Find the new priming event (empty data with different ID than first) + let newPrimingEvent = events2.first { event in + (event.data.isEmpty || event.data.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty) + && event.id != nil && event.id != firstPrimingEventId + } + #expect(newPrimingEvent != nil, "Should have a NEW priming event after replay for resumability") + } + + @Test("Priming events (empty data) are skipped during replay") + func primingEventsAreSkippedDuringReplay() async throws { + // Per MCP spec: "replay messages that would have been sent after the last event ID" + // Priming events have empty data and are NOT messages - they should be skipped during replay. + // This aligns with Python SDK behavior which skips None messages during replay. + + let eventStore = InMemoryEventStore() + + // Store events directly: priming (empty), message, priming (empty), message + // This simulates an edge case where multiple priming events might exist + let primingEvent1 = try await eventStore.storeEvent(streamId: "stream-A", message: Data()) + let message1 = """ + {"jsonrpc":"2.0","method":"test","params":{"msg":"first"}} + """.data(using: .utf8)! + _ = try await eventStore.storeEvent(streamId: "stream-A", message: message1) + + // Hypothetical second priming event (shouldn't happen in practice, but test the safeguard) + _ = try await eventStore.storeEvent(streamId: "stream-A", message: Data()) + + let message2 = """ + {"jsonrpc":"2.0","method":"test","params":{"msg":"second"}} + """.data(using: .utf8)! + _ = try await eventStore.storeEvent(streamId: "stream-A", message: message2) + + // Replay events after the first priming event + actor MessageCollector { + var messages: [Data] = [] + func add(_ msg: Data) { messages.append(msg) } + func get() -> [Data] { messages } + } + let collector = MessageCollector() + + let streamId = try await eventStore.replayEventsAfter(primingEvent1) { _, message in + await collector.add(message) + } + + #expect(streamId == "stream-A") + let messages = await collector.get() + + // Should only get the two actual messages, not the priming events + #expect(messages.count == 2, "Should only replay actual messages, not priming events (empty data)") + + // Verify both messages are actual JSON-RPC messages, not empty + for message in messages { + #expect(!message.isEmpty, "Replayed message should not be empty (priming events should be skipped)") + let text = String(data: message, encoding: .utf8) ?? "" + #expect(text.contains("jsonrpc"), "Replayed message should be a JSON-RPC message") + } + } + + @Test("GET without Last-Event-ID opens fresh stream") + func getWithoutLastEventIdOpensFreshStream() async throws { + let eventStore = InMemoryEventStore() + let sessionId = "test-session-\(UUID().uuidString)" + + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { sessionId }, + eventStore: eventStore + ) + ) + try await transport.connect() + + // Initialize + let initRequest = TestPayloads.postRequest(body: Self.initializeMessage) + _ = await transport.handleRequest(initRequest) + + // Store some events in the event store manually + _ = try await eventStore.storeEvent( + streamId: "_GET_stream", + message: """ + {"jsonrpc":"2.0","method":"old/notification","params":{}} + """.data(using: .utf8)! + ) + + // Open SSE stream WITHOUT Last-Event-ID + let getRequest = TestPayloads.getRequest(sessionId: sessionId) // No lastEventId + let getResponse = await transport.handleRequest(getRequest) + + #expect(getResponse.statusCode == 200) + #expect(getResponse.stream != nil) + // Should open a fresh stream, not replay old events + } +} diff --git a/Tests/MCPTests/RootsTests.swift b/Tests/MCPTests/RootsTests.swift new file mode 100644 index 00000000..ade1061a --- /dev/null +++ b/Tests/MCPTests/RootsTests.swift @@ -0,0 +1,622 @@ +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("Roots Tests") +struct RootsTests { + @Test("Root encoding and decoding") + func testRootCoding() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let root = Root( + uri: "file:///home/user/projects/myproject", + name: "My Project" + ) + + let data = try encoder.encode(root) + let decoded = try decoder.decode(Root.self, from: data) + + #expect(decoded.uri == "file:///home/user/projects/myproject") + #expect(decoded.name == "My Project") + } + + @Test("Root with metadata encoding and decoding") + func testRootWithMetadataCoding() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let root = Root( + uri: "file:///workspace/repo", + name: "Repository", + _meta: ["version": "1.0", "type": "git"] + ) + + let data = try encoder.encode(root) + let decoded = try decoder.decode(Root.self, from: data) + + #expect(decoded.uri == "file:///workspace/repo") + #expect(decoded.name == "Repository") + #expect(decoded._meta?["version"]?.stringValue == "1.0") + #expect(decoded._meta?["type"]?.stringValue == "git") + } + + @Test("Root without optional fields") + func testRootWithoutOptionalFields() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let root = Root(uri: "file:///path/to/root") + + let data = try encoder.encode(root) + let decoded = try decoder.decode(Root.self, from: data) + + #expect(decoded.uri == "file:///path/to/root") + #expect(decoded.name == nil) + #expect(decoded._meta == nil) + } + + @Test("Root URI must start with file://") + func testRootURIPrecondition() { + // Valid URIs should work + _ = Root(uri: "file:///valid/path") + _ = Root(uri: "file:///") + _ = Root(uri: "file:///C:/Users/test") + + // Note: Invalid URIs will cause a precondition failure, + // which cannot be tested directly in Swift Testing. + // The precondition is enforced at runtime. + } + + @Test("Root decoding fails for invalid URI") + func testRootDecodingFailsForInvalidURI() throws { + let decoder = JSONDecoder() + + // http:// URI should fail + let httpJSON = """ + {"uri": "http://example.com/path", "name": "Invalid"} + """.data(using: .utf8)! + + #expect(throws: DecodingError.self) { + _ = try decoder.decode(Root.self, from: httpJSON) + } + + // https:// URI should fail + let httpsJSON = """ + {"uri": "https://example.com/path", "name": "Invalid"} + """.data(using: .utf8)! + + #expect(throws: DecodingError.self) { + _ = try decoder.decode(Root.self, from: httpsJSON) + } + + // No protocol should fail + let noProtocolJSON = """ + {"uri": "/path/to/file", "name": "Invalid"} + """.data(using: .utf8)! + + #expect(throws: DecodingError.self) { + _ = try decoder.decode(Root.self, from: noProtocolJSON) + } + } + + @Test("Root Hashable conformance") + func testRootHashable() { + let root1 = Root(uri: "file:///path/a", name: "A") + let root2 = Root(uri: "file:///path/a", name: "A") + let root3 = Root(uri: "file:///path/b", name: "B") + + #expect(root1 == root2) + #expect(root1 != root3) + + var set = Set() + set.insert(root1) + set.insert(root2) + #expect(set.count == 1) + } + + @Test("ListRoots request encoding") + func testListRootsRequestEncoding() throws { + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + + let request = ListRoots.request(id: .number(1)) + + let data = try encoder.encode(request) + let json = String(data: data, encoding: .utf8)! + + #expect(json.contains("\"method\":\"roots/list\"")) + #expect(json.contains("\"id\":1")) + #expect(json.contains("\"jsonrpc\":\"2.0\"")) + } + + @Test("ListRoots result encoding and decoding") + func testListRootsResultCoding() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let roots = [ + Root(uri: "file:///home/user/project1", name: "Project 1"), + Root(uri: "file:///home/user/project2", name: "Project 2"), + ] + + let result = ListRoots.Result(roots: roots) + + let data = try encoder.encode(result) + let decoded = try decoder.decode(ListRoots.Result.self, from: data) + + #expect(decoded.roots.count == 2) + #expect(decoded.roots[0].uri == "file:///home/user/project1") + #expect(decoded.roots[0].name == "Project 1") + #expect(decoded.roots[1].uri == "file:///home/user/project2") + #expect(decoded.roots[1].name == "Project 2") + } + + @Test("ListRoots result with metadata") + func testListRootsResultWithMetadata() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let result = ListRoots.Result( + roots: [Root(uri: "file:///path")], + _meta: ["cursor": "next-page-token"] + ) + + let data = try encoder.encode(result) + let decoded = try decoder.decode(ListRoots.Result.self, from: data) + + #expect(decoded.roots.count == 1) + #expect(decoded._meta?["cursor"]?.stringValue == "next-page-token") + } + + @Test("ListRoots result empty roots") + func testListRootsResultEmptyRoots() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let result = ListRoots.Result(roots: []) + + let data = try encoder.encode(result) + let decoded = try decoder.decode(ListRoots.Result.self, from: data) + + #expect(decoded.roots.isEmpty) + } + + @Test("RootsListChangedNotification name") + func testRootsListChangedNotificationName() { + #expect(RootsListChangedNotification.name == "notifications/roots/list_changed") + } + + @Test("RootsListChangedNotification encoding") + func testRootsListChangedNotificationEncoding() throws { + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + + let notification = RootsListChangedNotification.message(.init()) + + let data = try encoder.encode(notification) + let json = String(data: data, encoding: .utf8)! + + #expect(json.contains("\"method\":\"notifications/roots/list_changed\"")) + #expect(json.contains("\"jsonrpc\":\"2.0\"")) + } + + @Test("Client capabilities include roots") + func testClientCapabilitiesIncludeRoots() throws { + let capabilities = Client.Capabilities( + roots: .init(listChanged: true) + ) + + #expect(capabilities.roots != nil) + #expect(capabilities.roots?.listChanged == true) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(capabilities) + let decoded = try decoder.decode(Client.Capabilities.self, from: data) + + #expect(decoded.roots != nil) + #expect(decoded.roots?.listChanged == true) + } + + @Test("Client capabilities roots without listChanged") + func testClientCapabilitiesRootsWithoutListChanged() throws { + let capabilities = Client.Capabilities( + roots: .init() + ) + + #expect(capabilities.roots != nil) + #expect(capabilities.roots?.listChanged == 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.roots != nil) + #expect(decoded.roots?.listChanged == nil) + } + + @Test("Root requiredURIPrefix constant") + func testRootRequiredURIPrefix() { + #expect(Root.requiredURIPrefix == "file://") + } + + @Test("Root JSON format matches MCP spec") + func testRootJSONFormat() throws { + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + + let root = Root( + uri: "file:///home/user/projects/myproject", + name: "My Project" + ) + + let data = try encoder.encode(root) + let json = String(data: data, encoding: .utf8)! + + // Verify JSON structure matches MCP spec + #expect(json.contains("\"uri\":\"file:///home/user/projects/myproject\"")) + #expect(json.contains("\"name\":\"My Project\"")) + } + + @Test("ListRoots result decodes from TypeScript SDK format") + func testDecodesFromTypeScriptFormat() throws { + let decoder = JSONDecoder() + + let json = """ + { + "roots": [ + { + "uri": "file:///home/user/project", + "name": "My Project" + } + ] + } + """.data(using: .utf8)! + + let result = try decoder.decode(ListRoots.Result.self, from: json) + + #expect(result.roots.count == 1) + #expect(result.roots[0].uri == "file:///home/user/project") + #expect(result.roots[0].name == "My Project") + } + + @Test("ListRoots result decodes from Python SDK format") + func testDecodesFromPythonFormat() throws { + let decoder = JSONDecoder() + + // Python SDK may include _meta + let json = """ + { + "roots": [ + { + "uri": "file:///users/fake/test", + "name": "Test Root 1" + }, + { + "uri": "file:///users/fake/test/2", + "name": "Test Root 2" + } + ], + "_meta": {} + } + """.data(using: .utf8)! + + let result = try decoder.decode(ListRoots.Result.self, from: json) + + #expect(result.roots.count == 2) + #expect(result.roots[0].uri == "file:///users/fake/test") + #expect(result.roots[0].name == "Test Root 1") + #expect(result.roots[1].uri == "file:///users/fake/test/2") + #expect(result.roots[1].name == "Test Root 2") + } +} + +@Suite("Roots Integration Tests") +struct RootsIntegrationTests { + @Test( + .timeLimit(.minutes(1)) + ) + func testRootsCapabilitiesNegotiation() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger( + label: "mcp.test.roots", + 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 + ) + + let server = Server( + name: "RootsTestServer", + version: "1.0.0", + capabilities: .init() + ) + + // Client with roots capability + let client = Client( + name: "RootsTestClient", + version: "1.0" + ) + await client.setCapabilities(.init(roots: .init(listChanged: true))) + + // Register roots handler (required since we declared the capability) + await client.withRootsHandler { + [Root(uri: "file:///test/path")] + } + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Verify roots capability is working by attempting to list roots + // (if server couldn't see the capability, listRoots would throw) + let roots = try await server.listRoots() + #expect(roots.count == 1) + #expect(roots[0].uri == "file:///test/path") + + await server.stop() + await client.disconnect() + try? clientToServerRead.close() + try? clientToServerWrite.close() + try? serverToClientRead.close() + try? serverToClientWrite.close() + } + + @Test( + .timeLimit(.minutes(1)) + ) + func testServerListRootsFromClient() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger( + label: "mcp.test.roots.list", + 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 + ) + + let server = Server( + name: "RootsTestServer", + version: "1.0.0", + capabilities: .init() + ) + + // Client with roots capability and handler + let expectedRoots = [ + Root(uri: "file:///home/user/project1", name: "Project 1"), + Root(uri: "file:///home/user/project2", name: "Project 2"), + ] + + let client = Client( + name: "RootsTestClient", + version: "1.0" + ) + await client.setCapabilities(.init(roots: .init(listChanged: true))) + await client.withRootsHandler { + expectedRoots + } + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Server requests roots from client + let roots = try await server.listRoots() + + #expect(roots.count == 2) + #expect(roots[0].uri == "file:///home/user/project1") + #expect(roots[0].name == "Project 1") + #expect(roots[1].uri == "file:///home/user/project2") + #expect(roots[1].name == "Project 2") + + await server.stop() + await client.disconnect() + try? clientToServerRead.close() + try? clientToServerWrite.close() + try? serverToClientRead.close() + try? serverToClientWrite.close() + } + + @Test( + .timeLimit(.minutes(1)) + ) + func testServerListRootsFailsWithoutCapability() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger( + label: "mcp.test.roots.nocap", + 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 + ) + + let server = Server( + name: "RootsTestServer", + version: "1.0.0", + capabilities: .init() + ) + + // Client WITHOUT roots capability + let client = Client( + name: "RootsTestClient", + version: "1.0" + ) + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Server should fail to request roots since client doesn't have capability + await #expect(throws: MCPError.self) { + _ = try await server.listRoots() + } + + await server.stop() + await client.disconnect() + try? clientToServerRead.close() + try? clientToServerWrite.close() + try? serverToClientRead.close() + try? serverToClientWrite.close() + } + + @Test( + .timeLimit(.minutes(1)) + ) + func testClientSendRootsChanged() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger( + label: "mcp.test.roots.changed", + 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 + ) + + // Use a continuation to wait for notification + let notificationExpectation = AsyncStream.makeStream(of: Void.self) + + let server = Server( + name: "RootsTestServer", + version: "1.0.0", + capabilities: .init() + ) + await server.onNotification(RootsListChangedNotification.self) { _ in + notificationExpectation.continuation.yield() + notificationExpectation.continuation.finish() + } + + // Client with roots.listChanged capability + let client = Client( + name: "RootsTestClient", + version: "1.0" + ) + await client.setCapabilities(.init(roots: .init(listChanged: true))) + await client.withRootsHandler { + [Root(uri: "file:///path")] + } + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Client sends roots changed notification + try await client.sendRootsChanged() + + // Wait for notification to be processed (with timeout via test time limit) + var notificationReceived = false + for await _ in notificationExpectation.stream { + notificationReceived = true + } + + #expect(notificationReceived == true) + + await server.stop() + await client.disconnect() + try? clientToServerRead.close() + try? clientToServerWrite.close() + try? serverToClientRead.close() + try? serverToClientWrite.close() + } + + @Test( + .timeLimit(.minutes(1)) + ) + func testClientSendRootsChangedFailsWithoutCapability() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger( + label: "mcp.test.roots.changed.nocap", + 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 + ) + + let server = Server( + name: "RootsTestServer", + version: "1.0.0", + capabilities: .init() + ) + + // Client WITH roots capability but WITHOUT listChanged + let client = Client( + name: "RootsTestClient", + version: "1.0" + ) + await client.setCapabilities(.init(roots: .init())) // No listChanged + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Client should fail to send roots changed notification + await #expect(throws: MCPError.self) { + try await client.sendRootsChanged() + } + + await server.stop() + await client.disconnect() + try? clientToServerRead.close() + try? clientToServerWrite.close() + try? serverToClientRead.close() + try? serverToClientWrite.close() + } +} diff --git a/Tests/MCPTests/RoundtripTests.swift b/Tests/MCPTests/RoundtripTests.swift index 5cd3f865..58606d82 100644 --- a/Tests/MCPTests/RoundtripTests.swift +++ b/Tests/MCPTests/RoundtripTests.swift @@ -44,18 +44,22 @@ struct RoundtripTests { tools: .init() ) ) - await server.withMethodHandler(ListTools.self) { _ in + await server.withRequestHandler(ListTools.self) { _, _ in return ListTools.Result(tools: [ Tool( name: "add", description: "Adds two numbers together", inputSchema: [ - "a": ["type": "integer", "description": "The first number"], - "a": ["type": "integer", "description": "The second number"], + "type": "object", + "properties": [ + "a": ["type": "integer", "description": "The first number"], + "b": ["type": "integer", "description": "The second number"] + ], + "required": ["a", "b"] ]) ]) } - await server.withMethodHandler(CallTool.self) { request in + await server.withRequestHandler(CallTool.self) { request, _ in guard request.name == "add" else { return CallTool.Result(content: [.text("Invalid tool name")], isError: true) } @@ -71,7 +75,7 @@ struct RoundtripTests { } // Add resource handlers to server - await server.withMethodHandler(ListResources.self) { _ in + await server.withRequestHandler(ListResources.self) { _, _ in return ListResources.Result(resources: [ Resource( name: "Example Text", @@ -88,7 +92,7 @@ struct RoundtripTests { ]) } - await server.withMethodHandler(ReadResource.self) { request in + await server.withRequestHandler(ReadResource.self) { request, _ in guard request.uri == "test://example.txt" else { return ReadResource.Result(contents: [.text("Resource not found", uri: request.uri)] ) diff --git a/Tests/MCPTests/SamplingTests.swift b/Tests/MCPTests/SamplingTests.swift index 26129ac2..72159eed 100644 --- a/Tests/MCPTests/SamplingTests.swift +++ b/Tests/MCPTests/SamplingTests.swift @@ -26,7 +26,7 @@ struct SamplingTests { let decodedTextMessage = try decoder.decode(Sampling.Message.self, from: textData) #expect(decodedTextMessage.role == .user) - if case .text(let text) = decodedTextMessage.content { + if case .text(let text, _, _) = decodedTextMessage.content.first { #expect(text == "Hello, world!") } else { #expect(Bool(false), "Expected text content") @@ -40,12 +40,27 @@ struct SamplingTests { let decodedImageMessage = try decoder.decode(Sampling.Message.self, from: imageData) #expect(decodedImageMessage.role == .assistant) - if case .image(let data, let mimeType) = decodedImageMessage.content { + if case .image(let data, let mimeType, _, _) = decodedImageMessage.content.first { #expect(data == "base64imagedata") #expect(mimeType == "image/png") } else { #expect(Bool(false), "Expected image content") } + + // Test audio content + let audioMessage: Sampling.Message = .user( + .audio(data: "base64audiodata", mimeType: "audio/wav")) + + let audioData = try encoder.encode(audioMessage) + let decodedAudioMessage = try decoder.decode(Sampling.Message.self, from: audioData) + + #expect(decodedAudioMessage.role == .user) + if case .audio(let data, let mimeType, _, _) = decodedAudioMessage.content.first { + #expect(data == "base64audiodata") + #expect(mimeType == "audio/wav") + } else { + #expect(Bool(false), "Expected audio content") + } } @Test("ModelPreferences encoding and decoding") @@ -93,13 +108,33 @@ struct SamplingTests { let encoder = JSONEncoder() let decoder = JSONDecoder() - let reasons: [Sampling.StopReason] = [.endTurn, .stopSequence, .maxTokens] + // Test standard stop reasons + let reasons: [Sampling.StopReason] = [.endTurn, .stopSequence, .maxTokens, .toolUse] for reason in reasons { let data = try encoder.encode(reason) let decoded = try decoder.decode(Sampling.StopReason.self, from: data) #expect(decoded == reason) } + + // Test "refusal" stop reason (part of MCP spec) + let refusalReason = Sampling.StopReason(rawValue: "refusal") + let refusalData = try encoder.encode(refusalReason) + let decodedRefusal = try decoder.decode(Sampling.StopReason.self, from: refusalData) + #expect(decodedRefusal.rawValue == "refusal") + + // Test "other" stop reason (part of MCP spec) + let otherReason = Sampling.StopReason(rawValue: "other") + let otherData = try encoder.encode(otherReason) + let decodedOther = try decoder.decode(Sampling.StopReason.self, from: otherData) + #expect(decodedOther.rawValue == "other") + + // Test custom/provider-specific stop reason + let customReason = Sampling.StopReason(rawValue: "customProviderReason") + let customData = try encoder.encode(customReason) + let decodedCustom = try decoder.decode(Sampling.StopReason.self, from: customData) + #expect(decodedCustom == customReason) + #expect(decodedCustom.rawValue == "customProviderReason") } @Test("CreateMessage request parameters") @@ -145,7 +180,7 @@ struct SamplingTests { #expect(decoded.metadata?["provider"]?.stringValue == "test") } - @Test("CreateMessage result") + @Test("CreateMessage result (without tools)") func testCreateMessageResult() throws { let encoder = JSONEncoder() let decoder = JSONDecoder() @@ -164,13 +199,130 @@ struct SamplingTests { #expect(decoded.stopReason == .endTurn) #expect(decoded.role == .assistant) - if case .text(let text) = decoded.content { + // Content is now SamplingContent (single block), not an array + 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 result decodes array content (MCP spec compatibility)") + func testCreateMessageResultDecodesArrayContent() throws { + let decoder = JSONDecoder() + + // MCP spec allows content to be either single or array. + // Some clients may send content as array even for non-tool requests. + let jsonWithArrayContent = """ + { + "model": "claude-4-sonnet", + "stopReason": "endTurn", + "role": "assistant", + "content": [{"type": "text", "text": "Response from array format"}] + } + """.data(using: .utf8)! + + let decoded = try decoder.decode(CreateSamplingMessage.Result.self, from: jsonWithArrayContent) + + #expect(decoded.model == "claude-4-sonnet") + #expect(decoded.stopReason == .endTurn) + #expect(decoded.role == .assistant) + + if case .text(let text, _, _) = decoded.content { + #expect(text == "Response from array format") + } else { + #expect(Bool(false), "Expected text content") + } + } + + @Test("CreateMessage result decodes single content") + func testCreateMessageResultDecodesSingleContent() throws { + let decoder = JSONDecoder() + + // Single object content (common format) + let jsonWithSingleContent = """ + { + "model": "claude-4-sonnet", + "role": "assistant", + "content": {"type": "text", "text": "Response from single format"} + } + """.data(using: .utf8)! + + let decoded = try decoder.decode(CreateSamplingMessage.Result.self, from: jsonWithSingleContent) + + if case .text(let text, _, _) = decoded.content { + #expect(text == "Response from single format") + } else { + #expect(Bool(false), "Expected text content") + } + } + + @Test("CreateMessage result with tools") + func testCreateMessageResultWithTools() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + // Test with tool use content + let toolUse = ToolUseContent( + name: "get_weather", + id: "call-123", + input: ["city": "Paris"] + ) + let result = CreateSamplingMessageWithTools.Result( + model: "claude-4-sonnet", + stopReason: .toolUse, + role: .assistant, + content: .toolUse(toolUse) + ) + + let data = try encoder.encode(result) + let decoded = try decoder.decode(CreateSamplingMessageWithTools.Result.self, from: data) + + #expect(decoded.model == "claude-4-sonnet") + #expect(decoded.stopReason == .toolUse) + #expect(decoded.role == .assistant) + #expect(decoded.content.count == 1) + + if case .toolUse(let content) = decoded.content.first { + #expect(content.name == "get_weather") + #expect(content.id == "call-123") + } else { + #expect(Bool(false), "Expected tool use content") + } + } + + @Test("CreateMessage result with parallel tool calls") + func testCreateMessageResultWithParallelToolCalls() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + // Test with multiple tool use content (parallel calls) + let toolUse1 = ToolUseContent(name: "get_weather", id: "call-1", input: ["city": "Paris"]) + let toolUse2 = ToolUseContent(name: "get_time", id: "call-2", input: ["city": "Paris"]) + + let result = CreateSamplingMessageWithTools.Result( + model: "claude-4-sonnet", + stopReason: .toolUse, + role: .assistant, + content: [.toolUse(toolUse1), .toolUse(toolUse2)] + ) + + let data = try encoder.encode(result) + let decoded = try decoder.decode(CreateSamplingMessageWithTools.Result.self, from: data) + + #expect(decoded.model == "claude-4-sonnet") + #expect(decoded.stopReason == .toolUse) + #expect(decoded.content.count == 2) + + if case .toolUse(let content1) = decoded.content[0], + case .toolUse(let content2) = decoded.content[1] { + #expect(content1.name == "get_weather") + #expect(content2.name == "get_time") + } else { + #expect(Bool(false), "Expected tool use content") + } + } + @Test("CreateMessage request creation") func testCreateMessageRequest() throws { let messages: [Sampling.Message] = [ @@ -189,9 +341,9 @@ struct SamplingTests { #expect(request.params.maxTokens == 100) } - @Test("Server capabilities include sampling") - func testServerCapabilitiesIncludeSampling() throws { - let capabilities = Server.Capabilities( + @Test("Client capabilities include sampling") + func testClientCapabilitiesIncludeSampling() throws { + let capabilities = Client.Capabilities( sampling: .init() ) @@ -201,84 +353,79 @@ struct SamplingTests { let decoder = JSONDecoder() let data = try encoder.encode(capabilities) - let decoded = try decoder.decode(Server.Capabilities.self, from: data) + let decoded = try decoder.decode(Client.Capabilities.self, from: data) #expect(decoded.sampling != nil) } - @Test("Client capabilities include sampling") - func testClientCapabilitiesIncludeSampling() throws { + @Test("Client capabilities with sampling.tools") + func testClientCapabilitiesWithSamplingTools() throws { let capabilities = Client.Capabilities( - sampling: .init() + sampling: .init(tools: .init()) ) - #expect(capabilities.sampling != nil) + #expect(capabilities.sampling?.tools != nil) + #expect(capabilities.sampling?.context == nil) let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] let decoder = JSONDecoder() let data = try encoder.encode(capabilities) - let decoded = try decoder.decode(Client.Capabilities.self, from: data) + let json = String(data: data, encoding: .utf8)! - #expect(decoded.sampling != nil) + #expect(json.contains("\"tools\":{}")) + + let decoded = try decoder.decode(Client.Capabilities.self, from: data) + #expect(decoded.sampling?.tools != nil) + #expect(decoded.sampling?.context == nil) } - @Test("Client sampling handler registration") - func testClientSamplingHandlerRegistration() async throws { - let client = Client(name: "TestClient", version: "1.0") + @Test("Client capabilities with sampling.context") + func testClientCapabilitiesWithSamplingContext() throws { + let capabilities = Client.Capabilities( + sampling: .init(context: .init()) + ) + + #expect(capabilities.sampling?.context != nil) + #expect(capabilities.sampling?.tools == nil) - // 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") - ) - } + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + let decoder = JSONDecoder() + + let data = try encoder.encode(capabilities) + let json = String(data: data, encoding: .utf8)! + + #expect(json.contains("\"context\":{}")) - // Should return self for method chaining - #expect(handlerClient === client) + let decoded = try decoder.decode(Client.Capabilities.self, from: data) + #expect(decoded.sampling?.context != nil) + #expect(decoded.sampling?.tools == nil) } - @Test("Server sampling request method") - func testServerSamplingRequestMethod() async throws { - let transport = MockTransport() - let server = Server( - name: "TestServer", - version: "1.0", - capabilities: .init(sampling: .init()) + @Test("Client capabilities with sampling.tools and sampling.context") + func testClientCapabilitiesWithSamplingToolsAndContext() throws { + let capabilities = Client.Capabilities( + sampling: .init(context: .init(), tools: .init()) ) - try await server.start(transport: transport) + #expect(capabilities.sampling?.tools != nil) + #expect(capabilities.sampling?.context != nil) - // Test that server can attempt to request sampling - let messages: [Sampling.Message] = [ - .user("Test message") - ] + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + let decoder = JSONDecoder() - 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)") - } + let data = try encoder.encode(capabilities) + let json = String(data: data, encoding: .utf8)! - await server.stop() + #expect(json.contains("\"tools\":{}")) + #expect(json.contains("\"context\":{}")) + + let decoded = try decoder.decode(Client.Capabilities.self, from: data) + #expect(decoded.sampling?.tools != nil) + #expect(decoded.sampling?.context != nil) } @Test("Sampling message content JSON format") @@ -287,7 +434,7 @@ struct SamplingTests { encoder.outputFormatting = [.sortedKeys] // Test text content JSON format - let textContent: Sampling.Message.Content = .text("Hello") + let textContent: Sampling.Message.ContentBlock = .text("Hello") let textData = try encoder.encode(textContent) let textJSON = String(data: textData, encoding: .utf8)! @@ -295,7 +442,7 @@ struct SamplingTests { #expect(textJSON.contains("\"text\":\"Hello\"")) // Test image content JSON format - let imageContent: Sampling.Message.Content = .image( + let imageContent: Sampling.Message.ContentBlock = .image( data: "base64data", mimeType: "image/png") let imageData = try encoder.encode(imageContent) let imageJSON = String(data: imageData, encoding: .utf8)! @@ -303,6 +450,32 @@ struct SamplingTests { #expect(imageJSON.contains("\"type\":\"image\"")) #expect(imageJSON.contains("\"data\":\"base64data\"")) #expect(imageJSON.contains("\"mimeType\":\"image\\/png\"")) + + // Test audio content JSON format + let audioContent: Sampling.Message.ContentBlock = .audio( + data: "base64audiodata", mimeType: "audio/wav") + let audioData = try encoder.encode(audioContent) + let audioJSON = String(data: audioData, encoding: .utf8)! + + #expect(audioJSON.contains("\"type\":\"audio\"")) + #expect(audioJSON.contains("\"data\":\"base64audiodata\"")) + #expect(audioJSON.contains("\"mimeType\":\"audio\\/wav\"")) + } + + @Test("Backwards compatibility: Sampling.Message.Content alias") + func testContentTypeAliasBackwardsCompatibility() throws { + // Test that the deprecated Content type alias still works + let content: Sampling.Message.ContentBlock = .text("Hello") + + if case .text(let text, _, _) = content { + #expect(text == "Hello") + } else { + #expect(Bool(false), "Expected text content") + } + + // Verify it's the same type as ContentBlock + let block: Sampling.Message.ContentBlock = content + #expect(block == content) } @Test("UnitInterval in Sampling.ModelPreferences") @@ -335,7 +508,7 @@ struct SamplingTests { // Test user message factory method let userMessage: Sampling.Message = .user("Hello, world!") #expect(userMessage.role == .user) - if case .text(let text) = userMessage.content { + if case .text(let text, _, _) = userMessage.content.first { #expect(text == "Hello, world!") } else { #expect(Bool(false), "Expected text content") @@ -344,7 +517,7 @@ struct SamplingTests { // Test assistant message factory method let assistantMessage: Sampling.Message = .assistant("Hi there!") #expect(assistantMessage.role == .assistant) - if case .text(let text) = assistantMessage.content { + if case .text(let text, _, _) = assistantMessage.content.first { #expect(text == "Hi there!") } else { #expect(Bool(false), "Expected text content") @@ -354,7 +527,7 @@ struct SamplingTests { let imageMessage: Sampling.Message = .user( .image(data: "base64data", mimeType: "image/png")) #expect(imageMessage.role == .user) - if case .image(let data, let mimeType) = imageMessage.content { + if case .image(let data, let mimeType, _, _) = imageMessage.content.first { #expect(data == "base64data") #expect(mimeType == "image/png") } else { @@ -365,9 +538,9 @@ struct SamplingTests { @Test("Content ExpressibleByStringLiteral") func testContentExpressibleByStringLiteral() throws { // Test string literal assignment - let content: Sampling.Message.Content = "Hello from string literal" + let content: Sampling.Message.ContentBlock = "Hello from string literal" - if case .text(let text) = content { + if case .text(let text, _, _) = content { #expect(text == "Hello from string literal") } else { #expect(Bool(false), "Expected text content") @@ -375,7 +548,7 @@ struct SamplingTests { // Test in message creation let message: Sampling.Message = .user("Direct string literal") - if case .text(let text) = message.content { + if case .text(let text, _, _) = message.content.first { #expect(text == "Direct string literal") } else { #expect(Bool(false), "Expected text content") @@ -401,10 +574,10 @@ struct SamplingTests { let location = "San Francisco" // Test string interpolation - let content: Sampling.Message.Content = + let content: Sampling.Message.ContentBlock = "Hello \(userName), the temperature in \(location) is \(temperature)°F" - if case .text(let text) = content { + if case .text(let text, _, _) = content { #expect(text == "Hello Alice, the temperature in San Francisco is 72°F") } else { #expect(Bool(false), "Expected text content") @@ -413,7 +586,7 @@ struct SamplingTests { // Test in message creation with interpolation let message = Sampling.Message.user( "Welcome \(userName)! Today's weather in \(location) is \(temperature)°F") - if case .text(let text) = message.content { + if case .text(let text, _, _) = message.content.first { #expect(text == "Welcome Alice! Today's weather in San Francisco is 72°F") } else { #expect(Bool(false), "Expected text content") @@ -425,7 +598,7 @@ struct SamplingTests { let listMessage: Sampling.Message = .assistant( "You have \(count) items: \(items.joined(separator: ", "))") - if case .text(let text) = listMessage.content { + if case .text(let text, _, _) = listMessage.content.first { #expect(text == "You have 3 items: apples, bananas, oranges") } else { #expect(Bool(false), "Expected text content") @@ -442,7 +615,7 @@ struct SamplingTests { let userMessage: Sampling.Message = .user( "Hi, I'm \(customerName) and I have an issue with order \(orderNumber)") #expect(userMessage.role == .user) - if case .text(let text) = userMessage.content { + if case .text(let text, _, _) = userMessage.content.first { #expect(text == "Hi, I'm Bob and I have an issue with order ORD-12345") } else { #expect(Bool(false), "Expected text content") @@ -453,7 +626,7 @@ struct SamplingTests { "Hello \(customerName), I can help you with your \(issueType) issue for order \(orderNumber)" ) #expect(assistantMessage.role == .assistant) - if case .text(let text) = assistantMessage.content { + if case .text(let text, _, _) = assistantMessage.content.first { #expect( text == "Hello Bob, I can help you with your delivery delay issue for order ORD-12345" @@ -475,7 +648,7 @@ struct SamplingTests { #expect(conversation.count == 4) // Verify interpolated content - if case .text(let text) = conversation[2].content { + if case .text(let text, _, _) = conversation[2].content.first { #expect(text == "I have an issue with order ORD-12345 - it's a delivery delay") } else { #expect(Bool(false), "Expected text content") @@ -519,10 +692,10 @@ struct SamplingTests { #expect(mixedContent.count == 4) // Verify content types - if case .text = mixedContent[0].content, - case .image = mixedContent[1].content, - case .text = mixedContent[2].content, - case .text = mixedContent[3].content + if case .text = mixedContent[0].content.first, + case .image = mixedContent[1].content.first, + case .text = mixedContent[2].content.first, + case .text = mixedContent[3].content.first { // All content types are correct } else { @@ -543,6 +716,252 @@ struct SamplingTests { } } +@Suite("Sampling Message Validation Tests") +struct SamplingMessageValidationTests { + @Test("validateToolUseResultMessages passes for empty messages") + func testEmptyMessages() throws { + let messages: [Sampling.Message] = [] + #expect(throws: Never.self) { + try Sampling.Message.validateToolUseResultMessages(messages) + } + } + + @Test("validateToolUseResultMessages passes for simple text messages") + func testSimpleTextMessages() throws { + let messages: [Sampling.Message] = [ + .user("Hello"), + .assistant("Hi there!"), + ] + #expect(throws: Never.self) { + try Sampling.Message.validateToolUseResultMessages(messages) + } + } + + @Test("validateToolUseResultMessages passes for valid tool_use then tool_result") + func testValidToolUseAndResult() throws { + let toolUseContent = ToolUseContent( + name: "get_weather", + id: "tool-123", + input: ["city": "Paris"] + ) + let toolResultContent = ToolResultContent( + toolUseId: "tool-123", + content: [.text("Sunny, 72°F")] + ) + + let messages: [Sampling.Message] = [ + .user("What's the weather?"), + .assistant(.toolUse(toolUseContent)), + Sampling.Message(role: .user, content: [.toolResult(toolResultContent)]), + ] + + #expect(throws: Never.self) { + try Sampling.Message.validateToolUseResultMessages(messages) + } + } + + @Test("validateToolUseResultMessages fails when tool_result mixed with other content") + func testToolResultMixedWithOtherContent() throws { + let toolResultContent = ToolResultContent( + toolUseId: "tool-123", + content: [.text("Result")] + ) + + let messages: [Sampling.Message] = [ + .user("Hello"), + .assistant(.toolUse(ToolUseContent(name: "test", id: "tool-123", input: [:]))), + Sampling.Message(role: .user, content: [ + .text("Some text"), // Mixed with tool_result - invalid! + .toolResult(toolResultContent), + ]), + ] + + #expect(throws: MCPError.self) { + try Sampling.Message.validateToolUseResultMessages(messages) + } + } + + @Test("validateToolUseResultMessages fails when tool_result without preceding tool_use") + func testToolResultWithoutPrecedingToolUse() throws { + let toolResultContent = ToolResultContent( + toolUseId: "tool-123", + content: [.text("Result")] + ) + + let messages: [Sampling.Message] = [ + .user("Hello"), + .assistant("Let me help you"), // No tool_use here + Sampling.Message(role: .user, content: [.toolResult(toolResultContent)]), + ] + + #expect(throws: MCPError.self) { + try Sampling.Message.validateToolUseResultMessages(messages) + } + } + + @Test("validateToolUseResultMessages fails when tool_result is first message") + func testToolResultAsFirstMessage() throws { + let toolResultContent = ToolResultContent( + toolUseId: "tool-123", + content: [.text("Result")] + ) + + let messages: [Sampling.Message] = [ + Sampling.Message(role: .user, content: [.toolResult(toolResultContent)]), + ] + + #expect(throws: MCPError.self) { + try Sampling.Message.validateToolUseResultMessages(messages) + } + } + + @Test("validateToolUseResultMessages fails when tool IDs don't match") + func testMismatchedToolIds() throws { + let toolUseContent = ToolUseContent( + name: "get_weather", + id: "tool-123", + input: [:] + ) + let toolResultContent = ToolResultContent( + toolUseId: "tool-456", // Different ID! + content: [.text("Result")] + ) + + let messages: [Sampling.Message] = [ + .user("Hello"), + .assistant(.toolUse(toolUseContent)), + Sampling.Message(role: .user, content: [.toolResult(toolResultContent)]), + ] + + #expect(throws: MCPError.self) { + try Sampling.Message.validateToolUseResultMessages(messages) + } + } + + @Test("validateToolUseResultMessages passes with multiple matching tool_use/tool_result") + func testMultipleMatchingToolUseAndResults() throws { + let toolUse1 = ToolUseContent(name: "get_weather", id: "tool-1", input: [:]) + let toolUse2 = ToolUseContent(name: "get_time", id: "tool-2", input: [:]) + let toolResult1 = ToolResultContent(toolUseId: "tool-1", content: [.text("Sunny")]) + let toolResult2 = ToolResultContent(toolUseId: "tool-2", content: [.text("3pm")]) + + let messages: [Sampling.Message] = [ + .user("What's the weather and time?"), + Sampling.Message(role: .assistant, content: [ + .toolUse(toolUse1), + .toolUse(toolUse2), + ]), + Sampling.Message(role: .user, content: [ + .toolResult(toolResult1), + .toolResult(toolResult2), + ]), + ] + + #expect(throws: Never.self) { + try Sampling.Message.validateToolUseResultMessages(messages) + } + } + + @Test("validateToolUseResultMessages fails with partial tool_result match") + func testPartialToolResultMatch() throws { + let toolUse1 = ToolUseContent(name: "get_weather", id: "tool-1", input: [:]) + let toolUse2 = ToolUseContent(name: "get_time", id: "tool-2", input: [:]) + // Only providing result for tool-1, missing tool-2 + let toolResult1 = ToolResultContent(toolUseId: "tool-1", content: [.text("Sunny")]) + + let messages: [Sampling.Message] = [ + .user("What's the weather and time?"), + Sampling.Message(role: .assistant, content: [ + .toolUse(toolUse1), + .toolUse(toolUse2), + ]), + Sampling.Message(role: .user, content: [ + .toolResult(toolResult1), + ]), + ] + + #expect(throws: MCPError.self) { + try Sampling.Message.validateToolUseResultMessages(messages) + } + } + + @Test("validateToolUseResultMessages fails with extra tool_results not matching tool_use") + func testExtraToolResultsNotMatchingToolUse() throws { + // tool_use has [1], but tool_result has [1, 2] - extra result for non-existent tool + let toolUse = ToolUseContent(name: "get_weather", id: "tool-1", input: [:]) + let toolResult1 = ToolResultContent(toolUseId: "tool-1", content: [.text("Sunny")]) + let toolResult2 = ToolResultContent(toolUseId: "tool-2", content: [.text("Extra")]) + + let messages: [Sampling.Message] = [ + .user("What's the weather?"), + .assistant(.toolUse(toolUse)), + Sampling.Message(role: .user, content: [ + .toolResult(toolResult1), + .toolResult(toolResult2), // Extra result with no matching tool_use + ]), + ] + + #expect(throws: MCPError.self) { + try Sampling.Message.validateToolUseResultMessages(messages) + } + } + + @Test("validateToolUseResultMessages fails when text message follows tool_use") + func testTextMessageAfterToolUse() throws { + // After tool_use, you MUST provide tool_result - can't just send text + let toolUse = ToolUseContent(name: "get_weather", id: "tool-1", input: [:]) + + let messages: [Sampling.Message] = [ + .user("What's the weather?"), + .assistant(.toolUse(toolUse)), + .user("Thanks!"), // Invalid: should be tool_result, not text + ] + + #expect(throws: MCPError.self) { + try Sampling.Message.validateToolUseResultMessages(messages) + } + } + + @Test("validateToolUseResultMessages passes for conversation continuing after tool results") + func testConversationContinuesAfterToolResults() throws { + // Valid flow: tool_use → tool_result → text → text (conversation continues normally) + let toolUse = ToolUseContent(name: "get_weather", id: "tool-1", input: [:]) + let toolResult = ToolResultContent(toolUseId: "tool-1", content: [.text("Sunny")]) + + let messages: [Sampling.Message] = [ + .user("What's the weather?"), + .assistant(.toolUse(toolUse)), + Sampling.Message(role: .user, content: [.toolResult(toolResult)]), + .assistant("The weather is sunny!"), + .user("Great, thanks!"), // Valid: conversation continues after tool cycle + ] + + #expect(throws: Never.self) { + try Sampling.Message.validateToolUseResultMessages(messages) + } + } + + @Test("validateToolUseResultMessages passes when tool_use has text alongside") + func testToolUseWithTextInAssistantMessage() throws { + // Assistant can have both text and tool_use in the same message + let toolUse = ToolUseContent(name: "get_weather", id: "tool-1", input: [:]) + let toolResult = ToolResultContent(toolUseId: "tool-1", content: [.text("Sunny")]) + + let messages: [Sampling.Message] = [ + .user("What's the weather?"), + Sampling.Message(role: .assistant, content: [ + .text("Let me check the weather for you."), + .toolUse(toolUse), + ]), + Sampling.Message(role: .user, content: [.toolResult(toolResult)]), + ] + + #expect(throws: Never.self) { + try Sampling.Message.validateToolUseResultMessages(messages) + } + } +} + @Suite("Sampling Integration Tests") struct SamplingIntegrationTests { @Test( @@ -568,12 +987,11 @@ struct SamplingIntegrationTests { logger: logger ) - // Server with sampling capability + // Server (sampling is a client capability, not server) let server = Server( name: "SamplingTestServer", version: "1.0.0", capabilities: .init( - sampling: .init(), // Enable sampling tools: .init() ) ) @@ -595,98 +1013,6 @@ struct SamplingIntegrationTests { 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] = [ - .user("Analyze the following data and provide insights:"), - .user("Sales data: Q1: $100k, Q2: $150k, Q3: $200k, Q4: $180k"), - .user("Marketing data: Q1: $50k, Q2: $75k, Q3: $100k, Q4: $90k"), - ] - - 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)) ) @@ -709,7 +1035,7 @@ struct SamplingIntegrationTests { 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 { + if case .text(let text, _, _) = decodedTextMessage.content.first { #expect(text == "What do you see in this data?") } else { #expect(Bool(false), "Expected text content") @@ -719,7 +1045,7 @@ struct SamplingIntegrationTests { 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 { + if case .image(let data, let mimeType, _, _) = decodedImageMessage.content.first { #expect(data.contains("iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ")) #expect(mimeType == "image/png") } else { @@ -784,48 +1110,6 @@ struct SamplingIntegrationTests { #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] = [ - .user("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)) ) @@ -968,3 +1252,592 @@ struct SamplingIntegrationTests { #expect(decodedCreative.modelPreferences?.costPriority?.doubleValue == 0.4) } } + +@Suite("Client Sampling Parameters Tests") +struct ClientSamplingParametersTests { + @Test("ClientSamplingParameters init without tools") + func testClientSamplingParametersInitWithoutTools() throws { + let params = ClientSamplingParameters( + messages: [.user("Hello")], + maxTokens: 100 + ) + + #expect(params.messages.count == 1) + #expect(params.maxTokens == 100) + #expect(params.hasTools == false) + #expect(params.tools == nil) + } + + @Test("ClientSamplingParameters init with tools") + func testClientSamplingParametersInitWithTools() throws { + let tool = Tool( + name: "get_weather", + description: "Get weather", + inputSchema: .object([:]) + ) + + let params = ClientSamplingParameters( + messages: [.user("What's the weather?")], + maxTokens: 200, + tools: [tool], + toolChoice: ToolChoice(mode: .auto) + ) + + #expect(params.messages.count == 1) + #expect(params.maxTokens == 200) + #expect(params.hasTools == true) + #expect(params.tools?.count == 1) + #expect(params.tools?.first?.name == "get_weather") + #expect(params.toolChoice?.mode == .auto) + } + + @Test("ClientSamplingParameters encoding and decoding") + func testClientSamplingParametersCoding() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let tool = Tool( + name: "test_tool", + description: "A test tool", + inputSchema: .object([:]) + ) + + let params = ClientSamplingParameters( + messages: [.user("Hello"), .assistant("Hi there!")], + modelPreferences: ModelPreferences( + hints: [ModelPreferences.Hint(name: "claude-4")], + costPriority: 0.5 + ), + systemPrompt: "You are helpful", + temperature: 0.7, + maxTokens: 150, + stopSequences: ["STOP"], + tools: [tool], + toolChoice: ToolChoice(mode: .required) + ) + + let data = try encoder.encode(params) + let decoded = try decoder.decode(ClientSamplingParameters.self, from: data) + + #expect(decoded.messages.count == 2) + #expect(decoded.modelPreferences?.hints?.first?.name == "claude-4") + #expect(decoded.modelPreferences?.costPriority?.doubleValue == 0.5) + #expect(decoded.systemPrompt == "You are helpful") + #expect(decoded.temperature == 0.7) + #expect(decoded.maxTokens == 150) + #expect(decoded.stopSequences?.first == "STOP") + #expect(decoded.tools?.count == 1) + #expect(decoded.tools?.first?.name == "test_tool") + #expect(decoded.toolChoice?.mode == .required) + #expect(decoded.hasTools == true) + } + + @Test("ClientSamplingParameters hasTools with empty tools array") + func testClientSamplingParametersHasToolsEmpty() throws { + let params = ClientSamplingParameters( + messages: [.user("Hello")], + maxTokens: 100, + tools: [] // Empty array + ) + + // Empty array should be treated as no tools + #expect(params.hasTools == false) + } + + @Test("ClientSamplingRequest result type matches CreateSamplingMessageWithTools.Result") + func testClientSamplingRequestResultType() throws { + // Verify the type alias works correctly + let result: ClientSamplingRequest.Result = CreateSamplingMessageWithTools.Result( + model: "test-model", + stopReason: .endTurn, + role: .assistant, + content: .text("Hello") + ) + + #expect(result.model == "test-model") + #expect(result.stopReason == .endTurn) + #expect(result.content.count == 1) + } + + @Test("ClientSamplingParameters decodes from JSON without tools") + func testClientSamplingParametersDecodesWithoutTools() throws { + let decoder = JSONDecoder() + + let json = """ + { + "messages": [{"role": "user", "content": {"type": "text", "text": "Hello"}}], + "maxTokens": 100 + } + """.data(using: .utf8)! + + let params = try decoder.decode(ClientSamplingParameters.self, from: json) + + #expect(params.messages.count == 1) + #expect(params.maxTokens == 100) + #expect(params.hasTools == false) + #expect(params.tools == nil) + } + + @Test("ClientSamplingParameters decodes from JSON with tools") + func testClientSamplingParametersDecodesWithTools() throws { + let decoder = JSONDecoder() + + let json = """ + { + "messages": [{"role": "user", "content": {"type": "text", "text": "Hello"}}], + "maxTokens": 100, + "tools": [{"name": "test", "inputSchema": {"type": "object"}}], + "toolChoice": {"mode": "auto"} + } + """.data(using: .utf8)! + + let params = try decoder.decode(ClientSamplingParameters.self, from: json) + + #expect(params.messages.count == 1) + #expect(params.maxTokens == 100) + #expect(params.hasTools == true) + #expect(params.tools?.count == 1) + #expect(params.tools?.first?.name == "test") + #expect(params.toolChoice?.mode == .auto) + } +} + +@Suite("Sampling JSON Format Verification Tests") +struct SamplingJSONFormatTests { + @Test("ToolChoice encodes correctly") + func testToolChoiceEncoding() throws { + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + + let autoChoice = ToolChoice(mode: .auto) + let requiredChoice = ToolChoice(mode: .required) + let noneChoice = ToolChoice(mode: ToolChoice.Mode.none) // Explicit to avoid ambiguity with Optional.none + let nilChoice = ToolChoice() + + #expect(String(data: try encoder.encode(autoChoice), encoding: .utf8) == "{\"mode\":\"auto\"}") + #expect(String(data: try encoder.encode(requiredChoice), encoding: .utf8) == "{\"mode\":\"required\"}") + #expect(String(data: try encoder.encode(noneChoice), encoding: .utf8) == "{\"mode\":\"none\"}") + #expect(String(data: try encoder.encode(nilChoice), encoding: .utf8) == "{}") + } + + @Test("ToolChoice decodes all modes correctly") + func testToolChoiceDecoding() throws { + let decoder = JSONDecoder() + + let auto = try decoder.decode(ToolChoice.self, from: "{\"mode\":\"auto\"}".data(using: .utf8)!) + let required = try decoder.decode(ToolChoice.self, from: "{\"mode\":\"required\"}".data(using: .utf8)!) + let none = try decoder.decode(ToolChoice.self, from: "{\"mode\":\"none\"}".data(using: .utf8)!) + let empty = try decoder.decode(ToolChoice.self, from: "{}".data(using: .utf8)!) + + #expect(auto.mode == .auto) + #expect(required.mode == .required) + #expect(none.mode == ToolChoice.Mode.none) // Explicit to avoid ambiguity with Optional.none + #expect(empty.mode == nil) + } + + @Test("Sampling result single content encodes as object not array") + func testSingleContentEncodesAsObject() throws { + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + + let result = CreateSamplingMessage.Result( + model: "test", + role: .assistant, + content: .text("Hello") + ) + + let json = String(data: try encoder.encode(result), encoding: .utf8)! + + // Content should be an object, not an array + #expect(json.contains("\"content\":{\"text\":\"Hello\",\"type\":\"text\"}")) + #expect(!json.contains("\"content\":[")) + } + + @Test("Sampling result with tools single content encodes as object") + func testToolsResultSingleContentEncodesAsObject() throws { + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + + let result = CreateSamplingMessageWithTools.Result( + model: "test", + role: .assistant, + content: .text("Hello") + ) + + let json = String(data: try encoder.encode(result), encoding: .utf8)! + + // Single content should be an object, not an array + #expect(json.contains("\"content\":{\"text\":\"Hello\",\"type\":\"text\"}")) + #expect(!json.contains("\"content\":[")) + } + + @Test("Sampling result with tools multiple content encodes as array") + func testToolsResultMultipleContentEncodesAsArray() throws { + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + + let toolUse1 = ToolUseContent(name: "tool1", id: "1", input: [:]) + let toolUse2 = ToolUseContent(name: "tool2", id: "2", input: [:]) + + let result = CreateSamplingMessageWithTools.Result( + model: "test", + role: .assistant, + content: [.toolUse(toolUse1), .toolUse(toolUse2)] + ) + + let json = String(data: try encoder.encode(result), encoding: .utf8)! + + // Multiple content should be an array + #expect(json.contains("\"content\":[")) + } + + @Test("Sampling message single content encodes as object") + func testMessageSingleContentEncodesAsObject() throws { + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + + let message: Sampling.Message = .user("Hello") + + let json = String(data: try encoder.encode(message), encoding: .utf8)! + + // Single content should be an object + #expect(json.contains("\"content\":{\"text\":\"Hello\",\"type\":\"text\"}")) + #expect(!json.contains("\"content\":[")) + } + + @Test("Sampling message multiple content encodes as array") + func testMessageMultipleContentEncodesAsArray() throws { + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + + let message = Sampling.Message(role: .user, content: [ + .text("Hello"), + .image(data: "abc", mimeType: "image/png"), + ]) + + let json = String(data: try encoder.encode(message), encoding: .utf8)! + + // Multiple content should be an array + #expect(json.contains("\"content\":[")) + } + + @Test("StopReason encodes as raw string") + func testStopReasonEncoding() throws { + let encoder = JSONEncoder() + + #expect(String(data: try encoder.encode(StopReason.endTurn), encoding: .utf8) == "\"endTurn\"") + #expect(String(data: try encoder.encode(StopReason.stopSequence), encoding: .utf8) == "\"stopSequence\"") + #expect(String(data: try encoder.encode(StopReason.maxTokens), encoding: .utf8) == "\"maxTokens\"") + #expect(String(data: try encoder.encode(StopReason.toolUse), encoding: .utf8) == "\"toolUse\"") + + // Custom stop reason + let custom = StopReason(rawValue: "customReason") + #expect(String(data: try encoder.encode(custom), encoding: .utf8) == "\"customReason\"") + } + + @Test("ToolUseContent encodes with correct structure") + func testToolUseContentEncoding() throws { + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + + let toolUse = ToolUseContent( + name: "get_weather", + id: "call-123", + input: ["city": "Paris"] + ) + + let json = String(data: try encoder.encode(toolUse), encoding: .utf8)! + + #expect(json.contains("\"type\":\"tool_use\"")) + #expect(json.contains("\"name\":\"get_weather\"")) + #expect(json.contains("\"id\":\"call-123\"")) + #expect(json.contains("\"input\":{\"city\":\"Paris\"}")) + } + + @Test("ToolResultContent encodes with correct structure") + func testToolResultContentEncoding() throws { + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + + let toolResult = ToolResultContent( + toolUseId: "call-123", + content: [.text("Sunny, 72°F")], + isError: false + ) + + let json = String(data: try encoder.encode(toolResult), encoding: .utf8)! + + #expect(json.contains("\"type\":\"tool_result\"")) + #expect(json.contains("\"toolUseId\":\"call-123\"")) + #expect(json.contains("\"isError\":false")) + } + + @Test("CreateMessage result decodes from TypeScript SDK format") + func testDecodesFromTypeScriptFormat() throws { + let decoder = JSONDecoder() + + // Format matching TypeScript SDK output + let json = """ + { + "model": "claude-4-sonnet", + "stopReason": "endTurn", + "role": "assistant", + "content": { + "type": "text", + "text": "Hello from TypeScript SDK" + } + } + """.data(using: .utf8)! + + let result = try decoder.decode(CreateSamplingMessage.Result.self, from: json) + + #expect(result.model == "claude-4-sonnet") + #expect(result.stopReason == .endTurn) + #expect(result.role == .assistant) + if case .text(let text, _, _) = result.content { + #expect(text == "Hello from TypeScript SDK") + } else { + #expect(Bool(false), "Expected text content") + } + } + + @Test("CreateMessage result decodes from Python SDK format with array") + func testDecodesFromPythonFormatWithArray() throws { + let decoder = JSONDecoder() + + // Python SDK may send array content + let json = """ + { + "model": "claude-4-sonnet", + "stopReason": "toolUse", + "role": "assistant", + "content": [ + { + "type": "tool_use", + "name": "get_weather", + "id": "toolu_123", + "input": {"city": "Paris"} + } + ] + } + """.data(using: .utf8)! + + let result = try decoder.decode(CreateSamplingMessageWithTools.Result.self, from: json) + + #expect(result.model == "claude-4-sonnet") + #expect(result.stopReason == .toolUse) + #expect(result.content.count == 1) + if case .toolUse(let toolUse) = result.content.first { + #expect(toolUse.name == "get_weather") + #expect(toolUse.id == "toolu_123") + } else { + #expect(Bool(false), "Expected tool use content") + } + } +} + +@Suite("Server Sampling Capability Validation Tests") +struct ServerSamplingCapabilityValidationTests { + @Test( + "Server throws when tools provided without sampling.tools capability", + .timeLimit(.minutes(1)) + ) + func testServerThrowsWhenToolsWithoutCapability() async throws { + // Set up server + let server = Server( + name: "TestServer", + version: "1.0.0", + capabilities: .init() + ) + + // Set up client without sampling.tools capability (just basic sampling) + let client = Client( + name: "TestClient", + version: "1.0" + ) + await client.setCapabilities(Client.Capabilities(sampling: .init())) + + // Connect via in-memory transport + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Attempt to call createMessageWithTools should fail because client + // doesn't have sampling.tools capability + let params = CreateSamplingMessageWithTools.Parameters( + messages: [.user("Hello")], + maxTokens: 100, + tools: [ + Tool(name: "test_tool", inputSchema: .object([:])) + ], + toolChoice: nil + ) + + await #expect(throws: MCPError.self) { + _ = try await server.createMessageWithTools(params) + } + + await server.stop() + await client.disconnect() + } + + @Test( + "Server throws when client lacks sampling capability entirely", + .timeLimit(.minutes(1)) + ) + func testServerThrowsWithoutSamplingCapability() async throws { + // Set up server (sampling is a client capability, not server) + let server = Server( + name: "TestServer", + version: "1.0.0", + capabilities: .init() + ) + + // Set up client without sampling capability + let client = Client( + name: "TestClient", + version: "1.0" + ) + // Don't set any sampling capability + + // Connect via in-memory transport + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Attempt to call createMessage should fail because client + // doesn't have sampling capability + let params = SamplingParameters( + messages: [.user("Hello")], + maxTokens: 100 + ) + + await #expect(throws: MCPError.self) { + _ = try await server.createMessage(params) + } + + await server.stop() + await client.disconnect() + } + + @Test( + "Server succeeds when client has sampling capability", + .timeLimit(.minutes(1)) + ) + func testServerSucceedsWithSamplingCapability() async throws { + // Set up server (sampling is a client capability, not server) + let server = Server( + name: "TestServer", + version: "1.0.0", + capabilities: .init() + ) + + // Set up client WITH sampling capability + let client = Client( + name: "TestClient", + version: "1.0" + ) + await client.setCapabilities(Client.Capabilities(sampling: .init())) + + // Set up client sampling handler + await client.withSamplingHandler { _, _ in + return ClientSamplingRequest.Result( + model: "test-model", + stopReason: .endTurn, + role: .assistant, + content: [.text("Hello from client!")] + ) + } + + // Connect via in-memory transport + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Attempt to call createMessage should succeed + let params = SamplingParameters( + messages: [.user("Hello")], + maxTokens: 100 + ) + + let result = try await server.createMessage(params) + + #expect(result.model == "test-model") + #expect(result.stopReason == .endTurn) + if case .text(let text, _, _) = result.content { + #expect(text == "Hello from client!") + } else { + #expect(Bool(false), "Expected text content") + } + + await server.stop() + await client.disconnect() + } + + @Test( + "Server succeeds with tools when client has sampling.tools capability", + .timeLimit(.minutes(1)) + ) + func testServerSucceedsWithToolsCapability() async throws { + // Set up server (sampling is a client capability, not server) + let server = Server( + name: "TestServer", + version: "1.0.0", + capabilities: .init() + ) + + // Set up client WITH sampling.tools capability + let client = Client( + name: "TestClient", + version: "1.0" + ) + await client.setCapabilities(Client.Capabilities(sampling: .init(tools: .init()))) + + // Set up client sampling handler + await client.withSamplingHandler { _, _ in + return ClientSamplingRequest.Result( + model: "test-model", + stopReason: .toolUse, + role: .assistant, + content: [.toolUse(ToolUseContent( + name: "get_weather", + id: "call-123", + input: ["city": "Paris"] + ))] + ) + } + + // Connect via in-memory transport + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Attempt to call createMessageWithTools should succeed + let params = CreateSamplingMessageWithTools.Parameters( + messages: [.user("What's the weather?")], + maxTokens: 100, + tools: [ + Tool(name: "get_weather", inputSchema: .object([:])) + ], + toolChoice: ToolChoice(mode: .auto) + ) + + let result = try await server.createMessageWithTools(params) + + #expect(result.model == "test-model") + #expect(result.stopReason == .toolUse) + #expect(result.content.count == 1) + if case .toolUse(let toolUse) = result.content.first { + #expect(toolUse.name == "get_weather") + } else { + #expect(Bool(false), "Expected tool use content") + } + + await server.stop() + await client.disconnect() + } +} diff --git a/Tests/MCPTests/ServerTests.swift b/Tests/MCPTests/ServerTests.swift index 9bc9c01a..d65be872 100644 --- a/Tests/MCPTests/ServerTests.swift +++ b/Tests/MCPTests/ServerTests.swift @@ -39,7 +39,8 @@ struct ServerTests { try await server.start(transport: transport) // Wait for message processing and response - try await Task.sleep(for: .milliseconds(100)) + let received = await transport.waitForSentMessageCount(1) + #expect(received, "Timed out waiting for initialize response") #expect(await transport.sentMessages.count == 1) @@ -198,4 +199,50 @@ struct ServerTests { await server.stop() await transport.disconnect() } + + @Test("Invalid JSON-RPC message returns error") + func testInvalidJsonRpcMessageReturnsError() async throws { + let transport = MockTransport() + let server = Server(name: "TestServer", version: "1.0") + + try await server.start(transport: transport) + + // Initialize first + try await transport.queue( + request: Initialize.request( + .init( + protocolVersion: Version.latest, + capabilities: .init(), + clientInfo: .init(name: "TestClient", version: "1.0") + ) + ) + ) + + // Wait for init response + let initReceived = await transport.waitForSentMessageCount(1) + #expect(initReceived, "Timed out waiting for init response") + await transport.clearMessages() + + // Send invalid JSON-RPC message (missing jsonrpc field) + // This tests that the server properly validates incoming messages + let invalidMessage = #"{"method":"ping","id":"1"}"# + await transport.queueRaw(invalidMessage) + + // Wait for error response with polling instead of fixed sleep + let errorReceived = await transport.waitForSentMessage { message in + message.contains("error") + } + #expect(errorReceived, "Timed out waiting for error response") + + let messages = await transport.sentMessages + #expect(messages.count >= 1) + + // Should get an error response + if let response = messages.first { + #expect(response.contains("error")) + } + + await server.stop() + await transport.disconnect() + } } diff --git a/Tests/MCPTests/SessionManagerTests.swift b/Tests/MCPTests/SessionManagerTests.swift new file mode 100644 index 00000000..25118fd0 --- /dev/null +++ b/Tests/MCPTests/SessionManagerTests.swift @@ -0,0 +1,239 @@ +import Testing + +@testable import MCP + +/// Tests for SessionManager - a thread-safe session storage helper. +/// +/// These tests follow the TypeScript SDK's pattern for session management. +@Suite("Session Manager Tests") +struct SessionManagerTests { + + // MARK: - Basic Operations + + @Test("Initialization creates empty manager") + func initialization() async { + let manager = SessionManager() + let count = await manager.activeSessionCount + #expect(count == 0) + } + + @Test("Max sessions configuration") + func maxSessionsConfiguration() async { + let manager = SessionManager(maxSessions: 5) + let max = await manager.maxSessions + #expect(max == 5) + } + + @Test("Store and retrieve transport") + func storeAndRetrieveTransport() async throws { + let manager = SessionManager() + let transport = HTTPServerTransport() + + await manager.store(transport, forSessionId: "test-session-1") + + let retrieved = await manager.transport(forSessionId: "test-session-1") + #expect(retrieved != nil) + + let count = await manager.activeSessionCount + #expect(count == 1) + } + + @Test("Transport not found returns nil") + func transportNotFound() async { + let manager = SessionManager() + + let retrieved = await manager.transport(forSessionId: "nonexistent") + #expect(retrieved == nil) + } + + @Test("Remove transport") + func removeTransport() async throws { + let manager = SessionManager() + let transport = HTTPServerTransport() + + await manager.store(transport, forSessionId: "test-session") + + var count = await manager.activeSessionCount + #expect(count == 1) + + await manager.remove("test-session") + + count = await manager.activeSessionCount + #expect(count == 0) + + let retrieved = await manager.transport(forSessionId: "test-session") + #expect(retrieved == nil) + } + + @Test("Active session IDs") + func activeSessionIds() async throws { + let manager = SessionManager() + + await manager.store(HTTPServerTransport(), forSessionId: "session-a") + await manager.store(HTTPServerTransport(), forSessionId: "session-b") + await manager.store(HTTPServerTransport(), forSessionId: "session-c") + + let ids = await manager.activeSessionIds + #expect(ids.sorted() == ["session-a", "session-b", "session-c"]) + } + + @Test("Close all sessions") + func closeAll() async throws { + let manager = SessionManager() + + await manager.store(HTTPServerTransport(), forSessionId: "session-1") + await manager.store(HTTPServerTransport(), forSessionId: "session-2") + + var count = await manager.activeSessionCount + #expect(count == 2) + + await manager.closeAll() + + count = await manager.activeSessionCount + #expect(count == 0) + } + + // MARK: - Capacity Limits + + @Test("Capacity check") + func capacityCheck() async { + let manager = SessionManager(maxSessions: 2) + + // Initially can add + var canAdd = await manager.canAddSession() + #expect(canAdd == true) + + // Add two sessions + await manager.store(HTTPServerTransport(), forSessionId: "session-1") + await manager.store(HTTPServerTransport(), forSessionId: "session-2") + + // Now at capacity + canAdd = await manager.canAddSession() + #expect(canAdd == false) + + // Remove one + await manager.remove("session-1") + + // Can add again + canAdd = await manager.canAddSession() + #expect(canAdd == true) + } + + @Test("Unlimited capacity") + func unlimitedCapacity() async { + let manager = SessionManager() // No maxSessions + + // Add many sessions + for i in 0..<100 { + await manager.store(HTTPServerTransport(), forSessionId: "session-\(i)") + } + + // Should still be able to add + let canAdd = await manager.canAddSession() + #expect(canAdd == true) + + let count = await manager.activeSessionCount + #expect(count == 100) + } + + // MARK: - Session Cleanup + + @Test("Cleanup stale sessions") + func cleanUpStaleSessions() async throws { + let manager = SessionManager() + + // Store some sessions + await manager.store(HTTPServerTransport(), forSessionId: "old-session") + + // Wait a small amount to ensure the session activity time is in the past + try await Task.sleep(for: .milliseconds(10)) + + // Clean up with zero timeout - should remove all sessions since they're now "stale" + let removed = await manager.cleanUpStaleSessions(olderThan: .zero) + #expect(removed == 1) + + let count = await manager.activeSessionCount + #expect(count == 0) + } + + @Test("Recent session not cleaned") + func recentSessionNotCleaned() async throws { + let manager = SessionManager() + + await manager.store(HTTPServerTransport(), forSessionId: "recent-session") + + // Clean up with long timeout - should not remove recent session + let removed = await manager.cleanUpStaleSessions(olderThan: .seconds(3600)) + #expect(removed == 0) + + let count = await manager.activeSessionCount + #expect(count == 1) + } + + // MARK: - Multi-Client Simulation + + @Test("Multiple clients sequential") + func multipleClientsSequential() async throws { + let manager = SessionManager() + + // Simulate 10 clients connecting sequentially + for i in 0..<10 { + let transport = HTTPServerTransport() + await manager.store(transport, forSessionId: "session-\(i)") + } + + let count = await manager.activeSessionCount + #expect(count == 10) + + let ids = await manager.activeSessionIds + #expect(ids.count == 10) + } + + @Test("Multiple clients concurrent") + func multipleClientsConcurrent() async throws { + let manager = SessionManager() + + // Simulate 10 clients connecting concurrently + await withTaskGroup(of: Void.self) { group in + for i in 0..<10 { + group.addTask { + let transport = HTTPServerTransport() + await manager.store(transport, forSessionId: "concurrent-session-\(i)") + } + } + } + + let count = await manager.activeSessionCount + #expect(count == 10) + } + + @Test("Concurrent access and removal") + func concurrentAccessAndRemoval() async throws { + let manager = SessionManager() + + // Pre-populate with sessions + for i in 0..<20 { + await manager.store(HTTPServerTransport(), forSessionId: "session-\(i)") + } + + // Concurrently access and remove sessions + await withTaskGroup(of: Void.self) { group in + // Readers + for i in 0..<10 { + group.addTask { + _ = await manager.transport(forSessionId: "session-\(i)") + } + } + + // Removers + for i in 10..<20 { + group.addTask { + await manager.remove("session-\(i)") + } + } + } + + let count = await manager.activeSessionCount + #expect(count == 10) // Only the first 10 should remain + } +} diff --git a/Tests/MCPTests/StdioTransportTests.swift b/Tests/MCPTests/StdioTransportTests.swift index e395fbe0..05a47d12 100644 --- a/Tests/MCPTests/StdioTransportTests.swift +++ b/Tests/MCPTests/StdioTransportTests.swift @@ -9,6 +9,8 @@ import Testing @preconcurrency import SystemPackage #endif +// MARK: - Basic Tests + @Suite("Stdio Transport Tests") struct StdioTransportTests { @Test("Connection") @@ -125,3 +127,314 @@ struct StdioTransportTests { await transport.disconnect() } } + +// MARK: - Multiple Message Tests (mirrors TypeScript server/stdio.test.ts) + +@Suite("Stdio Transport Multiple Message Tests") +struct StdioTransportMultipleMessageTests { + @Test("Receive multiple messages") + func testReceiveMultipleMessages() async throws { + let (input, writer) = try FileDescriptor.pipe() + let (_, output) = try FileDescriptor.pipe() + let transport = StdioTransport(input: input, output: output, logger: nil) + try await transport.connect() + + // Write multiple JSON-RPC messages (mirrors TypeScript test_stdio_client) + let messages = [ + #"{"jsonrpc":"2.0","id":1,"method":"ping"}"#, + #"{"jsonrpc":"2.0","method":"notifications/initialized"}"#, + ] + + for message in messages { + try writer.writeAll((message + "\n").data(using: .utf8)!) + } + try writer.close() + + // Receive and verify all messages + let stream = await transport.receive() + var receivedMessages: [String] = [] + + for try await data in stream { + if let message = String(data: data, encoding: .utf8) { + receivedMessages.append(message) + } + } + + #expect(receivedMessages.count == 2) + #expect(receivedMessages[0] == messages[0]) + #expect(receivedMessages[1] == messages[1]) + + await transport.disconnect() + } + + @Test("Send multiple messages") + func testSendMultipleMessages() async throws { + let (reader, output) = try FileDescriptor.pipe() + let (input, _) = try FileDescriptor.pipe() + let transport = StdioTransport(input: input, output: output, logger: nil) + try await transport.connect() + + // Send multiple messages + let messages = [ + #"{"jsonrpc":"2.0","id":1,"method":"ping"}"#, + #"{"jsonrpc":"2.0","id":2,"result":{}}"#, + ] + + for message in messages { + try await transport.send(message.data(using: .utf8)!) + } + + // Read all output at once + var buffer = [UInt8](repeating: 0, count: 4096) + let bytesRead = try buffer.withUnsafeMutableBufferPointer { pointer in + try reader.read(into: UnsafeMutableRawBufferPointer(pointer)) + } + + let receivedOutput = String(data: Data(buffer[.. server communication + let (transportInput, clientWriter) = try FileDescriptor.pipe() + let (clientReader, transportOutput) = try FileDescriptor.pipe() + + let transport = StdioTransport( + input: transportInput, output: transportOutput, logger: nil) + try await transport.connect() + + // Client sends a request + let request = #"{"jsonrpc":"2.0","id":1,"method":"ping"}"# + try clientWriter.writeAll((request + "\n").data(using: .utf8)!) + + // Transport receives the request + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + let receivedRequest = try await iterator.next() + #expect(receivedRequest == request.data(using: .utf8)!) + + // Transport sends a response back + let response = #"{"jsonrpc":"2.0","id":1,"result":{"status":"ok"}}"# + try await transport.send(response.data(using: .utf8)!) + + // Client reads the response + var buffer = [UInt8](repeating: 0, count: 4096) + let bytesRead = try buffer.withUnsafeMutableBufferPointer { pointer in + try clientReader.read(into: UnsafeMutableRawBufferPointer(pointer)) + } + let receivedResponse = String( + data: Data(buffer[.. String? { sessionId } + } + let store = SessionStore() + let expectedSessionId = "test-session-123" + + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { expectedSessionId }, + onSessionInitialized: { sessionId in + await store.set(sessionId) + } + ) + ) + try await transport.connect() + + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + + let response = await transport.handleRequest(initRequest) + + #expect(response.statusCode == 200) + #expect(response.headers[HTTPHeader.sessionId] == expectedSessionId) + + let generatedSessionId = await store.get() + #expect(generatedSessionId == expectedSessionId) + + let actualSessionId = await transport.sessionId + #expect(actualSessionId == expectedSessionId) + } + + @Test("Stateful mode rejects invalid session ID") + func statefulModeRejectsInvalidSessionId() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { "valid-session" }) + ) + try await transport.connect() + + // First initialize + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // Then try with wrong session ID + let badRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + HTTPHeader.sessionId: "wrong-session", + ], + body: + #"{"jsonrpc":"2.0","method":"tools/list","id":"2"}"#.data(using: .utf8) + ) + + let response = await transport.handleRequest(badRequest) + #expect(response.statusCode == 404) + } + + @Test("Stateful mode requires session ID after init") + func statefulModeRequiresSessionIdAfterInit() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { "test-session" }) + ) + try await transport.connect() + + // Initialize + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // Try without session ID + let noSessionRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + #"{"jsonrpc":"2.0","method":"tools/list","id":"2"}"#.data(using: .utf8) + ) + + let response = await transport.handleRequest(noSessionRequest) + #expect(response.statusCode == 400) + } + + // MARK: - GET Request Handling + + @Test("GET requires Accept header") + func getRequiresAcceptHeader() async throws { + let transport = HTTPServerTransport() + try await transport.connect() + + let request = HTTPRequest(method: "GET", headers: [:]) + let response = await transport.handleRequest(request) + + #expect(response.statusCode == 406) + } + + @Test("GET returns SSE stream") + func getReturnsSSEStream() async throws { + let transport = HTTPServerTransport() + try await transport.connect() + + // Initialize first (stateless mode - no session required) + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + let request = HTTPRequest(method: "GET", headers: [HTTPHeader.accept: "text/event-stream"]) + let response = await transport.handleRequest(request) + + #expect(response.statusCode == 200) + #expect(response.headers[HTTPHeader.contentType] == "text/event-stream") + #expect(response.stream != nil) + } + + @Test("GET rejects multiple SSE streams") + func getRejectsMultipleSSEStreams() async throws { + let transport = HTTPServerTransport() + try await transport.connect() + + // Initialize + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // First GET + let request1 = HTTPRequest(method: "GET", headers: [HTTPHeader.accept: "text/event-stream"]) + let response1 = await transport.handleRequest(request1) + #expect(response1.statusCode == 200) + + // Second GET - should fail + let request2 = HTTPRequest(method: "GET", headers: [HTTPHeader.accept: "text/event-stream"]) + let response2 = await transport.handleRequest(request2) + #expect(response2.statusCode == 409) + } + + // MARK: - DELETE Request Handling + + @Test("DELETE closes session") + func deleteClosesSession() async throws { + actor ClosedState { + var closed = false + func markClosed() { closed = true } + func isClosed() -> Bool { closed } + } + let state = ClosedState() + + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { "test-session" }, + onSessionClosed: { _ in await state.markClosed() } + ) + ) + try await transport.connect() + + // Initialize + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // DELETE + let deleteRequest = HTTPRequest( + method: "DELETE", + headers: [HTTPHeader.sessionId: "test-session"] + ) + let response = await transport.handleRequest(deleteRequest) + + #expect(response.statusCode == 200) + let sessionClosed = await state.isClosed() + #expect(sessionClosed == true) + } + + @Test("DELETE in stateless mode returns 405") + func deleteInStatelessModeReturns405() async throws { + // Stateless mode (no session ID generator) + let transport = HTTPServerTransport() + try await transport.connect() + + // Initialize first + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // DELETE should return 405 in stateless mode (no session to terminate) + let deleteRequest = HTTPRequest(method: "DELETE", headers: [:]) + let response = await transport.handleRequest(deleteRequest) + + #expect(response.statusCode == 405) + #expect(response.headers[HTTPHeader.allow] == "GET, POST") + + // Verify it's a proper JSON-RPC error + if let body = response.body, let text = String(data: body, encoding: .utf8) { + #expect(text.contains("Session management is not enabled") || text.contains("Method Not Allowed")) + } + } + + @Test("Session terminated - requests return 404") + func sessionTerminatedRequestsReturn404() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { "test-session" }) + ) + try await transport.connect() + + // Initialize + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // DELETE to terminate session + let deleteRequest = HTTPRequest( + method: "DELETE", + headers: [HTTPHeader.sessionId: "test-session"] + ) + let deleteResponse = await transport.handleRequest(deleteRequest) + #expect(deleteResponse.statusCode == 200) + + // Try to use the terminated session - should return 404 + let postRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + HTTPHeader.sessionId: "test-session", + ], + body: #"{"jsonrpc":"2.0","method":"tools/list","id":"2"}"#.data(using: .utf8) + ) + let postResponse = await transport.handleRequest(postRequest) + #expect(postResponse.statusCode == 404) + + // GET with terminated session should also return 404 + let getRequest = HTTPRequest( + method: "GET", + headers: [ + HTTPHeader.accept: "text/event-stream", + HTTPHeader.sessionId: "test-session", + ] + ) + let getResponse = await transport.handleRequest(getRequest) + #expect(getResponse.statusCode == 404) + } + + // MARK: - Unsupported Methods + + @Test("Unsupported method returns 405") + func unsupportedMethodReturns405() async throws { + let transport = HTTPServerTransport() + try await transport.connect() + + let request = HTTPRequest(method: "PUT", headers: [:]) + let response = await transport.handleRequest(request) + + #expect(response.statusCode == 405) + #expect(response.headers[HTTPHeader.allow] == "GET, POST, DELETE") + } + + // MARK: - Protocol Version Validation + + @Test("Rejects unsupported protocol version") + func rejectsUnsupportedProtocolVersion() async throws { + let transport = HTTPServerTransport() + try await transport.connect() + + // Initialize first + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // Then try with unsupported version + let badRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + HTTPHeader.protocolVersion: "1999-01-01", + ], + body: + #"{"jsonrpc":"2.0","method":"tools/list","id":"2"}"#.data(using: .utf8) + ) + + let response = await transport.handleRequest(badRequest) + #expect(response.statusCode == 400) + } + + // MARK: - JSON Response Mode + + @Test("JSON response mode only requires application/json Accept header") + func jsonResponseModeOnlyRequiresJsonAcceptHeader() async throws { + let transport = HTTPServerTransport( + options: .init(enableJsonResponse: true) + ) + try await transport.connect() + + // Should succeed with only application/json (no text/event-stream required) + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json", // Only JSON, no SSE + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + + // Start a task to send the response + Task { + try await Task.sleep(for: .milliseconds(50)) + let responseData = + TestPayloads.initializeResult() + .data(using: .utf8)! + try await transport.send(responseData, relatedRequestId: .string("1")) + } + + let response = await transport.handleRequest(initRequest) + #expect(response.statusCode == 200) + #expect(response.headers[HTTPHeader.contentType] == "application/json") + } + + @Test("JSON response mode rejects request without application/json Accept header") + func jsonResponseModeRejectsWithoutJsonAccept() async throws { + let transport = HTTPServerTransport( + options: .init(enableJsonResponse: true) + ) + try await transport.connect() + + // Should fail without application/json + let request = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "text/plain", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + + let response = await transport.handleRequest(request) + #expect(response.statusCode == 406) + } + + @Test("SSE mode requires both Accept types") + func sseModeRequiresBothAcceptTypes() async throws { + let transport = HTTPServerTransport( + options: .init(enableJsonResponse: false) // SSE mode (default) + ) + try await transport.connect() + + // Should fail with only application/json (missing text/event-stream) + let request = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json", // Missing text/event-stream + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + + let response = await transport.handleRequest(request) + #expect(response.statusCode == 406) + if let body = response.body, let text = String(data: body, encoding: .utf8) { + #expect(text.contains("both")) + } + } + + @Test("JSON response mode returns JSON") + func jsonResponseModeReturnsJson() async throws { + let transport = HTTPServerTransport( + options: .init(enableJsonResponse: true) + ) + try await transport.connect() + + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + + // Start a task to send the response + Task { + // Wait a bit for the request to be processed + try await Task.sleep(for: .milliseconds(50)) + + // Send a response + let responseData = + TestPayloads.initializeResult() + .data(using: .utf8)! + try await transport.send(responseData, relatedRequestId: .string("1")) + } + + let response = await transport.handleRequest(initRequest) + + #expect(response.statusCode == 200) + #expect(response.headers[HTTPHeader.contentType] == "application/json") + #expect(response.body != nil) + #expect(response.stream == nil) + } + + // MARK: - Multiple Initialize Rejection + + @Test("Rejects double initialize") + func rejectsDoubleInitialize() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { UUID().uuidString }) + ) + try await transport.connect() + + // First initialize + let initRequest1 = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + let response1 = await transport.handleRequest(initRequest1) + #expect(response1.statusCode == 200) + + let sessionId = response1.headers[HTTPHeader.sessionId]! + + // Second initialize - should fail + let initRequest2 = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + HTTPHeader.sessionId: sessionId, + ], + body: + TestPayloads.initializeRequest(id: "2") + .data(using: .utf8) + ) + let response2 = await transport.handleRequest(initRequest2) + #expect(response2.statusCode == 400) + } + + // MARK: - DNS Rebinding Protection + + @Test("DNS rebinding protection allows valid host") + func dnsRebindingProtectionAllowsValidHost() async throws { + let transport = HTTPServerTransport( + options: .init(security: .forLocalhost(port: 8080)) + ) + try await transport.connect() + + let request = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + HTTPHeader.host: "localhost:8080", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + + let response = await transport.handleRequest(request) + #expect(response.statusCode == 200) + } + + @Test("DNS rebinding protection rejects invalid host") + func dnsRebindingProtectionRejectsInvalidHost() async throws { + let transport = HTTPServerTransport( + options: .init(security: .forLocalhost(port: 8080)) + ) + try await transport.connect() + + let request = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + HTTPHeader.host: "evil.com:8080", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + + let response = await transport.handleRequest(request) + // 421 Misdirected Request for invalid Host + #expect(response.statusCode == 421) + } + + @Test("DNS rebinding protection rejects missing host") + func dnsRebindingProtectionRejectsMissingHost() async throws { + let transport = HTTPServerTransport( + options: .init(security: .forLocalhost(port: 8080)) + ) + try await transport.connect() + + let request = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + // No Host header + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + + let response = await transport.handleRequest(request) + // 421 Misdirected Request for missing Host + #expect(response.statusCode == 421) + } + + @Test("DNS rebinding protection allows wildcard port") + func dnsRebindingProtectionAllowsWildcardPort() async throws { + let transport = HTTPServerTransport( + options: .init(security: .forLocalhost()) // Wildcard port + ) + try await transport.connect() + + // Test with various ports + for port in [8080, 3000, 9999] { + let request = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + HTTPHeader.host: "127.0.0.1:\(port)", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + + let response = await transport.handleRequest(request) + // First request succeeds (200), subsequent ones fail (400) because already initialized + #expect(response.statusCode == 200 || response.statusCode == 400) + } + } + + @Test("DNS rebinding protection allows valid origin") + func dnsRebindingProtectionAllowsValidOrigin() async throws { + let transport = HTTPServerTransport( + options: .init(security: .forLocalhost(port: 8080)) + ) + try await transport.connect() + + let request = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + HTTPHeader.host: "localhost:8080", + HTTPHeader.origin: "http://localhost:8080", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + + let response = await transport.handleRequest(request) + #expect(response.statusCode == 200) + } + + @Test("DNS rebinding protection rejects invalid origin") + func dnsRebindingProtectionRejectsInvalidOrigin() async throws { + let transport = HTTPServerTransport( + options: .init(security: .forLocalhost(port: 8080)) + ) + try await transport.connect() + + let request = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + HTTPHeader.host: "localhost:8080", + HTTPHeader.origin: "http://evil.com", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + + let response = await transport.handleRequest(request) + #expect(response.statusCode == 403) + } + + @Test("DNS rebinding protection allows request without origin") + func dnsRebindingProtectionAllowsRequestWithoutOrigin() async throws { + // Non-browser clients (like curl) don't send Origin header + let transport = HTTPServerTransport( + options: .init(security: .forLocalhost(port: 8080)) + ) + try await transport.connect() + + let request = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + HTTPHeader.host: "localhost:8080", + // No Origin header + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + + let response = await transport.handleRequest(request) + #expect(response.statusCode == 200) + } + + @Test("DNS rebinding protection disabled by default") + func dnsRebindingProtectionDisabledByDefault() async throws { + // Without security settings, any host should be allowed + let transport = HTTPServerTransport() + try await transport.connect() + + let request = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + HTTPHeader.host: "any-host.com:8080", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + + let response = await transport.handleRequest(request) + #expect(response.statusCode == 200) + } + + @Test("forBindAddress auto-enables for localhost") + func forBindAddressAutoEnablesForLocalhost() { + // Should return security settings for localhost addresses + #expect(TransportSecuritySettings.forBindAddress(host: "127.0.0.1", port: 8080) != nil) + #expect(TransportSecuritySettings.forBindAddress(host: "localhost", port: 3000) != nil) + #expect(TransportSecuritySettings.forBindAddress(host: "::1", port: 9000) != nil) + + // Should return nil for other addresses (no protection needed) + #expect(TransportSecuritySettings.forBindAddress(host: "0.0.0.0", port: 8080) == nil) + #expect(TransportSecuritySettings.forBindAddress(host: "192.168.1.1", port: 8080) == nil) + } + + @Test("DNS rebinding protection rejects invalid host on GET") + func dnsRebindingProtectionRejectsInvalidHostOnGet() async throws { + let transport = HTTPServerTransport( + options: .init(security: .forLocalhost(port: 8080)) + ) + try await transport.connect() + + // Initialize first with valid host + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + HTTPHeader.host: "localhost:8080", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // GET with invalid host + let getRequest = HTTPRequest( + method: "GET", + headers: [ + HTTPHeader.accept: "text/event-stream", + HTTPHeader.host: "evil.com:8080", + ] + ) + let response = await transport.handleRequest(getRequest) + #expect(response.statusCode == 421) + } + + // MARK: - Protocol Version on GET/DELETE + + @Test("Rejects unsupported protocol version on GET") + func rejectsUnsupportedProtocolVersionOnGet() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { "test-session" }) + ) + try await transport.connect() + + // Initialize first + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // GET with unsupported protocol version + let getRequest = HTTPRequest( + method: "GET", + headers: [ + HTTPHeader.accept: "text/event-stream", + HTTPHeader.sessionId: "test-session", + HTTPHeader.protocolVersion: "1999-01-01", + ] + ) + let response = await transport.handleRequest(getRequest) + #expect(response.statusCode == 400) + } + + @Test("Rejects unsupported protocol version on DELETE") + func rejectsUnsupportedProtocolVersionOnDelete() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { "test-session" }) + ) + try await transport.connect() + + // Initialize first + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // DELETE with unsupported protocol version + let deleteRequest = HTTPRequest( + method: "DELETE", + headers: [ + HTTPHeader.sessionId: "test-session", + HTTPHeader.protocolVersion: "1999-01-01", + ] + ) + let response = await transport.handleRequest(deleteRequest) + #expect(response.statusCode == 400) + } + + @Test("Accepts requests without protocol version header") + func acceptsRequestsWithoutProtocolVersionHeader() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { "test-session" }) + ) + try await transport.connect() + + // Initialize + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // Request without protocol version header should work + let request = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + HTTPHeader.sessionId: "test-session", + // No protocol version header + ], + body: #"{"jsonrpc":"2.0","method":"tools/list","id":"2"}"#.data(using: .utf8) + ) + let response = await transport.handleRequest(request) + #expect(response.statusCode == 200) + } + + // MARK: - Session Closed Callback Edge Cases + + @Test("Session closed callback not called for invalid session") + func sessionClosedCallbackNotCalledForInvalidSession() async throws { + actor CallbackState { + var called = false + func markCalled() { called = true } + func wasCalled() -> Bool { called } + } + let state = CallbackState() + + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { "valid-session" }, + onSessionClosed: { _ in await state.markCalled() } + ) + ) + try await transport.connect() + + // Initialize + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // Try to DELETE with invalid session ID + let deleteRequest = HTTPRequest( + method: "DELETE", + headers: [HTTPHeader.sessionId: "invalid-session"] + ) + let response = await transport.handleRequest(deleteRequest) + + #expect(response.statusCode == 404) + let called = await state.wasCalled() + #expect(called == false) // Callback should NOT be called for invalid session + } + + @Test("DELETE without callback works") + func deleteWithoutCallbackWorks() async throws { + // No onSessionClosed callback provided + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { "test-session" }) + ) + try await transport.connect() + + // Initialize + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // DELETE should work without callback + let deleteRequest = HTTPRequest( + method: "DELETE", + headers: [HTTPHeader.sessionId: "test-session"] + ) + let response = await transport.handleRequest(deleteRequest) + #expect(response.statusCode == 200) + } + + // MARK: - Batch Request Handling + + @Test("Batch initialize request rejected") + func batchInitializeRequestRejected() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { UUID().uuidString }) + ) + try await transport.connect() + + // Batch with initialize messages should be rejected + let batchRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.batchRequest([ + TestPayloads.initializeRequest(id: "1"), + TestPayloads.initializeRequest(id: "2", clientName: "test2"), + ]) + .data(using: .utf8) + ) + let response = await transport.handleRequest(batchRequest) + #expect(response.statusCode == 400) + } + + // MARK: - Uninitialized Server Handling + + @Test("Rejects requests to uninitialized server") + func rejectsRequestsToUninitializedServer() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { "test-session" }) + ) + try await transport.connect() + + // Send a non-initialize request without first initializing + let request = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + HTTPHeader.sessionId: "test-session", + ], + body: #"{"jsonrpc":"2.0","method":"tools/list","id":"1"}"#.data(using: .utf8) + ) + + let response = await transport.handleRequest(request) + #expect(response.statusCode == 400) + + // Verify it's a JSON-RPC error with "not initialized" + if let body = response.body, let text = String(data: body, encoding: .utf8) { + #expect(text.lowercased().contains("not initialized")) + } + } + + @Test("Rejects GET requests to uninitialized server") + func rejectsGetRequestsToUninitializedServer() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { "test-session" }) + ) + try await transport.connect() + + // Send a GET request without first initializing + let request = HTTPRequest( + method: "GET", + headers: [ + HTTPHeader.accept: "text/event-stream", + HTTPHeader.sessionId: "test-session", + ] + ) + + let response = await transport.handleRequest(request) + #expect(response.statusCode == 400) + + // Verify it's a JSON-RPC error with "not initialized" + if let body = response.body, let text = String(data: body, encoding: .utf8) { + #expect(text.lowercased().contains("not initialized")) + } + } + + // MARK: - Stateless Mode + + @Test("Stateless mode accepts requests with any session ID") + func statelessModeAcceptsAnySessionId() async throws { + // No session ID generator = stateless mode + let transport = HTTPServerTransport() + try await transport.connect() + + // Initialize first + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + let initResponse = await transport.handleRequest(initRequest) + #expect(initResponse.statusCode == 200) + + // In stateless mode, requests with different session IDs should work + for sessionId in ["session-1", "session-2", "random-id", ""] { + var headers = [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ] + if !sessionId.isEmpty { + headers[HTTPHeader.sessionId] = sessionId + } + + let request = HTTPRequest( + method: "POST", + headers: headers, + body: #"{"jsonrpc":"2.0","method":"notifications/initialized"}"#.data(using: .utf8) + ) + + let response = await transport.handleRequest(request) + // Notifications should return 202 regardless of session ID in stateless mode + #expect(response.statusCode == 202) + } + } + + @Test("Stateless mode allows request without session ID") + func statelessModeAllowsRequestWithoutSessionId() async throws { + let transport = HTTPServerTransport() + try await transport.connect() + + // Initialize first + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // Request without session ID should work in stateless mode + let request = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + // No session ID header + ], + body: #"{"jsonrpc":"2.0","method":"notifications/initialized"}"#.data(using: .utf8) + ) + + let response = await transport.handleRequest(request) + #expect(response.statusCode == 202) + } + + @Test("Stateless mode rejects requests before initialization") + func statelessModeRejectsRequestsBeforeInit() async throws { + // No session ID generator = stateless mode + let transport = HTTPServerTransport() + try await transport.connect() + + // Try to send request without initializing first + let request = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: #"{"jsonrpc":"2.0","method":"tools/list","id":"1"}"#.data(using: .utf8) + ) + + let response = await transport.handleRequest(request) + #expect(response.statusCode == 400) + + // Verify error message mentions initialization + if let body = response.body, let text = String(data: body, encoding: .utf8) { + #expect(text.lowercased().contains("not initialized")) + } + } + + @Test("Stateless mode rejects GET requests before initialization") + func statelessModeRejectsGetBeforeInit() async throws { + let transport = HTTPServerTransport() + try await transport.connect() + + // Try GET without initializing first + let request = HTTPRequest( + method: "GET", + headers: [HTTPHeader.accept: "text/event-stream"] + ) + + let response = await transport.handleRequest(request) + #expect(response.statusCode == 400) + + if let body = response.body, let text = String(data: body, encoding: .utf8) { + #expect(text.lowercased().contains("not initialized")) + } + } + + // MARK: - JSON-RPC Error Code Validation + + @Test("Parse error returns -32700") + func parseErrorReturns32700() async throws { + let transport = HTTPServerTransport() + try await transport.connect() + + // Initialize first + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // Send invalid JSON + let request = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: "not valid json".data(using: .utf8) + ) + + let response = await transport.handleRequest(request) + #expect(response.statusCode == 400) + if let body = response.body, let text = String(data: body, encoding: .utf8) { + #expect(text.contains("\(ErrorCode.parseError)")) + } + } + + @Test("Invalid request returns -32600") + func invalidRequestReturns32600() async throws { + let transport = HTTPServerTransport() + try await transport.connect() + + // Initialize first + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // Send message missing jsonrpc version (invalid JSON-RPC format) + let request = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: #"{"method":"tools/list","params":{},"id":"2"}"#.data(using: .utf8) + ) + + let response = await transport.handleRequest(request) + #expect(response.statusCode == 400) + if let body = response.body, let text = String(data: body, encoding: .utf8) { + #expect(text.contains("\(ErrorCode.invalidRequest)")) + } + } + + @Test("Empty body returns parse error -32700") + func emptyBodyReturnsParseError() async throws { + let transport = HTTPServerTransport() + try await transport.connect() + + // Initialize first + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // Send empty body + let request = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: nil + ) + + let response = await transport.handleRequest(request) + #expect(response.statusCode == 400) + if let body = response.body, let text = String(data: body, encoding: .utf8) { + #expect(text.contains("\(ErrorCode.parseError)")) + #expect(text.contains("Empty request body")) + } + } + + @Test("Valid JSON but not JSON-RPC array returns parse error") + func validJsonNotJsonRpcReturnsParseError() async throws { + let transport = HTTPServerTransport() + try await transport.connect() + + // Initialize first + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // Send valid JSON but not a JSON-RPC message (not an object or array of objects) + let request = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: #"["string", 123, true]"#.data(using: .utf8) + ) + + let response = await transport.handleRequest(request) + #expect(response.statusCode == 400) + if let body = response.body, let text = String(data: body, encoding: .utf8) { + #expect(text.contains("\(ErrorCode.parseError)")) + } + } + + @Test("Wrong jsonrpc version returns invalid request error") + func wrongJsonRpcVersionReturnsInvalidRequest() async throws { + let transport = HTTPServerTransport() + try await transport.connect() + + // Initialize first + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // Send message with wrong jsonrpc version + let request = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: #"{"jsonrpc":"1.0","method":"tools/list","params":{},"id":"2"}"#.data(using: .utf8) + ) + + let response = await transport.handleRequest(request) + #expect(response.statusCode == 400) + if let body = response.body, let text = String(data: body, encoding: .utf8) { + #expect(text.contains("\(ErrorCode.invalidRequest)")) + #expect(text.lowercased().contains("jsonrpc")) + } + } + + // MARK: - Session ID Validation Tests (per Python/TypeScript SDK patterns) + + @Test("Valid session IDs accepted") + func validSessionIdsAccepted() async throws { + // Valid session IDs: visible ASCII (0x21-0x7E) + let validSessionIds = [ + "test-session-id", + "1234567890", + "session!@#$%^&*()_+-=[]{}|;:,.<>?/", + "~", // 0x7E + "!", // 0x21 + UUID().uuidString, + ] + + for validId in validSessionIds { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { validId }) + ) + try await transport.connect() + + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + + let response = await transport.handleRequest(initRequest) + #expect(response.statusCode == 200, "Session ID '\(validId)' should be accepted") + #expect(response.headers[HTTPHeader.sessionId] == validId) + } + } + + @Test("Invalid session IDs rejected - space") + func invalidSessionIdWithSpaceRejected() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { "session with space" }) + ) + try await transport.connect() + + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + + let response = await transport.handleRequest(initRequest) + #expect(response.statusCode == 500) // Internal error for invalid generated ID + } + + @Test("Invalid session IDs rejected - control characters") + func invalidSessionIdWithControlCharsRejected() async throws { + let invalidIds = [ + "session\twith\ttab", // Tab (0x09) + "session\nwith\nnewline", // Newline (0x0A) + "session\rwith\rcarriage", // Carriage return (0x0D) + "session\u{7F}with\u{7F}del", // DEL (0x7F) + "session\u{00}with\u{00}null", // NULL (0x00) + ] + + for invalidId in invalidIds { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { invalidId }) + ) + try await transport.connect() + + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + + let response = await transport.handleRequest(initRequest) + #expect(response.statusCode == 500, "Session ID with control chars should be rejected") + } + } + + @Test("Invalid session IDs rejected - empty string") + func invalidSessionIdEmptyRejected() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { "" }) + ) + try await transport.connect() + + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + + let response = await transport.handleRequest(initRequest) + #expect(response.statusCode == 500) // Empty session ID is invalid + } + + // MARK: - GET Priming Events Tests (per Python/TypeScript SDK patterns) + + @Test("GET stream receives priming event with event store") + func getStreamReceivesPrimingEventWithEventStore() async throws { + let eventStore = InMemoryEventStore() + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { "test-session" }, + eventStore: eventStore + ) + ) + try await transport.connect() + + // Initialize with protocol version that supports priming events + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest(protocolVersion: Version.v2025_11_25) + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // GET request should receive priming event + let getRequest = HTTPRequest( + method: "GET", + headers: [ + HTTPHeader.accept: "text/event-stream", + HTTPHeader.sessionId: "test-session", + HTTPHeader.protocolVersion: Version.v2025_11_25, + ] + ) + let response = await transport.handleRequest(getRequest) + + #expect(response.statusCode == 200) + #expect(response.stream != nil) + + // Verify priming event was stored + let eventCount = await eventStore.eventCount + #expect(eventCount >= 1) // At least one priming event + } + + @Test("GET stream does not receive priming event for old protocol version") + func getStreamNoPrimingEventForOldProtocol() async throws { + let eventStore = InMemoryEventStore() + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { "test-session" }, + eventStore: eventStore + ) + ) + try await transport.connect() + + // Initialize with old protocol version + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // GET request should NOT receive priming event for old protocol + let getRequest = HTTPRequest( + method: "GET", + headers: [ + HTTPHeader.accept: "text/event-stream", + HTTPHeader.sessionId: "test-session", + HTTPHeader.protocolVersion: Version.v2024_11_05, + ] + ) + let response = await transport.handleRequest(getRequest) + + #expect(response.statusCode == 200) + + // No priming event for old protocol + let eventCount = await eventStore.eventCount + #expect(eventCount == 0) + } + + @Test("GET stream does not receive priming event without event store") + func getStreamNoPrimingEventWithoutEventStore() async throws { + // No event store configured + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { "test-session" }) + ) + try await transport.connect() + + // Initialize + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest(protocolVersion: Version.v2025_11_25) + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // GET request should still work but no priming event + let getRequest = HTTPRequest( + method: "GET", + headers: [ + HTTPHeader.accept: "text/event-stream", + HTTPHeader.sessionId: "test-session", + ] + ) + let response = await transport.handleRequest(getRequest) + + #expect(response.statusCode == 200) + #expect(response.stream != nil) + } + + @Test("Priming event includes retry interval when configured") + func primingEventIncludesRetryInterval() async throws { + let eventStore = InMemoryEventStore() + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { "test-session" }, + eventStore: eventStore, + retryInterval: 5000 // 5 seconds + ) + ) + try await transport.connect() + + // Initialize with new protocol version + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest(protocolVersion: Version.v2025_11_25) + .data(using: .utf8) + ) + let response = await transport.handleRequest(initRequest) + + #expect(response.statusCode == 200) + #expect(response.stream != nil) + + // Read first event from stream to check for retry field + if let stream = response.stream { + var receivedData = Data() + for try await chunk in stream { + receivedData.append(chunk) + break // Just get first chunk (priming event) + } + let eventString = String(data: receivedData, encoding: .utf8) ?? "" + #expect(eventString.contains("retry: 5000")) + } + } + + // MARK: - Resumability Tests (per Python/TypeScript SDK patterns) + + @Test("Event store stores events with unique IDs") + func eventStoreStoresEventsWithUniqueIds() async throws { + let eventStore = InMemoryEventStore() + + let id1 = try await eventStore.storeEvent(streamId: "stream-1", message: Data("msg1".utf8)) + let id2 = try await eventStore.storeEvent(streamId: "stream-1", message: Data("msg2".utf8)) + let id3 = try await eventStore.storeEvent(streamId: "stream-2", message: Data("msg3".utf8)) + + #expect(id1 != id2) + #expect(id2 != id3) + #expect(id1 != id3) + + let eventCount = await eventStore.eventCount + #expect(eventCount == 3) + } + + @Test("Event store replays events after last event ID") + func eventStoreReplaysEventsAfterLastId() async throws { + actor MessageCollector { + var messages: [String] = [] + func append(_ text: String) { messages.append(text) } + func getMessages() -> [String] { messages } + } + let collector = MessageCollector() + let eventStore = InMemoryEventStore() + + let id1 = try await eventStore.storeEvent(streamId: "stream-1", message: Data("msg1".utf8)) + _ = try await eventStore.storeEvent(streamId: "stream-1", message: Data("msg2".utf8)) + _ = try await eventStore.storeEvent(streamId: "stream-1", message: Data("msg3".utf8)) + + let replayedStreamId = try await eventStore.replayEventsAfter(id1) { _, message in + if let text = String(data: message, encoding: .utf8) { + await collector.append(text) + } + } + + let replayedMessages = await collector.getMessages() + #expect(replayedStreamId == "stream-1") + #expect(replayedMessages.count == 2) // msg2 and msg3, not msg1 + #expect(replayedMessages.contains("msg2")) + #expect(replayedMessages.contains("msg3")) + #expect(!replayedMessages.contains("msg1")) + } + + @Test("Event store only replays events from same stream") + func eventStoreOnlyReplaysFromSameStream() async throws { + actor MessageCollector { + var messages: [String] = [] + func append(_ text: String) { messages.append(text) } + func getMessages() -> [String] { messages } + } + let collector = MessageCollector() + let eventStore = InMemoryEventStore() + + let id1 = try await eventStore.storeEvent(streamId: "stream-1", message: Data("stream1-msg1".utf8)) + _ = try await eventStore.storeEvent(streamId: "stream-2", message: Data("stream2-msg1".utf8)) + _ = try await eventStore.storeEvent(streamId: "stream-1", message: Data("stream1-msg2".utf8)) + + _ = try await eventStore.replayEventsAfter(id1) { _, message in + if let text = String(data: message, encoding: .utf8) { + await collector.append(text) + } + } + + let replayedMessages = await collector.getMessages() + // Should only replay stream-1 messages after id1 + #expect(replayedMessages.count == 1) + #expect(replayedMessages.contains("stream1-msg2")) + #expect(!replayedMessages.contains("stream2-msg1")) + } + + @Test("GET with Last-Event-ID replays events") + func getWithLastEventIdReplaysEvents() async throws { + let eventStore = InMemoryEventStore() + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { "test-session" }, + eventStore: eventStore + ) + ) + try await transport.connect() + + // Initialize + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest(protocolVersion: Version.v2025_11_25) + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // Store some events directly + let eventId = try await eventStore.storeEvent( + streamId: "_GET_stream", + message: Data(#"{"jsonrpc":"2.0","method":"test"}"#.utf8) + ) + + // GET with Last-Event-ID should trigger replay + let getRequest = HTTPRequest( + method: "GET", + headers: [ + HTTPHeader.accept: "text/event-stream", + HTTPHeader.sessionId: "test-session", + HTTPHeader.lastEventId: eventId, + ] + ) + let response = await transport.handleRequest(getRequest) + + // Should return 200 with stream (replay mode) + #expect(response.statusCode == 200) + #expect(response.stream != nil) + } + + // MARK: - Protocol Version Negotiation Tests + + @Test("Protocol version stored after initialization") + func protocolVersionStoredAfterInit() async throws { + let eventStore = InMemoryEventStore() + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { "test-session" }, + eventStore: eventStore + ) + ) + try await transport.connect() + + // Initialize with 2025-11-25 + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest(protocolVersion: Version.v2025_11_25) + .data(using: .utf8) + ) + let initResponse = await transport.handleRequest(initRequest) + #expect(initResponse.statusCode == 200) + + // Priming event should be stored (only for >= 2025-11-25) + let eventCount = await eventStore.eventCount + #expect(eventCount >= 1) + } + + // MARK: - Cache-Control Header Tests + + @Test("POST SSE response has correct Cache-Control header") + func postSseResponseHasCorrectCacheControl() async throws { + let transport = HTTPServerTransport() + try await transport.connect() + + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + + let response = await transport.handleRequest(initRequest) + #expect(response.statusCode == 200) + #expect(response.headers[HTTPHeader.cacheControl] == "no-cache, no-transform") + } + + @Test("GET SSE response has correct Cache-Control header") + func getSseResponseHasCorrectCacheControl() async throws { + let transport = HTTPServerTransport() + try await transport.connect() + + // Initialize first + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + let getRequest = HTTPRequest( + method: "GET", + headers: [HTTPHeader.accept: "text/event-stream"] + ) + let response = await transport.handleRequest(getRequest) + + #expect(response.statusCode == 200) + #expect(response.headers[HTTPHeader.cacheControl] == "no-cache, no-transform") + } + + // MARK: - Session Callback Tests (per Python/TypeScript SDK patterns) + + @Test("Session initialized callback called with session ID") + func sessionInitializedCallbackCalled() async throws { + actor SessionTracker { + var sessionId: String? + func set(_ id: String) { sessionId = id } + func get() -> String? { sessionId } + } + let tracker = SessionTracker() + + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { "callback-test-session" }, + onSessionInitialized: { sessionId in + await tracker.set(sessionId) + } + ) + ) + try await transport.connect() + + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + + _ = await transport.handleRequest(initRequest) + + let capturedId = await tracker.get() + #expect(capturedId == "callback-test-session") + } + + @Test("Session closed callback called on DELETE") + func sessionClosedCallbackCalledOnDelete() async throws { + actor SessionTracker { + var closedSessionId: String? + func setClosed(_ id: String) { closedSessionId = id } + func getClosed() -> String? { closedSessionId } + } + let tracker = SessionTracker() + + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { "close-test-session" }, + onSessionClosed: { sessionId in + await tracker.setClosed(sessionId) + } + ) + ) + try await transport.connect() + + // Initialize + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // DELETE + let deleteRequest = HTTPRequest( + method: "DELETE", + headers: [HTTPHeader.sessionId: "close-test-session"] + ) + _ = await transport.handleRequest(deleteRequest) + + let closedId = await tracker.getClosed() + #expect(closedId == "close-test-session") + } + + @Test("Session closed callback not invoked for invalid session DELETE") + func sessionClosedCallbackNotInvokedForInvalidSessionDelete() async throws { + actor CallCounter { + var count = 0 + func increment() { count += 1 } + func getCount() -> Int { count } + } + let counter = CallCounter() + + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { "valid-session" }, + onSessionClosed: { _ in + await counter.increment() + } + ) + ) + try await transport.connect() + + // Initialize + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // DELETE with wrong session ID + let deleteRequest = HTTPRequest( + method: "DELETE", + headers: [HTTPHeader.sessionId: "wrong-session"] + ) + let response = await transport.handleRequest(deleteRequest) + + #expect(response.statusCode == 404) + let callCount = await counter.getCount() + #expect(callCount == 0) // Callback should NOT be called + } + + // MARK: - Terminated State Tests + + @Test("Terminated stateless transport returns 404") + func terminatedStatelessTransportReturns404() async throws { + // Stateless mode + let transport = HTTPServerTransport() + try await transport.connect() + + // Initialize + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // Manually disconnect/close the transport + await transport.disconnect() + + // Any subsequent request should return 404 + let postRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: #"{"jsonrpc":"2.0","method":"tools/list","id":"2"}"#.data(using: .utf8) + ) + let postResponse = await transport.handleRequest(postRequest) + #expect(postResponse.statusCode == 404) + + // GET should also return 404 + let getRequest = HTTPRequest( + method: "GET", + headers: [HTTPHeader.accept: "text/event-stream"] + ) + let getResponse = await transport.handleRequest(getRequest) + #expect(getResponse.statusCode == 404) + + // Even initialize should return 404 after termination + let reInitRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest(id: "3") + .data(using: .utf8) + ) + let reInitResponse = await transport.handleRequest(reInitRequest) + #expect(reInitResponse.statusCode == 404) + } + + @Test("Terminated stateful transport returns 404 for all requests") + func terminatedStatefulTransportReturns404ForAll() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { "test-session" }) + ) + try await transport.connect() + + // Initialize + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // Close the transport directly (simulating server shutdown) + await transport.close() + + // Any request should now return 404 - even with correct session ID + let postRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + HTTPHeader.sessionId: "test-session", + ], + body: #"{"jsonrpc":"2.0","method":"tools/list","id":"2"}"#.data(using: .utf8) + ) + let postResponse = await transport.handleRequest(postRequest) + #expect(postResponse.statusCode == 404) + + // Verify the error message indicates termination + if let body = postResponse.body, let text = String(data: body, encoding: .utf8) { + #expect(text.contains("terminated")) + } + } + + // MARK: - Stream Resumability Tests (Per MCP Spec) + + @Test("POST-initiated stream can be resumed via GET with Last-Event-ID") + func postInitiatedStreamCanBeResumedViaGet() async throws { + // Per spec: "This mechanism applies regardless of how the original SSE stream + // was initiated—even if a stream was originally started for a specific client + // request (via HTTP POST), the client will resume it via HTTP GET." + let eventStore = InMemoryEventStore() + let transport = HTTPServerTransport( + options: .init( + sessionIdGenerator: { "test-session" }, + eventStore: eventStore + ) + ) + try await transport.connect() + + // Initialize with protocol version 2025-11-25 (supports priming events) + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest(id: "init", protocolVersion: Version.v2025_11_25) + .data(using: .utf8) + ) + let initResponse = await transport.handleRequest(initRequest) + #expect(initResponse.statusCode == 200) + + // Send a POST request that starts an SSE stream + let postRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + HTTPHeader.sessionId: "test-session", + HTTPHeader.protocolVersion: Version.v2025_11_25, + ], + body: #"{"jsonrpc":"2.0","method":"tools/list","id":"req-1"}"#.data(using: .utf8) + ) + let postResponse = await transport.handleRequest(postRequest) + #expect(postResponse.statusCode == 200) + #expect(postResponse.stream != nil) + + // The event store should have events from the POST stream (priming event at minimum) + let eventCount = await eventStore.eventCount + #expect(eventCount >= 1, "Event store should have at least one event from POST stream") + + // Get an event ID from the store to simulate client reconnection + // We need to manually store an event to get an ID for the POST stream + let testStreamId = "test-stream-\(UUID().uuidString)" + let eventId = try await eventStore.storeEvent( + streamId: testStreamId, + message: #"{"jsonrpc":"2.0","result":{"tools":[]},"id":"req-1"}"#.data(using: .utf8)! + ) + + // Client reconnects via GET with Last-Event-ID (even though original was POST) + let getRequest = HTTPRequest( + method: "GET", + headers: [ + HTTPHeader.accept: "text/event-stream", + HTTPHeader.sessionId: "test-session", + HTTPHeader.protocolVersion: Version.v2025_11_25, + HTTPHeader.lastEventId: eventId, + ] + ) + let getResponse = await transport.handleRequest(getRequest) + + // Should return 200 with stream (resumption mode) + #expect(getResponse.statusCode == 200) + #expect(getResponse.stream != nil) + #expect(getResponse.headers[HTTPHeader.contentType] == "text/event-stream") + } + + @Test("Cross-stream event isolation during replay") + func crossStreamEventIsolationDuringReplay() async throws { + // Per spec: "Server MUST NOT replay messages delivered on a different stream" + let eventStore = InMemoryEventStore() + + // Store events for two different streams + let stream1Id = "stream-1" + let stream2Id = "stream-2" + + let event1_1 = try await eventStore.storeEvent( + streamId: stream1Id, + message: #"{"jsonrpc":"2.0","result":"stream1-msg1","id":"1"}"#.data(using: .utf8)! + ) + _ = try await eventStore.storeEvent( + streamId: stream2Id, + message: #"{"jsonrpc":"2.0","result":"stream2-msg1","id":"2"}"#.data(using: .utf8)! + ) + _ = try await eventStore.storeEvent( + streamId: stream1Id, + message: #"{"jsonrpc":"2.0","result":"stream1-msg2","id":"3"}"#.data(using: .utf8)! + ) + _ = try await eventStore.storeEvent( + streamId: stream2Id, + message: #"{"jsonrpc":"2.0","result":"stream2-msg2","id":"4"}"#.data(using: .utf8)! + ) + + // Replay events after event1_1 (should only get stream1 events) + actor MessageCollector { + var messages: [String] = [] + func append(_ text: String) { messages.append(text) } + func getMessages() -> [String] { messages } + } + let collector = MessageCollector() + + let replayedStreamId = try await eventStore.replayEventsAfter(event1_1) { _, message in + if let text = String(data: message, encoding: .utf8) { + await collector.append(text) + } + } + + // Should only replay stream-1 events + #expect(replayedStreamId == stream1Id) + let replayedMessages = await collector.getMessages() + #expect(replayedMessages.count == 1, "Should only replay one event from stream-1") + #expect(replayedMessages.contains { $0.contains("stream1-msg2") }, "Should contain stream1-msg2") + #expect(!replayedMessages.contains { $0.contains("stream2") }, "Should NOT contain any stream-2 events") + } + + @Test("Client sending JSON-RPC response returns 202 Accepted") + func clientSendingResponseReturns202() async throws { + // Per spec: For JSON-RPC response or notification input, + // server returns 202 Accepted with no body + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { "test-session" }) + ) + try await transport.connect() + + // Initialize + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // Client sends a JSON-RPC response (e.g., in reply to a sampling request from server) + // A response has "result" or "error" but no "method" + let responseRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + HTTPHeader.sessionId: "test-session", + ], + body: #"{"jsonrpc":"2.0","result":{"content":[{"type":"text","text":"Sample response"}]},"id":"server-req-1"}"#.data(using: .utf8) + ) + let response = await transport.handleRequest(responseRequest) + + // Per spec, responses should return 202 Accepted + #expect(response.statusCode == 202) + #expect(response.body == nil || response.body?.isEmpty == true, "202 response should have no body") + } + + @Test("Client sending JSON-RPC error response returns 202 Accepted") + func clientSendingErrorResponseReturns202() async throws { + // Per spec: For JSON-RPC response (including error responses) input, + // server returns 202 Accepted + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { "test-session" }) + ) + try await transport.connect() + + // Initialize + let initRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + ], + body: + TestPayloads.initializeRequest() + .data(using: .utf8) + ) + _ = await transport.handleRequest(initRequest) + + // Client sends a JSON-RPC error response (e.g., rejecting a sampling request) + let errorResponseRequest = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + HTTPHeader.sessionId: "test-session", + ], + body: #"{"jsonrpc":"2.0","error":{"code":-32600,"message":"User rejected sampling request"},"id":"server-req-1"}"#.data(using: .utf8) + ) + let response = await transport.handleRequest(errorResponseRequest) + + // Per spec, error responses should also return 202 Accepted + #expect(response.statusCode == 202) + } + + @Test("Priming events are not replayed during resumption") + func primingEventsNotReplayedDuringResumption() async throws { + // Per spec and InMemoryEventStore design: priming events (empty data) + // should be stored but NOT replayed as regular messages + let eventStore = InMemoryEventStore() + let streamId = "test-stream" + + // Store a priming event (empty data) + let primingEventId = try await eventStore.storeEvent(streamId: streamId, message: Data()) + + // Store a regular message + _ = try await eventStore.storeEvent( + streamId: streamId, + message: #"{"jsonrpc":"2.0","result":"test","id":"1"}"#.data(using: .utf8)! + ) + + // Replay from before priming event (using a fake "before" ID approach) + // We'll replay from the priming event itself + actor MessageCollector { + var messages: [Data] = [] + func append(_ data: Data) { messages.append(data) } + func getMessages() -> [Data] { messages } + } + let collector = MessageCollector() + + _ = try await eventStore.replayEventsAfter(primingEventId) { _, message in + await collector.append(message) + } + + let replayedMessages = await collector.getMessages() + + // Should only have 1 message (the regular one), not the priming event + #expect(replayedMessages.count == 1, "Should only replay regular messages, not priming events") + #expect(!replayedMessages.contains { $0.isEmpty }, "Should not contain empty (priming) events") + } +} diff --git a/Tests/MCPTests/TaskTests.swift b/Tests/MCPTests/TaskTests.swift new file mode 100644 index 00000000..d46de8ec --- /dev/null +++ b/Tests/MCPTests/TaskTests.swift @@ -0,0 +1,1187 @@ +import Foundation +import Testing + +@testable import MCP + +// MARK: - Task Type Tests + +@Suite("Task Type Tests") +struct TaskTypeTests { + + // MARK: - TaskStatus Tests + + @Test( + "TaskStatus raw values match spec", + arguments: [ + (TaskStatus.working, "working"), + (TaskStatus.inputRequired, "input_required"), + (TaskStatus.completed, "completed"), + (TaskStatus.failed, "failed"), + (TaskStatus.cancelled, "cancelled"), + ] + ) + func taskStatusRawValues(testCase: (status: TaskStatus, rawValue: String)) { + #expect(testCase.status.rawValue == testCase.rawValue) + } + + @Test( + "TaskStatus.isTerminal returns correct values", + arguments: [ + (TaskStatus.working, false), + (TaskStatus.inputRequired, false), + (TaskStatus.completed, true), + (TaskStatus.failed, true), + (TaskStatus.cancelled, true), + ] + ) + func taskStatusIsTerminal(testCase: (status: TaskStatus, isTerminal: Bool)) { + #expect(testCase.status.isTerminal == testCase.isTerminal) + } + + @Test("isTerminalStatus helper function matches TaskStatus.isTerminal") + func isTerminalStatusHelperFunction() { + #expect(isTerminalStatus(.working) == false) + #expect(isTerminalStatus(.inputRequired) == false) + #expect(isTerminalStatus(.completed) == true) + #expect(isTerminalStatus(.failed) == true) + #expect(isTerminalStatus(.cancelled) == true) + } + + @Test("TaskStatus encodes and decodes correctly") + func taskStatusEncodingDecoding() throws { + let statuses: [TaskStatus] = [.working, .inputRequired, .completed, .failed, .cancelled] + + for status in statuses { + let data = try JSONEncoder().encode(status) + let decoded = try JSONDecoder().decode(TaskStatus.self, from: data) + #expect(decoded == status) + } + } + + // MARK: - MCPTask Tests + + @Test("MCPTask encoding and decoding with all fields") + func mcpTaskFullEncodingDecoding() throws { + let task = MCPTask( + taskId: "task-123", + status: .working, + ttl: 60000, + createdAt: "2024-01-15T10:30:00Z", + lastUpdatedAt: "2024-01-15T10:30:05Z", + pollInterval: 1000, + statusMessage: "Processing..." + ) + + let data = try JSONEncoder().encode(task) + let decoded = try JSONDecoder().decode(MCPTask.self, from: data) + + #expect(decoded.taskId == "task-123") + #expect(decoded.status == .working) + #expect(decoded.ttl == 60000) + #expect(decoded.createdAt == "2024-01-15T10:30:00Z") + #expect(decoded.lastUpdatedAt == "2024-01-15T10:30:05Z") + #expect(decoded.pollInterval == 1000) + #expect(decoded.statusMessage == "Processing...") + } + + @Test("MCPTask with nil ttl encodes as null (per spec requirement)") + func mcpTaskNilTtlEncodesAsNull() throws { + let task = MCPTask( + taskId: "task-123", + status: .working, + ttl: nil, + createdAt: "2024-01-15T10:30:00Z", + lastUpdatedAt: "2024-01-15T10:30:05Z" + ) + + let data = try JSONEncoder().encode(task) + let jsonString = String(data: data, encoding: .utf8)! + + // Per MCP spec, ttl must always be present (encoded as null when nil) + #expect(jsonString.contains("\"ttl\":null")) + } + + @Test("MCPTask decodes ttl as null correctly") + func mcpTaskDecodesNullTtl() throws { + let jsonString = """ + { + "taskId": "task-123", + "status": "working", + "ttl": null, + "createdAt": "2024-01-15T10:30:00Z", + "lastUpdatedAt": "2024-01-15T10:30:05Z" + } + """ + + let data = jsonString.data(using: .utf8)! + let task = try JSONDecoder().decode(MCPTask.self, from: data) + + #expect(task.ttl == nil) + } + + @Test("MCPTask with optional fields omitted") + func mcpTaskOptionalFieldsOmitted() throws { + let task = MCPTask( + taskId: "task-123", + status: .completed, + ttl: 30000, + createdAt: "2024-01-15T10:30:00Z", + lastUpdatedAt: "2024-01-15T10:30:10Z" + ) + + let data = try JSONEncoder().encode(task) + let jsonString = String(data: data, encoding: .utf8)! + + // Optional fields should not be present + #expect(!jsonString.contains("pollInterval")) + #expect(!jsonString.contains("statusMessage")) + } + + // MARK: - TaskMetadata Tests + + @Test("TaskMetadata encoding and decoding") + func taskMetadataEncodingDecoding() throws { + let metadata = TaskMetadata(ttl: 60000) + + let data = try JSONEncoder().encode(metadata) + let decoded = try JSONDecoder().decode(TaskMetadata.self, from: data) + + #expect(decoded.ttl == 60000) + } + + @Test("TaskMetadata with nil ttl") + func taskMetadataNilTtl() throws { + let metadata = TaskMetadata(ttl: nil) + + let data = try JSONEncoder().encode(metadata) + let decoded = try JSONDecoder().decode(TaskMetadata.self, from: data) + + #expect(decoded.ttl == nil) + } + + // MARK: - RelatedTaskMetadata Tests + + @Test("RelatedTaskMetadata encoding and decoding") + func relatedTaskMetadataEncodingDecoding() throws { + let metadata = RelatedTaskMetadata(taskId: "task-456") + + let data = try JSONEncoder().encode(metadata) + let decoded = try JSONDecoder().decode(RelatedTaskMetadata.self, from: data) + + #expect(decoded.taskId == "task-456") + } + + // MARK: - Metadata Key Tests + + @Test("relatedTaskMetaKey has correct value") + func relatedTaskMetaKeyValue() { + #expect(relatedTaskMetaKey == "io.modelcontextprotocol/related-task") + } + + @Test("modelImmediateResponseKey has correct value") + func modelImmediateResponseKeyValue() { + #expect(modelImmediateResponseKey == "io.modelcontextprotocol/model-immediate-response") + } +} + +// MARK: - CreateTaskResult Tests + +@Suite("CreateTaskResult Tests") +struct CreateTaskResultTests { + + @Test("CreateTaskResult encoding and decoding") + func createTaskResultEncodingDecoding() throws { + let task = MCPTask( + taskId: "task-123", + status: .working, + ttl: 60000, + createdAt: "2024-01-15T10:30:00Z", + lastUpdatedAt: "2024-01-15T10:30:00Z", + pollInterval: 1000 + ) + + let result = CreateTaskResult(task: task) + + let data = try JSONEncoder().encode(result) + let decoded = try JSONDecoder().decode(CreateTaskResult.self, from: data) + + #expect(decoded.task.taskId == "task-123") + #expect(decoded.task.status == .working) + #expect(decoded.task.ttl == 60000) + #expect(decoded.task.pollInterval == 1000) + } + + @Test("CreateTaskResult with model immediate response") + func createTaskResultWithModelImmediateResponse() throws { + let task = MCPTask( + taskId: "task-123", + status: .working, + ttl: nil, + createdAt: "2024-01-15T10:30:00Z", + lastUpdatedAt: "2024-01-15T10:30:00Z" + ) + + let result = CreateTaskResult(task: task, modelImmediateResponse: "Starting task...") + + let data = try JSONEncoder().encode(result) + let decoded = try JSONDecoder().decode(CreateTaskResult.self, from: data) + + #expect(decoded._meta?[modelImmediateResponseKey]?.stringValue == "Starting task...") + } + + @Test("CreateTaskResult with _meta") + func createTaskResultWithMeta() throws { + let task = MCPTask( + taskId: "task-123", + status: .working, + ttl: nil, + createdAt: "2024-01-15T10:30:00Z", + lastUpdatedAt: "2024-01-15T10:30:00Z" + ) + + let meta: [String: Value] = [ + "custom": .string("value"), + modelImmediateResponseKey: .string("Processing your request..."), + ] + + let result = CreateTaskResult(task: task, _meta: meta) + + let data = try JSONEncoder().encode(result) + let decoded = try JSONDecoder().decode(CreateTaskResult.self, from: data) + + #expect(decoded._meta?["custom"]?.stringValue == "value") + #expect(decoded._meta?[modelImmediateResponseKey]?.stringValue == "Processing your request...") + } +} + +// MARK: - GetTask Tests + +@Suite("GetTask Method Tests") +struct GetTaskMethodTests { + + @Test("GetTask.name is correct") + func getTaskMethodName() { + #expect(GetTask.name == "tasks/get") + } + + @Test("GetTask.Parameters encoding and decoding") + func getTaskParametersEncodingDecoding() throws { + let params = GetTask.Parameters(taskId: "task-123") + + let data = try JSONEncoder().encode(params) + let decoded = try JSONDecoder().decode(GetTask.Parameters.self, from: data) + + #expect(decoded.taskId == "task-123") + } + + @Test("GetTask.Result encoding and decoding") + func getTaskResultEncodingDecoding() throws { + let result = GetTask.Result( + taskId: "task-123", + status: .completed, + ttl: 60000, + createdAt: "2024-01-15T10:30:00Z", + lastUpdatedAt: "2024-01-15T10:30:30Z", + pollInterval: 1000, + statusMessage: "Done" + ) + + let data = try JSONEncoder().encode(result) + let decoded = try JSONDecoder().decode(GetTask.Result.self, from: data) + + #expect(decoded.taskId == "task-123") + #expect(decoded.status == .completed) + #expect(decoded.ttl == 60000) + #expect(decoded.pollInterval == 1000) + #expect(decoded.statusMessage == "Done") + } + + @Test("GetTask.Result from MCPTask") + func getTaskResultFromMCPTask() throws { + let task = MCPTask( + taskId: "task-456", + status: .failed, + ttl: nil, + createdAt: "2024-01-15T10:30:00Z", + lastUpdatedAt: "2024-01-15T10:30:15Z", + statusMessage: "Connection timeout" + ) + + let result = GetTask.Result(task: task) + + #expect(result.taskId == task.taskId) + #expect(result.status == task.status) + #expect(result.statusMessage == "Connection timeout") + } + + @Test("GetTask.Result ttl encodes as null when nil") + func getTaskResultTtlEncodesAsNull() throws { + let result = GetTask.Result( + taskId: "task-123", + status: .working, + ttl: nil, + createdAt: "2024-01-15T10:30:00Z", + lastUpdatedAt: "2024-01-15T10:30:00Z" + ) + + let data = try JSONEncoder().encode(result) + let jsonString = String(data: data, encoding: .utf8)! + + #expect(jsonString.contains("\"ttl\":null")) + } +} + +// MARK: - GetTaskPayload Tests + +@Suite("GetTaskPayload Method Tests") +struct GetTaskPayloadMethodTests { + + @Test("GetTaskPayload.name is correct") + func getTaskPayloadMethodName() { + #expect(GetTaskPayload.name == "tasks/result") + } + + @Test("GetTaskPayload.Parameters encoding and decoding") + func getTaskPayloadParametersEncodingDecoding() throws { + let params = GetTaskPayload.Parameters(taskId: "task-123") + + let data = try JSONEncoder().encode(params) + let decoded = try JSONDecoder().decode(GetTaskPayload.Parameters.self, from: data) + + #expect(decoded.taskId == "task-123") + } + + @Test("GetTaskPayload.Result with extraFields (flattened result)") + func getTaskPayloadResultWithExtraFields() throws { + // Simulate a tools/call result flattened into extraFields + let extraFields: [String: Value] = [ + "content": .array([ + .object([ + "type": .string("text"), + "text": .string("Hello, world!"), + ]) + ]), + "isError": .bool(false), + ] + + let meta: [String: Value] = [ + relatedTaskMetaKey: .object(["taskId": .string("task-123")]) + ] + + let result = GetTaskPayload.Result(_meta: meta, extraFields: extraFields) + + let data = try JSONEncoder().encode(result) + let decoded = try JSONDecoder().decode(GetTaskPayload.Result.self, from: data) + + #expect(decoded._meta?[relatedTaskMetaKey] != nil) + #expect(decoded.extraFields?["isError"]?.boolValue == false) + } + + @Test("GetTaskPayload.Result fromResultValue convenience initializer") + func getTaskPayloadResultFromResultValue() throws { + let resultValue: Value = .object([ + "content": .array([.object(["type": .string("text"), "text": .string("Result")])]), + "isError": .bool(false), + ]) + + let result = GetTaskPayload.Result(fromResultValue: resultValue) + + #expect(result.extraFields?["isError"]?.boolValue == false) + } +} + +// MARK: - ListTasks Tests + +@Suite("ListTasks Method Tests") +struct ListTasksMethodTests { + + @Test("ListTasks.name is correct") + func listTasksMethodName() { + #expect(ListTasks.name == "tasks/list") + } + + @Test("ListTasks.Parameters encoding and decoding with cursor") + func listTasksParametersWithCursor() throws { + let params = ListTasks.Parameters(cursor: "page-2-token") + + let data = try JSONEncoder().encode(params) + let decoded = try JSONDecoder().decode(ListTasks.Parameters.self, from: data) + + #expect(decoded.cursor == "page-2-token") + } + + @Test("ListTasks.Parameters empty initializer") + func listTasksParametersEmpty() throws { + let params = ListTasks.Parameters() + + #expect(params.cursor == nil) + #expect(params._meta == nil) + } + + @Test("ListTasks.Result encoding and decoding") + func listTasksResultEncodingDecoding() throws { + let tasks = [ + MCPTask( + taskId: "task-1", + status: .completed, + ttl: nil, + createdAt: "2024-01-15T10:00:00Z", + lastUpdatedAt: "2024-01-15T10:05:00Z" + ), + MCPTask( + taskId: "task-2", + status: .working, + ttl: 60000, + createdAt: "2024-01-15T10:10:00Z", + lastUpdatedAt: "2024-01-15T10:10:00Z" + ), + ] + + let result = ListTasks.Result(tasks: tasks, nextCursor: "page-2") + + let data = try JSONEncoder().encode(result) + let decoded = try JSONDecoder().decode(ListTasks.Result.self, from: data) + + #expect(decoded.tasks.count == 2) + #expect(decoded.tasks[0].taskId == "task-1") + #expect(decoded.tasks[1].taskId == "task-2") + #expect(decoded.nextCursor == "page-2") + } + + @Test("ListTasks.Result without nextCursor indicates end of pagination") + func listTasksResultWithoutNextCursor() throws { + let result = ListTasks.Result(tasks: [], nextCursor: nil) + + let data = try JSONEncoder().encode(result) + let jsonString = String(data: data, encoding: .utf8)! + + #expect(!jsonString.contains("nextCursor")) + } +} + +// MARK: - CancelTask Tests + +@Suite("CancelTask Method Tests") +struct CancelTaskMethodTests { + + @Test("CancelTask.name is correct") + func cancelTaskMethodName() { + #expect(CancelTask.name == "tasks/cancel") + } + + @Test("CancelTask.Parameters encoding and decoding") + func cancelTaskParametersEncodingDecoding() throws { + let params = CancelTask.Parameters(taskId: "task-123") + + let data = try JSONEncoder().encode(params) + let decoded = try JSONDecoder().decode(CancelTask.Parameters.self, from: data) + + #expect(decoded.taskId == "task-123") + } + + @Test("CancelTask.Result encoding and decoding") + func cancelTaskResultEncodingDecoding() throws { + let result = CancelTask.Result( + taskId: "task-123", + status: .cancelled, + ttl: nil, + createdAt: "2024-01-15T10:30:00Z", + lastUpdatedAt: "2024-01-15T10:30:45Z", + statusMessage: "Cancelled by user" + ) + + let data = try JSONEncoder().encode(result) + let decoded = try JSONDecoder().decode(CancelTask.Result.self, from: data) + + #expect(decoded.taskId == "task-123") + #expect(decoded.status == .cancelled) + #expect(decoded.statusMessage == "Cancelled by user") + } + + @Test("CancelTask.Result from MCPTask") + func cancelTaskResultFromMCPTask() throws { + let task = MCPTask( + taskId: "task-456", + status: .cancelled, + ttl: nil, + createdAt: "2024-01-15T10:30:00Z", + lastUpdatedAt: "2024-01-15T10:30:30Z" + ) + + let result = CancelTask.Result(task: task) + + #expect(result.taskId == task.taskId) + #expect(result.status == .cancelled) + } +} + +// MARK: - TaskStatusNotification Tests + +@Suite("TaskStatusNotification Tests") +struct TaskStatusNotificationTests { + + @Test("TaskStatusNotification.name is correct") + func taskStatusNotificationName() { + #expect(TaskStatusNotification.name == "notifications/tasks/status") + } + + @Test("TaskStatusNotification.Parameters encoding and decoding") + func taskStatusNotificationParametersEncodingDecoding() throws { + let params = TaskStatusNotification.Parameters( + taskId: "task-123", + status: .inputRequired, + ttl: 60000, + createdAt: "2024-01-15T10:30:00Z", + lastUpdatedAt: "2024-01-15T10:30:10Z", + pollInterval: 500, + statusMessage: "Waiting for user input" + ) + + let data = try JSONEncoder().encode(params) + let decoded = try JSONDecoder().decode(TaskStatusNotification.Parameters.self, from: data) + + #expect(decoded.taskId == "task-123") + #expect(decoded.status == .inputRequired) + #expect(decoded.ttl == 60000) + #expect(decoded.pollInterval == 500) + #expect(decoded.statusMessage == "Waiting for user input") + } + + @Test("TaskStatusNotification.Parameters from MCPTask") + func taskStatusNotificationFromMCPTask() { + let task = MCPTask( + taskId: "task-789", + status: .completed, + ttl: nil, + createdAt: "2024-01-15T10:30:00Z", + lastUpdatedAt: "2024-01-15T10:31:00Z" + ) + + let params = TaskStatusNotification.Parameters(task: task) + + #expect(params.taskId == task.taskId) + #expect(params.status == task.status) + #expect(params.createdAt == task.createdAt) + #expect(params.lastUpdatedAt == task.lastUpdatedAt) + } +} + +// MARK: - Server Capabilities Tests + +@Suite("Server Tasks Capabilities Tests") +struct ServerTasksCapabilitiesTests { + + @Test("Server.Capabilities.Tasks encoding and decoding") + func serverTasksCapabilitiesEncodingDecoding() throws { + let capabilities = Server.Capabilities.Tasks( + list: .init(), + cancel: .init(), + requests: .init(tools: .init(call: .init())) + ) + + let data = try JSONEncoder().encode(capabilities) + let decoded = try JSONDecoder().decode(Server.Capabilities.Tasks.self, from: data) + + #expect(decoded.list != nil) + #expect(decoded.cancel != nil) + #expect(decoded.requests?.tools?.call != nil) + } + + @Test("Server.Capabilities.Tasks.full() creates complete capability") + func serverTasksCapabilitiesFull() throws { + let capabilities = Server.Capabilities.Tasks.full() + + #expect(capabilities.list != nil) + #expect(capabilities.cancel != nil) + #expect(capabilities.requests?.tools?.call != nil) + } + + @Test("hasTaskAugmentedToolsCall helper") + func hasTaskAugmentedToolsCallHelper() { + // No capabilities + #expect(hasTaskAugmentedToolsCall(nil) == false) + + // Empty capabilities + #expect(hasTaskAugmentedToolsCall(Server.Capabilities()) == false) + + // Tasks without requests + let capsNoRequests = Server.Capabilities(tasks: .init(list: .init())) + #expect(hasTaskAugmentedToolsCall(capsNoRequests) == false) + + // Full task support + let capsFull = Server.Capabilities(tasks: .full()) + #expect(hasTaskAugmentedToolsCall(capsFull) == true) + } + + @Test("requireTaskAugmentedToolsCall throws when not supported") + func requireTaskAugmentedToolsCallThrows() throws { + #expect(throws: MCPError.self) { + try requireTaskAugmentedToolsCall(nil) + } + + #expect(throws: MCPError.self) { + try requireTaskAugmentedToolsCall(Server.Capabilities()) + } + + // Should not throw with full support + try requireTaskAugmentedToolsCall(Server.Capabilities(tasks: .full())) + } +} + +// MARK: - Client Capabilities Tests + +@Suite("Client Tasks Capabilities Tests") +struct ClientTasksCapabilitiesTests { + + @Test("Client.Capabilities.Tasks encoding and decoding") + func clientTasksCapabilitiesEncodingDecoding() throws { + let capabilities = Client.Capabilities.Tasks( + list: .init(), + cancel: .init(), + requests: .init( + sampling: .init(createMessage: .init()), + elicitation: .init(create: .init()) + ) + ) + + let data = try JSONEncoder().encode(capabilities) + let decoded = try JSONDecoder().decode(Client.Capabilities.Tasks.self, from: data) + + #expect(decoded.list != nil) + #expect(decoded.cancel != nil) + #expect(decoded.requests?.sampling?.createMessage != nil) + #expect(decoded.requests?.elicitation?.create != nil) + } + + @Test("Client.Capabilities.Tasks.full() creates complete capability") + func clientTasksCapabilitiesFull() throws { + let capabilities = Client.Capabilities.Tasks.full() + + #expect(capabilities.list != nil) + #expect(capabilities.cancel != nil) + #expect(capabilities.requests?.sampling?.createMessage != nil) + #expect(capabilities.requests?.elicitation?.create != nil) + } + + @Test("hasTaskAugmentedElicitation helper") + func hasTaskAugmentedElicitationHelper() { + #expect(hasTaskAugmentedElicitation(nil) == false) + #expect(hasTaskAugmentedElicitation(Client.Capabilities()) == false) + + let capsWithElicitation = Client.Capabilities( + tasks: .init(requests: .init(elicitation: .init(create: .init()))) + ) + #expect(hasTaskAugmentedElicitation(capsWithElicitation) == true) + } + + @Test("hasTaskAugmentedSampling helper") + func hasTaskAugmentedSamplingHelper() { + #expect(hasTaskAugmentedSampling(nil) == false) + #expect(hasTaskAugmentedSampling(Client.Capabilities()) == false) + + let capsWithSampling = Client.Capabilities( + tasks: .init(requests: .init(sampling: .init(createMessage: .init()))) + ) + #expect(hasTaskAugmentedSampling(capsWithSampling) == true) + } + + @Test("requireTaskAugmentedElicitation throws when not supported") + func requireTaskAugmentedElicitationThrows() throws { + #expect(throws: MCPError.self) { + try requireTaskAugmentedElicitation(nil) + } + + // Should not throw with support + let caps = Client.Capabilities(tasks: .full()) + try requireTaskAugmentedElicitation(caps) + } + + @Test("requireTaskAugmentedSampling throws when not supported") + func requireTaskAugmentedSamplingThrows() throws { + #expect(throws: MCPError.self) { + try requireTaskAugmentedSampling(nil) + } + + // Should not throw with support + let caps = Client.Capabilities(tasks: .full()) + try requireTaskAugmentedSampling(caps) + } +} + +// MARK: - InMemoryTaskStore Tests + +@Suite("InMemoryTaskStore Tests") +struct InMemoryTaskStoreTests { + + @Test("createTask creates task with working status") + func createTaskCreatesWorkingTask() async throws { + let store = InMemoryTaskStore() + let metadata = TaskMetadata(ttl: 60000) + + let task = try await store.createTask(metadata: metadata, taskId: nil) + + #expect(task.status == .working) + #expect(task.ttl == 60000) + #expect(!task.taskId.isEmpty) + } + + @Test("createTask with custom taskId") + func createTaskWithCustomId() async throws { + let store = InMemoryTaskStore() + let metadata = TaskMetadata(ttl: nil) + + let task = try await store.createTask(metadata: metadata, taskId: "custom-id-123") + + #expect(task.taskId == "custom-id-123") + } + + @Test("createTask throws on duplicate taskId") + func createTaskThrowsOnDuplicate() async throws { + let store = InMemoryTaskStore() + let metadata = TaskMetadata() + + _ = try await store.createTask(metadata: metadata, taskId: "task-1") + + await #expect(throws: MCPError.self) { + _ = try await store.createTask(metadata: metadata, taskId: "task-1") + } + } + + @Test("getTask returns created task") + func getTaskReturnsCreatedTask() async throws { + let store = InMemoryTaskStore() + let created = try await store.createTask(metadata: TaskMetadata(), taskId: "task-123") + + let retrieved = await store.getTask(taskId: "task-123") + + #expect(retrieved?.taskId == created.taskId) + #expect(retrieved?.status == created.status) + } + + @Test("getTask returns nil for non-existent task") + func getTaskReturnsNilForNonExistent() async { + let store = InMemoryTaskStore() + + let result = await store.getTask(taskId: "non-existent") + + #expect(result == nil) + } + + @Test("updateTask changes status") + func updateTaskChangesStatus() async throws { + let store = InMemoryTaskStore() + _ = try await store.createTask(metadata: TaskMetadata(), taskId: "task-123") + + let updated = try await store.updateTask(taskId: "task-123", status: .completed, statusMessage: "Done") + + #expect(updated.status == .completed) + #expect(updated.statusMessage == "Done") + } + + @Test("updateTask throws when transitioning from terminal status") + func updateTaskThrowsFromTerminalStatus() async throws { + let store = InMemoryTaskStore() + _ = try await store.createTask(metadata: TaskMetadata(), taskId: "task-123") + + // Complete the task + _ = try await store.updateTask(taskId: "task-123", status: .completed, statusMessage: nil) + + // Try to update again - should throw + await #expect(throws: MCPError.self) { + _ = try await store.updateTask(taskId: "task-123", status: .working, statusMessage: nil) + } + } + + @Test("updateTask throws for non-existent task") + func updateTaskThrowsForNonExistent() async { + let store = InMemoryTaskStore() + + await #expect(throws: MCPError.self) { + _ = try await store.updateTask(taskId: "non-existent", status: .completed, statusMessage: nil) + } + } + + @Test("storeResult and getResult work correctly") + func storeAndGetResult() async throws { + let store = InMemoryTaskStore() + _ = try await store.createTask(metadata: TaskMetadata(), taskId: "task-123") + + let result: Value = .object(["data": .string("test result")]) + try await store.storeResult(taskId: "task-123", result: result) + + let retrieved = await store.getResult(taskId: "task-123") + + #expect(retrieved?.objectValue?["data"]?.stringValue == "test result") + } + + @Test("getResult returns nil when no result stored") + func getResultReturnsNilWhenNoResult() async throws { + let store = InMemoryTaskStore() + _ = try await store.createTask(metadata: TaskMetadata(), taskId: "task-123") + + let result = await store.getResult(taskId: "task-123") + + #expect(result == nil) + } + + @Test("listTasks returns all tasks") + func listTasksReturnsAllTasks() async throws { + let store = InMemoryTaskStore() + _ = try await store.createTask(metadata: TaskMetadata(), taskId: "task-1") + _ = try await store.createTask(metadata: TaskMetadata(), taskId: "task-2") + _ = try await store.createTask(metadata: TaskMetadata(), taskId: "task-3") + + let (tasks, _) = await store.listTasks(cursor: nil) + + #expect(tasks.count == 3) + } + + @Test("listTasks pagination works correctly") + func listTasksPagination() async throws { + let store = InMemoryTaskStore(pageSize: 2) + + // Create 5 tasks + for i in 1...5 { + _ = try await store.createTask(metadata: TaskMetadata(), taskId: "task-\(i)") + } + + // First page + let (page1, cursor1) = await store.listTasks(cursor: nil) + #expect(page1.count == 2) + #expect(cursor1 != nil) + + // Second page + let (page2, cursor2) = await store.listTasks(cursor: cursor1) + #expect(page2.count == 2) + #expect(cursor2 != nil) + + // Third page + let (page3, cursor3) = await store.listTasks(cursor: cursor2) + #expect(page3.count == 1) + #expect(cursor3 == nil) + } + + @Test("deleteTask removes task") + func deleteTaskRemovesTask() async throws { + let store = InMemoryTaskStore() + _ = try await store.createTask(metadata: TaskMetadata(), taskId: "task-123") + + let deleted = await store.deleteTask(taskId: "task-123") + #expect(deleted == true) + + let result = await store.getTask(taskId: "task-123") + #expect(result == nil) + } + + @Test("deleteTask returns false for non-existent task") + func deleteTaskReturnsFalseForNonExistent() async { + let store = InMemoryTaskStore() + + let deleted = await store.deleteTask(taskId: "non-existent") + + #expect(deleted == false) + } + + @Test("waitForUpdate and notifyUpdate work together") + func waitForUpdateAndNotifyUpdate() async throws { + let store = InMemoryTaskStore() + _ = try await store.createTask(metadata: TaskMetadata(), taskId: "task-123") + + // Start waiting in a separate task + let waitTask = Task { + try await store.waitForUpdate(taskId: "task-123") + return true + } + + // Give the wait a moment to start + try await Task.sleep(for: .milliseconds(50)) + + // Notify update + await store.notifyUpdate(taskId: "task-123") + + // Wait should complete + let result = try await waitTask.value + #expect(result == true) + } +} + +// MARK: - InMemoryTaskMessageQueue Tests + +@Suite("InMemoryTaskMessageQueue Tests") +struct InMemoryTaskMessageQueueTests { + + @Test("enqueue and dequeue work correctly") + func enqueueAndDequeue() async throws { + let queue = InMemoryTaskMessageQueue() + + let message = QueuedMessage.notification( + try JSONEncoder().encode(["test": "data"]), + timestamp: Date() + ) + + try await queue.enqueue(taskId: "task-123", message: message, maxSize: nil) + + let dequeued = await queue.dequeue(taskId: "task-123") + #expect(dequeued != nil) + + // Queue should now be empty + let empty = await queue.dequeue(taskId: "task-123") + #expect(empty == nil) + } + + @Test("enqueue respects maxSize") + func enqueueRespectsMaxSize() async throws { + let queue = InMemoryTaskMessageQueue() + + let message = QueuedMessage.notification(Data(), timestamp: Date()) + + try await queue.enqueue(taskId: "task-123", message: message, maxSize: 1) + + // Second enqueue should fail + await #expect(throws: MCPError.self) { + try await queue.enqueue(taskId: "task-123", message: message, maxSize: 1) + } + } + + @Test("dequeueAll returns all messages") + func dequeueAllReturnsAllMessages() async throws { + let queue = InMemoryTaskMessageQueue() + + for i in 0..<3 { + let message = QueuedMessage.notification( + try JSONEncoder().encode(["index": i]), + timestamp: Date() + ) + try await queue.enqueue(taskId: "task-123", message: message, maxSize: nil) + } + + let all = await queue.dequeueAll(taskId: "task-123") + #expect(all.count == 3) + + // Queue should now be empty + let empty = await queue.isEmpty(taskId: "task-123") + #expect(empty == true) + } + + @Test("isEmpty returns correct value") + func isEmptyReturnsCorrectValue() async throws { + let queue = InMemoryTaskMessageQueue() + + #expect(await queue.isEmpty(taskId: "task-123") == true) + + let message = QueuedMessage.notification(Data(), timestamp: Date()) + try await queue.enqueue(taskId: "task-123", message: message, maxSize: nil) + + #expect(await queue.isEmpty(taskId: "task-123") == false) + + _ = await queue.dequeue(taskId: "task-123") + + #expect(await queue.isEmpty(taskId: "task-123") == true) + } + + @Test("enqueueWithResolver stores resolver") + func enqueueWithResolverStoresResolver() async throws { + let queue = InMemoryTaskMessageQueue() + let resolver = Resolver() + + let message = QueuedMessage.request(Data(), timestamp: Date()) + let queuedRequest = QueuedRequestWithResolver( + message: message, + resolver: resolver, + originalRequestId: .string("req-1") + ) + + try await queue.enqueueWithResolver(taskId: "task-123", request: queuedRequest, maxSize: nil) + + // Resolver should be retrievable + let retrieved = await queue.getResolver(forRequestId: .string("req-1")) + #expect(retrieved != nil) + } + + @Test("removeResolver removes and returns resolver") + func removeResolverRemovesAndReturns() async throws { + let queue = InMemoryTaskMessageQueue() + let resolver = Resolver() + + let message = QueuedMessage.request(Data(), timestamp: Date()) + let queuedRequest = QueuedRequestWithResolver( + message: message, + resolver: resolver, + originalRequestId: .string("req-1") + ) + + try await queue.enqueueWithResolver(taskId: "task-123", request: queuedRequest, maxSize: nil) + + let removed = await queue.removeResolver(forRequestId: .string("req-1")) + #expect(removed != nil) + + // Should no longer be retrievable + let notFound = await queue.getResolver(forRequestId: .string("req-1")) + #expect(notFound == nil) + } +} + +// MARK: - Resolver Tests + +@Suite("Resolver Tests") +struct ResolverTests { + + @Test("setResult and wait work correctly") + func setResultAndWait() async throws { + let resolver = Resolver() + + // Set result in background + Task { + await resolver.setResult(.string("success")) + } + + let result = try await resolver.wait() + #expect(result.stringValue == "success") + } + + @Test("setError and wait throws correctly") + func setErrorAndWaitThrows() async throws { + let resolver = Resolver() + + // Set error in background + Task { + await resolver.setError(MCPError.internalError("test error")) + } + + await #expect(throws: MCPError.self) { + _ = try await resolver.wait() + } + } + + @Test("isDone returns correct value") + func isDoneReturnsCorrectValue() async { + let resolver = Resolver() + + #expect(await resolver.isDone == false) + + await resolver.setResult(.string("done")) + + #expect(await resolver.isDone == true) + } + + @Test("setResult is idempotent") + func setResultIsIdempotent() async throws { + let resolver = Resolver() + + await resolver.setResult(.string("first")) + await resolver.setResult(.string("second")) // Should be ignored + + let result = try await resolver.wait() + #expect(result.stringValue == "first") + } +} + +// MARK: - QueuedMessage Tests + +@Suite("QueuedMessage Tests") +struct QueuedMessageTests { + + @Test("QueuedMessage.request stores data and timestamp") + func queuedMessageRequest() { + let data = Data("test".utf8) + let timestamp = Date() + let message = QueuedMessage.request(data, timestamp: timestamp) + + #expect(message.data == data) + #expect(message.timestamp == timestamp) + } + + @Test("QueuedMessage.notification stores data and timestamp") + func queuedMessageNotification() { + let data = Data("notification".utf8) + let timestamp = Date() + let message = QueuedMessage.notification(data, timestamp: timestamp) + + #expect(message.data == data) + #expect(message.timestamp == timestamp) + } + + @Test("QueuedMessage.response stores data and timestamp") + func queuedMessageResponse() { + let data = Data("response".utf8) + let timestamp = Date() + let message = QueuedMessage.response(data, timestamp: timestamp) + + #expect(message.data == data) + #expect(message.timestamp == timestamp) + } + + @Test("QueuedMessage.error stores data and timestamp") + func queuedMessageError() { + let data = Data("error".utf8) + let timestamp = Date() + let message = QueuedMessage.error(data, timestamp: timestamp) + + #expect(message.data == data) + #expect(message.timestamp == timestamp) + } +} + +// MARK: - JSON Round-Trip Tests + +@Suite("Task JSON Round-Trip Tests") +struct TaskJSONRoundTripTests { + + @Test("Complete task workflow JSON encoding") + func completeTaskWorkflowJSON() throws { + // 1. Create task with metadata + let createParams = CallTool.Parameters( + name: "long_running_tool", + arguments: ["input": .string("data")], + task: TaskMetadata(ttl: 60000) + ) + + let createData = try JSONEncoder().encode(createParams) + let decodedCreate = try JSONDecoder().decode(CallTool.Parameters.self, from: createData) + #expect(decodedCreate.task?.ttl == 60000) + + // 2. Create task result + let task = MCPTask( + taskId: "task-abc123", + status: .working, + ttl: 60000, + createdAt: "2024-01-15T10:30:00Z", + lastUpdatedAt: "2024-01-15T10:30:00Z", + pollInterval: 1000 + ) + let createResult = CreateTaskResult(task: task, modelImmediateResponse: "Starting...") + + let resultData = try JSONEncoder().encode(createResult) + let decodedResult = try JSONDecoder().decode(CreateTaskResult.self, from: resultData) + #expect(decodedResult.task.taskId == "task-abc123") + #expect(decodedResult._meta?[modelImmediateResponseKey]?.stringValue == "Starting...") + + // 3. Task status notification + let notification = TaskStatusNotification.Parameters( + taskId: "task-abc123", + status: .inputRequired, + ttl: 60000, + createdAt: "2024-01-15T10:30:00Z", + lastUpdatedAt: "2024-01-15T10:30:05Z", + statusMessage: "Waiting for input" + ) + + let notificationData = try JSONEncoder().encode(notification) + let decodedNotification = try JSONDecoder().decode( + TaskStatusNotification.Parameters.self, from: notificationData) + #expect(decodedNotification.status == .inputRequired) + + // 4. Get task result + let payloadResult = GetTaskPayload.Result( + _meta: [relatedTaskMetaKey: .object(["taskId": .string("task-abc123")])], + extraFields: [ + "content": .array([.object(["type": .string("text"), "text": .string("Result")])]), + "isError": .bool(false), + ] + ) + + let payloadData = try JSONEncoder().encode(payloadResult) + let decodedPayload = try JSONDecoder().decode(GetTaskPayload.Result.self, from: payloadData) + #expect(decodedPayload.extraFields?["isError"]?.boolValue == false) + } +} diff --git a/Tests/MCPTests/ToolTests.swift b/Tests/MCPTests/ToolTests.swift index b08963b3..ba9945df 100644 --- a/Tests/MCPTests/ToolTests.swift +++ b/Tests/MCPTests/ToolTests.swift @@ -11,6 +11,7 @@ struct ToolTests { name: "test_tool", description: "A test tool", inputSchema: .object([ + "type": .string("object"), "properties": .object([ "param1": .string("Test parameter") ]) @@ -102,6 +103,7 @@ struct ToolTests { name: "calculate", description: "Performs calculations", inputSchema: .object([ + "type": .string("object"), "properties": .object([ "expression": .string("Mathematical expression to evaluate") ]) @@ -131,7 +133,7 @@ struct ToolTests { var tool = Tool( name: "test_tool", description: "Test tool description", - inputSchema: [:] + inputSchema: ["type": "object"] ) do { @@ -164,7 +166,7 @@ struct ToolTests { let tool = Tool( name: "test_tool", description: "Test tool description", - inputSchema: [:], + inputSchema: ["type": "object"], annotations: nil ) @@ -184,9 +186,10 @@ struct ToolTests { name: "test_tool", description: "Test tool description", inputSchema: .object([ + "type": .string("object"), "properties": .object([ "param1": .string("String parameter"), - "param2": .int(42), + "param2": .int(42) ]) ]) ) @@ -211,7 +214,7 @@ struct ToolTests { let data = try encoder.encode(content) let decoded = try decoder.decode(Tool.Content.self, from: data) - if case .text(let text) = decoded { + if case .text(let text, _, _) = decoded { #expect(text == "Hello, world!") } else { #expect(Bool(false), "Expected text content") @@ -220,22 +223,16 @@ struct ToolTests { @Test("Image content encoding and decoding") func testToolContentImageEncoding() throws { - let content = Tool.Content.image( - data: "base64data", - mimeType: "image/png", - metadata: ["width": "100", "height": "100"] - ) + let content = Tool.Content.image(data: "base64data", mimeType: "image/png") let encoder = JSONEncoder() let decoder = JSONDecoder() let data = try encoder.encode(content) let decoded = try decoder.decode(Tool.Content.self, from: data) - if case .image(let data, let mimeType, let metadata) = decoded { + if case .image(let data, let mimeType, _, _) = decoded { #expect(data == "base64data") #expect(mimeType == "image/png") - #expect(metadata?["width"] == "100") - #expect(metadata?["height"] == "100") } else { #expect(Bool(false), "Expected image content") } @@ -254,10 +251,10 @@ struct ToolTests { let data = try encoder.encode(content) let decoded = try decoder.decode(Tool.Content.self, from: data) - if case .resource(let uri, let mimeType, let text) = decoded { - #expect(uri == "file://test.txt") - #expect(mimeType == "text/plain") - #expect(text == "Sample text") + if case .resource(let resourceContent, _, _) = decoded { + #expect(resourceContent.uri == "file://test.txt") + #expect(resourceContent.mimeType == "text/plain") + #expect(resourceContent.text == "Sample text") } else { #expect(Bool(false), "Expected resource content") } @@ -275,7 +272,7 @@ struct ToolTests { let data = try encoder.encode(content) let decoded = try decoder.decode(Tool.Content.self, from: data) - if case .audio(let data, let mimeType) = decoded { + if case .audio(let data, let mimeType, _, _) = decoded { #expect(data == "base64audiodata") #expect(mimeType == "audio/wav") } else { @@ -325,8 +322,8 @@ struct ToolTests { @Test("ListTools result validation") func testListToolsResult() throws { let tools = [ - Tool(name: "tool1", description: "First tool", inputSchema: [:]), - Tool(name: "tool2", description: "Second tool", inputSchema: [:]), + Tool(name: "tool1", description: "First tool", inputSchema: ["type": "object"]), + Tool(name: "tool2", description: "Second tool", inputSchema: ["type": "object"]), ] let result = ListTools.Result(tools: tools, nextCursor: "next_page") @@ -360,7 +357,7 @@ struct ToolTests { #expect(result.content.count == 2) #expect(result.isError == nil) - if case .text(let text) = result.content[0] { + if case .text(let text, _, _) = result.content[0] { #expect(text == "Result 1") } else { #expect(Bool(false), "Expected text content") @@ -374,7 +371,7 @@ struct ToolTests { #expect(errorResult.content.count == 1) #expect(errorResult.isError == true) - if case .text(let text) = errorResult.content[0] { + if case .text(let text, _, _) = errorResult.content[0] { #expect(text == "Error message") } else { #expect(Bool(false), "Expected error text content") @@ -395,7 +392,7 @@ struct ToolTests { let anyRequest = try JSONDecoder().decode(AnyRequest.self, from: jsonData) - let handler = TypedRequestHandler { request in + let handler = TypedRequestHandler { request, _ in #expect(request.method == ListTools.name) #expect(request.id == 1) #expect(request.params.cursor == nil) @@ -403,12 +400,20 @@ struct ToolTests { let testTool = Tool( name: "test_tool", description: "Test tool for verification", - inputSchema: [:] + inputSchema: ["type": "object"] ) return ListTools.response(id: request.id, result: ListTools.Result(tools: [testTool])) } - let response = try await handler(anyRequest) + // Create a dummy context for testing + let dummyContext = Server.RequestHandlerContext( + sendNotification: { _ in }, + sendMessage: { _ in }, + sendData: { _ in }, + sessionId: nil, + shouldSendLogMessage: { _ in true } + ) + let response = try await handler(anyRequest, context: dummyContext) if case .success(let value) = response.result { let encoder = JSONEncoder() @@ -422,21 +427,801 @@ struct ToolTests { #expect(Bool(false), "Expected success result") } } -} @Test("Tool with missing description") func testToolWithMissingDescription() throws { let jsonString = """ { "name": "test_tool", - "inputSchema": {} + "inputSchema": {"type": "object"} } """ let jsonData = jsonString.data(using: .utf8)! - + let tool = try JSONDecoder().decode(Tool.self, from: jsonData) - + #expect(tool.name == "test_tool") #expect(tool.description == nil) - #expect(tool.inputSchema == [:]) - } \ No newline at end of file + #expect(tool.inputSchema == ["type": "object"]) + } + + // MARK: - Tool with outputSchema + + @Test("Tool with outputSchema encoding and decoding") + func testToolWithOutputSchema() throws { + let outputSchema: Value = [ + "type": "object", + "properties": [ + "result": ["type": "integer"] + ], + "required": ["result"] + ] + + let tool = Tool( + name: "calculate", + description: "Performs calculations", + inputSchema: ["type": "object"], + outputSchema: outputSchema + ) + + #expect(tool.outputSchema != nil) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(tool) + let decoded = try decoder.decode(Tool.self, from: data) + + #expect(decoded.outputSchema == outputSchema) + + // Verify JSON contains outputSchema + let jsonString = String(data: data, encoding: .utf8)! + #expect(jsonString.contains("\"outputSchema\"")) + } + + @Test("CallTool result with structuredContent") + func testCallToolResultWithStructuredContent() throws { + let structuredContent: Value = [ + "name": "John", + "age": 30 + ] + + let result = CallTool.Result( + content: [.text("User data")], + structuredContent: structuredContent + ) + + #expect(result.structuredContent == structuredContent) + #expect(result.isError == nil) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(result) + let decoded = try decoder.decode(CallTool.Result.self, from: data) + + #expect(decoded.structuredContent == structuredContent) + } + + // MARK: - Tool.Execution Tests + + @Test("Tool.Execution with taskSupport encoding and decoding") + func testToolExecutionWithTaskSupport() throws { + let execution = Tool.Execution(taskSupport: .required) + #expect(execution.taskSupport == .required) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(execution) + let decoded = try decoder.decode(Tool.Execution.self, from: data) + + #expect(decoded.taskSupport == .required) + + let jsonString = String(data: data, encoding: .utf8)! + #expect(jsonString.contains("\"taskSupport\":\"required\"")) + } + + @Test("Tool with execution property encoding and decoding") + func testToolWithExecution() throws { + let tool = Tool( + name: "long_running_task", + description: "A task that takes a long time", + inputSchema: ["type": "object"], + execution: Tool.Execution(taskSupport: .optional) + ) + + #expect(tool.execution?.taskSupport == .optional) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(tool) + let decoded = try decoder.decode(Tool.self, from: data) + + #expect(decoded.execution?.taskSupport == .optional) + + let jsonString = String(data: data, encoding: .utf8)! + #expect(jsonString.contains("\"execution\"")) + #expect(jsonString.contains("\"taskSupport\":\"optional\"")) + } + + @Test( + "Tool.Execution.TaskSupport enum values", + arguments: [ + (Tool.Execution.TaskSupport.forbidden, "forbidden"), + (Tool.Execution.TaskSupport.optional, "optional"), + (Tool.Execution.TaskSupport.required, "required") + ] + ) + func testTaskSupportEnumValues(testCase: (value: Tool.Execution.TaskSupport, rawValue: String)) throws { + #expect(testCase.value.rawValue == testCase.rawValue) + + let execution = Tool.Execution(taskSupport: testCase.value) + let encoder = JSONEncoder() + let data = try encoder.encode(execution) + let jsonString = String(data: data, encoding: .utf8)! + + #expect(jsonString.contains("\"\(testCase.rawValue)\"")) + } + + @Test("Tool.Execution with nil taskSupport") + func testToolExecutionWithNilTaskSupport() throws { + let execution = Tool.Execution(taskSupport: nil) + #expect(execution.taskSupport == nil) + + let encoder = JSONEncoder() + let data = try encoder.encode(execution) + + // Empty execution should encode as empty object + let jsonString = String(data: data, encoding: .utf8)! + #expect(jsonString == "{}") + } + + // MARK: - Tool with Title, Icons, _meta Tests + + @Test("Tool with top-level title property") + func testToolWithTitle() throws { + let tool = Tool( + name: "calculate", + title: "Calculator Tool", + description: "Performs calculations", + inputSchema: ["type": "object"] + ) + + #expect(tool.title == "Calculator Tool") + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(tool) + let decoded = try decoder.decode(Tool.self, from: data) + + #expect(decoded.title == "Calculator Tool") + + let jsonString = String(data: data, encoding: .utf8)! + #expect(jsonString.contains("\"title\":\"Calculator Tool\"")) + } + + @Test("Tool with icons") + func testToolWithIcons() throws { + let icons = [ + Icon(src: "https://example.com/icon.png", mimeType: "image/png", sizes: ["48x48"], theme: .light), + Icon(src: "https://example.com/icon-dark.png", mimeType: "image/png", sizes: ["48x48"], theme: .dark) + ] + + let tool = Tool( + name: "visual_tool", + description: "A tool with icons", + inputSchema: ["type": "object"], + icons: icons + ) + + #expect(tool.icons?.count == 2) + #expect(tool.icons?[0].theme == .light) + #expect(tool.icons?[1].theme == .dark) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(tool) + let decoded = try decoder.decode(Tool.self, from: data) + + #expect(decoded.icons?.count == 2) + #expect(decoded.icons?[0].src == "https://example.com/icon.png") + #expect(decoded.icons?[1].src == "https://example.com/icon-dark.png") + } + + @Test("Tool with _meta") + func testToolWithMeta() throws { + let meta: [String: Value] = [ + "vendor": .string("example"), + "version": .int(1), + "experimental": .bool(true) + ] + + let tool = Tool( + name: "meta_tool", + description: "A tool with metadata", + inputSchema: ["type": "object"], + _meta: meta + ) + + #expect(tool._meta?["vendor"]?.stringValue == "example") + #expect(tool._meta?["version"]?.intValue == 1) + #expect(tool._meta?["experimental"]?.boolValue == true) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(tool) + let decoded = try decoder.decode(Tool.self, from: data) + + #expect(decoded._meta?["vendor"]?.stringValue == "example") + #expect(decoded._meta?["version"]?.intValue == 1) + } + + @Test("Tool with all properties") + func testToolWithAllProperties() throws { + let tool = Tool( + name: "full_tool", + title: "Full Featured Tool", + description: "A tool with all properties", + inputSchema: [ + "type": "object", + "properties": [ + "input": ["type": "string"] + ] + ], + outputSchema: [ + "type": "object", + "properties": [ + "result": ["type": "integer"] + ] + ], + _meta: ["custom": .string("value")], + icons: [Icon(src: "https://example.com/icon.svg", mimeType: "image/svg+xml")], + execution: Tool.Execution(taskSupport: .optional), + annotations: Tool.Annotations( + title: "Annotated Title", + readOnlyHint: true, + destructiveHint: false, + idempotentHint: true, + openWorldHint: false + ) + ) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(tool) + let decoded = try decoder.decode(Tool.self, from: data) + + #expect(decoded.name == "full_tool") + #expect(decoded.title == "Full Featured Tool") + #expect(decoded.description == "A tool with all properties") + #expect(decoded.outputSchema != nil) + #expect(decoded._meta?["custom"]?.stringValue == "value") + #expect(decoded.icons?.count == 1) + #expect(decoded.execution?.taskSupport == .optional) + #expect(decoded.annotations.title == "Annotated Title") + #expect(decoded.annotations.readOnlyHint == true) + } + + // MARK: - ResourceLink Content Tests + + @Test("ResourceLink content encoding and decoding") + func testResourceLinkContent() throws { + let resourceLink = ResourceLink( + name: "data.json", + title: "Data File", + uri: "file:///data/output.json", + description: "Output data file", + mimeType: "application/json", + size: 1024 + ) + + let content = Tool.Content.resourceLink(resourceLink) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(content) + let decoded = try decoder.decode(Tool.Content.self, from: data) + + if case .resourceLink(let link) = decoded { + #expect(link.name == "data.json") + #expect(link.title == "Data File") + #expect(link.uri == "file:///data/output.json") + #expect(link.description == "Output data file") + #expect(link.mimeType == "application/json") + #expect(link.size == 1024) + } else { + #expect(Bool(false), "Expected resourceLink content") + } + + let jsonString = String(data: data, encoding: .utf8)! + #expect(jsonString.contains("\"type\":\"resource_link\"")) + } + + @Test("ResourceLink with icons and annotations") + func testResourceLinkWithIconsAndAnnotations() throws { + let resourceLink = ResourceLink( + name: "report.pdf", + uri: "file:///reports/report.pdf", + mimeType: "application/pdf", + annotations: Annotations(audience: [.assistant], priority: 0.8), + icons: [Icon(src: "https://example.com/pdf.png", mimeType: "image/png")] + ) + + let content = Tool.Content.resourceLink(resourceLink) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(content) + let decoded = try decoder.decode(Tool.Content.self, from: data) + + if case .resourceLink(let link) = decoded { + #expect(link.annotations?.audience == [.assistant]) + #expect(link.annotations?.priority == 0.8) + #expect(link.icons?.count == 1) + } else { + #expect(Bool(false), "Expected resourceLink content") + } + } + + // MARK: - Content with Annotations and _meta Tests + + @Test("Text content with annotations and _meta") + func testTextContentWithAnnotationsAndMeta() throws { + let annotations = Annotations(audience: [.user, .assistant], priority: 0.9) + let meta: [String: Value] = ["source": .string("calculation")] + + let content = Tool.Content.text("Result: 42", annotations: annotations, _meta: meta) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(content) + let decoded = try decoder.decode(Tool.Content.self, from: data) + + if case .text(let text, let decodedAnnotations, let decodedMeta) = decoded { + #expect(text == "Result: 42") + #expect(decodedAnnotations?.audience == [.user, .assistant]) + #expect(decodedAnnotations?.priority == 0.9) + #expect(decodedMeta?["source"]?.stringValue == "calculation") + } else { + #expect(Bool(false), "Expected text content") + } + } + + @Test("Image content with annotations") + func testImageContentWithAnnotations() throws { + let annotations = Annotations(audience: [.user]) + + let content = Tool.Content.image( + data: "base64imagedata", + mimeType: "image/png", + annotations: annotations, + _meta: nil + ) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(content) + let decoded = try decoder.decode(Tool.Content.self, from: data) + + if case .image(_, _, let decodedAnnotations, _) = decoded { + #expect(decodedAnnotations?.audience == [.user]) + } else { + #expect(Bool(false), "Expected image content") + } + } + + @Test("Resource content with annotations") + func testResourceContentWithAnnotations() throws { + let annotations = Annotations(priority: 0.5) + let resourceContent = Resource.Content.text("File contents", uri: "file:///test.txt", mimeType: "text/plain") + + let content = Tool.Content.resource(resource: resourceContent, annotations: annotations, _meta: nil) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(content) + let decoded = try decoder.decode(Tool.Content.self, from: data) + + if case .resource(_, let decodedAnnotations, _) = decoded { + #expect(decodedAnnotations?.priority == 0.5) + } else { + #expect(Bool(false), "Expected resource content") + } + } +} + +// MARK: - Tool Name Validation Tests + +@Suite("Tool Name Validation Tests") +struct ToolNameValidationTests { + + // MARK: - Valid Names + + @Test( + "Accepts valid tool names", + arguments: [ + "getUser", + "get_user_profile", + "user-profile-update", + "admin.tools.list", + "DATA_EXPORT_v2.1", + "a", + String(repeating: "a", count: 128) + ] + ) + func acceptsValidNames(toolName: String) throws { + let result = validateToolName(toolName) + #expect(result.isValid == true) + #expect(result.warnings.isEmpty) + } + + // MARK: - Invalid Names + + @Test("Rejects empty name") + func rejectsEmptyName() throws { + let result = validateToolName("") + #expect(result.isValid == false) + #expect(result.warnings.contains { $0.contains("cannot be empty") }) + } + + @Test("Rejects name exceeding max length") + func rejectsNameExceedingMaxLength() throws { + let longName = String(repeating: "a", count: 129) + let result = validateToolName(longName) + #expect(result.isValid == false) + #expect(result.warnings.contains { $0.contains("exceeds maximum length of 128 characters") }) + #expect(result.warnings.contains { $0.contains("current: 129") }) + } + + @Test( + "Rejects names with invalid characters", + arguments: [ + ("get user profile", " "), + ("get,user,profile", ","), + ("user/profile/update", "/"), + ("user@domain.com", "@") + ] + ) + func rejectsInvalidCharacters(testCase: (toolName: String, invalidChar: String)) throws { + let result = validateToolName(testCase.toolName) + #expect(result.isValid == false) + #expect(result.warnings.contains { $0.contains("invalid characters") }) + } + + @Test("Rejects multiple invalid characters") + func rejectsMultipleInvalidChars() throws { + let result = validateToolName("user name@domain,com") + #expect(result.isValid == false) + let warningWithChars = result.warnings.first { $0.contains("invalid characters") } + #expect(warningWithChars != nil) + } + + @Test("Rejects unicode characters") + func rejectsUnicodeCharacters() throws { + let result = validateToolName("user-ñame") // n with tilde + #expect(result.isValid == false) + } + + // MARK: - Warnings for Problematic Patterns + + @Test("Warns on leading dash") + func warnsOnLeadingDash() throws { + let result = validateToolName("-get-user") + #expect(result.isValid == true) + #expect(result.warnings.contains { $0.contains("starts or ends with a dash") }) + } + + @Test("Warns on trailing dash") + func warnsOnTrailingDash() throws { + let result = validateToolName("get-user-") + #expect(result.isValid == true) + #expect(result.warnings.contains { $0.contains("starts or ends with a dash") }) + } + + @Test("Warns on leading dot") + func warnsOnLeadingDot() throws { + let result = validateToolName(".get.user") + #expect(result.isValid == true) + #expect(result.warnings.contains { $0.contains("starts or ends with a dot") }) + } + + @Test("Warns on trailing dot") + func warnsOnTrailingDot() throws { + let result = validateToolName("get.user.") + #expect(result.isValid == true) + #expect(result.warnings.contains { $0.contains("starts or ends with a dot") }) + } + + // MARK: - Edge Cases + + @Test("Handles only dots") + func handlesOnlyDots() throws { + let result = validateToolName("...") + #expect(result.isValid == true) + #expect(result.warnings.contains { $0.contains("starts or ends with a dot") }) + } + + @Test("Handles only dashes") + func handlesOnlyDashes() throws { + let result = validateToolName("---") + #expect(result.isValid == true) + #expect(result.warnings.contains { $0.contains("starts or ends with a dash") }) + } + + @Test("Rejects only slashes") + func rejectsOnlySlashes() throws { + let result = validateToolName("///") + #expect(result.isValid == false) + #expect(result.warnings.contains { $0.contains("invalid characters") }) + } + + @Test("Rejects mixed valid and invalid characters") + func rejectsMixedValidInvalid() throws { + let result = validateToolName("user@name123") + #expect(result.isValid == false) + #expect(result.warnings.contains { $0.contains("invalid characters") }) + } + + // MARK: - validateAndWarnToolName + + @Test("validateAndWarnToolName returns true for valid name") + func validateAndWarnReturnsTrue() throws { + let isValid = validateAndWarnToolName("valid-tool-name") + #expect(isValid == true) + } + + @Test("validateAndWarnToolName returns false for invalid name") + func validateAndWarnReturnsFalse() throws { + #expect(validateAndWarnToolName("") == false) + #expect(validateAndWarnToolName(String(repeating: "a", count: 129)) == false) + #expect(validateAndWarnToolName("invalid name") == false) + } +} + +// MARK: - Unicode Tool Tests + +@Suite("Unicode Tool Tests") +struct UnicodeToolTests { + + /// Test strings with various Unicode characters (matching Python SDK) + static let unicodeTestStrings: [String: String] = [ + "cyrillic": "Слой хранилища, где располагаются", + "cyrillic_short": "Привет мир", + "chinese": "你好世界 - 这是一个测试", + "japanese": "こんにちは世界 - これはテストです", + "korean": "안녕하세요 세계 - 이것은 테스트입니다", + "arabic": "مرحبا بالعالم - هذا اختبار", + "hebrew": "שלום עולם - זה מבחן", + "greek": "Γεια σου κόσμε - αυτό είναι δοκιμή", + "emoji": "Hello 👋 World 🌍 - Testing 🧪 Unicode ✨", + "math": "∑ ∫ √ ∞ ≠ ≤ ≥ ∈ ∉ ⊆ ⊇", + "accented": "Café, naïve, résumé, piñata, Zürich", + "mixed": "Hello世界🌍Привет안녕مرحباשלום", + "special": "Line\nbreak\ttab\r\nCRLF", + "quotes": #"«French» „German" "English" 「Japanese」"#, + "currency": "€100 £50 ¥1000 ₹500 ₽200 ¢99" + ] + + @Test("Tool with Unicode description encodes and decodes correctly") + func unicodeDescriptionEncodingDecoding() throws { + let tool = Tool( + name: "echo_unicode", + description: "🔤 Echo Unicode text - Hello 👋 World 🌍 - Testing 🧪 Unicode ✨", + inputSchema: ["type": "object"] + ) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(tool) + let decoded = try decoder.decode(Tool.self, from: data) + + #expect(decoded.description == tool.description) + #expect(decoded.description?.contains("🔤") == true) + #expect(decoded.description?.contains("👋") == true) + } + + @Test( + "Unicode text in tool call arguments roundtrips correctly", + arguments: Array(unicodeTestStrings.keys) + ) + func unicodeArgumentsRoundtrip(testKey: String) throws { + let testString = Self.unicodeTestStrings[testKey]! + + let params = CallTool.Parameters( + name: "echo_unicode", + arguments: ["text": .string(testString)] + ) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(params) + let decoded = try decoder.decode(CallTool.Parameters.self, from: data) + + #expect(decoded.arguments?["text"]?.stringValue == testString) + } + + @Test("Unicode text in tool result content roundtrips correctly") + func unicodeResultContentRoundtrip() throws { + for (testName, testString) in Self.unicodeTestStrings { + let result = CallTool.Result( + content: [.text("Echo: \(testString)")] + ) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(result) + let decoded = try decoder.decode(CallTool.Result.self, from: data) + + if case .text(let text, _, _) = decoded.content[0] { + #expect(text == "Echo: \(testString)", "Failed for \(testName)") + } else { + Issue.record("Expected text content for \(testName)") + } + } + } + + @Test("Mixed Unicode content types roundtrip correctly") + func mixedUnicodeContentRoundtrip() throws { + let cyrillic = Self.unicodeTestStrings["cyrillic"]! + let mixed = Self.unicodeTestStrings["mixed"]! + + let result = CallTool.Result( + content: [ + .text(cyrillic), + .text(mixed) + ], + structuredContent: [ + "message": .string(mixed), + "data": .object([ + "text": .string(cyrillic) + ]) + ] + ) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(result) + let decoded = try decoder.decode(CallTool.Result.self, from: data) + + #expect(decoded.content.count == 2) + + if case .text(let text1, _, _) = decoded.content[0] { + #expect(text1 == cyrillic) + } + + if case .text(let text2, _, _) = decoded.content[1] { + #expect(text2 == mixed) + } + + #expect(decoded.structuredContent?.objectValue?["message"]?.stringValue == mixed) + } +} + +// MARK: - Tool Pagination Tests + +@Suite("Tool Pagination Tests") +struct ToolPaginationTests { + + @Test("ListTools cursor parameter encodes correctly") + func cursorParameterEncoding() throws { + let testCursor = "test-cursor-123" + let params = ListTools.Parameters(cursor: testCursor) + + let encoder = JSONEncoder() + let data = try encoder.encode(params) + let jsonString = String(data: data, encoding: .utf8)! + + #expect(jsonString.contains("\"cursor\":\"test-cursor-123\"")) + } + + @Test("ListTools result with nextCursor encodes correctly") + func resultWithNextCursor() throws { + let tools = [ + Tool(name: "tool1", inputSchema: ["type": "object"]), + Tool(name: "tool2", inputSchema: ["type": "object"]) + ] + let result = ListTools.Result(tools: tools, nextCursor: "next-page-token") + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(result) + let decoded = try decoder.decode(ListTools.Result.self, from: data) + + #expect(decoded.tools.count == 2) + #expect(decoded.nextCursor == "next-page-token") + } + + @Test("ListTools result without nextCursor indicates end of pagination") + func resultWithoutNextCursor() throws { + let tools = [ + Tool(name: "final_tool", inputSchema: ["type": "object"]) + ] + let result = ListTools.Result(tools: tools, nextCursor: nil) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(result) + let decoded = try decoder.decode(ListTools.Result.self, from: data) + + #expect(decoded.nextCursor == nil) + + // Verify null cursor is not included in JSON + let jsonString = String(data: data, encoding: .utf8)! + #expect(!jsonString.contains("nextCursor")) + } + + @Test("ListTools request with cursor decodes correctly") + func requestWithCursorDecoding() throws { + let jsonString = """ + {"jsonrpc":"2.0","id":"page-2","method":"tools/list","params":{"cursor":"page-1-token"}} + """ + let jsonData = jsonString.data(using: .utf8)! + + let decoded = try JSONDecoder().decode(Request.self, from: jsonData) + + #expect(decoded.id == "page-2") + #expect(decoded.params.cursor == "page-1-token") + } + + @Test("Simulated multi-page tool listing") + func simulatedMultiPageToolListing() throws { + // Simulate a server that returns 100 tools across multiple pages + let allTools = (0..<100).map { i in + Tool(name: "tool_\(i)", inputSchema: ["type": "object"]) + } + + let pageSize = 10 + var collectedTools: [Tool] = [] + var currentCursor: String? = nil + + // Simulate pagination + for pageIndex in 0..<10 { + let startIndex = pageIndex * pageSize + let endIndex = min(startIndex + pageSize, allTools.count) + let pageTools = Array(allTools[startIndex...Continuation? + + let name: String + + init(name: String, logger: Logger = Logger(label: "mcp.test.tracking-transport")) { + self.name = name + self.logger = logger + } + + public func connect() async throws { + isConnected = true + } + + public func disconnect() async { + isConnected = false + dataStreamContinuation?.finish() + dataStreamContinuation = nil + } + + public func send(_ data: Data) async throws { + sentMessages.append(SentMessage(data: data, relatedRequestId: nil)) + } + + public func send(_ data: Data, relatedRequestId: RequestId?) async throws { + sentMessages.append(SentMessage(data: data, relatedRequestId: relatedRequestId)) + } + + public func receive() -> AsyncThrowingStream { + return AsyncThrowingStream { continuation in + dataStreamContinuation = continuation + for message in dataToReceive { + continuation.yield(message) + } + dataToReceive.removeAll() + } + } + + func queue(data: Data) { + if let continuation = dataStreamContinuation { + continuation.yield(data) + } else { + dataToReceive.append(data) + } + } + + func queue(request: Request) throws { + queue(data: try encoder.encode(request)) + } + + func clearMessages() { + sentMessages.removeAll() + dataToReceive.removeAll() + } + } + + @Test("Response goes to correct transport when connection changes during request handling") + func testResponseGoesToCorrectTransportWhenConnectionChanges() async throws { + // Create two transports (simulating two clients) + let transportA = TrackingMockTransport(name: "TransportA") + let transportB = TrackingMockTransport(name: "TransportB") + + // Create a server with a custom handler that we can control + let server = Server(name: "TestServer", version: "1.0") + + // Use a continuation to control when the handler completes + actor HandlerControl { + var handlerContinuation: CheckedContinuation? + var handlerWasCalled = false + + func waitForSignal() async { + await withCheckedContinuation { continuation in + handlerContinuation = continuation + } + } + + func signalHandler() { + handlerContinuation?.resume() + handlerContinuation = nil + } + + func markCalled() { + handlerWasCalled = true + } + } + + let control = HandlerControl() + + // 1. Start server with transport A + try await server.start(transport: transportA) + + // Register a custom ping handler that waits for our signal AFTER start + // (start() registers default handlers, so we must override after) + // This simulates a slow handler to create a timing window + await server.withRequestHandler(Ping.self) { [control] _, _ in + await control.markCalled() + // Wait for signal before returning + await control.waitForSignal() + return Empty() + } + + // Wait for server to be ready + try await Task.sleep(for: .milliseconds(50)) + + // Initialize the server first + try await transportA.queue( + request: Initialize.request( + .init( + protocolVersion: Version.latest, + capabilities: .init(), + clientInfo: .init(name: "TestClient", version: "1.0") + ) + )) + + // Wait for initialization + try await Task.sleep(for: .milliseconds(100)) + await transportA.clearMessages() + + // 2. Send a ping request from transport A (with ID 100) + let pingJSON = """ + {"jsonrpc":"2.0","id":100,"method":"ping","params":{}} + """ + await transportA.queue(data: pingJSON.data(using: .utf8)!) + + // Wait for the handler to be called + try await Task.sleep(for: .milliseconds(50)) + + // Verify handler was called + let wasCalled = await control.handlerWasCalled + #expect(wasCalled == true, "Handler should have been called") + + // 3. While A's request is processing, switch to transport B + // This simulates another client connecting and reassigning server.connection + try await server.start(transport: transportB) + + // Give time for connection switch + try await Task.sleep(for: .milliseconds(50)) + + // 4. Complete A's request by signaling the handler + await control.signalHandler() + + // Wait for response to be sent + try await Task.sleep(for: .milliseconds(100)) + + // 5. Verify the response went to transport A (not transport B) + let messagesA = await transportA.sentMessages + let messagesB = await transportB.sentMessages + + #expect(messagesA.count >= 1, "Transport A should have received the response") + + // Find the response with ID 100 + let responseToA = messagesA.first { msg in + if let str = msg.asString { + return str.contains("\"id\":100") || str.contains("\"id\": 100") + } + return false + } + #expect(responseToA != nil, "Transport A should have received response for request 100") + + // Verify the related request ID was passed + if let response = responseToA { + #expect(response.relatedRequestId == .number(100), "Response should have relatedRequestId set to 100") + } + + // Transport B should NOT have received the response for request 100 + let responseToB = messagesB.first { msg in + if let str = msg.asString { + return str.contains("\"id\":100") || str.contains("\"id\": 100") + } + return false + } + #expect(responseToB == nil, "Transport B should NOT have received response for request 100") + + // Cleanup + await server.stop() + await transportA.disconnect() + await transportB.disconnect() + } + + @Test("Simple request routes correctly with relatedRequestId") + func testSimpleRequestRoutesCorrectly() async throws { + let transportA = TrackingMockTransport(name: "TransportA") + + let server = Server(name: "TestServer", version: "1.0") + + // Start server with transport A + try await server.start(transport: transportA) + try await Task.sleep(for: .milliseconds(50)) + + // Initialize + try await transportA.queue( + request: Initialize.request( + .init( + protocolVersion: Version.latest, + capabilities: .init(), + clientInfo: .init(name: "TestClient", version: "1.0") + ) + )) + try await Task.sleep(for: .milliseconds(100)) + await transportA.clearMessages() + + // Send ping request from A with a specific ID + let pingJSON = """ + {"jsonrpc":"2.0","id":42,"method":"ping","params":{}} + """ + await transportA.queue(data: pingJSON.data(using: .utf8)!) + + // Wait for response + try await Task.sleep(for: .milliseconds(100)) + + // Verify response went to transport A + let messagesA = await transportA.sentMessages + let pingResponse = messagesA.first { msg in + if let str = msg.asString { + return str.contains("\"id\":42") || str.contains("\"id\": 42") + } + return false + } + #expect(pingResponse != nil, "Transport A should have received ping response") + + // Verify relatedRequestId was set + if let response = pingResponse { + #expect(response.relatedRequestId == .number(42), "Response should have relatedRequestId set to 42") + } + + // Cleanup + await server.stop() + await transportA.disconnect() + } + + @Test("Batch response goes to correct transport") + func testBatchResponseGoesToCorrectTransport() async throws { + let transportA = TrackingMockTransport(name: "TransportA") + + let server = Server(name: "TestServer", version: "1.0") + + // Start server + try await server.start(transport: transportA) + try await Task.sleep(for: .milliseconds(50)) + + // Initialize + try await transportA.queue( + request: Initialize.request( + .init( + protocolVersion: Version.latest, + capabilities: .init(), + clientInfo: .init(name: "TestClient", version: "1.0") + ) + )) + try await Task.sleep(for: .milliseconds(100)) + await transportA.clearMessages() + + // Send a batch request + let batchJSON = """ + [ + {"jsonrpc":"2.0","id":1,"method":"ping","params":{}}, + {"jsonrpc":"2.0","id":2,"method":"ping","params":{}} + ] + """ + let batchData = batchJSON.data(using: .utf8)! + await transportA.queue(data: batchData) + + // Wait for response + try await Task.sleep(for: .milliseconds(100)) + + // Verify batch response went to transport A + let messagesA = await transportA.sentMessages + #expect(messagesA.count >= 1, "Transport A should have received the batch response") + + let batchResponse = messagesA.first { msg in + if let str = msg.asString { + // Batch response should be an array containing both IDs + return str.hasPrefix("[") && str.contains("\"id\":1") && str.contains("\"id\":2") + } + return false + } + #expect(batchResponse != nil, "Transport A should have received batch response with both request IDs") + + // Cleanup + await server.stop() + await transportA.disconnect() + } + + @Test("Error response goes to correct transport with relatedRequestId") + func testErrorResponseGoesToCorrectTransport() async throws { + let transportA = TrackingMockTransport(name: "TransportA") + + let server = Server(name: "TestServer", version: "1.0") + + // Start server + try await server.start(transport: transportA) + try await Task.sleep(for: .milliseconds(50)) + + // Initialize + try await transportA.queue( + request: Initialize.request( + .init( + protocolVersion: Version.latest, + capabilities: .init(), + clientInfo: .init(name: "TestClient", version: "1.0") + ) + )) + try await Task.sleep(for: .milliseconds(100)) + await transportA.clearMessages() + + // Send a request for an unknown method + let unknownMethodJSON = """ + {"jsonrpc":"2.0","id":99,"method":"unknown/method","params":{}} + """ + await transportA.queue(data: unknownMethodJSON.data(using: .utf8)!) + + // Wait for response + try await Task.sleep(for: .milliseconds(100)) + + // Verify error response went to transport A + let messagesA = await transportA.sentMessages + let errorResponse = messagesA.first { msg in + if let str = msg.asString { + return str.contains("\"id\":99") && str.contains("\"error\"") + } + return false + } + #expect(errorResponse != nil, "Transport A should have received error response") + + // Verify relatedRequestId was set even for error responses + if let response = errorResponse { + #expect(response.relatedRequestId == .number(99), "Error response should have relatedRequestId set to 99") + } + + // Cleanup + await server.stop() + await transportA.disconnect() + } + + @Test("Notifications sent via context include relatedRequestId") + func testNotificationsSentViaContextIncludeRelatedRequestId() async throws { + let transportA = TrackingMockTransport(name: "TransportA") + + let server = Server(name: "TestServer", version: "1.0") + + // Track when handler sends notification + actor NotificationTracker { + var notificationSent = false + func markSent() { notificationSent = true } + } + let tracker = NotificationTracker() + + // Start server with transport A + try await server.start(transport: transportA) + + // Register a handler that sends a notification using the context + await server.withRequestHandler(Ping.self) { [tracker] _, context in + // Send a notification mid-execution using the context + // The notification should include the relatedRequestId + try await context.sendNotification(ToolListChangedNotification()) + await tracker.markSent() + return Empty() + } + + try await Task.sleep(for: .milliseconds(50)) + + // Initialize + try await transportA.queue( + request: Initialize.request( + .init( + protocolVersion: Version.latest, + capabilities: .init(), + clientInfo: .init(name: "TestClient", version: "1.0") + ) + )) + try await Task.sleep(for: .milliseconds(100)) + await transportA.clearMessages() + + // Send ping with a specific ID + let pingJSON = """ + {"jsonrpc":"2.0","id":42,"method":"ping","params":{}} + """ + await transportA.queue(data: pingJSON.data(using: .utf8)!) + + // Wait for handler to execute + try await Task.sleep(for: .milliseconds(200)) + + // Check that notification was sent + let wasSent = await tracker.notificationSent + #expect(wasSent == true, "Handler should have sent a notification") + + // Verify the notification was sent with the correct relatedRequestId + let messagesA = await transportA.sentMessages + let notification = messagesA.first { msg in + if let str = msg.asString { + return str.contains("\"method\":\"notifications/tools/list_changed\"") + } + return false + } + #expect(notification != nil, "Transport A should have received the notification") + + // Verify the notification has the relatedRequestId of the original request + if let notif = notification { + #expect(notif.relatedRequestId == .number(42), "Notification should have relatedRequestId matching the request (42)") + } + + // Cleanup + await server.stop() + await transportA.disconnect() + } +} diff --git a/Tests/MCPTests/VersioningTests.swift b/Tests/MCPTests/VersioningTests.swift index d1896b53..b40300c4 100644 --- a/Tests/MCPTests/VersioningTests.swift +++ b/Tests/MCPTests/VersioningTests.swift @@ -13,9 +13,9 @@ struct VersioningTests { @Test("Client requests older supported version") func testClientRequestsOlderSupportedVersion() { - let clientVersion = "2024-11-05" + let clientVersion = Version.v2024_11_05 let negotiatedVersion = Version.negotiate(clientRequestedVersion: clientVersion) - #expect(negotiatedVersion == "2024-11-05") + #expect(negotiatedVersion == Version.v2024_11_05) } @Test("Client requests unsupported version") @@ -41,13 +41,15 @@ struct VersioningTests { @Test("Server's supported versions correctly defined") func testServerSupportedVersions() { - #expect(Version.supported.contains("2025-03-26")) - #expect(Version.supported.contains("2024-11-05")) - #expect(Version.supported.count == 2) + #expect(Version.supported.contains(Version.v2025_11_25)) + #expect(Version.supported.contains(Version.v2025_06_18)) + #expect(Version.supported.contains(Version.v2025_03_26)) + #expect(Version.supported.contains(Version.v2024_11_05)) + #expect(Version.supported.count == 4) } @Test("Server's latest version is correct") func testServerLatestVersion() { - #expect(Version.latest == "2025-03-26") + #expect(Version.latest == Version.v2025_11_25) } } From a89979ec4c7ae5863234810d9565c73173f0e823 Mon Sep 17 00:00:00 2001 From: Anthony DePasquale Date: Sun, 4 Jan 2026 11:41:22 +0100 Subject: [PATCH 4/8] Split long client and server files into smaller extensions --- .../HTTPClientTransport+Types.swift | 42 + .../Base/Transports/HTTPClientTransport.swift | 38 - Sources/MCP/Client/Client+Batching.swift | 164 ++ .../MCP/Client/Client+MessageHandling.swift | 447 ++++ .../MCP/Client/Client+ProtocolMethods.swift | 142 ++ Sources/MCP/Client/Client+Registration.swift | 229 +++ Sources/MCP/Client/Client+Requests.swift | 471 +++++ Sources/MCP/Client/Client.swift | 1821 +---------------- .../MCP/Server/Server+ClientRequests.swift | 359 ++++ .../MCP/Server/Server+RequestHandling.swift | 366 ++++ Sources/MCP/Server/Server+Sending.swift | 66 + Sources/MCP/Server/Server.swift | 820 +------- 12 files changed, 2371 insertions(+), 2594 deletions(-) create mode 100644 Sources/MCP/Base/Transports/HTTPClientTransport+Types.swift create mode 100644 Sources/MCP/Client/Client+Batching.swift create mode 100644 Sources/MCP/Client/Client+MessageHandling.swift create mode 100644 Sources/MCP/Client/Client+ProtocolMethods.swift create mode 100644 Sources/MCP/Client/Client+Registration.swift create mode 100644 Sources/MCP/Client/Client+Requests.swift create mode 100644 Sources/MCP/Server/Server+ClientRequests.swift create mode 100644 Sources/MCP/Server/Server+RequestHandling.swift create mode 100644 Sources/MCP/Server/Server+Sending.swift diff --git a/Sources/MCP/Base/Transports/HTTPClientTransport+Types.swift b/Sources/MCP/Base/Transports/HTTPClientTransport+Types.swift new file mode 100644 index 00000000..d76f3416 --- /dev/null +++ b/Sources/MCP/Base/Transports/HTTPClientTransport+Types.swift @@ -0,0 +1,42 @@ +import Foundation + +// Types extracted from HTTPClientTransport.swift +// - HTTPReconnectionOptions + +/// Configuration options for reconnection behavior of the HTTPClientTransport. +/// +/// These options control how the transport handles SSE stream disconnections +/// and reconnection attempts. +public struct HTTPReconnectionOptions: Sendable { + /// Initial delay between reconnection attempts in seconds. + /// Default is 1.0 second. + public var initialReconnectionDelay: TimeInterval + + /// Maximum delay between reconnection attempts in seconds. + /// Default is 30.0 seconds. + public var maxReconnectionDelay: TimeInterval + + /// Factor by which the reconnection delay increases after each attempt. + /// Default is 1.5. + public var reconnectionDelayGrowFactor: Double + + /// Maximum number of reconnection attempts before giving up. + /// Default is 2. + public var maxRetries: Int + + /// Creates reconnection options with default values. + public init( + initialReconnectionDelay: TimeInterval = 1.0, + maxReconnectionDelay: TimeInterval = 30.0, + reconnectionDelayGrowFactor: Double = 1.5, + maxRetries: Int = 2 + ) { + self.initialReconnectionDelay = initialReconnectionDelay + self.maxReconnectionDelay = maxReconnectionDelay + self.reconnectionDelayGrowFactor = reconnectionDelayGrowFactor + self.maxRetries = maxRetries + } + + /// Default reconnection options. + public static let `default` = HTTPReconnectionOptions() +} diff --git a/Sources/MCP/Base/Transports/HTTPClientTransport.swift b/Sources/MCP/Base/Transports/HTTPClientTransport.swift index dde41daf..214486bf 100644 --- a/Sources/MCP/Base/Transports/HTTPClientTransport.swift +++ b/Sources/MCP/Base/Transports/HTTPClientTransport.swift @@ -9,44 +9,6 @@ import Logging import FoundationNetworking #endif -/// Configuration options for reconnection behavior of the HTTPClientTransport. -/// -/// These options control how the transport handles SSE stream disconnections -/// and reconnection attempts. -public struct HTTPReconnectionOptions: Sendable { - /// Initial delay between reconnection attempts in seconds. - /// Default is 1.0 second. - public var initialReconnectionDelay: TimeInterval - - /// Maximum delay between reconnection attempts in seconds. - /// Default is 30.0 seconds. - public var maxReconnectionDelay: TimeInterval - - /// Factor by which the reconnection delay increases after each attempt. - /// Default is 1.5. - public var reconnectionDelayGrowFactor: Double - - /// Maximum number of reconnection attempts before giving up. - /// Default is 2. - public var maxRetries: Int - - /// Creates reconnection options with default values. - public init( - initialReconnectionDelay: TimeInterval = 1.0, - maxReconnectionDelay: TimeInterval = 30.0, - reconnectionDelayGrowFactor: Double = 1.5, - maxRetries: Int = 2 - ) { - self.initialReconnectionDelay = initialReconnectionDelay - self.maxReconnectionDelay = maxReconnectionDelay - self.reconnectionDelayGrowFactor = reconnectionDelayGrowFactor - self.maxRetries = maxRetries - } - - /// Default reconnection options. - public static let `default` = HTTPReconnectionOptions() -} - /// An implementation of the MCP Streamable HTTP transport protocol for clients. /// /// This transport implements the [Streamable HTTP transport](https://modelcontextprotocol.io/specification/2025-03-26/basic/transports#streamable-http) diff --git a/Sources/MCP/Client/Client+Batching.swift b/Sources/MCP/Client/Client+Batching.swift new file mode 100644 index 00000000..427b893c --- /dev/null +++ b/Sources/MCP/Client/Client+Batching.swift @@ -0,0 +1,164 @@ +import Foundation + +extension Client { + // MARK: - Batching + + /// A batch of requests. + /// + /// Objects of this type are passed as an argument to the closure + /// of the ``Client/withBatch(_:)`` method. + public actor Batch { + unowned let client: Client + var requests: [AnyRequest] = [] + + init(client: Client) { + self.client = client + } + + /// Adds a request to the batch and prepares its expected response task. + /// The actual sending happens when the `withBatch` scope completes. + /// - Returns: A `Task` that will eventually produce the result or throw an error. + public func addRequest(_ request: Request) async throws -> Task< + M.Result, Swift.Error + > { + requests.append(try AnyRequest(request)) + + // Create stream for receiving the response + let (stream, continuation) = AsyncThrowingStream.makeStream() + + // Clean up pending request if caller cancels (e.g., task cancelled) + // and send CancelledNotification to server per MCP spec + let requestId = request.id + continuation.onTermination = { @Sendable [weak client] termination in + Task { + guard let client else { return } + await client.cleanUpPendingRequest(id: requestId) + + // Per MCP spec: send notifications/cancelled when cancelling a request + // Only send if the stream was cancelled (not finished normally) + if case .cancelled = termination { + await client.sendCancellationNotification( + requestId: requestId, + reason: "Client cancelled the batch request" + ) + } + } + } + + // Register the pending request + await client.addPendingRequest(id: request.id, continuation: continuation) + + // Return a Task that waits for the response via the stream + return Task { + for try await result in stream { + return result + } + throw MCPError.internalError("No response received") + } + } + } + + /// Executes multiple requests in a single batch. + /// + /// This method allows you to group multiple MCP requests together, + /// which are then sent to the server as a single JSON array. + /// The server processes these requests and sends back a corresponding + /// JSON array of responses. + /// + /// Within the `body` closure, use the provided `Batch` actor to add + /// requests using `batch.addRequest(_:)`. Each call to `addRequest` + /// returns a `Task` handle representing the asynchronous operation + /// for that specific request's result. + /// + /// It's recommended to collect these `Task` handles into an array + /// within the `body` closure`. After the `withBatch` method returns + /// (meaning the batch request has been sent), you can then process + /// the results by awaiting each `Task` in the collected array. + /// + /// Example 1: Batching multiple tool calls and collecting typed tasks: + /// ```swift + /// // Array to hold the task handles for each tool call + /// var toolTasks: [Task] = [] + /// try await client.withBatch { batch in + /// for i in 0..<10 { + /// toolTasks.append( + /// try await batch.addRequest( + /// CallTool.request(.init(name: "square", arguments: ["n": i])) + /// ) + /// ) + /// } + /// } + /// + /// // Process results after the batch is sent + /// print("Processing \(toolTasks.count) tool results...") + /// for (index, task) in toolTasks.enumerated() { + /// do { + /// let result = try await task.value + /// print("\(index): \(result.content)") + /// } catch { + /// print("\(index) failed: \(error)") + /// } + /// } + /// ``` + /// + /// Example 2: Batching different request types and awaiting individual tasks: + /// ```swift + /// // Declare optional task variables beforehand + /// var pingTask: Task? + /// var promptTask: Task? + /// + /// try await client.withBatch { batch in + /// // Assign the tasks within the batch closure + /// pingTask = try await batch.addRequest(Ping.request()) + /// promptTask = try await batch.addRequest(GetPrompt.request(.init(name: "greeting"))) + /// } + /// + /// // Await the results after the batch is sent + /// do { + /// if let pingTask = pingTask { + /// try await pingTask.value // Await ping result (throws if ping failed) + /// print("Ping successful") + /// } + /// if let promptTask = promptTask { + /// let promptResult = try await promptTask.value // Await prompt result + /// print("Prompt description: \(promptResult.description ?? "None")") + /// } + /// } catch { + /// print("Error processing batch results: \(error)") + /// } + /// ``` + /// + /// - Parameter body: An asynchronous closure that takes a `Batch` object as input. + /// Use this object to add requests to the batch. + /// - Throws: `MCPError.internalError` if the client is not connected. + /// Can also rethrow errors from the `body` closure or from sending the batch request. + public func withBatch(body: @escaping (Batch) async throws -> Void) async throws { + guard let connection = connection else { + throw MCPError.internalError("Client connection not initialized") + } + + // Create Batch actor, passing self (Client) + let batch = Batch(client: self) + + // Populate the batch actor by calling the user's closure. + try await body(batch) + + // Get the collected requests from the batch actor + let requests = await batch.requests + + // Check if there are any requests to send + guard !requests.isEmpty else { + await logger?.debug("Batch requested but no requests were added.") + return // Nothing to send + } + + await logger?.debug( + "Sending batch request", metadata: ["count": "\(requests.count)"]) + + // Encode the array of AnyMethod requests into a single JSON payload + let data = try encoder.encode(requests) + try await connection.send(data) + + // Responses will be handled asynchronously by the message loop and handleBatchResponse/handleResponse. + } +} diff --git a/Sources/MCP/Client/Client+MessageHandling.swift b/Sources/MCP/Client/Client+MessageHandling.swift new file mode 100644 index 00000000..2e57baea --- /dev/null +++ b/Sources/MCP/Client/Client+MessageHandling.swift @@ -0,0 +1,447 @@ +import Foundation + +extension Client { + // MARK: - Message Handling + + func handleResponse(_ response: Response) async { + await logger?.trace( + "Processing response", + metadata: ["id": "\(response.id)"]) + + // Check for task-augmented response BEFORE resuming the request. + // Per MCP spec 2025-11-25: progress tokens continue for task lifetime. + // If this is a CreateTaskResult, we need to keep the progress handler alive. + if case .success(let value) = response.result, + case .object(let resultObject) = value { + checkForTaskResponse(response: response, value: resultObject) + } + + // Attempt to remove the pending request using the response ID. + // Resume with the response only if it hadn't yet been removed. + if let removedRequest = self.removePendingRequest(id: response.id) { + // If we successfully removed it, resume its continuation. + switch response.result { + case .success(let value): + removedRequest.resume(returning: value) + case .failure(let error): + removedRequest.resume(throwing: error) + } + } else { + // Request was already removed (e.g., by send error handler or disconnect). + // Log this, but it's not an error in race condition scenarios. + await logger?.warning( + "Attempted to handle response for already removed request", + metadata: ["id": "\(response.id)"] + ) + } + } + + /// Check if a response is a task-augmented response (CreateTaskResult). + /// + /// If the response contains a `task` object with `taskId`, this is a task-augmented + /// response. Per MCP spec, progress notifications can continue until the task reaches + /// terminal status, so we migrate the progress handler from request tracking to task tracking. + /// + /// This matches the TypeScript SDK pattern where task progress tokens are kept alive + /// until the task completes. + func checkForTaskResponse(response: Response, value: [String: Value]) { + // Check if we have a progress token for this request + guard let progressToken = requestProgressTokens[response.id] else { return } + + // Check if response has task.taskId (CreateTaskResult pattern) + // This mirrors TypeScript's check: result.task?.taskId + guard let taskValue = value["task"], + case .object(let taskObject) = taskValue, + let taskIdValue = taskObject["taskId"], + case .string(let taskId) = taskIdValue else { + // Not a task response - clean up request tracking + // (the progress callback itself is cleaned up in send() after receiving result) + requestProgressTokens.removeValue(forKey: response.id) + return + } + + // This is a task-augmented response! + // Migrate progress token from request tracking to task tracking. + // This keeps the progress handler alive until the task completes. + taskProgressTokens[taskId] = progressToken + requestProgressTokens.removeValue(forKey: response.id) + + Task { + await logger?.debug( + "Keeping progress handler alive for task", + metadata: [ + "taskId": "\(taskId)", + "progressToken": "\(progressToken)", + ] + ) + } + } + + /// Clean up the progress handler for a completed task. + /// + /// Call this method when a task reaches terminal status (completed, failed, cancelled) + /// to remove the progress callback and timeout controller. + /// + /// ## Example + /// + /// ```swift + /// // Register task status notification handler + /// await client.onNotification(TaskStatusNotification.self) { message in + /// if message.params.status.isTerminal { + /// await client.cleanupTaskProgressHandler(taskId: message.params.taskId) + /// } + /// } + /// ``` + /// + /// - Parameter taskId: The ID of the task that completed. + public func cleanUpTaskProgressHandler(taskId: String) { + guard let progressToken = taskProgressTokens.removeValue(forKey: taskId) else { return } + + progressCallbacks.removeValue(forKey: progressToken) + timeoutControllers.removeValue(forKey: progressToken) + + Task { + await logger?.debug( + "Cleaned up progress handler for completed task", + metadata: ["taskId": "\(taskId)"] + ) + } + } + + func handleMessage(_ message: Message) async { + await logger?.trace( + "Processing notification", + metadata: ["method": "\(message.method)"]) + + // Check if this is a progress notification and invoke any registered callback + if message.method == ProgressNotification.name { + await handleProgressNotification(message) + } + + // Check if this is a task status notification and clean up progress handlers + // for terminal task statuses (per MCP spec, progress tokens are valid until terminal status) + if message.method == TaskStatusNotification.name { + await handleTaskStatusNotification(message) + } + + // Find notification handlers for this method + guard let handlers = notificationHandlers[message.method] else { return } + + // Convert notification parameters to concrete type and call handlers + for handler in handlers { + do { + try await handler(message) + } catch { + await logger?.error( + "Error handling notification", + metadata: [ + "method": "\(message.method)", + "error": "\(error)", + ]) + } + } + } + + /// Handle a progress notification by invoking any registered callback. + func handleProgressNotification(_ message: Message) async { + do { + // Decode as ProgressNotification.Parameters + let paramsData = try encoder.encode(message.params) + let params = try decoder.decode(ProgressNotification.Parameters.self, from: paramsData) + + // Look up the callback for this token + guard let callback = progressCallbacks[params.progressToken] else { + // TypeScript SDK logs an error for unknown progress tokens + await logger?.warning( + "Received progress notification for unknown token", + metadata: ["progressToken": "\(params.progressToken)"]) + return + } + + // Signal the timeout controller if one exists for this token + // This allows resetTimeoutOnProgress to work + if let timeoutController = timeoutControllers[params.progressToken] { + await timeoutController.signalProgress() + } + + // Invoke the callback + let progress = Progress( + value: params.progress, + total: params.total, + message: params.message + ) + await callback(progress) + } catch { + await logger?.warning( + "Failed to decode progress notification", + metadata: ["error": "\(error)"]) + } + } + + /// Handle a task status notification by cleaning up progress handlers for terminal tasks. + /// + /// Per MCP spec 2025-11-25: progress tokens continue throughout task lifetime until terminal status. + /// This method automatically cleans up progress handlers when a task reaches completed, failed, or cancelled. + func handleTaskStatusNotification(_ message: Message) async { + do { + // Decode as TaskStatusNotification.Parameters + let paramsData = try encoder.encode(message.params) + let params = try decoder.decode(TaskStatusNotification.Parameters.self, from: paramsData) + + // If the task reached a terminal status, clean up its progress handler + if params.status.isTerminal { + cleanUpTaskProgressHandler(taskId: params.taskId) + } + } catch { + // Don't log errors for task status notifications - they may not be task-related + // and the user may not have registered a handler for them + } + } + + /// Handle an incoming request from the server (bidirectional communication). + /// + /// This enables server→client requests such as sampling, roots, and elicitation. + /// + /// ## Task-Augmented Request Handling + /// + /// For `sampling/createMessage` and `elicitation/create` requests, this method + /// checks for a `task` field in the request params. If present, it routes to + /// the task-augmented handler (which returns `CreateTaskResult`) instead of + /// the normal handler. + /// + /// This follows the Python SDK pattern of storing task-augmented handlers + /// separately and checking at dispatch time, rather than the TypeScript pattern + /// of wrapping handlers at registration time. The Python pattern was chosen + /// because: + /// - It allows handlers to be registered in any order without losing task-awareness + /// - It keeps task logic separate from normal handler logic + /// - It's more explicit about which handler is called for which request type + func handleIncomingRequest(_ request: Request) async { + await logger?.trace( + "Processing incoming request from server", + metadata: [ + "method": "\(request.method)", + "id": "\(request.id)", + ]) + + // Validate elicitation mode against client capabilities + // Per spec: Client MUST return -32602 if server requests unsupported mode + if request.method == Elicit.name { + if let modeError = await validateElicitationMode(request) { + await sendResponse(modeError) + return + } + } + + // Check for task-augmented sampling/elicitation requests first + // This matches the Python SDK pattern where task detection happens at dispatch time + if let taskResponse = await handleTaskAugmentedRequest(request) { + await sendResponse(taskResponse) + return + } + + // Find handler for method name + guard let handler = requestHandlers[request.method] else { + await logger?.warning( + "No handler registered for server request", + metadata: ["method": "\(request.method)"]) + + // Send error response + let response = AnyMethod.response( + id: request.id, + error: MCPError.methodNotFound("Client has no handler for: \(request.method)") + ) + await sendResponse(response) + return + } + + // Execute the handler and send response + do { + let response = try await handler(request) + + // Check cancellation before sending response (per MCP spec: + // "Receivers of a cancellation notification SHOULD... Not send a response + // for the cancelled request") + if Task.isCancelled { + await logger?.debug( + "Server request cancelled, suppressing response", + metadata: ["id": "\(request.id)"] + ) + return + } + + await sendResponse(response) + } catch { + // Also check cancellation on error path - don't send error response if cancelled + if Task.isCancelled { + await logger?.debug( + "Server request cancelled during error handling, suppressing response", + metadata: ["id": "\(request.id)"] + ) + return + } + + await logger?.error( + "Error handling server request", + metadata: [ + "method": "\(request.method)", + "error": "\(error)", + ]) + let errorResponse = AnyMethod.response( + id: request.id, + error: (error as? MCPError) ?? MCPError.internalError(error.localizedDescription) + ) + await sendResponse(errorResponse) + } + } + + /// Validate that an elicitation request uses a mode supported by client capabilities. + /// + /// Per MCP spec: Client MUST return -32602 (Invalid params) if server sends + /// an elicitation/create request with a mode not declared in client capabilities. + /// + /// - Parameter request: The incoming elicitation request + /// - Returns: An error response if mode is unsupported, nil if valid + func validateElicitationMode(_ request: Request) async -> Response? { + do { + let paramsData = try encoder.encode(request.params) + let params = try decoder.decode(Elicit.Parameters.self, from: paramsData) + + switch params { + case .form: + // Form mode requires form capability + if capabilities.elicitation?.form == nil { + return Response( + id: request.id, + error: .invalidParams("Client does not support form elicitation mode") + ) + } + case .url: + // URL mode requires url capability + if capabilities.elicitation?.url == nil { + return Response( + id: request.id, + error: .invalidParams("Client does not support URL elicitation mode") + ) + } + } + } catch { + // If we can't decode the params, let the normal handler deal with it + await logger?.warning( + "Failed to decode elicitation params for mode validation", + metadata: ["error": "\(error)"]) + } + + return nil + } + + /// Check if a request is task-augmented and handle it if so. + /// + /// - Parameter request: The incoming request + /// - Returns: A response if the request was task-augmented and handled, nil otherwise + func handleTaskAugmentedRequest(_ request: Request) async -> Response? { + do { + // Check for task-augmented sampling request + if request.method == CreateSamplingMessage.name, + let taskHandler = taskAugmentedSamplingHandler { + let paramsData = try encoder.encode(request.params) + let params = try decoder.decode(CreateSamplingMessage.Parameters.self, from: paramsData) + + if let taskMetadata = params.task { + let result = try await taskHandler(params, taskMetadata) + let resultData = try encoder.encode(result) + let resultValue = try decoder.decode(Value.self, from: resultData) + return Response(id: request.id, result: resultValue) + } + } + + // Check for task-augmented elicitation request + if request.method == Elicit.name, + let taskHandler = taskAugmentedElicitationHandler { + let paramsData = try encoder.encode(request.params) + let params = try decoder.decode(Elicit.Parameters.self, from: paramsData) + + let taskMetadata: TaskMetadata? = switch params { + case .form(let formParams): formParams.task + case .url(let urlParams): urlParams.task + } + + if let taskMetadata { + let result = try await taskHandler(params, taskMetadata) + let resultData = try encoder.encode(result) + let resultValue = try decoder.decode(Value.self, from: resultData) + return Response(id: request.id, result: resultValue) + } + } + } catch let error as MCPError { + return Response(id: request.id, error: error) + } catch { + return Response(id: request.id, error: MCPError.internalError(error.localizedDescription)) + } + + // Not a task-augmented request + return nil + } + + /// Send a response back to the server. + func sendResponse(_ response: Response) async { + guard let connection = connection else { + await logger?.warning("Cannot send response - client not connected") + return + } + + do { + let responseData = try encoder.encode(response) + try await connection.send(responseData) + } catch { + await logger?.error( + "Failed to send response to server", + metadata: ["error": "\(error)"]) + } + } + + // MARK: - + + /// Validate the server capabilities. + /// Throws an error if the client is configured to be strict and the capability is not supported. + func validateServerCapability( + _ keyPath: KeyPath, + _ name: String + ) + throws + { + if configuration.strict { + guard let capabilities = serverCapabilities else { + throw MCPError.methodNotFound("Server capabilities not initialized") + } + guard capabilities[keyPath: keyPath] != nil else { + throw MCPError.methodNotFound("\(name) is not supported by the server") + } + } + } + + // Add handler for batch responses + func handleBatchResponse(_ responses: [AnyResponse]) async { + await logger?.trace("Processing batch response", metadata: ["count": "\(responses.count)"]) + for response in responses { + // Attempt to remove the pending request. + // If successful, pendingRequest contains the request. + if let pendingRequest = self.removePendingRequest(id: response.id) { + // If we successfully removed it, handle the response using the pending request. + switch response.result { + case .success(let value): + pendingRequest.resume(returning: value) + case .failure(let error): + pendingRequest.resume(throwing: error) + } + } else { + // If removal failed, it means the request ID was not found (or already handled). + // Log a warning. + await logger?.warning( + "Received response in batch for unknown or already handled request ID", + metadata: ["id": "\(response.id)"] + ) + } + } + } +} diff --git a/Sources/MCP/Client/Client+ProtocolMethods.swift b/Sources/MCP/Client/Client+ProtocolMethods.swift new file mode 100644 index 00000000..dbb23517 --- /dev/null +++ b/Sources/MCP/Client/Client+ProtocolMethods.swift @@ -0,0 +1,142 @@ +import Foundation + +extension Client { + // MARK: - Prompts + + public func getPrompt(name: String, arguments: [String: String]? = nil) async throws + -> (description: String?, messages: [Prompt.Message]) + { + try validateServerCapability(\.prompts, "Prompts") + let request = GetPrompt.request(.init(name: name, arguments: arguments)) + let result = try await send(request) + return (description: result.description, messages: result.messages) + } + + public func listPrompts(cursor: String? = nil) async throws + -> (prompts: [Prompt], nextCursor: String?) + { + try validateServerCapability(\.prompts, "Prompts") + let request: Request + if let cursor = cursor { + request = ListPrompts.request(.init(cursor: cursor)) + } else { + request = ListPrompts.request(.init()) + } + let result = try await send(request) + return (prompts: result.prompts, nextCursor: result.nextCursor) + } + + // MARK: - Resources + + public func readResource(uri: String) async throws -> [Resource.Content] { + try validateServerCapability(\.resources, "Resources") + let request = ReadResource.request(.init(uri: uri)) + let result = try await send(request) + return result.contents + } + + public func listResources(cursor: String? = nil) async throws -> ( + resources: [Resource], nextCursor: String? + ) { + try validateServerCapability(\.resources, "Resources") + let request: Request + if let cursor = cursor { + request = ListResources.request(.init(cursor: cursor)) + } else { + request = ListResources.request(.init()) + } + let result = try await send(request) + return (resources: result.resources, nextCursor: result.nextCursor) + } + + public func subscribeToResource(uri: String) async throws { + try validateServerCapability(\.resources?.subscribe, "Resource subscription") + let request = ResourceSubscribe.request(.init(uri: uri)) + _ = try await send(request) + } + + public func unsubscribeFromResource(uri: String) async throws { + try validateServerCapability(\.resources?.subscribe, "Resource subscription") + let request = ResourceUnsubscribe.request(.init(uri: uri)) + _ = try await send(request) + } + + public func listResourceTemplates(cursor: String? = nil) async throws -> ( + templates: [Resource.Template], nextCursor: String? + ) { + try validateServerCapability(\.resources, "Resources") + let request: Request + if let cursor = cursor { + request = ListResourceTemplates.request(.init(cursor: cursor)) + } else { + request = ListResourceTemplates.request(.init()) + } + let result = try await send(request) + return (templates: result.templates, nextCursor: result.nextCursor) + } + + // MARK: - Tools + + public func listTools(cursor: String? = nil) async throws -> ( + tools: [Tool], nextCursor: String? + ) { + try validateServerCapability(\.tools, "Tools") + let request: Request + if let cursor = cursor { + request = ListTools.request(.init(cursor: cursor)) + } else { + request = ListTools.request(.init()) + } + let result = try await send(request) + return (tools: result.tools, nextCursor: result.nextCursor) + } + + public func callTool(name: String, arguments: [String: Value]? = nil) async throws -> ( + content: [Tool.Content], structuredContent: Value?, isError: Bool? + ) { + try validateServerCapability(\.tools, "Tools") + let request = CallTool.request(.init(name: name, arguments: arguments)) + let result = try await send(request) + // TODO: Add client-side output validation against the tool's outputSchema. + // TypeScript and Python SDKs cache tool outputSchemas from listTools() and + // validate structuredContent when receiving tool results. + return (content: result.content, structuredContent: result.structuredContent, isError: result.isError) + } + + // MARK: - Completions + + /// Request completion suggestions from the server. + /// + /// Completions provide autocomplete suggestions for prompt arguments or resource + /// template URI parameters. + /// + /// - Parameters: + /// - ref: A reference to the prompt or resource template to get completions for. + /// - argument: The argument being completed, including its name and partial value. + /// - context: Optional additional context with previously-resolved argument values. + /// - Returns: The completion suggestions from the server. + public func complete( + ref: CompletionReference, + argument: CompletionArgument, + context: CompletionContext? = nil + ) async throws -> CompletionSuggestions { + try validateServerCapability(\.completions, "Completions") + let request = Complete.request(.init(ref: ref, argument: argument, context: context)) + let result = try await send(request) + return result.completion + } + + // MARK: - Logging + + /// Set the minimum log level for messages from the server. + /// + /// After calling this method, the server should only send log messages + /// at the specified level or higher (more severe). + /// + /// - Parameter level: The minimum log level to receive. + public func setLoggingLevel(_ level: LoggingLevel) async throws { + try validateServerCapability(\.logging, "Logging") + let request = SetLoggingLevel.request(.init(level: level)) + _ = try await send(request) + } +} diff --git a/Sources/MCP/Client/Client+Registration.swift b/Sources/MCP/Client/Client+Registration.swift new file mode 100644 index 00000000..bab2886e --- /dev/null +++ b/Sources/MCP/Client/Client+Registration.swift @@ -0,0 +1,229 @@ +import Foundation + +extension Client { + // MARK: - Handler Registration + + /// Register a handler for a notification + @discardableResult + public func onNotification( + _ type: N.Type, + handler: @escaping @Sendable (Message) async throws -> Void + ) async -> Self { + notificationHandlers[N.name, default: []].append(TypedNotificationHandler(handler)) + return self + } + + /// Send a notification to the server + public func notify(_ notification: Message) async throws { + guard let connection else { + throw MCPError.internalError("Client connection not initialized") + } + + let notificationData = try encoder.encode(notification) + try await connection.send(notificationData) + } + + /// Send a progress notification to the server. + /// + /// This is a convenience method for sending progress notifications from the client + /// to the server. This enables bidirectional progress reporting where clients can + /// inform servers about their own progress (e.g., during client-side processing). + /// + /// ## Example + /// + /// ```swift + /// // Client reports its own progress to the server + /// try await client.sendProgressNotification( + /// token: .string("client-task-123"), + /// progress: 50.0, + /// total: 100.0, + /// message: "Processing client-side data..." + /// ) + /// ``` + /// + /// - Parameters: + /// - token: The progress token to associate with this notification + /// - progress: The current progress value (should increase monotonically) + /// - total: The total progress value, if known + /// - message: An optional human-readable message describing current progress + public func sendProgressNotification( + token: ProgressToken, + progress: Double, + total: Double? = nil, + message: String? = nil + ) async throws { + try await notify(ProgressNotification.message(.init( + progressToken: token, + progress: progress, + total: total, + message: message + ))) + } + + /// Send a notification that the list of available roots has changed. + /// + /// Servers that receive this notification should request an updated + /// list of roots via the roots/list request. + /// + /// - Throws: `MCPError.invalidRequest` if the client has not declared + /// the `roots.listChanged` capability. + public func sendRootsChanged() async throws { + guard capabilities.roots?.listChanged == true else { + throw MCPError.invalidRequest( + "Client does not support roots.listChanged capability") + } + try await notify(RootsListChangedNotification.message(.init())) + } + + /// Register a handler for server→client requests. + /// + /// This enables bidirectional communication where the server can send requests + /// to the client (e.g., sampling, roots, elicitation). + /// + /// - Parameters: + /// - type: The method type to handle + /// - handler: The handler function that receives parameters and returns a result + /// - Returns: Self for chaining + @discardableResult + public func withRequestHandler( + _ type: M.Type, + handler: @escaping @Sendable (M.Parameters) async throws -> M.Result + ) -> Self { + requestHandlers[M.name] = TypedClientRequestHandler(handler) + return self + } + + /// Register a handler for `roots/list` requests from the server. + /// + /// When the server requests the list of roots, this handler will be called + /// to provide the available filesystem directories. + /// + /// - Important: The client must have declared `roots` capability during initialization. + /// + /// - Parameter handler: A closure that returns the list of available roots. + /// - Returns: Self for chaining. + /// - Precondition: `capabilities.roots` must be non-nil. + @discardableResult + public func withRootsHandler( + _ handler: @escaping @Sendable () async throws -> [Root] + ) -> Self { + precondition( + capabilities.roots != nil, + "Cannot register roots handler: Client does not have roots capability" + ) + return withRequestHandler(ListRoots.self) { _ in + ListRoots.Result(roots: try await handler()) + } + } + + /// Register a handler for `sampling/createMessage` requests from the server. + /// + /// When the server requests a sampling completion, this handler will be called + /// to generate the LLM response. + /// + /// The handler receives parameters that may or may not include tools. Check `params.hasTools` + /// to determine if tool use is enabled for this request. + /// + /// - Important: The client must have declared `sampling` capability during initialization. + /// + /// ## Example + /// + /// ```swift + /// client.withSamplingHandler { params in + /// // Call your LLM with the messages + /// let response = try await llm.complete( + /// messages: params.messages, + /// tools: params.tools, // May be nil + /// maxTokens: params.maxTokens + /// ) + /// + /// return ClientSamplingRequest.Result( + /// model: "gpt-4", + /// stopReason: .endTurn, + /// role: .assistant, + /// content: .text(response.text) + /// ) + /// } + /// ``` + /// + /// - Parameter handler: A closure that receives sampling parameters and returns the result. + /// - Returns: Self for chaining. + /// - Precondition: `capabilities.sampling` must be non-nil. + @discardableResult + public func withSamplingHandler( + _ handler: @escaping @Sendable (ClientSamplingRequest.Parameters) async throws -> ClientSamplingRequest.Result + ) -> Self { + precondition( + capabilities.sampling != nil, + "Cannot register sampling handler: Client does not have sampling capability" + ) + return withRequestHandler(ClientSamplingRequest.self, handler: handler) + } + + /// Register a handler for `elicitation/create` requests from the server. + /// + /// When the server requests user input via elicitation, this handler will be called + /// to collect the input and return the result. + /// + /// - Important: The client must have declared `elicitation` capability during initialization. + /// + /// - Parameter handler: A closure that receives elicitation parameters and returns the result. + /// - Returns: Self for chaining. + /// - Precondition: `capabilities.elicitation` must be non-nil. + @discardableResult + public func withElicitationHandler( + _ handler: @escaping @Sendable (Elicit.Parameters) async throws -> Elicit.Result + ) -> Self { + precondition( + capabilities.elicitation != nil, + "Cannot register elicitation handler: Client does not have elicitation capability" + ) + return withRequestHandler(Elicit.self, handler: handler) + } + + /// Internal method to set a request handler box directly. + /// + /// This is used by task-augmented handlers that need to return different result types + /// based on whether the request has a `task` field. + /// + /// - Important: This is an internal API that may change without notice. + internal func _setRequestHandler(method: String, handler: ClientRequestHandlerBox) { + requestHandlers[method] = handler + } + + /// Internal method to get an existing request handler box. + /// + /// This is used to retrieve the existing handler before wrapping it with + /// a task-aware handler that preserves the normal handler as a fallback. + /// + /// - Important: This is an internal API that may change without notice. + internal func _getRequestHandler(method: String) -> ClientRequestHandlerBox? { + requestHandlers[method] + } + + /// Internal method to set the task-augmented sampling handler. + /// + /// This handler is called when the server sends a `sampling/createMessage` request + /// with a `task` field. The handler should return `CreateTaskResult` instead of + /// the normal sampling result. + /// + /// - Important: This is an internal API that may change without notice. + internal func _setTaskAugmentedSamplingHandler( + _ handler: @escaping ExperimentalClientTaskHandlers.TaskAugmentedSamplingHandler + ) { + taskAugmentedSamplingHandler = handler + } + + /// Internal method to set the task-augmented elicitation handler. + /// + /// This handler is called when the server sends an `elicitation/create` request + /// with a `task` field. The handler should return `CreateTaskResult` instead of + /// the normal elicitation result. + /// + /// - Important: This is an internal API that may change without notice. + internal func _setTaskAugmentedElicitationHandler( + _ handler: @escaping ExperimentalClientTaskHandlers.TaskAugmentedElicitationHandler + ) { + taskAugmentedElicitationHandler = handler + } +} diff --git a/Sources/MCP/Client/Client+Requests.swift b/Sources/MCP/Client/Client+Requests.swift new file mode 100644 index 00000000..e96c3d2f --- /dev/null +++ b/Sources/MCP/Client/Client+Requests.swift @@ -0,0 +1,471 @@ +import Foundation + +extension Client { + // MARK: - Request Options + + /// Options that can be given per request. + /// + /// Similar to TypeScript SDK's `RequestOptions`, this allows configuring + /// timeout behavior for individual requests, including progress-aware timeouts. + public struct RequestOptions: Sendable { + /// The default request timeout (60 seconds), matching TypeScript SDK. + public static let defaultTimeout: Duration = .seconds(60) + + /// A timeout for this request. + /// + /// If exceeded, the request will be cancelled and an `MCPError.requestTimeout` + /// will be thrown. A `CancelledNotification` will also be sent to the server. + /// + /// If `nil`, no timeout is applied (the request can wait indefinitely). + /// Default is `nil` to match existing behavior. + public var timeout: Duration? + + /// If `true`, receiving a progress notification resets the timeout clock. + /// + /// This is useful for long-running operations that send periodic progress updates. + /// As long as the server keeps sending progress, the request won't time out. + /// + /// When combined with `maxTotalTimeout`, this allows both: + /// - Per-interval timeout that resets on progress + /// - Overall hard limit that prevents infinite waiting + /// + /// Default is `false`. + /// + /// - Note: Only effective when `timeout` is set and the request uses `onProgress`. + public var resetTimeoutOnProgress: Bool + + /// Maximum total time to wait for the request, regardless of progress. + /// + /// When `resetTimeoutOnProgress` is `true`, this provides a hard upper limit + /// on the total wait time. Even if progress notifications keep arriving, + /// the request will be cancelled if this limit is exceeded. + /// + /// If `nil`, there's no maximum total timeout (only the regular `timeout` + /// applies, potentially reset by progress). + /// + /// - Note: Only effective when both `timeout` and `resetTimeoutOnProgress` are set. + public var maxTotalTimeout: Duration? + + /// Creates request options with the specified configuration. + /// + /// - Parameters: + /// - timeout: The timeout duration, or `nil` for no timeout. + /// - resetTimeoutOnProgress: Whether to reset the timeout when progress is received. + /// - maxTotalTimeout: Maximum total time to wait regardless of progress. + public init( + timeout: Duration? = nil, + resetTimeoutOnProgress: Bool = false, + maxTotalTimeout: Duration? = nil + ) { + self.timeout = timeout + self.resetTimeoutOnProgress = resetTimeoutOnProgress + self.maxTotalTimeout = maxTotalTimeout + } + + /// Request options with the default timeout (60 seconds). + public static let withDefaultTimeout = RequestOptions(timeout: defaultTimeout) + + /// Request options with no timeout. + public static let noTimeout = RequestOptions(timeout: nil) + } + + // MARK: - Requests + + /// Send a request and receive its response. + /// + /// This method sends a request without a timeout. For timeout support, + /// use `send(_:options:)` instead. + public func send(_ request: Request) async throws -> M.Result { + try await send(request, options: nil) + } + + /// Send a request and receive its response with options. + /// + /// - Parameters: + /// - request: The request to send. + /// - options: Options for this request, including timeout configuration. + /// - Returns: The response result. + /// - Throws: `MCPError.requestTimeout` if the timeout is exceeded. + public func send( + _ request: Request, + options: RequestOptions? + ) async throws -> M.Result { + guard let connection = connection else { + throw MCPError.internalError("Client connection not initialized") + } + + let requestData = try encoder.encode(request) + + // Create stream for receiving the response + let (stream, continuation) = AsyncThrowingStream.makeStream() + + // Track whether we've timed out (for the onTermination handler) + let requestId = request.id + let timeout = options?.timeout + + // Clean up pending request if caller cancels (e.g., task cancelled or timeout) + // and send CancelledNotification to server per MCP spec + continuation.onTermination = { @Sendable [weak self] termination in + Task { + guard let self else { return } + await self.cleanUpPendingRequest(id: requestId) + + // Per MCP spec: send notifications/cancelled when cancelling a request + // Only send if the stream was cancelled (not finished normally) + if case .cancelled = termination { + let reason = if let timeout { + "Request timed out after \(timeout)" + } else { + "Client cancelled the request" + } + await self.sendCancellationNotification( + requestId: requestId, + reason: reason + ) + } + } + } + + // Add the pending request before attempting to send + addPendingRequest(id: request.id, continuation: continuation) + + // Send the request data + do { + try await connection.send(requestData) + } catch { + // If send fails, remove the pending request and rethrow + if removePendingRequest(id: request.id) != nil { + continuation.finish(throwing: error) + } + throw error + } + + // Wait for response with optional timeout + if let timeout { + // Use withTimeout pattern for cancellation-aware timeout + return try await withThrowingTaskGroup(of: M.Result.self) { group in + // Add the main task that waits for the response + group.addTask { + for try await result in stream { + return result + } + throw MCPError.internalError("No response received") + } + + // Add the timeout task + group.addTask { + try await Task.sleep(for: timeout) + throw MCPError.requestTimeout(timeout: timeout, message: "Request timed out") + } + + // Return whichever completes first + guard let result = try await group.next() else { + throw MCPError.internalError("No response received") + } + + // Cancel the other task + group.cancelAll() + + return result + } + } else { + // No timeout - wait indefinitely for response + for try await result in stream { + return result + } + + // Stream closed without yielding a response + throw MCPError.internalError("No response received") + } + } + + /// Send a request with a progress callback. + /// + /// This method automatically sets up progress tracking by: + /// 1. Generating a unique progress token based on the request ID + /// 2. Injecting the token into the request's `_meta.progressToken` + /// 3. Invoking the callback when progress notifications are received + /// + /// The callback is automatically cleaned up when the request completes. + /// + /// ## Example + /// + /// ```swift + /// let result = try await client.send( + /// CallTool.request(.init(name: "slow_operation", arguments: ["steps": 5])), + /// onProgress: { progress in + /// print("Progress: \(progress.value)/\(progress.total ?? 0) - \(progress.message ?? "")") + /// } + /// ) + /// ``` + /// + /// - Parameters: + /// - request: The request to send + /// - onProgress: A callback invoked when progress notifications are received + /// - Returns: The response result + public func send( + _ request: Request, + onProgress: @escaping ProgressCallback + ) async throws -> M.Result { + try await send(request, options: nil, onProgress: onProgress) + } + + /// Send a request with options and a progress callback. + /// + /// - Parameters: + /// - request: The request to send. + /// - options: Options for this request, including timeout configuration. + /// - onProgress: A callback invoked when progress notifications are received. + /// - Returns: The response result. + /// - Throws: `MCPError.requestTimeout` if the timeout is exceeded. + public func send( + _ request: Request, + options: RequestOptions?, + onProgress: @escaping ProgressCallback + ) async throws -> M.Result { + guard let connection = connection else { + throw MCPError.internalError("Client connection not initialized") + } + + // Generate a progress token from the request ID + let progressToken: ProgressToken = switch request.id { + case .number(let n): .integer(n) + case .string(let s): .string(s) + } + + // Encode the request, inject progressToken into _meta, then re-encode + let requestData = try encoder.encode(request) + var requestDict = try decoder.decode([String: Value].self, from: requestData) + + // Ensure params exists and inject _meta.progressToken + var params = requestDict["params"]?.objectValue ?? [:] + var meta = params["_meta"]?.objectValue ?? [:] + meta["progressToken"] = switch progressToken { + case .string(let s): .string(s) + case .integer(let n): .int(n) + } + params["_meta"] = .object(meta) + requestDict["params"] = .object(params) + + let modifiedRequestData = try encoder.encode(requestDict) + + // Register the progress callback and track the request → token mapping + // (used to detect task-augmented responses and keep progress handlers alive) + progressCallbacks[progressToken] = onProgress + requestProgressTokens[request.id] = progressToken + + // Create timeout controller if resetTimeoutOnProgress is enabled + let timeoutController: TimeoutController? + if let timeout = options?.timeout, options?.resetTimeoutOnProgress == true { + let controller = TimeoutController( + timeout: timeout, + resetOnProgress: true, + maxTotalTimeout: options?.maxTotalTimeout + ) + timeoutControllers[progressToken] = controller + timeoutController = controller + } else { + timeoutController = nil + } + + // Create stream for receiving the response + let (stream, continuation) = AsyncThrowingStream.makeStream() + + let requestId = request.id + let timeout = options?.timeout + continuation.onTermination = { @Sendable [weak self] termination in + Task { + guard let self else { return } + await self.cleanUpPendingRequest(id: requestId) + await self.removeRequestProgressToken(id: requestId) + await self.removeProgressCallback(token: progressToken) + await self.removeTimeoutController(token: progressToken) + + if case .cancelled = termination { + let reason = if let timeout { + "Request timed out after \(timeout)" + } else { + "Client cancelled the request" + } + await self.sendCancellationNotification( + requestId: requestId, + reason: reason + ) + } + } + } + + // Add the pending request before attempting to send + addPendingRequest(id: request.id, continuation: continuation) + + // Send the modified request data + do { + try await connection.send(modifiedRequestData) + } catch { + if removePendingRequest(id: request.id) != nil { + continuation.finish(throwing: error) + } + removeRequestProgressToken(id: request.id) + removeProgressCallback(token: progressToken) + removeTimeoutController(token: progressToken) + throw error + } + + // Wait for response with optional timeout + if let timeout { + // Use TimeoutController if resetTimeoutOnProgress is enabled + if let controller = timeoutController { + return try await withThrowingTaskGroup(of: M.Result.self) { group in + group.addTask { + for try await result in stream { + return result + } + throw MCPError.internalError("No response received") + } + + group.addTask { + try await controller.waitForTimeout() + throw MCPError.internalError("Unreachable - timeout should throw") + } + + guard let result = try await group.next() else { + throw MCPError.internalError("No response received") + } + + group.cancelAll() + await controller.cancel() + removeProgressCallback(token: progressToken) + removeTimeoutController(token: progressToken) + return result + } + } else { + // Simple timeout without progress-aware reset + return try await withThrowingTaskGroup(of: M.Result.self) { group in + group.addTask { + for try await result in stream { + return result + } + throw MCPError.internalError("No response received") + } + + group.addTask { + try await Task.sleep(for: timeout) + throw MCPError.requestTimeout(timeout: timeout, message: "Request timed out") + } + + guard let result = try await group.next() else { + throw MCPError.internalError("No response received") + } + + group.cancelAll() + removeProgressCallback(token: progressToken) + return result + } + } + } else { + for try await result in stream { + removeProgressCallback(token: progressToken) + removeTimeoutController(token: progressToken) + return result + } + + removeProgressCallback(token: progressToken) + removeTimeoutController(token: progressToken) + throw MCPError.internalError("No response received") + } + } + + /// Remove a progress callback for the given token. + /// + /// If the token is being tracked for a task (task-augmented response), the callback + /// is NOT removed. This keeps progress handlers alive until the task completes. + private func removeProgressCallback(token: ProgressToken) { + // Check if this token is being tracked for a task + // If so, don't remove the callback - it needs to stay alive until task completes + let isTaskProgressToken = taskProgressTokens.values.contains(token) + if isTaskProgressToken { + return + } + progressCallbacks.removeValue(forKey: token) + } + + /// Remove a timeout controller for the given token. + /// + /// If the token is being tracked for a task (task-augmented response), the controller + /// is NOT removed. This keeps timeout tracking alive until the task completes. + private func removeTimeoutController(token: ProgressToken) { + // Check if this token is being tracked for a task + // If so, don't remove the controller - it needs to stay alive until task completes + let isTaskProgressToken = taskProgressTokens.values.contains(token) + if isTaskProgressToken { + return + } + timeoutControllers.removeValue(forKey: token) + } + + /// Remove the request → progress token mapping for the given request ID. + private func removeRequestProgressToken(id: RequestId) { + requestProgressTokens.removeValue(forKey: id) + } + + func addPendingRequest( + id: RequestId, + continuation: AsyncThrowingStream.Continuation + ) { + pendingRequests[id] = AnyPendingRequest(continuation: continuation) + } + + func removePendingRequest(id: RequestId) -> AnyPendingRequest? { + return pendingRequests.removeValue(forKey: id) + } + + /// Removes a pending request without returning it. + /// Used by onTermination handlers when the request has been cancelled. + func cleanUpPendingRequest(id: RequestId) { + pendingRequests.removeValue(forKey: id) + } + + /// Send a CancelledNotification to the server for a cancelled request. + /// + /// Per MCP spec: "When a party wants to cancel an in-progress request, it sends + /// a `notifications/cancelled` notification containing the ID of the request to cancel." + /// + /// This is called when a client Task waiting for a response is cancelled. + /// The notification is sent on a best-effort basis - failures are logged but not thrown. + func sendCancellationNotification(requestId: RequestId, reason: String?) async { + guard let connection = connection else { + await logger?.debug( + "Cannot send cancellation notification - connection is nil", + metadata: ["requestId": "\(requestId)"] + ) + return + } + + let notification = CancelledNotification.message(.init( + requestId: requestId, + reason: reason + )) + + do { + let notificationData = try encoder.encode(notification) + try await connection.send(notificationData) + await logger?.debug( + "Sent cancellation notification", + metadata: [ + "requestId": "\(requestId)", + "reason": "\(reason ?? "none")", + ] + ) + } catch { + // Log but don't throw - cancellation notification is best-effort + // per MCP spec's fire-and-forget nature of notifications + await logger?.debug( + "Failed to send cancellation notification", + metadata: [ + "requestId": "\(requestId)", + "error": "\(error)", + ] + ) + } + } +} diff --git a/Sources/MCP/Client/Client.swift b/Sources/MCP/Client/Client.swift index 9f11a3af..c65f507a 100644 --- a/Sources/MCP/Client/Client.swift +++ b/Sources/MCP/Client/Client.swift @@ -273,16 +273,16 @@ public actor Client { } /// The connection to the server - private var connection: (any Transport)? + var connection: (any Transport)? /// The logger for the client - private var logger: Logger? { + var logger: Logger? { get async { await connection?.logger } } /// The client information - private let clientInfo: Client.Info + let clientInfo: Client.Info /// The client name public nonisolated var name: String { clientInfo.name } /// The client version @@ -307,32 +307,32 @@ public actor Client { } /// The server capabilities - private var serverCapabilities: Server.Capabilities? + var serverCapabilities: Server.Capabilities? /// The server version - private var serverVersion: String? + var serverVersion: String? /// The server instructions - private var instructions: String? + var instructions: String? /// A dictionary of type-erased notification handlers, keyed by method name - private var notificationHandlers: [String: [NotificationHandlerBox]] = [:] + var notificationHandlers: [String: [NotificationHandlerBox]] = [:] /// A dictionary of type-erased request handlers for server→client requests, keyed by method name - private var requestHandlers: [String: ClientRequestHandlerBox] = [:] + var requestHandlers: [String: ClientRequestHandlerBox] = [:] /// Task-augmented sampling handler (called when request has `task` field) - private var taskAugmentedSamplingHandler: ExperimentalClientTaskHandlers.TaskAugmentedSamplingHandler? + var taskAugmentedSamplingHandler: ExperimentalClientTaskHandlers.TaskAugmentedSamplingHandler? /// Task-augmented elicitation handler (called when request has `task` field) - private var taskAugmentedElicitationHandler: ExperimentalClientTaskHandlers.TaskAugmentedElicitationHandler? + var taskAugmentedElicitationHandler: ExperimentalClientTaskHandlers.TaskAugmentedElicitationHandler? /// The task for the message handling loop - private var task: Task? + var task: Task? /// In-flight server request handler Tasks, tracked by request ID. /// Used for protocol-level cancellation when CancelledNotification is received. - private var inFlightServerRequestTasks: [RequestId: Task] = [:] + var inFlightServerRequestTasks: [RequestId: Task] = [:] /// An error indicating a type mismatch when decoding a pending request - private struct TypeMismatchError: Swift.Error {} + struct TypeMismatchError: Swift.Error {} /// A type-erased pending request using AsyncThrowingStream for cancellation-aware waiting. - private struct AnyPendingRequest { + struct AnyPendingRequest { private let _yield: (Result) -> Void private let _finish: () -> Void @@ -377,31 +377,31 @@ public actor Client { } /// A dictionary of type-erased pending requests, keyed by request ID - private var pendingRequests: [RequestId: AnyPendingRequest] = [:] + var pendingRequests: [RequestId: AnyPendingRequest] = [:] /// Progress callbacks for requests, keyed by progress token. /// Used to invoke callbacks when progress notifications are received. - private var progressCallbacks: [ProgressToken: ProgressCallback] = [:] + var progressCallbacks: [ProgressToken: ProgressCallback] = [:] /// Timeout controllers for requests with progress-aware timeouts. /// Used to reset timeouts when progress notifications are received. - private var timeoutControllers: [ProgressToken: TimeoutController] = [:] + var timeoutControllers: [ProgressToken: TimeoutController] = [:] /// Mapping from request ID to progress token. /// Used to detect task-augmented responses and keep progress handlers alive. - private var requestProgressTokens: [RequestId: ProgressToken] = [:] + var requestProgressTokens: [RequestId: ProgressToken] = [:] /// Mapping from task ID to progress token. /// Keeps progress handlers alive for task-augmented requests until the task completes. /// Per MCP spec 2025-11-25: "For task-augmented requests, the progressToken provided /// in the original request MUST continue to be used for progress notifications /// throughout the task's lifetime, even after the CreateTaskResult has been returned." - private var taskProgressTokens: [String: ProgressToken] = [:] + var taskProgressTokens: [String: ProgressToken] = [:] // Add reusable JSON encoder/decoder - private let encoder = JSONEncoder() - private let decoder = JSONDecoder() + let encoder = JSONEncoder() + let decoder = JSONDecoder() /// Controls timeout behavior for a single request, supporting reset on progress. /// /// This actor manages the timeout state for requests that use `resetTimeoutOnProgress`. /// When progress is received, calling `signalProgress()` resets the timeout clock. - private actor TimeoutController { + actor TimeoutController { /// The per-interval timeout duration. let timeout: Duration /// Whether to reset timeout when progress is received. @@ -564,7 +564,7 @@ public actor Client { // When the receive loop exits unexpectedly (transport closed without // disconnect() being called), clean up pending requests. Task { - await self.cleanupPendingRequestsOnUnexpectedDisconnect() + await self.cleanUpPendingRequestsOnUnexpectedDisconnect() } } @@ -691,7 +691,7 @@ public actor Client { /// This is called from the receive loop's defer block when the transport closes /// without `disconnect()` being called (e.g., server process exits). We only /// clean up requests that haven't already been handled by `disconnect()`. - private func cleanupPendingRequestsOnUnexpectedDisconnect() async { + func cleanUpPendingRequestsOnUnexpectedDisconnect() async { guard !pendingRequests.isEmpty else { return } await logger?.debug( @@ -707,12 +707,12 @@ public actor Client { // MARK: - In-Flight Server Request Tracking (Protocol-Level Cancellation) /// Track an in-flight server request handler Task. - private func trackInFlightServerRequest(_ requestId: RequestId, task: Task) { + func trackInFlightServerRequest(_ requestId: RequestId, task: Task) { inFlightServerRequestTasks[requestId] = task } /// Remove an in-flight server request handler Task. - private func removeInFlightServerRequest(_ requestId: RequestId) { + func removeInFlightServerRequest(_ requestId: RequestId) { inFlightServerRequestTasks.removeValue(forKey: requestId) } @@ -720,7 +720,7 @@ public actor Client { /// /// Called when a CancelledNotification is received for a specific requestId. /// Per MCP spec, if the request is unknown or already completed, this is a no-op. - private func cancelInFlightServerRequest(_ requestId: RequestId, reason: String?) async { + func cancelInFlightServerRequest(_ requestId: RequestId, reason: String?) async { if let task = inFlightServerRequestTasks[requestId] { task.cancel() await logger?.debug( @@ -734,1751 +734,60 @@ public actor Client { // Per spec: MAY ignore if request is unknown - no error needed } - // MARK: - Registration - - /// Register a handler for a notification - @discardableResult - public func onNotification( - _ type: N.Type, - handler: @escaping @Sendable (Message) async throws -> Void - ) async -> Self { - notificationHandlers[N.name, default: []].append(TypedNotificationHandler(handler)) - return self - } - - /// Send a notification to the server - public func notify(_ notification: Message) async throws { - guard let connection = connection else { - throw MCPError.internalError("Client connection not initialized") - } - - let notificationData = try encoder.encode(notification) - try await connection.send(notificationData) - } - - /// Send a progress notification to the server. - /// - /// This is a convenience method for sending progress notifications from the client - /// to the server. This enables bidirectional progress reporting where clients can - /// inform servers about their own progress (e.g., during client-side processing). - /// - /// ## Example - /// - /// ```swift - /// // Client reports its own progress to the server - /// try await client.sendProgressNotification( - /// token: .string("client-task-123"), - /// progress: 50.0, - /// total: 100.0, - /// message: "Processing client-side data..." - /// ) - /// ``` - /// - /// - Parameters: - /// - token: The progress token to associate with this notification - /// - progress: The current progress value (should increase monotonically) - /// - total: The total progress value, if known - /// - message: An optional human-readable message describing current progress - public func sendProgressNotification( - token: ProgressToken, - progress: Double, - total: Double? = nil, - message: String? = nil - ) async throws { - try await notify(ProgressNotification.message(.init( - progressToken: token, - progress: progress, - total: total, - message: message - ))) - } - - /// Send a notification that the list of available roots has changed. - /// - /// Servers that receive this notification should request an updated - /// list of roots via the roots/list request. - /// - /// - Throws: `MCPError.invalidRequest` if the client has not declared - /// the `roots.listChanged` capability. - public func sendRootsChanged() async throws { - guard capabilities.roots?.listChanged == true else { - throw MCPError.invalidRequest( - "Client does not support roots.listChanged capability") - } - try await notify(RootsListChangedNotification.message(.init())) - } - - /// Register a handler for server→client requests. - /// - /// This enables bidirectional communication where the server can send requests - /// to the client (e.g., sampling, roots, elicitation). - /// - /// - Parameters: - /// - type: The method type to handle - /// - handler: The handler function that receives parameters and returns a result - /// - Returns: Self for chaining - @discardableResult - public func withRequestHandler( - _ type: M.Type, - handler: @escaping @Sendable (M.Parameters) async throws -> M.Result - ) -> Self { - requestHandlers[M.name] = TypedClientRequestHandler(handler) - return self - } - - /// Register a handler for `roots/list` requests from the server. - /// - /// When the server requests the list of roots, this handler will be called - /// to provide the available filesystem directories. - /// - /// - Important: The client must have declared `roots` capability during initialization. - /// - /// - Parameter handler: A closure that returns the list of available roots. - /// - Returns: Self for chaining. - /// - Precondition: `capabilities.roots` must be non-nil. - @discardableResult - public func withRootsHandler( - _ handler: @escaping @Sendable () async throws -> [Root] - ) -> Self { - precondition( - capabilities.roots != nil, - "Cannot register roots handler: Client does not have roots capability" - ) - return withRequestHandler(ListRoots.self) { _ in - ListRoots.Result(roots: try await handler()) - } - } - - /// Register a handler for `sampling/createMessage` requests from the server. - /// - /// When the server requests a sampling completion, this handler will be called - /// to generate the LLM response. - /// - /// The handler receives parameters that may or may not include tools. Check `params.hasTools` - /// to determine if tool use is enabled for this request. - /// - /// - Important: The client must have declared `sampling` capability during initialization. - /// - /// ## Example - /// - /// ```swift - /// client.withSamplingHandler { params in - /// // Call your LLM with the messages - /// let response = try await llm.complete( - /// messages: params.messages, - /// tools: params.tools, // May be nil - /// maxTokens: params.maxTokens - /// ) - /// - /// return ClientSamplingRequest.Result( - /// model: "gpt-4", - /// stopReason: .endTurn, - /// role: .assistant, - /// content: .text(response.text) - /// ) - /// } - /// ``` - /// - /// - Parameter handler: A closure that receives sampling parameters and returns the result. - /// - Returns: Self for chaining. - /// - Precondition: `capabilities.sampling` must be non-nil. - @discardableResult - public func withSamplingHandler( - _ handler: @escaping @Sendable (ClientSamplingRequest.Parameters) async throws -> ClientSamplingRequest.Result - ) -> Self { - precondition( - capabilities.sampling != nil, - "Cannot register sampling handler: Client does not have sampling capability" - ) - return withRequestHandler(ClientSamplingRequest.self, handler: handler) - } - - /// Register a handler for `elicitation/create` requests from the server. - /// - /// When the server requests user input via elicitation, this handler will be called - /// to collect the input and return the result. - /// - /// - Important: The client must have declared `elicitation` capability during initialization. - /// - /// - Parameter handler: A closure that receives elicitation parameters and returns the result. - /// - Returns: Self for chaining. - /// - Precondition: `capabilities.elicitation` must be non-nil. - @discardableResult - public func withElicitationHandler( - _ handler: @escaping @Sendable (Elicit.Parameters) async throws -> Elicit.Result - ) -> Self { - precondition( - capabilities.elicitation != nil, - "Cannot register elicitation handler: Client does not have elicitation capability" - ) - return withRequestHandler(Elicit.self, handler: handler) - } - - /// Internal method to set a request handler box directly. - /// - /// This is used by task-augmented handlers that need to return different result types - /// based on whether the request has a `task` field. - /// - /// - Important: This is an internal API that may change without notice. - internal func _setRequestHandler(method: String, handler: ClientRequestHandlerBox) { - requestHandlers[method] = handler - } - - /// Internal method to get an existing request handler box. - /// - /// This is used to retrieve the existing handler before wrapping it with - /// a task-aware handler that preserves the normal handler as a fallback. - /// - /// - Important: This is an internal API that may change without notice. - internal func _getRequestHandler(method: String) -> ClientRequestHandlerBox? { - requestHandlers[method] - } - - /// Internal method to set the task-augmented sampling handler. - /// - /// This handler is called when the server sends a `sampling/createMessage` request - /// with a `task` field. The handler should return `CreateTaskResult` instead of - /// the normal sampling result. - /// - /// - Important: This is an internal API that may change without notice. - internal func _setTaskAugmentedSamplingHandler( - _ handler: @escaping ExperimentalClientTaskHandlers.TaskAugmentedSamplingHandler - ) { - taskAugmentedSamplingHandler = handler - } - - /// Internal method to set the task-augmented elicitation handler. - /// - /// This handler is called when the server sends an `elicitation/create` request - /// with a `task` field. The handler should return `CreateTaskResult` instead of - /// the normal elicitation result. - /// - /// - Important: This is an internal API that may change without notice. - internal func _setTaskAugmentedElicitationHandler( - _ handler: @escaping ExperimentalClientTaskHandlers.TaskAugmentedElicitationHandler - ) { - taskAugmentedElicitationHandler = handler - } - - // MARK: - Request Options - - /// Options that can be given per request. - /// - /// Similar to TypeScript SDK's `RequestOptions`, this allows configuring - /// timeout behavior for individual requests, including progress-aware timeouts. - public struct RequestOptions: Sendable { - /// The default request timeout (60 seconds), matching TypeScript SDK. - public static let defaultTimeout: Duration = .seconds(60) - - /// A timeout for this request. - /// - /// If exceeded, the request will be cancelled and an `MCPError.requestTimeout` - /// will be thrown. A `CancelledNotification` will also be sent to the server. - /// - /// If `nil`, no timeout is applied (the request can wait indefinitely). - /// Default is `nil` to match existing behavior. - public var timeout: Duration? - - /// If `true`, receiving a progress notification resets the timeout clock. - /// - /// This is useful for long-running operations that send periodic progress updates. - /// As long as the server keeps sending progress, the request won't time out. - /// - /// When combined with `maxTotalTimeout`, this allows both: - /// - Per-interval timeout that resets on progress - /// - Overall hard limit that prevents infinite waiting - /// - /// Default is `false`. - /// - /// - Note: Only effective when `timeout` is set and the request uses `onProgress`. - public var resetTimeoutOnProgress: Bool - - /// Maximum total time to wait for the request, regardless of progress. - /// - /// When `resetTimeoutOnProgress` is `true`, this provides a hard upper limit - /// on the total wait time. Even if progress notifications keep arriving, - /// the request will be cancelled if this limit is exceeded. - /// - /// If `nil`, there's no maximum total timeout (only the regular `timeout` - /// applies, potentially reset by progress). - /// - /// - Note: Only effective when both `timeout` and `resetTimeoutOnProgress` are set. - public var maxTotalTimeout: Duration? - - /// Creates request options with the specified configuration. - /// - /// - Parameters: - /// - timeout: The timeout duration, or `nil` for no timeout. - /// - resetTimeoutOnProgress: Whether to reset the timeout when progress is received. - /// - maxTotalTimeout: Maximum total time to wait regardless of progress. - public init( - timeout: Duration? = nil, - resetTimeoutOnProgress: Bool = false, - maxTotalTimeout: Duration? = nil - ) { - self.timeout = timeout - self.resetTimeoutOnProgress = resetTimeoutOnProgress - self.maxTotalTimeout = maxTotalTimeout - } - - /// Request options with the default timeout (60 seconds). - public static let withDefaultTimeout = RequestOptions(timeout: defaultTimeout) - - /// Request options with no timeout. - public static let noTimeout = RequestOptions(timeout: nil) - } - - // MARK: - Requests - - /// Send a request and receive its response. - /// - /// This method sends a request without a timeout. For timeout support, - /// use `send(_:options:)` instead. - public func send(_ request: Request) async throws -> M.Result { - try await send(request, options: nil) - } - - /// Send a request and receive its response with options. - /// - /// - Parameters: - /// - request: The request to send. - /// - options: Options for this request, including timeout configuration. - /// - Returns: The response result. - /// - Throws: `MCPError.requestTimeout` if the timeout is exceeded. - public func send( - _ request: Request, - options: RequestOptions? - ) async throws -> M.Result { - guard let connection = connection else { - throw MCPError.internalError("Client connection not initialized") - } - - let requestData = try encoder.encode(request) - - // Create stream for receiving the response - let (stream, continuation) = AsyncThrowingStream.makeStream() - - // Track whether we've timed out (for the onTermination handler) - let requestId = request.id - let timeout = options?.timeout - - // Clean up pending request if caller cancels (e.g., task cancelled or timeout) - // and send CancelledNotification to server per MCP spec - continuation.onTermination = { @Sendable [weak self] termination in - Task { - guard let self else { return } - await self.cleanupPendingRequest(id: requestId) - - // Per MCP spec: send notifications/cancelled when cancelling a request - // Only send if the stream was cancelled (not finished normally) - if case .cancelled = termination { - let reason = if let timeout { - "Request timed out after \(timeout)" - } else { - "Client cancelled the request" - } - await self.sendCancellationNotification( - requestId: requestId, - reason: reason - ) - } - } - } - - // Add the pending request before attempting to send - addPendingRequest(id: request.id, continuation: continuation) - - // Send the request data - do { - try await connection.send(requestData) - } catch { - // If send fails, remove the pending request and rethrow - if removePendingRequest(id: request.id) != nil { - continuation.finish(throwing: error) - } - throw error - } - - // Wait for response with optional timeout - if let timeout { - // Use withTimeout pattern for cancellation-aware timeout - return try await withThrowingTaskGroup(of: M.Result.self) { group in - // Add the main task that waits for the response - group.addTask { - for try await result in stream { - return result - } - throw MCPError.internalError("No response received") - } - - // Add the timeout task - group.addTask { - try await Task.sleep(for: timeout) - throw MCPError.requestTimeout(timeout: timeout, message: "Request timed out") - } - - // Return whichever completes first - guard let result = try await group.next() else { - throw MCPError.internalError("No response received") - } - - // Cancel the other task - group.cancelAll() - - return result - } - } else { - // No timeout - wait indefinitely for response - for try await result in stream { - return result - } - - // Stream closed without yielding a response - throw MCPError.internalError("No response received") - } - } - - /// Send a request with a progress callback. - /// - /// This method automatically sets up progress tracking by: - /// 1. Generating a unique progress token based on the request ID - /// 2. Injecting the token into the request's `_meta.progressToken` - /// 3. Invoking the callback when progress notifications are received - /// - /// The callback is automatically cleaned up when the request completes. - /// - /// ## Example - /// - /// ```swift - /// let result = try await client.send( - /// CallTool.request(.init(name: "slow_operation", arguments: ["steps": 5])), - /// onProgress: { progress in - /// print("Progress: \(progress.value)/\(progress.total ?? 0) - \(progress.message ?? "")") - /// } - /// ) - /// ``` - /// - /// - Parameters: - /// - request: The request to send - /// - onProgress: A callback invoked when progress notifications are received - /// - Returns: The response result - public func send( - _ request: Request, - onProgress: @escaping ProgressCallback - ) async throws -> M.Result { - try await send(request, options: nil, onProgress: onProgress) - } - - /// Send a request with options and a progress callback. - /// - /// - Parameters: - /// - request: The request to send. - /// - options: Options for this request, including timeout configuration. - /// - onProgress: A callback invoked when progress notifications are received. - /// - Returns: The response result. - /// - Throws: `MCPError.requestTimeout` if the timeout is exceeded. - public func send( - _ request: Request, - options: RequestOptions?, - onProgress: @escaping ProgressCallback - ) async throws -> M.Result { - guard let connection = connection else { - throw MCPError.internalError("Client connection not initialized") - } - - // Generate a progress token from the request ID - let progressToken: ProgressToken = switch request.id { - case .number(let n): .integer(n) - case .string(let s): .string(s) - } - - // Encode the request, inject progressToken into _meta, then re-encode - let requestData = try encoder.encode(request) - var requestDict = try decoder.decode([String: Value].self, from: requestData) - - // Ensure params exists and inject _meta.progressToken - var params = requestDict["params"]?.objectValue ?? [:] - var meta = params["_meta"]?.objectValue ?? [:] - meta["progressToken"] = switch progressToken { - case .string(let s): .string(s) - case .integer(let n): .int(n) - } - params["_meta"] = .object(meta) - requestDict["params"] = .object(params) - - let modifiedRequestData = try encoder.encode(requestDict) - - // Register the progress callback and track the request → token mapping - // (used to detect task-augmented responses and keep progress handlers alive) - progressCallbacks[progressToken] = onProgress - requestProgressTokens[request.id] = progressToken - - // Create timeout controller if resetTimeoutOnProgress is enabled - let timeoutController: TimeoutController? - if let timeout = options?.timeout, options?.resetTimeoutOnProgress == true { - let controller = TimeoutController( - timeout: timeout, - resetOnProgress: true, - maxTotalTimeout: options?.maxTotalTimeout - ) - timeoutControllers[progressToken] = controller - timeoutController = controller - } else { - timeoutController = nil - } - - // Create stream for receiving the response - let (stream, continuation) = AsyncThrowingStream.makeStream() - - let requestId = request.id - let timeout = options?.timeout - continuation.onTermination = { @Sendable [weak self] termination in - Task { - guard let self else { return } - await self.cleanupPendingRequest(id: requestId) - await self.removeRequestProgressToken(id: requestId) - await self.removeProgressCallback(token: progressToken) - await self.removeTimeoutController(token: progressToken) - - if case .cancelled = termination { - let reason = if let timeout { - "Request timed out after \(timeout)" - } else { - "Client cancelled the request" - } - await self.sendCancellationNotification( - requestId: requestId, - reason: reason - ) - } - } - } - - // Add the pending request before attempting to send - addPendingRequest(id: request.id, continuation: continuation) - - // Send the modified request data - do { - try await connection.send(modifiedRequestData) - } catch { - if removePendingRequest(id: request.id) != nil { - continuation.finish(throwing: error) - } - removeRequestProgressToken(id: request.id) - removeProgressCallback(token: progressToken) - removeTimeoutController(token: progressToken) - throw error - } - - // Wait for response with optional timeout - if let timeout { - // Use TimeoutController if resetTimeoutOnProgress is enabled - if let controller = timeoutController { - return try await withThrowingTaskGroup(of: M.Result.self) { group in - group.addTask { - for try await result in stream { - return result - } - throw MCPError.internalError("No response received") - } - - group.addTask { - try await controller.waitForTimeout() - throw MCPError.internalError("Unreachable - timeout should throw") - } - - guard let result = try await group.next() else { - throw MCPError.internalError("No response received") - } - - group.cancelAll() - await controller.cancel() - removeProgressCallback(token: progressToken) - removeTimeoutController(token: progressToken) - return result - } - } else { - // Simple timeout without progress-aware reset - return try await withThrowingTaskGroup(of: M.Result.self) { group in - group.addTask { - for try await result in stream { - return result - } - throw MCPError.internalError("No response received") - } - - group.addTask { - try await Task.sleep(for: timeout) - throw MCPError.requestTimeout(timeout: timeout, message: "Request timed out") - } - - guard let result = try await group.next() else { - throw MCPError.internalError("No response received") - } - - group.cancelAll() - removeProgressCallback(token: progressToken) - return result - } - } - } else { - for try await result in stream { - removeProgressCallback(token: progressToken) - removeTimeoutController(token: progressToken) - return result - } - - removeProgressCallback(token: progressToken) - removeTimeoutController(token: progressToken) - throw MCPError.internalError("No response received") - } - } - - /// Remove a progress callback for the given token. - /// - /// If the token is being tracked for a task (task-augmented response), the callback - /// is NOT removed. This keeps progress handlers alive until the task completes. - private func removeProgressCallback(token: ProgressToken) { - // Check if this token is being tracked for a task - // If so, don't remove the callback - it needs to stay alive until task completes - let isTaskProgressToken = taskProgressTokens.values.contains(token) - if isTaskProgressToken { - return - } - progressCallbacks.removeValue(forKey: token) - } - - /// Remove a timeout controller for the given token. - /// - /// If the token is being tracked for a task (task-augmented response), the controller - /// is NOT removed. This keeps timeout tracking alive until the task completes. - private func removeTimeoutController(token: ProgressToken) { - // Check if this token is being tracked for a task - // If so, don't remove the controller - it needs to stay alive until task completes - let isTaskProgressToken = taskProgressTokens.values.contains(token) - if isTaskProgressToken { - return - } - timeoutControllers.removeValue(forKey: token) - } - - /// Remove the request → progress token mapping for the given request ID. - private func removeRequestProgressToken(id: RequestId) { - requestProgressTokens.removeValue(forKey: id) - } - - private func addPendingRequest( - id: RequestId, - continuation: AsyncThrowingStream.Continuation - ) { - pendingRequests[id] = AnyPendingRequest(continuation: continuation) - } - - private func removePendingRequest(id: RequestId) -> AnyPendingRequest? { - return pendingRequests.removeValue(forKey: id) - } - - /// Removes a pending request without returning it. - /// Used by onTermination handlers when the request has been cancelled. - private func cleanupPendingRequest(id: RequestId) { - pendingRequests.removeValue(forKey: id) - } - - /// Send a CancelledNotification to the server for a cancelled request. - /// - /// Per MCP spec: "When a party wants to cancel an in-progress request, it sends - /// a `notifications/cancelled` notification containing the ID of the request to cancel." - /// - /// This is called when a client Task waiting for a response is cancelled. - /// The notification is sent on a best-effort basis - failures are logged but not thrown. - private func sendCancellationNotification(requestId: RequestId, reason: String?) async { - guard let connection = connection else { - await logger?.debug( - "Cannot send cancellation notification - connection is nil", - metadata: ["requestId": "\(requestId)"] - ) - return - } - - let notification = CancelledNotification.message(.init( - requestId: requestId, - reason: reason - )) - - do { - let notificationData = try encoder.encode(notification) - try await connection.send(notificationData) - await logger?.debug( - "Sent cancellation notification", - metadata: [ - "requestId": "\(requestId)", - "reason": "\(reason ?? "none")", - ] - ) - } catch { - // Log but don't throw - cancellation notification is best-effort - // per MCP spec's fire-and-forget nature of notifications - await logger?.debug( - "Failed to send cancellation notification", - metadata: [ - "requestId": "\(requestId)", - "error": "\(error)", - ] - ) - } - } - - // MARK: - Batching - - /// A batch of requests. - /// - /// Objects of this type are passed as an argument to the closure - /// of the ``Client/withBatch(_:)`` method. - public actor Batch { - unowned let client: Client - var requests: [AnyRequest] = [] - - init(client: Client) { - self.client = client - } - - /// Adds a request to the batch and prepares its expected response task. - /// The actual sending happens when the `withBatch` scope completes. - /// - Returns: A `Task` that will eventually produce the result or throw an error. - public func addRequest(_ request: Request) async throws -> Task< - M.Result, Swift.Error - > { - requests.append(try AnyRequest(request)) - - // Create stream for receiving the response - let (stream, continuation) = AsyncThrowingStream.makeStream() - - // Clean up pending request if caller cancels (e.g., task cancelled) - // and send CancelledNotification to server per MCP spec - let requestId = request.id - continuation.onTermination = { @Sendable [weak client] termination in - Task { - guard let client else { return } - await client.cleanupPendingRequest(id: requestId) - - // Per MCP spec: send notifications/cancelled when cancelling a request - // Only send if the stream was cancelled (not finished normally) - if case .cancelled = termination { - await client.sendCancellationNotification( - requestId: requestId, - reason: "Client cancelled the batch request" - ) - } - } - } - - // Register the pending request - await client.addPendingRequest(id: request.id, continuation: continuation) - - // Return a Task that waits for the response via the stream - return Task { - for try await result in stream { - return result - } - throw MCPError.internalError("No response received") - } - } - } - - /// Executes multiple requests in a single batch. - /// - /// This method allows you to group multiple MCP requests together, - /// which are then sent to the server as a single JSON array. - /// The server processes these requests and sends back a corresponding - /// JSON array of responses. - /// - /// Within the `body` closure, use the provided `Batch` actor to add - /// requests using `batch.addRequest(_:)`. Each call to `addRequest` - /// returns a `Task` handle representing the asynchronous operation - /// for that specific request's result. - /// - /// It's recommended to collect these `Task` handles into an array - /// within the `body` closure`. After the `withBatch` method returns - /// (meaning the batch request has been sent), you can then process - /// the results by awaiting each `Task` in the collected array. - /// - /// Example 1: Batching multiple tool calls and collecting typed tasks: - /// ```swift - /// // Array to hold the task handles for each tool call - /// var toolTasks: [Task] = [] - /// try await client.withBatch { batch in - /// for i in 0..<10 { - /// toolTasks.append( - /// try await batch.addRequest( - /// CallTool.request(.init(name: "square", arguments: ["n": i])) - /// ) - /// ) - /// } - /// } - /// - /// // Process results after the batch is sent - /// print("Processing \(toolTasks.count) tool results...") - /// for (index, task) in toolTasks.enumerated() { - /// do { - /// let result = try await task.value - /// print("\(index): \(result.content)") - /// } catch { - /// print("\(index) failed: \(error)") - /// } - /// } - /// ``` - /// - /// Example 2: Batching different request types and awaiting individual tasks: - /// ```swift - /// // Declare optional task variables beforehand - /// var pingTask: Task? - /// var promptTask: Task? - /// - /// try await client.withBatch { batch in - /// // Assign the tasks within the batch closure - /// pingTask = try await batch.addRequest(Ping.request()) - /// promptTask = try await batch.addRequest(GetPrompt.request(.init(name: "greeting"))) - /// } - /// - /// // Await the results after the batch is sent - /// do { - /// if let pingTask = pingTask { - /// try await pingTask.value // Await ping result (throws if ping failed) - /// print("Ping successful") - /// } - /// if let promptTask = promptTask { - /// let promptResult = try await promptTask.value // Await prompt result - /// print("Prompt description: \(promptResult.description ?? "None")") - /// } - /// } catch { - /// print("Error processing batch results: \(error)") - /// } - /// ``` - /// - /// - Parameter body: An asynchronous closure that takes a `Batch` object as input. - /// Use this object to add requests to the batch. - /// - Throws: `MCPError.internalError` if the client is not connected. - /// Can also rethrow errors from the `body` closure or from sending the batch request. - public func withBatch(body: @escaping (Batch) async throws -> Void) async throws { - guard let connection = connection else { - throw MCPError.internalError("Client connection not initialized") - } - - // Create Batch actor, passing self (Client) - let batch = Batch(client: self) - - // Populate the batch actor by calling the user's closure. - try await body(batch) - - // Get the collected requests from the batch actor - let requests = await batch.requests - - // Check if there are any requests to send - guard !requests.isEmpty else { - await logger?.debug("Batch requested but no requests were added.") - return // Nothing to send - } - - await logger?.debug( - "Sending batch request", metadata: ["count": "\(requests.count)"]) - - // Encode the array of AnyMethod requests into a single JSON payload - let data = try encoder.encode(requests) - try await connection.send(data) - - // Responses will be handled asynchronously by the message loop and handleBatchResponse/handleResponse. - } - - // MARK: - Lifecycle - - /// Initialize the connection with the server. - /// - /// - Important: This method is deprecated. Initialization now happens automatically - /// when calling `connect(transport:)`. You should use that method instead. - /// - /// - Returns: The server's initialization response containing capabilities and server info - @available( - *, deprecated, - message: - "Initialization now happens automatically during connect. Use connect(transport:) instead." - ) - public func initialize() async throws -> Initialize.Result { - return try await _initialize() - } - - /// Internal initialization implementation - private func _initialize() async throws -> Initialize.Result { - let request = Initialize.request( - .init( - protocolVersion: Version.latest, - capabilities: capabilities, - clientInfo: clientInfo - )) - - let result = try await send(request) - - // Per MCP spec: "If the client does not support the version in the - // server's response, it SHOULD disconnect." - guard Version.supported.contains(result.protocolVersion) else { - await disconnect() - throw MCPError.invalidRequest( - "Server responded with unsupported protocol version: \(result.protocolVersion). " + - "Supported versions: \(Version.supported.sorted().joined(separator: ", "))" - ) - } - - self.serverCapabilities = result.capabilities - self.serverVersion = result.protocolVersion - self.instructions = result.instructions - - // HTTP transports must set the protocol version in headers after initialization - if let httpTransport = connection as? HTTPClientTransport { - await httpTransport.setProtocolVersion(result.protocolVersion) - } - - try await notify(InitializedNotification.message()) - - return result - } - - public func ping() async throws { - let request = Ping.request() - _ = try await send(request) - } - - // MARK: - Prompts - - public func getPrompt(name: String, arguments: [String: String]? = nil) async throws - -> (description: String?, messages: [Prompt.Message]) - { - try validateServerCapability(\.prompts, "Prompts") - let request = GetPrompt.request(.init(name: name, arguments: arguments)) - let result = try await send(request) - return (description: result.description, messages: result.messages) - } - - public func listPrompts(cursor: String? = nil) async throws - -> (prompts: [Prompt], nextCursor: String?) - { - try validateServerCapability(\.prompts, "Prompts") - let request: Request - if let cursor = cursor { - request = ListPrompts.request(.init(cursor: cursor)) - } else { - request = ListPrompts.request(.init()) - } - let result = try await send(request) - return (prompts: result.prompts, nextCursor: result.nextCursor) - } - - // MARK: - Resources - - public func readResource(uri: String) async throws -> [Resource.Content] { - try validateServerCapability(\.resources, "Resources") - let request = ReadResource.request(.init(uri: uri)) - let result = try await send(request) - return result.contents - } - - public func listResources(cursor: String? = nil) async throws -> ( - resources: [Resource], nextCursor: String? - ) { - try validateServerCapability(\.resources, "Resources") - let request: Request - if let cursor = cursor { - request = ListResources.request(.init(cursor: cursor)) - } else { - request = ListResources.request(.init()) - } - let result = try await send(request) - return (resources: result.resources, nextCursor: result.nextCursor) - } - - public func subscribeToResource(uri: String) async throws { - try validateServerCapability(\.resources?.subscribe, "Resource subscription") - let request = ResourceSubscribe.request(.init(uri: uri)) - _ = try await send(request) - } - - public func unsubscribeFromResource(uri: String) async throws { - try validateServerCapability(\.resources?.subscribe, "Resource subscription") - let request = ResourceUnsubscribe.request(.init(uri: uri)) - _ = try await send(request) - } - - public func listResourceTemplates(cursor: String? = nil) async throws -> ( - templates: [Resource.Template], nextCursor: String? - ) { - try validateServerCapability(\.resources, "Resources") - let request: Request - if let cursor = cursor { - request = ListResourceTemplates.request(.init(cursor: cursor)) - } else { - request = ListResourceTemplates.request(.init()) - } - let result = try await send(request) - return (templates: result.templates, nextCursor: result.nextCursor) - } - - // MARK: - Tools - - public func listTools(cursor: String? = nil) async throws -> ( - tools: [Tool], nextCursor: String? - ) { - try validateServerCapability(\.tools, "Tools") - let request: Request - if let cursor = cursor { - request = ListTools.request(.init(cursor: cursor)) - } else { - request = ListTools.request(.init()) - } - let result = try await send(request) - return (tools: result.tools, nextCursor: result.nextCursor) - } - - public func callTool(name: String, arguments: [String: Value]? = nil) async throws -> ( - content: [Tool.Content], structuredContent: Value?, isError: Bool? - ) { - try validateServerCapability(\.tools, "Tools") - let request = CallTool.request(.init(name: name, arguments: arguments)) - let result = try await send(request) - // TODO: Add client-side output validation against the tool's outputSchema. - // TypeScript and Python SDKs cache tool outputSchemas from listTools() and - // validate structuredContent when receiving tool results. - return (content: result.content, structuredContent: result.structuredContent, isError: result.isError) - } - - // MARK: - Completions - - /// Request completion suggestions from the server. - /// - /// Completions provide autocomplete suggestions for prompt arguments or resource - /// template URI parameters. - /// - /// - Parameters: - /// - ref: A reference to the prompt or resource template to get completions for. - /// - argument: The argument being completed, including its name and partial value. - /// - context: Optional additional context with previously-resolved argument values. - /// - Returns: The completion suggestions from the server. - public func complete( - ref: CompletionReference, - argument: CompletionArgument, - context: CompletionContext? = nil - ) async throws -> CompletionSuggestions { - try validateServerCapability(\.completions, "Completions") - let request = Complete.request(.init(ref: ref, argument: argument, context: context)) - let result = try await send(request) - return result.completion - } - - // MARK: - Logging + // MARK: - Lifecycle - /// Set the minimum log level for messages from the server. + /// Initialize the connection with the server. /// - /// After calling this method, the server should only send log messages - /// at the specified level or higher (more severe). + /// - Important: This method is deprecated. Initialization now happens automatically + /// when calling `connect(transport:)`. You should use that method instead. /// - /// - Parameter level: The minimum log level to receive. - public func setLoggingLevel(_ level: LoggingLevel) async throws { - try validateServerCapability(\.logging, "Logging") - let request = SetLoggingLevel.request(.init(level: level)) - _ = try await send(request) + /// - Returns: The server's initialization response containing capabilities and server info + @available( + *, deprecated, + message: + "Initialization now happens automatically during connect. Use connect(transport:) instead." + ) + public func initialize() async throws -> Initialize.Result { + return try await _initialize() } - // MARK: - Tasks (Experimental) - // Note: These methods are internal. Access via client.experimental.* - - func getTask(taskId: String) async throws -> GetTask.Result { - try validateServerCapability(\.tasks, "Tasks") - let request = GetTask.request(.init(taskId: taskId)) - return try await send(request) - } + /// Internal initialization implementation + func _initialize() async throws -> Initialize.Result { + let request = Initialize.request( + .init( + protocolVersion: Version.latest, + capabilities: capabilities, + clientInfo: clientInfo + )) - func listTasks(cursor: String? = nil) async throws -> (tasks: [MCPTask], nextCursor: String?) { - try validateServerCapability(\.tasks, "Tasks") - let request: Request - if let cursor { - request = ListTasks.request(.init(cursor: cursor)) - } else { - request = ListTasks.request(.init()) - } let result = try await send(request) - return (tasks: result.tasks, nextCursor: result.nextCursor) - } - - func cancelTask(taskId: String) async throws -> CancelTask.Result { - try validateServerCapability(\.tasks, "Tasks") - let request = CancelTask.request(.init(taskId: taskId)) - return try await send(request) - } - - func getTaskResult(taskId: String) async throws -> GetTaskPayload.Result { - try validateServerCapability(\.tasks, "Tasks") - let request = GetTaskPayload.request(.init(taskId: taskId)) - return try await send(request) - } - - /// Get the task result decoded as a specific type. - /// - /// This method retrieves the task result and decodes the `extraFields` as the specified type. - /// The `extraFields` contain the actual result payload (e.g., CallTool.Result fields). - func getTaskResultAs(taskId: String, type: T.Type) async throws -> T { - let result = try await getTaskResult(taskId: taskId) - - // The result's extraFields contain the actual result payload - // We need to encode them back to JSON and decode as the target type - guard let extraFields = result.extraFields else { - throw MCPError.invalidParams("Task result has no payload") - } - - // Convert extraFields to the target type - let encoder = JSONEncoder() - let decoder = JSONDecoder() - - // Encode the extraFields as JSON - let jsonData = try encoder.encode(extraFields) - - // Decode as the target type - return try decoder.decode(T.self, from: jsonData) - } - - func callToolAsTask( - name: String, - arguments: [String: Value]? = nil, - ttl: Int? = nil - ) async throws -> CreateTaskResult { - try validateServerCapability(\.tasks, "Tasks") - try validateServerCapability(\.tools, "Tools") - - let taskMetadata = TaskMetadata(ttl: ttl) - let request = CallTool.request(.init( - name: name, - arguments: arguments, - task: taskMetadata - )) - - // The server should return CreateTaskResult for task-augmented requests - // We need to decode as CreateTaskResult instead of CallTool.Result - guard let connection = connection else { - throw MCPError.internalError("Client connection not initialized") - } - - let requestData = try encoder.encode(request) - - // Create stream for receiving the response - let (stream, continuation) = AsyncThrowingStream.makeStream() - - let requestId = request.id - continuation.onTermination = { @Sendable [weak self] _ in - Task { await self?.cleanupPendingRequest(id: requestId) } - } - - addPendingRequest(id: request.id, continuation: continuation) - - do { - try await connection.send(requestData) - } catch { - if removePendingRequest(id: request.id) != nil { - continuation.finish(throwing: error) - } - throw error - } - - for try await result in stream { - return result - } - - throw MCPError.internalError("No response received") - } - - func pollTask(taskId: String) -> AsyncThrowingStream { - AsyncThrowingStream { continuation in - let pollingTask = Task { - do { - while !Task.isCancelled { - let task = try await self.getTask(taskId: taskId) - continuation.yield(task) - - if isTerminalStatus(task.status) { - continuation.finish() - return - } - - // Wait based on pollInterval (default 1 second) - let intervalMs = task.pollInterval ?? 1000 - try await Task.sleep(for: .milliseconds(intervalMs)) - } - // Task was cancelled - continuation.finish(throwing: CancellationError()) - } catch { - continuation.finish(throwing: error) - } - } - - // Cancel the polling task when the stream is terminated - continuation.onTermination = { _ in - pollingTask.cancel() - } - } - } - - func pollUntilTerminal(taskId: String) async throws -> GetTask.Result { - for try await status in pollTask(taskId: taskId) { - if isTerminalStatus(status.status) { - return status - } - } - // This shouldn't happen, but handle it gracefully - throw MCPError.internalError("Task polling ended unexpectedly") - } - - func callToolAsTaskAndWait( - name: String, - arguments: [String: Value]? = nil, - ttl: Int? = nil - ) async throws -> (content: [Tool.Content], isError: Bool?) { - // Start the task - let createResult = try await callToolAsTask(name: name, arguments: arguments, ttl: ttl) - let taskId = createResult.task.taskId - - // Wait for the result (uses blocking getTaskResult) - let payloadResult = try await getTaskResult(taskId: taskId) - - // Decode the result as CallTool.Result - // Per MCP spec, the result fields are flattened directly in the response (via extraFields) - guard let extraFields = payloadResult.extraFields else { - throw MCPError.internalError("Task completed but no result available") - } - - // Convert extraFields back to Value for decoding - let resultValue = Value.object(extraFields) - let resultData = try encoder.encode(resultValue) - let toolResult = try decoder.decode(CallTool.Result.self, from: resultData) - return (content: toolResult.content, isError: toolResult.isError) - } - - func callToolStream( - name: String, - arguments: [String: Value]? = nil, - ttl: Int? = nil - ) -> AsyncThrowingStream { - AsyncThrowingStream { continuation in - let streamTask = Task { - do { - // Step 1: Create the task - let createResult = try await self.callToolAsTask(name: name, arguments: arguments, ttl: ttl) - let task = createResult.task - continuation.yield(.taskCreated(task)) - - // Step 2: Poll for status updates until terminal - var lastStatus = task.status - var finalTask = task - - while !isTerminalStatus(lastStatus) { - // Wait based on pollInterval (default 1 second) - let intervalMs = finalTask.pollInterval ?? 1000 - try await Task.sleep(for: .milliseconds(intervalMs)) - - // Get updated status - let statusResult = try await self.getTask(taskId: task.taskId) - finalTask = MCPTask( - taskId: statusResult.taskId, - status: statusResult.status, - ttl: statusResult.ttl, - createdAt: statusResult.createdAt, - lastUpdatedAt: statusResult.lastUpdatedAt, - pollInterval: statusResult.pollInterval, - statusMessage: statusResult.statusMessage - ) - - // Only yield if status or message changed - if statusResult.status != lastStatus || statusResult.statusMessage != nil { - continuation.yield(.taskStatus(finalTask)) - } - lastStatus = statusResult.status - } - - // Step 3: Get the final result - if finalTask.status == .completed { - let payloadResult = try await self.getTaskResult(taskId: task.taskId) - - // Decode the result as CallTool.Result - if let extraFields = payloadResult.extraFields { - let resultValue = Value.object(extraFields) - let resultData = try self.encoder.encode(resultValue) - let toolResult = try self.decoder.decode(CallTool.Result.self, from: resultData) - continuation.yield(.result(toolResult)) - } else { - // No result available - return empty result - continuation.yield(.result(CallTool.Result(content: []))) - } - } else if finalTask.status == .failed { - let error = MCPError.internalError(finalTask.statusMessage ?? "Task failed") - continuation.yield(.error(error)) - } else if finalTask.status == .cancelled { - let error = MCPError.internalError("Task was cancelled") - continuation.yield(.error(error)) - } - - continuation.finish() - } catch let error as MCPError { - continuation.yield(.error(error)) - continuation.finish() - } catch { - let mcpError = MCPError.internalError(error.localizedDescription) - continuation.yield(.error(mcpError)) - continuation.finish() - } - } - - // Cancel the stream task if the stream is terminated - continuation.onTermination = { _ in - streamTask.cancel() - } - } - } - - // MARK: - - - private func handleResponse(_ response: Response) async { - await logger?.trace( - "Processing response", - metadata: ["id": "\(response.id)"]) - - // Check for task-augmented response BEFORE resuming the request. - // Per MCP spec 2025-11-25: progress tokens continue for task lifetime. - // If this is a CreateTaskResult, we need to keep the progress handler alive. - if case .success(let value) = response.result, - case .object(let resultObject) = value { - checkForTaskResponse(response: response, value: resultObject) - } - - // Attempt to remove the pending request using the response ID. - // Resume with the response only if it hadn't yet been removed. - if let removedRequest = self.removePendingRequest(id: response.id) { - // If we successfully removed it, resume its continuation. - switch response.result { - case .success(let value): - removedRequest.resume(returning: value) - case .failure(let error): - removedRequest.resume(throwing: error) - } - } else { - // Request was already removed (e.g., by send error handler or disconnect). - // Log this, but it's not an error in race condition scenarios. - await logger?.warning( - "Attempted to handle response for already removed request", - metadata: ["id": "\(response.id)"] - ) - } - } - - /// Check if a response is a task-augmented response (CreateTaskResult). - /// - /// If the response contains a `task` object with `taskId`, this is a task-augmented - /// response. Per MCP spec, progress notifications can continue until the task reaches - /// terminal status, so we migrate the progress handler from request tracking to task tracking. - /// - /// This matches the TypeScript SDK pattern where task progress tokens are kept alive - /// until the task completes. - private func checkForTaskResponse(response: Response, value: [String: Value]) { - // Check if we have a progress token for this request - guard let progressToken = requestProgressTokens[response.id] else { return } - - // Check if response has task.taskId (CreateTaskResult pattern) - // This mirrors TypeScript's check: result.task?.taskId - guard let taskValue = value["task"], - case .object(let taskObject) = taskValue, - let taskIdValue = taskObject["taskId"], - case .string(let taskId) = taskIdValue else { - // Not a task response - clean up request tracking - // (the progress callback itself is cleaned up in send() after receiving result) - requestProgressTokens.removeValue(forKey: response.id) - return - } - - // This is a task-augmented response! - // Migrate progress token from request tracking to task tracking. - // This keeps the progress handler alive until the task completes. - taskProgressTokens[taskId] = progressToken - requestProgressTokens.removeValue(forKey: response.id) - - Task { - await logger?.debug( - "Keeping progress handler alive for task", - metadata: [ - "taskId": "\(taskId)", - "progressToken": "\(progressToken)", - ] - ) - } - } - - /// Clean up the progress handler for a completed task. - /// - /// Call this method when a task reaches terminal status (completed, failed, cancelled) - /// to remove the progress callback and timeout controller. - /// - /// ## Example - /// - /// ```swift - /// // Register task status notification handler - /// await client.onNotification(TaskStatusNotification.self) { message in - /// if message.params.status.isTerminal { - /// await client.cleanupTaskProgressHandler(taskId: message.params.taskId) - /// } - /// } - /// ``` - /// - /// - Parameter taskId: The ID of the task that completed. - public func cleanupTaskProgressHandler(taskId: String) { - guard let progressToken = taskProgressTokens.removeValue(forKey: taskId) else { return } - - progressCallbacks.removeValue(forKey: progressToken) - timeoutControllers.removeValue(forKey: progressToken) - - Task { - await logger?.debug( - "Cleaned up progress handler for completed task", - metadata: ["taskId": "\(taskId)"] - ) - } - } - - private func handleMessage(_ message: Message) async { - await logger?.trace( - "Processing notification", - metadata: ["method": "\(message.method)"]) - - // Check if this is a progress notification and invoke any registered callback - if message.method == ProgressNotification.name { - await handleProgressNotification(message) - } - - // Check if this is a task status notification and clean up progress handlers - // for terminal task statuses (per MCP spec, progress tokens are valid until terminal status) - if message.method == TaskStatusNotification.name { - await handleTaskStatusNotification(message) - } - - // Find notification handlers for this method - guard let handlers = notificationHandlers[message.method] else { return } - - // Convert notification parameters to concrete type and call handlers - for handler in handlers { - do { - try await handler(message) - } catch { - await logger?.error( - "Error handling notification", - metadata: [ - "method": "\(message.method)", - "error": "\(error)", - ]) - } - } - } - - /// Handle a progress notification by invoking any registered callback. - private func handleProgressNotification(_ message: Message) async { - do { - // Decode as ProgressNotification.Parameters - let paramsData = try encoder.encode(message.params) - let params = try decoder.decode(ProgressNotification.Parameters.self, from: paramsData) - - // Look up the callback for this token - guard let callback = progressCallbacks[params.progressToken] else { - // TypeScript SDK logs an error for unknown progress tokens - await logger?.warning( - "Received progress notification for unknown token", - metadata: ["progressToken": "\(params.progressToken)"]) - return - } - - // Signal the timeout controller if one exists for this token - // This allows resetTimeoutOnProgress to work - if let timeoutController = timeoutControllers[params.progressToken] { - await timeoutController.signalProgress() - } - - // Invoke the callback - let progress = Progress( - value: params.progress, - total: params.total, - message: params.message - ) - await callback(progress) - } catch { - await logger?.warning( - "Failed to decode progress notification", - metadata: ["error": "\(error)"]) - } - } - - /// Handle a task status notification by cleaning up progress handlers for terminal tasks. - /// - /// Per MCP spec 2025-11-25: progress tokens continue throughout task lifetime until terminal status. - /// This method automatically cleans up progress handlers when a task reaches completed, failed, or cancelled. - private func handleTaskStatusNotification(_ message: Message) async { - do { - // Decode as TaskStatusNotification.Parameters - let paramsData = try encoder.encode(message.params) - let params = try decoder.decode(TaskStatusNotification.Parameters.self, from: paramsData) - - // If the task reached a terminal status, clean up its progress handler - if params.status.isTerminal { - cleanupTaskProgressHandler(taskId: params.taskId) - } - } catch { - // Don't log errors for task status notifications - they may not be task-related - // and the user may not have registered a handler for them - } - } - - /// Handle an incoming request from the server (bidirectional communication). - /// - /// This enables server→client requests such as sampling, roots, and elicitation. - /// - /// ## Task-Augmented Request Handling - /// - /// For `sampling/createMessage` and `elicitation/create` requests, this method - /// checks for a `task` field in the request params. If present, it routes to - /// the task-augmented handler (which returns `CreateTaskResult`) instead of - /// the normal handler. - /// - /// This follows the Python SDK pattern of storing task-augmented handlers - /// separately and checking at dispatch time, rather than the TypeScript pattern - /// of wrapping handlers at registration time. The Python pattern was chosen - /// because: - /// - It allows handlers to be registered in any order without losing task-awareness - /// - It keeps task logic separate from normal handler logic - /// - It's more explicit about which handler is called for which request type - private func handleIncomingRequest(_ request: Request) async { - await logger?.trace( - "Processing incoming request from server", - metadata: [ - "method": "\(request.method)", - "id": "\(request.id)", - ]) - - // Validate elicitation mode against client capabilities - // Per spec: Client MUST return -32602 if server requests unsupported mode - if request.method == Elicit.name { - if let modeError = await validateElicitationMode(request) { - await sendResponse(modeError) - return - } - } - - // Check for task-augmented sampling/elicitation requests first - // This matches the Python SDK pattern where task detection happens at dispatch time - if let taskResponse = await handleTaskAugmentedRequest(request) { - await sendResponse(taskResponse) - return - } - - // Find handler for method name - guard let handler = requestHandlers[request.method] else { - await logger?.warning( - "No handler registered for server request", - metadata: ["method": "\(request.method)"]) - // Send error response - let response = AnyMethod.response( - id: request.id, - error: MCPError.methodNotFound("Client has no handler for: \(request.method)") - ) - await sendResponse(response) - return - } - - // Execute the handler and send response - do { - let response = try await handler(request) - - // Check cancellation before sending response (per MCP spec: - // "Receivers of a cancellation notification SHOULD... Not send a response - // for the cancelled request") - if Task.isCancelled { - await logger?.debug( - "Server request cancelled, suppressing response", - metadata: ["id": "\(request.id)"] - ) - return - } - - await sendResponse(response) - } catch { - // Also check cancellation on error path - don't send error response if cancelled - if Task.isCancelled { - await logger?.debug( - "Server request cancelled during error handling, suppressing response", - metadata: ["id": "\(request.id)"] - ) - return - } - - await logger?.error( - "Error handling server request", - metadata: [ - "method": "\(request.method)", - "error": "\(error)", - ]) - let errorResponse = AnyMethod.response( - id: request.id, - error: (error as? MCPError) ?? MCPError.internalError(error.localizedDescription) + // Per MCP spec: "If the client does not support the version in the + // server's response, it SHOULD disconnect." + guard Version.supported.contains(result.protocolVersion) else { + await disconnect() + throw MCPError.invalidRequest( + "Server responded with unsupported protocol version: \(result.protocolVersion). " + + "Supported versions: \(Version.supported.sorted().joined(separator: ", "))" ) - await sendResponse(errorResponse) - } - } - - /// Validate that an elicitation request uses a mode supported by client capabilities. - /// - /// Per MCP spec: Client MUST return -32602 (Invalid params) if server sends - /// an elicitation/create request with a mode not declared in client capabilities. - /// - /// - Parameter request: The incoming elicitation request - /// - Returns: An error response if mode is unsupported, nil if valid - private func validateElicitationMode(_ request: Request) async -> Response? { - do { - let paramsData = try encoder.encode(request.params) - let params = try decoder.decode(Elicit.Parameters.self, from: paramsData) - - switch params { - case .form: - // Form mode requires form capability - if capabilities.elicitation?.form == nil { - return Response( - id: request.id, - error: .invalidParams("Client does not support form elicitation mode") - ) - } - case .url: - // URL mode requires url capability - if capabilities.elicitation?.url == nil { - return Response( - id: request.id, - error: .invalidParams("Client does not support URL elicitation mode") - ) - } - } - } catch { - // If we can't decode the params, let the normal handler deal with it - await logger?.warning( - "Failed to decode elicitation params for mode validation", - metadata: ["error": "\(error)"]) - } - - return nil - } - - /// Check if a request is task-augmented and handle it if so. - /// - /// - Parameter request: The incoming request - /// - Returns: A response if the request was task-augmented and handled, nil otherwise - private func handleTaskAugmentedRequest(_ request: Request) async -> Response? { - do { - // Check for task-augmented sampling request - if request.method == CreateSamplingMessage.name, - let taskHandler = taskAugmentedSamplingHandler { - let paramsData = try encoder.encode(request.params) - let params = try decoder.decode(CreateSamplingMessage.Parameters.self, from: paramsData) - - if let taskMetadata = params.task { - let result = try await taskHandler(params, taskMetadata) - let resultData = try encoder.encode(result) - let resultValue = try decoder.decode(Value.self, from: resultData) - return Response(id: request.id, result: resultValue) - } - } - - // Check for task-augmented elicitation request - if request.method == Elicit.name, - let taskHandler = taskAugmentedElicitationHandler { - let paramsData = try encoder.encode(request.params) - let params = try decoder.decode(Elicit.Parameters.self, from: paramsData) - - let taskMetadata: TaskMetadata? = switch params { - case .form(let formParams): formParams.task - case .url(let urlParams): urlParams.task - } - - if let taskMetadata { - let result = try await taskHandler(params, taskMetadata) - let resultData = try encoder.encode(result) - let resultValue = try decoder.decode(Value.self, from: resultData) - return Response(id: request.id, result: resultValue) - } - } - } catch let error as MCPError { - return Response(id: request.id, error: error) - } catch { - return Response(id: request.id, error: MCPError.internalError(error.localizedDescription)) } - // Not a task-augmented request - return nil - } - - /// Send a response back to the server. - private func sendResponse(_ response: Response) async { - guard let connection = connection else { - await logger?.warning("Cannot send response - client not connected") - return - } + self.serverCapabilities = result.capabilities + self.serverVersion = result.protocolVersion + self.instructions = result.instructions - do { - let responseData = try encoder.encode(response) - try await connection.send(responseData) - } catch { - await logger?.error( - "Failed to send response to server", - metadata: ["error": "\(error)"]) + // HTTP transports must set the protocol version in headers after initialization + if let httpTransport = connection as? HTTPClientTransport { + await httpTransport.setProtocolVersion(result.protocolVersion) } - } - // MARK: - + try await notify(InitializedNotification.message()) - /// Validate the server capabilities. - /// Throws an error if the client is configured to be strict and the capability is not supported. - private func validateServerCapability( - _ keyPath: KeyPath, - _ name: String - ) - throws - { - if configuration.strict { - guard let capabilities = serverCapabilities else { - throw MCPError.methodNotFound("Server capabilities not initialized") - } - guard capabilities[keyPath: keyPath] != nil else { - throw MCPError.methodNotFound("\(name) is not supported by the server") - } - } + return result } - // Add handler for batch responses - private func handleBatchResponse(_ responses: [AnyResponse]) async { - await logger?.trace("Processing batch response", metadata: ["count": "\(responses.count)"]) - for response in responses { - // Attempt to remove the pending request. - // If successful, pendingRequest contains the request. - if let pendingRequest = self.removePendingRequest(id: response.id) { - // If we successfully removed it, handle the response using the pending request. - switch response.result { - case .success(let value): - pendingRequest.resume(returning: value) - case .failure(let error): - pendingRequest.resume(throwing: error) - } - } else { - // If removal failed, it means the request ID was not found (or already handled). - // Log a warning. - await logger?.warning( - "Received response in batch for unknown or already handled request ID", - metadata: ["id": "\(response.id)"] - ) - } - } + public func ping() async throws { + let request = Ping.request() + _ = try await send(request) } } diff --git a/Sources/MCP/Server/Server+ClientRequests.swift b/Sources/MCP/Server/Server+ClientRequests.swift new file mode 100644 index 00000000..595c7acd --- /dev/null +++ b/Sources/MCP/Server/Server+ClientRequests.swift @@ -0,0 +1,359 @@ +import Foundation + +extension Server { + // MARK: - Server to Client Requests + + /// Send a request to the client and wait for a response. + /// + /// This enables bidirectional communication where the server can request + /// information from the client (e.g., roots, sampling, elicitation). + /// + /// - Parameter request: The request to send + /// - Returns: The result from the client + public func sendRequest(_ request: Request) async throws -> M.Result { + guard let connection = connection else { + throw MCPError.internalError("Server connection not initialized") + } + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + let requestData = try encoder.encode(request) + + // Create stream for receiving the response + let (stream, continuation) = AsyncThrowingStream.makeStream() + + // Clean up pending request if cancelled + let requestId = request.id + continuation.onTermination = { @Sendable [weak self] _ in + Task { await self?.cleanUpPendingRequest(id: requestId) } + } + + // Register the pending request + pendingRequests[request.id] = AnyServerPendingRequest(continuation: continuation) + + // Send the request + do { + try await connection.send(requestData) + } catch { + pendingRequests.removeValue(forKey: request.id) + continuation.finish(throwing: error) + throw error + } + + // Wait for response + for try await result in stream { + return result + } + + throw MCPError.internalError("No response received from client") + } + + func cleanUpPendingRequest(id: RequestId) { + pendingRequests.removeValue(forKey: id) + } + + // MARK: - In-Flight Request Tracking (Protocol-Level Cancellation) + + /// Track an in-flight request handler Task. + func trackInFlightRequest(_ requestId: RequestId, task: Task) { + inFlightHandlerTasks[requestId] = task + } + + /// Remove an in-flight request handler Task. + func removeInFlightRequest(_ requestId: RequestId) { + inFlightHandlerTasks.removeValue(forKey: requestId) + } + + /// Cancel an in-flight request handler Task. + /// + /// Called when a CancelledNotification is received for a specific requestId. + /// Per MCP spec, if the request is unknown or already completed, this is a no-op. + func cancelInFlightRequest(_ requestId: RequestId, reason: String?) async { + if let task = inFlightHandlerTasks[requestId] { + task.cancel() + await logger?.debug( + "Cancelled in-flight request", + metadata: [ + "id": "\(requestId)", + "reason": "\(reason ?? "none")", + ] + ) + } + // Per spec: MAY ignore if request is unknown - no error needed + } + + /// Generate a unique request ID for server→client requests. + func generateRequestId() -> RequestId { + let id = nextRequestId + nextRequestId += 1 + return .number(id) + } + + /// Request the list of roots from the client. + /// + /// Roots represent filesystem directories that the client has access to. + /// Servers can use this to understand the scope of files they can work with. + /// + /// - Throws: MCPError if the client doesn't support roots or if the request fails. + /// - Returns: The list of roots from the client. + public func listRoots() async throws -> [Root] { + // Check that client supports roots + guard clientCapabilities?.roots != nil else { + throw MCPError.invalidRequest("Client does not support roots capability") + } + + let request: Request = ListRoots.request(id: generateRequestId()) + let result = try await sendRequest(request) + return result.roots + } + + /// Request a sampling completion from the client (without tools). + /// + /// This enables servers to request LLM completions through the client, + /// allowing sophisticated agentic behaviors while maintaining security. + /// + /// The result will be a single content block (text, image, or audio). + /// For tool-enabled sampling, use `createMessageWithTools(_:)` instead. + /// + /// - Parameter params: The sampling parameters including messages, model preferences, etc. + /// - Throws: MCPError if the client doesn't support sampling or if the request fails. + /// - Returns: The sampling result from the client containing a single content block. + public func createMessage(_ params: CreateSamplingMessage.Parameters) async throws -> CreateSamplingMessage.Result { + // Check that client supports sampling + guard clientCapabilities?.sampling != nil else { + throw MCPError.invalidRequest("Client does not support sampling capability") + } + + let request: Request = CreateSamplingMessage.request(id: generateRequestId(), params) + return try await sendRequest(request) + } + + /// Request a sampling completion from the client with tool support. + /// + /// This enables servers to request LLM completions that may involve tool use. + /// The result may contain tool use content, and content can be an array for parallel tool calls. + /// + /// - Parameter params: The sampling parameters including messages, tools, and model preferences. + /// - Throws: MCPError if the client doesn't support sampling or tool capabilities. + /// - Returns: The sampling result from the client, which may include tool use content. + public func createMessageWithTools(_ params: CreateSamplingMessageWithTools.Parameters) async throws -> CreateSamplingMessageWithTools.Result { + // Check that client supports sampling + guard clientCapabilities?.sampling != nil else { + throw MCPError.invalidRequest("Client does not support sampling capability") + } + + // Check tools capability + guard clientCapabilities?.sampling?.tools != nil else { + throw MCPError.invalidRequest("Client does not support sampling tools capability") + } + + // Validate tool_use/tool_result message structure per MCP specification + try Sampling.Message.validateToolUseResultMessages(params.messages) + + let request: Request = CreateSamplingMessageWithTools.request(id: generateRequestId(), params) + return try await sendRequest(request) + } + + /// Request user input via elicitation from the client. + /// + /// Elicitation allows servers to request structured input from users through + /// the client, either via forms or external URLs (e.g., OAuth flows). + /// + /// - Parameter params: The elicitation parameters. + /// - Throws: MCPError if the client doesn't support elicitation or if the request fails. + /// - Returns: The elicitation result from the client. + public func elicit(_ params: Elicit.Parameters) async throws -> Elicit.Result { + // Check that client supports elicitation + guard clientCapabilities?.elicitation != nil else { + throw MCPError.invalidRequest("Client does not support elicitation capability") + } + + // Check mode-specific capabilities + switch params { + case .form: + guard clientCapabilities?.elicitation?.form != nil else { + throw MCPError.invalidRequest("Client does not support form elicitation") + } + case .url: + guard clientCapabilities?.elicitation?.url != nil else { + throw MCPError.invalidRequest("Client does not support URL elicitation") + } + } + + let request: Request = Elicit.request(id: generateRequestId(), params) + let result = try await sendRequest(request) + + // TODO: Add elicitation response validation against the requestedSchema. + // TypeScript SDK uses JSON Schema validators (AJV, CfWorker) to validate + // elicitation responses against the requestedSchema. Python SDK uses Pydantic. + // The ideal solution is to use the same JSON Schema validator for both + // elicitation and tool validation, for spec compliance and consistency. + + return result + } + + func checkInitialized() throws { + guard isInitialized else { + throw MCPError.invalidRequest("Server is not initialized") + } + } + + // MARK: - Client Task Polling (Server → Client) + + /// Get a task from the client. + /// + /// Internal method used by experimental server task features. + func getClientTask(taskId: String) async throws -> GetTask.Result { + guard clientCapabilities?.tasks != nil else { + throw MCPError.invalidRequest("Client does not support tasks capability") + } + + let request = GetTask.request(.init(taskId: taskId)) + return try await sendRequest(request) + } + + /// Get the result payload of a client task. + /// + /// Internal method used by experimental server task features. + func getClientTaskResult(taskId: String) async throws -> GetTaskPayload.Result { + guard clientCapabilities?.tasks != nil else { + throw MCPError.invalidRequest("Client does not support tasks capability") + } + + let request = GetTaskPayload.request(.init(taskId: taskId)) + return try await sendRequest(request) + } + + /// Get the task result decoded as a specific type. + /// + /// Internal method used by experimental server task features. + func getClientTaskResultAs(taskId: String, type: T.Type) async throws -> T { + let result = try await getClientTaskResult(taskId: taskId) + + // The result's extraFields contain the actual result payload + guard let extraFields = result.extraFields else { + throw MCPError.invalidParams("Task result has no payload") + } + + // Convert extraFields to the target type + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let jsonData = try encoder.encode(extraFields) + return try decoder.decode(T.self, from: jsonData) + } + + // MARK: - Task-Augmented Requests (Server → Client) + + /// Send a task-augmented elicitation request to the client. + /// + /// The client returns a `CreateTaskResult` instead of an `ElicitResult`. + /// Use client task polling to get the final result. + /// + /// Internal method used by experimental server task features. + func sendElicitAsTask(_ params: Elicit.Parameters) async throws -> CreateTaskResult { + // Check that client supports task-augmented elicitation + try requireTaskAugmentedElicitation(clientCapabilities) + + // Check mode-specific capabilities + switch params { + case .form: + guard clientCapabilities?.elicitation?.form != nil else { + throw MCPError.invalidRequest("Client does not support form elicitation") + } + case .url: + guard clientCapabilities?.elicitation?.url != nil else { + throw MCPError.invalidRequest("Client does not support URL elicitation") + } + } + + guard let connection else { + throw MCPError.internalError("Server connection not initialized") + } + + // Build the request + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + let request: Request = Elicit.request(id: generateRequestId(), params) + let requestData = try encoder.encode(request) + + // Create stream for receiving the response + let (stream, continuation) = AsyncThrowingStream.makeStream() + + let requestId = request.id + continuation.onTermination = { @Sendable [weak self] _ in + Task { await self?.cleanUpPendingRequest(id: requestId) } + } + + // Register the pending request + pendingRequests[requestId] = AnyServerPendingRequest(continuation: continuation) + + // Send the request + do { + try await connection.send(requestData) + } catch { + pendingRequests.removeValue(forKey: requestId) + continuation.finish(throwing: error) + throw error + } + + // Wait for single result + for try await result in stream { + return result + } + + throw MCPError.internalError("No response received") + } + + /// Send a task-augmented sampling request to the client. + /// + /// The client returns a `CreateTaskResult` instead of a `CreateSamplingMessage.Result`. + /// Use client task polling to get the final result. + /// + /// Internal method used by experimental server task features. + func sendCreateMessageAsTask(_ params: CreateSamplingMessage.Parameters) async throws -> CreateTaskResult { + // Check that client supports task-augmented sampling + try requireTaskAugmentedSampling(clientCapabilities) + + guard clientCapabilities?.sampling != nil else { + throw MCPError.invalidRequest("Client does not support sampling capability") + } + + guard let connection else { + throw MCPError.internalError("Server connection not initialized") + } + + // Build the request + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + let request = CreateSamplingMessage.request(id: generateRequestId(), params) + let requestData = try encoder.encode(request) + + // Create stream for receiving the response + let (stream, continuation) = AsyncThrowingStream.makeStream() + + let requestId = request.id + continuation.onTermination = { @Sendable [weak self] _ in + Task { await self?.cleanUpPendingRequest(id: requestId) } + } + + // Register the pending request + pendingRequests[requestId] = AnyServerPendingRequest(continuation: continuation) + + // Send the request + do { + try await connection.send(requestData) + } catch { + pendingRequests.removeValue(forKey: requestId) + continuation.finish(throwing: error) + throw error + } + + // Wait for single result + for try await result in stream { + return result + } + + throw MCPError.internalError("No response received") + } +} diff --git a/Sources/MCP/Server/Server+RequestHandling.swift b/Sources/MCP/Server/Server+RequestHandling.swift new file mode 100644 index 00000000..3d03a96c --- /dev/null +++ b/Sources/MCP/Server/Server+RequestHandling.swift @@ -0,0 +1,366 @@ +import Foundation + +extension Server { + // MARK: - Request Handling + + /// A JSON-RPC batch containing multiple requests and/or notifications + struct Batch: Sendable { + /// An item in a JSON-RPC batch + enum Item: Sendable { + case request(Request) + case notification(Message) + + } + + var items: [Item] + + init(items: [Item]) { + self.items = items + } + } + + /// Process a batch of requests and/or notifications + func handleBatch(_ batch: Batch) async throws { + // Capture the connection at batch start. + // This ensures all batch responses go to the correct client. + let capturedConnection = self.connection + + await logger?.trace("Processing batch request", metadata: ["size": "\(batch.items.count)"]) + + if batch.items.isEmpty { + // Empty batch is invalid according to JSON-RPC spec + let error = MCPError.invalidRequest("Batch array must not be empty") + let response = AnyMethod.response(id: .random, error: error) + // Use captured connection for error response + if let connection = capturedConnection { + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + let responseData = try encoder.encode(response) + try await connection.send(responseData) + } + return + } + + // Process each item in the batch and collect responses + var responses: [Response] = [] + + for item in batch.items { + do { + switch item { + case .request(let request): + // For batched requests, collect responses instead of sending immediately + if let response = try await handleRequest(request, sendResponse: false) { + responses.append(response) + } + + case .notification(let notification): + // Handle notification (no response needed) + try await handleMessage(notification) + } + } catch { + // Only add errors to response for requests (notifications don't have responses) + if case .request(let request) = item { + let mcpError = + error as? MCPError ?? MCPError.internalError(error.localizedDescription) + responses.append(AnyMethod.response(id: request.id, error: mcpError)) + } + } + } + + // Send collected responses if any (using captured connection) + if !responses.isEmpty { + guard let connection = capturedConnection else { + await logger?.warning("Cannot send batch response - connection was nil at batch start") + return + } + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + let responseData = try encoder.encode(responses) + + try await connection.send(responseData) + } + } + + // MARK: - Request and Message Handling + + /// Internal context for routing responses to the correct transport. + /// + /// When handling requests, we capture the current connection at request time. + /// This ensures that when the handler completes (which may be async), the response + /// is sent to the correct client even if `self.connection` has changed in the meantime. + /// + /// This pattern is critical for HTTP transports where multiple clients can connect + /// and the server's `connection` reference gets reassigned. + struct RequestContext { + /// The transport connection captured at request time + let capturedConnection: (any Transport)? + /// The ID of the request being handled + let requestId: RequestId + /// The session ID from the transport, if available. + /// + /// For HTTP transports with multiple concurrent clients, this identifies + /// the specific session. Used for per-session features like log levels. + let sessionId: String? + } + + /// Wrapper for encoding type-erased notifications as JSON-RPC messages. + private struct NotificationWrapper: Encodable { + let jsonrpc = "2.0" + let method: String + let params: Value + + init(notification: any Notification) { + self.method = type(of: notification).name + + // Encode the notification's params to Value + // Since Notification is Codable, we encode it and extract the params field + let encoder = JSONEncoder() + let decoder = JSONDecoder() + if let data = try? encoder.encode(notification), + let dict = try? decoder.decode([String: Value].self, from: data), + let params = dict["params"] { + self.params = params + } else { + self.params = .object([:]) + } + } + } + + /// Send a response using the captured request context. + /// + /// This ensures responses are routed to the correct client by: + /// 1. Using the connection that was active when the request was received + /// 2. Passing the request ID so multiplexing transports can route correctly + func send(_ response: Response, using context: RequestContext) async throws { + guard let connection = context.capturedConnection else { + await logger?.warning( + "Cannot send response - connection was nil at request time", + metadata: ["requestId": "\(context.requestId)"] + ) + return + } + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + + let responseData = try encoder.encode(response) + try await connection.send(responseData, relatedRequestId: context.requestId) + } + + /// Handle a request and either send the response immediately or return it + /// + /// - Parameters: + /// - request: The request to handle + /// - sendResponse: Whether to send the response immediately (true) or return it (false) + /// - Returns: The response when sendResponse is false + func handleRequest(_ request: Request, sendResponse: Bool = true) + async throws -> Response? + { + // Capture the connection and session ID at request time. + // This ensures responses go to the correct client even if self.connection + // changes while the handler is executing (e.g., another client connects). + let capturedConnection = self.connection + let context = RequestContext( + capturedConnection: capturedConnection, + requestId: request.id, + sessionId: await capturedConnection?.sessionId + ) + + // Check if this is a pre-processed error request (empty method) + if request.method.isEmpty && !sendResponse { + // This is a placeholder for an invalid request that couldn't be parsed in batch mode + return AnyMethod.response( + id: request.id, + error: MCPError.invalidRequest("Invalid batch item format") + ) + } + + await logger?.trace( + "Processing request", + metadata: [ + "method": "\(request.method)", + "id": "\(request.id)", + ]) + + if configuration.strict { + // The client SHOULD NOT send requests other than pings + // before the server has responded to the initialize request. + switch request.method { + case Initialize.name, Ping.name: + break + default: + try checkInitialized() + } + } + + // Find handler for method name + guard let handler = methodHandlers[request.method] else { + let error = MCPError.methodNotFound("Unknown method: \(request.method)") + let response = AnyMethod.response(id: request.id, error: error) + + if sendResponse { + try await send(response, using: context) + return nil + } + + return response + } + + // Create the public handler context with sendNotification capability + let handlerContext = RequestHandlerContext( + sendNotification: { [context] notification in + guard let connection = context.capturedConnection else { + throw MCPError.internalError("Cannot send notification - connection was nil at request time") + } + + // Wrap the notification in a JSON-RPC message structure + let wrapper = NotificationWrapper(notification: notification) + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + + let notificationData = try encoder.encode(wrapper) + try await connection.send(notificationData, relatedRequestId: context.requestId) + }, + sendMessage: { [context] message in + guard let connection = context.capturedConnection else { + throw MCPError.internalError("Cannot send notification - connection was nil at request time") + } + + // Message already encodes to JSON-RPC format with method and params + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + + let messageData = try encoder.encode(message) + try await connection.send(messageData, relatedRequestId: context.requestId) + }, + sendData: { [context] data in + guard let connection = context.capturedConnection else { + throw MCPError.internalError("Cannot send data - connection was nil at request time") + } + + // Send raw data (used for queued task messages) + try await connection.send(data, relatedRequestId: context.requestId) + }, + sessionId: context.sessionId, + shouldSendLogMessage: { [weak self, context] level in + guard let self else { return true } + return await self.shouldSendLogMessage(at: level, forSession: context.sessionId) + } + ) + + do { + // Handle request and get response + let response: Response = try await handler(request, context: handlerContext) + + // Check cancellation before sending response (per MCP spec: + // "Receivers of a cancellation notification SHOULD... Not send a response + // for the cancelled request") + if Task.isCancelled { + await logger?.debug( + "Request cancelled, suppressing response", + metadata: ["id": "\(request.id)"] + ) + return nil + } + + if sendResponse { + try await send(response, using: context) + return nil + } + + return response + } catch { + // Also check cancellation on error path - don't send error response if cancelled + if Task.isCancelled { + await logger?.debug( + "Request cancelled during error handling, suppressing response", + metadata: ["id": "\(request.id)"] + ) + return nil + } + + let mcpError = error as? MCPError ?? MCPError.internalError(error.localizedDescription) + let response: Response = AnyMethod.response(id: request.id, error: mcpError) + + if sendResponse { + try await send(response, using: context) + return nil + } + + return response + } + } + + func handleMessage(_ message: Message) async throws { + await logger?.trace( + "Processing notification", + metadata: ["method": "\(message.method)"]) + + if configuration.strict { + // Check initialization state unless this is an initialized notification + if message.method != InitializedNotification.name { + try checkInitialized() + } + } + + // Find notification handlers for this method + guard let handlers = notificationHandlers[message.method] else { return } + + // Convert notification parameters to concrete type and call handlers + for handler in handlers { + do { + try await handler(message) + } catch { + await logger?.error( + "Error handling notification", + metadata: [ + "method": "\(message.method)", + "error": "\(error)", + ]) + } + } + } + + /// Handle a response from the client (for server→client requests). + func handleClientResponse(_ response: Response) async { + await logger?.trace( + "Processing client response", + metadata: ["id": "\(response.id)"]) + + // Check response routers first (e.g., for task-related responses) + for router in responseRouters { + switch response.result { + case .success(let value): + if await router.routeResponse(requestId: response.id, response: value) { + await logger?.trace( + "Response routed via router", + metadata: ["id": "\(response.id)"]) + return + } + case .failure(let error): + if await router.routeError(requestId: response.id, error: error) { + await logger?.trace( + "Error routed via router", + metadata: ["id": "\(response.id)"]) + return + } + } + } + + // Fall back to normal pending request handling + if let pendingRequest = pendingRequests.removeValue(forKey: response.id) { + switch response.result { + case .success(let value): + pendingRequest.resume(returning: value) + case .failure(let error): + pendingRequest.resume(throwing: error) + } + } else { + await logger?.warning( + "Received response for unknown request", + metadata: ["id": "\(response.id)"]) + } + } +} diff --git a/Sources/MCP/Server/Server+Sending.swift b/Sources/MCP/Server/Server+Sending.swift new file mode 100644 index 00000000..4567aa4e --- /dev/null +++ b/Sources/MCP/Server/Server+Sending.swift @@ -0,0 +1,66 @@ +import Foundation + +extension Server { + // MARK: - Sending + + /// Send a response to a request + public func send(_ response: Response) async throws { + guard let connection else { + throw MCPError.internalError("Server connection not initialized") + } + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + + let responseData = try encoder.encode(response) + try await connection.send(responseData) + } + + /// Send a notification to connected clients + public func notify(_ notification: Message) async throws { + guard let connection = connection else { + throw MCPError.internalError("Server connection not initialized") + } + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + + let notificationData = try encoder.encode(notification) + try await connection.send(notificationData) + } + + /// Send a log message notification to connected clients. + /// + /// This method can be called outside of request handlers to send log messages + /// asynchronously. The message will only be sent if: + /// - The server has declared the `logging` capability + /// - The message's level is at or above the minimum level set by the session + /// + /// If the logging capability is not declared, this method silently returns without + /// sending (matching TypeScript SDK behavior). + /// + /// - Parameters: + /// - level: The severity level of the log message + /// - logger: An optional name for the logger producing the message + /// - data: The log message data (can be a string or structured data) + /// - sessionId: Optional session ID for per-session log level filtering. + /// If `nil`, the log level for the nil-session (default) is used. + public func sendLogMessage( + level: LoggingLevel, + logger: String? = nil, + data: Value, + sessionId: String? = nil + ) async throws { + // Check if logging capability is declared (matching TypeScript SDK behavior) + guard capabilities.logging != nil else { return } + + // Check if this message should be sent based on the session's log level + guard shouldSendLogMessage(at: level, forSession: sessionId) else { return } + + try await notify(LogMessageNotification.message(.init( + level: level, + logger: logger, + data: data + ))) + } +} diff --git a/Sources/MCP/Server/Server.swift b/Sources/MCP/Server/Server.swift index 3e3e01f0..773ac6e9 100644 --- a/Sources/MCP/Server/Server.swift +++ b/Sources/MCP/Server/Server.swift @@ -385,7 +385,7 @@ public actor Server { } /// A type-erased pending request for server→client requests (bidirectional communication). - private struct AnyServerPendingRequest { + struct AnyServerPendingRequest { private let _yield: (Result) -> Void private let _finish: () -> Void @@ -430,11 +430,11 @@ public actor Server { } /// Server information - private let serverInfo: Server.Info + let serverInfo: Server.Info /// The server connection - private var connection: (any Transport)? + var connection: (any Transport)? /// The server logger - private var logger: Logger? { + var logger: Logger? { get async { await connection?.logger } @@ -474,29 +474,29 @@ public actor Server { } /// Request handlers - private var methodHandlers: [String: RequestHandlerBox] = [:] + var methodHandlers: [String: RequestHandlerBox] = [:] /// Notification handlers - private var notificationHandlers: [String: [NotificationHandlerBox]] = [:] + var notificationHandlers: [String: [NotificationHandlerBox]] = [:] /// Pending requests sent from server to client (for bidirectional communication) - private var pendingRequests: [RequestId: AnyServerPendingRequest] = [:] + var pendingRequests: [RequestId: AnyServerPendingRequest] = [:] /// Counter for generating unique request IDs - private var nextRequestId = 0 + var nextRequestId = 0 /// Response routers for intercepting responses before normal handling - private var responseRouters: [any ResponseRouter] = [] + var responseRouters: [any ResponseRouter] = [] /// Whether the server is initialized - private var isInitialized = false + var isInitialized = false /// The client information - private var clientInfo: Client.Info? + var clientInfo: Client.Info? /// The client capabilities - private var clientCapabilities: Client.Capabilities? + var clientCapabilities: Client.Capabilities? /// The protocol version - private var protocolVersion: String? + var protocolVersion: String? /// The list of subscriptions - private var subscriptions: [String: Set] = [:] + var subscriptions: [String: Set] = [:] /// The task for the message handling loop - private var task: Task? + var task: Task? /// Per-session minimum log levels set by clients. /// /// For HTTP transports with multiple concurrent clients, each session can @@ -504,11 +504,11 @@ public actor Server { /// transports without session support like stdio). /// /// Log messages below a session's level will be filtered out for that session. - private var loggingLevels: [String?: LoggingLevel] = [:] + var loggingLevels: [String?: LoggingLevel] = [:] /// In-flight request handler Tasks, tracked by request ID. /// Used for protocol-level cancellation when CancelledNotification is received. - private var inFlightHandlerTasks: [RequestId: Task] = [:] + var inFlightHandlerTasks: [RequestId: Task] = [:] public init( name: String, @@ -740,787 +740,7 @@ public actor Server { return self } - // MARK: - Sending - - /// Send a response to a request - public func send(_ response: Response) async throws { - guard let connection = connection else { - throw MCPError.internalError("Server connection not initialized") - } - - let encoder = JSONEncoder() - encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] - - let responseData = try encoder.encode(response) - try await connection.send(responseData) - } - - /// Send a notification to connected clients - public func notify(_ notification: Message) async throws { - guard let connection = connection else { - throw MCPError.internalError("Server connection not initialized") - } - - let encoder = JSONEncoder() - encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] - - let notificationData = try encoder.encode(notification) - try await connection.send(notificationData) - } - - /// Send a log message notification to connected clients. - /// - /// This method can be called outside of request handlers to send log messages - /// asynchronously. The message will only be sent if: - /// - The server has declared the `logging` capability - /// - The message's level is at or above the minimum level set by the session - /// - /// If the logging capability is not declared, this method silently returns without - /// sending (matching TypeScript SDK behavior). - /// - /// - Parameters: - /// - level: The severity level of the log message - /// - logger: An optional name for the logger producing the message - /// - data: The log message data (can be a string or structured data) - /// - sessionId: Optional session ID for per-session log level filtering. - /// If `nil`, the log level for the nil-session (default) is used. - public func sendLogMessage( - level: LoggingLevel, - logger: String? = nil, - data: Value, - sessionId: String? = nil - ) async throws { - // Check if logging capability is declared (matching TypeScript SDK behavior) - guard capabilities.logging != nil else { return } - - // Check if this message should be sent based on the session's log level - guard shouldSendLogMessage(at: level, forSession: sessionId) else { return } - - try await notify(LogMessageNotification.message(.init( - level: level, - logger: logger, - data: data - ))) - } - - /// A JSON-RPC batch containing multiple requests and/or notifications - struct Batch: Sendable { - /// An item in a JSON-RPC batch - enum Item: Sendable { - case request(Request) - case notification(Message) - - } - - var items: [Item] - - init(items: [Item]) { - self.items = items - } - } - - /// Process a batch of requests and/or notifications - private func handleBatch(_ batch: Batch) async throws { - // Capture the connection at batch start. - // This ensures all batch responses go to the correct client. - let capturedConnection = self.connection - - await logger?.trace("Processing batch request", metadata: ["size": "\(batch.items.count)"]) - - if batch.items.isEmpty { - // Empty batch is invalid according to JSON-RPC spec - let error = MCPError.invalidRequest("Batch array must not be empty") - let response = AnyMethod.response(id: .random, error: error) - // Use captured connection for error response - if let connection = capturedConnection { - let encoder = JSONEncoder() - encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] - let responseData = try encoder.encode(response) - try await connection.send(responseData) - } - return - } - - // Process each item in the batch and collect responses - var responses: [Response] = [] - - for item in batch.items { - do { - switch item { - case .request(let request): - // For batched requests, collect responses instead of sending immediately - if let response = try await handleRequest(request, sendResponse: false) { - responses.append(response) - } - - case .notification(let notification): - // Handle notification (no response needed) - try await handleMessage(notification) - } - } catch { - // Only add errors to response for requests (notifications don't have responses) - if case .request(let request) = item { - let mcpError = - error as? MCPError ?? MCPError.internalError(error.localizedDescription) - responses.append(AnyMethod.response(id: request.id, error: mcpError)) - } - } - } - - // Send collected responses if any (using captured connection) - if !responses.isEmpty { - guard let connection = capturedConnection else { - await logger?.warning("Cannot send batch response - connection was nil at batch start") - return - } - - let encoder = JSONEncoder() - encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] - let responseData = try encoder.encode(responses) - - try await connection.send(responseData) - } - } - - // MARK: - Request and Message Handling - - /// Internal context for routing responses to the correct transport. - /// - /// When handling requests, we capture the current connection at request time. - /// This ensures that when the handler completes (which may be async), the response - /// is sent to the correct client even if `self.connection` has changed in the meantime. - /// - /// This pattern is critical for HTTP transports where multiple clients can connect - /// and the server's `connection` reference gets reassigned. - private struct RequestContext { - /// The transport connection captured at request time - let capturedConnection: (any Transport)? - /// The ID of the request being handled - let requestId: RequestId - /// The session ID from the transport, if available. - /// - /// For HTTP transports with multiple concurrent clients, this identifies - /// the specific session. Used for per-session features like log levels. - let sessionId: String? - } - - /// Wrapper for encoding type-erased notifications as JSON-RPC messages. - private struct NotificationWrapper: Encodable { - let jsonrpc = "2.0" - let method: String - let params: Value - - init(notification: any Notification) { - self.method = type(of: notification).name - - // Encode the notification's params to Value - // Since Notification is Codable, we encode it and extract the params field - let encoder = JSONEncoder() - let decoder = JSONDecoder() - if let data = try? encoder.encode(notification), - let dict = try? decoder.decode([String: Value].self, from: data), - let params = dict["params"] { - self.params = params - } else { - self.params = .object([:]) - } - } - } - - /// Send a response using the captured request context. - /// - /// This ensures responses are routed to the correct client by: - /// 1. Using the connection that was active when the request was received - /// 2. Passing the request ID so multiplexing transports can route correctly - private func send(_ response: Response, using context: RequestContext) async throws { - guard let connection = context.capturedConnection else { - await logger?.warning( - "Cannot send response - connection was nil at request time", - metadata: ["requestId": "\(context.requestId)"] - ) - return - } - - let encoder = JSONEncoder() - encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] - - let responseData = try encoder.encode(response) - try await connection.send(responseData, relatedRequestId: context.requestId) - } - - /// Handle a request and either send the response immediately or return it - /// - /// - Parameters: - /// - request: The request to handle - /// - sendResponse: Whether to send the response immediately (true) or return it (false) - /// - Returns: The response when sendResponse is false - private func handleRequest(_ request: Request, sendResponse: Bool = true) - async throws -> Response? - { - // Capture the connection and session ID at request time. - // This ensures responses go to the correct client even if self.connection - // changes while the handler is executing (e.g., another client connects). - let capturedConnection = self.connection - let context = RequestContext( - capturedConnection: capturedConnection, - requestId: request.id, - sessionId: await capturedConnection?.sessionId - ) - - // Check if this is a pre-processed error request (empty method) - if request.method.isEmpty && !sendResponse { - // This is a placeholder for an invalid request that couldn't be parsed in batch mode - return AnyMethod.response( - id: request.id, - error: MCPError.invalidRequest("Invalid batch item format") - ) - } - - await logger?.trace( - "Processing request", - metadata: [ - "method": "\(request.method)", - "id": "\(request.id)", - ]) - - if configuration.strict { - // The client SHOULD NOT send requests other than pings - // before the server has responded to the initialize request. - switch request.method { - case Initialize.name, Ping.name: - break - default: - try checkInitialized() - } - } - - // Find handler for method name - guard let handler = methodHandlers[request.method] else { - let error = MCPError.methodNotFound("Unknown method: \(request.method)") - let response = AnyMethod.response(id: request.id, error: error) - - if sendResponse { - try await send(response, using: context) - return nil - } - - return response - } - - // Create the public handler context with sendNotification capability - let handlerContext = RequestHandlerContext( - sendNotification: { [context] notification in - guard let connection = context.capturedConnection else { - throw MCPError.internalError("Cannot send notification - connection was nil at request time") - } - - // Wrap the notification in a JSON-RPC message structure - let wrapper = NotificationWrapper(notification: notification) - - let encoder = JSONEncoder() - encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] - - let notificationData = try encoder.encode(wrapper) - try await connection.send(notificationData, relatedRequestId: context.requestId) - }, - sendMessage: { [context] message in - guard let connection = context.capturedConnection else { - throw MCPError.internalError("Cannot send notification - connection was nil at request time") - } - - // Message already encodes to JSON-RPC format with method and params - let encoder = JSONEncoder() - encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] - - let messageData = try encoder.encode(message) - try await connection.send(messageData, relatedRequestId: context.requestId) - }, - sendData: { [context] data in - guard let connection = context.capturedConnection else { - throw MCPError.internalError("Cannot send data - connection was nil at request time") - } - - // Send raw data (used for queued task messages) - try await connection.send(data, relatedRequestId: context.requestId) - }, - sessionId: context.sessionId, - shouldSendLogMessage: { [weak self, context] level in - guard let self else { return true } - return await self.shouldSendLogMessage(at: level, forSession: context.sessionId) - } - ) - - do { - // Handle request and get response - let response: Response = try await handler(request, context: handlerContext) - - // Check cancellation before sending response (per MCP spec: - // "Receivers of a cancellation notification SHOULD... Not send a response - // for the cancelled request") - if Task.isCancelled { - await logger?.debug( - "Request cancelled, suppressing response", - metadata: ["id": "\(request.id)"] - ) - return nil - } - - if sendResponse { - try await send(response, using: context) - return nil - } - - return response - } catch { - // Also check cancellation on error path - don't send error response if cancelled - if Task.isCancelled { - await logger?.debug( - "Request cancelled during error handling, suppressing response", - metadata: ["id": "\(request.id)"] - ) - return nil - } - - let mcpError = error as? MCPError ?? MCPError.internalError(error.localizedDescription) - let response: Response = AnyMethod.response(id: request.id, error: mcpError) - - if sendResponse { - try await send(response, using: context) - return nil - } - - return response - } - } - - private func handleMessage(_ message: Message) async throws { - await logger?.trace( - "Processing notification", - metadata: ["method": "\(message.method)"]) - - if configuration.strict { - // Check initialization state unless this is an initialized notification - if message.method != InitializedNotification.name { - try checkInitialized() - } - } - - // Find notification handlers for this method - guard let handlers = notificationHandlers[message.method] else { return } - - // Convert notification parameters to concrete type and call handlers - for handler in handlers { - do { - try await handler(message) - } catch { - await logger?.error( - "Error handling notification", - metadata: [ - "method": "\(message.method)", - "error": "\(error)", - ]) - } - } - } - - /// Handle a response from the client (for server→client requests). - private func handleClientResponse(_ response: Response) async { - await logger?.trace( - "Processing client response", - metadata: ["id": "\(response.id)"]) - - // Check response routers first (e.g., for task-related responses) - for router in responseRouters { - switch response.result { - case .success(let value): - if await router.routeResponse(requestId: response.id, response: value) { - await logger?.trace( - "Response routed via router", - metadata: ["id": "\(response.id)"]) - return - } - case .failure(let error): - if await router.routeError(requestId: response.id, error: error) { - await logger?.trace( - "Error routed via router", - metadata: ["id": "\(response.id)"]) - return - } - } - } - - // Fall back to normal pending request handling - if let pendingRequest = pendingRequests.removeValue(forKey: response.id) { - switch response.result { - case .success(let value): - pendingRequest.resume(returning: value) - case .failure(let error): - pendingRequest.resume(throwing: error) - } - } else { - await logger?.warning( - "Received response for unknown request", - metadata: ["id": "\(response.id)"]) - } - } - - // MARK: - Server→Client Requests (Bidirectional Communication) - - /// Send a request to the client and wait for a response. - /// - /// This enables bidirectional communication where the server can request - /// information from the client (e.g., roots, sampling, elicitation). - /// - /// - Parameter request: The request to send - /// - Returns: The result from the client - public func sendRequest(_ request: Request) async throws -> M.Result { - guard let connection = connection else { - throw MCPError.internalError("Server connection not initialized") - } - - let encoder = JSONEncoder() - encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] - let requestData = try encoder.encode(request) - - // Create stream for receiving the response - let (stream, continuation) = AsyncThrowingStream.makeStream() - - // Clean up pending request if cancelled - let requestId = request.id - continuation.onTermination = { @Sendable [weak self] _ in - Task { await self?.cleanupPendingRequest(id: requestId) } - } - - // Register the pending request - pendingRequests[request.id] = AnyServerPendingRequest(continuation: continuation) - - // Send the request - do { - try await connection.send(requestData) - } catch { - pendingRequests.removeValue(forKey: request.id) - continuation.finish(throwing: error) - throw error - } - - // Wait for response - for try await result in stream { - return result - } - - throw MCPError.internalError("No response received from client") - } - - private func cleanupPendingRequest(id: RequestId) { - pendingRequests.removeValue(forKey: id) - } - - // MARK: - In-Flight Request Tracking (Protocol-Level Cancellation) - - /// Track an in-flight request handler Task. - private func trackInFlightRequest(_ requestId: RequestId, task: Task) { - inFlightHandlerTasks[requestId] = task - } - - /// Remove an in-flight request handler Task. - private func removeInFlightRequest(_ requestId: RequestId) { - inFlightHandlerTasks.removeValue(forKey: requestId) - } - - /// Cancel an in-flight request handler Task. - /// - /// Called when a CancelledNotification is received for a specific requestId. - /// Per MCP spec, if the request is unknown or already completed, this is a no-op. - private func cancelInFlightRequest(_ requestId: RequestId, reason: String?) async { - if let task = inFlightHandlerTasks[requestId] { - task.cancel() - await logger?.debug( - "Cancelled in-flight request", - metadata: [ - "id": "\(requestId)", - "reason": "\(reason ?? "none")", - ] - ) - } - // Per spec: MAY ignore if request is unknown - no error needed - } - - /// Generate a unique request ID for server→client requests. - private func generateRequestId() -> RequestId { - let id = nextRequestId - nextRequestId += 1 - return .number(id) - } - - /// Request the list of roots from the client. - /// - /// Roots represent filesystem directories that the client has access to. - /// Servers can use this to understand the scope of files they can work with. - /// - /// - Throws: MCPError if the client doesn't support roots or if the request fails. - /// - Returns: The list of roots from the client. - public func listRoots() async throws -> [Root] { - // Check that client supports roots - guard clientCapabilities?.roots != nil else { - throw MCPError.invalidRequest("Client does not support roots capability") - } - - let request: Request = ListRoots.request(id: generateRequestId()) - let result = try await sendRequest(request) - return result.roots - } - - /// Request a sampling completion from the client (without tools). - /// - /// This enables servers to request LLM completions through the client, - /// allowing sophisticated agentic behaviors while maintaining security. - /// - /// The result will be a single content block (text, image, or audio). - /// For tool-enabled sampling, use `createMessageWithTools(_:)` instead. - /// - /// - Parameter params: The sampling parameters including messages, model preferences, etc. - /// - Throws: MCPError if the client doesn't support sampling or if the request fails. - /// - Returns: The sampling result from the client containing a single content block. - public func createMessage(_ params: CreateSamplingMessage.Parameters) async throws -> CreateSamplingMessage.Result { - // Check that client supports sampling - guard clientCapabilities?.sampling != nil else { - throw MCPError.invalidRequest("Client does not support sampling capability") - } - - let request: Request = CreateSamplingMessage.request(id: generateRequestId(), params) - return try await sendRequest(request) - } - - /// Request a sampling completion from the client with tool support. - /// - /// This enables servers to request LLM completions that may involve tool use. - /// The result may contain tool use content, and content can be an array for parallel tool calls. - /// - /// - Parameter params: The sampling parameters including messages, tools, and model preferences. - /// - Throws: MCPError if the client doesn't support sampling or tool capabilities. - /// - Returns: The sampling result from the client, which may include tool use content. - public func createMessageWithTools(_ params: CreateSamplingMessageWithTools.Parameters) async throws -> CreateSamplingMessageWithTools.Result { - // Check that client supports sampling - guard clientCapabilities?.sampling != nil else { - throw MCPError.invalidRequest("Client does not support sampling capability") - } - - // Check tools capability - guard clientCapabilities?.sampling?.tools != nil else { - throw MCPError.invalidRequest("Client does not support sampling tools capability") - } - - // Validate tool_use/tool_result message structure per MCP specification - try Sampling.Message.validateToolUseResultMessages(params.messages) - - let request: Request = CreateSamplingMessageWithTools.request(id: generateRequestId(), params) - return try await sendRequest(request) - } - - /// Request user input via elicitation from the client. - /// - /// Elicitation allows servers to request structured input from users through - /// the client, either via forms or external URLs (e.g., OAuth flows). - /// - /// - Parameter params: The elicitation parameters. - /// - Throws: MCPError if the client doesn't support elicitation or if the request fails. - /// - Returns: The elicitation result from the client. - public func elicit(_ params: Elicit.Parameters) async throws -> Elicit.Result { - // Check that client supports elicitation - guard clientCapabilities?.elicitation != nil else { - throw MCPError.invalidRequest("Client does not support elicitation capability") - } - - // Check mode-specific capabilities - switch params { - case .form: - guard clientCapabilities?.elicitation?.form != nil else { - throw MCPError.invalidRequest("Client does not support form elicitation") - } - case .url: - guard clientCapabilities?.elicitation?.url != nil else { - throw MCPError.invalidRequest("Client does not support URL elicitation") - } - } - - let request: Request = Elicit.request(id: generateRequestId(), params) - let result = try await sendRequest(request) - - // TODO: Add elicitation response validation against the requestedSchema. - // TypeScript SDK uses JSON Schema validators (AJV, CfWorker) to validate - // elicitation responses against the requestedSchema. Python SDK uses Pydantic. - // The ideal solution is to use the same JSON Schema validator for both - // elicitation and tool validation, for spec compliance and consistency. - - return result - } - - private func checkInitialized() throws { - guard isInitialized else { - throw MCPError.invalidRequest("Server is not initialized") - } - } - - // MARK: - Client Task Polling (Server → Client) - - /// Get a task from the client. - /// - /// Internal method used by experimental server task features. - func getClientTask(taskId: String) async throws -> GetTask.Result { - guard clientCapabilities?.tasks != nil else { - throw MCPError.invalidRequest("Client does not support tasks capability") - } - - let request = GetTask.request(.init(taskId: taskId)) - return try await sendRequest(request) - } - - /// Get the result payload of a client task. - /// - /// Internal method used by experimental server task features. - func getClientTaskResult(taskId: String) async throws -> GetTaskPayload.Result { - guard clientCapabilities?.tasks != nil else { - throw MCPError.invalidRequest("Client does not support tasks capability") - } - - let request = GetTaskPayload.request(.init(taskId: taskId)) - return try await sendRequest(request) - } - - /// Get the task result decoded as a specific type. - /// - /// Internal method used by experimental server task features. - func getClientTaskResultAs(taskId: String, type: T.Type) async throws -> T { - let result = try await getClientTaskResult(taskId: taskId) - - // The result's extraFields contain the actual result payload - guard let extraFields = result.extraFields else { - throw MCPError.invalidParams("Task result has no payload") - } - - // Convert extraFields to the target type - let encoder = JSONEncoder() - let decoder = JSONDecoder() - - let jsonData = try encoder.encode(extraFields) - return try decoder.decode(T.self, from: jsonData) - } - - // MARK: - Task-Augmented Requests (Server → Client) - - /// Send a task-augmented elicitation request to the client. - /// - /// The client returns a `CreateTaskResult` instead of an `ElicitResult`. - /// Use client task polling to get the final result. - /// - /// Internal method used by experimental server task features. - func sendElicitAsTask(_ params: Elicit.Parameters) async throws -> CreateTaskResult { - // Check that client supports task-augmented elicitation - try requireTaskAugmentedElicitation(clientCapabilities) - - // Check mode-specific capabilities - switch params { - case .form: - guard clientCapabilities?.elicitation?.form != nil else { - throw MCPError.invalidRequest("Client does not support form elicitation") - } - case .url: - guard clientCapabilities?.elicitation?.url != nil else { - throw MCPError.invalidRequest("Client does not support URL elicitation") - } - } - - guard let connection else { - throw MCPError.internalError("Server connection not initialized") - } - - // Build the request - let encoder = JSONEncoder() - encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] - let request: Request = Elicit.request(id: generateRequestId(), params) - let requestData = try encoder.encode(request) - - // Create stream for receiving the response - let (stream, continuation) = AsyncThrowingStream.makeStream() - - let requestId = request.id - continuation.onTermination = { @Sendable [weak self] _ in - Task { await self?.cleanupPendingRequest(id: requestId) } - } - - // Register the pending request - pendingRequests[requestId] = AnyServerPendingRequest(continuation: continuation) - - // Send the request - do { - try await connection.send(requestData) - } catch { - pendingRequests.removeValue(forKey: requestId) - continuation.finish(throwing: error) - throw error - } - - // Wait for single result - for try await result in stream { - return result - } - - throw MCPError.internalError("No response received") - } - - /// Send a task-augmented sampling request to the client. - /// - /// The client returns a `CreateTaskResult` instead of a `CreateSamplingMessage.Result`. - /// Use client task polling to get the final result. - /// - /// Internal method used by experimental server task features. - func sendCreateMessageAsTask(_ params: CreateSamplingMessage.Parameters) async throws -> CreateTaskResult { - // Check that client supports task-augmented sampling - try requireTaskAugmentedSampling(clientCapabilities) - - guard clientCapabilities?.sampling != nil else { - throw MCPError.invalidRequest("Client does not support sampling capability") - } - - guard let connection else { - throw MCPError.internalError("Server connection not initialized") - } - - // Build the request - let encoder = JSONEncoder() - encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] - let request = CreateSamplingMessage.request(id: generateRequestId(), params) - let requestData = try encoder.encode(request) - - // Create stream for receiving the response - let (stream, continuation) = AsyncThrowingStream.makeStream() - - let requestId = request.id - continuation.onTermination = { @Sendable [weak self] _ in - Task { await self?.cleanupPendingRequest(id: requestId) } - } - - // Register the pending request - pendingRequests[requestId] = AnyServerPendingRequest(continuation: continuation) - - // Send the request - do { - try await connection.send(requestData) - } catch { - pendingRequests.removeValue(forKey: requestId) - continuation.finish(throwing: error) - throw error - } - - // Wait for single result - for try await result in stream { - return result - } - - throw MCPError.internalError("No response received") - } - - private func registerDefaultHandlers( + func registerDefaultHandlers( initializeHook: (@Sendable (Client.Info, Client.Capabilities) async throws -> Void)? ) { // Initialize @@ -1592,7 +812,7 @@ public actor Server { /// - Parameters: /// - level: The minimum log level to send. /// - sessionId: The session identifier, or `nil` for transports without sessions. - private func setLoggingLevel(_ level: LoggingLevel, forSession sessionId: String?) { + func setLoggingLevel(_ level: LoggingLevel, forSession sessionId: String?) { loggingLevels[sessionId] = level } @@ -1619,7 +839,7 @@ public actor Server { return level.isAtLeast(sessionLevel) } - private func setInitialState( + func setInitialState( clientInfo: Client.Info, clientCapabilities: Client.Capabilities, protocolVersion: String From c52e6f7f922548615b5edafbe05d6917928dd42d Mon Sep 17 00:00:00 2001 From: Anthony DePasquale Date: Sun, 4 Jan 2026 13:01:39 +0100 Subject: [PATCH 5/8] Add missing functionality and many fixes --- .../MCP/Client/Client+MessageHandling.swift | 38 +- Sources/MCP/Client/Client+Registration.swift | 37 +- Sources/MCP/Client/Client+Requests.swift | 55 ++ Sources/MCP/Client/Client.swift | 27 + .../MCP/Server/Server+ClientRequests.swift | 11 + .../MCP/Server/Server+RequestHandling.swift | 78 ++- Sources/MCP/Server/Server.swift | 207 ++++++- .../MCPTests/RequestHandlerContextTests.swift | 531 ++++++++++++++++++ Tests/MCPTests/ToolTests.swift | 5 +- 9 files changed, 977 insertions(+), 12 deletions(-) create mode 100644 Tests/MCPTests/RequestHandlerContextTests.swift diff --git a/Sources/MCP/Client/Client+MessageHandling.swift b/Sources/MCP/Client/Client+MessageHandling.swift index 2e57baea..9dee9db5 100644 --- a/Sources/MCP/Client/Client+MessageHandling.swift +++ b/Sources/MCP/Client/Client+MessageHandling.swift @@ -3,6 +3,24 @@ import Foundation extension Client { // MARK: - Message Handling + /// Extract `_meta` from request parameters if present. + /// + /// Since `AnyMethod.Parameters` is `Value`, we need to extract `_meta` manually. + private func extractMeta(from params: Value) -> RequestMeta? { + guard case .object(let dict) = params, + let metaValue = dict["_meta"] else { + return nil + } + // Decode the _meta value as RequestMeta + let encoder = JSONEncoder() + let decoder = JSONDecoder() + guard let data = try? encoder.encode(metaValue), + let meta = try? decoder.decode(RequestMeta.self, from: data) else { + return nil + } + return meta + } + func handleResponse(_ response: Response) async { await logger?.trace( "Processing response", @@ -255,9 +273,27 @@ extension Client { return } + // Create the request handler context + // This provides cancellation checking and notification sending to the handler + let requestMeta = extractMeta(from: request.params) + let context = RequestHandlerContext( + sendNotification: { [weak self] notification in + guard let self else { + throw MCPError.internalError("Client was deallocated") + } + guard let connection = await self.connection else { + throw MCPError.internalError("Cannot send notification - client not connected") + } + let notificationData = try self.encoder.encode(notification) + try await connection.send(notificationData) + }, + requestId: request.id, + _meta: requestMeta + ) + // Execute the handler and send response do { - let response = try await handler(request) + let response = try await handler(request, context: context) // Check cancellation before sending response (per MCP spec: // "Receivers of a cancellation notification SHOULD... Not send a response diff --git a/Sources/MCP/Client/Client+Registration.swift b/Sources/MCP/Client/Client+Registration.swift index bab2886e..2a5b4173 100644 --- a/Sources/MCP/Client/Client+Registration.swift +++ b/Sources/MCP/Client/Client+Registration.swift @@ -80,14 +80,32 @@ extension Client { /// This enables bidirectional communication where the server can send requests /// to the client (e.g., sampling, roots, elicitation). /// + /// The handler receives a `RequestHandlerContext` that provides: + /// - `isCancelled` and `checkCancellation()` for responding to cancellation + /// - `sendProgressNotification()` for reporting progress back to the server + /// + /// ## Example + /// + /// ```swift + /// client.withRequestHandler(CreateSamplingMessage.self) { params, context in + /// // Check for cancellation during long operations + /// try context.checkCancellation() + /// + /// // Process the request + /// let result = try await processRequest(params) + /// + /// return result + /// } + /// ``` + /// /// - Parameters: /// - type: The method type to handle - /// - handler: The handler function that receives parameters and returns a result + /// - handler: The handler function that receives parameters and context, returns a result /// - Returns: Self for chaining @discardableResult public func withRequestHandler( _ type: M.Type, - handler: @escaping @Sendable (M.Parameters) async throws -> M.Result + handler: @escaping @Sendable (M.Parameters, RequestHandlerContext) async throws -> M.Result ) -> Self { requestHandlers[M.name] = TypedClientRequestHandler(handler) return self @@ -111,7 +129,7 @@ extension Client { capabilities.roots != nil, "Cannot register roots handler: Client does not have roots capability" ) - return withRequestHandler(ListRoots.self) { _ in + return withRequestHandler(ListRoots.self) { _, _ in ListRoots.Result(roots: try await handler()) } } @@ -129,7 +147,10 @@ extension Client { /// ## Example /// /// ```swift - /// client.withSamplingHandler { params in + /// client.withSamplingHandler { params, context in + /// // Check for cancellation during long operations + /// try context.checkCancellation() + /// /// // Call your LLM with the messages /// let response = try await llm.complete( /// messages: params.messages, @@ -146,12 +167,12 @@ extension Client { /// } /// ``` /// - /// - Parameter handler: A closure that receives sampling parameters and returns the result. + /// - Parameter handler: A closure that receives sampling parameters and context, returns the result. /// - Returns: Self for chaining. /// - Precondition: `capabilities.sampling` must be non-nil. @discardableResult public func withSamplingHandler( - _ handler: @escaping @Sendable (ClientSamplingRequest.Parameters) async throws -> ClientSamplingRequest.Result + _ handler: @escaping @Sendable (ClientSamplingRequest.Parameters, RequestHandlerContext) async throws -> ClientSamplingRequest.Result ) -> Self { precondition( capabilities.sampling != nil, @@ -167,12 +188,12 @@ extension Client { /// /// - Important: The client must have declared `elicitation` capability during initialization. /// - /// - Parameter handler: A closure that receives elicitation parameters and returns the result. + /// - Parameter handler: A closure that receives elicitation parameters and context, returns the result. /// - Returns: Self for chaining. /// - Precondition: `capabilities.elicitation` must be non-nil. @discardableResult public func withElicitationHandler( - _ handler: @escaping @Sendable (Elicit.Parameters) async throws -> Elicit.Result + _ handler: @escaping @Sendable (Elicit.Parameters, RequestHandlerContext) async throws -> Elicit.Result ) -> Self { precondition( capabilities.elicitation != nil, diff --git a/Sources/MCP/Client/Client+Requests.swift b/Sources/MCP/Client/Client+Requests.swift index e96c3d2f..61b369bd 100644 --- a/Sources/MCP/Client/Client+Requests.swift +++ b/Sources/MCP/Client/Client+Requests.swift @@ -425,6 +425,61 @@ extension Client { pendingRequests.removeValue(forKey: id) } + // MARK: - Request Cancellation + + /// Cancel an in-flight request by its ID. + /// + /// This method cancels a pending request and sends a `CancelledNotification` to the server. + /// Use this when you need to cancel a request that was sent earlier but hasn't completed yet. + /// + /// Per MCP spec: "When a party wants to cancel an in-progress request, it sends a + /// `notifications/cancelled` notification containing the ID of the request to cancel." + /// + /// ## Example + /// + /// ```swift + /// // Create a request with a known ID + /// let requestId = RequestId.string("my-request-123") + /// let request = CallTool.request(id: requestId, .init(name: "slow_operation")) + /// + /// // Start the request in a separate task + /// Task { + /// do { + /// let result = try await client.send(request) + /// print("Result: \(result)") + /// } catch let error as MCPError where error.code == MCPError.Code.requestCancelled { + /// print("Request was cancelled") + /// } + /// } + /// + /// // Later, cancel it by ID + /// try await client.cancelRequest(requestId, reason: "User cancelled") + /// ``` + /// + /// - Parameters: + /// - id: The ID of the request to cancel. This must match the ID used when sending the request. + /// - reason: An optional human-readable reason for the cancellation, for logging/debugging. + /// - Throws: This method does not throw. Cancellation notifications are best-effort per the spec. + /// + /// - Note: If the request has already completed or is unknown, this is a no-op per the MCP spec. + /// - Note: The `initialize` request MUST NOT be cancelled per the MCP spec. + /// - Important: For task-augmented requests, use the `tasks/cancel` method instead. + public func cancelRequest(_ id: RequestId, reason: String? = nil) async { + // Remove and finish the pending request with cancellation error + if let pendingRequest = removePendingRequest(id: id) { + pendingRequest.resume(throwing: MCPError.requestCancelled(reason: reason)) + } + + // Clean up any progress-related state + if let progressToken = requestProgressTokens.removeValue(forKey: id) { + progressCallbacks.removeValue(forKey: progressToken) + timeoutControllers.removeValue(forKey: progressToken) + } + + // Send cancellation notification to server (best-effort) + await sendCancellationNotification(requestId: id, reason: reason) + } + /// Send a CancelledNotification to the server for a cancelled request. /// /// Per MCP spec: "When a party wants to cancel an in-progress request, it sends diff --git a/Sources/MCP/Client/Client.swift b/Sources/MCP/Client/Client.swift index c65f507a..a857224e 100644 --- a/Sources/MCP/Client/Client.swift +++ b/Sources/MCP/Client/Client.swift @@ -188,6 +188,33 @@ public actor Client { /// Use this to send notifications from within a request handler. let sendNotification: @Sendable (any NotificationMessageProtocol) async throws -> Void + /// The JSON-RPC ID of the request being handled. + /// + /// This can be useful for tracking, logging, or correlating messages. + /// It matches the TypeScript SDK's `extra.requestId`. + public let requestId: RequestId + + /// The request metadata from the `_meta` field, if present. + /// + /// Contains metadata like the progress token for progress notifications. + /// This matches the TypeScript SDK's `extra._meta` and Python's `ctx.meta`. + /// + /// ## Example + /// + /// ```swift + /// client.withRequestHandler(CreateSamplingMessage.self) { params, context in + /// if let progressToken = context._meta?.progressToken { + /// try await context.sendProgressNotification( + /// token: progressToken, + /// progress: 50, + /// total: 100 + /// ) + /// } + /// return result + /// } + /// ``` + public let _meta: RequestMeta? + // MARK: - Convenience Methods /// Send a progress notification to the server. diff --git a/Sources/MCP/Server/Server+ClientRequests.swift b/Sources/MCP/Server/Server+ClientRequests.swift index 595c7acd..caf235b3 100644 --- a/Sources/MCP/Server/Server+ClientRequests.swift +++ b/Sources/MCP/Server/Server+ClientRequests.swift @@ -52,6 +52,17 @@ extension Server { pendingRequests.removeValue(forKey: id) } + /// Register a pending request from a context's sendRequest call. + /// + /// This is used by RequestHandlerContext.sendRequest to register pending + /// requests that will be fulfilled when the client responds. + func registerContextRequest( + id: RequestId, + continuation: AsyncThrowingStream.Continuation + ) { + pendingRequests[id] = AnyServerPendingRequest(continuation: continuation) + } + // MARK: - In-Flight Request Tracking (Protocol-Level Cancellation) /// Track an in-flight request handler Task. diff --git a/Sources/MCP/Server/Server+RequestHandling.swift b/Sources/MCP/Server/Server+RequestHandling.swift index 3d03a96c..2be03fc9 100644 --- a/Sources/MCP/Server/Server+RequestHandling.swift +++ b/Sources/MCP/Server/Server+RequestHandling.swift @@ -102,6 +102,28 @@ extension Server { /// For HTTP transports with multiple concurrent clients, this identifies /// the specific session. Used for per-session features like log levels. let sessionId: String? + /// The request metadata from `_meta` field, if present. + /// + /// Contains the progress token and any additional metadata. + let meta: RequestMeta? + } + + /// Extract `_meta` from request parameters if present. + /// + /// Since `AnyMethod.Parameters` is `Value`, we need to extract `_meta` manually. + private func extractMeta(from params: Value) -> RequestMeta? { + guard case .object(let dict) = params, + let metaValue = dict["_meta"] else { + return nil + } + // Decode the _meta value as RequestMeta + let encoder = JSONEncoder() + let decoder = JSONDecoder() + guard let data = try? encoder.encode(metaValue), + let meta = try? decoder.decode(RequestMeta.self, from: data) else { + return nil + } + return meta } /// Wrapper for encoding type-erased notifications as JSON-RPC messages. @@ -161,10 +183,12 @@ extension Server { // This ensures responses go to the correct client even if self.connection // changes while the handler is executing (e.g., another client connects). let capturedConnection = self.connection + let requestMeta = extractMeta(from: request.params) let context = RequestContext( capturedConnection: capturedConnection, requestId: request.id, - sessionId: await capturedConnection?.sessionId + sessionId: await capturedConnection?.sessionId, + meta: requestMeta ) // Check if this is a pre-processed error request (empty method) @@ -244,9 +268,61 @@ extension Server { try await connection.send(data, relatedRequestId: context.requestId) }, sessionId: context.sessionId, + requestId: context.requestId, + _meta: context.meta, shouldSendLogMessage: { [weak self, context] level in guard let self else { return true } return await self.shouldSendLogMessage(at: level, forSession: context.sessionId) + }, + sendRequest: { [weak self, context] requestData in + guard let self else { + throw MCPError.internalError("Server reference lost") + } + guard let connection = context.capturedConnection else { + throw MCPError.internalError("Cannot send request - connection was nil at request time") + } + + // Parse the request to get its ID + guard let jsonObject = try? JSONSerialization.jsonObject(with: requestData) as? [String: Any], + let requestId = jsonObject["id"] else { + throw MCPError.invalidParams("Could not parse request ID") + } + + // Convert request ID to RequestId type + let typedRequestId: RequestId + if let numId = requestId as? Int { + typedRequestId = .number(numId) + } else if let strId = requestId as? String { + typedRequestId = .string(strId) + } else { + throw MCPError.invalidParams("Invalid request ID type") + } + + // Create stream for receiving the response + let (stream, continuation) = AsyncThrowingStream.makeStream() + + continuation.onTermination = { @Sendable [weak self] _ in + Task { await self?.cleanUpPendingRequest(id: typedRequestId) } + } + + // Register the pending request + await self.registerContextRequest(id: typedRequestId, continuation: continuation) + + // Send the request using captured connection + do { + try await connection.send(requestData, relatedRequestId: context.requestId) + } catch { + await self.cleanUpPendingRequest(id: typedRequestId) + continuation.finish(throwing: error) + throw error + } + + // Wait for response + for try await result in stream { + return result + } + + throw MCPError.internalError("No response received from client") } ) diff --git a/Sources/MCP/Server/Server.swift b/Sources/MCP/Server/Server.swift index 773ac6e9..b5bb9cf4 100644 --- a/Sources/MCP/Server/Server.swift +++ b/Sources/MCP/Server/Server.swift @@ -5,7 +5,75 @@ import struct Foundation.Date import class Foundation.JSONDecoder import class Foundation.JSONEncoder -/// Model Context Protocol server +/// Model Context Protocol server. +/// +/// ## Architecture: One Server per Client +/// +/// The Swift SDK uses a **one-Server-per-client** architecture, where each client +/// connection gets its own `Server` instance. This mirrors the TypeScript SDK's +/// design and differs from Python's shared-Server model. +/// +/// ### Comparison with Other SDKs +/// +/// **Python SDK (shared Server):** +/// ``` +/// ┌──────────────────────────────────────┐ +/// │ Server (ONE) │ +/// │ - Handler registry (shared) │ +/// │ - No connection state │ +/// └──────────────────────────────────────┘ +/// │ server.run() creates ↓ +/// ┌─────────────┐ ┌─────────────┐ +/// │ Session A │ │ Session B │ +/// │ (Transport) │ │ (Transport) │ +/// └─────────────┘ └─────────────┘ +/// ``` +/// +/// **Swift & TypeScript SDKs (per-client Server):** +/// ``` +/// ┌─────────────┐ ┌─────────────┐ +/// │ Server A │ │ Server B │ +/// │ (Handlers) │ │ (Handlers) │ +/// │ (Transport) │ │ (Transport) │ +/// └─────────────┘ └─────────────┘ +/// ``` +/// +/// ### Scalability Considerations +/// +/// The per-client model is appropriate for MCP's typical use cases: +/// - AI assistants connecting to tool servers (single-digit connections) +/// - IDE plugins and developer tools (tens of connections) +/// - Multi-user applications (hundreds of connections) +/// +/// Memory overhead per Server instance is minimal (a few KB for handler references +/// and state). For realistic MCP deployments, this scales well. +/// +/// For high-connection scenarios (10,000+), consider: +/// - Horizontal scaling with connection-time load balancing +/// - MCP's stateless mode for true per-request distribution +/// - The Python SDK's shared-Server pattern (requires architectural changes) +/// +/// ### Design Rationale +/// +/// The per-client model was chosen because it: +/// 1. Matches TypeScript SDK's official examples and patterns +/// 2. Provides complete isolation between client connections +/// 3. Simplifies reasoning about connection state +/// 4. Avoids complex session management code +/// +/// For HTTP transports, each session creates its own `(Server, HTTPServerTransport)` +/// pair, stored by session ID for request routing. +/// +/// ## API Design: Context vs Server Methods +/// +/// The `RequestHandlerContext` provides request-scoped capabilities: +/// - `requestId`, `_meta` - Request identification and metadata +/// - `sendNotification()` - Send notifications during handling +/// - `elicit()`, `elicitUrl()` - Request user input (matches Python's `ctx.elicit()`) +/// - `isCancelled` - Check for request cancellation +/// +/// Sampling is done via `server.createMessage()` (matches TypeScript), not through +/// the context. This design follows each reference SDK's conventions where appropriate. public actor Server { /// The server configuration public struct Configuration: Hashable, Codable, Sendable { @@ -207,12 +275,63 @@ public actor Server { /// For simple transports (stdio, single-connection), this is `nil`. public let sessionId: String? + /// The JSON-RPC ID of the request being handled. + /// + /// This can be useful for tracking, logging, or correlating messages. + /// It matches the TypeScript SDK's `extra.requestId`. + public let requestId: RequestId + + /// The request metadata from the `_meta` field, if present. + /// + /// Contains metadata like the progress token for progress notifications. + /// This matches the TypeScript SDK's `extra._meta` and Python's `ctx.meta`. + /// + /// ## Example + /// + /// ```swift + /// server.withRequestHandler(CallTool.self) { request, context in + /// if let progressToken = context._meta?.progressToken { + /// try await context.sendProgress(token: progressToken, progress: 50, total: 100) + /// } + /// return CallTool.Result(content: [.text("Done")]) + /// } + /// ``` + public let _meta: RequestMeta? + /// Check if a log message at the given level should be sent. /// /// This respects the minimum log level set by the client via `logging/setLevel`. /// Messages below the threshold will be silently dropped. let shouldSendLogMessage: @Sendable (LoggingLevel) async -> Bool + /// Send a request to the client and wait for a response. + /// + /// This enables bidirectional communication from within a request handler, + /// allowing servers to request information from the client (e.g., elicitation, + /// sampling) during request processing. + /// + /// This matches the TypeScript SDK's `extra.sendRequest()` functionality. + /// + /// ## Example + /// + /// ```swift + /// server.withRequestHandler(CallTool.self) { request, context in + /// // Request user input via elicitation + /// let result = try await context.elicit( + /// message: "Please confirm the operation", + /// requestedSchema: ElicitationSchema(properties: [ + /// "confirm": .boolean(BooleanSchema(title: "Confirm")) + /// ]) + /// ) + /// + /// if result.action == .accept { + /// // Process confirmed action + /// } + /// return CallTool.Result(content: [.text("Done")]) + /// } + /// ``` + let sendRequest: @Sendable (Data) async throws -> Value + // MARK: - Convenience Methods /// Send a progress notification to the client. @@ -326,6 +445,92 @@ public actor Server { try await sendMessage(TaskStatusNotification.message(.init(task: task))) } + // MARK: - Bidirectional Requests + + /// Request user input via form elicitation from the client. + /// + /// This enables servers to request structured input from users through + /// the client during request handling. The client presents a form based + /// on the provided schema and returns the user's response. + /// + /// This matches the TypeScript SDK's `extra.sendRequest({ method: 'elicitation/create', ... })` + /// and Python's `ctx.elicit()` functionality. + /// + /// ## Example + /// + /// ```swift + /// server.withRequestHandler(CallTool.self) { request, context in + /// let result = try await context.elicit( + /// message: "Please confirm the operation", + /// requestedSchema: ElicitationSchema(properties: [ + /// "confirm": .boolean(BooleanSchema(title: "Confirm")) + /// ]) + /// ) + /// + /// if result.action == .accept { + /// // User confirmed + /// } + /// return CallTool.Result(content: [.text("Done")]) + /// } + /// ``` + /// + /// - Parameters: + /// - message: The message to present to the user + /// - requestedSchema: The schema defining the form fields + /// - Returns: The elicitation result from the client + /// - Throws: MCPError if the request fails + public func elicit( + message: String, + requestedSchema: ElicitationSchema + ) async throws -> ElicitResult { + let params = ElicitRequestFormParams( + mode: "form", + message: message, + requestedSchema: requestedSchema + ) + let request = Elicit.request(id: .random, .form(params)) + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + let requestData = try encoder.encode(request) + + let responseValue = try await sendRequest(requestData) + let decoder = JSONDecoder() + let responseData = try encoder.encode(responseValue) + return try decoder.decode(ElicitResult.self, from: responseData) + } + + /// Request user interaction via URL-mode elicitation from the client. + /// + /// This enables servers to request out-of-band interactions through + /// external URLs (e.g., OAuth flows, credential collection). + /// + /// - Parameters: + /// - message: Human-readable explanation of why the interaction is needed + /// - url: The URL the user should navigate to + /// - elicitationId: Unique identifier for tracking this elicitation + /// - Returns: The elicitation result from the client + /// - Throws: MCPError if the request fails + public func elicitUrl( + message: String, + url: String, + elicitationId: String + ) async throws -> ElicitResult { + let params = ElicitRequestURLParams( + message: message, + elicitationId: elicitationId, + url: url + ) + let request = Elicit.request(id: .random, .url(params)) + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + let requestData = try encoder.encode(request) + + let responseValue = try await sendRequest(requestData) + let decoder = JSONDecoder() + let responseData = try encoder.encode(responseValue) + return try decoder.decode(ElicitResult.self, from: responseData) + } + // MARK: - Cancellation Checking /// Whether the request has been cancelled. diff --git a/Tests/MCPTests/RequestHandlerContextTests.swift b/Tests/MCPTests/RequestHandlerContextTests.swift new file mode 100644 index 00000000..aae258c4 --- /dev/null +++ b/Tests/MCPTests/RequestHandlerContextTests.swift @@ -0,0 +1,531 @@ +import Foundation +import Testing + +@testable import MCP + +/// Tests for RequestHandlerContext functionality. +/// +/// These tests verify that handlers have access to request context information +/// and can make bidirectional requests, matching the TypeScript SDK's +/// `RequestHandlerExtra` and Python SDK's `RequestContext` / `Context`. +/// +/// Based on: +/// - TypeScript: `packages/core/test/shared/protocol.test.ts` +/// - TypeScript: `test/integration/test/taskLifecycle.test.ts` +/// - Python: `tests/server/fastmcp/test_server.py` (test_context_injection) +/// - Python: `tests/server/fastmcp/test_elicitation.py` +/// - Python: `tests/issues/test_176_progress_token.py` + +// MARK: - Server RequestHandlerContext Tests + +@Suite("Server.RequestHandlerContext Tests") +struct ServerRequestHandlerContextTests { + + // MARK: - requestId Tests + + /// Test that handlers can access context.requestId. + /// Based on Python SDK's test_context_injection: `assert ctx.request_id is not None` + @Test("Handler can access context.requestId") + func testHandlerCanAccessRequestId() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + // Track the requestId received in handler + actor RequestIdTracker { + var receivedRequestId: RequestId? + func set(_ id: RequestId) { receivedRequestId = id } + } + let tracker = RequestIdTracker() + + let server = Server( + name: "TestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "test_tool", description: "Test", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, context in + // Handler accesses context.requestId - this is what we're testing + await tracker.set(context.requestId) + return CallTool.Result(content: [.text("Request ID: \(context.requestId)")]) + } + + let client = Client(name: "TestClient", version: "1.0.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "test_tool", arguments: [:]) + + // Verify handler received a valid requestId + let receivedId = await tracker.receivedRequestId + #expect(receivedId != nil, "Handler should have access to requestId") + + // Verify the response mentions the request ID + if case .text(let text, _, _) = result.content.first { + #expect(text.contains("Request ID:"), "Response should contain request ID") + } + + await client.disconnect() + } + + /// Test that context.requestId matches the actual JSON-RPC request ID. + @Test("context.requestId matches JSON-RPC request ID") + func testRequestIdMatchesJsonRpcId() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + actor RequestIdTracker { + var receivedRequestIds: [RequestId] = [] + func add(_ id: RequestId) { receivedRequestIds.append(id) } + } + let tracker = RequestIdTracker() + + let server = Server( + name: "TestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, context in + await tracker.add(context.requestId) + return ListTools.Result(tools: [ + Tool(name: "test_tool", description: "Test", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { _, context in + await tracker.add(context.requestId) + return CallTool.Result(content: [.text("OK")]) + } + + let client = Client(name: "TestClient", version: "1.0.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Make multiple requests + _ = try await client.send(ListTools.request()) + _ = try await client.callTool(name: "test_tool", arguments: [:]) + _ = try await client.callTool(name: "test_tool", arguments: [:]) + + let receivedIds = await tracker.receivedRequestIds + #expect(receivedIds.count == 3, "Should have received 3 request IDs") + + // Verify all IDs are unique (each request gets a unique ID) + let uniqueIds = Set(receivedIds.map { "\($0)" }) + #expect(uniqueIds.count == 3, "Each request should have a unique ID") + + await client.disconnect() + } + + // MARK: - _meta Tests + + /// Test that handlers can access context._meta when present. + /// Based on TypeScript SDK's `extra._meta` access tests. + @Test("Handler can access context._meta when present") + func testHandlerCanAccessMeta() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + actor MetaTracker { + var receivedMeta: RequestMeta? + func set(_ meta: RequestMeta?) { receivedMeta = meta } + } + let tracker = MetaTracker() + + let server = Server( + name: "TestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "test_tool", description: "Test", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { _, context in + // Handler accesses context._meta - this is what we're testing + await tracker.set(context._meta) + return CallTool.Result(content: [.text("OK")]) + } + + let client = Client(name: "TestClient", version: "1.0.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Call tool WITH _meta containing progressToken + _ = try await client.send( + CallTool.request(.init( + name: "test_tool", + arguments: [:], + _meta: RequestMeta(progressToken: .string("test-token-123")) + )) + ) + + let receivedMeta = await tracker.receivedMeta + #expect(receivedMeta != nil, "Handler should have access to _meta") + #expect(receivedMeta?.progressToken == .string("test-token-123"), "progressToken should match") + + await client.disconnect() + } + + /// Test that context._meta is nil when not provided in request. + @Test("context._meta is nil when not provided") + func testMetaIsNilWhenNotProvided() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + actor MetaTracker { + var receivedMeta: RequestMeta? = RequestMeta() // Initialize to non-nil + var wasSet = false + func set(_ meta: RequestMeta?) { + receivedMeta = meta + wasSet = true + } + } + let tracker = MetaTracker() + + let server = Server( + name: "TestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "test_tool", description: "Test", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { _, context in + await tracker.set(context._meta) + return CallTool.Result(content: [.text("OK")]) + } + + let client = Client(name: "TestClient", version: "1.0.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Call tool WITHOUT _meta + _ = try await client.callTool(name: "test_tool", arguments: [:]) + + let wasSet = await tracker.wasSet + let receivedMeta = await tracker.receivedMeta + #expect(wasSet, "Handler should have been called") + #expect(receivedMeta == nil, "context._meta should be nil when not provided") + + await client.disconnect() + } + + /// Test using context._meta?.progressToken as a convenience pattern. + /// Based on Python SDK's test_176_progress_token.py showing progressToken access via context. + @Test("context._meta?.progressToken convenience pattern") + func testProgressTokenFromContextMeta() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + actor ProgressTracker { + var updates: [(token: ProgressToken, progress: Double)] = [] + func add(token: ProgressToken, progress: Double) { + updates.append((token, progress)) + } + } + let progressTracker = ProgressTracker() + + let server = Server( + name: "TestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "progress_tool", description: "Test", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { _, context in + // Use context._meta?.progressToken instead of request._meta?.progressToken + // This is the convenience pattern we're testing + if let progressToken = context._meta?.progressToken { + try await context.sendProgress(token: progressToken, progress: 0.5, total: 1.0) + try await context.sendProgress(token: progressToken, progress: 1.0, total: 1.0) + } + return CallTool.Result(content: [.text("Done")]) + } + + let client = Client(name: "TestClient", version: "1.0.0") + + await client.onNotification(ProgressNotification.self) { message in + await progressTracker.add(token: message.params.progressToken, progress: message.params.progress) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Call tool with progressToken in _meta + _ = try await client.send( + CallTool.request(.init( + name: "progress_tool", + arguments: [:], + _meta: RequestMeta(progressToken: .string("ctx-token")) + )) + ) + + try await Task.sleep(for: .milliseconds(100)) + + let updates = await progressTracker.updates + #expect(updates.count == 2, "Should receive 2 progress notifications") + #expect(updates.allSatisfy { $0.token == .string("ctx-token") }, "All tokens should match") + + await client.disconnect() + } + + // MARK: - context.elicit() Tests + + /// Test that handlers can use context.elicit() for bidirectional elicitation. + /// Based on TypeScript SDK's `extra.sendRequest({ method: 'elicitation/create' })` tests + /// in test/integration/test/taskLifecycle.test.ts and test/integration/test/server.test.ts. + @Test("Handler can use context.elicit() for form elicitation") + func testContextElicitForm() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "TestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "askName", description: "Ask user for name", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { _, context in + // Use context.elicit() instead of server.elicit() + // This is the bidirectional request pattern from TypeScript's extra.sendRequest() + let result = try await context.elicit( + message: "What is your name?", + requestedSchema: ElicitationSchema( + properties: ["name": .string(StringSchema(title: "Name"))], + required: ["name"] + ) + ) + + if result.action == .accept, let name = result.content?["name"]?.stringValue { + return CallTool.Result(content: [.text("Hello, \(name)!")]) + } else { + return CallTool.Result(content: [.text("No name provided")], isError: true) + } + } + + let client = Client(name: "TestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(form: Client.Capabilities.Elicitation.Form()) + )) + + await client.withElicitationHandler { params, _ in + guard case .form(let formParams) = params else { + return ElicitResult(action: .decline) + } + #expect(formParams.message == "What is your name?") + return ElicitResult(action: .accept, content: ["name": .string("Bob")]) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "askName", arguments: [:]) + + #expect(result.isError == nil) + if case .text(let text, _, _) = result.content.first { + #expect(text == "Hello, Bob!") + } else { + Issue.record("Expected text content") + } + + await client.disconnect() + } + + /// Test that context.elicit() handles user decline. + @Test("context.elicit() handles user decline") + func testContextElicitDecline() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "TestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "confirm", description: "Confirm", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { _, context in + let result = try await context.elicit( + message: "Confirm?", + requestedSchema: ElicitationSchema( + properties: ["ok": .boolean(BooleanSchema(title: "OK"))] + ) + ) + + return switch result.action { + case .accept: CallTool.Result(content: [.text("Accepted")]) + case .decline: CallTool.Result(content: [.text("Declined")]) + case .cancel: CallTool.Result(content: [.text("Cancelled")]) + } + } + + let client = Client(name: "TestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(form: Client.Capabilities.Elicitation.Form()) + )) + + await client.withElicitationHandler { _, _ in + return ElicitResult(action: .decline) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "confirm", arguments: [:]) + + if case .text(let text, _, _) = result.content.first { + #expect(text == "Declined") + } else { + Issue.record("Expected text content") + } + + await client.disconnect() + } + + // MARK: - Sampling from Handlers + // Note: For sampling from within request handlers, use server.createMessage() which is + // thoroughly tested in SamplingTests.swift. The context provides elicit() and elicitUrl() + // convenience methods (tested above), matching Python's ctx.elicit() pattern. Sampling + // is done via the server directly, matching TypeScript's pattern where extra.sendRequest() + // is generic and server.createMessage() is the convenience method. +} + +// MARK: - Client RequestHandlerContext Tests + +@Suite("Client.RequestHandlerContext Tests") +struct ClientRequestHandlerContextTests { + + /// Test that client handlers can access context.requestId. + @Test("Client handler can access context.requestId") + func testClientHandlerCanAccessRequestId() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + actor RequestIdTracker { + var receivedRequestId: RequestId? + func set(_ id: RequestId) { receivedRequestId = id } + } + let tracker = RequestIdTracker() + + let server = Server( + name: "TestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "elicitTool", description: "Elicit", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] _, _ in + let result = try await server.elicit(.form(ElicitRequestFormParams( + message: "Test", + requestedSchema: ElicitationSchema(properties: ["x": .string(StringSchema())]) + ))) + return CallTool.Result(content: [.text("Action: \(result.action)")]) + } + + let client = Client(name: "TestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(form: Client.Capabilities.Elicitation.Form()) + )) + + await client.withElicitationHandler { _, context in + // Client handler accesses context.requestId + await tracker.set(context.requestId) + return ElicitResult(action: .accept, content: ["x": .string("test")]) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + _ = try await client.callTool(name: "elicitTool", arguments: [:]) + + let receivedId = await tracker.receivedRequestId + #expect(receivedId != nil, "Client handler should have access to requestId") + + await client.disconnect() + } + + /// Test that client handlers can access context._meta when present. + @Test("Client handler can access context._meta") + func testClientHandlerCanAccessMeta() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + actor MetaTracker { + var receivedMeta: RequestMeta? + func set(_ meta: RequestMeta?) { receivedMeta = meta } + } + let tracker = MetaTracker() + + let server = Server( + name: "TestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "elicitTool", description: "Elicit", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] _, _ in + let result = try await server.elicit(.form(ElicitRequestFormParams( + message: "Test", + requestedSchema: ElicitationSchema(properties: ["x": .string(StringSchema())]), + _meta: RequestMeta(progressToken: .string("server-token")) + ))) + return CallTool.Result(content: [.text("Action: \(result.action)")]) + } + + let client = Client(name: "TestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(form: Client.Capabilities.Elicitation.Form()) + )) + + await client.withElicitationHandler { _, context in + // Client handler accesses context._meta + await tracker.set(context._meta) + return ElicitResult(action: .accept, content: ["x": .string("test")]) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + _ = try await client.callTool(name: "elicitTool", arguments: [:]) + + let receivedMeta = await tracker.receivedMeta + #expect(receivedMeta != nil, "Client handler should have access to _meta") + #expect(receivedMeta?.progressToken == .string("server-token"), "progressToken should match") + + await client.disconnect() + } +} diff --git a/Tests/MCPTests/ToolTests.swift b/Tests/MCPTests/ToolTests.swift index ba9945df..618a821f 100644 --- a/Tests/MCPTests/ToolTests.swift +++ b/Tests/MCPTests/ToolTests.swift @@ -411,7 +411,10 @@ struct ToolTests { sendMessage: { _ in }, sendData: { _ in }, sessionId: nil, - shouldSendLogMessage: { _ in true } + requestId: .number(1), + _meta: nil, + shouldSendLogMessage: { _ in true }, + sendRequest: { _ in throw MCPError.internalError("Not implemented") } ) let response = try await handler(anyRequest, context: dummyContext) From 83a7af7d438ffeba7707bc58ea5b96e8f66f3595 Mon Sep 17 00:00:00 2001 From: Anthony DePasquale Date: Sun, 4 Jan 2026 14:55:06 +0100 Subject: [PATCH 6/8] Add per-message context to Transport layer --- Sources/MCP/Base/Transport.swift | 86 +++++++++++++- .../Base/Transports/HTTPClientTransport.swift | 20 ++-- .../HTTPServerTransport+Types.swift | 65 +++++++++++ .../Base/Transports/HTTPServerTransport.swift | 64 ++++++++--- .../Base/Transports/InMemoryTransport.swift | 12 +- .../Base/Transports/NetworkTransport.swift | 14 +-- .../MCP/Base/Transports/StdioTransport.swift | 12 +- Sources/MCP/Client/Client.swift | 6 +- .../MCP/Server/Server+RequestHandling.swift | 34 +++++- Sources/MCP/Server/Server.swift | 69 +++++++++++- Tests/MCPTests/HTTPClientTransportTests.swift | 30 ++--- Tests/MCPTests/Helpers/MockTransport.swift | 10 +- Tests/MCPTests/InMemoryTransportTests.swift | 24 ++-- Tests/MCPTests/NetworkTransportTests.swift | 10 +- .../MCPTests/RequestHandlerContextTests.swift | 105 ++++++++++++++++++ Tests/MCPTests/StdioTransportTests.swift | 32 +++--- Tests/MCPTests/ToolTests.swift | 3 + Tests/MCPTests/TransportSwitchingTests.swift | 10 +- 18 files changed, 491 insertions(+), 115 deletions(-) diff --git a/Sources/MCP/Base/Transport.swift b/Sources/MCP/Base/Transport.swift index 2c9057a2..0a3054cf 100644 --- a/Sources/MCP/Base/Transport.swift +++ b/Sources/MCP/Base/Transport.swift @@ -2,6 +2,81 @@ import Logging import struct Foundation.Data +// MARK: - Message Context Types + +/// Context information associated with a received message. +/// +/// This is the Swift equivalent of TypeScript's `MessageExtraInfo`, which is passed +/// via `onmessage(message, extra)`. It carries per-message context like authentication +/// info and SSE stream management callbacks. +/// +/// For simple transports (stdio, in-memory), context is typically `nil`. +/// For HTTP transports, context includes authentication info and SSE controls. +public struct MessageContext: Sendable { + /// Authentication information for this message's request. + /// + /// Contains validated access token information when using HTTP transports + /// with OAuth or other token-based authentication. Request handlers can + /// access this via `context.authInfo`. + public let authInfo: AuthInfo? + + /// Closes the SSE stream for this request, triggering client reconnection. + /// + /// Only available when using HTTPServerTransport with eventStore configured. + /// Use this to implement polling behavior during long-running operations. + public let closeSSEStream: (@Sendable () async -> Void)? + + /// Closes the standalone GET SSE stream, triggering client reconnection. + /// + /// Only available when using HTTPServerTransport with eventStore configured. + public let closeStandaloneSSEStream: (@Sendable () async -> Void)? + + public init( + authInfo: AuthInfo? = nil, + closeSSEStream: (@Sendable () async -> Void)? = nil, + closeStandaloneSSEStream: (@Sendable () async -> Void)? = nil + ) { + self.authInfo = authInfo + self.closeSSEStream = closeSSEStream + self.closeStandaloneSSEStream = closeStandaloneSSEStream + } +} + +/// A message received from a transport with optional context. +/// +/// This is the Swift equivalent of TypeScript's `onmessage(message, extra)` pattern, +/// adapted for Swift's `AsyncThrowingStream` approach. Each message carries its own +/// context, eliminating race conditions that would occur if context were stored +/// as mutable state on the transport. +/// +/// ## Example +/// +/// ```swift +/// for try await message in transport.receive() { +/// let data = message.data +/// if let authInfo = message.context?.authInfo { +/// // Handle authenticated request +/// } +/// } +/// ``` +public struct TransportMessage: Sendable { + /// The raw message data (JSON-RPC message). + public let data: Data + + /// Context associated with this message. + /// + /// Includes authentication info, SSE stream controls, and other per-message + /// context. For simple transports, this is `nil`. + public let context: MessageContext? + + public init(data: Data, context: MessageContext? = nil) { + self.data = data + self.context = context + } +} + +// MARK: - Transport Protocol + /// Protocol defining the transport layer for MCP communication public protocol Transport: Actor { var logger: Logger { get } @@ -36,8 +111,15 @@ public protocol Transport: Actor { /// - relatedRequestId: The ID of the request this message relates to (for response routing) func send(_ data: Data, relatedRequestId: RequestId?) async throws - /// Receives data in an async sequence - func receive() -> AsyncThrowingStream + /// Receives messages with optional context in an async sequence. + /// + /// Each message includes optional context (auth info, SSE closures, etc.) + /// that was associated with it at receive time. This pattern matches + /// TypeScript's `onmessage(message, extra)` callback approach. + /// + /// For simple transports, messages are yielded with `nil` context. + /// For HTTP transports, context includes authentication info and SSE controls. + func receive() -> AsyncThrowingStream } // MARK: - Default Implementation diff --git a/Sources/MCP/Base/Transports/HTTPClientTransport.swift b/Sources/MCP/Base/Transports/HTTPClientTransport.swift index 214486bf..886d73d0 100644 --- a/Sources/MCP/Base/Transports/HTTPClientTransport.swift +++ b/Sources/MCP/Base/Transports/HTTPClientTransport.swift @@ -75,8 +75,8 @@ public actor HTTPClientTransport: Transport { private let requestModifier: (URLRequest) -> URLRequest private var isConnected = false - private let messageStream: AsyncThrowingStream - private let messageContinuation: AsyncThrowingStream.Continuation + private let messageStream: AsyncThrowingStream + private let messageContinuation: AsyncThrowingStream.Continuation /// Stream for signaling when session ID is set private var sessionIDSignalStream: AsyncStream? @@ -150,7 +150,7 @@ public actor HTTPClientTransport: Transport { self.requestModifier = requestModifier // Create message stream - var continuation: AsyncThrowingStream.Continuation! + var continuation: AsyncThrowingStream.Continuation! self.messageStream = AsyncThrowingStream { continuation = $0 } self.messageContinuation = continuation @@ -447,10 +447,10 @@ public actor HTTPClientTransport: Transport { // Process response based on content type if contentType.contains("text/event-stream") { logger.warning("SSE responses aren't fully supported on Linux") - messageContinuation.yield(data) + messageContinuation.yield(TransportMessage(data: data)) } else if contentType.contains("application/json") { logger.trace("Received JSON response", metadata: ["size": "\(data.count)"]) - messageContinuation.yield(data) + messageContinuation.yield(TransportMessage(data: data)) } else if expectsContentType && !data.isEmpty { // Per MCP spec: requests MUST receive application/json or text/event-stream // Notifications expect 202 Accepted with no body, so unexpected content-type is ignored @@ -495,7 +495,7 @@ public actor HTTPClientTransport: Transport { buffer.append(byte) } logger.trace("Received JSON response", metadata: ["size": "\(buffer.count)"]) - messageContinuation.yield(buffer) + messageContinuation.yield(TransportMessage(data: buffer)) } else { // Collect data to check if response has content var buffer = Data() @@ -572,14 +572,14 @@ public actor HTTPClientTransport: Transport { /// Receives data in an async sequence /// - /// This returns an AsyncThrowingStream that emits Data objects representing + /// This returns an AsyncThrowingStream that emits TransportMessage objects representing /// each JSON-RPC message received from the server. This includes: /// /// - Direct responses to client requests /// - Server-initiated messages delivered via SSE streams /// - /// - Returns: An AsyncThrowingStream of Data objects - public func receive() -> AsyncThrowingStream { + /// - Returns: An AsyncThrowingStream of TransportMessage objects + public func receive() -> AsyncThrowingStream { return messageStream } @@ -875,7 +875,7 @@ public actor HTTPClientTransport: Transport { if processed.isResponse { receivedResponse = true } - messageContinuation.yield(processed.data) + messageContinuation.yield(TransportMessage(data: processed.data)) } } diff --git a/Sources/MCP/Base/Transports/HTTPServerTransport+Types.swift b/Sources/MCP/Base/Transports/HTTPServerTransport+Types.swift index 6adb29e0..1e8ff55b 100644 --- a/Sources/MCP/Base/Transports/HTTPServerTransport+Types.swift +++ b/Sources/MCP/Base/Transports/HTTPServerTransport+Types.swift @@ -1,12 +1,77 @@ import Foundation // Types extracted from HTTPServerTransport.swift +// - AuthInfo // - Options // - SecuritySettings // - EventStore protocol // - HTTPRequest // - HTTPResponse +// MARK: - Authentication + +/// Information about a validated access token. +/// +/// This struct contains authentication context that can be provided to request handlers +/// when using HTTP transports with OAuth or other token-based authentication. +/// +/// Matches the TypeScript SDK's `AuthInfo` interface. +/// +/// ## Example +/// +/// ```swift +/// server.withRequestHandler(CallTool.self) { params, context in +/// if let authInfo = context.authInfo { +/// print("Authenticated as: \(authInfo.clientId)") +/// print("Scopes: \(authInfo.scopes)") +/// } +/// return CallTool.Result(content: [.text("Done")]) +/// } +/// ``` +public struct AuthInfo: Hashable, Codable, Sendable { + /// The access token string. + public let token: String + + /// The client ID associated with this token. + public let clientId: String + + /// Scopes associated with this token. + public let scopes: [String] + + /// When the token expires (in seconds since epoch). + /// + /// If `nil`, the token does not expire or expiration is unknown. + public let expiresAt: Int? + + /// The RFC 8707 resource server identifier for which this token is valid. + /// + /// If set, this should match the MCP server's resource identifier (minus hash fragment). + public let resource: String? + + /// Additional data associated with the token. + /// + /// Use this for any additional data that needs to be attached to the auth info. + public let extra: [String: Value]? + + public init( + token: String, + clientId: String, + scopes: [String], + expiresAt: Int? = nil, + resource: String? = nil, + extra: [String: Value]? = nil + ) { + self.token = token + self.clientId = clientId + self.scopes = scopes + self.expiresAt = expiresAt + self.resource = resource + self.extra = extra + } +} + +// MARK: - Transport Options + /// Configuration options for HTTPServerTransport public struct HTTPServerTransportOptions: Sendable { /// Function that generates a session ID for the transport. diff --git a/Sources/MCP/Base/Transports/HTTPServerTransport.swift b/Sources/MCP/Base/Transports/HTTPServerTransport.swift index 69f9c872..ee691f8c 100644 --- a/Sources/MCP/Base/Transports/HTTPServerTransport.swift +++ b/Sources/MCP/Base/Transports/HTTPServerTransport.swift @@ -87,8 +87,8 @@ public actor HTTPServerTransport: Transport { private let standaloneSseStreamId = "_GET_stream" // Server receive stream (messages from HTTP clients go here) - private let serverStream: AsyncThrowingStream - private let serverContinuation: AsyncThrowingStream.Continuation + private let serverStream: AsyncThrowingStream + private let serverContinuation: AsyncThrowingStream.Continuation /// Closure called when the transport is closed public var onClose: (@Sendable () async -> Void)? @@ -111,7 +111,7 @@ public actor HTTPServerTransport: Transport { ) // Create server receive stream - var continuation: AsyncThrowingStream.Continuation! + var continuation: AsyncThrowingStream.Continuation! self.serverStream = AsyncThrowingStream { continuation = $0 } self.serverContinuation = continuation } @@ -206,7 +206,7 @@ public actor HTTPServerTransport: Transport { } /// Returns the stream of messages from HTTP clients. - public func receive() -> AsyncThrowingStream { + public func receive() -> AsyncThrowingStream { return serverStream } @@ -219,9 +219,11 @@ public actor HTTPServerTransport: Transport { /// - GET: Establish SSE stream for server-initiated notifications /// - DELETE: Terminate the session /// - /// - Parameter request: The incoming HTTP request + /// - Parameters: + /// - request: The incoming HTTP request + /// - authInfo: Authentication information for this request (from middleware) /// - Returns: An HTTP response - public func handleRequest(_ request: HTTPRequest) async -> HTTPResponse { + public func handleRequest(_ request: HTTPRequest, authInfo: AuthInfo? = nil) async -> HTTPResponse { // Check if transport has been terminated (applies to all modes) // Per spec: server MUST respond to requests after termination with 404 Not Found if terminated { @@ -239,7 +241,7 @@ public actor HTTPServerTransport: Transport { switch request.method.uppercased() { case "POST": - return await handlePostRequest(request) + return await handlePostRequest(request, authInfo: authInfo) case "GET": return await handleGetRequest(request) case "DELETE": @@ -256,7 +258,7 @@ public actor HTTPServerTransport: Transport { // MARK: - POST Request Handling - private func handlePostRequest(_ request: HTTPRequest) async -> HTTPResponse { + private func handlePostRequest(_ request: HTTPRequest, authInfo: AuthInfo?) async -> HTTPResponse { // Validate Accept header // Per spec: Client must accept both application/json and text/event-stream for SSE mode. // However, when JSON response mode is enabled, only application/json is required. @@ -419,7 +421,9 @@ public actor HTTPServerTransport: Transport { if !hasRequests { // Only notifications - yield to server and return 202 - serverContinuation.yield(body) + // Notifications don't need SSE closures since there's no response stream + let context = MessageContext(authInfo: authInfo) + serverContinuation.yield(TransportMessage(data: body, context: context)) return HTTPResponse(statusCode: 202, headers: sessionHeaders()) } @@ -434,7 +438,7 @@ public actor HTTPServerTransport: Transport { // Check if using JSON response mode if options.enableJsonResponse { - return await handleJsonResponseMode(streamId: streamId, requestIds: requestIds, body: body) + return await handleJsonResponseMode(streamId: streamId, requestIds: requestIds, body: body, authInfo: authInfo) } // SSE streaming mode @@ -443,14 +447,16 @@ public actor HTTPServerTransport: Transport { requestIds: requestIds, body: body, request: request, - messages: messages + messages: messages, + authInfo: authInfo ) } private func handleJsonResponseMode( streamId: String, requestIds: [RequestId], - body: Data + body: Data, + authInfo: AuthInfo? ) async -> HTTPResponse { // Create stream for receiving the response let (stream, continuation) = AsyncThrowingStream.makeStream() @@ -458,8 +464,9 @@ public actor HTTPServerTransport: Transport { let state = JsonStreamState(continuation: continuation) jsonStreamMapping[streamId] = state - // Yield the message to the server - serverContinuation.yield(body) + // JSON response mode doesn't have SSE streams to close + let context = MessageContext(authInfo: authInfo) + serverContinuation.yield(TransportMessage(data: body, context: context)) // Wait for response - this is cancellation-aware unlike withCheckedContinuation do { @@ -484,7 +491,8 @@ public actor HTTPServerTransport: Transport { requestIds: [RequestId], body: Data, request: HTTPRequest, - messages: [[String: Any]] + messages: [[String: Any]], + authInfo: AuthInfo? ) async -> HTTPResponse { let (stream, streamContinuation) = AsyncThrowingStream.makeStream() @@ -510,8 +518,30 @@ public actor HTTPServerTransport: Transport { // Write priming event if appropriate await writePrimingEvent(streamId: streamId, continuation: streamContinuation, protocolVersion: protocolVersion) - // Yield the message to the server - serverContinuation.yield(body) + // Create SSE closure for handlers to close this request's stream + // Use requestIds[0] for the primary request (batch requests share the same stream) + let closeSSEStreamClosure: (@Sendable () async -> Void)? = if let firstRequestId = requestIds.first { + { [weak self] in + await self?.closeSSEStream(for: firstRequestId) + } + } else { + nil + } + + // Create closure for standalone SSE stream + let closeStandaloneSSEStreamClosure: @Sendable () async -> Void = { [weak self] in + await self?.closeStandaloneSSEStream() + } + + // Create context with auth info and SSE closures + let context = MessageContext( + authInfo: authInfo, + closeSSEStream: closeSSEStreamClosure, + closeStandaloneSSEStream: closeStandaloneSSEStreamClosure + ) + + // Yield the message to the server with context + serverContinuation.yield(TransportMessage(data: body, context: context)) var headers = sessionHeaders() headers[HTTPHeader.contentType] = "text/event-stream" diff --git a/Sources/MCP/Base/Transports/InMemoryTransport.swift b/Sources/MCP/Base/Transports/InMemoryTransport.swift index 9234dfc4..195a85dd 100644 --- a/Sources/MCP/Base/Transports/InMemoryTransport.swift +++ b/Sources/MCP/Base/Transports/InMemoryTransport.swift @@ -24,7 +24,7 @@ public actor InMemoryTransport: Transport { // Message queues private var incomingMessages: [Data] = [] - private var messageContinuation: AsyncThrowingStream.Continuation? + private var messageContinuation: AsyncThrowingStream.Continuation? /// Creates a new in-memory transport /// @@ -168,7 +168,7 @@ public actor InMemoryTransport: Transport { logger.debug("Message received", metadata: ["size": "\(data.count)"]) if let continuation = messageContinuation { - continuation.yield(data) + continuation.yield(TransportMessage(data: data)) } else { // Queue message if stream not yet created incomingMessages.append(data) @@ -177,14 +177,14 @@ public actor InMemoryTransport: Transport { /// Receives messages from the paired transport /// - /// - Returns: An AsyncThrowingStream of Data objects representing messages - public func receive() -> AsyncThrowingStream { - return AsyncThrowingStream { continuation in + /// - Returns: An AsyncThrowingStream of TransportMessage objects representing messages + public func receive() -> AsyncThrowingStream { + return AsyncThrowingStream { continuation in self.messageContinuation = continuation // Deliver any queued messages for message in self.incomingMessages { - continuation.yield(message) + continuation.yield(TransportMessage(data: message)) } self.incomingMessages.removeAll() diff --git a/Sources/MCP/Base/Transports/NetworkTransport.swift b/Sources/MCP/Base/Transports/NetworkTransport.swift index 4a27cd18..9c1cdd9a 100644 --- a/Sources/MCP/Base/Transports/NetworkTransport.swift +++ b/Sources/MCP/Base/Transports/NetworkTransport.swift @@ -239,8 +239,8 @@ import Logging private var reconnectAttempt = 0 private var heartbeatTask: Task? private var lastHeartbeatTime: Date? - private let messageStream: AsyncThrowingStream - private let messageContinuation: AsyncThrowingStream.Continuation + private let messageStream: AsyncThrowingStream + private let messageContinuation: AsyncThrowingStream.Continuation // Connection is marked nonisolated(unsafe) to allow access from closures private nonisolated(unsafe) var connection: NetworkConnectionProtocol @@ -296,7 +296,7 @@ import Logging self.bufferConfig = bufferConfig // Create message stream - var continuation: AsyncThrowingStream.Continuation! + var continuation: AsyncThrowingStream.Continuation! self.messageStream = AsyncThrowingStream { continuation = $0 } self.messageContinuation = continuation } @@ -513,11 +513,11 @@ import Logging /// Receives data in an async sequence /// - /// This returns an AsyncThrowingStream that emits Data objects representing + /// This returns an AsyncThrowingStream that emits TransportMessage objects representing /// each JSON-RPC message received from the network connection. /// - /// - Returns: An AsyncThrowingStream of Data objects - public func receive() -> AsyncThrowingStream { + /// - Returns: An AsyncThrowingStream of TransportMessage objects + public func receive() -> AsyncThrowingStream { return messageStream } @@ -582,7 +582,7 @@ import Logging if !messageData.isEmpty { logger.debug( "Message received", metadata: ["size": "\(messageData.count)"]) - messageContinuation.yield(Data(messageData)) + messageContinuation.yield(TransportMessage(data: Data(messageData))) } } } catch let error as NWError { diff --git a/Sources/MCP/Base/Transports/StdioTransport.swift b/Sources/MCP/Base/Transports/StdioTransport.swift index b6ae1fae..bb21aa9d 100644 --- a/Sources/MCP/Base/Transports/StdioTransport.swift +++ b/Sources/MCP/Base/Transports/StdioTransport.swift @@ -54,8 +54,8 @@ import struct Foundation.Data public nonisolated let logger: Logger private var isConnected = false - private let messageStream: AsyncThrowingStream - private let messageContinuation: AsyncThrowingStream.Continuation + private let messageStream: AsyncThrowingStream + private let messageContinuation: AsyncThrowingStream.Continuation /// Creates a new stdio transport with the specified file descriptors /// @@ -77,7 +77,7 @@ import struct Foundation.Data factory: { _ in SwiftLogNoOpLogHandler() }) // Create message stream - var continuation: AsyncThrowingStream.Continuation! + var continuation: AsyncThrowingStream.Continuation! self.messageStream = AsyncThrowingStream { continuation = $0 } self.messageContinuation = continuation } @@ -172,7 +172,7 @@ import struct Foundation.Data if !messageData.isEmpty { logger.trace( "Message received", metadata: ["size": "\(messageData.count)"]) - messageContinuation.yield(Data(messageData)) + messageContinuation.yield(TransportMessage(data: Data(messageData))) } } } catch let error where MCPError.isResourceTemporarilyUnavailable(error) { @@ -240,8 +240,8 @@ import struct Foundation.Data /// or batches containing multiple requests/notifications encoded as JSON arrays. /// Each message is guaranteed to be a complete JSON object or array. /// - /// - Returns: An AsyncThrowingStream of Data objects representing JSON-RPC messages - public func receive() -> AsyncThrowingStream { + /// - Returns: An AsyncThrowingStream of TransportMessage objects representing JSON-RPC messages + public func receive() -> AsyncThrowingStream { return messageStream } } diff --git a/Sources/MCP/Client/Client.swift b/Sources/MCP/Client/Client.swift index a857224e..4b8fc240 100644 --- a/Sources/MCP/Client/Client.swift +++ b/Sources/MCP/Client/Client.swift @@ -597,9 +597,13 @@ public actor Client { do { let stream = await connection.receive() - for try await data in stream { + for try await transportMessage in stream { if Task.isCancelled { break } + // Extract the raw data from the transport message + // (Client doesn't use message context - authInfo and SSE closures are server-side only) + let data = transportMessage.data + // Attempt to decode data // Try decoding as a batch response first if let batchResponse = try? decoder.decode([AnyResponse].self, from: data) { diff --git a/Sources/MCP/Server/Server+RequestHandling.swift b/Sources/MCP/Server/Server+RequestHandling.swift index 2be03fc9..7194b038 100644 --- a/Sources/MCP/Server/Server+RequestHandling.swift +++ b/Sources/MCP/Server/Server+RequestHandling.swift @@ -20,7 +20,7 @@ extension Server { } /// Process a batch of requests and/or notifications - func handleBatch(_ batch: Batch) async throws { + func handleBatch(_ batch: Batch, messageContext: MessageContext? = nil) async throws { // Capture the connection at batch start. // This ensures all batch responses go to the correct client. let capturedConnection = self.connection @@ -49,7 +49,7 @@ extension Server { switch item { case .request(let request): // For batched requests, collect responses instead of sending immediately - if let response = try await handleRequest(request, sendResponse: false) { + if let response = try await handleRequest(request, sendResponse: false, messageContext: messageContext) { responses.append(response) } @@ -106,6 +106,18 @@ extension Server { /// /// Contains the progress token and any additional metadata. let meta: RequestMeta? + /// Authentication information, if available. + /// + /// Set by HTTP transports when OAuth or other authentication is in use. + let authInfo: AuthInfo? + /// Closure to close the SSE stream for this request. + /// + /// Only set by HTTP transports with SSE support. + let closeSSEStream: (@Sendable () async -> Void)? + /// Closure to close the standalone SSE stream. + /// + /// Only set by HTTP transports with SSE support. + let closeStandaloneSSEStream: (@Sendable () async -> Void)? } /// Extract `_meta` from request parameters if present. @@ -175,8 +187,9 @@ extension Server { /// - Parameters: /// - request: The request to handle /// - sendResponse: Whether to send the response immediately (true) or return it (false) + /// - messageContext: Optional context from the transport message (authInfo, SSE closures) /// - Returns: The response when sendResponse is false - func handleRequest(_ request: Request, sendResponse: Bool = true) + func handleRequest(_ request: Request, sendResponse: Bool = true, messageContext: MessageContext? = nil) async throws -> Response? { // Capture the connection and session ID at request time. @@ -184,11 +197,21 @@ extension Server { // changes while the handler is executing (e.g., another client connects). let capturedConnection = self.connection let requestMeta = extractMeta(from: request.params) + + // Extract context from transport message (set by HTTP transports with per-message context) + // This pattern aligns with TypeScript's onmessage(message, { authInfo, closeSSEStream, ... }) + let authInfo = messageContext?.authInfo + let closeSSEStream = messageContext?.closeSSEStream + let closeStandaloneSSEStream = messageContext?.closeStandaloneSSEStream + let context = RequestContext( capturedConnection: capturedConnection, requestId: request.id, sessionId: await capturedConnection?.sessionId, - meta: requestMeta + meta: requestMeta, + authInfo: authInfo, + closeSSEStream: closeSSEStream, + closeStandaloneSSEStream: closeStandaloneSSEStream ) // Check if this is a pre-processed error request (empty method) @@ -270,6 +293,9 @@ extension Server { sessionId: context.sessionId, requestId: context.requestId, _meta: context.meta, + authInfo: context.authInfo, + closeSSEStream: context.closeSSEStream, + closeStandaloneSSEStream: context.closeStandaloneSSEStream, shouldSendLogMessage: { [weak self, context] level in guard let self else { return true } return await self.shouldSendLogMessage(at: level, forSession: context.sessionId) diff --git a/Sources/MCP/Server/Server.swift b/Sources/MCP/Server/Server.swift index b5bb9cf4..829766f3 100644 --- a/Sources/MCP/Server/Server.swift +++ b/Sources/MCP/Server/Server.swift @@ -215,6 +215,15 @@ public actor Server { /// a long-running tool), it should use this context to ensure the notification is /// routed to the correct client, even if other clients have connected in the meantime. /// + /// This context provides: + /// - Request identification (`requestId`, `_meta`) + /// - Session tracking (`sessionId`) + /// - Authentication context (`authInfo`) + /// - Notification sending (`sendNotification`, `sendMessage`, `sendProgress`) + /// - Bidirectional requests (`elicit`, `elicitUrl`) + /// - Cancellation checking (`isCancelled`, `checkCancellation`) + /// - SSE stream management (`closeSSEStream`, `closeStandaloneSSEStream`) + /// /// Example: /// ```swift /// server.withRequestHandler(CallTool.self) { params, context in @@ -298,6 +307,54 @@ public actor Server { /// ``` public let _meta: RequestMeta? + /// Authentication information for this request. + /// + /// Contains validated access token information when using HTTP transports + /// with OAuth or other token-based authentication. + /// + /// This matches the TypeScript SDK's `extra.authInfo`. + /// + /// ## Example + /// + /// ```swift + /// server.withRequestHandler(CallTool.self) { params, context in + /// if let authInfo = context.authInfo { + /// print("Authenticated as: \(authInfo.clientId)") + /// print("Scopes: \(authInfo.scopes)") + /// + /// // Check if token has required scope + /// guard authInfo.scopes.contains("tools:execute") else { + /// throw MCPError.invalidRequest("Missing required scope") + /// } + /// } + /// return CallTool.Result(content: [.text("Done")]) + /// } + /// ``` + public let authInfo: AuthInfo? + + /// Closes the SSE stream for this request, triggering client reconnection. + /// + /// Only available when using StreamableHTTPServerTransport with eventStore configured. + /// Use this to implement polling behavior during long-running operations - + /// the client will reconnect after the retry interval specified in the priming event. + /// + /// This matches the TypeScript SDK's `extra.closeSSEStream()` and + /// Python's `ctx.close_sse_stream()`. + /// + /// - Note: This is `nil` when not using an HTTP/SSE transport. + public let closeSSEStream: (@Sendable () async -> Void)? + + /// Closes the standalone GET SSE stream, triggering client reconnection. + /// + /// Only available when using StreamableHTTPServerTransport with eventStore configured. + /// Use this to implement polling behavior for server-initiated notifications. + /// + /// This matches the TypeScript SDK's `extra.closeStandaloneSSEStream()` and + /// Python's `ctx.close_standalone_sse_stream()`. + /// + /// - Note: This is `nil` when not using an HTTP/SSE transport. + public let closeStandaloneSSEStream: (@Sendable () async -> Void)? + /// Check if a log message at the given level should be sent. /// /// This respects the minimum log level set by the client via `logging/setLevel`. @@ -747,9 +804,13 @@ public actor Server { task = Task { do { let stream = await transport.receive() - for try await data in stream { + for try await transportMessage in stream { if Task.isCancelled { break } // Check cancellation inside loop + // Extract the raw data and optional context from the transport message + let data = transportMessage.data + let messageContext = transportMessage.context + var requestID: RequestId? do { // Attempt to decode as batch first, then as individual request, response, or notification @@ -761,7 +822,7 @@ public actor Server { Task { [weak self] in guard let self else { return } do { - try await self.handleBatch(batch) + try await self.handleBatch(batch, messageContext: messageContext) } catch { await self.logger?.error( "Error handling batch", @@ -779,13 +840,13 @@ public actor Server { // can await a response while the message loop continues processing // incoming messages including that response. let requestId = request.id - let handlerTask = Task { [weak self] in + let handlerTask = Task { [weak self, messageContext] in guard let self else { return } defer { Task { await self.removeInFlightRequest(requestId) } } do { - _ = try await self.handleRequest(request, sendResponse: true) + _ = try await self.handleRequest(request, sendResponse: true, messageContext: messageContext) } catch { // handleRequest already sends error responses, so this // only catches errors from send() itself diff --git a/Tests/MCPTests/HTTPClientTransportTests.swift b/Tests/MCPTests/HTTPClientTransportTests.swift index 6ff4547b..6b62fe4a 100644 --- a/Tests/MCPTests/HTTPClientTransportTests.swift +++ b/Tests/MCPTests/HTTPClientTransportTests.swift @@ -188,7 +188,7 @@ import Testing var iterator = stream.makeAsyncIterator() let receivedData = try await iterator.next() - #expect(receivedData == responseData) + #expect(receivedData?.data == responseData) } @Test("Send and Receive Session ID", .httpClientTransportSetup) @@ -1085,7 +1085,7 @@ import Testing let expectedData = #"{"key":"value"}"#.data(using: .utf8)! let receivedData = try await iterator.next() - #expect(receivedData == expectedData) + #expect(receivedData?.data == expectedData) await transport.disconnect() } @@ -1147,7 +1147,7 @@ import Testing let expectedData = #"{"key":"value"}"#.data(using: .utf8)! let receivedData = try await iterator.next() - #expect(receivedData == expectedData) + #expect(receivedData?.data == expectedData) await transport.disconnect() } @@ -1566,7 +1566,7 @@ import Testing // Should only receive the actual message, not the priming event let expectedData = #"{"result":"ok"}"#.data(using: .utf8)! - #expect(receivedData == expectedData) + #expect(receivedData?.data == expectedData) await transport.disconnect() } @@ -1715,19 +1715,19 @@ import Testing // First: notification let msg1 = try await iterator.next() #expect(msg1 != nil) - let msg1String = String(data: msg1!, encoding: .utf8)! + let msg1String = String(data: msg1!.data, encoding: .utf8)! #expect(msg1String.contains("notifications/progress")) // Second: server request let msg2 = try await iterator.next() #expect(msg2 != nil) - let msg2String = String(data: msg2!, encoding: .utf8)! + let msg2String = String(data: msg2!.data, encoding: .utf8)! #expect(msg2String.contains("sampling/createMessage")) // Third: response let msg3 = try await iterator.next() #expect(msg3 != nil) - let msg3String = String(data: msg3!, encoding: .utf8)! + let msg3String = String(data: msg3!.data, encoding: .utf8)! #expect(msg3String.contains("\"result\"")) // The lastReceivedEventId should be evt-3 (last event with ID) @@ -1791,7 +1791,7 @@ import Testing let msg = try await iterator.next() #expect(msg != nil) - let msgString = String(data: msg!, encoding: .utf8)! + let msgString = String(data: msg!.data, encoding: .utf8)! #expect(msgString.contains("\"error\"")) #expect(msgString.contains("\(ErrorCode.invalidRequest)")) @@ -1860,7 +1860,7 @@ import Testing let msg = try await iterator.next() #expect(msg != nil) - let msgString = String(data: msg!, encoding: .utf8)! + let msgString = String(data: msg!.data, encoding: .utf8)! // The ID should be remapped to "original-req-42" #expect(msgString.contains("\"id\":\"original-req-42\"")) @@ -1926,7 +1926,7 @@ import Testing let msg = try await iterator.next() #expect(msg != nil) - let msgString = String(data: msg!, encoding: .utf8)! + let msgString = String(data: msg!.data, encoding: .utf8)! // The ID should be remapped to 42 (numeric) #expect(msgString.contains("\"id\":42")) @@ -1991,7 +1991,7 @@ import Testing let msg = try await iterator.next() #expect(msg != nil) - let msgString = String(data: msg!, encoding: .utf8)! + let msgString = String(data: msg!.data, encoding: .utf8)! // The ID should remain as "original-id" #expect(msgString.contains("\"id\":\"original-id\"")) @@ -2058,7 +2058,7 @@ import Testing let msg = try await iterator.next() #expect(msg != nil) - let msgString = String(data: msg!, encoding: .utf8)! + let msgString = String(data: msg!.data, encoding: .utf8)! // The ID should be remapped to "my-failed-request" #expect(msgString.contains("\"id\":\"my-failed-request\"")) @@ -2126,21 +2126,21 @@ import Testing // First: server request - ID should NOT be remapped let msg1 = try await iterator.next() #expect(msg1 != nil) - let msg1String = String(data: msg1!, encoding: .utf8)! + let msg1String = String(data: msg1!.data, encoding: .utf8)! #expect(msg1String.contains("\"id\":\"server-req-1\"")) // Original ID preserved #expect(msg1String.contains("sampling/createMessage")) // Second: notification - no ID field, should pass through unchanged let msg2 = try await iterator.next() #expect(msg2 != nil) - let msg2String = String(data: msg2!, encoding: .utf8)! + let msg2String = String(data: msg2!.data, encoding: .utf8)! #expect(msg2String.contains("notifications/progress")) #expect(!msg2String.contains("my-original-request")) // Third: response - ID SHOULD be remapped let msg3 = try await iterator.next() #expect(msg3 != nil) - let msg3String = String(data: msg3!, encoding: .utf8)! + let msg3String = String(data: msg3!.data, encoding: .utf8)! #expect(msg3String.contains("\"id\":\"my-original-request\"")) // Remapped ID #expect(!msg3String.contains("server-resp-id")) // Original ID replaced #expect(msg3String.contains("\"result\"")) diff --git a/Tests/MCPTests/Helpers/MockTransport.swift b/Tests/MCPTests/Helpers/MockTransport.swift index e35b8357..b0d9d176 100644 --- a/Tests/MCPTests/Helpers/MockTransport.swift +++ b/Tests/MCPTests/Helpers/MockTransport.swift @@ -30,7 +30,7 @@ actor MockTransport: Transport { private var dataToReceive: [Data] = [] private(set) var receivedMessages: [String] = [] - private var dataStreamContinuation: AsyncThrowingStream.Continuation? + private var dataStreamContinuation: AsyncThrowingStream.Continuation? var shouldFailConnect = false var shouldFailSend = false @@ -59,11 +59,11 @@ actor MockTransport: Transport { sentData.append(message) } - public func receive() -> AsyncThrowingStream { - return AsyncThrowingStream { continuation in + public func receive() -> AsyncThrowingStream { + return AsyncThrowingStream { continuation in dataStreamContinuation = continuation for message in dataToReceive { - continuation.yield(message) + continuation.yield(TransportMessage(data: message)) if let string = String(data: message, encoding: .utf8) { receivedMessages.append(string) } @@ -82,7 +82,7 @@ actor MockTransport: Transport { func queue(data: Data) { if let continuation = dataStreamContinuation { - continuation.yield(data) + continuation.yield(TransportMessage(data: data)) } else { dataToReceive.append(data) } diff --git a/Tests/MCPTests/InMemoryTransportTests.swift b/Tests/MCPTests/InMemoryTransportTests.swift index d8228e77..4a7ce5d4 100644 --- a/Tests/MCPTests/InMemoryTransportTests.swift +++ b/Tests/MCPTests/InMemoryTransportTests.swift @@ -64,8 +64,8 @@ struct InMemoryTransportTests { // Start receiving on server let serverReceiveTask = Task { var messages: [Data] = [] - for try await message in await serverTransport.receive() { - messages.append(message) + for try await transportMessage in await serverTransport.receive() { + messages.append(transportMessage.data) if messages.count >= 3 { break } @@ -105,8 +105,8 @@ struct InMemoryTransportTests { // Set up receivers let receive1Task = Task { var messages: [String] = [] - for try await data in await transport1.receive() { - if let message = String(data: data, encoding: .utf8) { + for try await transportMessage in await transport1.receive() { + if let message = String(data: transportMessage.data, encoding: .utf8) { messages.append(message) if messages.count >= 2 { break @@ -118,8 +118,8 @@ struct InMemoryTransportTests { let receive2Task = Task { var messages: [String] = [] - for try await data in await transport2.receive() { - if let message = String(data: data, encoding: .utf8) { + for try await transportMessage in await transport2.receive() { + if let message = String(data: transportMessage.data, encoding: .utf8) { messages.append(message) if messages.count >= 2 { break @@ -267,8 +267,8 @@ struct InMemoryTransportTests { let messages = await serverTransport.receive() var receivedMessages: [String] = [] - for try await data in messages { - if let message = String(data: data, encoding: .utf8) { + for try await transportMessage in messages { + if let message = String(data: transportMessage.data, encoding: .utf8) { receivedMessages.append(message) if receivedMessages.count >= 3 { break @@ -323,8 +323,8 @@ struct InMemoryTransportTests { // Start receiving let receiveTask = Task { - for try await data in await serverTransport.receive() { - return data + for try await transportMessage in await serverTransport.receive() { + return transportMessage.data } return Data() } @@ -351,8 +351,8 @@ struct InMemoryTransportTests { // Start receiving let receiveTask = Task { var messages: [String] = [] - for try await data in await serverTransport.receive() { - if let message = String(data: data, encoding: .utf8) { + for try await transportMessage in await serverTransport.receive() { + if let message = String(data: transportMessage.data, encoding: .utf8) { messages.append(message) if messages.count >= 10 { break diff --git a/Tests/MCPTests/NetworkTransportTests.swift b/Tests/MCPTests/NetworkTransportTests.swift index be14b0dc..7e7ce41e 100644 --- a/Tests/MCPTests/NetworkTransportTests.swift +++ b/Tests/MCPTests/NetworkTransportTests.swift @@ -335,7 +335,7 @@ import Testing #expect(received != nil) if let received = received { - #expect(received == message.data(using: .utf8)!) + #expect(received.data == message.data(using: .utf8)!) } await transport.disconnect() @@ -450,8 +450,8 @@ import Testing let stream = await transport.receive() var receiveCount = 0 - for try await data in stream { - if let receivedStr = String(data: data, encoding: .utf8) { + for try await transportMessage in stream { + if let receivedStr = String(data: transportMessage.data, encoding: .utf8) { #expect(messages.contains(receivedStr)) receiveCount += 1 @@ -554,7 +554,7 @@ import Testing let received = try await iterator.next() #expect(received != nil) if let received = received { - #expect(received == message.data(using: .utf8)!) + #expect(received.data == message.data(using: .utf8)!) } await transport.disconnect() @@ -583,7 +583,7 @@ import Testing let received = try await iterator.next() #expect(received != nil) if let received = received { - #expect(received.count == largeMessage.count) + #expect(received.data.count == largeMessage.count) } await transport.disconnect() diff --git a/Tests/MCPTests/RequestHandlerContextTests.swift b/Tests/MCPTests/RequestHandlerContextTests.swift index aae258c4..494815f2 100644 --- a/Tests/MCPTests/RequestHandlerContextTests.swift +++ b/Tests/MCPTests/RequestHandlerContextTests.swift @@ -414,6 +414,111 @@ struct ServerRequestHandlerContextTests { // convenience methods (tested above), matching Python's ctx.elicit() pattern. Sampling // is done via the server directly, matching TypeScript's pattern where extra.sendRequest() // is generic and server.createMessage() is the convenience method. + + // MARK: - authInfo Tests + + /// Test that context.authInfo is nil for non-HTTP transports. + /// Based on TypeScript SDK's `extra.authInfo` which is only populated for authenticated HTTP connections. + @Test("context.authInfo is nil for InMemoryTransport") + func testAuthInfoNilForInMemoryTransport() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + actor AuthInfoTracker { + var receivedAuthInfo: AuthInfo? + var wasChecked = false + func set(_ authInfo: AuthInfo?) { + receivedAuthInfo = authInfo + wasChecked = true + } + } + let tracker = AuthInfoTracker() + + let server = Server( + name: "TestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "test_tool", description: "Test", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { _, context in + // Handler accesses context.authInfo - should be nil for InMemoryTransport + await tracker.set(context.authInfo) + return CallTool.Result(content: [.text("OK")]) + } + + let client = Client(name: "TestClient", version: "1.0.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + _ = try await client.callTool(name: "test_tool", arguments: [:]) + + let wasChecked = await tracker.wasChecked + let receivedAuthInfo = await tracker.receivedAuthInfo + #expect(wasChecked, "Handler should have been called") + #expect(receivedAuthInfo == nil, "authInfo should be nil for InMemoryTransport") + + await client.disconnect() + } + + // MARK: - closeSSEStream Tests + + /// Test that context.closeSSEStream is nil for non-HTTP transports. + /// Based on TypeScript SDK's `extra.closeSSEStream` which is only available for HTTP/SSE transports. + @Test("context.closeSSEStream is nil for InMemoryTransport") + func testCloseSSEStreamNilForInMemoryTransport() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + actor StreamClosureTracker { + var closeSSEStreamWasNil = false + var closeStandaloneSSEStreamWasNil = false + func set(closeSSE: Bool, closeStandalone: Bool) { + closeSSEStreamWasNil = closeSSE + closeStandaloneSSEStreamWasNil = closeStandalone + } + } + let tracker = StreamClosureTracker() + + let server = Server( + name: "TestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "test_tool", description: "Test", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { _, context in + // Check that SSE stream closures are nil for InMemoryTransport + await tracker.set( + closeSSE: context.closeSSEStream == nil, + closeStandalone: context.closeStandaloneSSEStream == nil + ) + return CallTool.Result(content: [.text("OK")]) + } + + let client = Client(name: "TestClient", version: "1.0.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + _ = try await client.callTool(name: "test_tool", arguments: [:]) + + let closeSSEStreamWasNil = await tracker.closeSSEStreamWasNil + let closeStandaloneSSEStreamWasNil = await tracker.closeStandaloneSSEStreamWasNil + #expect(closeSSEStreamWasNil, "closeSSEStream should be nil for InMemoryTransport") + #expect(closeStandaloneSSEStreamWasNil, "closeStandaloneSSEStream should be nil for InMemoryTransport") + + await client.disconnect() + } } // MARK: - Client RequestHandlerContext Tests diff --git a/Tests/MCPTests/StdioTransportTests.swift b/Tests/MCPTests/StdioTransportTests.swift index 05a47d12..fe123f79 100644 --- a/Tests/MCPTests/StdioTransportTests.swift +++ b/Tests/MCPTests/StdioTransportTests.swift @@ -59,12 +59,12 @@ struct StdioTransportTests { try writer.close() // Start receiving messages - let stream: AsyncThrowingStream = await transport.receive() + let stream = await transport.receive() var iterator = stream.makeAsyncIterator() // Get first message let received = try await iterator.next() - #expect(received == #"{"key":"value"}"#.data(using: .utf8)!) + #expect(received?.data == #"{"key":"value"}"#.data(using: .utf8)!) await transport.disconnect() } @@ -81,7 +81,7 @@ struct StdioTransportTests { try writer.writeAll(invalidJSON.data(using: .utf8)!) try writer.close() - let stream: AsyncThrowingStream = await transport.receive() + let stream = await transport.receive() var iterator = stream.makeAsyncIterator() _ = try await iterator.next() @@ -154,8 +154,8 @@ struct StdioTransportMultipleMessageTests { let stream = await transport.receive() var receivedMessages: [String] = [] - for try await data in stream { - if let message = String(data: data, encoding: .utf8) { + for try await transportMessage in stream { + if let message = String(data: transportMessage.data, encoding: .utf8) { receivedMessages.append(message) } } @@ -226,7 +226,7 @@ struct StdioTransportMessageFramingTests { // Should receive the complete reassembled message let expectedMessage = #"{"jsonrpc":"2.0","id":1,"method":"ping"}"# - #expect(received == expectedMessage.data(using: .utf8)!) + #expect(received?.data == expectedMessage.data(using: .utf8)!) await transport.disconnect() } @@ -248,8 +248,8 @@ struct StdioTransportMessageFramingTests { let stream = await transport.receive() var receivedMessages: [String] = [] - for try await data in stream { - if let message = String(data: data, encoding: .utf8) { + for try await transportMessage in stream { + if let message = String(data: transportMessage.data, encoding: .utf8) { receivedMessages.append(message) } } @@ -277,8 +277,8 @@ struct StdioTransportMessageFramingTests { let stream = await transport.receive() var receivedMessages: [String] = [] - for try await data in stream { - if let message = String(data: data, encoding: .utf8) { + for try await transportMessage in stream { + if let message = String(data: transportMessage.data, encoding: .utf8) { receivedMessages.append(message) } } @@ -309,7 +309,7 @@ struct StdioTransportMessageFramingTests { var iterator = stream.makeAsyncIterator() let received = try await iterator.next() - #expect(received == largeMessage.data(using: .utf8)!) + #expect(received?.data == largeMessage.data(using: .utf8)!) await transport.disconnect() } @@ -337,7 +337,7 @@ struct StdioTransportBidirectionalTests { let stream = await transport.receive() var iterator = stream.makeAsyncIterator() let receivedRequest = try await iterator.next() - #expect(receivedRequest == request.data(using: .utf8)!) + #expect(receivedRequest?.data == request.data(using: .utf8)!) // Transport sends a response back let response = #"{"jsonrpc":"2.0","id":1,"result":{"status":"ok"}}"# @@ -377,8 +377,8 @@ struct StdioTransportEOFHandlingTests { let stream = await transport.receive() var messages: [Data] = [] - for try await data in stream { - messages.append(data) + for try await transportMessage in stream { + messages.append(transportMessage.data) } // Should have received exactly one message, then stream should end @@ -406,8 +406,8 @@ struct StdioTransportEOFHandlingTests { let stream = await transport.receive() var messages: [Data] = [] - for try await data in stream { - messages.append(data) + for try await transportMessage in stream { + messages.append(transportMessage.data) } // Should only receive the complete message, incomplete one is discarded diff --git a/Tests/MCPTests/ToolTests.swift b/Tests/MCPTests/ToolTests.swift index 618a821f..916331b5 100644 --- a/Tests/MCPTests/ToolTests.swift +++ b/Tests/MCPTests/ToolTests.swift @@ -413,6 +413,9 @@ struct ToolTests { sessionId: nil, requestId: .number(1), _meta: nil, + authInfo: nil, + closeSSEStream: nil, + closeStandaloneSSEStream: nil, shouldSendLogMessage: { _ in true }, sendRequest: { _ in throw MCPError.internalError("Not implemented") } ) diff --git a/Tests/MCPTests/TransportSwitchingTests.swift b/Tests/MCPTests/TransportSwitchingTests.swift index b13f3ca7..22a59007 100644 --- a/Tests/MCPTests/TransportSwitchingTests.swift +++ b/Tests/MCPTests/TransportSwitchingTests.swift @@ -34,7 +34,7 @@ struct TransportSwitchingTests { private(set) var sentMessages: [SentMessage] = [] private var dataToReceive: [Data] = [] - private var dataStreamContinuation: AsyncThrowingStream.Continuation? + private var dataStreamContinuation: AsyncThrowingStream.Continuation? let name: String @@ -61,11 +61,11 @@ struct TransportSwitchingTests { sentMessages.append(SentMessage(data: data, relatedRequestId: relatedRequestId)) } - public func receive() -> AsyncThrowingStream { - return AsyncThrowingStream { continuation in + public func receive() -> AsyncThrowingStream { + return AsyncThrowingStream { continuation in dataStreamContinuation = continuation for message in dataToReceive { - continuation.yield(message) + continuation.yield(TransportMessage(data: message)) } dataToReceive.removeAll() } @@ -73,7 +73,7 @@ struct TransportSwitchingTests { func queue(data: Data) { if let continuation = dataStreamContinuation { - continuation.yield(data) + continuation.yield(TransportMessage(data: data)) } else { dataToReceive.append(data) } From 291f895b130be71c6ef138afa2f162e33bc680e2 Mon Sep 17 00:00:00 2001 From: Anthony DePasquale Date: Sun, 4 Jan 2026 16:29:22 +0100 Subject: [PATCH 7/8] Groundwork for future OAuth support --- .../Base/Transports/HTTPClientTransport.swift | 11 ++ Sources/MCP/Base/Transports/OAuth.swift | 126 ++++++++++++++++++ 2 files changed, 137 insertions(+) create mode 100644 Sources/MCP/Base/Transports/OAuth.swift diff --git a/Sources/MCP/Base/Transports/HTTPClientTransport.swift b/Sources/MCP/Base/Transports/HTTPClientTransport.swift index 886d73d0..4a99803e 100644 --- a/Sources/MCP/Base/Transports/HTTPClientTransport.swift +++ b/Sources/MCP/Base/Transports/HTTPClientTransport.swift @@ -74,6 +74,9 @@ public actor HTTPClientTransport: Transport { /// Closure to modify requests before they are sent private let requestModifier: (URLRequest) -> URLRequest + /// OAuth provider for automatic token management (reserved for future OAuth implementation) + private let authProvider: (any OAuthClientProvider)? + private var isConnected = false private let messageStream: AsyncThrowingStream private let messageContinuation: AsyncThrowingStream.Continuation @@ -112,6 +115,10 @@ public actor HTTPClientTransport: Transport { /// - sseInitializationTimeout: Maximum time to wait for session ID before proceeding with SSE (default: 10 seconds) /// - reconnectionOptions: Configuration for reconnection behavior (default: .default) /// - requestModifier: Optional closure to customize requests before they are sent (default: no modification) + /// - authProvider: Optional OAuth provider for automatic token management. + /// When provided, the transport will use the provider to obtain Bearer tokens + /// and handle 401 responses. This parameter is reserved for future OAuth + /// implementation and is not currently used. /// - logger: Optional logger instance for transport events public init( endpoint: URL, @@ -120,6 +127,7 @@ public actor HTTPClientTransport: Transport { sseInitializationTimeout: TimeInterval = 10, reconnectionOptions: HTTPReconnectionOptions = .default, requestModifier: @escaping (URLRequest) -> URLRequest = { $0 }, + authProvider: (any OAuthClientProvider)? = nil, logger: Logger? = nil ) { self.init( @@ -129,6 +137,7 @@ public actor HTTPClientTransport: Transport { sseInitializationTimeout: sseInitializationTimeout, reconnectionOptions: reconnectionOptions, requestModifier: requestModifier, + authProvider: authProvider, logger: logger ) } @@ -140,6 +149,7 @@ public actor HTTPClientTransport: Transport { sseInitializationTimeout: TimeInterval = 10, reconnectionOptions: HTTPReconnectionOptions = .default, requestModifier: @escaping (URLRequest) -> URLRequest = { $0 }, + authProvider: (any OAuthClientProvider)? = nil, logger: Logger? = nil ) { self.endpoint = endpoint @@ -148,6 +158,7 @@ public actor HTTPClientTransport: Transport { self.sseInitializationTimeout = sseInitializationTimeout self.reconnectionOptions = reconnectionOptions self.requestModifier = requestModifier + self.authProvider = authProvider // Create message stream var continuation: AsyncThrowingStream.Continuation! diff --git a/Sources/MCP/Base/Transports/OAuth.swift b/Sources/MCP/Base/Transports/OAuth.swift new file mode 100644 index 00000000..11111d5b --- /dev/null +++ b/Sources/MCP/Base/Transports/OAuth.swift @@ -0,0 +1,126 @@ +import Foundation + +// MARK: - OAuth Support +// +// This file provides the foundational types for OAuth 2.0 support in HTTP transports. +// The full OAuth implementation (discovery, PKCE, token exchange, provider implementations) +// will be added later. +// +// Current status: +// - OAuthTokens: Complete (matches RFC 6749) +// - UnauthorizedContext: Complete (for 401 handling) +// - OAuthClientProvider: Protocol defined, no implementations yet +// - HTTPClientTransport.authProvider: Parameter added, not yet wired up + +// MARK: - OAuth Types + +/// OAuth 2.0 tokens for authenticated requests. +/// +/// This struct holds the tokens obtained through an OAuth 2.0 authorization flow, +/// matching the token response format defined in RFC 6749 Section 5.1. +public struct OAuthTokens: Sendable, Codable, Equatable { + /// The access token to use for Bearer authentication. + public let accessToken: String + + /// The type of token issued. Per RFC 6749, this is case-insensitive. + /// For MCP, this is always "Bearer". + public let tokenType: String + + /// The lifetime in seconds of the access token from when it was issued. + public let expiresIn: Int? + + /// The scope of the access token as a space-delimited string. + public let scope: String? + + /// The refresh token for obtaining new access tokens. + public let refreshToken: String? + + public init( + accessToken: String, + tokenType: String = "Bearer", + expiresIn: Int? = nil, + scope: String? = nil, + refreshToken: String? = nil + ) { + self.accessToken = accessToken + self.tokenType = tokenType + self.expiresIn = expiresIn + self.scope = scope + self.refreshToken = refreshToken + } + + private enum CodingKeys: String, CodingKey { + case accessToken = "access_token" + case tokenType = "token_type" + case expiresIn = "expires_in" + case scope + case refreshToken = "refresh_token" + } +} + +/// Context provided when the server returns a 401 Unauthorized response. +/// +/// Contains information extracted from the `WWW-Authenticate` header +/// that guides the OAuth authorization flow. +public struct UnauthorizedContext: Sendable { + /// The URL to the Protected Resource Metadata (RFC 9728). + public let resourceMetadataURL: URL? + + /// The scope requested by the server. + public let scope: String? + + /// The full `WWW-Authenticate` header value for custom parsing. + public let wwwAuthenticate: String? + + public init( + resourceMetadataURL: URL? = nil, + scope: String? = nil, + wwwAuthenticate: String? = nil + ) { + self.resourceMetadataURL = resourceMetadataURL + self.scope = scope + self.wwwAuthenticate = wwwAuthenticate + } +} + +// MARK: - OAuth Provider Protocol + +/// A provider for OAuth 2.0 authentication in HTTP client transports. +/// +/// The transport calls this protocol's methods to obtain Bearer tokens +/// and handle authorization failures. Implementations manage their own +/// token storage and refresh logic. +/// +/// ## Transport Integration +/// +/// When the transport has an `authProvider`: +/// 1. Before each request: calls `tokens()` to get the access token +/// 2. On 401 response: calls `handleUnauthorized(context:)` to re-authenticate +/// 3. Retries the request with the new token +/// +/// ## SDK-Provided Implementations +/// +/// The SDK will provide implementations for common flows: +/// - Authorization code flow with PKCE (interactive) +/// - Client credentials flow (machine-to-machine) +/// - Private key JWT authentication +/// +/// Custom implementations can be created for specialized needs. +public protocol OAuthClientProvider: Sendable { + /// Returns the current OAuth tokens, refreshing if necessary. + /// + /// Implementations should: + /// - Return cached tokens if still valid + /// - Refresh expired tokens using the refresh token + /// - Return `nil` if not authenticated (triggers `handleUnauthorized`) + func tokens() async throws -> OAuthTokens? + + /// Handles a 401 Unauthorized response by performing authorization. + /// + /// Called when the server rejects the request. The context contains + /// information from the `WWW-Authenticate` header to guide the flow. + /// + /// - Parameter context: Information from the 401 response + /// - Returns: New tokens after successful authorization + func handleUnauthorized(context: UnauthorizedContext) async throws -> OAuthTokens +} From d458f99ce82a013f1c879223df5ee11e7ad9f873 Mon Sep 17 00:00:00 2001 From: Anthony DePasquale Date: Sun, 4 Jan 2026 16:06:16 +0100 Subject: [PATCH 8/8] Add missing functionality and fixes --- .../HummingbirdIntegration/Sources/main.swift | 15 +- Examples/VaporIntegration/Sources/main.swift | 13 +- Sources/MCP/Base/Messages.swift | 77 +- Sources/MCP/Base/Progress.swift | 26 + Sources/MCP/Base/Transport.swift | 55 + .../Base/Transports/HTTPClientTransport.swift | 14 +- .../HTTPServerTransport+Types.swift | 75 + .../Base/Transports/HTTPServerTransport.swift | 26 +- .../Base/Transports/InMemoryEventStore.swift | 2 + .../Base/Transports/InMemoryTransport.swift | 40 +- .../Base/Transports/NetworkTransport.swift | 33 +- .../MCP/Base/Transports/StdioTransport.swift | 4 +- Sources/MCP/Client/Client+Batching.swift | 6 +- .../MCP/Client/Client+MessageHandling.swift | 9 +- .../MCP/Client/Client+ProtocolMethods.swift | 102 +- Sources/MCP/Client/Client+Registration.swift | 22 +- Sources/MCP/Client/Client+Requests.swift | 6 +- Sources/MCP/Client/Client+Tasks.swift | 16 +- Sources/MCP/Client/Client.swift | 37 +- .../ExperimentalClientFeatures.swift | 8 +- Sources/MCP/Extensions/Data+Extensions.swift | 2 +- .../Tasks/ServerTaskContext.swift | 61 +- .../Experimental/Tasks/TaskContext.swift | 8 +- .../Server/Experimental/Tasks/TaskStore.swift | 10 +- .../Experimental/Tasks/TaskSupport.swift | 3 +- .../MCP/Server/Server+ClientRequests.swift | 2 +- .../MCP/Server/Server+RequestHandling.swift | 57 +- Sources/MCP/Server/Server+Sending.swift | 2 +- Sources/MCP/Server/Server.swift | 98 +- Sources/MCP/Server/SessionManager.swift | 4 + Tests/MCPTests/ClientTests.swift | 257 ++++ Tests/MCPTests/CompletionTests.swift | 48 +- Tests/MCPTests/ConcurrentExecutionTests.swift | 296 ++++ Tests/MCPTests/ErrorHandlingTests.swift | 795 +++++++++++ Tests/MCPTests/HTTPClientTransportTests.swift | 41 +- Tests/MCPTests/HTTPIntegrationTests.swift | 124 ++ Tests/MCPTests/Helpers/TestPayloads.swift | 22 + .../ImplementationMetadataTests.swift | 627 +++++++++ Tests/MCPTests/InMemoryEventStoreTests.swift | 209 +++ .../MCPTests/IntegrationRoundtripTests.swift | 1249 +++++++++++++++++ Tests/MCPTests/PrimingEventsTests.swift | 69 +- .../MCPTests/RequestHandlerContextTests.swift | 871 ++++++++++++ Tests/MCPTests/ResourceTests.swift | 20 +- Tests/MCPTests/RootsTests.swift | 6 +- Tests/MCPTests/RoundtripTests.swift | 10 +- Tests/MCPTests/ServerTests.swift | 230 +++ Tests/MCPTests/SessionLifecycleTests.swift | 560 ++++++++ .../StreamableHTTPServerTransportTests.swift | 27 + Tests/MCPTests/TaskModeValidationTests.swift | 515 +++++++ Tests/MCPTests/TaskTests.swift | 22 +- Tests/MCPTests/ToolTests.swift | 1 + 51 files changed, 6541 insertions(+), 291 deletions(-) create mode 100644 Tests/MCPTests/ConcurrentExecutionTests.swift create mode 100644 Tests/MCPTests/ErrorHandlingTests.swift create mode 100644 Tests/MCPTests/ImplementationMetadataTests.swift create mode 100644 Tests/MCPTests/IntegrationRoundtripTests.swift create mode 100644 Tests/MCPTests/SessionLifecycleTests.swift create mode 100644 Tests/MCPTests/TaskModeValidationTests.swift diff --git a/Examples/HummingbirdIntegration/Sources/main.swift b/Examples/HummingbirdIntegration/Sources/main.swift index 0ea7243c..b1eb9f04 100644 --- a/Examples/HummingbirdIntegration/Sources/main.swift +++ b/Examples/HummingbirdIntegration/Sources/main.swift @@ -31,6 +31,12 @@ import Hummingbird import Logging import MCP +// MARK: - Configuration + +/// Server bind address - using localhost enables automatic DNS rebinding protection +let serverHost = "localhost" +let serverPort = 3000 + // MARK: - Server Setup /// Create the MCP server (ONE instance for all clients) @@ -141,8 +147,11 @@ func handlePost(request: Request, context: MCPRequestContext) async throws -> Re let newSessionId = UUID().uuidString // Create new transport with session callbacks + // Using forBindAddress auto-configures DNS rebinding protection for localhost let newTransport = HTTPServerTransport( - options: .init( + options: .forBindAddress( + host: serverHost, + port: serverPort, sessionIdGenerator: { newSessionId }, onSessionInitialized: { sessionId in logger.info("Session initialized: \(sessionId)") @@ -322,10 +331,10 @@ struct HummingbirdMCPExample { // Create and run application let app = Application( router: router, - configuration: .init(address: .hostname("localhost", port: 3000)) + configuration: .init(address: .hostname(serverHost, port: serverPort)) ) - logger.info("Starting MCP server on http://localhost:3000/mcp") + logger.info("Starting MCP server on http://\(serverHost):\(serverPort)/mcp") logger.info("Available tools: echo, add") try await app.run() diff --git a/Examples/VaporIntegration/Sources/main.swift b/Examples/VaporIntegration/Sources/main.swift index fd6985e6..6f05c6ab 100644 --- a/Examples/VaporIntegration/Sources/main.swift +++ b/Examples/VaporIntegration/Sources/main.swift @@ -29,6 +29,12 @@ import Foundation import MCP import Vapor +// MARK: - Configuration + +/// Server bind address - using localhost enables automatic DNS rebinding protection +let serverHost = "localhost" +let serverPort = 8080 + // MARK: - Server Setup /// Create the MCP server (ONE instance for all clients) @@ -124,8 +130,11 @@ func handlePost(_ req: Vapor.Request) async throws -> Vapor.Response { let newSessionId = UUID().uuidString // Create new transport with session callbacks + // Using forBindAddress auto-configures DNS rebinding protection for localhost let newTransport = HTTPServerTransport( - options: .init( + options: .forBindAddress( + host: serverHost, + port: serverPort, sessionIdGenerator: { newSessionId }, onSessionInitialized: { sessionId in req.logger.info("Session initialized: \(sessionId)") @@ -270,7 +279,7 @@ struct VaporMCPExample { "OK" } - app.logger.info("Starting MCP server on http://localhost:8080/mcp") + app.logger.info("Starting MCP server on http://\(serverHost):\(serverPort)/mcp") app.logger.info("Available tools: echo, add") try await app.execute() diff --git a/Sources/MCP/Base/Messages.swift b/Sources/MCP/Base/Messages.swift index bc568023..478c15a3 100644 --- a/Sources/MCP/Base/Messages.swift +++ b/Sources/MCP/Base/Messages.swift @@ -131,9 +131,18 @@ extension Request { if M.Parameters.self is NotRequired.Type { // For NotRequired parameters, use decodeIfPresent or init() - params = - (try container.decodeIfPresent(M.Parameters.self, forKey: .params) - ?? (M.Parameters.self as! NotRequired.Type).init() as! M.Parameters) + if let decoded = try container.decodeIfPresent(M.Parameters.self, forKey: .params) { + params = decoded + } else if let notRequiredType = M.Parameters.self as? NotRequired.Type, + let defaultValue = notRequiredType.init() as? M.Parameters + { + params = defaultValue + } else { + throw DecodingError.dataCorrupted( + DecodingError.Context( + codingPath: container.codingPath, + debugDescription: "Failed to create default NotRequired parameters")) + } } else if let value = try? container.decode(M.Parameters.self, forKey: .params) { // If params exists and can be decoded, use it params = value @@ -142,8 +151,8 @@ extension Request { { // If params is missing or explicitly null, use Empty for Empty parameters // or throw for non-Empty parameters - if M.Parameters.self == Empty.self { - params = Empty() as! M.Parameters + if let emptyValue = Empty() as? M.Parameters { + params = emptyValue } else { throw DecodingError.dataCorrupted( DecodingError.Context( @@ -172,14 +181,22 @@ extension AnyRequest { } } -/// A box for request handlers that can be type-erased +/// A box for request handlers that can be type-erased. +/// +/// This class uses `@unchecked Sendable` because Swift cannot automatically infer +/// `Sendable` for non-final classes. However, this is safe because: +/// - The only subclass (`TypedRequestHandler`) stores only an immutable `@Sendable` closure +/// - No mutable state exists in either class after initialization +/// - The closure is `let` and marked `@Sendable` class RequestHandlerBox: @unchecked Sendable { func callAsFunction(_ request: AnyRequest, context: Server.RequestHandlerContext) async throws -> AnyResponse { fatalError("Must override") } } -/// A typed request handler that can be used to handle requests of a specific type +/// A typed request handler that can be used to handle requests of a specific type. +/// +/// See `RequestHandlerBox` for why `@unchecked Sendable` is safe here. final class TypedRequestHandler: RequestHandlerBox, @unchecked Sendable { private let _handle: @Sendable (Request, Server.RequestHandlerContext) async throws -> Response @@ -359,9 +376,18 @@ public struct Message: NotificationMessageProtocol, Hashable, C if N.Parameters.self is NotRequired.Type { // For NotRequired parameters, use decodeIfPresent or init() - params = - (try container.decodeIfPresent(N.Parameters.self, forKey: .params) - ?? (N.Parameters.self as! NotRequired.Type).init() as! N.Parameters) + if let decoded = try container.decodeIfPresent(N.Parameters.self, forKey: .params) { + params = decoded + } else if let notRequiredType = N.Parameters.self as? NotRequired.Type, + let defaultValue = notRequiredType.init() as? N.Parameters + { + params = defaultValue + } else { + throw DecodingError.dataCorrupted( + DecodingError.Context( + codingPath: container.codingPath, + debugDescription: "Failed to create default NotRequired parameters")) + } } else if let value = try? container.decode(N.Parameters.self, forKey: .params) { // If params exists and can be decoded, use it params = value @@ -370,8 +396,8 @@ public struct Message: NotificationMessageProtocol, Hashable, C { // If params is missing or explicitly null, use Empty for Empty parameters // or throw for non-Empty parameters - if N.Parameters.self == Empty.self { - params = Empty() as! N.Parameters + if let emptyValue = Empty() as? N.Parameters { + params = emptyValue } else { throw DecodingError.dataCorrupted( DecodingError.Context( @@ -411,12 +437,20 @@ extension Notification { } } -/// A box for notification handlers that can be type-erased +/// A box for notification handlers that can be type-erased. +/// +/// This class uses `@unchecked Sendable` because Swift cannot automatically infer +/// `Sendable` for non-final classes. However, this is safe because: +/// - The only subclass (`TypedNotificationHandler`) stores only an immutable `@Sendable` closure +/// - No mutable state exists in either class after initialization +/// - The closure is `let` and marked `@Sendable` class NotificationHandlerBox: @unchecked Sendable { func callAsFunction(_ notification: Message) async throws {} } -/// A typed notification handler that can be used to handle notifications of a specific type +/// A typed notification handler that can be used to handle notifications of a specific type. +/// +/// See `NotificationHandlerBox` for why `@unchecked Sendable` is safe here. final class TypedNotificationHandler: NotificationHandlerBox, @unchecked Sendable { @@ -438,14 +472,22 @@ final class TypedNotificationHandler: NotificationHandlerBox, // MARK: - Client Request Handlers -/// A box for client request handlers that can be type-erased +/// A box for client request handlers that can be type-erased. +/// +/// This class uses `@unchecked Sendable` because Swift cannot automatically infer +/// `Sendable` for non-final classes. However, this is safe because: +/// - The only subclass (`TypedClientRequestHandler`) stores only an immutable `@Sendable` closure +/// - No mutable state exists in either class after initialization +/// - The closure is `let` and marked `@Sendable` class ClientRequestHandlerBox: @unchecked Sendable { func callAsFunction(_ request: AnyRequest, context: Client.RequestHandlerContext) async throws -> AnyResponse { fatalError("Must override") } } -/// A typed client request handler that can be used to handle requests of a specific type +/// A typed client request handler that can be used to handle requests of a specific type. +/// +/// See `ClientRequestHandlerBox` for why `@unchecked Sendable` is safe here. final class TypedClientRequestHandler: ClientRequestHandlerBox, @unchecked Sendable { private let _handle: @Sendable (M.Parameters, Client.RequestHandlerContext) async throws -> M.Result @@ -473,7 +515,8 @@ final class TypedClientRequestHandler: ClientRequestHandlerBox, @unch } catch let error as MCPError { return Response(id: typedRequest.id, error: error) } catch { - return Response(id: typedRequest.id, error: MCPError.internalError(error.localizedDescription)) + // Sanitize non-MCP errors to avoid leaking internal details + return Response(id: typedRequest.id, error: MCPError.internalError("An internal error occurred")) } } } diff --git a/Sources/MCP/Base/Progress.swift b/Sources/MCP/Base/Progress.swift index 41835731..3a4c53a6 100644 --- a/Sources/MCP/Base/Progress.swift +++ b/Sources/MCP/Base/Progress.swift @@ -26,6 +26,32 @@ public struct RequestMeta: Hashable, Codable, Sendable { self.additionalFields = additionalFields } + // MARK: - Convenience Accessors + + /// The related task ID, if present. + /// + /// Extracts the task ID from `_meta["io.modelcontextprotocol/related-task"].taskId`. + /// This matches the TypeScript SDK's `_meta[RELATED_TASK_META_KEY]?.taskId`. + /// + /// ## Example + /// + /// ```swift + /// if let taskId = context._meta?.relatedTaskId { + /// print("Request is part of task: \(taskId)") + /// } + /// ``` + /// + /// - Note: For the full `RelatedTaskMetadata` struct, use the experimental tasks API. + public var relatedTaskId: String? { + guard let metaValue = additionalFields?["io.modelcontextprotocol/related-task"], + case .object(let dict) = metaValue, + let taskIdValue = dict["taskId"], + let taskId = taskIdValue.stringValue else { + return nil + } + return taskId + } + private enum CodingKeys: String, CodingKey { case progressToken } diff --git a/Sources/MCP/Base/Transport.swift b/Sources/MCP/Base/Transport.swift index 0a3054cf..60378be1 100644 --- a/Sources/MCP/Base/Transport.swift +++ b/Sources/MCP/Base/Transport.swift @@ -4,6 +4,50 @@ import struct Foundation.Data // MARK: - Message Context Types +/// Information about the incoming HTTP request. +/// +/// This is the Swift equivalent of TypeScript's `RequestInfo` interface, which +/// provides access to HTTP request headers for request handlers. +/// +/// ## Example +/// +/// ```swift +/// server.withRequestHandler(CallTool.self) { params, context in +/// if let requestInfo = context.requestInfo { +/// // Access custom headers +/// if let customHeader = requestInfo.headers["X-Custom-Header"] { +/// print("Custom header: \(customHeader)") +/// } +/// } +/// return CallTool.Result(content: [.text("Done")]) +/// } +/// ``` +public struct RequestInfo: Hashable, Sendable { + /// The HTTP headers from the request. + /// + /// Header names are preserved as provided by the HTTP framework. + /// Use case-insensitive comparison when looking up headers. + public let headers: [String: String] + + public init(headers: [String: String]) { + self.headers = headers + } + + /// Get a header value (case-insensitive lookup). + /// + /// - Parameter name: The header name to look up + /// - Returns: The header value, or nil if not found + public func header(_ name: String) -> String? { + let lowercased = name.lowercased() + for (key, value) in headers { + if key.lowercased() == lowercased { + return value + } + } + return nil + } +} + /// Context information associated with a received message. /// /// This is the Swift equivalent of TypeScript's `MessageExtraInfo`, which is passed @@ -20,6 +64,15 @@ public struct MessageContext: Sendable { /// access this via `context.authInfo`. public let authInfo: AuthInfo? + /// Information about the incoming HTTP request. + /// + /// Contains HTTP headers from the original request. Only available for + /// HTTP transports. Request handlers can access this via `context.requestInfo`. + /// + /// This matches TypeScript SDK's `extra.requestInfo` and allows handlers + /// to inspect custom headers for authentication, client identification, etc. + public let requestInfo: RequestInfo? + /// Closes the SSE stream for this request, triggering client reconnection. /// /// Only available when using HTTPServerTransport with eventStore configured. @@ -33,10 +86,12 @@ public struct MessageContext: Sendable { public init( authInfo: AuthInfo? = nil, + requestInfo: RequestInfo? = nil, closeSSEStream: (@Sendable () async -> Void)? = nil, closeStandaloneSSEStream: (@Sendable () async -> Void)? = nil ) { self.authInfo = authInfo + self.requestInfo = requestInfo self.closeSSEStream = closeSSEStream self.closeStandaloneSSEStream = closeStandaloneSSEStream } diff --git a/Sources/MCP/Base/Transports/HTTPClientTransport.swift b/Sources/MCP/Base/Transports/HTTPClientTransport.swift index 4a99803e..11c2ad24 100644 --- a/Sources/MCP/Base/Transports/HTTPClientTransport.swift +++ b/Sources/MCP/Base/Transports/HTTPClientTransport.swift @@ -25,7 +25,15 @@ import Logging /// - Regular HTTP (`streaming=false`): Simple request/response pattern /// - Streaming HTTP with SSE (`streaming=true`): Enables server-to-client push messages /// -/// - Important: Server-Sent Events (SSE) functionality is not supported on Linux platforms. +/// ## Linux Platform Limitations +/// +/// SSE functionality is unavailable on Linux because `URLSession.AsyncBytes` is not yet +/// implemented in swift-corelibs-foundation (see [swift#57548](https://github.com/swiftlang/swift/issues/57548)). +/// +/// **What works:** HTTP POST requests and JSON responses (tool calls, resource reads, prompts). +/// +/// **What doesn't work:** Server-initiated push notifications, streaming responses, and +/// stream resumability. On Linux, set `streaming: false` to avoid warnings. /// /// ## Example Usage /// @@ -161,8 +169,8 @@ public actor HTTPClientTransport: Transport { self.authProvider = authProvider // Create message stream - var continuation: AsyncThrowingStream.Continuation! - self.messageStream = AsyncThrowingStream { continuation = $0 } + let (stream, continuation) = AsyncThrowingStream.makeStream() + self.messageStream = stream self.messageContinuation = continuation self.logger = diff --git a/Sources/MCP/Base/Transports/HTTPServerTransport+Types.swift b/Sources/MCP/Base/Transports/HTTPServerTransport+Types.swift index 1e8ff55b..a0e000a2 100644 --- a/Sources/MCP/Base/Transports/HTTPServerTransport+Types.swift +++ b/Sources/MCP/Base/Transports/HTTPServerTransport+Types.swift @@ -70,6 +70,16 @@ public struct AuthInfo: Hashable, Codable, Sendable { } } +extension AuthInfo: CustomStringConvertible { + /// Redacts the token to prevent accidental exposure in logs. + /// + /// The token is still accessible via the `token` property for legitimate use, + /// but this prevents it from appearing in string interpolation or print statements. + public var description: String { + "AuthInfo(clientId: \(clientId), scopes: \(scopes), token: [REDACTED])" + } +} + // MARK: - Transport Options /// Configuration options for HTTPServerTransport @@ -130,6 +140,71 @@ public struct HTTPServerTransportOptions: Sendable { self.retryInterval = retryInterval self.security = security } + + /// Creates options with automatic security configuration based on bind address. + /// + /// This factory method follows the same convention as the TypeScript and Python SDKs: + /// - For localhost addresses (`127.0.0.1`, `localhost`, `::1`), DNS rebinding protection + /// is automatically enabled with appropriate allowed hosts/origins. + /// - For other addresses (e.g., `0.0.0.0`), no automatic security is configured. + /// + /// ## Example + /// + /// ```swift + /// // Auto-configures DNS rebinding protection for localhost + /// let options = HTTPServerTransportOptions.forBindAddress( + /// host: "127.0.0.1", + /// port: 8080, + /// sessionIdGenerator: { UUID().uuidString } + /// ) + /// + /// // No automatic protection for 0.0.0.0 - configure manually if needed + /// let options = HTTPServerTransportOptions.forBindAddress( + /// host: "0.0.0.0", + /// port: 8080, + /// security: TransportSecuritySettings( + /// enableDnsRebindingProtection: true, + /// allowedHosts: ["myserver.local:8080"], + /// allowedOrigins: ["http://myserver.local:8080"] + /// ) + /// ) + /// ``` + /// + /// - Parameters: + /// - host: The host address the server will bind to + /// - port: The port number + /// - sessionIdGenerator: Function that generates session IDs (nil for stateless mode) + /// - onSessionInitialized: Called when a new session is initialized + /// - onSessionClosed: Called when a session is closed + /// - enableJsonResponse: If true, return JSON responses instead of SSE streams + /// - eventStore: Event store for resumability support + /// - retryInterval: Retry interval in milliseconds for SSE + /// - security: Override the auto-configured security settings + /// - Returns: Configured transport options + public static func forBindAddress( + host: String, + port: Int, + sessionIdGenerator: (@Sendable () -> String)? = nil, + onSessionInitialized: (@Sendable (String) async -> Void)? = nil, + onSessionClosed: (@Sendable (String) async -> Void)? = nil, + enableJsonResponse: Bool = false, + eventStore: EventStore? = nil, + retryInterval: Int? = nil, + security: TransportSecuritySettings? = nil + ) -> HTTPServerTransportOptions { + // Auto-configure security for localhost if not explicitly provided + let effectiveSecurity = security ?? TransportSecuritySettings.forBindAddress(host: host, port: port) + + return HTTPServerTransportOptions( + sessionIdGenerator: sessionIdGenerator, + onSessionInitialized: onSessionInitialized, + onSessionClosed: onSessionClosed, + enableJsonResponse: enableJsonResponse, + eventStore: eventStore, + retryInterval: retryInterval, + security: effectiveSecurity + ) + } } /// Security settings for DNS rebinding protection. diff --git a/Sources/MCP/Base/Transports/HTTPServerTransport.swift b/Sources/MCP/Base/Transports/HTTPServerTransport.swift index ee691f8c..1e93b990 100644 --- a/Sources/MCP/Base/Transports/HTTPServerTransport.swift +++ b/Sources/MCP/Base/Transports/HTTPServerTransport.swift @@ -111,8 +111,8 @@ public actor HTTPServerTransport: Transport { ) // Create server receive stream - var continuation: AsyncThrowingStream.Continuation! - self.serverStream = AsyncThrowingStream { continuation = $0 } + let (stream, continuation) = AsyncThrowingStream.makeStream() + self.serverStream = stream self.serverContinuation = continuation } @@ -169,7 +169,7 @@ public actor HTTPServerTransport: Transport { return } - guard let requestId = requestId else { return } + guard let requestId else { return } // Get the stream for this request guard let streamId = requestToStreamMapping[requestId] else { @@ -422,7 +422,8 @@ public actor HTTPServerTransport: Transport { if !hasRequests { // Only notifications - yield to server and return 202 // Notifications don't need SSE closures since there's no response stream - let context = MessageContext(authInfo: authInfo) + let requestInfo = RequestInfo(headers: request.headers) + let context = MessageContext(authInfo: authInfo, requestInfo: requestInfo) serverContinuation.yield(TransportMessage(data: body, context: context)) return HTTPResponse(statusCode: 202, headers: sessionHeaders()) } @@ -438,7 +439,7 @@ public actor HTTPServerTransport: Transport { // Check if using JSON response mode if options.enableJsonResponse { - return await handleJsonResponseMode(streamId: streamId, requestIds: requestIds, body: body, authInfo: authInfo) + return await handleJsonResponseMode(streamId: streamId, requestIds: requestIds, body: body, request: request, authInfo: authInfo) } // SSE streaming mode @@ -456,6 +457,7 @@ public actor HTTPServerTransport: Transport { streamId: String, requestIds: [RequestId], body: Data, + request: HTTPRequest, authInfo: AuthInfo? ) async -> HTTPResponse { // Create stream for receiving the response @@ -465,7 +467,8 @@ public actor HTTPServerTransport: Transport { jsonStreamMapping[streamId] = state // JSON response mode doesn't have SSE streams to close - let context = MessageContext(authInfo: authInfo) + let requestInfo = RequestInfo(headers: request.headers) + let context = MessageContext(authInfo: authInfo, requestInfo: requestInfo) serverContinuation.yield(TransportMessage(data: body, context: context)) // Wait for response - this is cancellation-aware unlike withCheckedContinuation @@ -533,9 +536,11 @@ public actor HTTPServerTransport: Transport { await self?.closeStandaloneSSEStream() } - // Create context with auth info and SSE closures + // Create context with auth info, request info, and SSE closures + let requestInfo = RequestInfo(headers: request.headers) let context = MessageContext( authInfo: authInfo, + requestInfo: requestInfo, closeSSEStream: closeSSEStreamClosure, closeStandaloneSSEStream: closeStandaloneSSEStreamClosure ) @@ -739,8 +744,7 @@ public actor HTTPServerTransport: Transport { } // Validate Host header (required when protection is enabled) - let hostHeader = request.header(HTTPHeader.host) - if hostHeader == nil { + guard let hostHeader = request.header(HTTPHeader.host) else { logger.warning("DNS rebinding protection: Missing Host header") // Use 421 Misdirected Request for Host header issues return createJsonErrorResponse( @@ -751,13 +755,13 @@ public actor HTTPServerTransport: Transport { } let hostMatches = security.allowedHosts.contains { pattern in - matchesHostPattern(hostHeader!, pattern: pattern) + matchesHostPattern(hostHeader, pattern: pattern) } if !hostMatches { logger.warning( "DNS rebinding protection: Host header rejected", - metadata: ["host": "\(hostHeader!)"] + metadata: ["host": "\(hostHeader)"] ) // Use 421 Misdirected Request for Host header issues return createJsonErrorResponse( diff --git a/Sources/MCP/Base/Transports/InMemoryEventStore.swift b/Sources/MCP/Base/Transports/InMemoryEventStore.swift index 74359554..8f4d76b4 100644 --- a/Sources/MCP/Base/Transports/InMemoryEventStore.swift +++ b/Sources/MCP/Base/Transports/InMemoryEventStore.swift @@ -45,6 +45,8 @@ import Foundation /// /// - **Not persistent**: Events are lost when the process restarts /// - **Single process**: Cannot be shared across multiple server instances +/// - **Unbounded stream count**: While events per stream are limited, the number of streams +/// is unbounded. Use middleware or infrastructure-level controls to limit connections. /// /// For production deployments, implement `EventStore` with a persistent backend like /// Redis, PostgreSQL, or another appropriate storage system. diff --git a/Sources/MCP/Base/Transports/InMemoryTransport.swift b/Sources/MCP/Base/Transports/InMemoryTransport.swift index 195a85dd..e8a41318 100644 --- a/Sources/MCP/Base/Transports/InMemoryTransport.swift +++ b/Sources/MCP/Base/Transports/InMemoryTransport.swift @@ -22,9 +22,9 @@ public actor InMemoryTransport: Transport { private var isConnected = false private var pairedTransport: InMemoryTransport? - // Message queues - private var incomingMessages: [Data] = [] - private var messageContinuation: AsyncThrowingStream.Continuation? + // Message stream + private let messageStream: AsyncThrowingStream + private let messageContinuation: AsyncThrowingStream.Continuation /// Creates a new in-memory transport /// @@ -36,6 +36,11 @@ public actor InMemoryTransport: Transport { label: "mcp.transport.in-memory", factory: { _ in SwiftLogNoOpLogHandler() } ) + + // Create message stream + let (stream, continuation) = AsyncThrowingStream.makeStream() + self.messageStream = stream + self.messageContinuation = continuation } /// Creates a connected pair of in-memory transports @@ -115,8 +120,7 @@ public actor InMemoryTransport: Transport { guard isConnected else { return } isConnected = false - messageContinuation?.finish() - messageContinuation = nil + messageContinuation.finish() // Notify paired transport of disconnection if let paired = pairedTransport { @@ -129,8 +133,7 @@ public actor InMemoryTransport: Transport { /// Handles disconnection from the paired transport private func handlePeerDisconnection() { if isConnected { - messageContinuation?.finish(throwing: MCPError.connectionClosed) - messageContinuation = nil + messageContinuation.finish(throwing: MCPError.connectionClosed) isConnected = false logger.info("Peer transport disconnected") } @@ -166,32 +169,13 @@ public actor InMemoryTransport: Transport { } logger.debug("Message received", metadata: ["size": "\(data.count)"]) - - if let continuation = messageContinuation { - continuation.yield(TransportMessage(data: data)) - } else { - // Queue message if stream not yet created - incomingMessages.append(data) - } + messageContinuation.yield(TransportMessage(data: data)) } /// Receives messages from the paired transport /// /// - Returns: An AsyncThrowingStream of TransportMessage objects representing messages public func receive() -> AsyncThrowingStream { - return AsyncThrowingStream { continuation in - self.messageContinuation = continuation - - // Deliver any queued messages - for message in self.incomingMessages { - continuation.yield(TransportMessage(data: message)) - } - self.incomingMessages.removeAll() - - // Check if already disconnected - if !self.isConnected { - continuation.finish() - } - } + return messageStream } } diff --git a/Sources/MCP/Base/Transports/NetworkTransport.swift b/Sources/MCP/Base/Transports/NetworkTransport.swift index 9c1cdd9a..3c15750a 100644 --- a/Sources/MCP/Base/Transports/NetworkTransport.swift +++ b/Sources/MCP/Base/Transports/NetworkTransport.swift @@ -242,7 +242,16 @@ import Logging private let messageStream: AsyncThrowingStream private let messageContinuation: AsyncThrowingStream.Continuation - // Connection is marked nonisolated(unsafe) to allow access from closures + /// The underlying network connection. + /// + /// This property uses `nonisolated(unsafe)` because `NWConnection` is not `Sendable`, + /// but its callback-based APIs (`stateUpdateHandler`, `send`, `receive`) run on `.main` + /// queue and need to access the connection from outside actor isolation. + /// + /// This is safe because: + /// - The connection reference is never reassigned after initialization + /// - `NWConnection` is designed to be thread-safe when used with a consistent queue + /// - All callbacks are dispatched to `.main` queue via `connection.start(queue: .main)` private nonisolated(unsafe) var connection: NetworkConnectionProtocol /// Logger instance for transport-related events @@ -296,8 +305,8 @@ import Logging self.bufferConfig = bufferConfig // Create message stream - var continuation: AsyncThrowingStream.Continuation! - self.messageStream = AsyncThrowingStream { continuation = $0 } + let (stream, continuation) = AsyncThrowingStream.makeStream() + self.messageStream = stream self.messageContinuation = continuation } @@ -382,7 +391,7 @@ import Logging // Start a new heartbeat task heartbeatTask = Task { [weak self] in - guard let self = self else { return } + guard let self else { return } // Initial delay before starting heartbeats try? await Task.sleep(for: .seconds(1)) @@ -413,7 +422,7 @@ import Logging // Try to send the heartbeat (without the newline delimiter used for normal messages) try await withCheckedThrowingContinuation { [weak self] (continuation: CheckedContinuation) in - guard let self = self else { + guard let self else { continuation.resume(throwing: MCPError.internalError("Transport deallocated")) return } @@ -423,7 +432,7 @@ import Logging contentContext: .defaultMessage, isComplete: true, completion: .contentProcessed { [weak self] error in - if let error = error { + if let error { continuation.resume(throwing: error) } else { Task { [weak self] in @@ -475,7 +484,7 @@ import Logging try await withCheckedThrowingContinuation { [weak self] (continuation: CheckedContinuation) in - guard let self = self else { + guard let self else { continuation.resume(throwing: MCPError.internalError("Transport deallocated")) return } @@ -485,9 +494,9 @@ import Logging contentContext: .defaultMessage, isComplete: true, completion: .contentProcessed { [weak self] error in - guard let self = self else { return } + guard let self else { return } - if let error = error { + if let error { self.logger.error("Send error: \(error)") // Schedule reconnection attempt if connection lost @@ -686,7 +695,7 @@ import Logging private func receiveData() async throws -> Data { try await withCheckedThrowingContinuation { [weak self] (continuation: CheckedContinuation) in - guard let self = self else { + guard let self else { continuation.resume(throwing: MCPError.internalError("Transport deallocated")) return } @@ -694,9 +703,9 @@ import Logging let maxLength = bufferConfig.maxReceiveBufferSize ?? Int.max connection.receive(minimumIncompleteLength: 1, maximumLength: maxLength) { [weak self] content, _, isComplete, error in - if let error = error { + if let error { continuation.resume(throwing: MCPError.transportError(error)) - } else if let content = content { + } else if let content { continuation.resume(returning: content) } else if isComplete { self?.logger.trace("Connection completed by peer") diff --git a/Sources/MCP/Base/Transports/StdioTransport.swift b/Sources/MCP/Base/Transports/StdioTransport.swift index bb21aa9d..7d9e49e9 100644 --- a/Sources/MCP/Base/Transports/StdioTransport.swift +++ b/Sources/MCP/Base/Transports/StdioTransport.swift @@ -77,8 +77,8 @@ import struct Foundation.Data factory: { _ in SwiftLogNoOpLogHandler() }) // Create message stream - var continuation: AsyncThrowingStream.Continuation! - self.messageStream = AsyncThrowingStream { continuation = $0 } + let (stream, continuation) = AsyncThrowingStream.makeStream() + self.messageStream = stream self.messageContinuation = continuation } diff --git a/Sources/MCP/Client/Client+Batching.swift b/Sources/MCP/Client/Client+Batching.swift index 427b893c..160f3248 100644 --- a/Sources/MCP/Client/Client+Batching.swift +++ b/Sources/MCP/Client/Client+Batching.swift @@ -115,11 +115,11 @@ extension Client { /// /// // Await the results after the batch is sent /// do { - /// if let pingTask = pingTask { + /// if let pingTask { /// try await pingTask.value // Await ping result (throws if ping failed) /// print("Ping successful") /// } - /// if let promptTask = promptTask { + /// if let promptTask { /// let promptResult = try await promptTask.value // Await prompt result /// print("Prompt description: \(promptResult.description ?? "None")") /// } @@ -133,7 +133,7 @@ extension Client { /// - Throws: `MCPError.internalError` if the client is not connected. /// Can also rethrow errors from the `body` closure or from sending the batch request. public func withBatch(body: @escaping (Batch) async throws -> Void) async throws { - guard let connection = connection else { + guard let connection else { throw MCPError.internalError("Client connection not initialized") } diff --git a/Sources/MCP/Client/Client+MessageHandling.swift b/Sources/MCP/Client/Client+MessageHandling.swift index 9dee9db5..83e65382 100644 --- a/Sources/MCP/Client/Client+MessageHandling.swift +++ b/Sources/MCP/Client/Client+MessageHandling.swift @@ -323,9 +323,10 @@ extension Client { "method": "\(request.method)", "error": "\(error)", ]) + // Error already logged above - sanitize for response let errorResponse = AnyMethod.response( id: request.id, - error: (error as? MCPError) ?? MCPError.internalError(error.localizedDescription) + error: (error as? MCPError) ?? MCPError.internalError("An internal error occurred") ) await sendResponse(errorResponse) } @@ -412,7 +413,9 @@ extension Client { } catch let error as MCPError { return Response(id: request.id, error: error) } catch { - return Response(id: request.id, error: MCPError.internalError(error.localizedDescription)) + // Log full error for debugging, but sanitize for response + await logger?.error("Task handler error", metadata: ["error": "\(error)"]) + return Response(id: request.id, error: MCPError.internalError("An internal error occurred")) } // Not a task-augmented request @@ -421,7 +424,7 @@ extension Client { /// Send a response back to the server. func sendResponse(_ response: Response) async { - guard let connection = connection else { + guard let connection else { await logger?.warning("Cannot send response - client not connected") return } diff --git a/Sources/MCP/Client/Client+ProtocolMethods.swift b/Sources/MCP/Client/Client+ProtocolMethods.swift index dbb23517..dbc56a99 100644 --- a/Sources/MCP/Client/Client+ProtocolMethods.swift +++ b/Sources/MCP/Client/Client+ProtocolMethods.swift @@ -3,104 +3,129 @@ import Foundation extension Client { // MARK: - Prompts + /// Get a prompt by name. + /// + /// - Parameters: + /// - name: The name of the prompt to retrieve. + /// - arguments: Optional arguments to pass to the prompt. + /// - Returns: The prompt result containing description and messages. public func getPrompt(name: String, arguments: [String: String]? = nil) async throws - -> (description: String?, messages: [Prompt.Message]) + -> GetPrompt.Result { try validateServerCapability(\.prompts, "Prompts") let request = GetPrompt.request(.init(name: name, arguments: arguments)) - let result = try await send(request) - return (description: result.description, messages: result.messages) + return try await send(request) } - public func listPrompts(cursor: String? = nil) async throws - -> (prompts: [Prompt], nextCursor: String?) - { + /// List available prompts from the server. + /// + /// - Parameter cursor: Optional cursor for pagination. + /// - Returns: The list result containing prompts and optional next cursor. + public func listPrompts(cursor: String? = nil) async throws -> ListPrompts.Result { try validateServerCapability(\.prompts, "Prompts") let request: Request - if let cursor = cursor { + if let cursor { request = ListPrompts.request(.init(cursor: cursor)) } else { request = ListPrompts.request(.init()) } - let result = try await send(request) - return (prompts: result.prompts, nextCursor: result.nextCursor) + return try await send(request) } // MARK: - Resources - public func readResource(uri: String) async throws -> [Resource.Content] { + /// Read a resource by URI. + /// + /// - Parameter uri: The URI of the resource to read. + /// - Returns: The read result containing resource contents. + public func readResource(uri: String) async throws -> ReadResource.Result { try validateServerCapability(\.resources, "Resources") let request = ReadResource.request(.init(uri: uri)) - let result = try await send(request) - return result.contents + return try await send(request) } - public func listResources(cursor: String? = nil) async throws -> ( - resources: [Resource], nextCursor: String? - ) { + /// List available resources from the server. + /// + /// - Parameter cursor: Optional cursor for pagination. + /// - Returns: The list result containing resources and optional next cursor. + public func listResources(cursor: String? = nil) async throws -> ListResources.Result { try validateServerCapability(\.resources, "Resources") let request: Request - if let cursor = cursor { + if let cursor { request = ListResources.request(.init(cursor: cursor)) } else { request = ListResources.request(.init()) } - let result = try await send(request) - return (resources: result.resources, nextCursor: result.nextCursor) + return try await send(request) } + /// Subscribe to updates for a resource. + /// + /// - Parameter uri: The URI of the resource to subscribe to. public func subscribeToResource(uri: String) async throws { try validateServerCapability(\.resources?.subscribe, "Resource subscription") let request = ResourceSubscribe.request(.init(uri: uri)) _ = try await send(request) } + /// Unsubscribe from updates for a resource. + /// + /// - Parameter uri: The URI of the resource to unsubscribe from. public func unsubscribeFromResource(uri: String) async throws { try validateServerCapability(\.resources?.subscribe, "Resource subscription") let request = ResourceUnsubscribe.request(.init(uri: uri)) _ = try await send(request) } - public func listResourceTemplates(cursor: String? = nil) async throws -> ( - templates: [Resource.Template], nextCursor: String? - ) { + /// List available resource templates from the server. + /// + /// - Parameter cursor: Optional cursor for pagination. + /// - Returns: The list result containing templates and optional next cursor. + public func listResourceTemplates(cursor: String? = nil) async throws + -> ListResourceTemplates.Result + { try validateServerCapability(\.resources, "Resources") let request: Request - if let cursor = cursor { + if let cursor { request = ListResourceTemplates.request(.init(cursor: cursor)) } else { request = ListResourceTemplates.request(.init()) } - let result = try await send(request) - return (templates: result.templates, nextCursor: result.nextCursor) + return try await send(request) } // MARK: - Tools - public func listTools(cursor: String? = nil) async throws -> ( - tools: [Tool], nextCursor: String? - ) { + /// List available tools from the server. + /// + /// - Parameter cursor: Optional cursor for pagination. + /// - Returns: The list result containing tools and optional next cursor. + public func listTools(cursor: String? = nil) async throws -> ListTools.Result { try validateServerCapability(\.tools, "Tools") let request: Request - if let cursor = cursor { + if let cursor { request = ListTools.request(.init(cursor: cursor)) } else { request = ListTools.request(.init()) } - let result = try await send(request) - return (tools: result.tools, nextCursor: result.nextCursor) + return try await send(request) } - public func callTool(name: String, arguments: [String: Value]? = nil) async throws -> ( - content: [Tool.Content], structuredContent: Value?, isError: Bool? - ) { + /// Call a tool by name. + /// + /// - Parameters: + /// - name: The name of the tool to call. + /// - arguments: Optional arguments to pass to the tool. + /// - Returns: The tool call result containing content, structured content, and error flag. + public func callTool(name: String, arguments: [String: Value]? = nil) async throws + -> CallTool.Result + { try validateServerCapability(\.tools, "Tools") let request = CallTool.request(.init(name: name, arguments: arguments)) - let result = try await send(request) // TODO: Add client-side output validation against the tool's outputSchema. // TypeScript and Python SDKs cache tool outputSchemas from listTools() and // validate structuredContent when receiving tool results. - return (content: result.content, structuredContent: result.structuredContent, isError: result.isError) + return try await send(request) } // MARK: - Completions @@ -114,16 +139,15 @@ extension Client { /// - ref: A reference to the prompt or resource template to get completions for. /// - argument: The argument being completed, including its name and partial value. /// - context: Optional additional context with previously-resolved argument values. - /// - Returns: The completion suggestions from the server. + /// - Returns: The completion result from the server. public func complete( ref: CompletionReference, argument: CompletionArgument, context: CompletionContext? = nil - ) async throws -> CompletionSuggestions { + ) async throws -> Complete.Result { try validateServerCapability(\.completions, "Completions") let request = Complete.request(.init(ref: ref, argument: argument, context: context)) - let result = try await send(request) - return result.completion + return try await send(request) } // MARK: - Logging diff --git a/Sources/MCP/Client/Client+Registration.swift b/Sources/MCP/Client/Client+Registration.swift index 2a5b4173..ff985a2a 100644 --- a/Sources/MCP/Client/Client+Registration.swift +++ b/Sources/MCP/Client/Client+Registration.swift @@ -118,19 +118,33 @@ extension Client { /// /// - Important: The client must have declared `roots` capability during initialization. /// - /// - Parameter handler: A closure that returns the list of available roots. + /// ## Example + /// + /// ```swift + /// client.withRootsHandler { context in + /// // Access request context if needed + /// print("Request ID: \(context.requestId)") + /// + /// return [ + /// Root(uri: "file:///home/user/project", name: "Project"), + /// Root(uri: "file:///home/user/docs", name: "Documents") + /// ] + /// } + /// ``` + /// + /// - Parameter handler: A closure that receives the request context and returns the list of available roots. /// - Returns: Self for chaining. /// - Precondition: `capabilities.roots` must be non-nil. @discardableResult public func withRootsHandler( - _ handler: @escaping @Sendable () async throws -> [Root] + _ handler: @escaping @Sendable (RequestHandlerContext) async throws -> [Root] ) -> Self { precondition( capabilities.roots != nil, "Cannot register roots handler: Client does not have roots capability" ) - return withRequestHandler(ListRoots.self) { _, _ in - ListRoots.Result(roots: try await handler()) + return withRequestHandler(ListRoots.self) { _, context in + ListRoots.Result(roots: try await handler(context)) } } diff --git a/Sources/MCP/Client/Client+Requests.swift b/Sources/MCP/Client/Client+Requests.swift index 61b369bd..82a702ff 100644 --- a/Sources/MCP/Client/Client+Requests.swift +++ b/Sources/MCP/Client/Client+Requests.swift @@ -90,7 +90,7 @@ extension Client { _ request: Request, options: RequestOptions? ) async throws -> M.Result { - guard let connection = connection else { + guard let connection else { throw MCPError.internalError("Client connection not initialized") } @@ -223,7 +223,7 @@ extension Client { options: RequestOptions?, onProgress: @escaping ProgressCallback ) async throws -> M.Result { - guard let connection = connection else { + guard let connection else { throw MCPError.internalError("Client connection not initialized") } @@ -488,7 +488,7 @@ extension Client { /// This is called when a client Task waiting for a response is cancelled. /// The notification is sent on a best-effort basis - failures are logged but not thrown. func sendCancellationNotification(requestId: RequestId, reason: String?) async { - guard let connection = connection else { + guard let connection else { await logger?.debug( "Cannot send cancellation notification - connection is nil", metadata: ["requestId": "\(requestId)"] diff --git a/Sources/MCP/Client/Client+Tasks.swift b/Sources/MCP/Client/Client+Tasks.swift index bfcefc87..12b87931 100644 --- a/Sources/MCP/Client/Client+Tasks.swift +++ b/Sources/MCP/Client/Client+Tasks.swift @@ -10,7 +10,7 @@ extension Client { return try await send(request) } - func listTasks(cursor: String? = nil) async throws -> (tasks: [MCPTask], nextCursor: String?) { + func listTasks(cursor: String? = nil) async throws -> ListTasks.Result { try validateServerCapability(\.tasks, "Tasks") let request: Request if let cursor { @@ -18,8 +18,7 @@ extension Client { } else { request = ListTasks.request(.init()) } - let result = try await send(request) - return (tasks: result.tasks, nextCursor: result.nextCursor) + return try await send(request) } func cancelTask(taskId: String) async throws -> CancelTask.Result { @@ -75,7 +74,7 @@ extension Client { // The server should return CreateTaskResult for task-augmented requests // We need to decode as CreateTaskResult instead of CallTool.Result - guard let connection = connection else { + guard let connection else { throw MCPError.internalError("Client connection not initialized") } @@ -152,7 +151,7 @@ extension Client { name: String, arguments: [String: Value]? = nil, ttl: Int? = nil - ) async throws -> (content: [Tool.Content], isError: Bool?) { + ) async throws -> CallTool.Result { // Start the task let createResult = try await callToolAsTask(name: name, arguments: arguments, ttl: ttl) let taskId = createResult.task.taskId @@ -169,8 +168,7 @@ extension Client { // Convert extraFields back to Value for decoding let resultValue = Value.object(extraFields) let resultData = try encoder.encode(resultValue) - let toolResult = try decoder.decode(CallTool.Result.self, from: resultData) - return (content: toolResult.content, isError: toolResult.isError) + return try decoder.decode(CallTool.Result.self, from: resultData) } func callToolStream( @@ -241,7 +239,9 @@ extension Client { continuation.yield(.error(error)) continuation.finish() } catch { - let mcpError = MCPError.internalError(error.localizedDescription) + // Log full error for debugging, but sanitize for stream consumer + await logger?.error("Task stream error", metadata: ["error": "\(error)"]) + let mcpError = MCPError.internalError("An internal error occurred") continuation.yield(.error(mcpError)) continuation.finish() } diff --git a/Sources/MCP/Client/Client.swift b/Sources/MCP/Client/Client.swift index 4b8fc240..2fcde84e 100644 --- a/Sources/MCP/Client/Client.swift +++ b/Sources/MCP/Client/Client.swift @@ -215,6 +215,30 @@ public actor Client { /// ``` public let _meta: RequestMeta? + /// The task ID for task-augmented requests, if present. + /// + /// This is a convenience property that extracts the task ID from the + /// `_meta["io.modelcontextprotocol/related-task"]` field. When a server + /// sends a task-augmented elicitation or sampling request, this property + /// will contain the associated task ID. + /// + /// This matches the TypeScript SDK's `extra.taskId` and aligns with + /// `Server.RequestHandlerContext.taskId`. + /// + /// ## Example + /// + /// ```swift + /// client.withElicitationHandler { params, context in + /// if let taskId = context.taskId { + /// print("Handling elicitation for task: \(taskId)") + /// } + /// return ElicitResult(action: .accept, content: [:]) + /// } + /// ``` + public var taskId: String? { + _meta?.relatedTaskId + } + // MARK: - Convenience Methods /// Send a progress notification to the server. @@ -536,9 +560,20 @@ public actor Client { public init( name: String, version: String, + title: String? = nil, + description: String? = nil, + icons: [Icon]? = nil, + websiteUrl: String? = nil, configuration: Configuration = .default ) { - self.clientInfo = Client.Info(name: name, version: version) + self.clientInfo = Client.Info( + name: name, + version: version, + title: title, + description: description, + icons: icons, + websiteUrl: websiteUrl + ) self.capabilities = Capabilities() self.configuration = configuration } diff --git a/Sources/MCP/Client/Experimental/ExperimentalClientFeatures.swift b/Sources/MCP/Client/Experimental/ExperimentalClientFeatures.swift index 4c341ea3..750de6de 100644 --- a/Sources/MCP/Client/Experimental/ExperimentalClientFeatures.swift +++ b/Sources/MCP/Client/Experimental/ExperimentalClientFeatures.swift @@ -97,9 +97,9 @@ public struct ExperimentalClientTasks: Sendable { /// List all tasks. /// /// - Parameter cursor: Optional pagination cursor - /// - Returns: Tuple of (tasks, nextCursor). nextCursor is nil if no more pages. + /// - Returns: The list result containing tasks and optional next cursor. /// - Throws: MCPError if the server doesn't support tasks - public func listTasks(cursor: String? = nil) async throws -> (tasks: [MCPTask], nextCursor: String?) { + public func listTasks(cursor: String? = nil) async throws -> ListTasks.Result { try await client.listTasks(cursor: cursor) } @@ -277,13 +277,13 @@ public struct ExperimentalClientTasks: Sendable { /// - name: The name of the tool to call /// - arguments: Optional arguments for the tool /// - ttl: Optional time-to-live in milliseconds for the task result - /// - Returns: The tool result (same as `callTool()`) + /// - Returns: The tool call result. /// - Throws: MCPError if the request fails or the task fails public func callToolAsTaskAndWait( name: String, arguments: [String: Value]? = nil, ttl: Int? = nil - ) async throws -> (content: [Tool.Content], isError: Bool?) { + ) async throws -> CallTool.Result { try await client.callToolAsTaskAndWait(name: name, arguments: arguments, ttl: ttl) } diff --git a/Sources/MCP/Extensions/Data+Extensions.swift b/Sources/MCP/Extensions/Data+Extensions.swift index 745b4a30..0a953f2d 100644 --- a/Sources/MCP/Extensions/Data+Extensions.swift +++ b/Sources/MCP/Extensions/Data+Extensions.swift @@ -55,7 +55,7 @@ extension Data { // Process MIME type var mimeType = mediatype.isEmpty ? "text/plain" : String(mediatype) - if let charset = charset, !charset.isEmpty, mimeType.starts(with: "text/") { + if let charset, !charset.isEmpty, mimeType.starts(with: "text/") { mimeType += ";charset=\(charset)" } diff --git a/Sources/MCP/Server/Experimental/Tasks/ServerTaskContext.swift b/Sources/MCP/Server/Experimental/Tasks/ServerTaskContext.swift index cac8bd27..2e532c06 100644 --- a/Sources/MCP/Server/Experimental/Tasks/ServerTaskContext.swift +++ b/Sources/MCP/Server/Experimental/Tasks/ServerTaskContext.swift @@ -1,4 +1,5 @@ import Foundation +import os // MARK: - Server Task Context @@ -39,9 +40,19 @@ import Foundation /// return CallTool.Result(content: [.text("Done!")]) /// } /// ``` -public final class ServerTaskContext: @unchecked Sendable { - /// The task this context is for. - public private(set) var task: MCPTask +public final class ServerTaskContext: Sendable { + /// Mutable state protected by a lock. + /// + /// This state may be accessed concurrently - for example, the task handler + /// reads `isCancelled` while another context calls `requestCancellation()`. + private struct State: Sendable { + var task: MCPTask + var isCancelled: Bool = false + var requestIdCounter: Int = 0 + } + + /// Lock-protected mutable state. + private let state: OSAllocatedUnfairLock /// The task store for persistence. private let store: any TaskStore @@ -55,17 +66,20 @@ public final class ServerTaskContext: @unchecked Sendable { /// Server reference for task-augmented requests (elicitAsTask, createMessageAsTask). private let server: Server? - /// Counter for generating request IDs. - private var requestIdCounter: Int = 0 - - /// Whether cancellation has been requested. - private var _isCancelled = false + /// The task this context is for. + public var task: MCPTask { + state.withLock { $0.task } + } /// Check if cancellation has been requested. - public var isCancelled: Bool { _isCancelled } + public var isCancelled: Bool { + state.withLock { $0.isCancelled } + } /// The task ID. - public var taskId: String { task.taskId } + public var taskId: String { + state.withLock { $0.task.taskId } + } /// Create a server task context. /// @@ -82,7 +96,7 @@ public final class ServerTaskContext: @unchecked Sendable { clientCapabilities: Client.Capabilities? = nil, server: Server? = nil ) { - self.task = task + self.state = OSAllocatedUnfairLock(initialState: State(task: task)) self.store = store self.queue = queue self.clientCapabilities = clientCapabilities @@ -91,8 +105,11 @@ public final class ServerTaskContext: @unchecked Sendable { /// Generate a unique request ID for queued requests. private func nextRequestId() -> RequestId { - requestIdCounter += 1 - return .string("task-\(taskId)-req-\(requestIdCounter)") + let counter = state.withLock { state -> Int in + state.requestIdCounter += 1 + return state.requestIdCounter + } + return .string("task-\(taskId)-req-\(counter)") } /// Request cancellation of the task. @@ -100,7 +117,7 @@ public final class ServerTaskContext: @unchecked Sendable { /// This sets the `isCancelled` flag but doesn't immediately stop execution. /// Task handlers should check this flag periodically and exit gracefully. public func requestCancellation() { - _isCancelled = true + state.withLock { $0.isCancelled = true } } /// Update the task status with a message. @@ -118,7 +135,7 @@ public final class ServerTaskContext: @unchecked Sendable { status: .working, statusMessage: message ) - task = updatedTask + state.withLock { $0.task = updatedTask } if notify { await sendStatusNotification() } @@ -139,7 +156,7 @@ public final class ServerTaskContext: @unchecked Sendable { status: .inputRequired, statusMessage: message ) - task = updatedTask + state.withLock { $0.task = updatedTask } if notify { await sendStatusNotification() } @@ -160,7 +177,7 @@ public final class ServerTaskContext: @unchecked Sendable { status: .completed, statusMessage: nil ) - task = updatedTask + state.withLock { $0.task = updatedTask } if notify { await sendStatusNotification() } @@ -196,7 +213,7 @@ public final class ServerTaskContext: @unchecked Sendable { status: .failed, statusMessage: error ) - task = updatedTask + state.withLock { $0.task = updatedTask } if notify { await sendStatusNotification() } @@ -204,12 +221,18 @@ public final class ServerTaskContext: @unchecked Sendable { /// Fail the task with an Error. /// + /// For security, non-MCP errors are sanitized to avoid leaking internal details. + /// Use ``fail(error:notify:)-6k8lh`` with a string message if you need to send + /// specific error information to clients. + /// /// - Parameters: /// - error: The error that caused the failure /// - notify: Whether to send a `TaskStatusNotification` to the client (default: true) /// - Throws: Error if the task cannot be updated public func fail(error: any Error, notify: Bool = true) async throws { - try await fail(error: error.localizedDescription, notify: notify) + // Sanitize non-MCP errors to avoid leaking internal details to clients + let message = (error as? MCPError)?.message ?? "An internal error occurred" + try await fail(error: message, notify: notify) } /// Send a task status notification to the client. diff --git a/Sources/MCP/Server/Experimental/Tasks/TaskContext.swift b/Sources/MCP/Server/Experimental/Tasks/TaskContext.swift index dff9913c..910b6080 100644 --- a/Sources/MCP/Server/Experimental/Tasks/TaskContext.swift +++ b/Sources/MCP/Server/Experimental/Tasks/TaskContext.swift @@ -163,10 +163,16 @@ public actor TaskContext { /// Fail the task with an Error. /// + /// For security, non-MCP errors are sanitized to avoid leaking internal details. + /// Use ``fail(error:)-swift.method`` with a string message if you need to send + /// specific error information to clients. + /// /// - Parameter error: The error that caused the failure /// - Throws: Error if the task cannot be updated public func fail(error: any Error) async throws { - try await fail(error: error.localizedDescription) + // Sanitize non-MCP errors to avoid leaking internal details to clients + let message = (error as? MCPError)?.message ?? "An internal error occurred" + try await fail(error: message) } /// Cancel the task. diff --git a/Sources/MCP/Server/Experimental/Tasks/TaskStore.swift b/Sources/MCP/Server/Experimental/Tasks/TaskStore.swift index bc17ab85..0f77783f 100644 --- a/Sources/MCP/Server/Experimental/Tasks/TaskStore.swift +++ b/Sources/MCP/Server/Experimental/Tasks/TaskStore.swift @@ -51,8 +51,8 @@ public protocol TaskStore: Sendable { /// List tasks with pagination. /// /// - Parameter cursor: Optional cursor for pagination - /// - Returns: Tuple of (tasks, nextCursor). nextCursor is nil if no more pages. - func listTasks(cursor: String?) async -> (tasks: [MCPTask], nextCursor: String?) + /// - Returns: The list result containing tasks and optional next cursor. + func listTasks(cursor: String?) async -> ListTasks.Result /// Delete a task. /// @@ -243,7 +243,7 @@ public actor InMemoryTaskStore: TaskStore { tasks[taskId]?.result } - public func listTasks(cursor: String?) async -> (tasks: [MCPTask], nextCursor: String?) { + public func listTasks(cursor: String?) async -> ListTasks.Result { cleanUpExpired() let allTaskIds = Array(tasks.keys).sorted() @@ -264,7 +264,7 @@ public actor InMemoryTaskStore: TaskStore { nil } - return (tasks: pageTasks, nextCursor: nextCursor) + return ListTasks.Result(tasks: pageTasks, nextCursor: nextCursor) } public func deleteTask(taskId: String) async -> Bool { @@ -293,6 +293,8 @@ public actor InMemoryTaskStore: TaskStore { let waiterId = UUID() try await withTaskCancellationHandler { + // Check early to avoid creating a waiter that will be immediately cancelled + try Task.checkCancellation() try await withCheckedThrowingContinuation { (continuation: CheckedContinuation) in waiters[taskId, default: []].append(Waiter(id: waiterId, continuation: continuation)) } diff --git a/Sources/MCP/Server/Experimental/Tasks/TaskSupport.swift b/Sources/MCP/Server/Experimental/Tasks/TaskSupport.swift index 93fff668..a348fba7 100644 --- a/Sources/MCP/Server/Experimental/Tasks/TaskSupport.swift +++ b/Sources/MCP/Server/Experimental/Tasks/TaskSupport.swift @@ -246,8 +246,7 @@ extension Server { // tasks/list - List all tasks withRequestHandler(ListTasks.self) { params, _ in - let (tasks, nextCursor) = await taskSupport.store.listTasks(cursor: params.cursor) - return ListTasks.Result(tasks: tasks, nextCursor: nextCursor) + await taskSupport.store.listTasks(cursor: params.cursor) } // tasks/cancel - Cancel a running task diff --git a/Sources/MCP/Server/Server+ClientRequests.swift b/Sources/MCP/Server/Server+ClientRequests.swift index caf235b3..7dbaee7d 100644 --- a/Sources/MCP/Server/Server+ClientRequests.swift +++ b/Sources/MCP/Server/Server+ClientRequests.swift @@ -11,7 +11,7 @@ extension Server { /// - Parameter request: The request to send /// - Returns: The result from the client public func sendRequest(_ request: Request) async throws -> M.Result { - guard let connection = connection else { + guard let connection else { throw MCPError.internalError("Server connection not initialized") } diff --git a/Sources/MCP/Server/Server+RequestHandling.swift b/Sources/MCP/Server/Server+RequestHandling.swift index 7194b038..6a7ee417 100644 --- a/Sources/MCP/Server/Server+RequestHandling.swift +++ b/Sources/MCP/Server/Server+RequestHandling.swift @@ -60,8 +60,13 @@ extension Server { } catch { // Only add errors to response for requests (notifications don't have responses) if case .request(let request) = item { + // Log full error for debugging, but sanitize for client response. + // Only log non-MCP errors since MCP errors are expected/user-facing. + if !(error is MCPError) { + await logger?.error("Error handling batch item", metadata: ["error": "\(error)"]) + } let mcpError = - error as? MCPError ?? MCPError.internalError(error.localizedDescription) + error as? MCPError ?? MCPError.internalError("An internal error occurred") responses.append(AnyMethod.response(id: request.id, error: mcpError)) } } @@ -110,6 +115,11 @@ extension Server { /// /// Set by HTTP transports when OAuth or other authentication is in use. let authInfo: AuthInfo? + /// Information about the incoming HTTP request. + /// + /// Contains HTTP headers from the original request. Only available for + /// HTTP transports. This matches TypeScript SDK's `extra.requestInfo`. + let requestInfo: RequestInfo? /// Closure to close the SSE stream for this request. /// /// Only set by HTTP transports with SSE support. @@ -199,8 +209,9 @@ extension Server { let requestMeta = extractMeta(from: request.params) // Extract context from transport message (set by HTTP transports with per-message context) - // This pattern aligns with TypeScript's onmessage(message, { authInfo, closeSSEStream, ... }) + // This pattern aligns with TypeScript's onmessage(message, { authInfo, requestInfo, closeSSEStream, ... }) let authInfo = messageContext?.authInfo + let requestInfo = messageContext?.requestInfo let closeSSEStream = messageContext?.closeSSEStream let closeStandaloneSSEStream = messageContext?.closeStandaloneSSEStream @@ -210,6 +221,7 @@ extension Server { sessionId: await capturedConnection?.sessionId, meta: requestMeta, authInfo: authInfo, + requestInfo: requestInfo, closeSSEStream: closeSSEStream, closeStandaloneSSEStream: closeStandaloneSSEStream ) @@ -230,14 +242,26 @@ extension Server { "id": "\(request.id)", ]) + // Check initialization state for strict mode (matches Python SDK behavior). + // We chose to align with Python (block at Server level) rather than TypeScript + // (block only at HTTP transport level) for consistent behavior across all transports. if configuration.strict { - // The client SHOULD NOT send requests other than pings - // before the server has responded to the initialize request. switch request.method { - case Initialize.name, Ping.name: - break - default: - try checkInitialized() + case Initialize.name, Ping.name: + // Always allow initialize and ping requests + break + default: + guard isInitialized else { + let error = MCPError.invalidRequest("Server is not initialized") + let response = AnyMethod.response(id: request.id, error: error) + + if sendResponse { + try await send(response, using: context) + return nil + } + + return response + } } } @@ -294,6 +318,7 @@ extension Server { requestId: context.requestId, _meta: context.meta, authInfo: context.authInfo, + requestInfo: context.requestInfo, closeSSEStream: context.closeSSEStream, closeStandaloneSSEStream: context.closeStandaloneSSEStream, shouldSendLogMessage: { [weak self, context] level in @@ -383,7 +408,11 @@ extension Server { return nil } - let mcpError = error as? MCPError ?? MCPError.internalError(error.localizedDescription) + // Log full error for debugging, but sanitize for client response + if !(error is MCPError) { + await logger?.error("Request handler error", metadata: ["error": "\(error)"]) + } + let mcpError = error as? MCPError ?? MCPError.internalError("An internal error occurred") let response: Response = AnyMethod.response(id: request.id, error: mcpError) if sendResponse { @@ -400,10 +429,14 @@ extension Server { "Processing notification", metadata: ["method": "\(message.method)"]) + // Check initialization state for strict mode (matches Python SDK behavior). + // For notifications (unlike requests), we log and ignore since no response is expected. if configuration.strict { - // Check initialization state unless this is an initialized notification - if message.method != InitializedNotification.name { - try checkInitialized() + if message.method != InitializedNotification.name && !isInitialized { + await logger?.warning( + "Ignoring notification before initialization", + metadata: ["method": "\(message.method)"]) + return } } diff --git a/Sources/MCP/Server/Server+Sending.swift b/Sources/MCP/Server/Server+Sending.swift index 4567aa4e..67dfb59e 100644 --- a/Sources/MCP/Server/Server+Sending.swift +++ b/Sources/MCP/Server/Server+Sending.swift @@ -18,7 +18,7 @@ extension Server { /// Send a notification to connected clients public func notify(_ notification: Message) async throws { - guard let connection = connection else { + guard let connection else { throw MCPError.internalError("Server connection not initialized") } diff --git a/Sources/MCP/Server/Server.swift b/Sources/MCP/Server/Server.swift index 829766f3..81269389 100644 --- a/Sources/MCP/Server/Server.swift +++ b/Sources/MCP/Server/Server.swift @@ -77,20 +77,35 @@ import class Foundation.JSONEncoder public actor Server { /// The server configuration public struct Configuration: Hashable, Codable, Sendable { - /// The default configuration. - public static let `default` = Configuration(strict: false) + /// The default configuration (strict mode enabled). + /// + /// This matches Python SDK behavior where the server rejects non-ping requests + /// before initialization at the session level. TypeScript SDK only enforces this + /// at the HTTP transport level, not at the server/session level. + /// + /// We chose to align with Python because: + /// - Consistent behavior across all transports (stdio, HTTP, in-memory) + /// - More defensive - prevents misbehaving clients from accessing functionality before init + /// - Better aligns with MCP spec intent (clients "SHOULD NOT" send requests before init) + /// - Ping is still allowed for health checks + public static let `default` = Configuration(strict: true) - /// The strict configuration. - public static let strict = Configuration(strict: true) + /// The lenient configuration (strict mode disabled). + /// + /// Use this for compatibility with non-compliant clients that send requests + /// before initialization. This matches TypeScript SDK's server-level behavior. + public static let lenient = Configuration(strict: false) - /// When strict mode is enabled, the server: + /// When strict mode is enabled (default), the server: /// - Requires clients to send an initialize request before any other requests - /// - Rejects all requests from uninitialized clients with a protocol error + /// - Allows ping requests before initialization (for health checks) + /// - Rejects all other requests from uninitialized clients with a protocol error + /// + /// The MCP specification says clients "SHOULD NOT" send requests other than + /// pings before initialization. Strict mode enforces this at the server level. /// - /// While the MCP specification requires clients to initialize the connection - /// before sending other requests, some implementations may not follow this. - /// Disabling strict mode allows the server to be more lenient with non-compliant - /// clients, though this may lead to undefined behavior. + /// Set to `false` for lenient behavior that allows requests before initialization. + /// This may be useful for non-compliant clients but can lead to undefined behavior. public var strict: Bool } @@ -307,6 +322,27 @@ public actor Server { /// ``` public let _meta: RequestMeta? + /// The task ID for task-augmented requests, if present. + /// + /// This is a convenience property that extracts the task ID from the + /// `_meta["io.modelcontextprotocol/related-task"]` field. + /// + /// This matches the TypeScript SDK's `extra.taskId`. + /// + /// ## Example + /// + /// ```swift + /// server.withRequestHandler(CallTool.self) { params, context in + /// if let taskId = context.taskId { + /// print("Handling request as part of task: \(taskId)") + /// } + /// return CallTool.Result(content: [.text("Done")]) + /// } + /// ``` + public var taskId: String? { + _meta?.relatedTaskId + } + /// Authentication information for this request. /// /// Contains validated access token information when using HTTP transports @@ -332,6 +368,28 @@ public actor Server { /// ``` public let authInfo: AuthInfo? + /// Information about the incoming HTTP request. + /// + /// Contains HTTP headers from the original request. Only available for + /// HTTP transports. + /// + /// This matches the TypeScript SDK's `extra.requestInfo`. + /// + /// ## Example + /// + /// ```swift + /// server.withRequestHandler(CallTool.self) { params, context in + /// if let requestInfo = context.requestInfo { + /// // Access custom headers + /// if let apiVersion = requestInfo.header("X-API-Version") { + /// print("Client API version: \(apiVersion)") + /// } + /// } + /// return CallTool.Result(content: [.text("Done")]) + /// } + /// ``` + public let requestInfo: RequestInfo? + /// Closes the SSE stream for this request, triggering client reconnection. /// /// Only available when using StreamableHTTPServerTransport with eventStore configured. @@ -775,11 +833,22 @@ public actor Server { public init( name: String, version: String, + title: String? = nil, + description: String? = nil, + icons: [Icon]? = nil, + websiteUrl: String? = nil, instructions: String? = nil, capabilities: Server.Capabilities = .init(), configuration: Configuration = .default ) { - self.serverInfo = Server.Info(name: name, version: version) + self.serverInfo = Server.Info( + name: name, + version: version, + title: title, + description: description, + icons: icons, + websiteUrl: websiteUrl + ) self.capabilities = capabilities self.configuration = configuration self.instructions = instructions @@ -878,10 +947,11 @@ public actor Server { // handles it internally. Message handling code won't throw EAGAIN. await logger?.error( "Error processing message", metadata: ["error": "\(error)"]) + // Sanitize non-MCP errors to avoid leaking internal details to clients let response = AnyMethod.response( id: requestID ?? .random, error: error as? MCPError - ?? MCPError.internalError(error.localizedDescription) + ?? MCPError.internalError("An internal error occurred") ) try? await send(response) } @@ -908,7 +978,7 @@ public actor Server { task?.cancel() task = nil - if let connection = connection { + if let connection { await connection.disconnect() } connection = nil @@ -1011,7 +1081,7 @@ public actor Server { ) { // Initialize withRequestHandler(Initialize.self) { [weak self] params, _ in - guard let self = self else { + guard let self else { throw MCPError.internalError("Server was deallocated") } diff --git a/Sources/MCP/Server/SessionManager.swift b/Sources/MCP/Server/SessionManager.swift index 62d84159..0e833675 100644 --- a/Sources/MCP/Server/SessionManager.swift +++ b/Sources/MCP/Server/SessionManager.swift @@ -91,6 +91,10 @@ public actor SessionManager { /// Stores a transport for a session ID. /// + /// - Important: If ``maxSessions`` is configured, check ``canAddSession()`` before + /// calling this method for new sessions. This method does not enforce capacity limits, + /// allowing flexibility in how applications handle capacity (reject, queue, evict oldest, etc.). + /// /// - Parameters: /// - transport: The transport to store /// - sessionId: The session ID diff --git a/Tests/MCPTests/ClientTests.swift b/Tests/MCPTests/ClientTests.swift index 5fa39701..6d3274b0 100644 --- a/Tests/MCPTests/ClientTests.swift +++ b/Tests/MCPTests/ClientTests.swift @@ -944,4 +944,261 @@ struct ClientTests { await client.disconnect() } + + // MARK: - Initialization Request Tests + // Based on TypeScript SDK: should initialize with matching protocol version + // Based on Python SDK: test_client_session_initialize + + @Test("Client sends latest protocol version in initialize request") + func testClientSendsLatestProtocolVersion() async throws { + // TypeScript SDK: should initialize with matching protocol version + // Python SDK: test_client_session_version_negotiation_success + // Verifies that the client sends the latest protocol version in its initialize request + let transport = MockTransport() + let client = Client(name: "TestClient", version: "1.0") + + // Set up a task to handle the initialize response + let initTask = Task { + try await Task.sleep(for: .milliseconds(10)) + if let lastMessage = await transport.sentMessages.last, + let data = lastMessage.data(using: .utf8), + let request = try? JSONDecoder().decode(Request.self, from: data) + { + // Verify the client sent the latest protocol version + #expect(request.params.protocolVersion == Version.latest) + + let response = Initialize.response( + id: request.id, + result: .init( + protocolVersion: Version.latest, + capabilities: .init(), + serverInfo: .init(name: "TestServer", version: "1.0"), + instructions: nil + ) + ) + try await transport.queue(response: response) + } + } + + defer { initTask.cancel() } + + let result = try await client.connect(transport: transport) + #expect(result.protocolVersion == Version.latest) + + await client.disconnect() + } + + @Test("Client info is correctly sent in initialize request") + func testClientInfoSentInInitializeRequest() async throws { + // Python SDK: test_client_session_custom_client_info, test_client_session_default_client_info + // Verifies that the client's name and version are correctly included in the initialize request + let transport = MockTransport() + let clientName = "CustomTestClient" + let clientVersion = "2.3.4" + let client = Client(name: clientName, version: clientVersion) + + // Set up a task to handle the initialize response and verify client info + let initTask = Task { + try await Task.sleep(for: .milliseconds(10)) + if let lastMessage = await transport.sentMessages.last, + let data = lastMessage.data(using: .utf8), + let request = try? JSONDecoder().decode(Request.self, from: data) + { + // Verify the client info in the request + #expect(request.params.clientInfo.name == clientName) + #expect(request.params.clientInfo.version == clientVersion) + + let response = Initialize.response( + id: request.id, + result: .init( + protocolVersion: Version.latest, + capabilities: .init(), + serverInfo: .init(name: "TestServer", version: "1.0"), + instructions: nil + ) + ) + try await transport.queue(response: response) + } + } + + defer { initTask.cancel() } + + try await client.connect(transport: transport) + await client.disconnect() + } + + @Test("Client capabilities are sent in initialize request") + func testClientCapabilitiesSentInInitializeRequest() async throws { + // Python SDK: test_client_capabilities_default, test_client_capabilities_with_custom_callbacks + // Verifies that client capabilities are correctly included in the initialize request + let transport = MockTransport() + let client = Client(name: "TestClient", version: "1.0") + + // Set client capabilities with roots and sampling + await client.setCapabilities(.init( + sampling: .init(), + roots: .init(listChanged: true) + )) + + // Set up a task to handle the initialize response and verify capabilities + let initTask = Task { + try await Task.sleep(for: .milliseconds(10)) + if let lastMessage = await transport.sentMessages.last, + let data = lastMessage.data(using: .utf8), + let request = try? JSONDecoder().decode(Request.self, from: data) + { + // Verify the client capabilities in the request + #expect(request.params.capabilities.sampling != nil) + #expect(request.params.capabilities.roots != nil) + #expect(request.params.capabilities.roots?.listChanged == true) + + let response = Initialize.response( + id: request.id, + result: .init( + protocolVersion: Version.latest, + capabilities: .init(), + serverInfo: .init(name: "TestServer", version: "1.0"), + instructions: nil + ) + ) + try await transport.queue(response: response) + } + } + + defer { initTask.cancel() } + + try await client.connect(transport: transport) + await client.disconnect() + } + + @Test("Server capabilities accessible after initialization") + func testServerCapabilitiesAccessibleAfterInit() async throws { + // Python SDK: test_get_server_capabilities + // Verifies that getServerCapabilities() returns nil before connect and is populated after + let transport = MockTransport() + let client = Client(name: "TestClient", version: "1.0") + + // Before connect, capabilities should be nil + #expect(await client.getServerCapabilities() == nil) + + // Create server capabilities with various features + let serverCapabilities = Server.Capabilities( + logging: .init(), + prompts: .init(listChanged: true), + resources: .init(subscribe: true, listChanged: true), + tools: .init(listChanged: false) + ) + + // Set up a task to handle the initialize response + let initTask = Task { + try await Task.sleep(for: .milliseconds(10)) + if let lastMessage = await transport.sentMessages.last, + let data = lastMessage.data(using: .utf8), + let request = try? JSONDecoder().decode(Request.self, from: data) + { + let response = Initialize.response( + id: request.id, + result: .init( + protocolVersion: Version.latest, + capabilities: serverCapabilities, + serverInfo: .init(name: "TestServer", version: "1.0"), + instructions: nil + ) + ) + try await transport.queue(response: response) + } + } + + defer { initTask.cancel() } + + try await client.connect(transport: transport) + + // After connect, capabilities should be populated + let capabilities = await client.getServerCapabilities() + #expect(capabilities != nil) + #expect(capabilities?.prompts?.listChanged == true) + #expect(capabilities?.resources?.subscribe == true) + #expect(capabilities?.resources?.listChanged == true) + #expect(capabilities?.tools?.listChanged == false) + #expect(capabilities?.logging != nil) + + await client.disconnect() + } + + @Test("Instructions from server accessible in initialize result") + func testInstructionsAccessibleFromInitializeResult() async throws { + // TypeScript SDK: should initialize with matching protocol version (checks getInstructions()) + // Python SDK: test_client_session_initialize (checks result.instructions) + // Verifies that instructions from the server's response are accessible + let transport = MockTransport() + let client = Client(name: "TestClient", version: "1.0") + let serverInstructions = "These are the server instructions for the client." + + // Set up a task to handle the initialize response with instructions + let initTask = Task { + try await Task.sleep(for: .milliseconds(10)) + if let lastMessage = await transport.sentMessages.last, + let data = lastMessage.data(using: .utf8), + let request = try? JSONDecoder().decode(Request.self, from: data) + { + let response = Initialize.response( + id: request.id, + result: .init( + protocolVersion: Version.latest, + capabilities: .init(), + serverInfo: .init(name: "TestServer", version: "1.0"), + instructions: serverInstructions + ) + ) + try await transport.queue(response: response) + } + } + + defer { initTask.cancel() } + + // The result from connect contains the instructions + let result = try await client.connect(transport: transport) + #expect(result.instructions == serverInstructions) + + await client.disconnect() + } + + @Test("Server info accessible in initialize result") + func testServerInfoAccessibleFromInitializeResult() async throws { + // TypeScript SDK: should connect new client to old, supported server version (checks getServerVersion()) + // Python SDK: test_client_session_initialize (checks result.serverInfo) + // Verifies that server info from the response is accessible + let transport = MockTransport() + let client = Client(name: "TestClient", version: "1.0") + let serverName = "CustomMCPServer" + let serverVersion = "3.2.1" + + // Set up a task to handle the initialize response + let initTask = Task { + try await Task.sleep(for: .milliseconds(10)) + if let lastMessage = await transport.sentMessages.last, + let data = lastMessage.data(using: .utf8), + let request = try? JSONDecoder().decode(Request.self, from: data) + { + let response = Initialize.response( + id: request.id, + result: .init( + protocolVersion: Version.latest, + capabilities: .init(), + serverInfo: .init(name: serverName, version: serverVersion), + instructions: nil + ) + ) + try await transport.queue(response: response) + } + } + + defer { initTask.cancel() } + + let result = try await client.connect(transport: transport) + #expect(result.serverInfo.name == serverName) + #expect(result.serverInfo.version == serverVersion) + + await client.disconnect() + } } diff --git a/Tests/MCPTests/CompletionTests.swift b/Tests/MCPTests/CompletionTests.swift index 8dfaa0fa..8d6d8ad4 100644 --- a/Tests/MCPTests/CompletionTests.swift +++ b/Tests/MCPTests/CompletionTests.swift @@ -518,7 +518,7 @@ struct CompletionTests { let receivedContext = await received.getContext() #expect(receivedContext != nil) #expect(receivedContext?.arguments?["previous"] == "value") - #expect(result.values == ["test-completion"]) + #expect(result.completion.values == ["test-completion"]) // Verify the ref and argument were received correctly let receivedRef = await received.getRef() @@ -570,7 +570,7 @@ struct CompletionTests { ) #expect(await received.wasContextNil()) - #expect(result.values == ["no-context-completion"]) + #expect(result.completion.values == ["no-context-completion"]) await client.disconnect() await server.stop() @@ -636,8 +636,8 @@ struct CompletionTests { ref: .resource(ResourceTemplateReference(uri: "db://{database}/{table}")), argument: CompletionArgument(name: "database", value: "") ) - #expect(dbResult.values.contains("users_db")) - #expect(dbResult.values.contains("products_db")) + #expect(dbResult.completion.values.contains("users_db")) + #expect(dbResult.completion.values.contains("products_db")) // Then complete table with database context let tableResult = try await client.complete( @@ -645,7 +645,7 @@ struct CompletionTests { argument: CompletionArgument(name: "table", value: ""), context: CompletionContext(arguments: ["database": "users_db"]) ) - #expect(tableResult.values == ["users", "sessions", "permissions"]) + #expect(tableResult.completion.values == ["users", "sessions", "permissions"]) // Different database gives different tables let tableResult2 = try await client.complete( @@ -653,7 +653,7 @@ struct CompletionTests { argument: CompletionArgument(name: "table", value: ""), context: CompletionContext(arguments: ["database": "products_db"]) ) - #expect(tableResult2.values == ["products", "categories", "inventory"]) + #expect(tableResult2.completion.values == ["products", "categories", "inventory"]) await client.disconnect() await server.stop() @@ -717,7 +717,7 @@ struct CompletionTests { argument: CompletionArgument(name: "table", value: ""), context: CompletionContext(arguments: ["database": "test_db"]) ) - #expect(result.values == ["users", "orders", "products"]) + #expect(result.completion.values == ["users", "orders", "products"]) await client.disconnect() await server.stop() @@ -763,7 +763,7 @@ struct CompletionTests { argument: CompletionArgument(name: "language", value: "py") ) - #expect(result.values == ["python"]) + #expect(result.completion.values == ["python"]) // Request completion with "j" prefix let result2 = try await client.complete( @@ -771,9 +771,9 @@ struct CompletionTests { argument: CompletionArgument(name: "language", value: "j") ) - #expect(result2.values.contains("javascript")) - #expect(result2.values.contains("java")) - #expect(result2.values.allSatisfy { $0.hasPrefix("j") }) + #expect(result2.completion.values.contains("javascript")) + #expect(result2.completion.values.contains("java")) + #expect(result2.completion.values.allSatisfy { $0.hasPrefix("j") }) await client.disconnect() await server.stop() @@ -921,10 +921,10 @@ struct CompletionTests { argument: CompletionArgument(name: "repo", value: ""), context: CompletionContext(arguments: ["owner": "modelcontextprotocol"]) ) - #expect(result1.values.contains("python-sdk")) - #expect(result1.values.contains("typescript-sdk")) - #expect(result1.values.contains("specification")) - #expect(result1.total == 3) + #expect(result1.completion.values.contains("python-sdk")) + #expect(result1.completion.values.contains("typescript-sdk")) + #expect(result1.completion.values.contains("specification")) + #expect(result1.completion.total == 3) // Test with microsoft owner let result2 = try await client.complete( @@ -932,16 +932,16 @@ struct CompletionTests { argument: CompletionArgument(name: "repo", value: ""), context: CompletionContext(arguments: ["owner": "microsoft"]) ) - #expect(result2.values.contains("vscode")) - #expect(result2.values.contains("typescript")) - #expect(result2.values.contains("playwright")) + #expect(result2.completion.values.contains("vscode")) + #expect(result2.completion.values.contains("typescript")) + #expect(result2.completion.values.contains("playwright")) // Test with no context let result3 = try await client.complete( ref: .resource(ResourceTemplateReference(uri: "github://repos/{owner}/{repo}")), argument: CompletionArgument(name: "repo", value: "") ) - #expect(result3.values == ["repo1", "repo2", "repo3"]) + #expect(result3.completion.values == ["repo1", "repo2", "repo3"]) await client.disconnect() await server.stop() @@ -1003,7 +1003,7 @@ struct CompletionTests { argument: CompletionArgument(name: "name", value: "A"), context: CompletionContext(arguments: ["department": "engineering"]) ) - #expect(result1.values == ["Alice"]) + #expect(result1.completion.values == ["Alice"]) // Test with sales department let result2 = try await client.complete( @@ -1011,7 +1011,7 @@ struct CompletionTests { argument: CompletionArgument(name: "name", value: "D"), context: CompletionContext(arguments: ["department": "sales"]) ) - #expect(result2.values == ["David"]) + #expect(result2.completion.values == ["David"]) // Test with marketing department let result3 = try await client.complete( @@ -1019,15 +1019,15 @@ struct CompletionTests { argument: CompletionArgument(name: "name", value: "G"), context: CompletionContext(arguments: ["department": "marketing"]) ) - #expect(result3.values == ["Grace"]) + #expect(result3.completion.values == ["Grace"]) // Test with no context let result4 = try await client.complete( ref: .prompt(PromptReference(name: "team-greeting")), argument: CompletionArgument(name: "name", value: "U") ) - #expect(result4.values.contains("Unknown1")) - #expect(result4.values.contains("Unknown2")) + #expect(result4.completion.values.contains("Unknown1")) + #expect(result4.completion.values.contains("Unknown2")) await client.disconnect() await server.stop() diff --git a/Tests/MCPTests/ConcurrentExecutionTests.swift b/Tests/MCPTests/ConcurrentExecutionTests.swift new file mode 100644 index 00000000..18088403 --- /dev/null +++ b/Tests/MCPTests/ConcurrentExecutionTests.swift @@ -0,0 +1,296 @@ +import Testing + +@testable import MCP + +/// Tests that verify server handlers execute concurrently. +/// +/// These tests are based on Python SDK's `test_188_concurrency.py`: +/// - `test_messages_are_executed_concurrently_tools` +/// - `test_messages_are_executed_concurrently_tools_and_resources` +/// +/// The pattern uses coordination primitives (events) to prove concurrent execution: +/// 1. First handler starts and waits on an event +/// 2. Second handler starts (only possible if handlers run concurrently) +/// 3. Second handler signals the event +/// 4. First handler completes +/// +/// If handlers ran sequentially, the first handler would block forever +/// waiting for an event that the second handler (which never starts) should signal. +@Suite("Concurrent Execution Tests") +struct ConcurrentExecutionTests { + + // MARK: - Helper Types + + /// An actor that allows async coordination between concurrent tasks. + /// Similar to Python's anyio.Event(). + private actor AsyncEvent { + private var signaled = false + private var waiters: [CheckedContinuation] = [] + + func signal() { + signaled = true + for waiter in waiters { + waiter.resume() + } + waiters.removeAll() + } + + func wait() async { + if signaled { return } + await withCheckedContinuation { continuation in + waiters.append(continuation) + } + } + + var isSignaled: Bool { signaled } + } + + /// An actor that tracks the order of events for verification. + private actor CallOrderTracker { + private var order: [String] = [] + + func append(_ event: String) { + order.append(event) + } + + var events: [String] { order } + } + + // MARK: - Concurrent Tool Execution Tests + + /// Tests that tool calls execute concurrently on the server. + /// + /// Based on Python SDK's `test_messages_are_executed_concurrently_tools`. + /// + /// Pattern: + /// - "sleep" tool starts and waits on an event + /// - "trigger" tool starts (proves concurrency), waits for sleep to start, then signals + /// - Both tools complete + /// + /// If execution were sequential, the sleep tool would block forever. + @Test("Tool calls execute concurrently on server") + func toolCallsExecuteConcurrently() async throws { + let event = AsyncEvent() + let toolStarted = AsyncEvent() + let callOrder = CallOrderTracker() + + let server = Server( + name: "ConcurrentToolServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "sleep", description: "Waits for event", inputSchema: ["type": "object"]), + Tool(name: "trigger", description: "Triggers the event", inputSchema: ["type": "object"]), + ]) + } + + await server.withRequestHandler(CallTool.self) { request, _ in + if request.name == "sleep" { + await callOrder.append("waiting_for_event") + await toolStarted.signal() + await event.wait() + await callOrder.append("tool_end") + return CallTool.Result(content: [.text("done")]) + } else if request.name == "trigger" { + // Wait for sleep tool to start before signaling + await toolStarted.wait() + await callOrder.append("trigger_started") + await event.signal() + await callOrder.append("trigger_end") + return CallTool.Result(content: [.text("triggered")]) + } + return CallTool.Result(content: [.text("unknown")]) + } + + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + let client = Client(name: "ConcurrentTestClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Start the sleep tool (will wait on event) + let sleepTask = Task { + try await client.send(CallTool.request(.init(name: "sleep", arguments: nil))) + } + + // Start the trigger tool (will signal the event) + let triggerTask = Task { + try await client.send(CallTool.request(.init(name: "trigger", arguments: nil))) + } + + // Wait for both to complete + _ = try await sleepTask.value + _ = try await triggerTask.value + + // Verify the order proves concurrent execution + let events = await callOrder.events + #expect( + events == ["waiting_for_event", "trigger_started", "trigger_end", "tool_end"], + "Expected concurrent execution order, but got: \(events)" + ) + } + + /// Tests that tool and resource handlers execute concurrently. + /// + /// Based on Python SDK's `test_messages_are_executed_concurrently_tools_and_resources`. + /// + /// Pattern: + /// - "sleep" tool starts and waits on an event + /// - resource read starts (proves concurrency), signals the event + /// - Both complete + @Test("Tool and resource calls execute concurrently on server") + func toolAndResourceCallsExecuteConcurrently() async throws { + let event = AsyncEvent() + let toolStarted = AsyncEvent() + let callOrder = CallOrderTracker() + + let server = Server( + name: "ConcurrentMixedServer", + version: "1.0.0", + capabilities: .init(resources: .init(), tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "sleep", description: "Waits for event", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, _ in + if request.name == "sleep" { + await callOrder.append("waiting_for_event") + await toolStarted.signal() + await event.wait() + await callOrder.append("tool_end") + return CallTool.Result(content: [.text("done")]) + } + return CallTool.Result(content: [.text("unknown")]) + } + + await server.withRequestHandler(ListResources.self) { _, _ in + ListResources.Result(resources: [ + Resource(name: "Slow Resource", uri: "test://slow_resource") + ]) + } + + await server.withRequestHandler(ReadResource.self) { request, _ in + if request.uri == "test://slow_resource" { + // Wait for tool to start before signaling + await toolStarted.wait() + await event.signal() + await callOrder.append("resource_end") + return ReadResource.Result(contents: [ + .text("slow", uri: "test://slow_resource") + ]) + } + return ReadResource.Result(contents: []) + } + + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + let client = Client(name: "ConcurrentMixedClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Start the sleep tool (will wait on event) + let sleepTask = Task { + try await client.send(CallTool.request(.init(name: "sleep", arguments: nil))) + } + + // Start the resource read (will signal the event) + let resourceTask = Task { + try await client.send(ReadResource.request(.init(uri: "test://slow_resource"))) + } + + // Wait for both to complete + _ = try await sleepTask.value + _ = try await resourceTask.value + + // Verify the order proves concurrent execution + let events = await callOrder.events + #expect( + events == ["waiting_for_event", "resource_end", "tool_end"], + "Expected concurrent execution order, but got: \(events)" + ) + } + + /// Tests that multiple concurrent tool calls all execute in parallel. + /// + /// Pattern: Start N tools that all wait on a shared event, then signal it once. + /// If sequential, only the first would run and block forever. + @Test("Multiple concurrent tool calls all execute in parallel") + func multipleConcurrentToolCallsExecuteInParallel() async throws { + let event = AsyncEvent() + let startedCount = StartedCounter() + let expectedConcurrency = 5 + + let server = Server( + name: "ParallelToolServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "wait_tool", description: "Waits for event", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { _, _ in + // Track that this handler started + await startedCount.increment() + + // Wait for the event + await event.wait() + return CallTool.Result(content: [.text("done")]) + } + + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + let client = Client(name: "ParallelTestClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Start multiple tool calls concurrently + let tasks = (0..= 1) + + // Should be an error response, not a crash + if let response = messages.first { + #expect(response.contains("error"), "Should return an error response") + } + + // Verify server is still alive - send a valid ping request + await transport.clearMessages() + + let validPingRequest = """ + {"jsonrpc":"2.0","id":"ping-1","method":"ping","params":{}} + """ + await transport.queueRaw(validPingRequest) + + // Wait for ping response + let pingReceived = await transport.waitForSentMessage { message in + message.contains("ping-1") + } + #expect(pingReceived, "Server should still be responsive after malformed request") + + let pingMessages = await transport.sentMessages + #expect(pingMessages.count >= 1) + if let pingResponse = pingMessages.first { + #expect(pingResponse.contains("result"), "Ping should succeed after malformed request") + } + + await server.stop() + await transport.disconnect() + } + + /// Test that multiple concurrent malformed requests don't crash the server. + /// + /// Based on Python SDK's test_multiple_concurrent_malformed_requests. + @Test(.timeLimit(.minutes(1))) + func multipleConcurrentMalformedRequestsDontCrashServer() async throws { + let transport = MockTransport() + let server = Server( + name: "ConcurrentMalformedServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + try await server.start(transport: transport) + + // Send multiple malformed requests + for i in 0..<10 { + let malformedRequest = """ + {"jsonrpc":"2.0","id":"malformed-\(i)","method":"initialize"} + """ + await transport.queueRaw(malformedRequest) + } + + // Wait for responses using proper synchronization + let received = await transport.waitForSentMessageCount(10, timeout: .seconds(5)) + #expect(received, "Should receive responses for all malformed requests") + + // Should receive error responses for all requests + let messages = await transport.sentMessages + + // All responses should be errors + for message in messages { + #expect(message.contains("error"), "Each response should be an error") + } + + await server.stop() + await transport.disconnect() + } + + /// Test that server remains functional after malformed input. + /// + /// This tests the "recovery" aspect - after receiving malformed input, + /// the server should still be able to process valid requests. + @Test(.timeLimit(.minutes(1))) + func serverRecoveriesAfterMalformedInput() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.malformed-recovery") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let server = Server( + name: "RecoveryServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "test_tool", inputSchema: ["type": "object"]) + ]) + } + + let client = Client(name: "RecoveryClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // First, verify normal operation works + let tools1 = try await client.send(ListTools.request(.init())) + #expect(tools1.tools.count == 1) + + // Now send malformed data directly to the server transport + // We'll use a separate pipe for this to simulate malformed client + // Skip this part for now as it requires raw transport access + + // Verify server is still functional after init + let tools2 = try await client.send(ListTools.request(.init())) + #expect(tools2.tools.count == 1) + #expect(tools2.tools.first?.name == "test_tool") + } + + /// Test that missing jsonrpc field returns parse error. + @Test(.timeLimit(.minutes(1))) + func missingJsonRpcFieldReturnsError() async throws { + let transport = MockTransport() + let server = Server(name: "ParseErrorServer", version: "1.0.0") + + try await server.start(transport: transport) + + // Send message missing required "jsonrpc" field + let invalidMessage = """ + {"method":"ping","id":"parse-test"} + """ + await transport.queueRaw(invalidMessage) + + // Wait for response + let responseReceived = await transport.waitForSentMessage { message in + message.contains("parse-test") || message.contains("error") + } + #expect(responseReceived, "Should receive an error response") + + let messages = await transport.sentMessages + #expect(messages.count >= 1) + if let response = messages.first { + #expect(response.contains("error"), "Missing jsonrpc should return error") + } + + await server.stop() + await transport.disconnect() + } + + /// Test that completely invalid JSON returns parse error. + @Test(.timeLimit(.minutes(1))) + func invalidJsonReturnsParseError() async throws { + let transport = MockTransport() + let server = Server(name: "InvalidJsonServer", version: "1.0.0") + + try await server.start(transport: transport) + + // Send completely invalid JSON + let invalidJson = "not json at all" + await transport.queueRaw(invalidJson) + + // Wait for error response + try await Task.sleep(for: .milliseconds(100)) + + // Server may not respond to completely invalid messages, but it shouldn't crash + // The key is that subsequent valid messages still work + + // Initialize properly and verify server is still functional + try await transport.queue( + request: Initialize.request( + .init( + protocolVersion: Version.latest, + capabilities: .init(), + clientInfo: .init(name: "TestClient", version: "1.0") + ) + ) + ) + + let initReceived = await transport.waitForSentMessage { message in + message.contains("serverInfo") + } + #expect(initReceived, "Server should still handle valid requests after invalid JSON") + + await server.stop() + await transport.disconnect() + } +} + +// MARK: - Server Resilience Tests +// Based on Python SDK: tests/server/test_lowlevel_exception_handling.py + +/// Tests for server exception handling and resilience. +/// +/// These tests verify that the server properly handles exceptions in request handlers +/// without crashing, and can continue processing subsequent requests. +@Suite("Server Resilience") +struct ServerResilienceTests { + + /// Test that exceptions in request handlers are properly converted to error responses. + /// + /// Based on Python SDK's test_exception_handling_with_raise_exceptions_true. + @Test(.timeLimit(.minutes(1))) + func exceptionInHandlerReturnsErrorResponse() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.exception-handling") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let server = Server( + name: "ExceptionServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "failing_tool", inputSchema: ["type": "object"]) + ]) + } + + // Register a tool handler that throws an exception + await server.withRequestHandler(CallTool.self) { request, _ in + guard request.name == "failing_tool" else { + return CallTool.Result(content: [.text("Unknown tool")], isError: true) + } + // Throw an MCPError to simulate a handler failure + throw MCPError.internalError("Simulated handler failure") + } + + let client = Client(name: "ExceptionClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Call the failing tool - should get an error response, not a crash + do { + _ = try await client.send( + CallTool.request(.init(name: "failing_tool", arguments: [:])) + ) + Issue.record("Expected tool call to throw an error") + } catch let error as MCPError { + // Should receive the internal error + if case .internalError(let message) = error { + #expect(message?.contains("Simulated handler failure") == true) + } else { + // Other error types are also acceptable + } + } + + // Verify server is still functional after exception + let tools = try await client.send(ListTools.request(.init())) + #expect(tools.tools.count == 1) + #expect(tools.tools.first?.name == "failing_tool") + } + + /// Test that normal message handling is not affected by exceptions. + /// + /// Based on Python SDK's test_normal_message_handling_not_affected. + @Test(.timeLimit(.minutes(1))) + func normalMessageHandlingNotAffectedByExceptions() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.normal-handling") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let callCounter = AtomicCounter() + + let server = Server( + name: "NormalHandlingServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "good_tool", inputSchema: ["type": "object"]), + Tool(name: "bad_tool", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, _ in + _ = await callCounter.increment() + + if request.name == "bad_tool" { + throw MCPError.internalError("Bad tool failed") + } + return CallTool.Result(content: [.text("Success!")]) + } + + let client = Client(name: "NormalClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Call good tool - should succeed + let result1 = try await client.send( + CallTool.request(.init(name: "good_tool", arguments: [:])) + ) + if case .text(let text, _, _) = result1.content.first { + #expect(text == "Success!") + } + + // Call bad tool - should fail but not crash server + do { + _ = try await client.send( + CallTool.request(.init(name: "bad_tool", arguments: [:])) + ) + } catch { + // Expected + } + + // Call good tool again - should still work + let result2 = try await client.send( + CallTool.request(.init(name: "good_tool", arguments: [:])) + ) + if case .text(let text, _, _) = result2.content.first { + #expect(text == "Success!") + } + + // Verify all calls were processed + let count = await callCounter.value + #expect(count == 3, "All three tool calls should have been processed") + } + + /// Test multiple different exception types are handled gracefully. + /// + /// Based on Python SDK's test_exception_handling_with_raise_exceptions_false + /// which tests ValueError, RuntimeError, KeyError, and generic Exception. + @Test(.timeLimit(.minutes(1))) + func multipleExceptionTypesHandledGracefully() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.multiple-exceptions") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let server = Server( + name: "MultipleExceptionsServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "invalid_params_tool", inputSchema: ["type": "object"]), + Tool(name: "internal_error_tool", inputSchema: ["type": "object"]), + Tool(name: "resource_not_found_tool", inputSchema: ["type": "object"]), + Tool(name: "method_not_found_tool", inputSchema: ["type": "object"]), + Tool(name: "good_tool", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, _ in + switch request.name { + case "invalid_params_tool": + throw MCPError.invalidParams("Missing required parameter") + case "internal_error_tool": + throw MCPError.internalError("Something went wrong internally") + case "resource_not_found_tool": + throw MCPError.resourceNotFound(uri: "file:///missing.txt") + case "method_not_found_tool": + throw MCPError.methodNotFound("Unknown method") + case "good_tool": + return CallTool.Result(content: [.text("Works!")]) + default: + return CallTool.Result(content: [.text("Unknown")], isError: true) + } + } + + let client = Client(name: "MultipleExceptionsClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Test each error type + let errorTools = [ + ("invalid_params_tool", ErrorCode.invalidParams), + ("internal_error_tool", ErrorCode.internalError), + ("resource_not_found_tool", ErrorCode.resourceNotFound), + ("method_not_found_tool", ErrorCode.methodNotFound) + ] + + for (toolName, expectedCode) in errorTools { + do { + _ = try await client.send( + CallTool.request(.init(name: toolName, arguments: [:])) + ) + Issue.record("Expected \(toolName) to throw an error") + } catch let error as MCPError { + #expect(error.code == expectedCode, "\(toolName) should return error code \(expectedCode)") + } + } + + // Verify server still works after all the exceptions + let result = try await client.send( + CallTool.request(.init(name: "good_tool", arguments: [:])) + ) + if case .text(let text, _, _) = result.content.first { + #expect(text == "Works!") + } + + // Also verify list tools still works + let tools = try await client.send(ListTools.request(.init())) + #expect(tools.tools.count == 5) + } +} + +// MARK: - Timeout and Server Responsiveness Tests +// Based on Python SDK: tests/issues/test_88_random_error.py + +/// Tests for timeout handling and server responsiveness after timeouts. +/// +/// These tests verify that when a client request times out: +/// 1. The server task stays alive +/// 2. The server can still handle new requests +/// 3. The client can make new requests +/// 4. No resources are leaked +@Suite("Timeout and Server Responsiveness") +struct TimeoutServerResponsivenessTests { + + /// Test that server remains responsive after a client request times out. + /// + /// Based on Python SDK's test_notification_validation_error. + /// Uses per-request timeouts to avoid race conditions. + @Test(.timeLimit(.minutes(1))) + func serverRemainsResponsiveAfterTimeout() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.timeout-responsiveness") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let requestCount = AtomicCounter() + let slowRequestLock = AsyncEvent() + + let server = Server( + name: "TimeoutResponsivenessServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "slow", description: "A slow tool", inputSchema: ["type": "object"]), + Tool(name: "fast", description: "A fast tool", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, _ in + let count = await requestCount.increment() + + if request.name == "slow" { + // Wait for the lock - this should timeout + await slowRequestLock.wait() + return CallTool.Result(content: [.text("slow \(count)")]) + } else if request.name == "fast" { + return CallTool.Result(content: [.text("fast \(count)")]) + } + return CallTool.Result(content: [.text("unknown \(count)")]) + } + + let client = Client(name: "TimeoutClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // First call should work (fast operation, no timeout) + let result1 = try await client.send( + CallTool.request(.init(name: "fast", arguments: [:])) + ) + if case .text(let text, _, _) = result1.content.first { + #expect(text == "fast 1") + } + + // Second call should timeout (slow operation with minimal timeout) + do { + _ = try await client.send( + CallTool.request(.init(name: "slow", arguments: [:])), + options: .init(timeout: .milliseconds(10)) + ) + Issue.record("Expected slow tool to timeout") + } catch let error as MCPError { + if case .requestTimeout = error { + // Expected + } else { + // Cancellation error is also acceptable + } + } + + // Release the slow request to avoid hanging processes + await slowRequestLock.signal() + + // Third call should work (fast operation, no timeout) + // This proves the server is still responsive + let result3 = try await client.send( + CallTool.request(.init(name: "fast", arguments: [:])) + ) + if case .text(let text, _, _) = result3.content.first { + #expect(text == "fast 3", "Third call should succeed after timeout") + } + + // Verify all requests were processed by the server + let finalCount = await requestCount.value + #expect(finalCount >= 3, "Server should have processed at least 3 requests") + } + + /// Test multiple sequential requests after a timeout. + /// + /// This verifies that the server doesn't get into a bad state after a timeout. + @Test(.timeLimit(.minutes(1))) + func multipleRequestsAfterTimeout() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.multiple-after-timeout") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let server = Server( + name: "MultipleAfterTimeoutServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "configurable", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, _ in + let delay = request.arguments?["delay"]?.doubleValue ?? 0.0 + if delay > 0 { + try? await Task.sleep(for: .seconds(delay)) + } + return CallTool.Result(content: [.text("Completed after \(delay)s")]) + } + + let client = Client(name: "MultipleClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // First, trigger a timeout + do { + _ = try await client.send( + CallTool.request(.init(name: "configurable", arguments: ["delay": .double(10.0)])), + options: .init(timeout: .milliseconds(10)) + ) + } catch { + // Expected timeout + } + + // Now send multiple sequential requests - all should succeed + for i in 1...5 { + let result = try await client.send( + CallTool.request(.init(name: "configurable", arguments: ["delay": .double(0.0)])) + ) + if case .text(let text, _, _) = result.content.first { + #expect(text == "Completed after 0.0s", "Request \(i) should succeed") + } + } + } + + /// Test that concurrent requests where one times out don't affect other requests. + /// + /// Based on Python SDK's approach of testing timeout isolation. + @Test(.timeLimit(.minutes(1))) + func timeoutDoesntAffectConcurrentRequests() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.concurrent-timeout") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let slowRequestLock = AsyncEvent() + + let server = Server( + name: "ConcurrentTimeoutServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "slow_tool", inputSchema: ["type": "object"]), + Tool(name: "fast_tool", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, _ in + if request.name == "slow_tool" { + // Wait indefinitely until signaled + await slowRequestLock.wait() + return CallTool.Result(content: [.text("slow completed")]) + } else if request.name == "fast_tool" { + return CallTool.Result(content: [.text("fast completed")]) + } + return CallTool.Result(content: [.text("unknown")]) + } + + let client = Client(name: "ConcurrentClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Start a slow request that will timeout (using a very short timeout) + let slowTask = Task { + try await client.send( + CallTool.request(.init(name: "slow_tool", arguments: [:])), + options: .init(timeout: .milliseconds(50)) + ) + } + + // Give the slow request a moment to start + try await Task.sleep(for: .milliseconds(10)) + + // Start a fast request concurrently - this should succeed even though slow is pending + let fastTask = Task { + try await client.send( + CallTool.request(.init(name: "fast_tool", arguments: [:])) + ) + } + + // Fast request should succeed + let fastResult = try await fastTask.value + if case .text(let text, _, _) = fastResult.content.first { + #expect(text == "fast completed") + } + + // Slow request should timeout (or be cancelled) + do { + _ = try await slowTask.value + // If it didn't throw, it might have completed before timeout kicked in + // In concurrent scenarios, this is acceptable + } catch { + // Expected - timeout or cancellation + } + + // Release the slow request lock to clean up server resources + await slowRequestLock.signal() + } +} + +// MARK: - Helper Types + +/// An actor for thread-safe counting. +private actor AtomicCounter { + private var count = 0 + + func increment() -> Int { + count += 1 + return count + } + + var value: Int { count } +} + +/// An actor for async event signaling. +private actor AsyncEvent { + private var signaled = false + + func signal() { + signaled = true + } + + func wait() async { + while !signaled { + try? await Task.sleep(for: .milliseconds(10)) + } + } +} diff --git a/Tests/MCPTests/HTTPClientTransportTests.swift b/Tests/MCPTests/HTTPClientTransportTests.swift index 6b62fe4a..67355eba 100644 --- a/Tests/MCPTests/HTTPClientTransportTests.swift +++ b/Tests/MCPTests/HTTPClientTransportTests.swift @@ -1675,21 +1675,6 @@ import Testing logger: nil ) - // Set up handler for initial POST - await MockURLProtocol.requestHandlerStorage.setHandler { - [testEndpoint] (request: URLRequest) in - let response = HTTPURLResponse( - url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", - headerFields: [ - HTTPHeader.contentType: "text/plain", - HTTPHeader.sessionId: "test-session-notifications", - ])! - return (response, Data()) - } - - try await transport.connect() - try await transport.send(Data()) - // SSE stream with: // 1. A notification (has method, no id) - should NOT stop reconnection // 2. A server request (has method AND id) - should NOT stop reconnection @@ -1698,17 +1683,33 @@ import Testing let sseWithMixedMessages = "id: evt-1\ndata: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/progress\",\"params\":{\"progress\":50}}\n\nid: evt-2\ndata: {\"jsonrpc\":\"2.0\",\"method\":\"sampling/createMessage\",\"id\":\"server-req-1\",\"params\":{}}\n\nid: evt-3\ndata: {\"jsonrpc\":\"2.0\",\"id\":1,\"result\":{\"status\":\"ok\"}}\n\n" let sseData = sseWithMixedMessages.data(using: .utf8)! + // Set up a combined handler for both POST and SSE GET requests + // This avoids the race condition where the SSE GET fires before the handler is set await MockURLProtocol.requestHandlerStorage.setHandler { [testEndpoint, sseData] (request: URLRequest) in - let response = HTTPURLResponse( - url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", - headerFields: [HTTPHeader.contentType: "text/event-stream"])! - return (response, sseData) + if request.httpMethod == "GET" { + // SSE request + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [HTTPHeader.contentType: "text/event-stream"])! + return (response, sseData) + } else { + // Initial POST request + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [ + HTTPHeader.contentType: "text/plain", + HTTPHeader.sessionId: "test-session-notifications", + ])! + return (response, Data()) + } } - try await Task.sleep(for: .milliseconds(200)) + try await transport.connect() + try await transport.send(Data()) // Verify all three messages were received + // The iterator.next() calls will wait for messages to be available let stream = await transport.receive() var iterator = stream.makeAsyncIterator() diff --git a/Tests/MCPTests/HTTPIntegrationTests.swift b/Tests/MCPTests/HTTPIntegrationTests.swift index 786cd099..1d4b04c8 100644 --- a/Tests/MCPTests/HTTPIntegrationTests.swift +++ b/Tests/MCPTests/HTTPIntegrationTests.swift @@ -533,4 +533,128 @@ struct HTTPIntegrationTests { let trackedSessionId = await tracker.get() #expect(trackedSessionId == "callback-test-session") } + + // MARK: - Method Not Allowed Tests + + /// Tests that unsupported HTTP methods (PUT, PATCH) are rejected with 405. + /// + /// Based on Python SDK's `test_method_not_allowed`: + /// - PUT method should be rejected + /// - PATCH method should be rejected + @Test("Reject unsupported HTTP methods with 405") + func rejectUnsupportedHttpMethods() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { "method-test-session" }) + ) + try await transport.connect() + + // Initialize first to get a valid session + let initRequest = TestPayloads.postRequest(body: Self.initializeMessage) + let initResponse = await transport.handleRequest(initRequest) + #expect(initResponse.statusCode == 200) + + // Test PUT method + let putRequest = TestPayloads.customMethodRequest( + method: "PUT", + body: Self.initializeMessage, + sessionId: "method-test-session" + ) + let putResponse = await transport.handleRequest(putRequest) + #expect(putResponse.statusCode == 405, "PUT method should be rejected with 405 Method Not Allowed") + + // Test PATCH method + let patchRequest = TestPayloads.customMethodRequest( + method: "PATCH", + body: Self.initializeMessage, + sessionId: "method-test-session" + ) + let patchResponse = await transport.handleRequest(patchRequest) + #expect(patchResponse.statusCode == 405, "PATCH method should be rejected with 405 Method Not Allowed") + } + + // MARK: - Session Termination Tests + + /// Tests that requests to a terminated session fail with appropriate error. + /// + /// Based on Python SDK's `test_session_termination`: + /// 1. Initialize session + /// 2. Terminate session with DELETE + /// 3. Subsequent requests should fail with 404 "Session terminated" + @Test("Requests to terminated session fail with 404") + func requestsToTerminatedSessionFail() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { "terminated-session" }) + ) + try await transport.connect() + + // Initialize the session + let initRequest = TestPayloads.postRequest(body: Self.initializeMessage) + let initResponse = await transport.handleRequest(initRequest) + #expect(initResponse.statusCode == 200) + + // Make a successful request to confirm session is working + let pingRequest = TestPayloads.postRequest( + body: TestPayloads.pingRequest(), + sessionId: "terminated-session" + ) + let pingResponse = await transport.handleRequest(pingRequest) + #expect(pingResponse.statusCode == 200) + + // Terminate the session with DELETE + let deleteRequest = TestPayloads.deleteRequest(sessionId: "terminated-session") + let deleteResponse = await transport.handleRequest(deleteRequest) + #expect(deleteResponse.statusCode == 200) + + // Attempt to use the terminated session - should fail + let afterDeleteRequest = TestPayloads.postRequest( + body: TestPayloads.pingRequest(), + sessionId: "terminated-session" + ) + let afterDeleteResponse = await transport.handleRequest(afterDeleteRequest) + #expect(afterDeleteResponse.statusCode == 404, "Request to terminated session should return 404") + + // Verify the error message mentions session termination + if let body = afterDeleteResponse.body, let text = String(data: body, encoding: .utf8) { + #expect( + text.lowercased().contains("terminated") || text.lowercased().contains("session"), + "Error message should indicate session termination" + ) + } + } + + // MARK: - Backwards Compatibility Tests + + /// Tests that server accepts requests without protocol version header for backwards compatibility. + /// + /// Based on Python SDK's `test_server_backwards_compatibility_no_protocol_version`: + /// Older clients may not send the mcp-protocol-version header, and the server + /// should still accept their requests. + @Test("Backwards compatibility - accept requests without protocol version header") + func backwardsCompatibilityNoProtocolVersion() async throws { + let transport = HTTPServerTransport( + options: .init(sessionIdGenerator: { "compat-session" }) + ) + try await transport.connect() + + // Initialize the session (with protocol version) + let initRequest = TestPayloads.postRequest(body: Self.initializeMessage) + let initResponse = await transport.handleRequest(initRequest) + #expect(initResponse.statusCode == 200) + + // Make a request WITHOUT the protocol version header + let requestWithoutVersion = HTTPRequest( + method: "POST", + headers: [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + HTTPHeader.sessionId: "compat-session", + // Note: NO protocolVersion header + ], + body: TestPayloads.pingRequest().data(using: .utf8) + ) + let response = await transport.handleRequest(requestWithoutVersion) + + // Should succeed for backwards compatibility + #expect(response.statusCode == 200, "Server should accept requests without protocol version header for backwards compatibility") + } } diff --git a/Tests/MCPTests/Helpers/TestPayloads.swift b/Tests/MCPTests/Helpers/TestPayloads.swift index 29a22259..4639e40d 100644 --- a/Tests/MCPTests/Helpers/TestPayloads.swift +++ b/Tests/MCPTests/Helpers/TestPayloads.swift @@ -196,4 +196,26 @@ extension TestPayloads { body: nil ) } + + /// Creates an HTTP request with a custom method (for testing unsupported methods). + static func customMethodRequest( + method: String, + body: String? = nil, + sessionId: String? = nil, + protocolVersion: String = defaultVersion + ) -> HTTPRequest { + var headers = [ + HTTPHeader.accept: "application/json, text/event-stream", + HTTPHeader.contentType: "application/json", + HTTPHeader.protocolVersion: protocolVersion, + ] + if let sessionId { + headers[HTTPHeader.sessionId] = sessionId + } + return HTTPRequest( + method: method, + headers: headers, + body: body?.data(using: .utf8) + ) + } } diff --git a/Tests/MCPTests/ImplementationMetadataTests.swift b/Tests/MCPTests/ImplementationMetadataTests.swift new file mode 100644 index 00000000..f2de9ef0 --- /dev/null +++ b/Tests/MCPTests/ImplementationMetadataTests.swift @@ -0,0 +1,627 @@ +import Foundation +import Testing + +@testable import MCP + +// MARK: - Server.Info Metadata Tests + +@Suite("Server.Info Metadata Tests") +struct ServerInfoMetadataTests { + + @Test("Server.Info with all fields encodes correctly") + func testServerInfoWithAllFields() throws { + let icons = [ + Icon( + src: "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==", + mimeType: "image/png", + sizes: ["1x1"] + ) + ] + + let info = Server.Info( + name: "test-server", + version: "1.0.0", + title: "Test Server Display Name", + description: "A test server for unit testing", + icons: icons, + websiteUrl: "https://example.com" + ) + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + + let data = try encoder.encode(info) + + // Decode and verify fields directly (avoids JSON escaping issues) + let decoder = JSONDecoder() + let decoded = try decoder.decode(Server.Info.self, from: data) + + #expect(decoded.name == "test-server") + #expect(decoded.version == "1.0.0") + #expect(decoded.title == "Test Server Display Name") + #expect(decoded.description == "A test server for unit testing") + #expect(decoded.websiteUrl == "https://example.com") + #expect(decoded.icons?.count == 1) + #expect(decoded.icons?[0].mimeType == "image/png") + } + + @Test("Server.Info with only required fields encodes correctly") + func testServerInfoWithRequiredFieldsOnly() throws { + let info = Server.Info(name: "basic-server", version: "0.1.0") + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + + let data = try encoder.encode(info) + let json = String(data: data, encoding: .utf8)! + + #expect(json.contains("\"name\":\"basic-server\"")) + #expect(json.contains("\"version\":\"0.1.0\"")) + // Optional fields should not be present when nil + #expect(!json.contains("\"title\"")) + #expect(!json.contains("\"description\"")) + #expect(!json.contains("\"icons\"")) + #expect(!json.contains("\"websiteUrl\"")) + } + + @Test("Server.Info roundtrips correctly with all fields") + func testServerInfoRoundtrip() throws { + let original = Server.Info( + name: "roundtrip-server", + version: "2.0.0", + title: "Roundtrip Test Server", + description: "Testing roundtrip encoding", + icons: [ + Icon(src: "https://example.com/icon.png", mimeType: "image/png", sizes: ["48x48"], theme: .light), + Icon(src: "https://example.com/icon-dark.png", mimeType: "image/png", sizes: ["48x48"], theme: .dark) + ], + websiteUrl: "https://example.com/server" + ) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(original) + let decoded = try decoder.decode(Server.Info.self, from: data) + + #expect(decoded.name == original.name) + #expect(decoded.version == original.version) + #expect(decoded.title == original.title) + #expect(decoded.description == original.description) + #expect(decoded.websiteUrl == original.websiteUrl) + #expect(decoded.icons?.count == 2) + #expect(decoded.icons?[0].theme == .light) + #expect(decoded.icons?[1].theme == .dark) + } + + @Test("Server.Info decodes from TypeScript SDK format") + func testServerInfoDecodesFromTypeScriptFormat() throws { + // Format matching TypeScript SDK title.test.ts + let json = """ + { + "name": "test-server", + "version": "1.0.0", + "title": "Test Server Display Name" + } + """ + + let decoder = JSONDecoder() + let info = try decoder.decode(Server.Info.self, from: json.data(using: .utf8)!) + + #expect(info.name == "test-server") + #expect(info.version == "1.0.0") + #expect(info.title == "Test Server Display Name") + } + + @Test("Server.Info decodes from Python SDK format with icons") + func testServerInfoDecodesFromPythonFormat() throws { + // Format matching Python SDK test_1338_icons_and_metadata.py + let json = """ + { + "name": "TestServer", + "version": "1.0.0", + "websiteUrl": "https://example.com", + "icons": [ + { + "src": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==", + "mimeType": "image/png", + "sizes": ["1x1"] + } + ] + } + """ + + let decoder = JSONDecoder() + let info = try decoder.decode(Server.Info.self, from: json.data(using: .utf8)!) + + #expect(info.name == "TestServer") + #expect(info.version == "1.0.0") + #expect(info.websiteUrl == "https://example.com") + #expect(info.icons?.count == 1) + #expect(info.icons?[0].mimeType == "image/png") + #expect(info.icons?[0].sizes == ["1x1"]) + } +} + +// MARK: - Client.Info Metadata Tests + +@Suite("Client.Info Metadata Tests") +struct ClientInfoMetadataTests { + + @Test("Client.Info with all fields encodes correctly") + func testClientInfoWithAllFields() throws { + let info = Client.Info( + name: "test-client", + version: "1.0.0", + title: "Test Client Display Name", + description: "A test client for unit testing", + icons: [Icon(src: "https://example.com/client-icon.png", mimeType: "image/png")], + websiteUrl: "https://example.com/client" + ) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(info) + + // Decode and verify fields directly (avoids JSON escaping issues) + let decoded = try decoder.decode(Client.Info.self, from: data) + + #expect(decoded.name == "test-client") + #expect(decoded.version == "1.0.0") + #expect(decoded.title == "Test Client Display Name") + #expect(decoded.description == "A test client for unit testing") + #expect(decoded.websiteUrl == "https://example.com/client") + #expect(decoded.icons?.count == 1) + } + + @Test("Client.Info with only required fields encodes correctly") + func testClientInfoWithRequiredFieldsOnly() throws { + let info = Client.Info(name: "basic-client", version: "0.1.0") + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + + let data = try encoder.encode(info) + let json = String(data: data, encoding: .utf8)! + + #expect(json.contains("\"name\":\"basic-client\"")) + #expect(json.contains("\"version\":\"0.1.0\"")) + #expect(!json.contains("\"title\"")) + #expect(!json.contains("\"description\"")) + #expect(!json.contains("\"icons\"")) + #expect(!json.contains("\"websiteUrl\"")) + } + + @Test("Client.Info decodes from Python SDK format") + func testClientInfoDecodesFromPythonFormat() throws { + // Format matching Python SDK test_session.py custom_client_info + let json = """ + { + "name": "test-client", + "version": "1.2.3" + } + """ + + let decoder = JSONDecoder() + let info = try decoder.decode(Client.Info.self, from: json.data(using: .utf8)!) + + #expect(info.name == "test-client") + #expect(info.version == "1.2.3") + } +} + +// MARK: - Icon Type Tests + +@Suite("Icon Type Tests") +struct IconTypeTests { + + @Test("Icon with all fields encodes correctly") + func testIconWithAllFields() throws { + let icon = Icon( + src: "https://example.com/icon.png", + mimeType: "image/png", + sizes: ["48x48", "96x96"], + theme: .light + ) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(icon) + + // Decode and verify fields directly (avoids JSON escaping issues) + let decoded = try decoder.decode(Icon.self, from: data) + + #expect(decoded.src == "https://example.com/icon.png") + #expect(decoded.mimeType == "image/png") + #expect(decoded.sizes == ["48x48", "96x96"]) + #expect(decoded.theme == .light) + } + + @Test("Icon with only src encodes correctly") + func testIconWithOnlySrc() throws { + let icon = Icon(src: "https://example.com/icon.svg") + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(icon) + + // Decode and verify - optional fields should be nil + let decoded = try decoder.decode(Icon.self, from: data) + + #expect(decoded.src == "https://example.com/icon.svg") + #expect(decoded.mimeType == nil) + #expect(decoded.sizes == nil) + #expect(decoded.theme == nil) + } + + @Test("Icon with data URI encodes correctly") + func testIconWithDataUri() throws { + let dataUri = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==" + let icon = Icon(src: dataUri, mimeType: "image/png", sizes: ["1x1"]) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(icon) + let decoded = try decoder.decode(Icon.self, from: data) + + #expect(decoded.src == dataUri) + #expect(decoded.mimeType == "image/png") + #expect(decoded.sizes == ["1x1"]) + } + + @Test("Icon dark theme encodes correctly") + func testIconDarkTheme() throws { + let icon = Icon(src: "https://example.com/dark-icon.png", theme: .dark) + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + + let data = try encoder.encode(icon) + let json = String(data: data, encoding: .utf8)! + + #expect(json.contains("\"theme\":\"dark\"")) + } + + @Test("Multiple icons roundtrip correctly") + func testMultipleIconsRoundtrip() throws { + // Based on Python SDK test_multiple_icons test + let icons = [ + Icon(src: "data:image/png;base64,icon1", mimeType: "image/png", sizes: ["16x16"]), + Icon(src: "data:image/png;base64,icon2", mimeType: "image/png", sizes: ["32x32"]), + Icon(src: "data:image/png;base64,icon3", mimeType: "image/png", sizes: ["64x64"]) + ] + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(icons) + let decoded = try decoder.decode([Icon].self, from: data) + + #expect(decoded.count == 3) + #expect(decoded[0].sizes == ["16x16"]) + #expect(decoded[1].sizes == ["32x32"]) + #expect(decoded[2].sizes == ["64x64"]) + } +} + +// MARK: - Server Initialization with Metadata Tests + +@Suite("Server Initialization with Metadata Tests") +struct ServerInitializationMetadataTests { + + @Test("Server can be initialized with all metadata fields") + func testServerInitWithAllMetadata() async throws { + let icons = [ + Icon( + src: "https://example.com/server-icon.png", + mimeType: "image/png", + sizes: ["48x48"] + ) + ] + + let server = Server( + name: "metadata-server", + version: "1.0.0", + title: "Metadata Test Server", + description: "A server with full metadata", + icons: icons, + websiteUrl: "https://example.com" + ) + + // Verify the server info has all the fields + #expect(server.name == "metadata-server") + #expect(server.version == "1.0.0") + + // Access the serverInfo through the internal property + let serverInfo = await server.serverInfo + #expect(serverInfo.title == "Metadata Test Server") + #expect(serverInfo.description == "A server with full metadata") + #expect(serverInfo.websiteUrl == "https://example.com") + #expect(serverInfo.icons?.count == 1) + #expect(serverInfo.icons?[0].src == "https://example.com/server-icon.png") + } + + @Test("Server without optional metadata fields") + func testServerInitWithoutOptionalMetadata() async throws { + // Based on Python SDK test_no_icons_or_website test + let server = Server( + name: "basic-server", + version: "0.1.0" + ) + + #expect(server.name == "basic-server") + #expect(server.version == "0.1.0") + + let serverInfo = await server.serverInfo + #expect(serverInfo.title == nil) + #expect(serverInfo.description == nil) + #expect(serverInfo.icons == nil) + #expect(serverInfo.websiteUrl == nil) + } +} + +// MARK: - Client Initialization with Metadata Tests + +@Suite("Client Initialization with Metadata Tests") +struct ClientInitializationMetadataTests { + + @Test("Client can be initialized with all metadata fields") + func testClientInitWithAllMetadata() async throws { + let client = Client( + name: "metadata-client", + version: "1.0.0", + title: "Metadata Test Client", + description: "A client with full metadata", + icons: [Icon(src: "https://example.com/client-icon.png")], + websiteUrl: "https://example.com/client" + ) + + #expect(client.name == "metadata-client") + #expect(client.version == "1.0.0") + + // Access the clientInfo + let clientInfo = await client.clientInfo + #expect(clientInfo.title == "Metadata Test Client") + #expect(clientInfo.description == "A client with full metadata") + #expect(clientInfo.websiteUrl == "https://example.com/client") + #expect(clientInfo.icons?.count == 1) + } + + @Test("Client without optional metadata fields") + func testClientInitWithoutOptionalMetadata() async throws { + let client = Client( + name: "basic-client", + version: "0.1.0" + ) + + #expect(client.name == "basic-client") + #expect(client.version == "0.1.0") + + let clientInfo = await client.clientInfo + #expect(clientInfo.title == nil) + #expect(clientInfo.description == nil) + #expect(clientInfo.icons == nil) + #expect(clientInfo.websiteUrl == nil) + } +} + +// MARK: - Initialize Result with Metadata Tests + +@Suite("Initialize Result Metadata Tests") +struct InitializeResultMetadataTests { + + @Test("Initialize result with server title in serverInfo") + func testInitializeResultWithServerTitle() throws { + // Based on TypeScript SDK title.test.ts "should support serverInfo with title" + let result = Initialize.Result( + protocolVersion: Version.latest, + capabilities: Server.Capabilities(), + serverInfo: Server.Info( + name: "test-server", + version: "1.0.0", + title: "Test Server Display Name" + ) + ) + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + + let data = try encoder.encode(result) + let json = String(data: data, encoding: .utf8)! + + #expect(json.contains("\"name\":\"test-server\"")) + #expect(json.contains("\"version\":\"1.0.0\"")) + #expect(json.contains("\"title\":\"Test Server Display Name\"")) + } + + @Test("Initialize result with all metadata fields decodes correctly") + func testInitializeResultWithAllMetadata() throws { + let json = """ + { + "protocolVersion": "2025-11-25", + "capabilities": {}, + "serverInfo": { + "name": "full-metadata-server", + "version": "2.0.0", + "title": "Full Metadata Server", + "description": "Server with all metadata fields", + "websiteUrl": "https://example.com", + "icons": [ + { + "src": "https://example.com/icon.png", + "mimeType": "image/png", + "sizes": ["48x48"], + "theme": "light" + } + ] + }, + "instructions": "Use this server for testing." + } + """ + + let decoder = JSONDecoder() + let result = try decoder.decode(Initialize.Result.self, from: json.data(using: .utf8)!) + + #expect(result.protocolVersion == Version.v2025_11_25) + #expect(result.serverInfo.name == "full-metadata-server") + #expect(result.serverInfo.version == "2.0.0") + #expect(result.serverInfo.title == "Full Metadata Server") + #expect(result.serverInfo.description == "Server with all metadata fields") + #expect(result.serverInfo.websiteUrl == "https://example.com") + #expect(result.serverInfo.icons?.count == 1) + #expect(result.serverInfo.icons?[0].theme == .light) + #expect(result.instructions == "Use this server for testing.") + } + + @Test("Initialize result from Python SDK format") + func testInitializeResultFromPythonFormat() throws { + // Based on Python SDK test_session.py test_client_session_initialize + let json = """ + { + "protocolVersion": "2025-11-25", + "capabilities": { + "logging": null, + "resources": null, + "tools": null, + "experimental": null, + "prompts": null + }, + "serverInfo": { + "name": "mock-server", + "version": "0.1.0" + }, + "instructions": "The server instructions." + } + """ + + let decoder = JSONDecoder() + let result = try decoder.decode(Initialize.Result.self, from: json.data(using: .utf8)!) + + #expect(result.serverInfo.name == "mock-server") + #expect(result.serverInfo.version == "0.1.0") + #expect(result.instructions == "The server instructions.") + } +} + +// MARK: - Integration: Metadata in Initialize Flow + +@Suite("Metadata in Initialize Integration Tests") +struct MetadataInitializeIntegrationTests { + + @Test("Server metadata is returned in initialize result") + func testServerMetadataInInitializeResult() async throws { + // Based on TypeScript SDK title.test.ts "should support serverInfo with title" + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "test-server", + version: "1.0.0", + title: "Test Server Display Name", + description: "A server for integration testing", + websiteUrl: "https://example.com" + ) + + try await server.start(transport: serverTransport) + + let client = Client(name: "test-client", version: "1.0.0") + let initResult = try await client.connect(transport: clientTransport) + + // Verify server metadata is in the initialize result + #expect(initResult.serverInfo.name == "test-server") + #expect(initResult.serverInfo.version == "1.0.0") + #expect(initResult.serverInfo.title == "Test Server Display Name") + #expect(initResult.serverInfo.description == "A server for integration testing") + #expect(initResult.serverInfo.websiteUrl == "https://example.com") + + await client.disconnect() + await server.stop() + } + + @Test("Server with icons metadata is returned correctly") + func testServerWithIconsMetadataInInitializeResult() async throws { + // Based on Python SDK test_1338_icons_and_metadata.py + let testIcon = Icon( + src: "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==", + mimeType: "image/png", + sizes: ["1x1"] + ) + + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "icon-server", + version: "1.0.0", + icons: [testIcon], + websiteUrl: "https://example.com" + ) + + try await server.start(transport: serverTransport) + + let client = Client(name: "test-client", version: "1.0.0") + let initResult = try await client.connect(transport: clientTransport) + + // Verify icons are in the initialize result + #expect(initResult.serverInfo.icons?.count == 1) + #expect(initResult.serverInfo.icons?[0].src == testIcon.src) + #expect(initResult.serverInfo.icons?[0].mimeType == "image/png") + #expect(initResult.serverInfo.icons?[0].sizes == ["1x1"]) + #expect(initResult.serverInfo.websiteUrl == "https://example.com") + + await client.disconnect() + await server.stop() + } + + @Test("Server instructions are included in initialize result") + func testServerInstructionsInInitializeResult() async throws { + // Based on Python SDK test_client_session_initialize + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let instructions = "This server provides tools for testing. Use the 'test' tool for validation." + let server = Server( + name: "instructions-server", + version: "1.0.0", + instructions: instructions + ) + + try await server.start(transport: serverTransport) + + let client = Client(name: "test-client", version: "1.0.0") + let initResult = try await client.connect(transport: clientTransport) + + #expect(initResult.instructions == instructions) + + await client.disconnect() + await server.stop() + } + + @Test("Server without optional metadata still initializes correctly") + func testServerWithoutOptionalMetadata() async throws { + // Based on Python SDK test_no_icons_or_website + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "basic-server", + version: "1.0.0" + ) + + try await server.start(transport: serverTransport) + + let client = Client(name: "test-client", version: "1.0.0") + let initResult = try await client.connect(transport: clientTransport) + + #expect(initResult.serverInfo.name == "basic-server") + #expect(initResult.serverInfo.version == "1.0.0") + #expect(initResult.serverInfo.title == nil) + #expect(initResult.serverInfo.description == nil) + #expect(initResult.serverInfo.icons == nil) + #expect(initResult.serverInfo.websiteUrl == nil) + #expect(initResult.instructions == nil) + + await client.disconnect() + await server.stop() + } +} diff --git a/Tests/MCPTests/InMemoryEventStoreTests.swift b/Tests/MCPTests/InMemoryEventStoreTests.swift index 1bd79f5f..eda8b65f 100644 --- a/Tests/MCPTests/InMemoryEventStoreTests.swift +++ b/Tests/MCPTests/InMemoryEventStoreTests.swift @@ -461,4 +461,213 @@ struct InMemoryEventStoreTests { streamCount = await store.streamCount #expect(streamCount == 1) } + + // MARK: - Priming Events + + @Test("Priming events (empty data) are stored but skipped during replay") + func primingEventsSkippedDuringReplay() async throws { + let store = InMemoryEventStore() + + // Store a regular event + let msg1 = #"{"jsonrpc":"2.0","result":"first","id":"1"}"#.data(using: .utf8)! + let eventId1 = try await store.storeEvent(streamId: "stream", message: msg1) + + // Store a priming event (empty data) + let primingEventId = try await store.storeEvent(streamId: "stream", message: Data()) + + // Store another regular event + let msg2 = #"{"jsonrpc":"2.0","result":"second","id":"2"}"#.data(using: .utf8)! + _ = try await store.storeEvent(streamId: "stream", message: msg2) + + // All three events should be stored + let count = await store.eventCount + #expect(count == 3) + + // Priming event should have a valid stream ID + let primingStreamId = await store.streamIdForEventId(primingEventId) + #expect(primingStreamId == "stream") + + // Replay from first event - should only get the second regular event + // (priming event should be skipped) + actor MessageCollector { + var messages: [Data] = [] + func add(_ msg: Data) { messages.append(msg) } + func get() -> [Data] { messages } + } + let collector = MessageCollector() + + _ = try await store.replayEventsAfter(eventId1) { _, message in + await collector.add(message) + } + + let replayedMessages = await collector.get() + #expect(replayedMessages.count == 1) // Only the non-priming event + #expect(replayedMessages[0] == msg2) // The second regular message + } + + @Test("Replay events in strict chronological order") + func replayEventsInChronologicalOrder() async throws { + let store = InMemoryEventStore() + + // Store events with explicit ordering in their content + var eventIds: [String] = [] + for i in 0..<10 { + let message = #"{"order":\#(i)}"#.data(using: .utf8)! + let eventId = try await store.storeEvent(streamId: "stream", message: message) + eventIds.append(eventId) + } + + // Replay from the first event + actor OrderCollector { + var orders: [Int] = [] + func add(_ order: Int) { orders.append(order) } + func get() -> [Int] { orders } + } + let collector = OrderCollector() + + _ = try await store.replayEventsAfter(eventIds[0]) { _, message in + if let json = try? JSONSerialization.jsonObject(with: message) as? [String: Any], + let order = json["order"] as? Int + { + await collector.add(order) + } + } + + let replayedOrders = await collector.get() + + // Should have events 1-9 (after event 0) + #expect(replayedOrders.count == 9) + + // Verify strict chronological ordering + for i in 0.. Int { count } + } + let counter = Counter() + + let streamId = try await store.replayEventsAfter(lastEventId) { _, _ in + await counter.increment() + } + + #expect(streamId == "stream") + let replayedCount = await counter.value() + #expect(replayedCount == 0) // Nothing to replay after the most recent event + } + + @Test("Replay returns correct stream ID") + func replayReturnsCorrectStreamId() async throws { + let store = InMemoryEventStore() + + // Store events on different streams + let msg1 = Data("test".utf8) + let eventIdStream1 = try await store.storeEvent(streamId: "stream-alpha", message: msg1) + let eventIdStream2 = try await store.storeEvent(streamId: "stream-beta", message: msg1) + + // Add more events to both streams + _ = try await store.storeEvent(streamId: "stream-alpha", message: msg1) + _ = try await store.storeEvent(streamId: "stream-beta", message: msg1) + + // Replay from stream-alpha should return "stream-alpha" + let streamId1 = try await store.replayEventsAfter(eventIdStream1) { _, _ in } + #expect(streamId1 == "stream-alpha") + + // Replay from stream-beta should return "stream-beta" + let streamId2 = try await store.replayEventsAfter(eventIdStream2) { _, _ in } + #expect(streamId2 == "stream-beta") + } + + // MARK: - Edge Cases + + @Test("Store and replay with special characters in stream ID") + func storeAndReplayWithSpecialStreamId() async throws { + let store = InMemoryEventStore() + + // Test with various special characters that might be in a stream ID + let specialStreamId = "stream-with-dashes_and_underscores.and.dots" + let message = #"{"test":"value"}"#.data(using: .utf8)! + + let eventId1 = try await store.storeEvent(streamId: specialStreamId, message: message) + _ = try await store.storeEvent(streamId: specialStreamId, message: message) + + // Verify stream ID can be retrieved + let retrievedStreamId = await store.streamIdForEventId(eventId1) + #expect(retrievedStreamId == specialStreamId) + + // Verify replay works + actor Counter { + var count = 0 + func increment() { count += 1 } + func value() -> Int { count } + } + let counter = Counter() + + let replayedStreamId = try await store.replayEventsAfter(eventId1) { _, _ in + await counter.increment() + } + + #expect(replayedStreamId == specialStreamId) + #expect(await counter.value() == 1) + } + + @Test("Multiple streams interleaved storage and replay") + func multipleStreamsInterleavedStorageAndReplay() async throws { + let store = InMemoryEventStore() + + // Interleave storage across multiple streams + let msg1 = #"{"stream":"1","seq":1}"#.data(using: .utf8)! + let eventIdS1_1 = try await store.storeEvent(streamId: "stream-1", message: msg1) + + let msg2 = #"{"stream":"2","seq":1}"#.data(using: .utf8)! + _ = try await store.storeEvent(streamId: "stream-2", message: msg2) + + let msg3 = #"{"stream":"1","seq":2}"#.data(using: .utf8)! + _ = try await store.storeEvent(streamId: "stream-1", message: msg3) + + let msg4 = #"{"stream":"2","seq":2}"#.data(using: .utf8)! + _ = try await store.storeEvent(streamId: "stream-2", message: msg4) + + let msg5 = #"{"stream":"1","seq":3}"#.data(using: .utf8)! + _ = try await store.storeEvent(streamId: "stream-1", message: msg5) + + // Replay from stream-1's first event + actor MessageCollector { + var sequences: [Int] = [] + func add(_ seq: Int) { sequences.append(seq) } + func get() -> [Int] { sequences } + } + let collector = MessageCollector() + + let streamId = try await store.replayEventsAfter(eventIdS1_1) { _, message in + if let json = try? JSONSerialization.jsonObject(with: message) as? [String: Any], + let stream = json["stream"] as? String, + let seq = json["seq"] as? Int + { + // Verify only stream-1 messages are replayed + #expect(stream == "1") + await collector.add(seq) + } + } + + #expect(streamId == "stream-1") + let sequences = await collector.get() + #expect(sequences == [2, 3]) // Only stream-1 events after the first one + } } diff --git a/Tests/MCPTests/IntegrationRoundtripTests.swift b/Tests/MCPTests/IntegrationRoundtripTests.swift new file mode 100644 index 00000000..e4a9f29b --- /dev/null +++ b/Tests/MCPTests/IntegrationRoundtripTests.swift @@ -0,0 +1,1249 @@ +import Foundation +import Testing + +@testable import MCP + +/// Integration roundtrip tests that verify full client-server communication flows. +/// +/// These tests are based on the Python SDK's `tests/server/fastmcp/test_integration.py` +/// and TypeScript SDK's integration tests. They test complete roundtrip scenarios +/// including callbacks, notifications, and bi-directional communication. +/// +/// Key test patterns covered: +/// - `test_basic_prompts` - GetPrompt with argument substitution +/// - `test_tool_progress` - Progress notifications during tool execution +/// - `test_sampling` - Server requesting LLM sampling from client +/// - `test_elicitation` - Server requesting user input from client +/// - `test_notifications` - Logging and list change notifications + +@Suite("Integration Roundtrip Tests") +struct IntegrationRoundtripTests { + + // MARK: - Basic Tools Roundtrip Tests + + /// Tests basic tool functionality with list and call operations. + /// + /// Based on Python SDK's `test_basic_tools`: + /// 1. Client lists tools + /// 2. Client calls tools with arguments + /// 3. Server executes and returns results + @Test("Basic tools roundtrip - list and call") + func testBasicToolsRoundtrip() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ToolServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + // Register tools + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool( + name: "sum", + description: "Adds two numbers", + inputSchema: [ + "type": "object", + "properties": [ + "a": ["type": "integer", "description": "First number"], + "b": ["type": "integer", "description": "Second number"], + ], + "required": ["a", "b"], + ] + ), + Tool( + name: "get_weather", + description: "Gets weather for a city", + inputSchema: [ + "type": "object", + "properties": [ + "city": ["type": "string", "description": "City name"] + ], + "required": ["city"], + ] + ), + ]) + } + + await server.withRequestHandler(CallTool.self) { request, _ in + switch request.name { + case "sum": + let a = request.arguments?["a"]?.intValue ?? 0 + let b = request.arguments?["b"]?.intValue ?? 0 + return CallTool.Result(content: [.text("\(a + b)")]) + + case "get_weather": + let city = request.arguments?["city"]?.stringValue ?? "Unknown" + return CallTool.Result(content: [ + .text("Weather in \(city): 22°C, Sunny") + ]) + + default: + return CallTool.Result( + content: [.text("Unknown tool: \(request.name)")], + isError: true + ) + } + } + + try await server.start(transport: serverTransport) + + let client = Client(name: "ToolTestClient", version: "1.0.0") + let initResult = try await client.connect(transport: clientTransport) + + // Verify tools capability + #expect(initResult.capabilities.tools != nil) + + // Test listing tools + let toolsResult = try await client.listTools() + #expect(toolsResult.tools.count == 2) + + let sumTool = toolsResult.tools.first { $0.name == "sum" } + #expect(sumTool != nil) + #expect(sumTool?.description == "Adds two numbers") + + let weatherTool = toolsResult.tools.first { $0.name == "get_weather" } + #expect(weatherTool != nil) + + // Test sum tool + let sumResult = try await client.callTool( + name: "sum", + arguments: ["a": 5, "b": 3] + ) + #expect(sumResult.content.count == 1) + if case .text(let text, _, _) = sumResult.content[0] { + #expect(text == "8") + } else { + Issue.record("Expected text content") + } + + // Test weather tool + let weatherResult = try await client.callTool( + name: "get_weather", + arguments: ["city": "London"] + ) + #expect(weatherResult.content.count == 1) + if case .text(let text, _, _) = weatherResult.content[0] { + #expect(text.contains("Weather in London")) + #expect(text.contains("22°C")) + } else { + Issue.record("Expected text content") + } + + await client.disconnect() + await server.stop() + } + + /// Tests that calling an unknown tool returns an error. + /// + /// Based on TypeScript SDK's integration tests for error handling. + @Test("Unknown tool returns error") + func testUnknownToolReturnsError() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ToolServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "known_tool", inputSchema: ["type": "object"]) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, _ in + if request.name == "known_tool" { + return CallTool.Result(content: [.text("OK")]) + } + return CallTool.Result( + content: [.text("Unknown tool: \(request.name)")], + isError: true + ) + } + + try await server.start(transport: serverTransport) + + let client = Client(name: "ToolTestClient", version: "1.0.0") + _ = try await client.connect(transport: clientTransport) + + // Call unknown tool + let result = try await client.callTool( + name: "nonexistent_tool", + arguments: [:] + ) + + #expect(result.isError == true) + if case .text(let text, _, _) = result.content[0] { + #expect(text.contains("Unknown tool")) + } + + await client.disconnect() + await server.stop() + } + + // MARK: - Basic Resources Roundtrip Tests + + /// Tests basic resource functionality with list and read operations. + /// + /// Based on Python SDK's `test_basic_resources`: + /// 1. Client lists resources + /// 2. Client reads resources + /// 3. Server returns resource contents + @Test("Basic resources roundtrip - list and read") + func testBasicResourcesRoundtrip() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ResourceServer", + version: "1.0.0", + capabilities: .init(resources: .init()) + ) + + // Register resources + await server.withRequestHandler(ListResources.self) { _, _ in + ListResources.Result(resources: [ + Resource( + name: "readme", + uri: "file://documents/readme", + description: "Project readme file", + mimeType: "text/plain" + ), + Resource( + name: "settings", + uri: "config://settings", + description: "Application settings", + mimeType: "application/json" + ), + ]) + } + + await server.withRequestHandler(ReadResource.self) { request, _ in + switch request.uri { + case "file://documents/readme": + return ReadResource.Result(contents: [ + .text("# Project Readme\n\nContent of readme file.", uri: request.uri) + ]) + + case "config://settings": + let settingsJSON = """ + {"theme": "dark", "language": "en", "notifications": true} + """ + return ReadResource.Result(contents: [ + .text(settingsJSON, uri: request.uri, mimeType: "application/json") + ]) + + default: + throw MCPError.resourceNotFound(uri: request.uri) + } + } + + try await server.start(transport: serverTransport) + + let client = Client(name: "ResourceTestClient", version: "1.0.0") + let initResult = try await client.connect(transport: clientTransport) + + // Verify resources capability + #expect(initResult.capabilities.resources != nil) + + // Test listing resources + let resourcesResult = try await client.listResources() + #expect(resourcesResult.resources.count == 2) + + let readme = resourcesResult.resources.first { $0.name == "readme" } + #expect(readme != nil) + #expect(readme?.uri == "file://documents/readme") + #expect(readme?.mimeType == "text/plain") + + let settings = resourcesResult.resources.first { $0.name == "settings" } + #expect(settings != nil) + #expect(settings?.uri == "config://settings") + + // Test reading readme resource + let readmeResult = try await client.readResource(uri: "file://documents/readme") + #expect(readmeResult.contents.count == 1) + let readmeContent = readmeResult.contents[0] + #expect(readmeContent.text?.contains("Project Readme") == true) + #expect(readmeContent.text?.contains("Content of readme") == true) + #expect(readmeContent.uri == "file://documents/readme") + + // Test reading settings resource + let settingsResult = try await client.readResource(uri: "config://settings") + #expect(settingsResult.contents.count == 1) + let settingsContent = settingsResult.contents[0] + #expect(settingsContent.text?.contains("\"theme\": \"dark\"") == true) + #expect(settingsContent.text?.contains("\"language\": \"en\"") == true) + #expect(settingsContent.mimeType == "application/json") + + await client.disconnect() + await server.stop() + } + + // MARK: - Prompts Roundtrip Tests + + /// Tests full getPrompt roundtrip with argument substitution. + /// + /// Based on Python SDK's `test_basic_prompts`: + /// 1. Client lists prompts + /// 2. Client gets a prompt with arguments + /// 3. Server substitutes arguments into prompt messages + /// 4. Client receives the formatted prompt + @Test("GetPrompt roundtrip with argument substitution") + func testGetPromptRoundtrip() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "PromptServer", + version: "1.0.0", + capabilities: .init(prompts: .init()) + ) + + // Register prompts + await server.withRequestHandler(ListPrompts.self) { _, _ in + ListPrompts.Result(prompts: [ + Prompt( + name: "review_code", + description: "Reviews code and provides feedback", + arguments: [ + Prompt.Argument(name: "code", description: "The code to review", required: true) + ] + ), + Prompt( + name: "debug_error", + description: "Helps debug an error", + arguments: [ + Prompt.Argument(name: "error", description: "The error message", required: true) + ] + ), + ]) + } + + // Handle getPrompt with argument substitution + await server.withRequestHandler(GetPrompt.self) { request, _ in + switch request.name { + case "review_code": + let code = request.arguments?["code"] ?? "" + return GetPrompt.Result( + description: "Code review prompt", + messages: [ + .user("Please review this code:\n\n\(code)") + ] + ) + + case "debug_error": + let error = request.arguments?["error"] ?? "" + return GetPrompt.Result( + description: "Debug error prompt", + messages: [ + .user("I'm seeing this error:"), + .user(.text(error)), + .assistant("I'll help debug that error. Let me analyze it."), + ] + ) + + default: + throw MCPError.invalidParams("Unknown prompt: \(request.name)") + } + } + + try await server.start(transport: serverTransport) + + let client = Client(name: "PromptTestClient", version: "1.0.0") + let initResult = try await client.connect(transport: clientTransport) + + // Verify prompts capability + #expect(initResult.capabilities.prompts != nil) + + // Test listing prompts + let promptsList = try await client.listPrompts() + #expect(promptsList.prompts.count == 2) + let reviewPrompt = promptsList.prompts.first { $0.name == "review_code" } + #expect(reviewPrompt != nil) + #expect(reviewPrompt?.arguments?.first?.name == "code") + + // Test review_code prompt with argument substitution + let codeToReview = "def hello():\n print('Hello')" + let reviewResult = try await client.getPrompt( + name: "review_code", + arguments: ["code": codeToReview] + ) + #expect(reviewResult.messages.count == 1) + if case .text(let text, _, _) = reviewResult.messages[0].content { + #expect(text.contains("Please review this code:")) + #expect(text.contains("def hello():")) + } else { + Issue.record("Expected text content") + } + + // Test debug_error prompt with multi-message response + let errorMessage = "TypeError: 'NoneType' object is not subscriptable" + let debugResult = try await client.getPrompt( + name: "debug_error", + arguments: ["error": errorMessage] + ) + #expect(debugResult.messages.count == 3) + #expect(debugResult.messages[0].role == .user) + #expect(debugResult.messages[1].role == .user) + #expect(debugResult.messages[2].role == .assistant) + + if case .text(let text, _, _) = debugResult.messages[0].content { + #expect(text.contains("I'm seeing this error:")) + } else { + Issue.record("Expected text content for first message") + } + + if case .text(let text, _, _) = debugResult.messages[1].content { + #expect(text.contains("TypeError")) + } else { + Issue.record("Expected text content for second message") + } + + if case .text(let text, _, _) = debugResult.messages[2].content { + #expect(text.contains("I'll help debug")) + } else { + Issue.record("Expected text content for third message") + } + + await client.disconnect() + await server.stop() + } + + // MARK: - Progress Notifications Roundtrip Tests + + /// Tests progress notifications during tool execution. + /// + /// Based on Python SDK's `test_tool_progress`: + /// 1. Client calls a long-running tool with a progress token + /// 2. Server sends progress notifications during execution + /// 3. Client receives and tracks progress updates + @Test("Progress notifications during tool execution") + func testProgressNotificationsDuringToolExecution() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ProgressServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool( + name: "long_running_task", + description: "A task that reports progress", + inputSchema: [ + "type": "object", + "properties": [ + "task_name": ["type": "string"], + "steps": ["type": "integer"], + ], + ] + ) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, context in + let taskName = request.arguments?["task_name"]?.stringValue ?? "Task" + let steps = request.arguments?["steps"]?.intValue ?? 3 + + // Send progress notifications for each step + for step in 1...steps { + let progress = Double(step) / Double(steps) + let message = "Step \(step)/\(steps): Processing..." + + // Send progress notification using sendMessage + try await context.sendMessage(ProgressNotification.message(.init( + progressToken: .string("progress-token"), + progress: progress, + total: 1.0, + message: message + ))) + + // Simulate work + try? await Task.sleep(for: .milliseconds(10)) + } + + return CallTool.Result(content: [ + .text("Task '\(taskName)' completed after \(steps) steps") + ]) + } + + try await server.start(transport: serverTransport) + + let client = Client(name: "ProgressTestClient", version: "1.0.0") + + // Track progress updates received by the client + let clientProgressUpdates = ClientProgressUpdates() + + await client.onNotification(ProgressNotification.self) { [clientProgressUpdates] message in + await clientProgressUpdates.append( + progress: message.params.progress, + total: message.params.total, + message: message.params.message + ) + } + + _ = try await client.connect(transport: clientTransport) + + // Call tool + let result = try await client.callTool( + name: "long_running_task", + arguments: ["task_name": "Test Task", "steps": 3] + ) + + // Give notifications time to be processed + try await Task.sleep(for: .milliseconds(100)) + + // Verify tool completed successfully + #expect(result.content.count == 1) + if case .text(let text, _, _) = result.content[0] { + #expect(text.contains("Test Task")) + #expect(text.contains("completed")) + } else { + Issue.record("Expected text content") + } + + // Verify progress updates were received + let updates = await clientProgressUpdates.updates + #expect(updates.count == 3) + + // Verify progress values + for (index, update) in updates.enumerated() { + let expectedProgress = Double(index + 1) / 3.0 + #expect(abs(update.progress - expectedProgress) < 0.01) + #expect(update.total == 1.0) + #expect(update.message?.contains("Step \(index + 1)/3") == true) + } + + await client.disconnect() + await server.stop() + } + + // MARK: - Sampling Roundtrip Tests + + /// Tests server requesting LLM sampling from client during tool execution. + /// + /// Based on Python SDK's `test_sampling`: + /// 1. Client calls a tool that needs LLM assistance + /// 2. Server sends CreateSamplingMessage request to client + /// 3. Client's sampling callback processes the request + /// 4. Server receives the LLM response and completes the tool + @Test("Sampling roundtrip - server requests LLM from client") + func testSamplingRoundtrip() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "SamplingServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool( + name: "generate_poem", + description: "Generates a poem using LLM", + inputSchema: [ + "type": "object", + "properties": [ + "topic": ["type": "string"] + ], + ] + ) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] request, _ in + let topic = request.arguments?["topic"]?.stringValue ?? "nature" + + // Server requests LLM sampling from client + let samplingResult = try await server.createMessage( + CreateSamplingMessage.Parameters( + messages: [.user("Write a short poem about \(topic)")], + systemPrompt: "You are a creative poet.", + maxTokens: 100 + ) + ) + + // Return the LLM response + if case .text(let text, _, _) = samplingResult.content { + return CallTool.Result(content: [.text("Generated poem:\n\(text)")]) + } else { + return CallTool.Result(content: [.text("Failed to generate poem")]) + } + } + + try await server.start(transport: serverTransport) + + // Create client with sampling capability + let samplingCallbackInvoked = SamplingCallbackTracker() + + let client = Client(name: "SamplingTestClient", version: "1.0.0") + + // Set client capabilities + await client.setCapabilities(Client.Capabilities( + sampling: Client.Capabilities.Sampling() + )) + + // Set up sampling callback that simulates LLM response + await client.withRequestHandler(CreateSamplingMessage.self) { [samplingCallbackInvoked] params, _ in + await samplingCallbackInvoked.record(params: params) + + // Return simulated LLM response + return CreateSamplingMessage.Result( + model: "test-model", + stopReason: .endTurn, + role: .assistant, + content: .text("This is a simulated LLM response for testing") + ) + } + + _ = try await client.connect(transport: clientTransport) + + // Call the tool that triggers sampling + let result = try await client.callTool( + name: "generate_poem", + arguments: ["topic": "nature"] + ) + + // Verify sampling callback was invoked + let invocations = await samplingCallbackInvoked.invocations + #expect(invocations.count == 1) + #expect(invocations[0].messages.count == 1) + if case .text(let text, _, _) = invocations[0].messages[0].content.first { + #expect(text.contains("poem")) + #expect(text.contains("nature")) + } + #expect(invocations[0].systemPrompt == "You are a creative poet.") + + // Verify tool returned the LLM response + #expect(result.content.count == 1) + if case .text(let text, _, _) = result.content[0] { + #expect(text.contains("simulated LLM response")) + } else { + Issue.record("Expected text content") + } + + await client.disconnect() + await server.stop() + } + + // MARK: - Elicitation Roundtrip Tests + + /// Tests server requesting user input from client during tool execution. + /// + /// Based on Python SDK's `test_elicitation`: + /// 1. Client calls a tool that needs user input + /// 2. Server sends elicitation request to client + /// 3. Client's elicitation callback processes the request + /// 4. Server receives the user response and completes the tool + @Test("Elicitation roundtrip - server requests user input from client") + func testElicitationRoundtrip() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ElicitationServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool( + name: "book_table", + description: "Books a restaurant table", + inputSchema: [ + "type": "object", + "properties": [ + "date": ["type": "string"], + "time": ["type": "string"], + "party_size": ["type": "integer"], + ], + ] + ) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] request, _ in + let date = request.arguments?["date"]?.stringValue ?? "" + let time = request.arguments?["time"]?.stringValue ?? "" + let partySize = request.arguments?["party_size"]?.intValue ?? 2 + + // Simulate date unavailable - request alternative from user + if date == "2024-12-25" { + let elicitResult = try await server.elicit(ElicitRequestParams.form(ElicitRequestFormParams( + message: "No tables available for \(date). Would you like to try an alternative date?", + requestedSchema: ElicitationSchema( + properties: [ + "checkAlternative": .boolean(BooleanSchema( + title: "Try Alternative", + description: "Would you like to try an alternative date?" + )), + "alternativeDate": .string(StringSchema( + title: "Alternative Date", + description: "Enter an alternative date" + )), + ], + required: ["checkAlternative"] + ) + ))) + + if elicitResult.action == .accept, + let checkAlt = elicitResult.content?["checkAlternative"]?.boolValue, + checkAlt, + let altDate = elicitResult.content?["alternativeDate"]?.stringValue + { + return CallTool.Result(content: [ + .text("[SUCCESS] Booked for \(altDate) at \(time) for \(partySize) guests") + ]) + } else { + return CallTool.Result(content: [ + .text("[CANCELLED] Booking cancelled by user") + ]) + } + } + + // Date is available + return CallTool.Result(content: [ + .text("[SUCCESS] Booked for \(date) at \(time) for \(partySize) guests") + ]) + } + + try await server.start(transport: serverTransport) + + // Create client with elicitation capability + let elicitationCallbackInvoked = ElicitationCallbackTracker() + + let client = Client(name: "ElicitationTestClient", version: "1.0.0") + + // Set client capabilities + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(form: Client.Capabilities.Elicitation.Form()) + )) + + // Set up elicitation callback + await client.withElicitationHandler { [elicitationCallbackInvoked] params, _ in + await elicitationCallbackInvoked.record(params: params) + + // Simulate user accepting and providing alternative date + if case .form(let formParams) = params { + if formParams.message.contains("No tables available") { + return ElicitResult( + action: .accept, + content: [ + "checkAlternative": .bool(true), + "alternativeDate": .string("2024-12-26"), + ] + ) + } + } + + return ElicitResult(action: .decline) + } + + _ = try await client.connect(transport: clientTransport) + + // Test booking with unavailable date (triggers elicitation) + let result1 = try await client.callTool( + name: "book_table", + arguments: [ + "date": "2024-12-25", + "time": "19:00", + "party_size": 4, + ] + ) + + // Verify elicitation was invoked + let invocations = await elicitationCallbackInvoked.invocations + #expect(invocations.count == 1) + + // Verify booking succeeded with alternative date + #expect(result1.content.count == 1) + if case .text(let text, _, _) = result1.content[0] { + #expect(text.contains("[SUCCESS]")) + #expect(text.contains("2024-12-26")) + } else { + Issue.record("Expected text content") + } + + // Test booking with available date (no elicitation) + let result2 = try await client.callTool( + name: "book_table", + arguments: [ + "date": "2024-12-20", + "time": "20:00", + "party_size": 2, + ] + ) + + // Verify no additional elicitation was triggered + let invocationsAfter = await elicitationCallbackInvoked.invocations + #expect(invocationsAfter.count == 1) // Still just 1 + + // Verify booking succeeded directly + if case .text(let text, _, _) = result2.content[0] { + #expect(text.contains("[SUCCESS]")) + #expect(text.contains("2024-12-20")) + } else { + Issue.record("Expected text content") + } + + await client.disconnect() + await server.stop() + } + + // MARK: - Logging Notifications Roundtrip Tests + + /// Tests logging notifications at different levels during tool execution. + /// + /// Based on Python SDK's `test_notifications`: + /// 1. Client calls a tool that generates log messages + /// 2. Server sends log notifications at various levels + /// 3. Client receives and collects the log messages + @Test("Logging notifications during tool execution") + func testLoggingNotificationsDuringToolExecution() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "LoggingServer", + version: "1.0.0", + capabilities: .init(logging: .init(), tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool( + name: "process_data", + description: "Processes data and logs progress", + inputSchema: [ + "type": "object", + "properties": [ + "data": ["type": "string"] + ], + ] + ) + ]) + } + + await server.withRequestHandler(CallTool.self) { request, context in + let data = request.arguments?["data"]?.stringValue ?? "" + + // Send log messages at different levels using sendMessage + try await context.sendMessage(LogMessageNotification.message(.init( + level: .debug, + logger: "process", + data: .string("Starting to process data") + ))) + try await context.sendMessage(LogMessageNotification.message(.init( + level: .info, + logger: "process", + data: .string("Processing: \(data)") + ))) + try await context.sendMessage(LogMessageNotification.message(.init( + level: .warning, + logger: "process", + data: .string("Data contains special characters") + ))) + try await context.sendMessage(LogMessageNotification.message(.init( + level: .error, + logger: "process", + data: .string("Simulated error for testing") + ))) + + return CallTool.Result(content: [.text("Processed: \(data)")]) + } + + try await server.start(transport: serverTransport) + + // Create client and track notifications + let logCollector = LogCollector() + + let client = Client(name: "LoggingTestClient", version: "1.0.0") + + // Set up notification handler + await client.onNotification(LogMessageNotification.self) { [logCollector] message in + await logCollector.append(message.params) + } + + _ = try await client.connect(transport: clientTransport) + + // Call tool that generates log messages + let result = try await client.callTool( + name: "process_data", + arguments: ["data": "test_data"] + ) + + // Verify tool completed + #expect(result.content.count == 1) + if case .text(let text, _, _) = result.content[0] { + #expect(text.contains("Processed: test_data")) + } + + // Give notifications time to arrive + try await Task.sleep(for: .milliseconds(100)) + + // Verify log messages at different levels + let logs = await logCollector.logs + #expect(logs.count >= 4) + + let levels = Set(logs.map { $0.level }) + #expect(levels.contains(.debug)) + #expect(levels.contains(.info)) + #expect(levels.contains(.warning)) + #expect(levels.contains(.error)) + + await client.disconnect() + await server.stop() + } + + // MARK: - Resource List Changed Notification Test + + /// Tests resource list changed notifications. + /// + /// Based on Python SDK's `test_notifications` (resource notifications part). + @Test("Resource list changed notification during tool execution") + func testResourceListChangedNotification() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ResourceNotificationServer", + version: "1.0.0", + capabilities: .init(resources: .init(listChanged: true), tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool( + name: "create_resource", + description: "Creates a new resource", + inputSchema: ["type": "object"] + ) + ]) + } + + await server.withRequestHandler(ListResources.self) { _, _ in + ListResources.Result(resources: [ + Resource(name: "Initial Resource", uri: "test://initial") + ]) + } + + await server.withRequestHandler(CallTool.self) { _, context in + // Notify that resource list has changed + try await context.sendNotification(ResourceListChangedNotification()) + + return CallTool.Result(content: [.text("Resource created")]) + } + + try await server.start(transport: serverTransport) + + // Track resource list changed notifications + let notificationReceived = NotificationTracker() + + let client = Client(name: "ResourceNotificationClient", version: "1.0.0") + + await client.onNotification(ResourceListChangedNotification.self) { [notificationReceived] _ in + await notificationReceived.recordNotification() + } + + _ = try await client.connect(transport: clientTransport) + + // Call tool that triggers notification + _ = try await client.callTool(name: "create_resource", arguments: nil) + + // Give notification time to arrive + try await Task.sleep(for: .milliseconds(100)) + + // Verify notification was received + let received = await notificationReceived.wasNotified + #expect(received) + + await client.disconnect() + await server.stop() + } + + // MARK: - Tool List Changed With Refresh Tests + + /// Tests that client receives tool list changed notification and can refresh the list. + /// + /// Based on TypeScript SDK's `should handle tool list changed notification with auto refresh`: + /// 1. Client connects and lists tools + /// 2. Server sends tool list changed notification (triggered via a tool call) + /// 3. Client receives notification and re-fetches tools + /// 4. Client sees updated tool list + @Test("Tool list changed notification with refresh") + func testToolListChangedNotificationWithRefresh() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + // Track available tools (mutable) + let toolRegistry = ToolRegistry() + + let server = Server( + name: "DynamicToolServer", + version: "1.0.0", + capabilities: .init(tools: .init(listChanged: true)) + ) + + // Dynamic tool list handler + await server.withRequestHandler(ListTools.self) { [toolRegistry] _, _ in + let tools = await toolRegistry.getTools() + return ListTools.Result(tools: tools) + } + + await server.withRequestHandler(CallTool.self) { [toolRegistry] request, context in + switch request.name { + case "add_tool": + // Add the new tool and send notification + let toolName = request.arguments?["name"]?.stringValue ?? "unnamed" + await toolRegistry.addTool(Tool( + name: toolName, + description: "Dynamically added", + inputSchema: ["type": "object"] + )) + try await context.sendToolListChanged() + return CallTool.Result(content: [.text("Added tool: \(toolName)")]) + default: + return CallTool.Result(content: [.text("Called: \(request.name)")]) + } + } + + try await server.start(transport: serverTransport) + + // Add initial tools + await toolRegistry.addTool(Tool( + name: "initial_tool", + description: "Initial tool", + inputSchema: ["type": "object"] + )) + await toolRegistry.addTool(Tool( + name: "add_tool", + description: "Tool that adds a new tool", + inputSchema: ["type": "object", "properties": ["name": ["type": "string"]]] + )) + + let client = Client(name: "DynamicToolClient", version: "1.0.0") + + // Track notifications + let notificationReceived = ToolListChangedTracker() + + await client.onNotification(ToolListChangedNotification.self) { [notificationReceived] _ in + await notificationReceived.recordNotification() + } + + _ = try await client.connect(transport: clientTransport) + + // Initial list should have 2 tools + let initialToolsResult = try await client.listTools() + #expect(initialToolsResult.tools.count == 2) + #expect(initialToolsResult.tools.contains { $0.name == "initial_tool" }) + #expect(initialToolsResult.tools.contains { $0.name == "add_tool" }) + + // Call the add_tool which adds a new tool and sends notification + _ = try await client.callTool(name: "add_tool", arguments: ["name": "new_tool"]) + + // Wait for notification to be processed + try await Task.sleep(for: .milliseconds(100)) + + // Verify notification was received + let received = await notificationReceived.wasNotified + #expect(received) + + // Client refreshes tool list after notification + let refreshedToolsResult = try await client.listTools() + #expect(refreshedToolsResult.tools.count == 3) + #expect(refreshedToolsResult.tools.contains { $0.name == "initial_tool" }) + #expect(refreshedToolsResult.tools.contains { $0.name == "add_tool" }) + #expect(refreshedToolsResult.tools.contains { $0.name == "new_tool" }) + + await client.disconnect() + await server.stop() + } + + /// Tests that client receives prompt list changed notification and can refresh the list. + /// + /// Based on TypeScript SDK's `should handle prompt list changed notification with auto refresh`. + @Test("Prompt list changed notification with refresh") + func testPromptListChangedNotificationWithRefresh() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + // Track available prompts (mutable) + let promptRegistry = PromptRegistry() + + let server = Server( + name: "DynamicPromptServer", + version: "1.0.0", + capabilities: .init(prompts: .init(listChanged: true), tools: .init()) + ) + + // Dynamic prompt list handler + await server.withRequestHandler(ListPrompts.self) { [promptRegistry] _, _ in + let prompts = await promptRegistry.getPrompts() + return ListPrompts.Result(prompts: prompts) + } + + await server.withRequestHandler(GetPrompt.self) { request, _ in + GetPrompt.Result(description: nil, messages: [.user("Prompt: \(request.name)")]) + } + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "add_prompt", inputSchema: ["type": "object", "properties": ["name": ["type": "string"]]]) + ]) + } + + await server.withRequestHandler(CallTool.self) { [promptRegistry] request, context in + guard request.name == "add_prompt" else { + return CallTool.Result(content: [.text("Unknown tool")], isError: true) + } + // Add the new prompt and send notification + let promptName = request.arguments?["name"]?.stringValue ?? "unnamed" + await promptRegistry.addPrompt(Prompt( + name: promptName, + description: "Dynamically added" + )) + try await context.sendPromptListChanged() + return CallTool.Result(content: [.text("Added prompt: \(promptName)")]) + } + + try await server.start(transport: serverTransport) + + // Add initial prompt + await promptRegistry.addPrompt(Prompt( + name: "initial_prompt", + description: "Initial prompt" + )) + + let client = Client(name: "DynamicPromptClient", version: "1.0.0") + + // Track notifications + let notificationReceived = PromptListChangedTracker() + + await client.onNotification(PromptListChangedNotification.self) { [notificationReceived] _ in + await notificationReceived.recordNotification() + } + + _ = try await client.connect(transport: clientTransport) + + // Initial list should have 1 prompt + let initialPrompts = try await client.listPrompts() + #expect(initialPrompts.prompts.count == 1) + #expect(initialPrompts.prompts[0].name == "initial_prompt") + + // Call tool that adds a new prompt and sends notification + _ = try await client.callTool(name: "add_prompt", arguments: ["name": "new_prompt"]) + + // Wait for notification to be processed + try await Task.sleep(for: .milliseconds(100)) + + // Verify notification was received + let received = await notificationReceived.wasNotified + #expect(received) + + // Client refreshes prompt list after notification + let refreshedPrompts = try await client.listPrompts() + #expect(refreshedPrompts.prompts.count == 2) + #expect(refreshedPrompts.prompts.contains { $0.name == "initial_prompt" }) + #expect(refreshedPrompts.prompts.contains { $0.name == "new_prompt" }) + + await client.disconnect() + await server.stop() + } +} + +// MARK: - Test Helpers + +/// Tracks progress updates received by the client. +private actor ClientProgressUpdates { + struct Update { + let progress: Double + let total: Double? + let message: String? + } + + var updates: [Update] = [] + + func append(progress: Double, total: Double?, message: String?) { + updates.append(Update(progress: progress, total: total, message: message)) + } +} + +/// Tracks sampling callback invocations. +private actor SamplingCallbackTracker { + var invocations: [CreateSamplingMessage.Parameters] = [] + + func record(params: CreateSamplingMessage.Parameters) { + invocations.append(params) + } +} + +/// Tracks elicitation callback invocations. +private actor ElicitationCallbackTracker { + var invocations: [ElicitRequestParams] = [] + + func record(params: ElicitRequestParams) { + invocations.append(params) + } +} + +/// Collects log notifications. +private actor LogCollector { + var logs: [LogMessageNotification.Parameters] = [] + + func append(_ notification: LogMessageNotification.Parameters) { + logs.append(notification) + } +} + +/// Tracks whether a notification was received. +private actor NotificationTracker { + var wasNotified = false + + func recordNotification() { + wasNotified = true + } +} + +/// Registry for dynamically adding tools in tests. +private actor ToolRegistry { + var tools: [Tool] = [] + + func addTool(_ tool: Tool) { + tools.append(tool) + } + + func getTools() -> [Tool] { + tools + } +} + +/// Registry for dynamically adding prompts in tests. +private actor PromptRegistry { + var prompts: [Prompt] = [] + + func addPrompt(_ prompt: Prompt) { + prompts.append(prompt) + } + + func getPrompts() -> [Prompt] { + prompts + } +} + +/// Tracks whether a tool list changed notification was received. +private actor ToolListChangedTracker { + var wasNotified = false + + func recordNotification() { + wasNotified = true + } +} + +/// Tracks whether a prompt list changed notification was received. +private actor PromptListChangedTracker { + var wasNotified = false + + func recordNotification() { + wasNotified = true + } +} diff --git a/Tests/MCPTests/PrimingEventsTests.swift b/Tests/MCPTests/PrimingEventsTests.swift index aeb6da25..e6b5977f 100644 --- a/Tests/MCPTests/PrimingEventsTests.swift +++ b/Tests/MCPTests/PrimingEventsTests.swift @@ -75,6 +75,41 @@ struct PrimingEventsTests { return receivedData } + /// Initialize the server via HTTP and wait for the initialize response. + /// + /// Per MCP spec lifecycle, clients must wait for the initialize response before + /// sending other requests. This helper sends the initialize request and reads + /// from the SSE stream until the response arrives, ensuring the Server has + /// fully processed the initialization. + /// + /// Note: When an event store is configured with protocol version >= 2025-11-25, + /// a priming event (with empty data) is sent first. This helper reads enough + /// chunks to receive the actual initialize response. + func initializeAndWaitForResponse( + transport: HTTPServerTransport, + protocolVersion: String = Version.latest + ) async throws { + let initRequest = TestPayloads.initializeRequest(protocolVersion: protocolVersion) + let initResponse = await transport.handleRequest( + TestPayloads.postRequest(body: initRequest, protocolVersion: protocolVersion) + ) + + guard initResponse.statusCode == 200 else { + throw MCPError.internalError("Initialize failed with status \(initResponse.statusCode)") + } + + // Wait for the actual initialize response on the SSE stream + // This ensures the Server has processed the initialize request + // Read up to 2 chunks to handle priming events (which come first with empty data) + if let stream = initResponse.stream { + let data = try await readFromStream(stream, maxChunks: 2, timeout: .seconds(2)) + let text = String(data: data, encoding: .utf8) ?? "" + guard text.contains("serverInfo") || text.contains("protocolVersion") else { + throw MCPError.internalError("Did not receive initialize response: \(text)") + } + } + } + /// Creates a configured MCP Server with tools for testing func createTestServer() -> Server { let server = Server( @@ -144,11 +179,9 @@ struct PrimingEventsTests { ) try await server.start(transport: transport) - // Initialize with latest supported version (2025-03-26) + // Initialize and wait for the response (per MCP spec lifecycle) // Note: Priming events require >= 2025-11-25 which is not yet supported - let initRequest = TestPayloads.initializeRequest(protocolVersion: Version.v2025_03_26) - let initResponse = await transport.handleRequest(TestPayloads.postRequest(body: initRequest, protocolVersion: Version.v2025_03_26)) - #expect(initResponse.statusCode == 200) + try await initializeAndWaitForResponse(transport: transport, protocolVersion: Version.v2025_03_26) // Send a tool call request let toolCallRequest = """ @@ -156,7 +189,7 @@ struct PrimingEventsTests { """ let response = await transport.handleRequest(TestPayloads.postRequest(body: toolCallRequest, sessionId: sessionId, protocolVersion: Version.v2025_03_26)) - #expect(response.statusCode == 200) + #expect(response.statusCode == 200, "Expected 200 but got \(response.statusCode)") if let stream = response.stream { let data = try await readFromStream(stream, maxChunks: 2) @@ -164,10 +197,10 @@ struct PrimingEventsTests { // Priming events have empty data - current versions won't have them #expect(!text.contains("data: \n\n"), "Should NOT have empty priming event for current protocol versions") - #expect(text.contains("Hello, Test!") || text.contains("result"), "Should contain tool result") + #expect(text.contains("Hello, Test!") || text.contains("result"), "Should contain tool result. Actual: \(text)") } else if let body = response.body { let text = String(data: body, encoding: .utf8) ?? "" - #expect(text.contains("Hello, Test!") || text.contains("result"), "Should contain tool result") + #expect(text.contains("Hello, Test!") || text.contains("result"), "Should contain tool result. Actual: \(text)") } } @@ -189,9 +222,8 @@ struct PrimingEventsTests { ) try await server.start(transport: transport) - // Initialize - let initRequest = TestPayloads.initializeRequest() - _ = await transport.handleRequest(TestPayloads.postRequest(body: initRequest)) + // Initialize and wait for the response (per MCP spec lifecycle) + try await initializeAndWaitForResponse(transport: transport) // Send a tool call request let toolCallRequest = """ @@ -234,9 +266,8 @@ struct PrimingEventsTests { ) try await server.start(transport: transport) - // Initialize with OLD protocol version (< 2025-11-25) - let initRequest = TestPayloads.initializeRequest() - _ = await transport.handleRequest(TestPayloads.postRequest(body: initRequest, protocolVersion: Version.v2024_11_05)) + // Initialize with OLD protocol version (< 2025-11-25) and wait for response + try await initializeAndWaitForResponse(transport: transport, protocolVersion: Version.v2024_11_05) // Send a tool call request let toolCallRequest = """ @@ -276,9 +307,8 @@ struct PrimingEventsTests { ) try await server.start(transport: transport) - // Initialize - let initRequest = TestPayloads.initializeRequest() - _ = await transport.handleRequest(TestPayloads.postRequest(body: initRequest)) + // Initialize and wait for the response (per MCP spec lifecycle) + try await initializeAndWaitForResponse(transport: transport) // Send a tool call request let toolCallRequest = """ @@ -330,7 +360,7 @@ struct PrimingEventsTests { ]) } - await server.withRequestHandler(CallTool.self) { request, context in + await server.withRequestHandler(CallTool.self) { request, _ in if request.name == "slow-tool" { // Simulate slow operation try? await Task.sleep(for: .milliseconds(500)) @@ -341,9 +371,8 @@ struct PrimingEventsTests { try await server.start(transport: transport) - // Initialize - let initRequest = TestPayloads.initializeRequest() - _ = await transport.handleRequest(TestPayloads.postRequest(body: initRequest)) + // Initialize and wait for the response (per MCP spec lifecycle) + try await initializeAndWaitForResponse(transport: transport) // The closeSSEStream method exists and is callable // We can't fully test the stream closure without complex async coordination diff --git a/Tests/MCPTests/RequestHandlerContextTests.swift b/Tests/MCPTests/RequestHandlerContextTests.swift index 494815f2..07ce1577 100644 --- a/Tests/MCPTests/RequestHandlerContextTests.swift +++ b/Tests/MCPTests/RequestHandlerContextTests.swift @@ -3,6 +3,124 @@ import Testing @testable import MCP +// MARK: - RequestInfo Tests + +@Suite("RequestInfo Tests") +struct RequestInfoTests { + + @Test("RequestInfo stores headers") + func testRequestInfoStoresHeaders() { + let headers = ["Content-Type": "application/json", "X-Custom": "test"] + let requestInfo = RequestInfo(headers: headers) + + #expect(requestInfo.headers == headers) + } + + @Test("RequestInfo.header() performs case-insensitive lookup") + func testRequestInfoHeaderCaseInsensitive() { + let headers = ["Content-Type": "application/json", "X-Custom-Header": "value"] + let requestInfo = RequestInfo(headers: headers) + + // Case-insensitive lookup + #expect(requestInfo.header("content-type") == "application/json") + #expect(requestInfo.header("CONTENT-TYPE") == "application/json") + #expect(requestInfo.header("Content-Type") == "application/json") + #expect(requestInfo.header("x-custom-header") == "value") + #expect(requestInfo.header("X-CUSTOM-HEADER") == "value") + } + + @Test("RequestInfo.header() returns nil for missing headers") + func testRequestInfoHeaderMissing() { + let requestInfo = RequestInfo(headers: ["Content-Type": "application/json"]) + + #expect(requestInfo.header("X-Missing") == nil) + #expect(requestInfo.header("Authorization") == nil) + } + + @Test("RequestInfo is Hashable and Sendable") + func testRequestInfoIsHashableAndSendable() { + let requestInfo1 = RequestInfo(headers: ["X-Test": "value"]) + let requestInfo2 = RequestInfo(headers: ["X-Test": "value"]) + let requestInfo3 = RequestInfo(headers: ["X-Test": "other"]) + + // Hashable + #expect(requestInfo1 == requestInfo2) + #expect(requestInfo1 != requestInfo3) + + // Sendable (compilation test - this compiles if Sendable) + let _: @Sendable () -> RequestInfo = { requestInfo1 } + } +} + +// MARK: - RequestMeta.relatedTaskId Tests + +@Suite("RequestMeta.relatedTaskId Tests") +struct RequestMetaRelatedTaskIdTests { + + @Test("relatedTaskId extracts task ID from additionalFields") + func testRelatedTaskIdExtractsTaskId() { + let meta = RequestMeta(additionalFields: [ + "io.modelcontextprotocol/related-task": .object(["taskId": .string("task-abc123")]) + ]) + + #expect(meta.relatedTaskId == "task-abc123") + } + + @Test("relatedTaskId returns nil when no related task metadata") + func testRelatedTaskIdNilWhenNoMetadata() { + let meta = RequestMeta() + + #expect(meta.relatedTaskId == nil) + } + + @Test("relatedTaskId returns nil when additionalFields is nil") + func testRelatedTaskIdNilWhenAdditionalFieldsNil() { + let meta = RequestMeta(progressToken: .string("token")) + + #expect(meta.relatedTaskId == nil) + } + + @Test("relatedTaskId returns nil when key is missing") + func testRelatedTaskIdNilWhenKeyMissing() { + let meta = RequestMeta(additionalFields: [ + "other-key": .object(["taskId": .string("task-123")]) + ]) + + #expect(meta.relatedTaskId == nil) + } + + @Test("relatedTaskId returns nil when taskId is not a string") + func testRelatedTaskIdNilWhenTaskIdNotString() { + let meta = RequestMeta(additionalFields: [ + "io.modelcontextprotocol/related-task": .object(["taskId": .double(123)]) + ]) + + #expect(meta.relatedTaskId == nil) + } + + @Test("relatedTaskId returns nil when value is not an object") + func testRelatedTaskIdNilWhenNotObject() { + let meta = RequestMeta(additionalFields: [ + "io.modelcontextprotocol/related-task": .string("not-an-object") + ]) + + #expect(meta.relatedTaskId == nil) + } + + @Test("relatedTaskId works with progressToken also set") + func testRelatedTaskIdWithProgressToken() { + let meta = RequestMeta( + progressToken: .string("progress-123"), + additionalFields: [ + "io.modelcontextprotocol/related-task": .object(["taskId": .string("task-xyz")]) + ] + ) + + #expect(meta.progressToken == .string("progress-123")) + #expect(meta.relatedTaskId == "task-xyz") + } +} + /// Tests for RequestHandlerContext functionality. /// /// These tests verify that handlers have access to request context information @@ -466,6 +584,165 @@ struct ServerRequestHandlerContextTests { await client.disconnect() } + // MARK: - requestInfo Tests + + /// Test that context.requestInfo is nil for non-HTTP transports. + /// Based on TypeScript SDK's `extra.requestInfo` which is only populated for HTTP connections. + @Test("context.requestInfo is nil for InMemoryTransport") + func testRequestInfoNilForInMemoryTransport() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + actor RequestInfoTracker { + var receivedRequestInfo: RequestInfo? + var wasChecked = false + func set(_ requestInfo: RequestInfo?) { + receivedRequestInfo = requestInfo + wasChecked = true + } + } + let tracker = RequestInfoTracker() + + let server = Server( + name: "TestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "test_tool", description: "Test", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { _, context in + // Handler accesses context.requestInfo - should be nil for InMemoryTransport + await tracker.set(context.requestInfo) + return CallTool.Result(content: [.text("OK")]) + } + + let client = Client(name: "TestClient", version: "1.0.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + _ = try await client.callTool(name: "test_tool", arguments: [:]) + + let wasChecked = await tracker.wasChecked + let receivedRequestInfo = await tracker.receivedRequestInfo + #expect(wasChecked, "Handler should have been called") + #expect(receivedRequestInfo == nil, "requestInfo should be nil for InMemoryTransport") + + await client.disconnect() + } + + // MARK: - taskId Tests + + /// Test that context.taskId extracts task ID from _meta when present. + /// Based on TypeScript SDK's `extra.taskId` which is extracted from `_meta[RELATED_TASK_META_KEY]`. + @Test("context.taskId extracts task ID from _meta") + func testContextTaskIdFromMeta() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + actor TaskIdTracker { + var receivedTaskId: String? + var wasChecked = false + func set(_ taskId: String?) { + receivedTaskId = taskId + wasChecked = true + } + } + let tracker = TaskIdTracker() + + let server = Server( + name: "TestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "task_tool", description: "Test", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { _, context in + // Handler accesses context.taskId - convenience property + await tracker.set(context.taskId) + return CallTool.Result(content: [.text("OK")]) + } + + let client = Client(name: "TestClient", version: "1.0.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Call tool with related task metadata in _meta + _ = try await client.send( + CallTool.request(.init( + name: "task_tool", + arguments: [:], + _meta: RequestMeta(additionalFields: [ + "io.modelcontextprotocol/related-task": .object(["taskId": .string("test-task-123")]) + ]) + )) + ) + + let wasChecked = await tracker.wasChecked + let receivedTaskId = await tracker.receivedTaskId + #expect(wasChecked, "Handler should have been called") + #expect(receivedTaskId == "test-task-123", "context.taskId should extract task ID from _meta") + + await client.disconnect() + } + + /// Test that context.taskId is nil when no related task metadata. + @Test("context.taskId is nil when no related task metadata") + func testContextTaskIdNilWhenNoTaskMeta() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + actor TaskIdTracker { + var receivedTaskId: String? = "initial" // Initialize to non-nil + var wasChecked = false + func set(_ taskId: String?) { + receivedTaskId = taskId + wasChecked = true + } + } + let tracker = TaskIdTracker() + + let server = Server( + name: "TestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "test_tool", description: "Test", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { _, context in + await tracker.set(context.taskId) + return CallTool.Result(content: [.text("OK")]) + } + + let client = Client(name: "TestClient", version: "1.0.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Call tool WITHOUT related task metadata + _ = try await client.callTool(name: "test_tool", arguments: [:]) + + let wasChecked = await tracker.wasChecked + let receivedTaskId = await tracker.receivedTaskId + #expect(wasChecked, "Handler should have been called") + #expect(receivedTaskId == nil, "context.taskId should be nil when no related task metadata") + + await client.disconnect() + } + // MARK: - closeSSEStream Tests /// Test that context.closeSSEStream is nil for non-HTTP transports. @@ -579,6 +856,129 @@ struct ClientRequestHandlerContextTests { await client.disconnect() } + /// Test that client handlers can access context.taskId when present. + /// This matches the server's context.taskId convenience property. + @Test("Client handler can access context.taskId") + func testClientHandlerCanAccessTaskId() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + actor TaskIdTracker { + var receivedTaskId: String? + var wasChecked = false + func set(_ id: String?) { + receivedTaskId = id + wasChecked = true + } + } + let tracker = TaskIdTracker() + + let server = Server( + name: "TestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "elicitTool", description: "Elicit", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] _, _ in + // Send elicitation with task metadata + let result = try await server.elicit(.form(ElicitRequestFormParams( + message: "Test", + requestedSchema: ElicitationSchema(properties: ["x": .string(StringSchema())]), + _meta: RequestMeta(additionalFields: [ + "io.modelcontextprotocol/related-task": .object(["taskId": .string("client-task-456")]) + ]) + ))) + return CallTool.Result(content: [.text("Action: \(result.action)")]) + } + + let client = Client(name: "TestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(form: Client.Capabilities.Elicitation.Form()) + )) + + await client.withElicitationHandler { _, context in + // Client handler accesses context.taskId - convenience property matching server context + await tracker.set(context.taskId) + return ElicitResult(action: .accept, content: ["x": .string("test")]) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + _ = try await client.callTool(name: "elicitTool", arguments: [:]) + + let wasChecked = await tracker.wasChecked + let receivedTaskId = await tracker.receivedTaskId + #expect(wasChecked, "Client handler should have been called") + #expect(receivedTaskId == "client-task-456", "context.taskId should extract task ID from _meta") + + await client.disconnect() + } + + /// Test that client context.taskId is nil when no related task metadata. + @Test("Client context.taskId is nil when no related task metadata") + func testClientContextTaskIdNilWhenNoTaskMeta() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + actor TaskIdTracker { + var receivedTaskId: String? = "initial" // Initialize to non-nil + var wasChecked = false + func set(_ id: String?) { + receivedTaskId = id + wasChecked = true + } + } + let tracker = TaskIdTracker() + + let server = Server( + name: "TestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "elicitTool", description: "Elicit", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] _, _ in + // Send elicitation WITHOUT task metadata + let result = try await server.elicit(.form(ElicitRequestFormParams( + message: "Test", + requestedSchema: ElicitationSchema(properties: ["x": .string(StringSchema())]) + ))) + return CallTool.Result(content: [.text("Action: \(result.action)")]) + } + + let client = Client(name: "TestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(form: Client.Capabilities.Elicitation.Form()) + )) + + await client.withElicitationHandler { _, context in + await tracker.set(context.taskId) + return ElicitResult(action: .accept, content: ["x": .string("test")]) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + _ = try await client.callTool(name: "elicitTool", arguments: [:]) + + let wasChecked = await tracker.wasChecked + let receivedTaskId = await tracker.receivedTaskId + #expect(wasChecked, "Client handler should have been called") + #expect(receivedTaskId == nil, "context.taskId should be nil when no related task metadata") + + await client.disconnect() + } + /// Test that client handlers can access context._meta when present. @Test("Client handler can access context._meta") func testClientHandlerCanAccessMeta() async throws { @@ -634,3 +1034,474 @@ struct ClientRequestHandlerContextTests { await client.disconnect() } } + +// MARK: - Additional Context Tests from Python/TypeScript SDKs + +/// Additional tests based on patterns from Python and TypeScript SDKs. +/// These tests ensure feature parity across SDK implementations. +@Suite("Additional RequestHandlerContext Tests") +struct AdditionalRequestHandlerContextTests { + + // MARK: - context.elicitUrl() Tests + + /// Test that handlers can use context.elicitUrl() for URL elicitation. + /// Based on Python SDK's ctx.session.elicit_url() and ctx.elicit_url() tests. + @Test("Handler can use context.elicitUrl() for URL elicitation") + func testContextElicitUrl() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "TestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "authorize", description: "Authorize access", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { _, context in + // Use context.elicitUrl() instead of server.elicit() + // This is the convenience method pattern from Python's ctx.elicit_url() + let result = try await context.elicitUrl( + message: "Please authorize access to files", + url: "https://example.com/oauth/authorize", + elicitationId: "file-auth-123" + ) + + return switch result.action { + case .accept: CallTool.Result(content: [.text("Authorized")]) + case .decline: CallTool.Result(content: [.text("Declined")]) + case .cancel: CallTool.Result(content: [.text("Cancelled")]) + } + } + + let client = Client(name: "TestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(url: Client.Capabilities.Elicitation.URL()) + )) + + await client.withElicitationHandler { params, _ in + guard case .url(let urlParams) = params else { + return ElicitResult(action: .decline) + } + #expect(urlParams.message == "Please authorize access to files") + #expect(urlParams.elicitationId == "file-auth-123") + #expect(urlParams.url == "https://example.com/oauth/authorize") + return ElicitResult(action: .accept) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "authorize", arguments: [:]) + + #expect(result.isError == nil) + if case .text(let text, _, _) = result.content.first { + #expect(text == "Authorized") + } else { + Issue.record("Expected text content") + } + + await client.disconnect() + } + + /// Test that context.elicitUrl() handles user decline. + @Test("context.elicitUrl() handles user decline") + func testContextElicitUrlDecline() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "TestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "authorize", description: "Authorize", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { _, context in + let result = try await context.elicitUrl( + message: "Authorize?", + url: "https://example.com/oauth", + elicitationId: "auth-decline-test" + ) + + return switch result.action { + case .accept: CallTool.Result(content: [.text("Authorized")]) + case .decline: CallTool.Result(content: [.text("Declined")]) + case .cancel: CallTool.Result(content: [.text("Cancelled")]) + } + } + + let client = Client(name: "TestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(url: Client.Capabilities.Elicitation.URL()) + )) + + await client.withElicitationHandler { _, _ in + return ElicitResult(action: .decline) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "authorize", arguments: [:]) + + if case .text(let text, _, _) = result.content.first { + #expect(text == "Declined") + } else { + Issue.record("Expected text content") + } + + await client.disconnect() + } + + // MARK: - context.elicit() Cancel Action Test + + /// Test that context.elicit() handles cancel action. + /// Based on TypeScript SDK's cancel action tests. + @Test("context.elicit() handles cancel action") + func testContextElicitCancel() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "TestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "confirm", description: "Confirm action", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { _, context in + let result = try await context.elicit( + message: "Confirm this action?", + requestedSchema: ElicitationSchema( + properties: ["confirm": .boolean(BooleanSchema(title: "Confirm"))] + ) + ) + + return switch result.action { + case .accept: CallTool.Result(content: [.text("Accepted")]) + case .decline: CallTool.Result(content: [.text("Declined")]) + case .cancel: CallTool.Result(content: [.text("Cancelled")]) + } + } + + let client = Client(name: "TestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(form: Client.Capabilities.Elicitation.Form()) + )) + + await client.withElicitationHandler { _, _ in + return ElicitResult(action: .cancel) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "confirm", arguments: [:]) + + if case .text(let text, _, _) = result.content.first { + #expect(text == "Cancelled") + } else { + Issue.record("Expected text content") + } + + await client.disconnect() + } + + // MARK: - Multiple Sequential Elicitation Requests + + /// Test multiple sequential elicitation requests within a single handler. + /// Based on TypeScript SDK's test for handling multiple sequential elicitation requests. + @Test("Handler can make multiple sequential elicitation requests") + func testMultipleSequentialElicitationRequests() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "TestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "wizard", description: "Multi-step wizard", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { _, context in + // First elicitation - get name + let nameResult = try await context.elicit( + message: "What is your name?", + requestedSchema: ElicitationSchema( + properties: ["name": .string(StringSchema(title: "Name"))], + required: ["name"] + ) + ) + + guard nameResult.action == .accept, + let name = nameResult.content?["name"]?.stringValue else { + return CallTool.Result(content: [.text("Name step failed")], isError: true) + } + + // Second elicitation - get age + let ageResult = try await context.elicit( + message: "What is your age?", + requestedSchema: ElicitationSchema( + properties: ["age": .number(NumberSchema(isInteger: true, title: "Age"))], + required: ["age"] + ) + ) + + guard ageResult.action == .accept, + let age = ageResult.content?["age"]?.intValue else { + return CallTool.Result(content: [.text("Age step failed")], isError: true) + } + + // Third elicitation - get city + let cityResult = try await context.elicit( + message: "What is your city?", + requestedSchema: ElicitationSchema( + properties: ["city": .string(StringSchema(title: "City"))], + required: ["city"] + ) + ) + + guard cityResult.action == .accept, + let city = cityResult.content?["city"]?.stringValue else { + return CallTool.Result(content: [.text("City step failed")], isError: true) + } + + return CallTool.Result(content: [.text("Hello \(name), age \(age), from \(city)!")]) + } + + let client = Client(name: "TestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(form: Client.Capabilities.Elicitation.Form()) + )) + + actor RequestCounter { + var count = 0 + func increment() { count += 1 } + } + let counter = RequestCounter() + + await client.withElicitationHandler { params, _ in + await counter.increment() + guard case .form(let formParams) = params else { + return ElicitResult(action: .decline) + } + + if formParams.message.contains("name") { + return ElicitResult(action: .accept, content: ["name": .string("Alice")]) + } else if formParams.message.contains("age") { + return ElicitResult(action: .accept, content: ["age": .int(30)]) + } else if formParams.message.contains("city") { + return ElicitResult(action: .accept, content: ["city": .string("New York")]) + } + return ElicitResult(action: .decline) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "wizard", arguments: [:]) + + let requestCount = await counter.count + #expect(requestCount == 3, "Should have made 3 elicitation requests") + + if case .text(let text, _, _) = result.content.first { + #expect(text == "Hello Alice, age 30, from New York!") + } else { + Issue.record("Expected text content") + } + + await client.disconnect() + } + + // MARK: - Sampling Handler Context Access Tests + + /// Test that sampling handler can access context.requestId. + /// Based on Python SDK's sampling callback context access patterns. + @Test("Sampling handler can access context.requestId") + func testSamplingHandlerCanAccessRequestId() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + actor RequestIdTracker { + var receivedRequestId: RequestId? + func set(_ id: RequestId) { receivedRequestId = id } + } + let tracker = RequestIdTracker() + + let server = Server( + name: "TestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "askLLM", description: "Ask LLM", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] _, _ in + let params = SamplingParameters( + messages: [.user("Hello")], + maxTokens: 100 + ) + let result = try await server.createMessage(params) + return CallTool.Result(content: [.text("LLM said: \(result.model)")]) + } + + let client = Client(name: "TestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities(sampling: .init())) + + await client.withSamplingHandler { _, context in + // Sampling handler accesses context.requestId + await tracker.set(context.requestId) + return ClientSamplingRequest.Result( + model: "test-model", + stopReason: .endTurn, + role: .assistant, + content: [.text("Hello from LLM")] + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + _ = try await client.callTool(name: "askLLM", arguments: [:]) + + let receivedId = await tracker.receivedRequestId + #expect(receivedId != nil, "Sampling handler should have access to requestId") + + await client.disconnect() + } + + /// Test that sampling handler can access context._meta when present. + @Test("Sampling handler can access context._meta") + func testSamplingHandlerCanAccessMeta() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + actor MetaTracker { + var receivedMeta: RequestMeta? + func set(_ meta: RequestMeta?) { receivedMeta = meta } + } + let tracker = MetaTracker() + + let server = Server( + name: "TestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "askLLM", description: "Ask LLM", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] _, _ in + var params = SamplingParameters( + messages: [.user("Hello")], + maxTokens: 100 + ) + params._meta = RequestMeta(progressToken: .string("sampling-token-123")) + let result = try await server.createMessage(params) + return CallTool.Result(content: [.text("LLM said: \(result.model)")]) + } + + let client = Client(name: "TestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities(sampling: .init())) + + await client.withSamplingHandler { _, context in + // Sampling handler accesses context._meta + await tracker.set(context._meta) + return ClientSamplingRequest.Result( + model: "test-model", + stopReason: .endTurn, + role: .assistant, + content: [.text("Hello from LLM")] + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + _ = try await client.callTool(name: "askLLM", arguments: [:]) + + let receivedMeta = await tracker.receivedMeta + #expect(receivedMeta != nil, "Sampling handler should have access to _meta") + #expect(receivedMeta?.progressToken == .string("sampling-token-123"), "progressToken should match") + + await client.disconnect() + } + + // MARK: - Roots Handler Context Access Tests + + /// Test that roots handler can access context.requestId. + /// Based on Python SDK's list_roots callback context access patterns. + @Test("Roots handler can access context.requestId") + func testRootsHandlerCanAccessRequestId() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + actor RequestIdTracker { + var receivedRequestId: RequestId? + func set(_ id: RequestId) { receivedRequestId = id } + } + let tracker = RequestIdTracker() + + let server = Server( + name: "TestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + ListTools.Result(tools: [ + Tool(name: "getRoots", description: "Get roots", inputSchema: [:]) + ]) + } + + await server.withRequestHandler(CallTool.self) { [server] _, _ in + let roots = try await server.listRoots() + return CallTool.Result(content: [.text("Found \(roots.count) roots")]) + } + + let client = Client(name: "TestClient", version: "1.0.0") + await client.setCapabilities(Client.Capabilities(roots: .init(listChanged: true))) + + // Use withRootsHandler with context parameter + await client.withRootsHandler { context in + // Roots handler accesses context.requestId + await tracker.set(context.requestId) + return [Root(uri: "file:///test/path", name: "Test Root")] + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let result = try await client.callTool(name: "getRoots", arguments: [:]) + + let receivedId = await tracker.receivedRequestId + #expect(receivedId != nil, "Roots handler should have access to requestId") + + if case .text(let text, _, _) = result.content.first { + #expect(text == "Found 1 roots") + } + + await client.disconnect() + } +} diff --git a/Tests/MCPTests/ResourceTests.swift b/Tests/MCPTests/ResourceTests.swift index b75fdbdb..45b848d6 100644 --- a/Tests/MCPTests/ResourceTests.swift +++ b/Tests/MCPTests/ResourceTests.swift @@ -743,15 +743,15 @@ struct ResourceTests { #expect(templates.templates[1].uriTemplate == "users://{user_id}/profile") // Read a resource using a templated URI - let greetingContent = try await client.readResource(uri: "greeting://World") - #expect(greetingContent.count == 1) - #expect(greetingContent[0].text == "Hello, World!") - #expect(greetingContent[0].mimeType == "text/plain") + let greetingResult = try await client.readResource(uri: "greeting://World") + #expect(greetingResult.contents.count == 1) + #expect(greetingResult.contents[0].text == "Hello, World!") + #expect(greetingResult.contents[0].mimeType == "text/plain") // Read another templated resource - let profileContent = try await client.readResource(uri: "users://123/profile") - #expect(profileContent.count == 1) - #expect(profileContent[0].text == "Profile for user 123") + let profileResult = try await client.readResource(uri: "users://123/profile") + #expect(profileResult.contents.count == 1) + #expect(profileResult.contents[0].text == "Profile for user 123") await client.disconnect() await server.stop() @@ -807,9 +807,9 @@ struct ResourceTests { #expect(resources.resources[1].mimeType == "application/json; charset=utf-8") // Read resource and verify MIME type is preserved - let content = try await client.readResource(uri: "ui://widget") - #expect(content.count == 1) - #expect(content[0].mimeType == "text/html;profile=mcp-app") + let readResult = try await client.readResource(uri: "ui://widget") + #expect(readResult.contents.count == 1) + #expect(readResult.contents[0].mimeType == "text/html;profile=mcp-app") await client.disconnect() await server.stop() diff --git a/Tests/MCPTests/RootsTests.swift b/Tests/MCPTests/RootsTests.swift index ade1061a..a619b09b 100644 --- a/Tests/MCPTests/RootsTests.swift +++ b/Tests/MCPTests/RootsTests.swift @@ -364,7 +364,7 @@ struct RootsIntegrationTests { await client.setCapabilities(.init(roots: .init(listChanged: true))) // Register roots handler (required since we declared the capability) - await client.withRootsHandler { + await client.withRootsHandler { _ in [Root(uri: "file:///test/path")] } @@ -425,7 +425,7 @@ struct RootsIntegrationTests { version: "1.0" ) await client.setCapabilities(.init(roots: .init(listChanged: true))) - await client.withRootsHandler { + await client.withRootsHandler { _ in expectedRoots } @@ -542,7 +542,7 @@ struct RootsIntegrationTests { version: "1.0" ) await client.setCapabilities(.init(roots: .init(listChanged: true))) - await client.withRootsHandler { + await client.withRootsHandler { _ in [Root(uri: "file:///path")] } diff --git a/Tests/MCPTests/RoundtripTests.swift b/Tests/MCPTests/RoundtripTests.swift index 58606d82..11e83c3f 100644 --- a/Tests/MCPTests/RoundtripTests.swift +++ b/Tests/MCPTests/RoundtripTests.swift @@ -147,9 +147,9 @@ struct RoundtripTests { } let listToolsTask = Task { - let (tools, _) = try await client.listTools() - #expect(tools.count == 1) - #expect(tools[0].name == "add") + let result = try await client.listTools() + #expect(result.tools.count == 1) + #expect(result.tools[0].name == "add") } let callToolTask = Task { @@ -188,8 +188,8 @@ struct RoundtripTests { // Test reading a resource let readResourceTask = Task { let result = try await client.readResource(uri: "test://example.txt") - #expect(result.count == 1) - #expect(result[0] == .text("Hello, World!", uri: "test://example.txt")) + #expect(result.contents.count == 1) + #expect(result.contents[0] == .text("Hello, World!", uri: "test://example.txt")) } try await withThrowingTaskGroup(of: Void.self) { group in diff --git a/Tests/MCPTests/ServerTests.swift b/Tests/MCPTests/ServerTests.swift index d65be872..5347408e 100644 --- a/Tests/MCPTests/ServerTests.swift +++ b/Tests/MCPTests/ServerTests.swift @@ -245,4 +245,234 @@ struct ServerTests { await server.stop() await transport.disconnect() } + + // MARK: - Ping Before Initialization Tests + // Based on Python SDK: test_ping_request_before_initialization + + @Test("Ping request allowed before initialization") + func testPingRequestAllowedBeforeInitialization() async throws { + // Per MCP spec, ping requests should be allowed before initialization + // This is important for health checks and connection verification + let transport = MockTransport() + let server = Server(name: "TestServer", version: "1.0") + + try await server.start(transport: transport) + + // Send ping request BEFORE sending initialize request + let pingRequest = """ + {"jsonrpc":"2.0","method":"ping","id":42} + """ + await transport.queueRaw(pingRequest) + + // Wait for ping response + let pingReceived = await transport.waitForSentMessage { message in + message.contains("\"id\":42") || message.contains("\"id\": 42") + } + #expect(pingReceived, "Timed out waiting for ping response") + + let messages = await transport.sentMessages + #expect(messages.count >= 1, "Should have received a response") + + // Verify we got a successful response (not an error about not being initialized) + if let response = messages.first { + #expect(response.contains("\"result\""), "Should have a result, not an error") + #expect(!response.contains("\"error\""), "Should not be an error response") + } + + await server.stop() + await transport.disconnect() + } + + // MARK: - Requests Before Initialization Behavior (Server level) + // + // MCP spec says clients "SHOULD NOT" send requests before initialization. + // Swift SDK aligns with Python SDK behavior: blocks at Server level for all transports. + // + // SDK behavior comparison: + // - Python: Blocks non-ping requests at session level (all transports) + // - TypeScript: Server allows requests; HTTP transport blocks for session management + // - Swift: Blocks non-ping requests at Server level (all transports) - matches Python + // + // We chose Python's approach for consistency across transports and better spec alignment. + + @Test("Server blocks non-ping requests before initialization (default strict mode)") + func testServerBlocksRequestsBeforeInitialization() async throws { + // MCP spec (lifecycle.mdx) says: + // "The client SHOULD NOT send requests other than pings before the server + // has responded to the initialize request." + // + // Swift SDK enforces this at the Server level (like Python), not just at + // HTTP transport level (like TypeScript). This provides consistent behavior + // across all transports. + + let transport = MockTransport() + let server = Server( + name: "TestServer", + version: "1.0", + capabilities: .init(prompts: .init()) + // Uses default configuration which has strict: true + ) + + // Register a prompts handler + await server.withRequestHandler(ListPrompts.self) { _, _ in + ListPrompts.Result(prompts: []) + } + + try await server.start(transport: transport) + + // Send prompts/list request BEFORE initialize + let promptsRequest = """ + {"jsonrpc":"2.0","method":"prompts/list","id":1} + """ + await transport.queueRaw(promptsRequest) + + // Wait for response + let responseReceived = await transport.waitForSentMessage { message in + message.contains("\"id\":1") || message.contains("\"id\": 1") + } + #expect(responseReceived, "Timed out waiting for response") + + let messages = await transport.sentMessages + #expect(messages.count >= 1, "Should have received a response") + + // Server should reject the request with an error + if let response = messages.first { + #expect(response.contains("\"error\""), "Should be an error response") + #expect( + response.contains("not initialized") || response.contains("Server is not initialized"), + "Error should indicate initialization required: \(response)" + ) + } + + await server.stop() + await transport.disconnect() + } + + @Test("Server allows requests before initialization in lenient mode") + func testServerAllowsRequestsBeforeInitializationInLenientMode() async throws { + // Lenient mode matches TypeScript SDK's server-level behavior + let transport = MockTransport() + let server = Server( + name: "TestServer", + version: "1.0", + capabilities: .init(prompts: .init()), + configuration: .lenient + ) + + // Register a prompts handler + await server.withRequestHandler(ListPrompts.self) { _, _ in + ListPrompts.Result(prompts: []) + } + + try await server.start(transport: transport) + + // Send prompts/list request BEFORE initialize + let promptsRequest = """ + {"jsonrpc":"2.0","method":"prompts/list","id":1} + """ + await transport.queueRaw(promptsRequest) + + // Wait for response + let responseReceived = await transport.waitForSentMessage { message in + message.contains("\"id\":1") || message.contains("\"id\": 1") + } + #expect(responseReceived, "Timed out waiting for response") + + let messages = await transport.sentMessages + #expect(messages.count >= 1, "Should have received a response") + + // Lenient mode: processes the request and returns result + if let response = messages.first { + #expect(response.contains("\"result\""), "Lenient mode should process requests before init") + } + + await server.stop() + await transport.disconnect() + } + + // MARK: - Protocol Version Negotiation Tests + // Based on Python SDK: test_server_session_initialize_with_older_protocol_version + + @Test("Server responds with client's requested protocol version when supported") + func testServerRespondsWithClientRequestedProtocolVersion() async throws { + // When a client requests an older but supported protocol version, + // the server should respond with that version, not the latest + let transport = MockTransport() + let server = Server(name: "TestServer", version: "1.0") + + try await server.start(transport: transport) + + // Client requests older supported version + let olderVersion = Version.v2024_11_05 + try await transport.queue( + request: Initialize.request( + .init( + protocolVersion: olderVersion, + capabilities: .init(), + clientInfo: .init(name: "OlderClient", version: "1.0") + ) + ) + ) + + // Wait for response + let received = await transport.waitForSentMessageCount(1) + #expect(received, "Timed out waiting for initialize response") + + let messages = await transport.sentMessages + #expect(messages.count >= 1) + + // Verify the server responded with the requested protocol version + if let response = messages.first { + #expect(response.contains("serverInfo")) + #expect(response.contains("protocolVersion")) + // Server should echo back the client's requested version + #expect( + response.contains("\"\(olderVersion)\""), + "Server should respond with client's requested version \(olderVersion), got: \(response)" + ) + } + + await server.stop() + await transport.disconnect() + } + + @Test("Server defaults to latest version for unsupported client version") + func testServerDefaultsToLatestForUnsupportedVersion() async throws { + // When a client requests an unsupported version, server should use latest + let transport = MockTransport() + let server = Server(name: "TestServer", version: "1.0") + + try await server.start(transport: transport) + + // Client requests unsupported version + let unsupportedVersion = "2023-01-01" + let request = Initialize.request( + .init( + protocolVersion: unsupportedVersion, + capabilities: .init(), + clientInfo: .init(name: "OldClient", version: "1.0") + ) + ) + try await transport.queue(request: request) + + // Wait for response + let received = await transport.waitForSentMessageCount(1) + #expect(received, "Timed out waiting for initialize response") + + let messages = await transport.sentMessages + #expect(messages.count >= 1) + + // Verify the server responded with the latest version (negotiation fallback) + if let response = messages.first { + #expect(response.contains("serverInfo")) + #expect(response.contains("protocolVersion")) + #expect( + response.contains("\"\(Version.latest)\""), + "Server should fall back to latest version for unsupported client version, got: \(response)" + ) + } + + await server.stop() + await transport.disconnect() + } } diff --git a/Tests/MCPTests/SessionLifecycleTests.swift b/Tests/MCPTests/SessionLifecycleTests.swift new file mode 100644 index 00000000..b670f663 --- /dev/null +++ b/Tests/MCPTests/SessionLifecycleTests.swift @@ -0,0 +1,560 @@ +import Foundation +import Logging +import Testing + +#if canImport(System) + import System +#else + @preconcurrency import SystemPackage +#endif + +@testable import MCP + +// MARK: - Session Lifecycle Tests + +/// Tests for session lifecycle functionality in MCP. +/// +/// These tests cover session initialization, client info handling, capability +/// exposure, and race conditions that can occur during the initialization flow. +/// +/// Note: This is distinct from SessionManagerTests.swift which tests the +/// SessionManager actor for HTTP server session storage/retrieval. +/// +/// Based on Python SDK tests: +/// - tests/client/test_session.py +/// - tests/server/test_session.py +/// - tests/server/test_session_race_condition.py +@Suite("Session Lifecycle Tests") +struct SessionLifecycleTests { + + // MARK: - Request Immediately After Initialize Response Tests + + @Suite("Request immediately after initialize response (race condition)") + struct InitializeRaceConditionTests { + + /// Test that requests are accepted immediately after initialize response. + /// + /// This reproduces the race condition in stateful HTTP mode where: + /// 1. Client sends InitializeRequest + /// 2. Server responds with InitializeResult + /// 3. Client immediately sends tools/list (before server receives InitializedNotification) + /// 4. Without fix: Server rejects with "Received request before initialization was complete" + /// 5. With fix: Server accepts and processes the request + /// + /// This test simulates the HTTP transport behavior where InitializedNotification + /// may arrive in a separate POST request after other requests. + /// + /// Based on Python SDK: tests/server/test_session_race_condition.py + @Test(.timeLimit(.minutes(1))) + func requestImmediatelyAfterInitializeResponse() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (_, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.race-condition") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + + let toolsListSuccess = ToolsListSuccessTracker() + + // Set up server with tools capability + let server = Server( + name: "RaceConditionTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + await toolsListSuccess.markSuccess() + return ListTools.Result(tools: [ + Tool( + name: "example_tool", + description: "An example tool", + inputSchema: ["type": "object"] + ) + ]) + } + + // Start the server + try await server.start(transport: serverTransport) + + // Wait for server to be ready + try await Task.sleep(for: .milliseconds(50)) + + // Simulate client behavior manually (like HTTP transport race condition) + let encoder = JSONEncoder() + + // Step 1: Send Initialize request + let initRequest = Request( + id: .number(1), + method: Initialize.name, + params: Initialize.Parameters( + protocolVersion: Version.latest, + capabilities: .init(), + clientInfo: .init(name: "race-condition-client", version: "1.0") + ) + ) + + let initData = try encoder.encode(initRequest) + _ = try clientToServerWrite.writeAll(initData) + _ = try clientToServerWrite.writeAll("\n".data(using: .utf8)!) + + // Wait for and read the initialize response + try await Task.sleep(for: .milliseconds(100)) + + // Step 2: Immediately send tools/list BEFORE InitializedNotification + // This is the race condition scenario + let toolsListRequest = Request( + id: .number(2), + method: ListTools.name, + params: ListTools.Parameters() + ) + + let toolsData = try encoder.encode(toolsListRequest) + _ = try clientToServerWrite.writeAll(toolsData) + _ = try clientToServerWrite.writeAll("\n".data(using: .utf8)!) + + // Wait for tools/list to be processed + try await Task.sleep(for: .milliseconds(200)) + + // Step 3: Now send InitializedNotification + let initializedNotification = InitializedNotification.message(.init()) + let notifData = try encoder.encode(initializedNotification) + _ = try clientToServerWrite.writeAll(notifData) + _ = try clientToServerWrite.writeAll("\n".data(using: .utf8)!) + + // Give time for all messages to be processed + try await Task.sleep(for: .milliseconds(100)) + + // Verify tools/list succeeded (race condition was handled correctly) + let success = await toolsListSuccess.wasSuccessful + #expect(success, "tools/list should succeed immediately after initialize response, before InitializedNotification") + + // Clean up + await server.stop() + } + + /// Test that server in lenient mode accepts requests before initialized notification. + /// + /// In lenient mode, the server should accept any request after receiving + /// the initialize request, without waiting for InitializedNotification. + @Test(.timeLimit(.minutes(1))) + func lenientModeAcceptsEarlyRequests() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.lenient-mode") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + // Set up server in lenient mode (default) + let server = Server( + name: "LenientModeServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + return ListTools.Result(tools: [ + Tool(name: "test_tool", inputSchema: ["type": "object"]) + ]) + } + + let client = Client(name: "LenientModeClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Make a request - should succeed in lenient mode + let tools = try await client.send(ListTools.request(.init())) + #expect(tools.tools.count == 1) + #expect(tools.tools.first?.name == "test_tool") + + await client.disconnect() + await server.stop() + } + } + + // MARK: - Client Info Tests + + @Suite("Client info handling") + struct ClientInfoTests { + + /// Test that custom client info is properly sent during initialization. + /// + /// Based on Python SDK: test_client_session_custom_client_info + @Test(.timeLimit(.minutes(1))) + func customClientInfoSentDuringInitialization() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.custom-client-info") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let receivedClientInfo = ReceivedClientInfoTracker() + + let server = Server( + name: "ClientInfoTestServer", + version: "1.0.0", + capabilities: .init() + ) + + // Use custom client info + let customName = "custom-test-client" + let customVersion = "2.3.4" + let client = Client(name: customName, version: customVersion) + + // Track received client info via initialize hook (trailing closure on start) + try await server.start(transport: serverTransport) { clientInfo, _ in + await receivedClientInfo.set(clientInfo) + } + _ = try await client.connect(transport: clientTransport) + + // Give time for hook to be called + try await Task.sleep(for: .milliseconds(100)) + + // Verify the custom client info was received + let info = await receivedClientInfo.info + #expect(info != nil, "Server should have received client info") + #expect(info?.name == customName, "Server should receive custom client name") + #expect(info?.version == customVersion, "Server should receive custom client version") + + await client.disconnect() + await server.stop() + } + + /// Test that default client info is properly sent during initialization. + /// + /// Based on Python SDK: test_client_session_default_client_info + @Test(.timeLimit(.minutes(1))) + func defaultClientInfoSentDuringInitialization() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.default-client-info") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let receivedClientInfo = ReceivedClientInfoTracker() + + let server = Server( + name: "DefaultClientInfoServer", + version: "1.0.0", + capabilities: .init() + ) + + // Use minimal client (name and version are required in Swift) + let client = Client(name: "test-app", version: "1.0") + + // Track received client info via initialize hook (trailing closure on start) + try await server.start(transport: serverTransport) { clientInfo, _ in + await receivedClientInfo.set(clientInfo) + } + _ = try await client.connect(transport: clientTransport) + + // Give time for hook to be called + try await Task.sleep(for: .milliseconds(100)) + + // Verify client info was received and has expected values + let info = await receivedClientInfo.info + #expect(info != nil, "Server should have received client info") + #expect(info?.name == "test-app") + #expect(info?.version == "1.0") + + await client.disconnect() + await server.stop() + } + } + + // MARK: - Server Capabilities Tests + + @Suite("Server capabilities") + struct ServerCapabilitiesTests { + + /// Test that getServerCapabilities returns nil before init and capabilities after. + /// + /// Based on Python SDK: test_get_server_capabilities + @Test(.timeLimit(.minutes(1))) + func getServerCapabilitiesBeforeAndAfterInit() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.server-capabilities") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + // Server with various capabilities enabled + let server = Server( + name: "CapabilitiesTestServer", + version: "1.0.0", + capabilities: .init( + logging: .init(), + prompts: .init(listChanged: true), + resources: .init(subscribe: true, listChanged: true), + tools: .init(listChanged: false) + ) + ) + + // Register minimal handlers so capabilities are advertised + await server.withRequestHandler(ListPrompts.self) { _, _ in + return ListPrompts.Result(prompts: []) + } + await server.withRequestHandler(ListResources.self) { _, _ in + return ListResources.Result(resources: []) + } + await server.withRequestHandler(ListTools.self) { _, _ in + return ListTools.Result(tools: []) + } + + let client = Client(name: "CapabilitiesTestClient", version: "1.0") + + // Check capabilities before connection - should be nil + let capabilitiesBeforeConnect = await client.serverCapabilities + #expect(capabilitiesBeforeConnect == nil, "Capabilities should be nil before connect") + + try await server.start(transport: serverTransport) + + // Connect and verify capabilities + let initResult = try await client.connect(transport: clientTransport) + + // Check capabilities after connection - should be populated + let capabilitiesAfterConnect = await client.serverCapabilities + #expect(capabilitiesAfterConnect != nil, "Capabilities should be set after connect") + + // Verify specific capabilities + #expect(capabilitiesAfterConnect?.logging != nil, "Logging capability should be present") + #expect(capabilitiesAfterConnect?.prompts != nil, "Prompts capability should be present") + #expect(capabilitiesAfterConnect?.prompts?.listChanged == true, "Prompts listChanged should be true") + #expect(capabilitiesAfterConnect?.resources != nil, "Resources capability should be present") + #expect(capabilitiesAfterConnect?.resources?.subscribe == true, "Resources subscribe should be true") + #expect(capabilitiesAfterConnect?.tools != nil, "Tools capability should be present") + #expect(capabilitiesAfterConnect?.tools?.listChanged == false, "Tools listChanged should be false") + + // Verify init result matches + #expect(initResult.capabilities.logging != nil) + #expect(initResult.capabilities.prompts != nil) + #expect(initResult.capabilities.resources != nil) + #expect(initResult.capabilities.tools != nil) + + await client.disconnect() + await server.stop() + + // After disconnect, capabilities should still be available (cached) + // This matches the Swift SDK behavior where we cache the last known capabilities + let capabilitiesAfterDisconnect = await client.serverCapabilities + #expect(capabilitiesAfterDisconnect != nil, "Capabilities remain cached after disconnect") + } + } + + // MARK: - In-Flight Request Tracking Tests + + @Suite("In-flight request tracking") + struct InFlightRequestTrackingTests { + + /// Test that in-flight request tracking is cleared after request completes. + /// + /// This verifies that the internal tracking of pending requests is properly + /// cleaned up after responses are received, preventing memory leaks. + /// + /// Based on Python SDK: test_in_flight_requests_cleared_after_completion + @Test(.timeLimit(.minutes(1))) + func inFlightRequestsClearedAfterCompletion() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.in-flight-cleanup") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let server = Server( + name: "InFlightTestServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + await server.withRequestHandler(ListTools.self) { _, _ in + return ListTools.Result(tools: [ + Tool(name: "test_tool", inputSchema: ["type": "object"]) + ]) + } + + let client = Client(name: "InFlightTestClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Send multiple requests and verify they complete + for i in 1...5 { + let tools = try await client.send(ListTools.request(.init())) + #expect(tools.tools.count == 1, "Request \(i) should succeed") + } + + // The client should still be functional after all requests complete + // This indirectly verifies that in-flight tracking was cleaned up + let finalTools = try await client.send(ListTools.request(.init())) + #expect(finalTools.tools.first?.name == "test_tool") + + await client.disconnect() + await server.stop() + } + + /// Test that multiple concurrent requests are properly tracked and cleaned up. + @Test(.timeLimit(.minutes(1))) + func concurrentRequestsTrackedAndCleanedUp() async throws { + let (clientToServerRead, clientToServerWrite) = try FileDescriptor.pipe() + let (serverToClientRead, serverToClientWrite) = try FileDescriptor.pipe() + + var logger = Logger(label: "mcp.test.concurrent-in-flight") + logger.logLevel = .warning + + let serverTransport = StdioTransport( + input: clientToServerRead, + output: serverToClientWrite, + logger: logger + ) + let clientTransport = StdioTransport( + input: serverToClientRead, + output: clientToServerWrite, + logger: logger + ) + + let server = Server( + name: "ConcurrentInFlightServer", + version: "1.0.0", + capabilities: .init(tools: .init()) + ) + + let callCount = CallCountTracker() + + await server.withRequestHandler(ListTools.self) { _, _ in + return ListTools.Result(tools: []) + } + + await server.withRequestHandler(CallTool.self) { request, _ in + // Small delay to ensure requests overlap + let delay = request.arguments?["delay"]?.doubleValue ?? 0.05 + try? await Task.sleep(for: .seconds(delay)) + await callCount.increment() + return CallTool.Result(content: [.text("Done")]) + } + + let client = Client(name: "ConcurrentInFlightClient", version: "1.0") + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Send multiple concurrent requests + await withTaskGroup(of: Void.self) { group in + for i in 0..<10 { + group.addTask { + let delay = Double(i % 3) * 0.01 + 0.01 + _ = try? await client.send( + CallTool.request(.init( + name: "test", + arguments: ["delay": .double(delay)] + )) + ) + } + } + } + + // All 10 requests should have been processed + let count = await callCount.count + #expect(count == 10, "All concurrent requests should complete") + + // Client should still be functional + let tools = try await client.send(ListTools.request(.init())) + #expect(tools.tools.isEmpty) + + await client.disconnect() + await server.stop() + } + } +} + +// MARK: - Helper Actors + +private actor ToolsListSuccessTracker { + private var _success = false + + func markSuccess() { + _success = true + } + + var wasSuccessful: Bool { _success } +} + +private actor ReceivedClientInfoTracker { + private var _info: Client.Info? + + func set(_ info: Client.Info) { + _info = info + } + + var info: Client.Info? { _info } +} + +private actor CallCountTracker { + private var _count = 0 + + func increment() { + _count += 1 + } + + var count: Int { _count } +} diff --git a/Tests/MCPTests/StreamableHTTPServerTransportTests.swift b/Tests/MCPTests/StreamableHTTPServerTransportTests.swift index 47ab5b4a..85476a21 100644 --- a/Tests/MCPTests/StreamableHTTPServerTransportTests.swift +++ b/Tests/MCPTests/StreamableHTTPServerTransportTests.swift @@ -957,6 +957,33 @@ struct HTTPServerTransportTests { #expect(TransportSecuritySettings.forBindAddress(host: "192.168.1.1", port: 8080) == nil) } + @Test("HTTPServerTransportOptions.forBindAddress auto-configures security") + func optionsForBindAddressAutoConfiguresSecurity() { + // For localhost, should auto-enable DNS rebinding protection + let localhostOptions = HTTPServerTransportOptions.forBindAddress(host: "127.0.0.1", port: 8080) + #expect(localhostOptions.security != nil) + #expect(localhostOptions.security?.enableDnsRebindingProtection == true) + #expect(localhostOptions.security?.allowedHosts.contains("127.0.0.1:8080") == true) + + // For 0.0.0.0, should not auto-enable security + let wildcardOptions = HTTPServerTransportOptions.forBindAddress(host: "0.0.0.0", port: 8080) + #expect(wildcardOptions.security == nil) + + // Should allow explicit security override + let customSecurity = TransportSecuritySettings( + enableDnsRebindingProtection: true, + allowedHosts: ["custom.local:8080"], + allowedOrigins: ["http://custom.local:8080"] + ) + let overriddenOptions = HTTPServerTransportOptions.forBindAddress( + host: "127.0.0.1", + port: 8080, + security: customSecurity + ) + #expect(overriddenOptions.security?.allowedHosts.contains("custom.local:8080") == true) + #expect(overriddenOptions.security?.allowedHosts.contains("127.0.0.1:8080") == false) + } + @Test("DNS rebinding protection rejects invalid host on GET") func dnsRebindingProtectionRejectsInvalidHostOnGet() async throws { let transport = HTTPServerTransport( diff --git a/Tests/MCPTests/TaskModeValidationTests.swift b/Tests/MCPTests/TaskModeValidationTests.swift new file mode 100644 index 00000000..3cb774c6 --- /dev/null +++ b/Tests/MCPTests/TaskModeValidationTests.swift @@ -0,0 +1,515 @@ +import Foundation +import Testing + +@testable import MCP + +/// Tests for task mode validation functions. +/// +/// These tests verify the `validateTaskMode` and `canUseToolWithTaskMode` functions +/// which are the Swift equivalents of Python SDK's `Experimental` class methods: +/// - `is_task` → checked via `isTaskRequest` parameter +/// - `validate_task_mode()` → `validateTaskMode(isTaskRequest:taskSupport:)` +/// - `validate_for_tool()` → `validateTaskMode(isTaskRequest:for:)` +/// - `can_use_tool()` → `canUseToolWithTaskMode(clientSupportsTask:taskSupport:)` +/// +/// Based on Python SDK's `tests/experimental/tasks/test_request_context.py` + +// MARK: - validateTaskMode Tests + +@Suite("validateTaskMode Tests") +struct ValidateTaskModeTests { + + // MARK: - Required Mode Tests + + @Test("REQUIRED mode with task request is valid") + func testRequiredWithTaskRequestIsValid() throws { + // Python: test_validate_task_mode_required_with_task_is_valid + #expect(throws: Never.self) { + try validateTaskMode(isTaskRequest: true, taskSupport: .required) + } + } + + @Test("REQUIRED mode without task request throws error") + func testRequiredWithoutTaskRequestThrows() throws { + // Python: test_validate_task_mode_required_without_task_returns_error + #expect(throws: MCPError.self) { + try validateTaskMode(isTaskRequest: false, taskSupport: .required) + } + } + + @Test("REQUIRED mode error contains descriptive message") + func testRequiredErrorMessage() throws { + // Python: test_validate_task_mode_required_without_task_raises_by_default + do { + try validateTaskMode(isTaskRequest: false, taskSupport: .required) + Issue.record("Expected MCPError to be thrown") + } catch let error as MCPError { + let description = String(describing: error) + #expect(description.contains("requires task-augmented")) + } + } + + // MARK: - Forbidden Mode Tests + + @Test("FORBIDDEN mode without task request is valid") + func testForbiddenWithoutTaskRequestIsValid() throws { + // Python: test_validate_task_mode_forbidden_without_task_is_valid + #expect(throws: Never.self) { + try validateTaskMode(isTaskRequest: false, taskSupport: .forbidden) + } + } + + @Test("FORBIDDEN mode with task request throws error") + func testForbiddenWithTaskRequestThrows() throws { + // Python: test_validate_task_mode_forbidden_with_task_returns_error + #expect(throws: MCPError.self) { + try validateTaskMode(isTaskRequest: true, taskSupport: .forbidden) + } + } + + @Test("FORBIDDEN mode error contains descriptive message") + func testForbiddenErrorMessage() throws { + // Python: test_validate_task_mode_forbidden_with_task_raises_by_default + do { + try validateTaskMode(isTaskRequest: true, taskSupport: .forbidden) + Issue.record("Expected MCPError to be thrown") + } catch let error as MCPError { + let description = String(describing: error) + #expect(description.contains("does not support task-augmented")) + } + } + + // MARK: - nil Mode (Treated as Forbidden) Tests + + @Test("nil mode treated as FORBIDDEN - task request throws") + func testNilModeTreatedAsForbidden() throws { + // Python: test_validate_task_mode_none_treated_as_forbidden + #expect(throws: MCPError.self) { + try validateTaskMode(isTaskRequest: true, taskSupport: nil) + } + } + + @Test("nil mode without task request is valid") + func testNilModeWithoutTaskRequestIsValid() throws { + #expect(throws: Never.self) { + try validateTaskMode(isTaskRequest: false, taskSupport: nil) + } + } + + // MARK: - Optional Mode Tests + + @Test("OPTIONAL mode with task request is valid") + func testOptionalWithTaskRequestIsValid() throws { + // Python: test_validate_task_mode_optional_with_task_is_valid + #expect(throws: Never.self) { + try validateTaskMode(isTaskRequest: true, taskSupport: .optional) + } + } + + @Test("OPTIONAL mode without task request is valid") + func testOptionalWithoutTaskRequestIsValid() throws { + // Python: test_validate_task_mode_optional_without_task_is_valid + #expect(throws: Never.self) { + try validateTaskMode(isTaskRequest: false, taskSupport: .optional) + } + } +} + +// MARK: - validateTaskMode for Tool Tests + +@Suite("validateTaskMode for Tool Tests") +struct ValidateTaskModeForToolTests { + + @Test("Tool with execution.taskSupport=required rejects non-task request") + func testToolWithRequiredRejectsNonTask() throws { + // Python: test_validate_for_tool_with_execution_required + let tool = Tool( + name: "test", + description: "test", + inputSchema: ["type": "object"], + execution: Tool.Execution(taskSupport: .required) + ) + + #expect(throws: MCPError.self) { + try validateTaskMode(isTaskRequest: false, for: tool) + } + } + + @Test("Tool with execution.taskSupport=required accepts task request") + func testToolWithRequiredAcceptsTask() throws { + let tool = Tool( + name: "test", + description: "test", + inputSchema: ["type": "object"], + execution: Tool.Execution(taskSupport: .required) + ) + + #expect(throws: Never.self) { + try validateTaskMode(isTaskRequest: true, for: tool) + } + } + + @Test("Tool without execution (nil) rejects task request") + func testToolWithoutExecutionRejectsTask() throws { + // Python: test_validate_for_tool_without_execution + let tool = Tool( + name: "test", + description: "test", + inputSchema: ["type": "object"], + execution: nil + ) + + #expect(throws: MCPError.self) { + try validateTaskMode(isTaskRequest: true, for: tool) + } + } + + @Test("Tool without execution (nil) accepts non-task request") + func testToolWithoutExecutionAcceptsNonTask() throws { + let tool = Tool( + name: "test", + description: "test", + inputSchema: ["type": "object"], + execution: nil + ) + + #expect(throws: Never.self) { + try validateTaskMode(isTaskRequest: false, for: tool) + } + } + + @Test("Tool with execution.taskSupport=optional accepts task request") + func testToolWithOptionalAcceptsTask() throws { + // Python: test_validate_for_tool_optional_with_task + let tool = Tool( + name: "test", + description: "test", + inputSchema: ["type": "object"], + execution: Tool.Execution(taskSupport: .optional) + ) + + #expect(throws: Never.self) { + try validateTaskMode(isTaskRequest: true, for: tool) + } + } + + @Test("Tool with execution.taskSupport=optional accepts non-task request") + func testToolWithOptionalAcceptsNonTask() throws { + let tool = Tool( + name: "test", + description: "test", + inputSchema: ["type": "object"], + execution: Tool.Execution(taskSupport: .optional) + ) + + #expect(throws: Never.self) { + try validateTaskMode(isTaskRequest: false, for: tool) + } + } + + @Test("Tool with execution but nil taskSupport treats as forbidden") + func testToolWithExecutionButNilTaskSupport() throws { + let tool = Tool( + name: "test", + description: "test", + inputSchema: ["type": "object"], + execution: Tool.Execution(taskSupport: nil) + ) + + // Task request should be rejected (nil = forbidden) + #expect(throws: MCPError.self) { + try validateTaskMode(isTaskRequest: true, for: tool) + } + + // Non-task request should be accepted + #expect(throws: Never.self) { + try validateTaskMode(isTaskRequest: false, for: tool) + } + } +} + +// MARK: - canUseToolWithTaskMode Tests + +@Suite("canUseToolWithTaskMode Tests") +struct CanUseToolWithTaskModeTests { + + @Test("REQUIRED mode with client task support returns true") + func testRequiredWithTaskSupport() { + // Python: test_can_use_tool_required_with_task_support + let result = canUseToolWithTaskMode(clientSupportsTask: true, taskSupport: .required) + #expect(result == true) + } + + @Test("REQUIRED mode without client task support returns false") + func testRequiredWithoutTaskSupport() { + // Python: test_can_use_tool_required_without_task_support + let result = canUseToolWithTaskMode(clientSupportsTask: false, taskSupport: .required) + #expect(result == false) + } + + @Test("OPTIONAL mode without client task support returns true") + func testOptionalWithoutTaskSupport() { + // Python: test_can_use_tool_optional_without_task_support + let result = canUseToolWithTaskMode(clientSupportsTask: false, taskSupport: .optional) + #expect(result == true) + } + + @Test("OPTIONAL mode with client task support returns true") + func testOptionalWithTaskSupport() { + let result = canUseToolWithTaskMode(clientSupportsTask: true, taskSupport: .optional) + #expect(result == true) + } + + @Test("FORBIDDEN mode without client task support returns true") + func testForbiddenWithoutTaskSupport() { + // Python: test_can_use_tool_forbidden_without_task_support + let result = canUseToolWithTaskMode(clientSupportsTask: false, taskSupport: .forbidden) + #expect(result == true) + } + + @Test("FORBIDDEN mode with client task support returns true") + func testForbiddenWithTaskSupport() { + let result = canUseToolWithTaskMode(clientSupportsTask: true, taskSupport: .forbidden) + #expect(result == true) + } + + @Test("nil mode (treated as FORBIDDEN) without client task support returns true") + func testNilModeWithoutTaskSupport() { + // Python: test_can_use_tool_none_without_task_support + let result = canUseToolWithTaskMode(clientSupportsTask: false, taskSupport: nil) + #expect(result == true) + } + + @Test("nil mode (treated as FORBIDDEN) with client task support returns true") + func testNilModeWithTaskSupport() { + let result = canUseToolWithTaskMode(clientSupportsTask: true, taskSupport: nil) + #expect(result == true) + } +} + +// MARK: - TaskMetadata and isTask Pattern Tests + +@Suite("Task Request Detection Tests") +struct TaskRequestDetectionTests { + + /// These tests verify the pattern for detecting if a request is task-augmented, + /// matching Python's `Experimental.is_task` property. + + @Test("Request with TaskMetadata is a task request") + func testRequestWithTaskMetadataIsTask() { + // Python: test_is_task_true_when_metadata_present + let metadata: TaskMetadata? = TaskMetadata(ttl: 60000) + let isTask = metadata != nil + #expect(isTask == true) + } + + @Test("Request without TaskMetadata is not a task request") + func testRequestWithoutTaskMetadataIsNotTask() { + // Python: test_is_task_false_when_no_metadata + let metadata: TaskMetadata? = nil + let isTask = metadata != nil + #expect(isTask == false) + } + + @Test("CallTool.Parameters with task metadata indicates task request") + func testCallToolWithTaskMetadataIsTask() { + let params = CallTool.Parameters( + name: "test_tool", + arguments: [:], + task: TaskMetadata(ttl: 60000) + ) + + let isTask = params.task != nil + #expect(isTask == true) + } + + @Test("CallTool.Parameters without task metadata indicates non-task request") + func testCallToolWithoutTaskMetadataIsNonTask() { + let params = CallTool.Parameters( + name: "test_tool", + arguments: [:], + task: nil + ) + + let isTask = params.task != nil + #expect(isTask == false) + } +} + +// MARK: - Client Task Capability Detection Tests + +@Suite("Client Task Capability Tests") +struct ClientTaskCapabilityTests { + + /// These tests verify the pattern for detecting if a client supports tasks, + /// matching Python's `Experimental.client_supports_tasks` property. + + @Test("Client with tasks capability supports tasks") + func testClientWithTasksCapability() { + // Python: test_client_supports_tasks_true + let capabilities = Client.Capabilities(tasks: .init()) + let supportsTask = capabilities.tasks != nil + #expect(supportsTask == true) + } + + @Test("Client without tasks capability does not support tasks") + func testClientWithoutTasksCapability() { + // Python: test_client_supports_tasks_false_no_tasks + let capabilities = Client.Capabilities() + let supportsTask = capabilities.tasks != nil + #expect(supportsTask == false) + } + + @Test("Nil capabilities means no task support") + func testNilCapabilities() { + // Python: test_client_supports_tasks_false_no_capabilities + let capabilities: Client.Capabilities? = nil + let supportsTask = capabilities?.tasks != nil + #expect(supportsTask == false) + } + + @Test("Server can check client task support") + func testServerCanCheckClientTaskSupport() { + // With tasks capability + let capsWithTasks = Client.Capabilities(tasks: .init()) + #expect(capsWithTasks.tasks != nil) + + // Without tasks capability + let capsWithoutTasks = Client.Capabilities(sampling: .init()) + #expect(capsWithoutTasks.tasks == nil) + } +} + +// MARK: - Integration Tests for Task Mode Validation + +@Suite("Task Mode Validation Integration Tests") +struct TaskModeValidationIntegrationTests { + + @Test("Tool invocation validation flow for required task tool") + func testToolInvocationValidationFlowRequired() throws { + // Simulate a tool that requires task-augmented invocation + let tool = Tool( + name: "long_running_analysis", + description: "A long-running analysis that requires task mode", + inputSchema: ["type": "object"], + execution: Tool.Execution(taskSupport: .required) + ) + + // Client without task support cannot use this tool + let clientCapsNoTasks = Client.Capabilities() + let canUse1 = canUseToolWithTaskMode( + clientSupportsTask: clientCapsNoTasks.tasks != nil, + taskSupport: tool.execution?.taskSupport + ) + #expect(canUse1 == false) + + // Client with task support can use this tool + let clientCapsWithTasks = Client.Capabilities(tasks: .init()) + let canUse2 = canUseToolWithTaskMode( + clientSupportsTask: clientCapsWithTasks.tasks != nil, + taskSupport: tool.execution?.taskSupport + ) + #expect(canUse2 == true) + + // Non-task request should be rejected + #expect(throws: MCPError.self) { + try validateTaskMode(isTaskRequest: false, for: tool) + } + + // Task request should be accepted + #expect(throws: Never.self) { + try validateTaskMode(isTaskRequest: true, for: tool) + } + } + + @Test("Tool invocation validation flow for forbidden task tool") + func testToolInvocationValidationFlowForbidden() throws { + // Simulate a tool that forbids task-augmented invocation + let tool = Tool( + name: "quick_lookup", + description: "A quick lookup that cannot be a task", + inputSchema: ["type": "object"], + execution: Tool.Execution(taskSupport: .forbidden) + ) + + // Any client can use this tool + let canUse = canUseToolWithTaskMode( + clientSupportsTask: false, + taskSupport: tool.execution?.taskSupport + ) + #expect(canUse == true) + + // Task request should be rejected + #expect(throws: MCPError.self) { + try validateTaskMode(isTaskRequest: true, for: tool) + } + + // Non-task request should be accepted + #expect(throws: Never.self) { + try validateTaskMode(isTaskRequest: false, for: tool) + } + } + + @Test("Tool invocation validation flow for optional task tool") + func testToolInvocationValidationFlowOptional() throws { + // Simulate a tool that optionally supports task-augmented invocation + let tool = Tool( + name: "flexible_processor", + description: "Can run as task or not", + inputSchema: ["type": "object"], + execution: Tool.Execution(taskSupport: .optional) + ) + + // Any client can use this tool + let canUseWithTask = canUseToolWithTaskMode( + clientSupportsTask: true, + taskSupport: tool.execution?.taskSupport + ) + #expect(canUseWithTask == true) + + let canUseWithoutTask = canUseToolWithTaskMode( + clientSupportsTask: false, + taskSupport: tool.execution?.taskSupport + ) + #expect(canUseWithoutTask == true) + + // Both task and non-task requests should be accepted + #expect(throws: Never.self) { + try validateTaskMode(isTaskRequest: true, for: tool) + } + #expect(throws: Never.self) { + try validateTaskMode(isTaskRequest: false, for: tool) + } + } + + @Test("Error codes match MCP spec") + func testErrorCodesMatchSpec() throws { + // Per MCP spec: METHOD_NOT_FOUND (-32601) for task mode violations + do { + try validateTaskMode(isTaskRequest: false, taskSupport: .required) + Issue.record("Expected error") + } catch let error as MCPError { + // MCPError.methodNotFound should be used + switch error { + case .methodNotFound: + // Correct error type + break + default: + Issue.record("Expected methodNotFound error, got \(error)") + } + } + + do { + try validateTaskMode(isTaskRequest: true, taskSupport: .forbidden) + Issue.record("Expected error") + } catch let error as MCPError { + switch error { + case .methodNotFound: + // Correct error type + break + default: + Issue.record("Expected methodNotFound error, got \(error)") + } + } + } +} diff --git a/Tests/MCPTests/TaskTests.swift b/Tests/MCPTests/TaskTests.swift index d46de8ec..ff7dc181 100644 --- a/Tests/MCPTests/TaskTests.swift +++ b/Tests/MCPTests/TaskTests.swift @@ -831,9 +831,9 @@ struct InMemoryTaskStoreTests { _ = try await store.createTask(metadata: TaskMetadata(), taskId: "task-2") _ = try await store.createTask(metadata: TaskMetadata(), taskId: "task-3") - let (tasks, _) = await store.listTasks(cursor: nil) + let result = await store.listTasks(cursor: nil) - #expect(tasks.count == 3) + #expect(result.tasks.count == 3) } @Test("listTasks pagination works correctly") @@ -846,19 +846,19 @@ struct InMemoryTaskStoreTests { } // First page - let (page1, cursor1) = await store.listTasks(cursor: nil) - #expect(page1.count == 2) - #expect(cursor1 != nil) + let page1Result = await store.listTasks(cursor: nil) + #expect(page1Result.tasks.count == 2) + #expect(page1Result.nextCursor != nil) // Second page - let (page2, cursor2) = await store.listTasks(cursor: cursor1) - #expect(page2.count == 2) - #expect(cursor2 != nil) + let page2Result = await store.listTasks(cursor: page1Result.nextCursor) + #expect(page2Result.tasks.count == 2) + #expect(page2Result.nextCursor != nil) // Third page - let (page3, cursor3) = await store.listTasks(cursor: cursor2) - #expect(page3.count == 1) - #expect(cursor3 == nil) + let page3Result = await store.listTasks(cursor: page2Result.nextCursor) + #expect(page3Result.tasks.count == 1) + #expect(page3Result.nextCursor == nil) } @Test("deleteTask removes task") diff --git a/Tests/MCPTests/ToolTests.swift b/Tests/MCPTests/ToolTests.swift index 916331b5..11cccb3d 100644 --- a/Tests/MCPTests/ToolTests.swift +++ b/Tests/MCPTests/ToolTests.swift @@ -414,6 +414,7 @@ struct ToolTests { requestId: .number(1), _meta: nil, authInfo: nil, + requestInfo: nil, closeSSEStream: nil, closeStandaloneSSEStream: nil, shouldSendLogMessage: { _ in true },