diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 889a1df..6e1fab0 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -2,27 +2,130 @@ name: CI on: push: - branches: [ "main" ] + branches: ["main"] pull_request: - branches: [ "main" ] + branches: ["main"] + +permissions: + contents: read + pull-requests: write jobs: test: - runs-on: macos-latest - + timeout-minutes: 10 strategy: matrix: + os: [macos-latest, ubuntu-latest] swift-version: - - ^6 + - 6.1.0 - name: Build and Test (Swift ${{ matrix.swift-version }}) + runs-on: ${{ matrix.os }} + name: Test (${{ matrix.os }}, Swift ${{ matrix.swift-version }}) steps: - - uses: actions/checkout@v4 - - uses: swift-actions/setup-swift@v2 - with: - swift-version: ${{ matrix.swift-version }} - - name: Build - run: swift build -v - - name: Run tests - run: swift test -v + - uses: actions/checkout@v4 + + - name: Prepare Linux apt dependencies + if: matrix.os == 'ubuntu-latest' + run: | + sudo -n apt-get update + sudo -n apt-get install -y --no-install-recommends libcurl4-openssl-dev + + - name: Setup Swift on Linux + if: matrix.os == 'ubuntu-latest' + uses: vapor/swiftly-action@v0.2 + with: + toolchain: ${{ matrix.swift-version }} + - name: Setup Swift on macOS + if: matrix.os == 'macos-latest' + uses: swift-actions/setup-swift@v2 + with: + swift-version: ${{ matrix.swift-version }} + + - name: Build + run: swift build -v + + - name: Run tests + run: swift test -v + + conformance: + timeout-minutes: 10 + runs-on: macos-latest + name: MCP Conformance Tests + + steps: + - uses: actions/checkout@v4 + + - name: Setup Swift + uses: swift-actions/setup-swift@v2 + with: + swift-version: 6.1.0 + + - name: Build Swift executables + run: | + swift build --product mcp-everything-client + swift build --product mcp-everything-server + + - name: Run client conformance tests + uses: modelcontextprotocol/conformance@v0.1.15 + with: + mode: client + command: '.build/debug/mcp-everything-client' + suite: 'core' + expected-failures: './conformance-baseline.yml' + node-version: '22' + + - name: Start server for testing + run: | + .build/debug/mcp-everything-server & + echo "SERVER_PID=$!" >> $GITHUB_ENV + sleep 3 + + - name: Run server conformance tests + uses: modelcontextprotocol/conformance@v0.1.15 + with: + mode: server + url: 'http://localhost:3001/mcp' + suite: 'core' + expected-failures: './conformance-baseline.yml' + node-version: '22' + + - name: Cleanup server + if: always() + run: | + if [ -n "$SERVER_PID" ]; then + kill $SERVER_PID 2>/dev/null || true + fi + + documentation: + name: Documentation + runs-on: macos-latest + timeout-minutes: 10 + steps: + - uses: actions/checkout@v4 + - uses: swift-actions/setup-swift@v2 + with: + swift-version: "6.1.0" + - name: Build Documentation + run: swift package generate-documentation --target MCP --warnings-as-errors + + static-linux-sdk-build: + name: Linux Static SDK Build (${{ matrix.swift-version }} - ${{ matrix.os }}) + strategy: + matrix: + os: [ubuntu-latest] + swift-version: + - 6.1.2 + runs-on: ${{ matrix.os }} + container: swift:${{ matrix.swift-version }}-noble + timeout-minutes: 30 + steps: + - name: Check out code + uses: actions/checkout@v4 + - name: Install static SDK and build + env: + BUILD_FLAGS: "" + run: | + which curl || (apt -q update && apt -yq install curl) + curl -s --retry 3 https://raw.githubusercontent.com/swiftlang/github-workflows/refs/heads/main/.github/workflows/scripts/install-and-build-with-sdk.sh | \ + bash -s -- --static --flags="$BUILD_FLAGS" ${{ matrix.swift-version }} diff --git a/.gitignore b/.gitignore index 0023a53..2d0abd0 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,4 @@ DerivedData/ .swiftpm/configuration/registries.json .swiftpm/xcode/package.xcworkspace/contents.xcworkspacedata .netrc +.index-build/ \ No newline at end of file diff --git a/.spi.yml b/.spi.yml new file mode 100644 index 0000000..f846581 --- /dev/null +++ b/.spi.yml @@ -0,0 +1,4 @@ +version: 1 +builder: + configs: + - documentation_targets: [MCP] \ No newline at end of file diff --git a/.vscode/extensions.json b/.vscode/extensions.json index baab102..6d8b1f0 100644 --- a/.vscode/extensions.json +++ b/.vscode/extensions.json @@ -1,3 +1,3 @@ { - "recommendations": ["swiftlang.swift-vscode"] + "recommendations": ["swiftlang.swift-vscode", "esbenp.prettier-vscode"] } diff --git a/.vscode/settings.json b/.vscode/settings.json index 757f266..fddc342 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -4,5 +4,12 @@ "editor.codeActionsOnSave": { "source.fixAll": "explicit", "source.organizeImports": "explicit" + }, + "[github-actions-workflow]": { + "editor.defaultFormatter": "esbenp.prettier-vscode" + }, + "[swift]": { + "editor.insertSpaces": true, + "editor.tabSize": 4 } } diff --git a/LICENSE b/LICENSE index e37fe81..4a93985 100644 --- a/LICENSE +++ b/LICENSE @@ -1,3 +1,11 @@ +The MCP project is undergoing a licensing transition from the MIT License to the Apache License, Version 2.0 ("Apache-2.0"). All new code and specification contributions to the project are licensed under Apache-2.0. Documentation contributions (excluding specifications) are licensed under CC-BY-4.0. + +Contributions for which relicensing consent has been obtained are licensed under Apache-2.0. Contributions made by authors who originally licensed their work under the MIT License and who have not yet granted explicit permission to relicense remain licensed under the MIT License. + +No rights beyond those granted by the applicable original license are conveyed for such contributions. + +--- + Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ @@ -48,9 +56,9 @@ "Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" + submitted to the Licensor for inclusion in the Work by the copyright + owner or by an individual or Legal Entity authorized to submit on behalf + of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, @@ -175,27 +183,34 @@ END OF TERMS AND CONDITIONS - APPENDIX: How to apply the Apache License to your work. +--- - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. +MIT License - Copyright 2025 Loopwork Limited +Copyright (c) 2024-2025 Model Context Protocol a Series of LF Projects, LLC. - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: - http://www.apache.org/licenses/LICENSE-2.0 +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. + +--- + +Creative Commons Attribution 4.0 International (CC-BY-4.0) + +Documentation in this project (excluding specifications) is licensed under +CC-BY-4.0. See https://creativecommons.org/licenses/by/4.0/legalcode for +the full license text. diff --git a/Package.resolved b/Package.resolved index eaee8ae..d21ebac 100644 --- a/Package.resolved +++ b/Package.resolved @@ -1,6 +1,51 @@ { - "originHash" : "cce36cb33302c2f1c28458e19b8439f736fc28106e4c6ea95d7992c74594c242", + "originHash" : "67fa9b6a0d1f334fd836cabb7af8c36110a6b88c5b49fa5a42a50125b18a056b", "pins" : [ + { + "identity" : "eventsource", + "kind" : "remoteSourceControl", + "location" : "https://github.com/mattt/eventsource.git", + "state" : { + "revision" : "e83f076811f32757305b8bf69ac92d05626ffdd7", + "version" : "1.1.0" + } + }, + { + "identity" : "swift-atomics", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-atomics.git", + "state" : { + "revision" : "b601256eab081c0f92f059e12818ac1d4f178ff7", + "version" : "1.3.0" + } + }, + { + "identity" : "swift-collections", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-collections.git", + "state" : { + "revision" : "7b847a3b7008b2dc2f47ca3110d8c782fb2e5c7e", + "version" : "1.3.0" + } + }, + { + "identity" : "swift-docc-plugin", + "kind" : "remoteSourceControl", + "location" : "https://github.com/swiftlang/swift-docc-plugin", + "state" : { + "branch" : "main", + "revision" : "e977f65879f82b375a044c8837597f690c067da6" + } + }, + { + "identity" : "swift-docc-symbolkit", + "kind" : "remoteSourceControl", + "location" : "https://github.com/swiftlang/swift-docc-symbolkit", + "state" : { + "revision" : "b45d1f2ed151d057b54504d653e0da5552844e34", + "version" : "1.0.0" + } + }, { "identity" : "swift-log", "kind" : "remoteSourceControl", @@ -10,6 +55,15 @@ "version" : "1.6.2" } }, + { + "identity" : "swift-nio", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-nio.git", + "state" : { + "revision" : "5e72fc102906ebe75a3487595a653e6f43725552", + "version" : "2.94.0" + } + }, { "identity" : "swift-system", "kind" : "remoteSourceControl", diff --git a/Package.swift b/Package.swift index daf946c..273715a 100644 --- a/Package.swift +++ b/Package.swift @@ -1,4 +1,4 @@ -// swift-tools-version:6.0 +// swift-tools-version:6.1 // The swift-tools-version declares the minimum version of Swift required to build this package. import PackageDescription @@ -6,34 +6,75 @@ import PackageDescription let package = Package( name: "mcp-swift-sdk", platforms: [ - .macOS(.v14), - .iOS(.v17), + .macOS("13.0"), + .macCatalyst("16.0"), + .iOS("16.0"), + .watchOS("9.0"), + .tvOS("16.0"), + .visionOS("1.0"), ], products: [ // Products define the executables and libraries a package produces, making them visible to other packages. .library( name: "MCP", - targets: ["MCP"]) + targets: ["MCP"]), + .executable( + name: "mcp-everything-server", + targets: ["MCPConformanceServer"]), + .executable( + name: "mcp-everything-client", + targets: ["MCPConformanceClient"]) ], dependencies: [ + .package(url: "https://github.com/swiftlang/swift-docc-plugin", branch: "main"), .package(url: "https://github.com/apple/swift-system.git", from: "1.0.0"), .package(url: "https://github.com/apple/swift-log.git", from: "1.5.0"), + .package(url: "https://github.com/mattt/eventsource.git", from: "1.1.0"), + .package(url: "https://github.com/apple/swift-nio.git", from: "2.65.0"), ], targets: [ - // Targets are the basic building blocks of a package, defining a module or a test suite. - // Targets can depend on other targets in this package and products from dependencies. .target( name: "MCP", dependencies: [ .product(name: "SystemPackage", package: "swift-system"), .product(name: "Logging", package: "swift-log"), - ]), + .product( + name: "EventSource", package: "eventsource", + condition: .when(platforms: [.macOS, .iOS, .tvOS, .visionOS, .watchOS, .macCatalyst])), + ], + swiftSettings: [ + .enableUpcomingFeature("StrictConcurrency") + ] + ), .testTarget( name: "MCPTests", dependencies: [ "MCP", .product(name: "SystemPackage", package: "swift-system"), .product(name: "Logging", package: "swift-log"), - ]), + .product( + name: "EventSource", package: "eventsource", + condition: .when(platforms: [.macOS, .iOS, .tvOS, .visionOS, .watchOS, .macCatalyst])), + ] + ), + .executableTarget( + name: "MCPConformanceServer", + dependencies: [ + "MCP", + .product(name: "Logging", package: "swift-log"), + .product(name: "NIOCore", package: "swift-nio"), + .product(name: "NIOPosix", package: "swift-nio"), + .product(name: "NIOHTTP1", package: "swift-nio"), + ], + path: "Sources/MCPConformance/Server" + ), + .executableTarget( + name: "MCPConformanceClient", + dependencies: [ + "MCP", + .product(name: "Logging", package: "swift-log"), + ], + path: "Sources/MCPConformance/Client" + ) ] ) diff --git a/Package@swift-6.0.swift b/Package@swift-6.0.swift new file mode 100644 index 0000000..5fdf2e4 --- /dev/null +++ b/Package@swift-6.0.swift @@ -0,0 +1,49 @@ +// swift-tools-version:6.0 +// The swift-tools-version declares the minimum version of Swift required to build this package. + +import PackageDescription + +// Base dependencies needed on all platforms +var dependencies: [Package.Dependency] = [ + .package(url: "https://github.com/apple/swift-system.git", from: "1.0.0"), + .package(url: "https://github.com/apple/swift-log.git", from: "1.5.0"), + .package(url: "https://github.com/mattt/eventsource.git", from: "1.1.0"), +] + +// Target dependencies needed on all platforms +var targetDependencies: [Target.Dependency] = [ + .product(name: "SystemPackage", package: "swift-system"), + .product(name: "Logging", package: "swift-log"), + .product( + name: "EventSource", package: "eventsource", + condition: .when(platforms: [.macOS, .iOS, .tvOS, .visionOS, .watchOS])), +] + +let package = Package( + name: "mcp-swift-sdk", + platforms: [ + .macOS("13.0"), + .macCatalyst("16.0"), + .iOS("16.0"), + .watchOS("9.0"), + .tvOS("16.0"), + .visionOS("1.0"), + ], + products: [ + // Products define the executables and libraries a package produces, making them visible to other packages. + .library( + name: "MCP", + targets: ["MCP"]) + ], + dependencies: dependencies, + targets: [ + // Targets are the basic building blocks of a package, defining a module or a test suite. + // Targets can depend on other targets in this package and products from dependencies. + .target( + name: "MCP", + dependencies: targetDependencies), + .testTarget( + name: "MCPTests", + dependencies: ["MCP"] + targetDependencies), + ] +) diff --git a/README.md b/README.md index 11a2c78..1bf1670 100644 --- a/README.md +++ b/README.md @@ -1,15 +1,70 @@ # MCP Swift SDK -Swift implementation of the [Model Context Protocol][mcp] (MCP). +Official Swift SDK for the [Model Context Protocol][mcp] (MCP). -> [!TIP] -> To see it in action, check out the [iMCP app](https://iMCP.app/?ref=mcp-swift-sdk). +## Overview + +The Model Context Protocol (MCP) defines a standardized way +for applications to communicate with AI and ML models. +This Swift SDK implements both client and server components +according to the [2025-11-25][mcp-spec-2025-11-25] (latest) version +of the MCP specification. + +## Table of contents + +- [Requirements](#requirements) +- [Installation](#installation) +- [Client Usage](#client-usage) + - [Basic Client Setup](#basic-client-setup) + - [Transport Options for Clients](#transport-options-for-clients) + - [Tools](#tools) + - [Resources](#resources) + - [Prompts](#prompts) + - [Completions](#completions) + - [Sampling](#sampling) + - [Elicitation](#elicitation) + - [Roots](#roots) + - [Logging](#logging) + - [Error Handling](#error-handling) + - [Cancellation](#cancellation) + - [Progress Tracking](#progress-tracking) + - [Advanced Client Features](#advanced-client-features) +- [Server Usage](#server-usage) + - [Basic Server Setup](#basic-server-setup) + - [Tools](#tools-1) + - [Resources](#resources-1) + - [Prompts](#prompts-1) + - [Completions](#completions-1) + - [Sampling](#sampling-1) + - [Elicitations](#elicitations) + - [Roots](#roots-1) + - [Logging](#logging-1) + - [Progress Tracking](#progress-tracking-1) + - [Initialize Hook](#initialize-hook) + - [HTTP Request Context in Handlers](#http-request-context-in-handlers) + - [Graceful Shutdown](#graceful-shutdown) +- [Transports](#transports) +- [Authentication](#authentication) + - [Client: Client Credentials Flow](#client-client-credentials-flow) + - [Client: Authorization Code Flow](#client-authorization-code-flow) + - [Client: Custom Token Provider](#client-custom-token-provider) + - [Client: Custom Token Storage](#client-custom-token-storage) + - [Client: private\_key\_jwt Authentication](#client-private_key_jwt-authentication) + - [Client: Endpoint Overrides](#client-endpoint-overrides) + - [Server: Serving Protected Resource Metadata](#server-serving-protected-resource-metadata) + - [Server: Validating Bearer Tokens](#server-validating-bearer-tokens) +- [Platform Availability](#platform-availability) +- [Debugging and Logging](#debugging-and-logging) +- [Additional Resources](#additional-resources) +- [Changelog](#changelog) +- [License](#license) ## Requirements -- Swift 6.0+ / Xcode 16+ -- macOS 14.0+ (Sonoma) -- iOS 17.0+ +- Swift 6.0+ (Xcode 16+) + +See the [Platform Availability](#platform-availability) section below +for platform-specific requirements. ## Installation @@ -19,11 +74,24 @@ Add the following to your `Package.swift` file: ```swift dependencies: [ - .package(url: "https://github.com/loopwork-ai/mcp-swift-sdk.git", from: "0.4.2") + .package(url: "https://github.com/modelcontextprotocol/swift-sdk.git", from: "0.11.0") ] ``` -## Usage +Then add the dependency to your target: + +```swift +.target( + name: "YourTarget", + dependencies: [ + .product(name: "MCP", package: "swift-sdk") + ] +) +``` + +## Client Usage + +The client component allows your application to connect to MCP servers. ### Basic Client Setup @@ -35,27 +103,584 @@ let client = Client(name: "MyApp", version: "1.0.0") // Create a transport and connect let transport = StdioTransport() -try await client.connect(transport: transport) +let result = try await client.connect(transport: transport) -// Initialize the connection -let result = try await client.initialize() +// Check server capabilities +if result.capabilities.tools != nil { + // Server supports tools (implicitly including tool calling if the 'tools' capability object is present) +} ``` +> [!NOTE] +> The `Client.connect(transport:)` method returns the initialization result. +> This return value is discardable, +> so you can ignore it if you don't need to check server capabilities. + +### Transport Options for Clients + +#### Stdio Transport + +For local subprocess communication: + +```swift +// Create a stdio transport (simplest option) +let transport = StdioTransport() +try await client.connect(transport: transport) +``` + +#### HTTP Transport + +For remote server communication: + +```swift +// Create a streaming HTTP transport +let transport = HTTPClientTransport( + endpoint: URL(string: "http://localhost:8080")!, + streaming: true // Enable Server-Sent Events for real-time updates +) +try await client.connect(transport: transport) +``` + +### Tools + +Tools represent functions that can be called by the client: + +```swift +// List available tools +let (tools, cursor) = try await client.listTools() +print("Available tools: \(tools.map { $0.name }.joined(separator: ", "))") + +// Call a tool with arguments and get the result +let (content, isError) = try await client.callTool( + name: "image-generator", + arguments: [ + "prompt": "A serene mountain landscape at sunset", + "style": "photorealistic", + "width": 1024, + "height": 768 + ] +) + +// Handle tool content +for item in content { + switch item { + case .text(let text): + print("Generated text: \(text)") + case .image(let data, let mimeType, let metadata): + if let width = metadata?["width"] as? Int, + let height = metadata?["height"] as? Int { + print("Generated \(width)x\(height) image of type \(mimeType)") + // Save or display the image data + } + case .audio(let data, let mimeType): + print("Received audio data of type \(mimeType)") + case .resource(let resource, _, _): + print("Received embedded resource: \(resource)") + case .resourceLink(let uri, let name, _, _, let mimeType, _): + print("Resource link: \(name) at \(uri), type: \(mimeType ?? "unknown")") + } +} +``` + +### Resources + +Resources represent data that can be accessed and potentially subscribed to: + +```swift +// List available resources +let (resources, nextCursor) = try await client.listResources() +print("Available resources: \(resources.map { $0.uri }.joined(separator: ", "))") + +// Read a resource +let contents = try await client.readResource(uri: "resource://example") +print("Resource content: \(contents)") + +// Subscribe to resource updates if supported +if result.capabilities.resources?.subscribe == true { + try await client.subscribeToResource(uri: "resource://example") + + // Register notification handler + await client.onNotification(ResourceUpdatedNotification.self) { message in + let uri = message.params.uri + print("Resource \(uri) updated with new content") + + // Fetch the updated resource content + let updatedContents = try await client.readResource(uri: uri) + print("Updated resource content received") + } +} +``` + +### Prompts + +Prompts represent templated conversation starters: + +```swift +// List available prompts +let (prompts, nextCursor) = try await client.listPrompts() +print("Available prompts: \(prompts.map { $0.name }.joined(separator: ", "))") + +// Get a prompt with arguments +let (description, messages) = try await client.getPrompt( + name: "customer-service", + arguments: [ + "customerName": "Alice", + "orderNumber": "ORD-12345", + "issue": "delivery delay" + ] +) + +// Use the prompt messages in your application +print("Prompt description: \(description)") +for message in messages { + if case .text(text: let text) = message.content { + print("\(message.role): \(text)") + } +} +``` + +### Completions + +Completions allow servers to provide autocompletion suggestions for prompt and resource template arguments as users type: + +```swift +// Request completions for a prompt argument +let completion = try await client.complete( + promptName: "code_review", + argumentName: "language", + argumentValue: "py" +) + +// Display suggestions to the user +for value in completion.values { + print("Suggestion: \(value)") +} + +if completion.hasMore == true { + print("More suggestions available (total: \(completion.total ?? 0))") +} +``` + +You can also provide context with already-resolved arguments: + +```swift +// First, user selects a language +let languageCompletion = try await client.complete( + promptName: "code_review", + argumentName: "language", + argumentValue: "py" +) +// User selects "python" + +// Then get framework suggestions based on the selected language +let frameworkCompletion = try await client.complete( + promptName: "code_review", + argumentName: "framework", + argumentValue: "fla", + context: ["language": .string("python")] +) +// Returns: ["flask"] +``` + +Completions work for resource templates as well: + +```swift +// Get path completions for a resource URI template +let pathCompletion = try await client.complete( + resourceURI: "file:///{path}", + argumentName: "path", + argumentValue: "/usr/" +) +// Returns: ["/usr/bin", "/usr/lib", "/usr/local"] +``` + +### Sampling + +Sampling allows servers to request LLM completions through the client, +enabling agentic behaviors while maintaining human-in-the-loop control. +Clients register a handler to process incoming sampling requests from servers. + +> [!TIP] +> Sampling requests flow from **server to client**, +> not client to server. +> This enables servers to request AI assistance +> while clients maintain control over model access and user approval. + +```swift +// Register a sampling handler in the client +await client.withSamplingHandler { parameters in + // Review the sampling request (human-in-the-loop step 1) + print("Server requests completion for: \(parameters.messages)") + + // Optionally modify the request based on user input + var messages = parameters.messages + if let systemPrompt = parameters.systemPrompt { + print("System prompt: \(systemPrompt)") + } + + // Sample from your LLM (this is where you'd call your AI service) + let completion = try await callYourLLMService( + messages: messages, + maxTokens: parameters.maxTokens, + temperature: parameters.temperature + ) + + // Review the completion (human-in-the-loop step 2) + print("LLM generated: \(completion)") + // User can approve, modify, or reject the completion here + + // Return the result to the server + return CreateSamplingMessage.Result( + model: "your-model-name", + stopReason: .endTurn, + role: .assistant, + content: .text(completion) + ) +} +``` + +### Elicitation + +Elicitation allows servers to request structured information directly from users through the client. +This is useful when servers need user input that wasn't provided in the original request, +such as credentials, configuration choices, or approval for sensitive operations. + +> [!TIP] +> Elicitation requests flow from **server to client**, +> similar to sampling. +> Clients must register a handler to respond to elicitation requests from servers. + +#### Client-Side: Handling Elicitation Requests + +Register an elicitation handler to respond to server requests: + +```swift +// Register an elicitation handler in the client +await client.withElicitationHandler { parameters in + switch parameters { + case .form(let form): + // Display the request to the user + print("Server requests: \(form.message)") + + // If a schema was provided, inspect it + if let schema = form.requestedSchema { + print("Required fields: \(schema.required ?? [])") + print("Schema: \(schema.properties)") + } + + // Present UI to collect user input + let userResponse = presentElicitationUI(form) + + // Return the user's response + if userResponse.accepted { + return CreateElicitation.Result( + action: .accept, + content: userResponse.data + ) + } else if userResponse.canceled { + return CreateElicitation.Result(action: .cancel) + } else { + return CreateElicitation.Result(action: .decline) + } + + case .url(let url): + // Direct the user to an external URL (e.g., for OAuth) + openURL(url.url) + return CreateElicitation.Result(action: .accept) + } +} +``` + +Common use cases for elicitation: + +- **Authentication**: Request credentials when needed rather than upfront +- **Confirmation**: Ask for user approval before sensitive operations +- **Configuration**: Collect preferences or settings during operation +- **Missing information**: Request additional details not provided initially + +### Roots + +Roots define the filesystem boundaries that a client exposes to servers. Servers discover roots by sending a `roots/list` request to the client; clients notify servers when the list changes. + +> [!TIP] +> To use roots, declare the `roots` capability when creating the client. + +```swift +let client = Client( + name: "MyApp", + version: "1.0.0", + capabilities: .init( + roots: .init(listChanged: true) + ) +) + +// Register a handler for roots/list requests from servers +await client.withRootsHandler { + return [ + Root(uri: "file:///Users/user/projects", name: "Projects"), + Root(uri: "file:///Users/user/documents", name: "Documents") + ] +} + +// Notify connected servers whenever roots change +try await client.notifyRootsChanged() +``` + +### Logging + +Clients can control server logging levels and receive structured log messages: + +```swift +// Set the minimum logging level +try await client.setLoggingLevel(.warning) + +// Register a handler for log messages from the server +await client.onNotification(LogMessageNotification.self) { message in + let level = message.params.level // LogLevel (debug, info, warning, etc.) + let logger = message.params.logger // Optional logger name + let data = message.params.data // Arbitrary JSON data + + // Display log message based on level + switch level { + case .error, .critical, .alert, .emergency: + print("❌ [\(logger ?? "server")] \(data)") + case .warning: + print("⚠️ [\(logger ?? "server")] \(data)") + default: + print("ℹ️ [\(logger ?? "server")] \(data)") + } +} +``` + +Log levels follow the standard syslog severity levels (RFC 5424): + +- **debug**: Detailed debugging information +- **info**: General informational messages +- **notice**: Normal but significant events +- **warning**: Warning conditions +- **error**: Error conditions +- **critical**: Critical conditions +- **alert**: Action must be taken immediately +- **emergency**: System is unusable + +### Error Handling + +Handle common client errors: + +```swift +do { + try await client.connect(transport: transport) + // Success +} catch let error as MCPError { + print("MCP Error: \(error.localizedDescription)") +} catch { + print("Unexpected error: \(error)") +} +``` + +### Cancellation + +Either side can cancel an in-progress request and handle incoming cancellations gracefully: + +#### Option 1: Convenience Methods with RequestContext Overload + +For common operations like tool calls, use the overloaded method that returns `RequestContext`: + +```swift +// Call a tool and get a context for cancellation +let context = try client.callTool( + name: "long-running-analysis", + arguments: ["data": largeDataset] +) + +// You can cancel the request at any time +try await client.cancelRequest(context.requestID, reason: "User cancelled") + +// Await the result (will throw CancellationError if cancelled) +do { + let result = try await context.value + print("Result: \(result.content)") +} catch is CancellationError { + print("Request was cancelled") +} +``` + +#### Option 2: Direct send() for Maximum Flexibility + +For full control or custom requests, use `send()` directly: + +```swift +// Create any request type +let request = CallTool.request(.init( + name: "long-running-analysis", + arguments: ["data": largeDataset] +)) + +// Send and get a context for cancellation tracking +let context: RequestContext = try client.send(request) + +// Cancel when needed +try await client.cancelRequest(context.requestID, reason: "Timeout") + +// Get the result +let result = try await context.value +let content = result.content +let isError = result.isError +``` + +### Progress tracking + +Clients can attach a progress token to a request and receive incremental progress updates for long-running operations: + +```swift +// Call a tool with progress tracking +let progressToken = ProgressToken.unique() + +// Register a notification handler to receive progress updates +await client.onNotification(ProgressNotification.self) { message in + let params = message.params + // Filter by your progress token + if params.progressToken == progressToken { + print("Progress: \(params.progress)/\(params.total ?? 0)") + if let message = params.message { + print("Status: \(message)") + } + } +} + +// Make the request with the progress token +let (content, isError) = try await client.callTool( + name: "long-running-tool", + arguments: ["input": "value"], + meta: Metadata(progressToken: progressToken) +) +``` + +### Advanced Client Features + +#### Strict vs Non-Strict Configuration + +Configure client behavior for capability checking: + +```swift +// Strict configuration - fail fast if a capability is missing +let strictClient = Client( + name: "StrictClient", + version: "1.0.0", + configuration: .strict +) + +// With strict configuration, calling a method for an unsupported capability +// will throw an error immediately without sending a request +do { + // This will throw an error if resources.list capability is not available + let resources = try await strictClient.listResources() +} catch let error as MCPError { + print("Capability not available: \(error.localizedDescription)") +} + +// Default (non-strict) configuration - attempt the request anyway +let client = Client( + name: "FlexibleClient", + version: "1.0.0", + configuration: .default +) + +// With default configuration, the client will attempt the request +// even if the capability wasn't advertised by the server +do { + let resources = try await client.listResources() +} catch let error as MCPError { + // Still handle the error if the server rejects the request + print("Server rejected request: \(error.localizedDescription)") +} +``` + +#### Request Batching + +Improve performance by sending multiple requests in a single batch: + +```swift +// Array to hold tool call tasks +var toolTasks: [Task] = [] + +// Send a batch of requests +try await client.withBatch { batch in + // Add multiple tool calls to the batch + for i in 0..<10 { + toolTasks.append( + try await batch.addRequest( + CallTool.request(.init(name: "square", arguments: ["n": Value(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)") + } +} +``` + +You can also batch different types of requests: + +```swift +// Declare task variables +var pingTask: Task? +var promptTask: Task? + +// Send a batch with different request types +try await client.withBatch { batch in + pingTask = try await batch.addRequest(Ping.request()) + promptTask = try await batch.addRequest( + GetPrompt.request(.init(name: "greeting")) + ) +} + +// Process individual results +do { + if let pingTask = pingTask { + try await pingTask.value + print("Ping successful") + } + + if let promptTask = promptTask { + let promptResult = try await promptTask.value + print("Prompt: \(promptResult.description ?? "None")") + } +} catch { + print("Error processing batch results: \(error)") +} +``` + +> [!NOTE] +> `Server` automatically handles batch requests from MCP clients. + +## Server Usage + +The server component allows your application to host model capabilities and respond to client requests. + ### Basic Server Setup ```swift import MCP -// Initialize the server with capabilities +// Create a server with given capabilities let server = Server( - name: "MyServer", + name: "MyModelServer", version: "1.0.0", capabilities: .init( - prompts: .init(), - resources: .init( - subscribe: true - ), - tools: .init() + completions: .init(), + logging: .init(), + prompts: .init(listChanged: true), + resources: .init(subscribe: true, listChanged: true), + tools: .init(listChanged: true) ) ) @@ -63,81 +688,948 @@ let server = Server( let transport = StdioTransport() try await server.start(transport: transport) -// Register method handlers -server.withMethodHandler(ReadResource.self) { params in - // Handle resource read request - let uri = params.uri - let content = [Resource.Content.text("Example content")] - return .init(contents: content) -} - -// Register notification handlers -server.onNotification(ResourceUpdatedNotification.self) { message in - // Handle resource update notification -} - -// Stop the server when done -await server.stop() +// Now register handlers for the capabilities you've enabled ``` -### Working with Tools +### Tools + +Register tool handlers to respond to client tool calls: ```swift -// List available tools -let tools = try await client.listTools() +// Register a tool list handler +await server.withMethodHandler(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") + ]) + ]) + ), + Tool( + name: "calculator", + description: "Perform calculations", + inputSchema: .object([ + "properties": .object([ + "expression": .string("Mathematical expression to evaluate") + ]) + ]) + ) + ] + return .init(tools: tools) +} -// Call a tool -let (content, isError) = try await client.callTool( - name: "example-tool", - arguments: ["key": "value"] -) +// Register a tool call handler +await server.withMethodHandler(CallTool.self) { params in + switch params.name { + case "weather": + let location = params.arguments?["location"]?.stringValue ?? "Unknown" + let units = params.arguments?["units"]?.stringValue ?? "metric" + let weatherData = getWeatherData(location: location, units: units) // Your implementation + return .init( + content: [.text("Weather for \(location): \(weatherData.temperature)°, \(weatherData.conditions)")], + isError: false + ) -// Handle tool content -for item in content { - switch item { - case .text(let text): - print(text) - case .image(let data, let mimeType, let metadata): - // Handle image data + case "calculator": + if let expression = params.arguments?["expression"]?.stringValue { + let result = evaluateExpression(expression) // Your implementation + return .init(content: [.text("\(result)")], isError: false) + } else { + return .init(content: [.text("Missing expression parameter")], isError: true) + } + + default: + return .init(content: [.text("Unknown tool")], isError: true) } } ``` -### Working with Resources +### Resources + +Implement resource handlers for data access: ```swift -// List available resources -let (resources, nextCursor) = try await client.listResources() +// Register a resource list handler +await server.withMethodHandler(ListResources.self) { params in + let resources = [ + Resource( + name: "Knowledge Base Articles", + uri: "resource://knowledge-base/articles", + description: "Collection of support articles and documentation" + ), + Resource( + name: "System Status", + uri: "resource://system/status", + description: "Current system operational status" + ) + ] + return .init(resources: resources, nextCursor: nil) +} -// Read a resource -let contents = try await client.readResource(uri: "resource://example") +// Register a resource read handler +await server.withMethodHandler(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)]) -// Subscribe to resource updates -try await client.subscribeToResource(uri: "resource://example") + case "resource://system/status": + let status = getCurrentSystemStatus() // Your implementation + let statusJson = """ + { + "status": "\(status.overall)", + "components": { + "database": "\(status.database)", + "api": "\(status.api)", + "model": "\(status.model)" + }, + "lastUpdated": "\(status.timestamp)" + } + """ + return .init(contents: [Resource.Content.text(statusJson, uri: params.uri, mimeType: "application/json")]) -// Handle resource updates -await client.onNotification(ResourceUpdatedNotification.self) { message in - let uri = message.params.uri - let content = message.params.content - // Handle the update + default: + throw MCPError.invalidParams("Unknown resource URI: \(params.uri)") + } +} + +// Register a resource subscribe handler +await server.withMethodHandler(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. + // addSubscription(clientID: /* some_client_identifier */, uri: params.uri) + print("Client subscribed to \(params.uri). Server needs to implement logic to track this subscription.") + return .init() } ``` -### Working with Prompts +### Prompts + +Implement prompt handlers: ```swift -// List available prompts -let (prompts, nextCursor) = try await client.listPrompts() +// Register a prompt list handler +await server.withMethodHandler(ListPrompts.self) { params in + let prompts = [ + Prompt( + name: "interview", + description: "Job interview conversation starter", + arguments: [ + .init(name: "position", description: "Job position", required: true), + .init(name: "company", description: "Company name", required: true), + .init(name: "interviewee", description: "Candidate name") + ] + ), + Prompt( + name: "customer-support", + description: "Customer support conversation starter", + arguments: [ + .init(name: "issue", description: "Customer issue", required: true), + .init(name: "product", description: "Product name", required: true) + ] + ) + ] + return .init(prompts: prompts, nextCursor: nil) +} -// Get a prompt with arguments -let (description, messages) = try await client.getPrompt( - name: "example-prompt", - arguments: ["key": "value"] +// Register a prompt get handler +await server.withMethodHandler(GetPrompt.self) { params in + switch params.name { + case "interview": + let position = params.arguments?["position"]?.stringValue ?? "Software Engineer" + let company = params.arguments?["company"]?.stringValue ?? "Acme Corp" + let interviewee = params.arguments?["interviewee"]?.stringValue ?? "Candidate" + + let description = "Job interview for \(position) position at \(company)" + let messages: [Prompt.Message] = [ + .user(.text(text: "You are an interviewer for the \(position) position at \(company).")), + .user(.text(text: "Hello, I'm \(interviewee) and I'm here for the \(position) interview.")), + .assistant(.text(text: "Hi \(interviewee), welcome to \(company)! I'd like to start by asking about your background and experience.")) + ] + + return .init(description: description, messages: messages) + + case "customer-support": + // Similar implementation for customer support prompt + + default: + throw MCPError.invalidParams("Unknown prompt name: \(params.name)") + } +} +``` + +### Completions + +Servers can provide autocompletion suggestions for prompt and resource template arguments: + +```swift +// Enable completions capability +let server = Server( + name: "MyServer", + version: "1.0.0", + capabilities: .init( + completions: .init(), + prompts: .init(listChanged: true) + ) +) + +// Register a completion handler +await server.withMethodHandler(Complete.self) { params in + // Get the argument being completed + let argumentName = params.argument.name + let currentValue = params.argument.value + + // Check which prompt or resource is being completed + switch params.ref { + case .prompt(let promptRef): + // Provide completions for a prompt argument + if promptRef.name == "code_review" && argumentName == "language" { + // Simple prefix matching + let allLanguages = ["python", "perl", "php", "javascript", "java", "swift"] + let matches = allLanguages.filter { $0.hasPrefix(currentValue.lowercased()) } + + return .init( + completion: .init( + values: Array(matches.prefix(100)), // Max 100 items + total: matches.count, + hasMore: matches.count > 100 + ) + ) + } + + case .resource(let resourceRef): + // Provide completions for a resource template argument + if resourceRef.uri == "file:///{path}" && argumentName == "path" { + // Return directory suggestions + let suggestions = try getDirectoryCompletions(for: currentValue) + return .init( + completion: .init( + values: suggestions, + total: suggestions.count, + hasMore: false + ) + ) + } + } + + // No completions available + return .init(completion: .init(values: [], total: 0, hasMore: false)) +} +``` + +You can also use context from already-resolved arguments: + +```swift +await server.withMethodHandler(Complete.self) { params in + // Access context from previous argument completions + if let context = params.context, + let language = context.arguments["language"]?.stringValue { + + // Provide framework suggestions based on selected language + if language == "python" { + let frameworks = ["flask", "django", "fastapi", "tornado"] + let matches = frameworks.filter { + $0.hasPrefix(params.argument.value.lowercased()) + } + return .init( + completion: .init(values: matches, total: matches.count, hasMore: false) + ) + } + } + + return .init(completion: .init(values: [], total: 0, hasMore: false)) +} +``` + +### Sampling + +Servers can request LLM completions from clients through sampling. This enables agentic behaviors where servers can ask for AI assistance while maintaining human oversight. + +```swift +// Enable sampling capability in server +let server = Server( + name: "MyModelServer", + version: "1.0.0", + capabilities: .init( + sampling: .init(), // Enable sampling capability + tools: .init(listChanged: true) + ) +) + +// Request sampling from the connected client +do { + let result = try await server.requestSampling( + messages: [ + .user("Analyze this data and suggest next steps") + ], + systemPrompt: "You are a helpful data analyst", + temperature: 0.7, + maxTokens: 150 + ) + + // Use the LLM completion in your server logic + print("LLM suggested: \(result.content)") + +} catch { + print("Sampling request failed: \(error)") +} +``` + +Sampling enables powerful agentic workflows: + +- **Decision-making**: Ask the LLM to choose between options +- **Content generation**: Request drafts for user approval +- **Data analysis**: Get AI insights on complex data +- **Multi-step reasoning**: Chain AI completions with tool calls + +### Elicitations + +Servers can request information from users through elicitation: + +```swift +// Ask the user to provide some additional information +let schema = Elicitation.RequestSchema( + title: "Additional Information Required", + description: "Please provide the following details to continue", + properties: [ + "name": .object([ + "type": .string("string"), + "description": .string("Your full name") + ]), + "confirmed": .object([ + "type": .string("boolean"), + "description": .string("Do you confirm the provided information?") + ]) + ], + required: ["name", "confirmed"] +) + +let result = try await server.requestElicitation( + message: "Some details are needed before proceeding", + requestedSchema: schema +) + +switch result.action { +case .accept: + if let content = result.content { + let name = content["name"]?.stringValue + let confirmed = content["confirmed"]?.boolValue + // Use the collected data... + } +case .decline: + throw MCPError.invalidRequest("User declined to provide information") +case .cancel: + throw MCPError.invalidRequest("Operation canceled by user") +} +``` + +For URL-based elicitation (e.g., OAuth flows), use the URL overload: + +```swift +let result = try await server.requestElicitation( + message: "Please sign in to continue", + url: "https://example.com/oauth/authorize?client_id=...", + elicitationId: UUID().uuidString ) ``` +### Roots + +Servers can request the list of filesystem roots that the client has exposed: + +```swift +// Request roots from the connected client +// (requires the client to declare the roots capability) +let roots = try await server.listRoots() +for root in roots { + print("Root: \(root.name ?? root.uri) at \(root.uri)") +} + +// React to root list changes +await server.onNotification(RootsListChangedNotification.self) { _ in + let updatedRoots = try await server.listRoots() + print("Roots updated: \(updatedRoots.map { $0.uri })") +} +``` + +### Logging + +Servers can send structured log messages to clients: + +```swift +// Enable logging capability +let server = Server( + name: "MyServer", + version: "1.0.0", + capabilities: .init( + logging: .init(), + tools: .init(listChanged: true) + ) +) + +// Send log messages at different severity levels +try await server.log( + level: .info, + logger: "database", + data: Value.object([ + "message": .string("Database connected successfully"), + "host": .string("localhost"), + "port": .int(5432) + ]) +) + +try await server.log( + level: .error, + logger: "api", + data: Value.object([ + "message": .string("Request failed"), + "statusCode": .int(500), + "error": .string("Internal server error") + ]) +) + +// You can also use codable types directly +struct ErrorLog: Codable { + let message: String + let code: Int + let timestamp: String +} + +let errorLog = ErrorLog( + message: "Operation failed", + code: 500, + timestamp: ISO8601DateFormatter().string(from: Date()) +) + +try await server.log(level: .error, logger: "operations", data: errorLog) +``` + +Clients can control which log levels they receive: + +```swift +// Register a handler for client's logging level preferences +await server.withMethodHandler(SetLoggingLevel.self) { params in + let minimumLevel = params.level + + // Store the client's preference and filter log messages accordingly + // (Implementation depends on your server architecture) + storeLogLevel(minimumLevel) + + return Empty() +} +``` + +### Progress Tracking + +Servers can send incremental progress notifications during long-running tool calls by reading the `progressToken` from the request metadata and sending `ProgressNotification` messages: + +```swift +await server.withMethodHandler(CallTool.self) { params in + // Read the progress token from request metadata + guard let token = params._meta?.progressToken else { + // No progress token provided — run without reporting progress + return .init(content: [.text("Done")], isError: false) + } + + // Report initial progress + let started = ProgressNotification.message( + .init(progressToken: token, progress: 0, total: 100) + ) + try await server.notify(started) + + // ... do work ... + + // Report intermediate progress + let halfway = ProgressNotification.message( + .init(progressToken: token, progress: 50, total: 100, message: "Halfway there") + ) + try await server.notify(halfway) + + // ... do more work ... + + // Report completion + let done = ProgressNotification.message( + .init(progressToken: token, progress: 100, total: 100, message: "Complete") + ) + try await server.notify(done) + + return .init(content: [.text("Done")], isError: false) +} +``` + +#### Initialize Hook + +Control client connections with an initialize hook: + +```swift +// Start the server with an initialize hook +try await server.start(transport: transport) { clientInfo, clientCapabilities in + // Validate client info + guard clientInfo.name != "BlockedClient" else { + throw MCPError.invalidRequest("This client is not allowed") + } + + // You can also inspect client capabilities + if clientCapabilities.sampling == nil { + print("Client does not support sampling") + } + + // Perform any server-side setup based on client info + print("Client \(clientInfo.name) v\(clientInfo.version) connected") + + // If the hook completes without throwing, initialization succeeds +} +``` + +### HTTP Request Context in Handlers + +When a server is connected over `StatefulHTTPServerTransport` or `StatelessHTTPServerTransport`, +method handlers can observe the originating HTTP request (headers, body, path, method) via +`Server.currentHandlerContext` — a task-local set automatically before each handler runs: + +```swift +await server.withMethodHandler(CallTool.self) { params in + let httpRequest = Server.currentHandlerContext?.httpContext + let authHeader = httpRequest?.headers["Authorization"] + // …use the header to scope the response, audit, etc. + return .init(content: [.text("ok")], isError: false) +} +``` + +`httpContext` is `nil` for transports that don't carry HTTP context (e.g. `StdioTransport`, +`InMemoryTransport`) and for handlers reached off the dispatch path. + +Task-locals are not inherited by `Task.detached`. If you spawn a detached task from a handler, +capture the context up front: + +```swift +let ctx = Server.currentHandlerContext +Task.detached { await doWork(with: ctx?.httpContext) } +``` + +Custom HTTP transports can opt in by conforming to `HTTPContextProviding` and returning the +`HTTPRequest` for a given JSON-RPC id while it is in flight. + +### Graceful Shutdown + +We recommend using +[Swift Service Lifecycle](https://github.com/swift-server/swift-service-lifecycle) +for managing startup and shutdown of services. + +First, add the dependency to your `Package.swift`: + +```swift +.package(url: "https://github.com/swift-server/swift-service-lifecycle.git", from: "2.3.0"), +``` + +Then implement the MCP server as a `Service`: + +```swift +import MCP +import ServiceLifecycle +import Logging + +struct MCPService: Service { + let server: Server + let transport: Transport + + init(server: Server, transport: Transport) { + self.server = server + self.transport = transport + } + + func run() async throws { + // Start the server + try await server.start(transport: transport) + + // Keep running until external cancellation + try await Task.sleep(for: .days(365 * 100)) // Effectively forever + } + + func shutdown() async throws { + // Gracefully shutdown the server + await server.stop() + } +} +``` + +Then use it in your application: + +```swift +import MCP +import ServiceLifecycle +import Logging + +let logger = Logger(label: "com.example.mcp-server") + +// Create the MCP server +let server = Server( + name: "MyModelServer", + version: "1.0.0", + capabilities: .init( + prompts: .init(listChanged: true), + resources: .init(subscribe: true, listChanged: true), + tools: .init(listChanged: true) + ) +) + +// Add handlers directly to the server +await server.withMethodHandler(ListTools.self) { _ in + // Your implementation + return .init(tools: [ + Tool(name: "example", description: "An example tool") + ]) +} + +await server.withMethodHandler(CallTool.self) { params in + // Your implementation + return .init(content: [.text("Tool result")], isError: false) +} + +// Create MCP service and other services +let transport = StdioTransport(logger: logger) +let mcpService = MCPService(server: server, transport: transport) +let databaseService = DatabaseService() // Your other services + +// Create service group with signal handling +let serviceGroup = ServiceGroup( + services: [mcpService, databaseService], + configuration: .init( + gracefulShutdownSignals: [.sigterm, .sigint] + ), + logger: logger +) + +// Run the service group - this blocks until shutdown +try await serviceGroup.run() +``` + +This approach has several benefits: + +- **Signal handling**: +Automatically traps SIGINT, SIGTERM and triggers graceful shutdown +- **Graceful shutdown**: +Properly shuts down your MCP server and other services +- **Timeout-based shutdown**: +Configurable shutdown timeouts to prevent hanging processes +- **Advanced service management**: +`[ServiceLifecycle](https://swiftpackageindex.com/swift-server/swift-service-lifecycle/documentation/servicelifecycle)` +also supports service dependencies, conditional services, +and other useful features. + +## Transports + +MCP's transport layer handles communication between clients and servers. +The Swift SDK provides multiple built-in transports: + + +| Transport | Description | Platforms | Best for | +| --------- | ----------- | --------- | -------- | +| [`StdioTransport`](/Sources/MCP/Base/Transports/StdioTransport.swift) | Implements [stdio transport](https://modelcontextprotocol.io/specification/2025-06-18/basic/transports#stdio) using standard input/output streams | Apple platforms, Linux with glibc | Local subprocesses, CLI tools | +| [`HTTPClientTransport`](/Sources/MCP/Base/Transports/HTTPClientTransport.swift) | Implements [Streamable HTTP transport](https://modelcontextprotocol.io/specification/2025-06-18/basic/transports#streamable-http) using Foundation's URL Loading System | All platforms with Foundation | Remote servers, web applications | +| [`StatelessHTTPServerTransport`](/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift) | HTTP server transport with simple request-response semantics; no session management or SSE streaming | All platforms with Foundation | Simple HTTP servers, serverless/edge functions | +| [`StatefulHTTPServerTransport`](/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerTransport.swift) | HTTP server transport with full session management and SSE streaming for server-initiated messages | All platforms with Foundation | Full-featured HTTP servers, streaming notifications | +| [`InMemoryTransport`](/Sources/MCP/Base/Transports/InMemoryTransport.swift) | Custom in-memory transport for direct communication within the same process | All platforms | Testing, debugging, same-process client-server communication | +| [`NetworkTransport`](/Sources/MCP/Base/Transports/NetworkTransport.swift) | Custom transport using Apple's Network framework for TCP/UDP connections | Apple platforms only | Low-level networking, custom protocols | + + +### Custom Transport Implementation + +You can implement a custom transport by conforming to the `Transport` protocol: + +```swift +import MCP +import Foundation + +public actor MyCustomTransport: Transport { + public nonisolated let logger: Logger + private var isConnected = false + private let messageStream: AsyncThrowingStream + private let messageContinuation: AsyncThrowingStream.Continuation + + public init(logger: Logger? = nil) { + self.logger = logger ?? Logger(label: "my.custom.transport") + + var continuation: AsyncThrowingStream.Continuation! + self.messageStream = AsyncThrowingStream { continuation = $0 } + self.messageContinuation = continuation + } + + public func connect() async throws { + // Implement your connection logic + isConnected = true + } + + public func disconnect() async { + // Implement your disconnection logic + isConnected = false + messageContinuation.finish() + } + + public func send(_ data: Data) async throws { + // Implement your message sending logic + } + + public func receive() -> AsyncThrowingStream { + return messageStream + } +} +``` + +## Authentication + +`HTTPClientTransport` supports OAuth 2.1 Bearer token authorization per the +[MCP authorization specification](https://modelcontextprotocol.io/specification/2025-11-25/basic/authorization). +When a server returns `401 Unauthorized` or `403 Forbidden`, the transport automatically: + +1. Discovers Protected Resource Metadata (RFC 9728) at `/.well-known/oauth-protected-resource` +2. Discovers Authorization Server Metadata (RFC 8414 / OIDC Discovery 1.0) +3. Registers the client dynamically (RFC 7591) if needed +4. Acquires a Bearer token using the configured grant flow (PKCE enforced) +5. Retries the original request with the token attached + +Authorization is opt-in and disabled by default. +Pass an `OAuthAuthorizer` to `HTTPClientTransport(authorizer:)` to enable it. + +### Client: Client Credentials Flow + +Machine-to-machine authentication using a pre-shared client secret: + +```swift +let config = OAuthConfiguration( + grantType: .clientCredentials, + authentication: .clientSecretBasic(clientID: "my-app", clientSecret: "s3cr3t") +) +let authorizer = OAuthAuthorizer(configuration: config) +let transport = HTTPClientTransport( + endpoint: URL(string: "https://api.example.com/mcp")!, + authorizer: authorizer +) +let client = Client(name: "MyClient", version: "1.0.0") +try await client.connect(transport: transport) +``` + +### Client: Authorization Code Flow + +Interactive, browser-based authentication with PKCE. +Implement `OAuthAuthorizationDelegate` to open the authorization URL and capture the redirect: + +```swift +struct MyAuthDelegate: OAuthAuthorizationDelegate { + func presentAuthorizationURL(_ url: URL) async throws -> URL { + // Open the URL in a browser/webview and wait for the callback redirect URI. + // The returned URL must include the authorization code and state parameters. + return try await openBrowserAndWaitForCallback(url) + } +} + +let config = OAuthConfiguration( + grantType: .authorizationCode, + authentication: .none(clientID: "my-app"), + authorizationDelegate: MyAuthDelegate() +) +let authorizer = OAuthAuthorizer(configuration: config) +let transport = HTTPClientTransport( + endpoint: URL(string: "https://api.example.com/mcp")!, + authorizer: authorizer +) +``` + +### Client: Custom Token Provider + +Supply an externally acquired token (e.g., from a system credential store) via `accessTokenProvider`. +The SDK calls this closure after discovery completes. Return `nil` to fall back to the configured grant flow: + +```swift +let config = OAuthConfiguration( + grantType: .clientCredentials, + authentication: .none(clientID: "my-app"), + accessTokenProvider: { context, session in + // context contains the discovered resource URI, token endpoint, scopes, etc. + return try await KeychainTokenStore.shared.loadToken(for: context.resource) + } +) +``` + +### Client: Custom Token Storage + +By default, tokens are stored in memory and lost when the process exits. +To persist tokens across sessions, implement `TokenStorage` and pass it to `OAuthAuthorizer`: + +```swift +final class KeychainTokenStorage: TokenStorage { + func save(_ token: OAuthAccessToken) { + // Encode and store token.value in the system Keychain + } + + func load() -> OAuthAccessToken? { + // Load and decode token from the Keychain + return nil + } + + func clear() { + // Delete from the Keychain + } +} + +let authorizer = OAuthAuthorizer( + configuration: config, + tokenStorage: KeychainTokenStorage() +) +``` + +### Client: `private_key_jwt` Authentication + +Authenticate to the token endpoint using an asymmetric key (RFC 7523). +The SDK provides a built-in ES256 helper for P-256 keys: + +```swift +let config = OAuthConfiguration( + grantType: .clientCredentials, + authentication: .privateKeyJWT( + clientID: "my-app", + assertionFactory: { tokenEndpoint, clientID in + try OAuthConfiguration.makePrivateKeyJWTAssertion( + clientID: clientID, + tokenEndpoint: tokenEndpoint, + privateKeyPEM: myEC256PrivateKeyPEM // PEM-encoded P-256 private key + ) + } + ) +) +``` + +### Client: Endpoint Overrides + +Skip automatic discovery by providing explicit endpoint URLs. +Useful when the server does not publish well-known metadata documents: + +```swift +let config = OAuthConfiguration( + grantType: .clientCredentials, + authentication: .clientSecretBasic(clientID: "app", clientSecret: "secret"), + endpointOverrides: OAuthConfiguration.EndpointOverrides( + tokenEndpoint: URL(string: "https://auth.example.com/oauth/token")! + ) +) +``` + +### Server: Serving Protected Resource Metadata + +Per the MCP authorization specification, servers **MUST** serve Protected Resource Metadata +at `/.well-known/oauth-protected-resource` so clients can discover authorization server endpoints. + +Use `ProtectedResourceMetadataValidator` as the first validator in your pipeline so that +unauthenticated discovery requests are handled before the bearer token check: + +```swift +let metadata = OAuthProtectedResourceServerMetadata( + resource: "https://api.example.com", + authorizationServers: [URL(string: "https://auth.example.com")!], + scopesSupported: ["read", "write"] +) +let metadataValidator = ProtectedResourceMetadataValidator(metadata: metadata) +``` + +### Server: Validating Bearer Tokens + +Use `BearerTokenValidator` to authenticate incoming requests. +Your `tokenValidator` closure **MUST** verify the token's `aud` claim to prevent +token substitution attacks where a token intended for another resource is replayed against your server: + +```swift +let resourceIdentifier = URL(string: "https://api.example.com")! + +let bearerValidator = BearerTokenValidator( + resourceMetadataURL: URL(string: "https://api.example.com/.well-known/oauth-protected-resource")!, + resourceIdentifier: resourceIdentifier, + tokenValidator: { token, request, context in + guard let claims = try? verifyAndDecodeJWT(token) else { + return .invalidToken(errorDescription: "Token verification failed") + } + // Pass audience and expiry to BearerTokenInfo; the SDK validates the + // audience claim against resourceIdentifier automatically. + return .valid(BearerTokenInfo( + audience: claims.audience, + expiresAt: claims.expiresAt + )) + } +) + +let pipeline = StandardValidationPipeline(validators: [ + metadataValidator, // serves /.well-known/oauth-protected-resource unauthenticated + bearerValidator, // validates Bearer tokens on all other requests + AcceptHeaderValidator(mode: .sseRequired), + ContentTypeValidator(), + SessionValidator(), +]) +``` + +## Platform Availability + +The Swift SDK has the following platform requirements: + + +| Platform | Minimum Version | +| ------------------ | ---------------------------------------------------------------------------------------- | +| macOS | 13.0+ | +| iOS / Mac Catalyst | 16.0+ | +| watchOS | 9.0+ | +| tvOS | 16.0+ | +| visionOS | 1.0+ | +| Linux | Distributions with `glibc` or `musl`, including Ubuntu, Debian, Fedora, and Alpine Linux | + + +While the core library works on any platform supporting Swift 6 +(including Linux), running a client or server requires a compatible transport. + +## Debugging and Logging + +Enable logging to help troubleshoot issues: + +```swift +import Logging +import MCP + +// Configure Logger +LoggingSystem.bootstrap { label in + var handler = StreamLogHandler.standardOutput(label: label) + handler.logLevel = .debug + return handler +} + +// Create logger +let logger = Logger(label: "com.example.mcp") + +// Pass to client/server +let client = Client(name: "MyApp", version: "1.0.0") + +// Pass to transport +let transport = StdioTransport(logger: logger) +``` + +## Additional Resources + +- [MCP Specification](https://modelcontextprotocol.io/specification/2025-11-25) +- [Protocol Documentation](https://modelcontextprotocol.io) +- [GitHub Repository](https://github.com/modelcontextprotocol/swift-sdk) + +## Changelog + +This project follows [Semantic Versioning](https://semver.org/). +For pre-1.0 releases, +minor version increments (0.X.0) may contain breaking changes. + +For details about changes in each release, +see the [GitHub Releases page](https://github.com/modelcontextprotocol/swift-sdk/releases). + ## License -This project is licensed under the Apache License, Version 2.0. +This project is licensed under Apache 2.0 for new contributions, with existing code under MIT. See the [LICENSE](LICENSE) file for details. [mcp]: https://modelcontextprotocol.io +[mcp-spec-2025-11-25]: https://modelcontextprotocol.io/specification/2025-11-25 diff --git a/SECURITY.md b/SECURITY.md new file mode 100644 index 0000000..5029242 --- /dev/null +++ b/SECURITY.md @@ -0,0 +1,21 @@ +# Security Policy + +Thank you for helping keep the Model Context Protocol and its ecosystem secure. + +## Reporting Security Issues + +If you discover a security vulnerability in this repository, please report it through +the [GitHub Security Advisory process](https://docs.github.com/en/code-security/security-advisories/guidance-on-reporting-and-writing-information-about-vulnerabilities/privately-reporting-a-security-vulnerability) +for this repository. + +Please **do not** report security vulnerabilities through public GitHub issues, discussions, +or pull requests. + +## What to Include + +To help us triage and respond quickly, please include: + +- A description of the vulnerability +- Steps to reproduce the issue +- The potential impact +- Any suggested fixes (optional) diff --git a/Sources/MCP/Base/Authorization/OAuthAuthorizationCodeFlow.swift b/Sources/MCP/Base/Authorization/OAuthAuthorizationCodeFlow.swift new file mode 100644 index 0000000..0c278f1 --- /dev/null +++ b/Sources/MCP/Base/Authorization/OAuthAuthorizationCodeFlow.swift @@ -0,0 +1,253 @@ +import Foundation + +#if canImport(FoundationNetworking) + import FoundationNetworking +#endif + +// MARK: - OAuthAuthorizationCodeFlowing Protocol + +/// Internal protocol for driving the OAuth 2.1 authorization code flow. +protocol OAuthAuthorizationCodeFlowing: Sendable { + func buildURL( + authorizationEndpoint: URL, + resource: URL, + redirectURI: URL, + clientID: String, + codeChallenge: String, + scopes: Set?, + state: String, + scopeSerializer: any OAuthScopeSelecting + ) throws -> URL + + func perform( + authorizationURL: URL, + redirectURI: URL, + state: String, + delegate: (any OAuthAuthorizationDelegate)?, + session: URLSession + ) async throws -> String +} + +// MARK: - No-Redirect Session Delegate + +final class OAuthNoRedirectSessionDelegate: NSObject, URLSessionTaskDelegate, + @unchecked Sendable +{ + func urlSession( + _ session: URLSession, + task: URLSessionTask, + willPerformHTTPRedirection response: HTTPURLResponse, + newRequest request: URLRequest, + completionHandler: @escaping (URLRequest?) -> Void + ) { + completionHandler(nil) + } +} + +// MARK: - Authorization Code Flow + +/// Handles the browser-facing steps of the OAuth 2.1 authorization_code flow. +/// +/// Builds the authorization request URL, drives the redirect (via delegate or direct HTTP), +/// and extracts the authorization code from the callback redirect URL. +public struct OAuthAuthorizationCodeFlow: Sendable { + + public init() {} + + /// Builds the authorization request URL. + /// + /// - Parameters: + /// - authorizationEndpoint: The AS authorization endpoint. + /// - resource: The RFC 8707 resource indicator. + /// - redirectURI: The redirect URI registered for this client. + /// - clientID: The OAuth client identifier. + /// - codeChallenge: The PKCE S256 code challenge. + /// - scopes: Optional scope set to request. + /// - state: The CSRF state nonce. + /// - scopeSerializer: Serializes the scope set to a space-separated string. + /// - Returns: The full authorization request URL with all query parameters. + public func buildURL( + authorizationEndpoint: URL, + resource: URL, + redirectURI: URL, + clientID: String, + codeChallenge: String, + scopes: Set?, + state: String, + scopeSerializer: any OAuthScopeSelecting + ) throws -> URL { + guard var components = URLComponents( + url: authorizationEndpoint, + resolvingAgainstBaseURL: false + ) else { + throw OAuthAuthorizationError.authorizationServerMetadataDiscoveryFailed + } + + var queryItems: [URLQueryItem] = [ + .init(name: OAuthParameterName.responseType, value: OAuthParameterName.code), + .init(name: OAuthParameterName.clientID, value: clientID), + .init(name: OAuthParameterName.redirectURI, value: redirectURI.absoluteString), + .init(name: OAuthParameterName.state, value: state), + .init(name: OAuthParameterName.resource, value: resource.absoluteString), + .init(name: OAuthParameterName.codeChallenge, value: codeChallenge), + .init( + name: OAuthParameterName.codeChallengeMethod, value: OAuthCodeChallengeMethod.s256), + ] + + if let scope = scopes.flatMap(scopeSerializer.serialize) { + queryItems.append(.init(name: OAuthParameterName.scope, value: scope)) + } + + components.queryItems = queryItems + guard let url = components.url else { + throw OAuthAuthorizationError.authorizationServerMetadataDiscoveryFailed + } + return url + } + + /// Drives the interactive authorization redirect and returns the authorization code. + /// + /// When a delegate is provided, presents the authorization URL and awaits the redirect. + /// Without a delegate, sends a GET request and captures the redirect `Location` header. + /// + /// - Parameters: + /// - authorizationURL: The full authorization request URL. + /// - redirectURI: The expected redirect URI base for validation. + /// - state: The CSRF state nonce to verify in the redirect. + /// - delegate: Optional user-facing delegate for browser-based flows. + /// - session: The `URLSession` used for the no-redirect path. + /// - Returns: The extracted authorization code. + public func perform( + authorizationURL: URL, + redirectURI: URL, + state: String, + delegate: (any OAuthAuthorizationDelegate)?, + session: URLSession + ) async throws -> String { + if let delegate { + let redirectURL = try await delegate.presentAuthorizationURL(authorizationURL) + return try extractCode( + from: redirectURL, + expectedRedirectURI: redirectURI, + expectedState: state + ) + } + + var request = URLRequest(url: authorizationURL) + request.httpMethod = "GET" + request.setValue( + "text/html, \(ContentType.json)", forHTTPHeaderField: HTTPHeaderName.accept) + + let noRedirectDelegate = OAuthNoRedirectSessionDelegate() + let noRedirectSession = URLSession( + configuration: session.configuration, + delegate: noRedirectDelegate, + delegateQueue: nil + ) + defer { noRedirectSession.invalidateAndCancel() } + + let (_, response) = try await noRedirectSession.data(for: request) + guard let httpResponse = response as? HTTPURLResponse else { + throw OAuthAuthorizationError.authorizationResponseMissingRedirectLocation + } + + guard (300..<400).contains(httpResponse.statusCode) else { + if httpResponse.statusCode >= 400 { + throw OAuthAuthorizationError.authorizationRequestFailed(statusCode: httpResponse.statusCode) + } + throw OAuthAuthorizationError.authorizationResponseMissingRedirectLocation + } + + guard let location = httpResponse.value(forHTTPHeaderField: HTTPHeaderName.location), + !location.isEmpty, + let redirectURL = URL(string: location) + else { + throw OAuthAuthorizationError.authorizationResponseMissingRedirectLocation + } + + return try extractCode( + from: redirectURL, + expectedRedirectURI: redirectURI, + expectedState: state + ) + } + + /// Extracts and validates the authorization code from the redirect URL. + /// + /// - Parameters: + /// - redirectURL: The redirect URL received from the authorization server. + /// - expectedRedirectURI: The redirect URI used in the authorization request. + /// - expectedState: The CSRF state nonce sent in the authorization request. + /// - Returns: The authorization code. + public func extractCode( + from redirectURL: URL, + expectedRedirectURI: URL, + expectedState: String + ) throws -> String { + guard + let redirectComponents = URLComponents(url: redirectURL, resolvingAgainstBaseURL: false), + let expectedComponents = URLComponents( + url: expectedRedirectURI, resolvingAgainstBaseURL: false) + else { + throw OAuthAuthorizationError.authorizationResponseRedirectMismatch( + expected: expectedRedirectURI.absoluteString, + actual: redirectURL.absoluteString + ) + } + + if normalizedRedirectBase(redirectComponents) != normalizedRedirectBase(expectedComponents) { + throw OAuthAuthorizationError.authorizationResponseRedirectMismatch( + expected: expectedRedirectURI.absoluteString, + actual: redirectURL.absoluteString + ) + } + + guard + let state = redirectComponents.queryItems?.first(where: { + $0.name == OAuthParameterName.state + })?.value, + !state.isEmpty + else { + throw OAuthAuthorizationError.authorizationResponseMissingState + } + + guard state == expectedState else { + throw OAuthAuthorizationError.authorizationResponseStateMismatch( + expected: expectedState, + actual: state + ) + } + + guard + let code = redirectComponents.queryItems?.first(where: { + $0.name == OAuthParameterName.code + })?.value, + !code.isEmpty + else { + throw OAuthAuthorizationError.authorizationResponseMissingCode + } + + return code + } + + // MARK: - Private Helpers + + private func normalizedRedirectBase(_ components: URLComponents) -> String { + let scheme = components.scheme?.lowercased() ?? "" + let host = components.host?.lowercased() ?? "" + let port: Int + if let explicitPort = components.port { + port = explicitPort + } else if scheme == OAuthURLScheme.https { + port = OAuthDefaultPort.https + } else if scheme == OAuthURLScheme.http { + port = OAuthDefaultPort.http + } else { + port = -1 + } + let path = components.path.isEmpty ? "/" : components.path + return "\(scheme)://\(host):\(port)\(path)" + } +} + +extension OAuthAuthorizationCodeFlow: OAuthAuthorizationCodeFlowing {} diff --git a/Sources/MCP/Base/Authorization/OAuthAuthorizer.swift b/Sources/MCP/Base/Authorization/OAuthAuthorizer.swift new file mode 100644 index 0000000..b2e5be3 --- /dev/null +++ b/Sources/MCP/Base/Authorization/OAuthAuthorizer.swift @@ -0,0 +1,873 @@ +import Foundation + +#if canImport(FoundationNetworking) + import FoundationNetworking +#endif + +#if canImport(CryptoKit) + import CryptoKit +#endif + +// MARK: - HTTPClientAuthorizer Protocol + +/// Abstraction used by ``HTTPClientTransport`` to handle OAuth authorization challenges. +/// +/// Implement this protocol to provide custom token acquisition strategies, +/// or use the built-in ``OAuthAuthorizer`` for a full OAuth 2.1 implementation. +/// +/// ``HTTPClientTransport`` calls these methods automatically when it receives +/// `401 Unauthorized` or `403 Forbidden` responses from the server. +public protocol HTTPClientAuthorizer: AnyObject, Sendable { + /// The maximum number of authorization retries permitted for a single request. + /// + /// ``HTTPClientTransport`` will not call ``handleChallenge(statusCode:headers:endpoint:operationKey:session:)`` + /// more than this many times for a single outgoing request. + var maxAuthorizationAttempts: Int { get } + + /// Validates that the MCP endpoint URL satisfies the security requirements for OAuth. + /// + /// Called once before the first request is sent. Throw ``OAuthAuthorizationError/insecureOAuthEndpoint(context:url:)`` + /// if the URL does not meet the requirements (e.g., non-HTTPS non-loopback). + /// - Parameter endpoint: The MCP endpoint URL to validate. + func validateEndpointSecurity(for endpoint: URL) throws + + /// Returns the `Authorization` header value to attach to the next request, if a valid token is available. + /// + /// - Parameter endpoint: The MCP endpoint being requested. + /// - Returns: A `"Bearer "` string, or `nil` if no valid token is cached. + func authorizationHeader(for endpoint: URL) -> String? + + /// Handles an authorization challenge received from the server and attempts to acquire a new token. + /// + /// Called by ``HTTPClientTransport`` when a `401` or `403` response is received. + /// The implementation should attempt to obtain a valid access token and store it + /// so that a subsequent call to ``authorizationHeader(for:)`` returns the new value. + /// + /// - Parameters: + /// - statusCode: The HTTP status code (401 or 403). + /// - headers: All response headers from the challenge response. + /// - endpoint: The MCP endpoint that returned the challenge. + /// - operationKey: An optional identifier for the MCP operation (e.g., the JSON-RPC method), + /// used to track step-up attempts per operation. + /// - session: The `URLSession` to use for discovery and token requests. + /// - Returns: `true` if a new token was acquired and the original request should be retried; + /// `false` if the challenge cannot be handled. + func handleChallenge( + statusCode: Int, + headers: [String: String], + endpoint: URL, + operationKey: String?, + session: URLSession + ) async throws -> Bool + + /// Proactively refreshes the access token if it is close to expiry. + /// + /// Called by ``HTTPClientTransport`` before sending each request and before opening + /// an SSE stream, allowing the token to be silently renewed without a 401 round-trip. + /// Implementations should swallow refresh errors — if refresh fails, the normal + /// ``handleChallenge(statusCode:headers:endpoint:operationKey:session:)`` path recovers. + /// + /// - Parameters: + /// - endpoint: The MCP endpoint about to be contacted. + /// - session: The `URLSession` to use for token refresh requests. + func prepareAuthorization(for endpoint: URL, session: URLSession) async throws +} + +extension HTTPClientAuthorizer { + public func prepareAuthorization(for endpoint: URL, session: URLSession) async throws {} +} + +// MARK: - OAuthAuthorizer + +/// Full OAuth 2.1 implementation of ``HTTPClientAuthorizer``. +/// +/// `OAuthAuthorizer` orchestrates the complete MCP authorization flow on behalf of an HTTP client: +/// +/// 1. **Protected Resource Metadata discovery** (RFC 9728) — fetches +/// `/.well-known/oauth-protected-resource` to locate the authorization server. +/// 2. **Authorization Server Metadata discovery** (RFC 8414 / OIDC Discovery 1.0) — fetches +/// `/.well-known/oauth-authorization-server` or `/.well-known/openid-configuration`. +/// 3. **Dynamic Client Registration** (RFC 7591) — registers the client if no credentials are +/// pre-configured and the AS advertises a registration endpoint. +/// 4. **Token acquisition** — performs the configured grant flow (`authorization_code` with PKCE, +/// or `client_credentials`), binding tokens to the resource indicator (RFC 8707). +/// 5. **Token refresh** — attempts a `refresh_token` grant before a full re-authorization. +/// 6. **Scope step-up** — handles `403 insufficient_scope` challenges by re-requesting with +/// the union of existing and required scopes. +/// +/// Pass an instance to `HTTPClientTransport(authorizer:)` to enable automatic authorization: +/// +/// ```swift +/// let config = OAuthConfiguration( +/// grantType: .clientCredentials, +/// authentication: .clientSecretBasic(clientID: "my-app", clientSecret: "s3cr3t") +/// ) +/// let authorizer = OAuthAuthorizer(configuration: config) +/// let transport = HTTPClientTransport(endpoint: serverURL, authorizer: authorizer) +/// ``` +/// +/// - Important: This type is `@unchecked Sendable`. All mutable state is accessed +/// exclusively through the `HTTPClientTransport` actor, which serializes every call. +/// Do **not** share a single `OAuthAuthorizer` instance across multiple transports — +/// doing so would violate the isolation contract and risk concurrent mutation. +public final class OAuthAuthorizer: HTTPClientAuthorizer, @unchecked Sendable { + + // MARK: - Mutable State + + private var configuration: OAuthConfiguration + private let tokenStorage: TokenStorage + private var selectedAuthorizationServer: URL? + private var protectedResourceMetadata: OAuthProtectedResourceMetadata? + private var authorizationServerMetadata: OAuthAuthorizationServerMetadata? + private var cachedProtectedResourceMetadataURL: URL? + private var stepUpAttempts: [String: Int] = [:] + private var clientRegistrationAttempted = false + private var clientSecretExpiresAt: Date? + + // MARK: - Composable Dependencies + + private let scopeSelector: any OAuthScopeSelecting + private let challengeParser: any OAuthWWWAuthenticateParsing + private let urlValidator: any OAuthURLValidating + private let discoveryClient: any OAuthDiscoveryFetching + private let tokenEndpointClient: any OAuthTokenRequesting + private let clientRegistrar: any OAuthClientRegistering + private let authCodeFlow: any OAuthAuthorizationCodeFlowing + + /// Creates an `OAuthAuthorizer` with the given configuration and optional injectable dependencies. + /// + /// - Parameters: + /// - configuration: OAuth 2.1 configuration controlling the grant type, authentication method, + /// endpoint discovery overrides, and retry policy. + /// - tokenStorage: Stores acquired access tokens. Defaults to ``InMemoryTokenStorage``, + /// which loses tokens when the process exits. Supply a Keychain-backed implementation + /// to persist tokens across sessions. + /// - scopeSelector: Strategy for selecting OAuth scopes from challenge and metadata hints. + /// Defaults to ``DefaultOAuthScopeSelector``. + /// - challengeParser: Parses `WWW-Authenticate: Bearer` challenge headers. + /// Defaults to ``DefaultOAuthWWWAuthenticateParser``. + /// - metadataDiscovery: Constructs well-known discovery URLs and validates resource URI matching. + /// Defaults to ``DefaultOAuthMetadataDiscovery``. + public convenience init( + configuration: OAuthConfiguration, + tokenStorage: TokenStorage = InMemoryTokenStorage(), + scopeSelector: any OAuthScopeSelecting = DefaultOAuthScopeSelector(), + challengeParser: any OAuthWWWAuthenticateParsing = DefaultOAuthWWWAuthenticateParser(), + metadataDiscovery: any OAuthMetadataDiscovering = DefaultOAuthMetadataDiscovery() + ) { + let urlValidator = OAuthURLValidator( + allowLoopbackHTTPForAuthorizationServer: + configuration.allowLoopbackHTTPAuthorizationServerEndpoints + ) + self.init( + configuration: configuration, + tokenStorage: tokenStorage, + scopeSelector: scopeSelector, + challengeParser: challengeParser, + urlValidator: urlValidator, + discoveryClient: OAuthDiscoveryClient( + metadataDiscovery: metadataDiscovery, + urlValidator: urlValidator + ), + tokenEndpointClient: OAuthTokenEndpointClient(urlValidator: urlValidator), + clientRegistrar: OAuthClientRegistrar(urlValidator: urlValidator), + authCodeFlow: OAuthAuthorizationCodeFlow() + ) + } + + init( + configuration: OAuthConfiguration, + tokenStorage: TokenStorage = InMemoryTokenStorage(), + scopeSelector: any OAuthScopeSelecting = DefaultOAuthScopeSelector(), + challengeParser: any OAuthWWWAuthenticateParsing = DefaultOAuthWWWAuthenticateParser(), + urlValidator: any OAuthURLValidating, + discoveryClient: any OAuthDiscoveryFetching, + tokenEndpointClient: any OAuthTokenRequesting, + clientRegistrar: any OAuthClientRegistering, + authCodeFlow: any OAuthAuthorizationCodeFlowing + ) { + self.configuration = configuration + self.tokenStorage = tokenStorage + self.scopeSelector = scopeSelector + self.challengeParser = challengeParser + self.urlValidator = urlValidator + self.discoveryClient = discoveryClient + self.tokenEndpointClient = tokenEndpointClient + self.clientRegistrar = clientRegistrar + self.authCodeFlow = authCodeFlow + } + + // MARK: - HTTPClientAuthorizer + + public var maxAuthorizationAttempts: Int { + configuration.retryPolicy.maxAuthorizationAttempts + } + + public func validateEndpointSecurity(for endpoint: URL) throws { + try urlValidator.validateHTTPSOrLoopback(endpoint, context: "MCP endpoint") + } + + public func authorizationHeader(for endpoint: URL) -> String? { + guard let accessToken = tokenStorage.load() else { return nil } + if let tokenAuthorizationServer = accessToken.authorizationServer, + let selectedAuthorizationServer, + !authorizationServersMatch(tokenAuthorizationServer, selectedAuthorizationServer) + { + tokenStorage.clear() + return nil + } + if accessToken.isExpired() { + return nil + } + return "\(OAuthTokenType.bearer) \(accessToken.value)" + } + + public func handleChallenge( + statusCode: Int, + headers: [String: String], + endpoint: URL, + operationKey: String? = nil, + session: URLSession + ) async throws -> Bool { + try validateEndpointSecurity(for: endpoint) + let challenge = challengeParser.parseBearer(from: headers) + + switch statusCode { + case 401: + if let refreshToken = tokenStorage.load()?.refreshToken { + tokenStorage.clear() + let metadata = try await discoverProtectedResourceMetadata( + endpoint: endpoint, + challenge: challenge, + session: session + ) + let asMetadata = try await resolveAuthorizationServerMetadata( + metadata: metadata, + session: session + ) + let resource = try canonicalResource(for: endpoint) + let requestedScopes = scopeSelector.selectScopes( + challengeScope: challenge?.scope, + scopesSupported: metadata.scopesSupported + ) + if try await refreshAccessToken( + refreshToken: refreshToken, + resource: resource, + requestedScopes: requestedScopes, + asMetadata: asMetadata, + session: session + ) { + return true + } + } else { + tokenStorage.clear() + } + + let metadata = try await discoverProtectedResourceMetadata( + endpoint: endpoint, + challenge: challenge, + session: session + ) + let requestedScopes = scopeSelector.selectScopes( + challengeScope: challenge?.scope, + scopesSupported: metadata.scopesSupported + ) + + let providerContext = try await makeAccessTokenProviderContext( + statusCode: statusCode, + endpoint: endpoint, + challenge: challenge, + metadata: metadata, + requestedScopes: requestedScopes, + session: session + ) + if let externalToken = try await fetchAccessTokenFromProvider( + context: providerContext, + session: session + ) { + storeExternalAccessToken( + externalToken, + requestedScopes: providerContext.requestedScopes, + authorizationServer: providerContext.authorizationServer + ) + return true + } + + try await acquireToken( + endpoint: endpoint, + metadata: metadata, + requestedScopes: requestedScopes, + session: session + ) + return true + + case 403: + guard challenge?.error?.lowercased() == "insufficient_scope" else { return false } + + let metadata = try await discoverProtectedResourceMetadata( + endpoint: endpoint, + challenge: challenge, + session: session + ) + let requiredScopes = + scopeSelector.selectScopes( + challengeScope: challenge?.scope, + scopesSupported: metadata.scopesSupported + ) ?? [] + + let existingScopes = tokenStorage.load()?.scopes ?? [] + let upgradedScopes = existingScopes.union(requiredScopes) + let resourceKey = try discoveryClient.metadataDiscovery.canonicalResourceURI( + from: endpoint + ).absoluteString + let operationAttemptKey = normalizedOperationKey(operationKey) + let attemptKey = + "\(resourceKey)|\(operationAttemptKey)|\(upgradedScopes.sorted().joined(separator: " "))" + let attempts = stepUpAttempts[attemptKey, default: 0] + guard attempts < configuration.retryPolicy.maxScopeUpgradeAttempts else { + return false + } + stepUpAttempts[attemptKey] = attempts + 1 + + let providerRequestedScopes = upgradedScopes.isEmpty ? nil : upgradedScopes + let providerContext = try await makeAccessTokenProviderContext( + statusCode: statusCode, + endpoint: endpoint, + challenge: challenge, + metadata: metadata, + requestedScopes: providerRequestedScopes, + session: session + ) + if let externalToken = try await fetchAccessTokenFromProvider( + context: providerContext, + session: session + ) { + storeExternalAccessToken( + externalToken, + requestedScopes: providerContext.requestedScopes, + authorizationServer: providerContext.authorizationServer + ) + return true + } + + try await acquireToken( + endpoint: endpoint, + metadata: metadata, + requestedScopes: upgradedScopes, + session: session + ) + return true + + default: + return false + } + } + + public func prepareAuthorization(for endpoint: URL, session: URLSession) async throws { + guard configuration.proactiveRefreshWindowSeconds > 0 else { return } + guard let token = tokenStorage.load() else { return } + guard token.isExpired(skewSeconds: configuration.proactiveRefreshWindowSeconds) else { + return + } + guard let refreshToken = token.refreshToken else { return } + guard let asMeta = authorizationServerMetadata, asMeta.tokenEndpoint != nil else { return } + + let resource: URL + do { + resource = try canonicalResource(for: endpoint) + } catch { + return + } + + let requestedScopes = token.scopes.isEmpty ? nil : token.scopes + _ = try? await refreshAccessToken( + refreshToken: refreshToken, + resource: resource, + requestedScopes: requestedScopes, + asMetadata: asMeta, + session: session + ) + } + + // MARK: - Discovery + + private func discoverProtectedResourceMetadata( + endpoint: URL, + challenge: OAuthBearerChallenge?, + session: URLSession + ) async throws -> OAuthProtectedResourceMetadata { + if let protectedResourceMetadata { + let incomingURL = challenge?.resourceMetadataURL + if let incomingURL, incomingURL != cachedProtectedResourceMetadataURL { + self.protectedResourceMetadata = nil + self.authorizationServerMetadata = nil + self.selectedAuthorizationServer = nil + self.cachedProtectedResourceMetadataURL = nil + } else { + return protectedResourceMetadata + } + } + + var candidates: [URL] = [] + + if let challengeURL = challenge?.resourceMetadataURL { + try urlValidator.validateHTTPSOrLoopback( + challengeURL, context: "Protected resource metadata URL") + if let host = URLComponents(url: challengeURL, resolvingAgainstBaseURL: false)?.host? + .lowercased(), urlValidator.isPrivateIPHost(host) + { + throw OAuthAuthorizationError.privateIPAddressBlocked( + context: "Protected resource metadata URL", + url: challengeURL.absoluteString + ) + } + candidates.append(challengeURL) + } + if let configuredURL = configuration.endpointOverrides.protectedResourceMetadataURL, + !candidates.contains(configuredURL) + { + try urlValidator.validateHTTPSOrLoopback( + configuredURL, + context: "Configured protected resource metadata URL" + ) + candidates.append(configuredURL) + } + + for fallback in discoveryClient.metadataDiscovery.protectedResourceMetadataURLs( + for: endpoint) + where !candidates.contains(fallback) { + candidates.append(fallback) + } + + let fallbackIssuer = try? discoveryClient.metadataDiscovery + .authorizationServerFallbackIssuer(from: endpoint) + let metadata = try await discoveryClient.fetchProtectedResourceMetadata( + candidates: candidates, fallbackIssuer: fallbackIssuer, session: session) + try validateProtectedResource(metadata: metadata, endpoint: endpoint) + + self.protectedResourceMetadata = metadata + self.cachedProtectedResourceMetadataURL = candidates.first + return metadata + } + + private func validateProtectedResource( + metadata: OAuthProtectedResourceMetadata, endpoint: URL + ) throws { + guard let resource = metadata.resource?.trimmingCharacters(in: .whitespacesAndNewlines), + !resource.isEmpty + else { + return + } + + guard let resourceURL = URL(string: resource) else { + throw OAuthAuthorizationError.invalidResourceURI( + "Protected resource metadata contains an invalid resource URI: \(resource)" + ) + } + + let expected = try discoveryClient.metadataDiscovery.canonicalResourceURI(from: endpoint) + let actual = try discoveryClient.metadataDiscovery.canonicalResourceURI(from: resourceURL) + guard discoveryClient.metadataDiscovery.protectedResourceMatches( + resource: actual, endpoint: expected) + else { + throw OAuthAuthorizationError.protectedResourceMismatch( + expected: expected.absoluteString, + actual: actual.absoluteString + ) + } + } + + private func resolveAuthorizationServerMetadata( + metadata: OAuthProtectedResourceMetadata, + session: URLSession + ) async throws -> OAuthAuthorizationServerMetadata { + if let cached = authorizationServerMetadata { + return cached + } + + let candidates: [URL] + if let override = configuration.endpointOverrides.authorizationServerURL { + try urlValidator.validateAuthorizationServer( + override, context: "Authorization server issuer") + candidates = [override] + } else if let selected = selectedAuthorizationServer { + candidates = [selected] + } else { + guard !metadata.authorizationServers.isEmpty else { + throw OAuthAuthorizationError.missingAuthorizationServer + } + candidates = metadata.authorizationServers + } + + let (server, asMetadata) = try await discoveryClient.fetchAuthorizationServerMetadata( + candidates: candidates, session: session) + self.selectedAuthorizationServer = server + self.authorizationServerMetadata = asMetadata + return asMetadata + } + + // MARK: - Token Acquisition + + private func acquireToken( + endpoint: URL, + metadata: OAuthProtectedResourceMetadata, + requestedScopes: Set?, + session: URLSession + ) async throws { + let asMetadata = try await resolveAuthorizationServerMetadata( + metadata: metadata, session: session) + try await maybeRegisterClient(asMetadata: asMetadata, session: session) + let resource = try canonicalResource(for: endpoint) + + switch configuration.grantType { + case .clientCredentials: + try await acquireTokenViaClientCredentials( + resource: resource, + requestedScopes: requestedScopes, + asMetadata: asMetadata, + session: session + ) + case .authorizationCode: + try await acquireTokenViaAuthorizationCode( + resource: resource, + requestedScopes: requestedScopes, + asMetadata: asMetadata, + session: session + ) + } + } + + private func acquireTokenViaClientCredentials( + resource: URL, + requestedScopes: Set?, + asMetadata: OAuthAuthorizationServerMetadata, + session: URLSession + ) async throws { + let tokenEndpoint = try resolveTokenEndpoint(asMetadata: asMetadata) + var bodyParameters: [String: String] = configuration.additionalTokenRequestParameters + bodyParameters[OAuthParameterName.grantType] = OAuthGrantTypeValue.clientCredentials + bodyParameters[OAuthParameterName.resource] = resource.absoluteString + if let scope = requestedScopes.flatMap(scopeSelector.serialize) { + bodyParameters[OAuthParameterName.scope] = scope + } + let decoded = try await tokenEndpointClient.request( + parameters: &bodyParameters, + endpoint: tokenEndpoint, + authentication: configuration.authentication, + session: session + ) + storeTokenResponse(decoded, requestedScopes: requestedScopes) + } + + private func acquireTokenViaAuthorizationCode( + resource: URL, + requestedScopes: Set?, + asMetadata: OAuthAuthorizationServerMetadata, + session: URLSession + ) async throws { + guard let authorizationEndpoint = asMetadata.authorizationEndpoint else { + throw OAuthAuthorizationError.tokenEndpointMissing + } + try urlValidator.validateAuthorizationServer( + authorizationEndpoint, context: "Authorization endpoint") + if let host = URLComponents(url: authorizationEndpoint, resolvingAgainstBaseURL: false)? + .host?.lowercased(), urlValidator.isPrivateIPHost(host) + { + throw OAuthAuthorizationError.privateIPAddressBlocked( + context: "Authorization endpoint", + url: authorizationEndpoint.absoluteString + ) + } + try urlValidator.validateRedirectURI(configuration.authorizationRedirectURI) + try PKCE.checkSupport(in: asMetadata) + + let verifier = PKCE.makeVerifier() + let challenge = try PKCE.makeChallenge(from: verifier) + let state = UUID().uuidString + + let authorizationURL = try authCodeFlow.buildURL( + authorizationEndpoint: authorizationEndpoint, + resource: resource, + redirectURI: configuration.authorizationRedirectURI, + clientID: configuration.authentication.clientID, + codeChallenge: challenge, + scopes: requestedScopes, + state: state, + scopeSerializer: scopeSelector + ) + + let authorizationCode = try await authCodeFlow.perform( + authorizationURL: authorizationURL, + redirectURI: configuration.authorizationRedirectURI, + state: state, + delegate: configuration.authorizationDelegate, + session: session + ) + + let tokenEndpoint = try resolveTokenEndpoint(asMetadata: asMetadata) + var bodyParameters: [String: String] = configuration.additionalTokenRequestParameters + bodyParameters[OAuthParameterName.grantType] = OAuthGrantTypeValue.authorizationCode + bodyParameters[OAuthParameterName.code] = authorizationCode + bodyParameters[OAuthParameterName.codeVerifier] = verifier + bodyParameters[OAuthParameterName.redirectURI] = + configuration.authorizationRedirectURI.absoluteString + bodyParameters[OAuthParameterName.resource] = resource.absoluteString + if let scope = requestedScopes.flatMap(scopeSelector.serialize) { + bodyParameters[OAuthParameterName.scope] = scope + } + + let decoded = try await tokenEndpointClient.request( + parameters: &bodyParameters, + endpoint: tokenEndpoint, + authentication: configuration.authentication, + session: session + ) + storeTokenResponse(decoded, requestedScopes: requestedScopes) + } + + // MARK: - Token Refresh + + private func refreshAccessToken( + refreshToken: String, + resource: URL, + requestedScopes: Set?, + asMetadata: OAuthAuthorizationServerMetadata, + session: URLSession + ) async throws -> Bool { + let tokenEndpoint: URL + do { + tokenEndpoint = try resolveTokenEndpoint(asMetadata: asMetadata) + } catch { + return false + } + + var bodyParameters: [String: String] = configuration.additionalTokenRequestParameters + bodyParameters[OAuthParameterName.grantType] = OAuthGrantTypeValue.refreshToken + bodyParameters[OAuthParameterName.refreshToken] = refreshToken + bodyParameters[OAuthParameterName.resource] = resource.absoluteString + if let scope = requestedScopes.flatMap(scopeSelector.serialize) { + bodyParameters[OAuthParameterName.scope] = scope + } + + do { + let decoded = try await tokenEndpointClient.request( + parameters: &bodyParameters, + endpoint: tokenEndpoint, + authentication: configuration.authentication, + session: session + ) + storeTokenResponse(decoded, requestedScopes: requestedScopes) + return true + } catch let error as OAuthAuthorizationError { + if case .tokenRequestFailed(let statusCode, _) = error, + (400..<500).contains(statusCode) + { + return false + } + throw error + } + } + + // MARK: - Client Registration + + private func maybeRegisterClient( + asMetadata: OAuthAuthorizationServerMetadata, + session: URLSession + ) async throws { + if let expiry = clientSecretExpiresAt, Date() >= expiry { + clientSecretExpiresAt = nil + clientRegistrationAttempted = false + configuration.authentication = .none(clientID: configuration.authentication.clientID) + } + + guard !clientRegistrationAttempted else { return } + guard case .none = configuration.authentication else { return } + + clientRegistrationAttempted = true + + if let (registration, updatedAuth) = try await clientRegistrar.register( + configuration: configuration, + asMetadata: asMetadata, + session: session + ) { + configuration.authentication = updatedAuth + if let expiresAt = registration.clientSecretExpiresAt, expiresAt > 0 { + clientSecretExpiresAt = Date(timeIntervalSince1970: Double(expiresAt)) + } + } + } + + // MARK: - State Helpers + + private func storeTokenResponse( + _ decoded: OAuthTokenResponse, + requestedScopes: Set? + ) { + let scopeSet: Set + if let scope = decoded.scope { + scopeSet = scopeSelector.parseScopeString(scope) + } else { + scopeSet = requestedScopes ?? [] + } + let expiresAt = decoded.expiresIn.map { Date().addingTimeInterval(TimeInterval($0)) } + tokenStorage.save(OAuthAccessToken( + value: decoded.accessToken, + tokenType: OAuthTokenType.bearer, + expiresAt: expiresAt, + scopes: scopeSet, + authorizationServer: selectedAuthorizationServer, + refreshToken: decoded.refreshToken, + clientID: nonEmptyClientID() + )) + } + + /// Returns the configured `client_id` or `nil` if the authorizer has not yet been assigned one. + private func nonEmptyClientID() -> String? { + let id = configuration.authentication.clientID + return id.isEmpty ? nil : id + } + + private func resolveTokenEndpoint( + asMetadata: OAuthAuthorizationServerMetadata + ) throws -> URL { + if let configuredEndpoint = configuration.endpointOverrides.tokenEndpoint { + try urlValidator.validateAuthorizationServer( + configuredEndpoint, context: "Configured token endpoint") + return configuredEndpoint + } + + guard let tokenEndpoint = asMetadata.tokenEndpoint else { + throw OAuthAuthorizationError.tokenEndpointMissing + } + try urlValidator.validateAuthorizationServer(tokenEndpoint, context: "Token endpoint") + if let host = URLComponents(url: tokenEndpoint, resolvingAgainstBaseURL: false)?.host? + .lowercased(), urlValidator.isPrivateIPHost(host) + { + throw OAuthAuthorizationError.privateIPAddressBlocked( + context: "Token endpoint", + url: tokenEndpoint.absoluteString + ) + } + return tokenEndpoint + } + + private func canonicalResource(for endpoint: URL) throws -> URL { + let endpointCanonical = try discoveryClient.metadataDiscovery.canonicalResourceURI( + from: endpoint) + + if let configuredResource = configuration.endpointOverrides.resource { + let configuredCanonical = try discoveryClient.metadataDiscovery.canonicalResourceURI( + from: configuredResource) + guard discoveryClient.metadataDiscovery.protectedResourceMatches( + resource: configuredCanonical, endpoint: endpointCanonical) + else { + throw OAuthAuthorizationError.protectedResourceMismatch( + expected: endpointCanonical.absoluteString, + actual: configuredCanonical.absoluteString + ) + } + return configuredCanonical + } + + if let prmResourceString = protectedResourceMetadata?.resource, + let prmResourceURL = URL(string: prmResourceString) + { + let prmCanonical = try discoveryClient.metadataDiscovery.canonicalResourceURI( + from: prmResourceURL) + guard discoveryClient.metadataDiscovery.protectedResourceMatches( + resource: prmCanonical, endpoint: endpointCanonical) + else { + throw OAuthAuthorizationError.protectedResourceMismatch( + expected: endpointCanonical.absoluteString, + actual: prmCanonical.absoluteString + ) + } + return prmCanonical + } + + return endpointCanonical + } + + private func authorizationServersMatch(_ lhs: URL, _ rhs: URL) -> Bool { + normalizedAuthorizationServer(lhs) == normalizedAuthorizationServer(rhs) + } + + private func normalizedAuthorizationServer(_ url: URL) -> URL? { + guard var components = URLComponents(url: url, resolvingAgainstBaseURL: false), + let scheme = components.scheme?.lowercased(), + let host = components.host?.lowercased(), + scheme == OAuthURLScheme.http || scheme == OAuthURLScheme.https + else { + return nil + } + components.scheme = scheme + components.host = host + components.query = nil + components.fragment = nil + if components.path == "/" { components.path = "" } + return components.url + } + + // MARK: - External Token Provider + + private func fetchAccessTokenFromProvider( + context: OAuthConfiguration.AccessTokenProviderContext, + session: URLSession + ) async throws -> String? { + guard let provider = configuration.accessTokenProvider else { return nil } + guard let token = try await provider(context, session), !token.isEmpty else { return nil } + return token + } + + private func storeExternalAccessToken( + _ token: String, + requestedScopes: Set?, + authorizationServer: URL? + ) { + tokenStorage.save(OAuthAccessToken( + value: token, + tokenType: OAuthTokenType.bearer, + expiresAt: nil, + scopes: requestedScopes ?? [], + authorizationServer: authorizationServer, + refreshToken: nil, + clientID: nonEmptyClientID() + )) + } + + private func makeAccessTokenProviderContext( + statusCode: Int, + endpoint: URL, + challenge: OAuthBearerChallenge?, + metadata: OAuthProtectedResourceMetadata, + requestedScopes: Set?, + session: URLSession + ) async throws -> OAuthConfiguration.AccessTokenProviderContext { + let asMetadata = try await resolveAuthorizationServerMetadata( + metadata: metadata, session: session) + let resource = try canonicalResource(for: endpoint) + let authorizationServer = configuration.endpointOverrides.authorizationServerURL + ?? selectedAuthorizationServer + ?? metadata.authorizationServers.first + let tokenEndpoint = configuration.endpointOverrides.tokenEndpoint ?? asMetadata.tokenEndpoint + + return OAuthConfiguration.AccessTokenProviderContext( + statusCode: statusCode, + endpoint: endpoint, + resource: resource, + authorizationServer: authorizationServer, + authorizationEndpoint: asMetadata.authorizationEndpoint, + tokenEndpoint: tokenEndpoint, + registrationEndpoint: asMetadata.registrationEndpoint, + challengedScope: challenge?.scope, + scopesSupported: metadata.scopesSupported, + requestedScopes: requestedScopes + ) + } + + private func normalizedOperationKey(_ operationKey: String?) -> String { + guard let operationKey else { return "" } + let normalized = operationKey.trimmingCharacters(in: .whitespacesAndNewlines) + return normalized.isEmpty ? "" : normalized + } +} diff --git a/Sources/MCP/Base/Authorization/OAuthClientRegistrar.swift b/Sources/MCP/Base/Authorization/OAuthClientRegistrar.swift new file mode 100644 index 0000000..6267421 --- /dev/null +++ b/Sources/MCP/Base/Authorization/OAuthClientRegistrar.swift @@ -0,0 +1,186 @@ +import Foundation + +#if canImport(FoundationNetworking) + import FoundationNetworking +#endif + +/// Internal protocol for OAuth dynamic client registration. +protocol OAuthClientRegistering: Sendable { + func register( + configuration: OAuthConfiguration, + asMetadata: OAuthAuthorizationServerMetadata, + session: URLSession + ) async throws -> ( + response: OAuthClientRegistrationResponse, + updatedAuthentication: OAuthConfiguration.TokenEndpointAuthentication + )? +} + +/// Stateless OAuth dynamic client registration logic. +/// +/// Handles Client ID Metadata Document (CIMD) detection and RFC 7591 dynamic registration. +/// State tracking (`clientRegistrationAttempted`, `clientSecretExpiresAt`) is the caller's responsibility. +struct OAuthClientRegistrar: Sendable { + let urlValidator: OAuthURLValidator + + init(urlValidator: OAuthURLValidator) { + self.urlValidator = urlValidator + } + + /// Attempts to register the client, if applicable. + /// + /// Returns `nil` if registration is not needed: + /// - Credentials are already configured (not `.none`) + /// - CIMD is in use and the server supports it (pre-registered) + /// - No registration endpoint is available and no CIMD mismatch error + /// + /// Throws if registration was attempted but failed (4xx, 5xx, or unexpected response). + func register( + configuration: OAuthConfiguration, + asMetadata: OAuthAuthorizationServerMetadata, + session: URLSession + ) async throws -> ( + response: OAuthClientRegistrationResponse, + updatedAuthentication: OAuthConfiguration.TokenEndpointAuthentication + )? { + guard case .none(let clientID) = configuration.authentication else { + return nil + } + + let hasClientIDMetadataDocument = isHTTPSURLWithPath(clientID) + let supportsClientIDMetadataDocument = asMetadata.clientIDMetadataDocumentSupported == true + + if supportsClientIDMetadataDocument, + clientIDLooksLikeURL(clientID), + !hasClientIDMetadataDocument + { + throw OAuthAuthorizationError.invalidClientIDMetadataURL(clientID) + } + + if hasClientIDMetadataDocument, supportsClientIDMetadataDocument { + return nil + } + + guard let registrationEndpoint = asMetadata.registrationEndpoint else { + if hasClientIDMetadataDocument && !supportsClientIDMetadataDocument { + throw OAuthAuthorizationError.cimdNotSupported(clientID: clientID) + } + return nil + } + try urlValidator.validateAuthorizationServer( + registrationEndpoint, context: "Client registration endpoint") + + var request = URLRequest(url: registrationEndpoint) + request.httpMethod = "POST" + request.setValue(ContentType.json, forHTTPHeaderField: HTTPHeaderName.contentType) + request.setValue(ContentType.json, forHTTPHeaderField: HTTPHeaderName.accept) + + let grantTypes: [String] + let responseTypes: [String] + switch configuration.grantType { + case .authorizationCode: + grantTypes = [OAuthGrantTypeValue.authorizationCode, OAuthGrantTypeValue.refreshToken] + responseTypes = [OAuthParameterName.code] + case .clientCredentials: + grantTypes = [OAuthGrantTypeValue.clientCredentials] + responseTypes = [] + } + + var registrationPayload: [String: Any] = [ + "client_name": configuration.clientName, + "grant_types": grantTypes, + "token_endpoint_auth_method": configuration.authentication.methodName, + ] + if !responseTypes.isEmpty { + registrationPayload["response_types"] = responseTypes + } + if configuration.grantType == .authorizationCode { + registrationPayload["redirect_uris"] = [ + configuration.authorizationRedirectURI.absoluteString + ] + } + + let httpBody = try JSONSerialization.data(withJSONObject: registrationPayload) + + request.httpBody = httpBody + + let (data, response) = try await session.data(for: request) + if let httpResponse = response as? HTTPURLResponse, + (400..<500).contains(httpResponse.statusCode) + { + let oauthError = + (try? JSONDecoder().decode(OAuthTokenErrorResponse.self, from: data))?.error + throw OAuthAuthorizationError.tokenRequestFailed( + statusCode: httpResponse.statusCode, + oauthError: oauthError + ) + } + guard let httpResponse = response as? HTTPURLResponse, + (200..<300).contains(httpResponse.statusCode) + else { + let statusCode = (response as? HTTPURLResponse)?.statusCode ?? 0 + let oauthError = + (try? JSONDecoder().decode(OAuthTokenErrorResponse.self, from: data))?.error + throw OAuthAuthorizationError.tokenRequestFailed( + statusCode: statusCode, + oauthError: oauthError + ) + } + + let registration = try JSONDecoder().decode(OAuthClientRegistrationResponse.self, from: data) + + let updatedAuth = OAuthClientRegistrar.updatedAuthentication( + from: registration, current: configuration.authentication) + return (response: registration, updatedAuthentication: updatedAuth) + } + + /// Derives the updated token endpoint authentication from a registration response. + /// + /// Updates the client ID (and secret, if issued) while preserving the authentication method. + static func updatedAuthentication( + from registration: OAuthClientRegistrationResponse, + current: OAuthConfiguration.TokenEndpointAuthentication + ) -> OAuthConfiguration.TokenEndpointAuthentication { + switch current { + case .none: + return .none(clientID: registration.clientID) + case .clientSecretBasic(_, let currentSecret): + return .clientSecretBasic( + clientID: registration.clientID, + clientSecret: registration.clientSecret ?? currentSecret + ) + case .clientSecretPost(_, let currentSecret): + return .clientSecretPost( + clientID: registration.clientID, + clientSecret: registration.clientSecret ?? currentSecret + ) + case .privateKeyJWT(_, let factory): + return .privateKeyJWT(clientID: registration.clientID, assertionFactory: factory) + } + } + + // MARK: - CIMD Helpers + + private func isHTTPSURLWithPath(_ value: String) -> Bool { + guard let url = URL(string: value), + let components = URLComponents(url: url, resolvingAgainstBaseURL: false) + else { + return false + } + return components.scheme?.lowercased() == OAuthURLScheme.https + && !components.path.isEmpty + && components.path != "/" + } + + private func clientIDLooksLikeURL(_ value: String) -> Bool { + guard let url = URL(string: value), + let scheme = url.scheme, + !scheme.isEmpty + else { + return false + } + return true + } +} + +extension OAuthClientRegistrar: OAuthClientRegistering {} diff --git a/Sources/MCP/Base/Authorization/OAuthConfiguration.swift b/Sources/MCP/Base/Authorization/OAuthConfiguration.swift new file mode 100644 index 0000000..a1c015f --- /dev/null +++ b/Sources/MCP/Base/Authorization/OAuthConfiguration.swift @@ -0,0 +1,415 @@ +import Foundation + +#if canImport(FoundationNetworking) + import FoundationNetworking +#endif + +#if canImport(CryptoKit) + import CryptoKit +#endif + +/// Configuration for OAuth 2.1 authorization used by ``HTTPClientTransport``. +/// +/// Authorization is optional and disabled by default. Configure this type and pass it +/// to `HTTPClientTransport(oauth:)` to enable automatic Bearer token acquisition for HTTP +/// transports. +/// +/// Supports both `authorization_code` (interactive, browser-based) and `client_credentials` +/// (machine-to-machine) grant types via the ``grantType`` property. +public struct OAuthConfiguration: Sendable { + /// The OAuth 2.1 grant type to use for token acquisition. + public enum GrantType: Sendable { + /// OAuth 2.1 authorization_code flow with PKCE. + case authorizationCode + + /// OAuth 2.1 client_credentials flow. + case clientCredentials + } + + /// How the client authenticates to the OAuth token endpoint. + public enum TokenEndpointAuthentication: Sendable, Equatable { + /// `client_secret_basic` authentication using the Authorization header. + case clientSecretBasic(clientID: String, clientSecret: String) + + /// `client_secret_post` authentication using form parameters. + case clientSecretPost(clientID: String, clientSecret: String) + + /// Public client authentication (`token_endpoint_auth_method=none`). + case none(clientID: String) + + /// `private_key_jwt` authentication. + /// + /// The built-in ``OAuthConfiguration/makePrivateKeyJWTAssertion(clientID:tokenEndpoint:privateKeyPEM:signingAlgorithm:audience:issuedAt:expiresIn:)`` + /// helper generates ES256 (P-256 ECDSA) assertions only. To use other algorithms + /// (e.g., RS256, ES384), provide a custom ``OAuthConfiguration/JWTAssertionFactory`` closure. + case privateKeyJWT(clientID: String, assertionFactory: JWTAssertionFactory) + + var clientID: String { + switch self { + case .clientSecretBasic(let clientID, _), .clientSecretPost(let clientID, _), + .none(let clientID), .privateKeyJWT(let clientID, _): + return clientID + } + } + + /// The token endpoint auth method name per RFC 7591. + var methodName: String { + switch self { + case .clientSecretBasic: return OAuthTokenEndpointAuthMethod.clientSecretBasic + case .clientSecretPost: return OAuthTokenEndpointAuthMethod.clientSecretPost + case .none: return OAuthTokenEndpointAuthMethod.none + case .privateKeyJWT: return OAuthTokenEndpointAuthMethod.privateKeyJWT + } + } + + func apply( + to request: inout URLRequest, + bodyParameters: inout [String: String], + tokenEndpoint: URL + ) async throws { + switch self { + case .clientSecretBasic(let clientID, let clientSecret): + let basic = Data("\(clientID):\(clientSecret)".utf8).base64EncodedString() + request.setValue("Basic \(basic)", forHTTPHeaderField: HTTPHeaderName.authorization) + case .clientSecretPost(let clientID, let clientSecret): + bodyParameters[OAuthParameterName.clientID] = clientID + bodyParameters[OAuthParameterName.clientSecret] = clientSecret + case .none(let clientID): + bodyParameters[OAuthParameterName.clientID] = clientID + case .privateKeyJWT(let clientID, let assertionFactory): + bodyParameters[OAuthParameterName.clientID] = clientID + bodyParameters[OAuthParameterName.clientAssertionType] = + OAuthClientAssertionType.jwtBearer + bodyParameters[OAuthParameterName.clientAssertion] = + try await assertionFactory(tokenEndpoint, clientID) + } + } + } + + /// Closure used to generate a `private_key_jwt` assertion. + public typealias JWTAssertionFactory = @Sendable (_ tokenEndpoint: URL, _ clientID: String) + async throws -> String + + /// Supported signing algorithms for SDK-generated `private_key_jwt` assertions. + public enum PrivateKeyJWTSigningAlgorithm: String, Sendable { + case ES256 + } + + /// Errors thrown while creating SDK-generated `private_key_jwt` assertions. + public enum PrivateKeyJWTAssertionError: LocalizedError, Sendable { + case invalidLifetime(TimeInterval) + case cryptographyUnavailable + + public var errorDescription: String? { + switch self { + case .invalidLifetime(let lifetime): + return "private_key_jwt assertion lifetime must be greater than zero seconds, got \(lifetime)" + case .cryptographyUnavailable: + return "private_key_jwt assertion signing requires CryptoKit support" + } + } + } + + /// Creates a signed `private_key_jwt` client assertion (RFC 7523). + /// + /// Use this helper to build a `JWTAssertionFactory` closure when configuring + /// ``TokenEndpointAuthentication/privateKeyJWT(clientID:assertionFactory:)``: + /// + /// ```swift + /// let factory: OAuthConfiguration.JWTAssertionFactory = { tokenEndpoint, clientID in + /// try OAuthConfiguration.makePrivateKeyJWTAssertion( + /// clientID: clientID, + /// tokenEndpoint: tokenEndpoint, + /// privateKeyPEM: myPEMString + /// ) + /// } + /// ``` + /// + /// The assertion is a compact-serialized JWS signed with the specified private key. + /// Only ES256 (P-256 ECDSA) is supported by this helper; the algorithm requires CryptoKit. + /// To use other algorithms, supply a custom ``JWTAssertionFactory`` closure directly to + /// ``TokenEndpointAuthentication/privateKeyJWT(clientID:assertionFactory:)``. + /// + /// - Parameters: + /// - clientID: The OAuth client identifier, used as both `iss` and `sub` claims. + /// - tokenEndpoint: Token endpoint URL, used as the default `aud` claim. + /// - privateKeyPEM: PEM-encoded EC private key for signing. + /// - signingAlgorithm: The JWS signing algorithm. Currently only `.ES256` is supported. + /// - audience: Explicit `aud` claim override. Defaults to `tokenEndpoint.absoluteString`. + /// - issuedAt: The `iat` claim. Defaults to `Date()`. + /// - expiresIn: Lifetime of the assertion in seconds. Defaults to 300 (5 minutes). Must be > 0. + /// - Returns: A compact-serialized JWT string (`header.payload.signature`). + /// - Throws: ``PrivateKeyJWTAssertionError/invalidLifetime(_:)`` if `expiresIn` ≤ 0, + /// or ``PrivateKeyJWTAssertionError/cryptographyUnavailable`` on platforms without CryptoKit. + public static func makePrivateKeyJWTAssertion( + clientID: String, + tokenEndpoint: URL, + privateKeyPEM: String, + signingAlgorithm: PrivateKeyJWTSigningAlgorithm = .ES256, + audience: String? = nil, + issuedAt: Date = Date(), + expiresIn: TimeInterval = 300 + ) throws -> String { + guard expiresIn > 0 else { + throw PrivateKeyJWTAssertionError.invalidLifetime(expiresIn) + } + + let header = try JSONSerialization.data(withJSONObject: [ + JWTClaimName.algorithm: signingAlgorithm.rawValue, + JWTClaimName.type: JWTClaimName.typeValue, + ]) + + let issuedAtUnix = Int(issuedAt.timeIntervalSince1970) + let lifetimeSeconds = max(1, Int(expiresIn.rounded(.down))) + let payload = try JSONSerialization.data(withJSONObject: [ + JWTClaimName.issuer: clientID, + JWTClaimName.subject: clientID, + JWTClaimName.audience: audience ?? tokenEndpoint.absoluteString, + JWTClaimName.issuedAt: issuedAtUnix, + JWTClaimName.expiration: issuedAtUnix + lifetimeSeconds, + JWTClaimName.jwtID: UUID().uuidString, + ]) + + let signingInput = "\(header.base64URLEncodedString()).\(payload.base64URLEncodedString())" + + #if canImport(CryptoKit) + switch signingAlgorithm { + case .ES256: + let privateKey = try P256.Signing.PrivateKey(pemRepresentation: privateKeyPEM) + let signature = try privateKey.signature(for: Data(signingInput.utf8)).rawRepresentation + return "\(signingInput).\(signature.base64URLEncodedString())" + } + #else + throw PrivateKeyJWTAssertionError.cryptographyUnavailable + #endif + } + + static func defaultAuthorizationRedirectURI() -> URL { + let port = Int.random(in: 49152...65535) + return URL(string: "http://\(OAuthLoopbackHost.ipv4):\(port)/callback")! + } + + // MARK: - Retry Policy + + /// Controls retry behavior for authorization challenges. + public struct RetryPolicy: Sendable { + /// Maximum number of authentication retries for a single MCP request. + public let maxAuthorizationAttempts: Int + + /// Maximum number of scope step-up attempts for a resource and operation. + public let maxScopeUpgradeAttempts: Int + + /// Creates a retry policy. + /// + /// Both values are clamped to a minimum of `1` to prevent infinite loops + /// or zero-retry configurations. + /// + /// - Parameters: + /// - maxAuthorizationAttempts: Maximum authorization retries per request. Defaults to 3. + /// - maxScopeUpgradeAttempts: Maximum scope step-up attempts per resource+operation. Defaults to 2. + public init( + maxAuthorizationAttempts: Int = 3, + maxScopeUpgradeAttempts: Int = 2 + ) { + self.maxAuthorizationAttempts = max(1, maxAuthorizationAttempts) + self.maxScopeUpgradeAttempts = max(1, maxScopeUpgradeAttempts) + } + + public static let `default` = RetryPolicy() + } + + // MARK: - Endpoint Overrides + + /// Optional endpoint overrides for discovery. + public struct EndpointOverrides: Sendable { + /// Optional override for the protected resource metadata URL. + public let protectedResourceMetadataURL: URL? + + /// Optional override for the authorization server issuer URL. + public let authorizationServerURL: URL? + + /// Optional override for the token endpoint. + public let tokenEndpoint: URL? + + /// Optional override for the resource indicator used in token requests. + public let resource: URL? + + public init( + protectedResourceMetadataURL: URL? = nil, + authorizationServerURL: URL? = nil, + tokenEndpoint: URL? = nil, + resource: URL? = nil + ) { + self.protectedResourceMetadataURL = protectedResourceMetadataURL + self.authorizationServerURL = authorizationServerURL + self.tokenEndpoint = tokenEndpoint + self.resource = resource + } + + public static let none = EndpointOverrides() + } + + // MARK: - Access Token Provider + + /// Context supplied to ``AccessTokenProvider`` after SDK discovery is complete. + public struct AccessTokenProviderContext: Sendable { + /// HTTP status code that triggered authorization handling (typically 401 or 403). + public let statusCode: Int + /// Target MCP endpoint URL. + public let endpoint: URL + /// Canonical RFC8707 resource URI for token audience binding. + public let resource: URL + /// Selected authorization server issuer URL, if discovered. + public let authorizationServer: URL? + /// Authorization endpoint from AS metadata, if available. + public let authorizationEndpoint: URL? + /// Token endpoint from AS metadata (or configuration override), if available. + public let tokenEndpoint: URL? + /// Dynamic registration endpoint from AS metadata, if available. + public let registrationEndpoint: URL? + /// `scope` value from the latest challenge header, when present. + public let challengedScope: String? + /// `scopes_supported` from protected resource metadata, when present. + public let scopesSupported: [String]? + /// Scope set selected by SDK for the pending authorization attempt. + public let requestedScopes: Set? + + public init( + statusCode: Int, + endpoint: URL, + resource: URL, + authorizationServer: URL?, + authorizationEndpoint: URL?, + tokenEndpoint: URL?, + registrationEndpoint: URL?, + challengedScope: String?, + scopesSupported: [String]?, + requestedScopes: Set? + ) { + self.statusCode = statusCode + self.endpoint = endpoint + self.resource = resource + self.authorizationServer = authorizationServer + self.authorizationEndpoint = authorizationEndpoint + self.tokenEndpoint = tokenEndpoint + self.registrationEndpoint = registrationEndpoint + self.challengedScope = challengedScope + self.scopesSupported = scopesSupported + self.requestedScopes = requestedScopes + } + } + + /// Optional provider for externally acquired access tokens. + public typealias AccessTokenProvider = @Sendable ( + _ context: AccessTokenProviderContext, + _ session: URLSession + ) async throws -> String? + + // MARK: - Properties + + /// The grant type used for token acquisition. + public let grantType: GrantType + + /// The configured token endpoint authentication method. + public var authentication: TokenEndpointAuthentication + + /// Controls retry behavior for authorization challenges. + public let retryPolicy: RetryPolicy + + /// Optional endpoint overrides for discovery. + public let endpointOverrides: EndpointOverrides + + /// Package-scoped compatibility option for local environments. + package var allowLoopbackHTTPAuthorizationServerEndpoints: Bool + + /// Redirect URI used for authorization requests. + public let authorizationRedirectURI: URL + + /// Additional form fields to include in token requests. + public let additionalTokenRequestParameters: [String: String] + + /// The `client_name` sent during dynamic client registration (RFC 7591). + /// Defaults to `"mcp-swift-sdk"`. Override with your application's name. + public let clientName: String + + /// Optional provider for externally acquired access tokens. + public let accessTokenProvider: AccessTokenProvider? + + /// Optional delegate for browser-based authorization code flows. + public let authorizationDelegate: (any OAuthAuthorizationDelegate)? + + /// How many seconds before token expiry ``OAuthAuthorizer`` proactively refreshes a token + /// when `prepareAuthorization(for:session:)` is called. + /// + /// Set to `0` to disable proactive refresh. Defaults to 60 seconds. + /// Must be greater than ``OAuthTokenExpirySkew/defaultSeconds`` (30 s) to have any effect, + /// since tokens within the default skew window are already treated as expired. + public let proactiveRefreshWindowSeconds: TimeInterval + + /// Creates an OAuth configuration. + /// + /// - Parameters: + /// - grantType: The OAuth 2.1 grant type. Defaults to `.clientCredentials`. + /// - authentication: How the client authenticates to the token endpoint. **Required.** + /// - retryPolicy: Controls how many retries are allowed for authorization challenges. + /// - endpointOverrides: Optional URL overrides that bypass automatic discovery. + /// - authorizationRedirectURI: Redirect URI for the `authorization_code` flow. + /// Defaults to a random loopback URI (`http://127.0.0.1:/callback`). + /// - clientName: The `client_name` sent during dynamic client registration. Defaults to `"mcp-swift-sdk"`. + /// - additionalTokenRequestParameters: Extra form fields appended to every token request. + /// - accessTokenProvider: Optional closure invoked after discovery, allowing the host app + /// to supply an externally acquired token (e.g., from a system credential store). + /// - authorizationDelegate: Optional delegate that presents the authorization URL to the + /// user for interactive `authorization_code` flows. + /// - proactiveRefreshWindowSeconds: Seconds before expiry at which a token is proactively + /// refreshed. Defaults to 60. Set to 0 to disable proactive refresh. + public init( + grantType: GrantType = .clientCredentials, + authentication: TokenEndpointAuthentication, + retryPolicy: RetryPolicy = .default, + endpointOverrides: EndpointOverrides = .none, + authorizationRedirectURI: URL? = nil, + clientName: String = "mcp-swift-sdk", + additionalTokenRequestParameters: [String: String] = [:], + accessTokenProvider: AccessTokenProvider? = nil, + authorizationDelegate: (any OAuthAuthorizationDelegate)? = nil, + proactiveRefreshWindowSeconds: TimeInterval = 60 + ) { + self.grantType = grantType + self.authentication = authentication + self.retryPolicy = retryPolicy + self.endpointOverrides = endpointOverrides + self.allowLoopbackHTTPAuthorizationServerEndpoints = false + self.authorizationRedirectURI = + authorizationRedirectURI ?? Self.defaultAuthorizationRedirectURI() + self.clientName = clientName + self.additionalTokenRequestParameters = additionalTokenRequestParameters + self.accessTokenProvider = accessTokenProvider + self.authorizationDelegate = authorizationDelegate + self.proactiveRefreshWindowSeconds = proactiveRefreshWindowSeconds + } +} + +extension OAuthConfiguration.TokenEndpointAuthentication { + public static func == (lhs: Self, rhs: Self) -> Bool { + switch (lhs, rhs) { + case (.clientSecretBasic(let li, let ls), .clientSecretBasic(let ri, let rs)): + return li == ri && ls == rs + case (.clientSecretPost(let li, let ls), .clientSecretPost(let ri, let rs)): + return li == ri && ls == rs + case (.none(let li), .none(let ri)): + return li == ri + case (.privateKeyJWT(let li, _), .privateKeyJWT(let ri, _)): + return li == ri + default: + return false + } + } +} + +/// Delegate that handles user-facing authorization steps for the authorization_code flow. +public protocol OAuthAuthorizationDelegate: Sendable { + /// Presents the authorization URL to the user and returns the redirect URL containing + /// the authorization code. + func presentAuthorizationURL(_ url: URL) async throws -> URL +} diff --git a/Sources/MCP/Base/Authorization/OAuthConstants.swift b/Sources/MCP/Base/Authorization/OAuthConstants.swift new file mode 100644 index 0000000..73cde0f --- /dev/null +++ b/Sources/MCP/Base/Authorization/OAuthConstants.swift @@ -0,0 +1,139 @@ +import Foundation + +// MARK: - OAuth Grant Types + +enum OAuthGrantTypeValue { + static let clientCredentials = "client_credentials" + static let authorizationCode = "authorization_code" + static let refreshToken = "refresh_token" +} + +// MARK: - OAuth Parameter Names + +enum OAuthParameterName { + static let grantType = "grant_type" + static let resource = "resource" + static let scope = "scope" + static let code = "code" + static let codeVerifier = "code_verifier" + static let codeChallenge = "code_challenge" + static let codeChallengeMethod = "code_challenge_method" + static let redirectURI = "redirect_uri" + static let responseType = "response_type" + static let clientID = "client_id" + static let clientSecret = "client_secret" + static let clientAssertion = "client_assertion" + static let clientAssertionType = "client_assertion_type" + static let refreshToken = "refresh_token" + static let state = "state" +} + +// MARK: - OAuth Well-Known Paths + +enum OAuthWellKnownPath { + static let protectedResource = "/.well-known/oauth-protected-resource" + static let authorizationServer = "/.well-known/oauth-authorization-server" + static let openIDConfiguration = "/.well-known/openid-configuration" +} + +// MARK: - OAuth Token Type + +enum OAuthTokenType { + static let bearer = "Bearer" +} + +// MARK: - OAuth Code Challenge Method + +enum OAuthCodeChallengeMethod { + static let s256 = "S256" +} + +// MARK: - OAuth Token Endpoint Auth Method + +enum OAuthTokenEndpointAuthMethod { + static let clientSecretBasic = "client_secret_basic" + static let clientSecretPost = "client_secret_post" + static let none = "none" + static let privateKeyJWT = "private_key_jwt" +} + +// MARK: - URL Scheme + +enum OAuthURLScheme { + static let http = "http" + static let https = "https" +} + +// MARK: - Default Ports + +enum OAuthDefaultPort { + static let http = 80 + static let https = 443 +} + +// MARK: - Loopback Hosts + +enum OAuthLoopbackHost { + static let localhost = "localhost" + static let ipv4 = "127.0.0.1" + static let ipv6 = "::1" + + static func isLoopback(_ host: String) -> Bool { + host == localhost || host == ipv4 || host == ipv6 + } +} + +// MARK: - Token Expiry Skew + +/// Clock skew tolerance applied when checking token expiry. +/// +/// ``OAuthAccessToken/isExpired(now:skewSeconds:)`` treats a token as expired +/// when `now + skewSeconds >= expiresAt`, giving the client a safety margin to +/// refresh tokens before they actually expire on the server. +public enum OAuthTokenExpirySkew { + /// Default clock skew buffer: 30 seconds. + public static let defaultSeconds: TimeInterval = 30 +} + +// MARK: - JWT Claim Names + +enum JWTClaimName { + static let algorithm = "alg" + static let type = "typ" + static let typeValue = "JWT" + static let issuer = "iss" + static let subject = "sub" + static let audience = "aud" + static let issuedAt = "iat" + static let expiration = "exp" + static let jwtID = "jti" +} + +// MARK: - Client Assertion Type + +enum OAuthClientAssertionType { + static let jwtBearer = "urn:ietf:params:oauth:client-assertion-type:jwt-bearer" +} + +// MARK: - HTTPHeaderName Extensions + +extension HTTPHeaderName { + static let location = "Location" +} + +// MARK: - ContentType Extensions + +extension ContentType { + static let formURLEncoded = "application/x-www-form-urlencoded" +} + +// MARK: - Data Base64URL Extension + +extension Data { + func base64URLEncodedString() -> String { + base64EncodedString() + .replacingOccurrences(of: "+", with: "-") + .replacingOccurrences(of: "/", with: "_") + .replacingOccurrences(of: "=", with: "") + } +} diff --git a/Sources/MCP/Base/Authorization/OAuthDiscovery.swift b/Sources/MCP/Base/Authorization/OAuthDiscovery.swift new file mode 100644 index 0000000..332a898 --- /dev/null +++ b/Sources/MCP/Base/Authorization/OAuthDiscovery.swift @@ -0,0 +1,361 @@ +import Foundation + +#if canImport(FoundationNetworking) + import FoundationNetworking +#endif + +// MARK: - Scope Selection Protocol + +/// Determines which OAuth scopes to request during token acquisition. +/// +/// ``OAuthAuthorizer`` uses this protocol to translate raw scope strings from +/// `WWW-Authenticate` challenges and Protected Resource Metadata into the +/// scope set passed to the token endpoint. +/// +/// Override the default ``DefaultOAuthScopeSelector`` to apply custom scope +/// filtering or transformation logic. +public protocol OAuthScopeSelecting { + /// Selects the scope set to request for a token. + /// + /// Priority order (highest first): + /// 1. `challengeScope` — the `scope` parameter from the `WWW-Authenticate` header. + /// 2. `scopesSupported` — the `scopes_supported` array from Protected Resource Metadata. + /// 3. `nil` — no scope restriction. + /// + /// - Parameters: + /// - challengeScope: Space-separated scope string from the Bearer challenge, or `nil`. + /// - scopesSupported: Array of supported scopes from the resource metadata, or `nil`. + /// - Returns: The set of scopes to request, or `nil` to omit the `scope` parameter entirely. + func selectScopes(challengeScope: String?, scopesSupported: [String]?) -> Set? + + /// Parses a space-separated OAuth scope string into individual scope tokens. + /// - Parameter scope: A scope string such as `"read write"`. + /// - Returns: A set of individual scope strings with whitespace-only tokens removed. + func parseScopeString(_ scope: String) -> Set + + /// Serializes a set of scopes into a space-separated string suitable for the `scope` parameter. + /// - Parameter scopes: The scope set to serialize. + /// - Returns: A sorted, space-separated string, or `nil` if the set is empty. + func serialize(_ scopes: Set) -> String? +} + +// MARK: - Default Scope Selector + +/// Default ``OAuthScopeSelecting`` implementation. +/// +/// Selects scopes in priority order: challenge scope > `scopes_supported` > `nil`. +/// Serializes scopes sorted alphabetically to produce deterministic `scope` parameters. +public struct DefaultOAuthScopeSelector: OAuthScopeSelecting { + public init() {} + + public func selectScopes(challengeScope: String?, scopesSupported: [String]?) -> Set? { + if let challengeScope { + let parsed = parseScopeString(challengeScope) + return parsed.isEmpty ? nil : parsed + } + + if let scopesSupported { + let parsed = Set(scopesSupported.filter { !$0.trimmingCharacters(in: .whitespaces).isEmpty }) + return parsed.isEmpty ? nil : parsed + } + + return nil + } + + public func parseScopeString(_ scope: String) -> Set { + Set( + scope + .split(whereSeparator: { $0.isWhitespace }) + .map(String.init) + .filter { !$0.isEmpty } + ) + } + + public func serialize(_ scopes: Set) -> String? { + guard !scopes.isEmpty else { return nil } + return scopes.sorted().joined(separator: " ") + } +} + +// MARK: - Metadata Discovery Protocol + +/// Builds discovery URLs and normalizes resource identifiers for OAuth metadata discovery. +/// +/// ``OAuthAuthorizer`` uses this protocol to construct the candidate URL list for both +/// Protected Resource Metadata (RFC 9728) and Authorization Server Metadata (RFC 8414 / OIDC) +/// discovery, and to normalize resource URIs for RFC 8707 audience binding. +/// +/// Override the default ``DefaultOAuthMetadataDiscovery`` to customise discovery URL +/// construction or resource-matching logic. +public protocol OAuthMetadataDiscovering: Sendable { + /// Returns candidate URLs for Protected Resource Metadata, ordered by priority. + /// + /// ``OAuthAuthorizer`` tries each URL in order and uses the first successful response. + /// + /// - Parameter endpoint: The MCP endpoint URL. + /// - Returns: An ordered list of discovery URLs (typically `/.well-known/oauth-protected-resource` + /// variants). + func protectedResourceMetadataURLs(for endpoint: URL) -> [URL] + + /// Returns candidate URLs for Authorization Server Metadata, ordered by priority. + /// + /// Covers RFC 8414 (`/.well-known/oauth-authorization-server`) and + /// OIDC Discovery 1.0 (`/.well-known/openid-configuration`), including + /// path-inserted variants for issuers with non-root paths. + /// + /// - Parameter issuer: The authorization server issuer URL. + /// - Returns: An ordered list of metadata discovery URLs. + func authorizationServerMetadataURLs(for issuer: URL) -> [URL] + + /// Derives the canonical RFC 8707 resource URI from an endpoint URL. + /// + /// The canonical form strips the query string, fragment, and trailing slash + /// while preserving the scheme, host, port, and path. + /// + /// - Parameter endpoint: The MCP endpoint URL. + /// - Returns: The canonical resource URI. + /// - Throws: ``OAuthAuthorizationError/invalidResourceURI(_:)`` if the URL does not + /// satisfy the HTTPS-or-loopback-HTTP requirement. + func canonicalResourceURI(from endpoint: URL) throws -> URL + + /// Derives a fallback authorization server issuer URL from an endpoint URL. + /// + /// Used when no authorization server is listed in Protected Resource Metadata. + /// Typically returns the scheme+host+port of the endpoint with an empty path. + /// + /// - Parameter endpoint: The MCP endpoint URL. + /// - Returns: A candidate issuer URL. + /// - Throws: ``OAuthAuthorizationError/invalidResourceURI(_:)`` if the URL is invalid. + func authorizationServerFallbackIssuer(from endpoint: URL) throws -> URL + + /// Returns `true` if `resource` is a prefix of `endpoint` in the URL hierarchy. + /// + /// A resource matches an endpoint when the two share the same scheme, host, and port, + /// and the endpoint path starts with the resource path. + /// + /// - Parameters: + /// - resource: The canonical resource URI from Protected Resource Metadata. + /// - endpoint: The canonical endpoint URI being requested. + /// - Returns: `true` if the endpoint falls within the resource's scope. + func protectedResourceMatches(resource: URL, endpoint: URL) -> Bool +} + +// MARK: - Default Metadata Discovery + +/// Default ``OAuthMetadataDiscovering`` implementation following RFC 9728 and RFC 8414. +/// +/// - Builds `/.well-known/oauth-protected-resource` URLs with and without the endpoint path suffix. +/// - Builds RFC 8414 and OIDC discovery URLs with path-inserted and path-appended variants. +/// - Canonicalises resource URIs by stripping query, fragment, and trailing slash. +/// - Matches resources using scheme/host/port equality and path prefix rules. +public struct DefaultOAuthMetadataDiscovery: OAuthMetadataDiscovering { + public init() {} + + public func protectedResourceMetadataURLs(for endpoint: URL) -> [URL] { + guard var components = URLComponents(url: endpoint, resolvingAgainstBaseURL: false), + let scheme = components.scheme?.lowercased(), + let host = components.host?.lowercased(), + Self.isSecureOAuthScheme(scheme: scheme, host: host) + else { + return [] + } + + components.query = nil + components.fragment = nil + + let endpointPath = components.path + let normalizedPath = endpointPath == "/" ? "" : endpointPath + + var urls: [URL] = [] + + var pathSpecific = components + pathSpecific.path = "\(OAuthWellKnownPath.protectedResource)\(normalizedPath)" + if let url = pathSpecific.url { + urls.append(url) + } + + var root = components + root.path = OAuthWellKnownPath.protectedResource + if let url = root.url { + urls.append(url) + } + + return urls + } + + public func authorizationServerMetadataURLs(for issuer: URL) -> [URL] { + guard var components = URLComponents(url: issuer, resolvingAgainstBaseURL: false), + let scheme = components.scheme?.lowercased(), + let host = components.host?.lowercased(), + Self.isSecureOAuthScheme(scheme: scheme, host: host) + else { + return [] + } + + components.query = nil + components.fragment = nil + + let path = components.path.trimmingCharacters(in: CharacterSet(charactersIn: "/")) + let hasPath = !path.isEmpty + + var urls: [URL] = [] + + if hasPath { + var oauthInserted = components + oauthInserted.path = "\(OAuthWellKnownPath.authorizationServer)/\(path)" + if let url = oauthInserted.url { + urls.append(url) + } + + var oidcInserted = components + oidcInserted.path = "\(OAuthWellKnownPath.openIDConfiguration)/\(path)" + if let url = oidcInserted.url { + urls.append(url) + } + + var oidcAppended = components + oidcAppended.path = "/\(path)\(OAuthWellKnownPath.openIDConfiguration)" + if let url = oidcAppended.url { + urls.append(url) + } + } else { + var oauth = components + oauth.path = OAuthWellKnownPath.authorizationServer + if let url = oauth.url { + urls.append(url) + } + + var oidc = components + oidc.path = OAuthWellKnownPath.openIDConfiguration + if let url = oidc.url { + urls.append(url) + } + } + + return urls + } + + public func canonicalResourceURI(from endpoint: URL) throws -> URL { + guard var components = URLComponents(url: endpoint, resolvingAgainstBaseURL: false), + let scheme = components.scheme?.lowercased(), + let host = components.host?.lowercased(), + Self.isSecureOAuthScheme(scheme: scheme, host: host) + else { + throw OAuthAuthorizationError.invalidResourceURI( + "Resource URI must use https or loopback http" + ) + } + + if components.fragment != nil { + throw OAuthAuthorizationError.invalidResourceURI("Resource URI must not contain a fragment") + } + + components.scheme = scheme + components.host = host + components.query = nil + components.fragment = nil + + if components.path == "/" { + components.path = "" + } + + guard let url = components.url else { + throw OAuthAuthorizationError.invalidResourceURI("Failed to normalize resource URI") + } + + return url + } + + public func authorizationServerFallbackIssuer(from endpoint: URL) throws -> URL { + guard var components = URLComponents(url: endpoint, resolvingAgainstBaseURL: false), + let scheme = components.scheme?.lowercased(), + let host = components.host?.lowercased(), + Self.isSecureOAuthScheme(scheme: scheme, host: host) + else { + throw OAuthAuthorizationError.invalidResourceURI( + "Resource URI must use https or loopback http" + ) + } + + components.scheme = scheme + components.host = host + components.path = "" + components.query = nil + components.fragment = nil + + guard let url = components.url else { + throw OAuthAuthorizationError.invalidResourceURI("Failed to derive issuer URI") + } + return url + } + + public func protectedResourceMatches(resource: URL, endpoint: URL) -> Bool { + guard let resourceComponents = URLComponents( + url: resource, + resolvingAgainstBaseURL: false + ), + let endpointComponents = URLComponents( + url: endpoint, + resolvingAgainstBaseURL: false + ) + else { + return false + } + + let resourceScheme = resourceComponents.scheme?.lowercased() + let endpointScheme = endpointComponents.scheme?.lowercased() + let resourceHost = resourceComponents.host?.lowercased() + let endpointHost = endpointComponents.host?.lowercased() + let resourcePort = resourceComponents.port ?? Self.defaultPort(for: resourceScheme) + let endpointPort = endpointComponents.port ?? Self.defaultPort(for: endpointScheme) + + guard resourceScheme == endpointScheme, + resourceHost == endpointHost, + resourcePort == endpointPort + else { + return false + } + + let resourcePath = Self.normalizedResourcePath(resourceComponents.path) + let endpointPath = Self.normalizedResourcePath(endpointComponents.path) + if resourcePath.isEmpty { + return true + } + if endpointPath == resourcePath { + return true + } + return endpointPath.hasPrefix(resourcePath + "/") + } + + private static func normalizedResourcePath(_ rawPath: String) -> String { + if rawPath.isEmpty || rawPath == "/" { + return "" + } + if rawPath.count > 1 && rawPath.hasSuffix("/") { + return String(rawPath.dropLast()) + } + return rawPath + } + + private static func defaultPort(for scheme: String?) -> Int? { + switch scheme?.lowercased() { + case OAuthURLScheme.http: + return OAuthDefaultPort.http + case OAuthURLScheme.https: + return OAuthDefaultPort.https + default: + return nil + } + } + + static func isSecureOAuthScheme(scheme: String, host: String) -> Bool { + if scheme == OAuthURLScheme.https { + return true + } + if scheme == OAuthURLScheme.http { + return OAuthLoopbackHost.isLoopback(host) + } + return false + } +} + diff --git a/Sources/MCP/Base/Authorization/OAuthDiscoveryClient.swift b/Sources/MCP/Base/Authorization/OAuthDiscoveryClient.swift new file mode 100644 index 0000000..8795ed7 --- /dev/null +++ b/Sources/MCP/Base/Authorization/OAuthDiscoveryClient.swift @@ -0,0 +1,142 @@ +import Foundation + +#if canImport(FoundationNetworking) + import FoundationNetworking +#endif + +/// Internal protocol for fetching OAuth discovery metadata. +protocol OAuthDiscoveryFetching: Sendable { + var metadataDiscovery: any OAuthMetadataDiscovering { get } + func fetchProtectedResourceMetadata(candidates: [URL], fallbackIssuer: URL?, session: URLSession) async throws -> OAuthProtectedResourceMetadata + func fetchAuthorizationServerMetadata(candidates: [URL], session: URLSession) async throws -> (server: URL, metadata: OAuthAuthorizationServerMetadata) +} + +/// Stateless OAuth metadata fetcher. +/// +/// Fetches Protected Resource Metadata (RFC 9728) and Authorization Server Metadata +/// (RFC 8414 / OIDC Discovery 1.0) from ordered candidate URL lists. +/// Cache management is the caller's responsibility. +struct OAuthDiscoveryClient: Sendable { + let metadataDiscovery: any OAuthMetadataDiscovering + let urlValidator: OAuthURLValidator + + init( + metadataDiscovery: any OAuthMetadataDiscovering, + urlValidator: OAuthURLValidator + ) { + self.metadataDiscovery = metadataDiscovery + self.urlValidator = urlValidator + } + + /// Fetches Protected Resource Metadata from the first candidate that returns a valid response. + /// + /// If all candidates fail and `fallbackIssuer` is provided, returns synthetic metadata + /// using that issuer as the authorization server — for servers that do not expose a + /// PRM document at any well-known path. + func fetchProtectedResourceMetadata( + candidates: [URL], + fallbackIssuer: URL?, + session: URLSession + ) async throws -> OAuthProtectedResourceMetadata { + let decoder = JSONDecoder() + for url in candidates { + var request = URLRequest(url: url) + request.httpMethod = "GET" + request.setValue(ContentType.json, forHTTPHeaderField: HTTPHeaderName.accept) + + do { + let (data, response) = try await session.data(for: request) + guard let httpResponse = response as? HTTPURLResponse, + (200..<300).contains(httpResponse.statusCode) + else { + continue + } + + let metadata = try decoder.decode(OAuthProtectedResourceMetadata.self, from: data) + guard !metadata.authorizationServers.isEmpty else { continue } + return metadata + } catch let error as OAuthAuthorizationError { + throw error + } catch { + continue + } + } + if let fallbackIssuer { + return OAuthProtectedResourceMetadata( + resource: nil, + authorizationServers: [fallbackIssuer], + scopesSupported: nil + ) + } + throw OAuthAuthorizationError.metadataDiscoveryFailed + } + + /// Fetches Authorization Server Metadata from candidates, preferring a response whose + /// `issuer` matches the candidate URL (RFC 8414 §3). If no candidate produces a matching + /// issuer, the first valid response is accepted and its own `issuer` is used as the + /// server identity — accommodating servers that serve metadata at one host but advertise + /// a different issuer. + func fetchAuthorizationServerMetadata( + candidates: [URL], + session: URLSession + ) async throws -> (server: URL, metadata: OAuthAuthorizationServerMetadata) { + let decoder = JSONDecoder() + var firstValid: (server: URL, metadata: OAuthAuthorizationServerMetadata)? + + for candidateServer in candidates { + guard (try? urlValidator.validateAuthorizationServer( + candidateServer, context: "Authorization server issuer")) != nil + else { + continue + } + if let host = URLComponents(url: candidateServer, resolvingAgainstBaseURL: false)? + .host?.lowercased(), urlValidator.isPrivateIPHost(host) + { + continue + } + + for metadataURL in metadataDiscovery.authorizationServerMetadataURLs( + for: candidateServer) + { + var request = URLRequest(url: metadataURL) + request.httpMethod = "GET" + request.setValue(ContentType.json, forHTTPHeaderField: HTTPHeaderName.accept) + + do { + let (data, response) = try await session.data(for: request) + guard let httpResponse = response as? HTTPURLResponse, + (200..<300).contains(httpResponse.statusCode) + else { + continue + } + + let asMetadata = try decoder.decode( + OAuthAuthorizationServerMetadata.self, from: data) + + // Prefer metadata whose issuer matches the candidate URL (RFC 8414 §3). + let issuerMatches = + asMetadata.issuer == nil + || asMetadata.issuer?.absoluteString.lowercased() + == candidateServer.absoluteString.lowercased() + if issuerMatches { + return (server: candidateServer, metadata: asMetadata) + } + // Keep as fallback in case no issuer-matching candidate is found. + if firstValid == nil { + let server = asMetadata.issuer ?? candidateServer + firstValid = (server: server, metadata: asMetadata) + } + } catch { + continue + } + } + } + + if let firstValid { + return firstValid + } + throw OAuthAuthorizationError.authorizationServerMetadataDiscoveryFailed + } +} + +extension OAuthDiscoveryClient: OAuthDiscoveryFetching {} diff --git a/Sources/MCP/Base/Authorization/OAuthErrors.swift b/Sources/MCP/Base/Authorization/OAuthErrors.swift new file mode 100644 index 0000000..2f8249c --- /dev/null +++ b/Sources/MCP/Base/Authorization/OAuthErrors.swift @@ -0,0 +1,203 @@ +import Foundation + +/// Errors thrown during OAuth 2.1 authorization by ``OAuthAuthorizer``. +/// +/// These errors surface when the authorizer is unable to complete the authorization flow, +/// either due to discovery failures, token exchange problems, security policy violations, +/// or authorization code flow issues. +public enum OAuthAuthorizationError: LocalizedError { + /// No authorization server URL could be found in the Protected Resource Metadata. + case missingAuthorizationServer + + /// All Protected Resource Metadata discovery candidates returned errors or invalid documents. + case metadataDiscoveryFailed + + /// All Authorization Server Metadata discovery candidates (RFC 8414 / OIDC) returned errors. + case authorizationServerMetadataDiscoveryFailed + + /// The authorization server metadata does not include a `token_endpoint`. + case tokenEndpointMissing + + /// The token endpoint returned a non-2xx HTTP response. + /// + /// - Parameters: + /// - statusCode: HTTP status code from the token endpoint. + /// - oauthError: The `error` field from the OAuth error response body, if present. + case tokenRequestFailed(statusCode: Int, oauthError: String?) + + /// The token response body is missing `access_token`, has an empty token, or specifies + /// a non-Bearer token type. + case tokenResponseInvalid + + /// A URL that must be a valid resource identifier (RFC 8707) failed validation. + case invalidResourceURI(String) + + /// The configured client ID looks like a URL but is not a valid HTTPS URL with a path, + /// which is required for Client ID Metadata Documents. + case invalidClientIDMetadataURL(String) + + /// The `resource` field in the Protected Resource Metadata does not match the requested endpoint. + /// + /// - Parameters: + /// - expected: The canonical URI derived from the endpoint. + /// - actual: The canonical URI derived from the metadata's `resource` field. + case protectedResourceMismatch(expected: String, actual: String) + + /// The authorization server does not support dynamic registration and no pre-registered + /// credentials were supplied. + case registrationInformationRequired + + /// The authorization server does not support Client ID Metadata Documents and no dynamic + /// registration endpoint is available. + /// + /// - Parameter clientID: The client ID URL that was provided as a CIMD URL. + case cimdNotSupported(clientID: String) + + /// A URL received from a discovery response resolves to a private or reserved IP address, + /// which is blocked to prevent SSRF attacks. + /// + /// - Parameters: + /// - context: Human-readable label identifying which URL was blocked. + /// - url: The blocked URL string. + case privateIPAddressBlocked(context: String, url: String) + + /// An endpoint URL used during the OAuth flow does not satisfy the HTTPS-or-loopback requirement. + /// + /// - Parameters: + /// - context: Human-readable label identifying which endpoint failed (e.g., `"Token endpoint"`). + /// - url: The offending URL string. + case insecureOAuthEndpoint(context: String, url: String) + + /// An authorization server endpoint does not satisfy the HTTPS-only requirement. + /// + /// - Parameters: + /// - context: Human-readable label identifying which endpoint failed. + /// - url: The offending URL string. + case insecureAuthorizationServerEndpoint(context: String, url: String) + + /// The redirect URI supplied for the `authorization_code` flow is not a valid HTTPS or + /// loopback HTTP URI, or it contains a fragment. + case invalidRedirectURI(String) + + /// The authorization endpoint returned an HTTP error response during the authorization code flow. + /// + /// - Parameter statusCode: HTTP status code from the authorization endpoint. + case authorizationRequestFailed(statusCode: Int) + + /// The authorization response did not include a `Location` redirect header. + case authorizationResponseMissingRedirectLocation + + /// The redirect URI in the authorization response does not match the expected redirect URI. + /// + /// - Parameters: + /// - expected: The redirect URI supplied in the authorization request. + /// - actual: The redirect URI received in the authorization response. + case authorizationResponseRedirectMismatch(expected: String, actual: String) + + /// The authorization response redirect URL is missing the `state` parameter. + case authorizationResponseMissingState + + /// The `state` in the authorization response does not match the one sent in the request. + /// + /// This may indicate a CSRF attack. + /// + /// - Parameters: + /// - expected: The `state` value sent in the authorization request. + /// - actual: The `state` value received in the authorization response. + case authorizationResponseStateMismatch(expected: String, actual: String) + + /// The authorization response redirect URL is missing the `code` parameter. + case authorizationResponseMissingCode + + /// The authorization server metadata does not include `code_challenge_methods_supported`, + /// which is required for PKCE (RFC 7636). + case pkceCodeChallengeMethodsMissing + + /// The authorization server does not advertise `S256` in `code_challenge_methods_supported`. + /// + /// The MCP specification mandates S256; plain PKCE is not accepted. + /// + /// - Parameter advertisedMethods: The methods listed in the server metadata. + case pkceS256NotSupported(advertisedMethods: [String]) + + /// PKCE S256 challenge generation is unavailable because CryptoKit is not present on this platform. + case pkceS256Unavailable + + /// The `issuer` field in the Authorization Server Metadata does not match the URL used to + /// resolve the metadata document, as required by RFC 8414 §3. + /// + /// - Parameters: + /// - expected: The issuer URL derived from the discovery candidate. + /// - actual: The `issuer` field value found in the metadata document. + case authorizationServerIssuerMismatch(expected: String, actual: String) + + public var errorDescription: String? { + switch self { + case .missingAuthorizationServer: + return "No authorization server was found in protected resource metadata" + case .metadataDiscoveryFailed: + return "Failed to discover protected resource metadata" + case .authorizationServerMetadataDiscoveryFailed: + return "Failed to discover authorization server metadata" + case .tokenEndpointMissing: + return "Authorization server metadata is missing token_endpoint" + case .tokenRequestFailed(let statusCode, let oauthError): + if let oauthError, !oauthError.isEmpty { + return "Token request failed with status \(statusCode) (oauth_error: \(oauthError))" + } + return "Token request failed with status \(statusCode)" + case .tokenResponseInvalid: + return "Token response is invalid" + case .invalidResourceURI(let detail): + return "Invalid resource URI: \(detail)" + case .invalidClientIDMetadataURL(let value): + return + "Client ID metadata document URL must use https and include a path: \(value)" + case .protectedResourceMismatch(let expected, let actual): + return + "Protected resource metadata resource mismatch. Expected \(expected), got \(actual)" + case .registrationInformationRequired: + return + "No supported client registration mechanism was available; provide pre-registered client credentials" + case .cimdNotSupported(let clientID): + return + "Authorization server does not support Client ID Metadata Documents; configure pre-registered credentials or ensure the server advertises client_id_metadata_document_supported: \(clientID)" + case .privateIPAddressBlocked(let context, let url): + return + "\(context) resolves to a private or reserved IP address which is blocked for SSRF protection: \(url)" + case .insecureOAuthEndpoint(let context, let url): + return "\(context) must use https or loopback http: \(url)" + case .insecureAuthorizationServerEndpoint(let context, let url): + return "\(context) must use https: \(url)" + case .invalidRedirectURI(let url): + return + "Redirect URI must use https or loopback http and must not include fragments: \(url)" + case .authorizationRequestFailed(let statusCode): + return "Authorization request failed with status \(statusCode)" + case .authorizationResponseMissingRedirectLocation: + return "Authorization response is missing redirect location" + case .authorizationResponseRedirectMismatch(let expected, let actual): + return + "Authorization response redirect URI mismatch. Expected \(expected), got \(actual)" + case .authorizationResponseMissingState: + return "Authorization response is missing state" + case .authorizationResponseStateMismatch(let expected, let actual): + return "Authorization response state mismatch. Expected \(expected), got \(actual)" + case .authorizationResponseMissingCode: + return "Authorization response is missing the authorization code" + case .pkceCodeChallengeMethodsMissing: + return + "Authorization server metadata must include code_challenge_methods_supported for PKCE" + case .pkceS256NotSupported(let advertisedMethods): + let methods = advertisedMethods.joined(separator: ", ") + return + "Authorization server metadata must support PKCE S256 (advertised: \(methods))" + case .pkceS256Unavailable: + return + "PKCE S256 code challenge generation is unavailable on this platform" + case .authorizationServerIssuerMismatch(let expected, let actual): + return + "Authorization server issuer mismatch. Expected \(expected), got \(actual)" + } + } +} diff --git a/Sources/MCP/Base/Authorization/OAuthModels.swift b/Sources/MCP/Base/Authorization/OAuthModels.swift new file mode 100644 index 0000000..28043b9 --- /dev/null +++ b/Sources/MCP/Base/Authorization/OAuthModels.swift @@ -0,0 +1,257 @@ +import Foundation + +#if canImport(FoundationNetworking) + import FoundationNetworking +#endif + +struct HTTPAuthenticationChallengeError: Error { + let statusCode: Int + let headers: [String: String] +} + +/// Parsed representation of a `WWW-Authenticate: Bearer` challenge header. +/// +/// Servers return this challenge in `401 Unauthorized` and `403 Forbidden` responses +/// to indicate that a Bearer token is required or that the presented token lacks +/// sufficient scope. +/// +/// The ``OAuthWWWAuthenticateParsing`` protocol produces instances of this type. +public struct OAuthBearerChallenge: Sendable { + /// Raw key-value parameters extracted from the `Bearer` challenge. + public let parameters: [String: String] + + /// Creates a challenge from raw parsed parameters. + /// - Parameter parameters: Key-value pairs from the `WWW-Authenticate` header. + public init(parameters: [String: String]) { + self.parameters = parameters + } + + /// The `resource_metadata` parameter, parsed as a URL. + /// + /// Points to the server's RFC 9728 Protected Resource Metadata document. + /// When present, ``OAuthAuthorizer`` uses this URL as the highest-priority + /// discovery candidate. + public var resourceMetadataURL: URL? { + guard let value = parameters["resource_metadata"] else { return nil } + return URL(string: value) + } + + /// The `scope` parameter from the Bearer challenge. + /// + /// Specifies the scopes required or recommended for this resource. + /// ``OAuthAuthorizer`` uses this value as the highest-priority scope hint. + public var scope: String? { + parameters["scope"] + } + + /// The `error` parameter from the Bearer challenge (e.g., `"invalid_token"`, `"insufficient_scope"`). + public var error: String? { + parameters["error"] + } + + /// The `error_description` parameter from the Bearer challenge. + public var errorDescription: String? { + parameters["error_description"] + } +} + +/// RFC9728 OAuth Protected Resource metadata (client-side, decode-only). +struct OAuthProtectedResourceMetadata: Decodable, Sendable, Equatable { + let resource: String? + let authorizationServers: [URL] + let scopesSupported: [String]? + + enum CodingKeys: String, CodingKey { + case resource + case authorizationServers = "authorization_servers" + case scopesSupported = "scopes_supported" + } +} + +/// RFC8414/OIDC authorization server metadata. +struct OAuthAuthorizationServerMetadata: Decodable, Sendable, Equatable { + let issuer: URL? + let authorizationEndpoint: URL? + let tokenEndpoint: URL? + let registrationEndpoint: URL? + let codeChallengeMethodsSupported: [String]? + let tokenEndpointAuthMethodsSupported: [String]? + let clientIDMetadataDocumentSupported: Bool? + + enum CodingKeys: String, CodingKey { + case issuer + case authorizationEndpoint = "authorization_endpoint" + case tokenEndpoint = "token_endpoint" + case registrationEndpoint = "registration_endpoint" + case codeChallengeMethodsSupported = "code_challenge_methods_supported" + case tokenEndpointAuthMethodsSupported = "token_endpoint_auth_methods_supported" + case clientIDMetadataDocumentSupported = "client_id_metadata_document_supported" + } +} + +struct OAuthTokenResponse: Decodable, Sendable, Equatable { + let accessToken: String + let tokenType: String + let expiresIn: Int? + let scope: String? + let refreshToken: String? + + enum CodingKeys: String, CodingKey { + case accessToken = "access_token" + case tokenType = "token_type" + case expiresIn = "expires_in" + case scope + case refreshToken = "refresh_token" + } +} + +struct OAuthTokenErrorResponse: Decodable { + let error: String +} + +/// An OAuth 2.1 access token and its associated metadata. +/// +/// Stored by ``TokenStorage`` and produced by ``OAuthAuthorizer`` after a successful +/// token request. Use ``isExpired(now:skewSeconds:)`` to check validity before use. +public struct OAuthAccessToken: Sendable, Codable { + /// The raw bearer token string for use in the `Authorization` header. + public let value: String + + /// The token type returned by the authorization server (should be `"Bearer"`). + public let tokenType: String + + /// The UTC date after which the token is considered expired, or `nil` if no expiry was specified. + public let expiresAt: Date? + + /// The set of OAuth scopes granted with this token. + public let scopes: Set + + /// The issuer URL of the authorization server that issued this token. + /// + /// Used to detect when the active authorization server changes between requests, + /// triggering a token invalidation. + public let authorizationServer: URL? + + /// The refresh token, if the authorization server issued one alongside the access token. + public let refreshToken: String? + + /// The OAuth `client_id` this token was issued to, if the authorizer had one at storage time. + /// + /// Captures the `client_id` held by the authorizer's ``OAuthConfiguration/authentication`` when + /// the token was saved — including identifiers assigned by Dynamic Client Registration + /// (RFC 7591). `nil` when no `client_id` was configured (for example, the placeholder used + /// before registration completes). + public let clientID: String? + + /// Creates a new access token record. + public init( + value: String, + tokenType: String, + expiresAt: Date?, + scopes: Set, + authorizationServer: URL?, + refreshToken: String?, + clientID: String? = nil + ) { + self.value = value + self.tokenType = tokenType + self.expiresAt = expiresAt + self.scopes = scopes + self.authorizationServer = authorizationServer + self.refreshToken = refreshToken + self.clientID = clientID + } + + /// Returns `true` if the token has expired or will expire within the skew window. + /// + /// - Parameters: + /// - now: The reference time to compare against. Defaults to `Date()`. + /// - skewSeconds: Clock skew buffer in seconds. Defaults to ``OAuthTokenExpirySkew/defaultSeconds`` (30 s). + /// Tokens are considered expired when `now + skewSeconds >= expiresAt`. + /// - Returns: `false` if `expiresAt` is `nil` (no expiry). + public func isExpired(now: Date = Date(), skewSeconds: TimeInterval = OAuthTokenExpirySkew.defaultSeconds) -> Bool { + guard let expiresAt else { return false } + return now.addingTimeInterval(skewSeconds) >= expiresAt + } +} + +/// Token introspection result passed from the caller's token validator to ``BearerTokenValidator``. +/// +/// The validator uses this to enforce expiry and audience checks before allowing a request through. +/// Produce an instance in your ``BearerTokenValidator/TokenValidator`` closure after verifying +/// the token's signature and extracting its claims. +public struct BearerTokenInfo: Sendable, Equatable { + /// Audience values from the token (`aud` JWT claim or introspection response). + /// + /// `nil` indicates an opaque token whose audience cannot be inspected; + /// ``BearerTokenValidator`` skips the audience check in that case. + public let audience: [String]? + + /// Scopes granted by the token. + public let scopes: Set? + + /// UTC date after which the token is considered expired. + /// + /// `nil` means no expiry information is available and the expiry check is skipped. + public let expiresAt: Date? + + public init( + audience: [String]? = nil, + scopes: Set? = nil, + expiresAt: Date? = nil + ) { + self.audience = audience + self.scopes = scopes + self.expiresAt = expiresAt + } +} + +struct OAuthClientRegistrationResponse: Decodable { + let clientID: String + let clientSecret: String? + let tokenEndpointAuthMethod: String? + /// Unix timestamp after which `clientSecret` is no longer valid, per RFC 7591 §3.2. + /// A value of `0` means the secret does not expire. + let clientSecretExpiresAt: Int? + + enum CodingKeys: String, CodingKey { + case clientID = "client_id" + case clientSecret = "client_secret" + case tokenEndpointAuthMethod = "token_endpoint_auth_method" + case clientSecretExpiresAt = "client_secret_expires_at" + } +} + +/// Server-side encodable RFC 9728 Protected Resource Metadata. +/// +/// Use this type to construct the metadata document that MCP servers **MUST** serve +/// at `/.well-known/oauth-protected-resource` per the MCP authorization specification. +/// +/// Pair with ``ProtectedResourceMetadataValidator`` to automatically serve this +/// document in the server's validation pipeline. +public struct OAuthProtectedResourceServerMetadata: Codable, Sendable { + /// The canonical resource identifier (RFC 8707). + public let resource: String + + /// One or more authorization server URLs that protect this resource. + public let authorizationServers: [URL] + + /// The scopes supported by this resource server. + public let scopesSupported: [String]? + + enum CodingKeys: String, CodingKey { + case resource + case authorizationServers = "authorization_servers" + case scopesSupported = "scopes_supported" + } + + public init( + resource: String, + authorizationServers: [URL], + scopesSupported: [String]? = nil + ) { + self.resource = resource + self.authorizationServers = authorizationServers + self.scopesSupported = scopesSupported + } +} diff --git a/Sources/MCP/Base/Authorization/OAuthTokenEndpointClient.swift b/Sources/MCP/Base/Authorization/OAuthTokenEndpointClient.swift new file mode 100644 index 0000000..3c61600 --- /dev/null +++ b/Sources/MCP/Base/Authorization/OAuthTokenEndpointClient.swift @@ -0,0 +1,90 @@ +import Foundation + +#if canImport(FoundationNetworking) + import FoundationNetworking +#endif + +/// Internal protocol for making OAuth token endpoint requests. +protocol OAuthTokenRequesting: Sendable { + func request( + parameters: inout [String: String], + endpoint: URL, + authentication: OAuthConfiguration.TokenEndpointAuthentication, + session: URLSession + ) async throws -> OAuthTokenResponse +} + +/// Stateless OAuth token endpoint HTTP client. +/// +/// Handles the low-level HTTP mechanics of making token requests. +struct OAuthTokenEndpointClient: Sendable { + let urlValidator: OAuthURLValidator + + init(urlValidator: OAuthURLValidator) { + self.urlValidator = urlValidator + } + + /// Makes a token request to the given endpoint. + func request( + parameters: inout [String: String], + endpoint: URL, + authentication: OAuthConfiguration.TokenEndpointAuthentication, + session: URLSession + ) async throws -> OAuthTokenResponse { + var urlRequest = URLRequest(url: endpoint) + urlRequest.httpMethod = "POST" + urlRequest.setValue(ContentType.formURLEncoded, forHTTPHeaderField: HTTPHeaderName.contentType) + urlRequest.setValue(ContentType.json, forHTTPHeaderField: HTTPHeaderName.accept) + + try await authentication.apply( + to: &urlRequest, + bodyParameters: ¶meters, + tokenEndpoint: endpoint + ) + urlRequest.httpBody = encodeForm(parameters) + + let (data, response) = try await session.data(for: urlRequest) + guard let httpResponse = response as? HTTPURLResponse else { + throw OAuthAuthorizationError.tokenRequestFailed(statusCode: -1, oauthError: nil) + } + + guard (200..<300).contains(httpResponse.statusCode) else { + let oauthError = + (try? JSONDecoder().decode(OAuthTokenErrorResponse.self, from: data))?.error + throw OAuthAuthorizationError.tokenRequestFailed( + statusCode: httpResponse.statusCode, + oauthError: oauthError + ) + } + + let decoded = try JSONDecoder().decode(OAuthTokenResponse.self, from: data) + guard !decoded.accessToken.isEmpty else { + throw OAuthAuthorizationError.tokenResponseInvalid + } + let tokenType = decoded.tokenType.trimmingCharacters(in: .whitespacesAndNewlines) + guard !tokenType.isEmpty, + tokenType.caseInsensitiveCompare(OAuthTokenType.bearer) == .orderedSame + else { + throw OAuthAuthorizationError.tokenResponseInvalid + } + return decoded + } + + private func encodeForm(_ params: [String: String]) -> Data { + let body = params + .sorted { $0.key < $1.key } + .map { key, value in "\(percentEncode(key))=\(percentEncode(value))" } + .joined(separator: "&") + return Data(body.utf8) + } + + private func percentEncode(_ string: String) -> String { + let allowed = CharacterSet( + charactersIn: + "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-._~" + ) + return string.addingPercentEncoding(withAllowedCharacters: allowed) ?? string + } +} + +extension OAuthTokenEndpointClient: OAuthTokenRequesting {} diff --git a/Sources/MCP/Base/Authorization/OAuthURLValidator.swift b/Sources/MCP/Base/Authorization/OAuthURLValidator.swift new file mode 100644 index 0000000..ddc180c --- /dev/null +++ b/Sources/MCP/Base/Authorization/OAuthURLValidator.swift @@ -0,0 +1,135 @@ +import Foundation + +#if canImport(FoundationNetworking) + import FoundationNetworking +#endif + +/// Internal protocol for URL security validation in OAuth flows. +protocol OAuthURLValidating: Sendable { + func validateHTTPSOrLoopback(_ url: URL, context: String) throws + func validateAuthorizationServer(_ url: URL, context: String) throws + func validateRedirectURI(_ url: URL) throws + func isPrivateIPHost(_ host: String) -> Bool +} + +/// URL security rules for OAuth endpoints. +/// +/// Validates that URLs used in the OAuth flow satisfy HTTPS requirements and SSRF protections. +/// Configured once and shared across the discovery, token, and registration components. +public struct OAuthURLValidator: Sendable { + + /// When `true`, loopback HTTP URLs are accepted for authorization server endpoints. + /// + /// This is a package-level compatibility option for local test environments. + public let allowLoopbackHTTPForAuthorizationServer: Bool + + public init(allowLoopbackHTTPForAuthorizationServer: Bool = false) { + self.allowLoopbackHTTPForAuthorizationServer = allowLoopbackHTTPForAuthorizationServer + } + + /// Validates that the URL uses HTTPS or loopback HTTP, and has no fragment. + /// + /// Used for MCP endpoints and protected resource metadata URLs. + public func validateHTTPSOrLoopback(_ url: URL, context: String) throws { + guard let components = URLComponents(url: url, resolvingAgainstBaseURL: false), + let scheme = components.scheme?.lowercased(), + let host = components.host?.lowercased(), + !host.isEmpty, + components.fragment == nil + else { + throw OAuthAuthorizationError.invalidResourceURI( + "Invalid \(context): \(url.absoluteString)" + ) + } + + guard scheme == OAuthURLScheme.https + || (scheme == OAuthURLScheme.http && OAuthLoopbackHost.isLoopback(host)) + else { + throw OAuthAuthorizationError.insecureOAuthEndpoint( + context: context, + url: url.absoluteString + ) + } + } + + /// Validates that the URL is an HTTPS authorization server endpoint. + /// + /// Loopback HTTP is permitted when `allowLoopbackHTTPForAuthorizationServer` is `true`. + public func validateAuthorizationServer(_ url: URL, context: String) throws { + guard let components = URLComponents(url: url, resolvingAgainstBaseURL: false), + let scheme = components.scheme?.lowercased(), + let host = components.host?.lowercased(), + !host.isEmpty, + components.fragment == nil + else { + throw OAuthAuthorizationError.invalidResourceURI( + "Invalid \(context): \(url.absoluteString)" + ) + } + + if allowLoopbackHTTPForAuthorizationServer, + scheme == OAuthURLScheme.http, + OAuthLoopbackHost.isLoopback(host) + { + return + } + + guard scheme == OAuthURLScheme.https else { + throw OAuthAuthorizationError.insecureAuthorizationServerEndpoint( + context: context, + url: url.absoluteString + ) + } + } + + /// Validates the redirect URI for authorization_code flows. + /// + /// Accepts HTTPS or loopback HTTP; rejects any URI containing a fragment. + public func validateRedirectURI(_ url: URL) throws { + guard let components = URLComponents(url: url, resolvingAgainstBaseURL: false), + let scheme = components.scheme?.lowercased(), + components.fragment == nil + else { + throw OAuthAuthorizationError.invalidRedirectURI(url.absoluteString) + } + + if scheme == OAuthURLScheme.https { return } + + if scheme == OAuthURLScheme.http, + let host = components.host?.lowercased(), + OAuthLoopbackHost.isLoopback(host) + { + return + } + + throw OAuthAuthorizationError.invalidRedirectURI(url.absoluteString) + } + + /// Returns `true` if `host` is a literal IPv4 or IPv6 address in a private or reserved range. + /// + /// Blocked ranges: + /// - IPv4: `10.0.0.0/8`, `172.16.0.0/12`, `192.168.0.0/16`, `169.254.0.0/16`, `100.64.0.0/10` + /// - IPv6: `fc00::/7` (ULA), `fe80::/10` (link-local) + /// + /// **Limitation**: only literal IP addresses are checked. DNS rebinding is not prevented here. + public func isPrivateIPHost(_ host: String) -> Bool { + let octets = host.split(separator: ".").compactMap { UInt8($0) } + if octets.count == 4 && !host.contains(":") { + let (a, b) = (octets[0], octets[1]) + return a == 10 // 10.0.0.0/8 + || (a == 172 && (16...31).contains(b)) // 172.16.0.0/12 + || (a == 192 && b == 168) // 192.168.0.0/16 + || (a == 169 && b == 254) // 169.254.0.0/16 (link-local / cloud metadata) + || (a == 100 && (64...127).contains(b)) // 100.64.0.0/10 (CGNAT) + } + let lower = host.lowercased() + if lower.hasPrefix("fc") || lower.hasPrefix("fd") { return true } // fc00::/7 (ULA) + if lower.hasPrefix("fe") && lower.count > 2 { + let idx = lower.index(lower.startIndex, offsetBy: 2) + if "89ab".contains(lower[idx]) { return true } // fe80::/10 (link-local) + } + return false + } +} + +extension OAuthURLValidator: OAuthURLValidating {} diff --git a/Sources/MCP/Base/Authorization/OAuthWWWAuthenticateParser.swift b/Sources/MCP/Base/Authorization/OAuthWWWAuthenticateParser.swift new file mode 100644 index 0000000..51f3c6f --- /dev/null +++ b/Sources/MCP/Base/Authorization/OAuthWWWAuthenticateParser.swift @@ -0,0 +1,189 @@ +import Foundation + +// MARK: - WWW-Authenticate Parsing Protocol + +/// Parses the `WWW-Authenticate` response header to extract Bearer challenge parameters. +/// +/// ``OAuthAuthorizer`` uses this protocol to inspect 401 and 403 responses from the server. +/// Override the default ``DefaultOAuthWWWAuthenticateParser`` to handle custom challenge formats. +public protocol OAuthWWWAuthenticateParsing { + /// Parses the `WWW-Authenticate` header from a response header dictionary. + /// + /// - Parameter headers: The full set of HTTP response headers (case-insensitive lookup is expected). + /// - Returns: An ``OAuthBearerChallenge`` if a `Bearer` scheme is found, or `nil` otherwise. + func parseBearer(from headers: [String: String]) -> OAuthBearerChallenge? +} + +// MARK: - Default Implementation + +/// Default ``OAuthWWWAuthenticateParsing`` implementation. +/// +/// Parses `WWW-Authenticate` headers following RFC 6750 §3, handling: +/// - Multiple challenge schemes in a single header value (e.g., `Bearer …, Basic …`) +/// - Quoted-string parameter values with backslash escaping +/// - Case-insensitive scheme and parameter key matching +public struct DefaultOAuthWWWAuthenticateParser: OAuthWWWAuthenticateParsing { + public init() {} + private let tokenCharacters = CharacterSet( + charactersIn: "!#$%&'*+-.^_`|~0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz" + ) + + public func parseBearer(from headers: [String: String]) -> OAuthBearerChallenge? { + guard let value = headers.first(where: { + $0.key.caseInsensitiveCompare(HTTPHeaderName.wwwAuthenticate) == .orderedSame + })?.value else { + return nil + } + return parseBearerHeader(value) + } + + private func parseBearerHeader(_ header: String) -> OAuthBearerChallenge? { + let trimmed = header.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { return nil } + guard let parametersPart = extractBearerParameters(from: trimmed) else { return nil } + + if parametersPart.isEmpty { + return OAuthBearerChallenge(parameters: [:]) + } + + let components = splitParameters(parametersPart) + var parameters: [String: String] = [:] + + for component in components { + let pair = component.split(separator: "=", maxSplits: 1) + guard pair.count == 2 else { continue } + let key = pair[0].trimmingCharacters(in: .whitespacesAndNewlines).lowercased() + var value = pair[1].trimmingCharacters(in: .whitespacesAndNewlines) + if value.hasPrefix("\"") && value.hasSuffix("\"") && value.count >= 2 { + value.removeFirst() + value.removeLast() + value = value.replacingOccurrences(of: "\\\"", with: "\"") + } + parameters[key] = value + } + + return OAuthBearerChallenge(parameters: parameters) + } + + private func extractBearerParameters(from header: String) -> String? { + let segments = splitParameters(header) + + for index in segments.indices { + let segment = segments[index].trimmingCharacters(in: .whitespacesAndNewlines) + guard isBearerChallengeStart(segment) else { continue } + + var parameters: [String] = [] + let initial = stripBearerScheme(from: segment) + if !initial.isEmpty { + parameters.append(initial) + } + + var nextIndex = segments.index(after: index) + while nextIndex < segments.endIndex { + let next = segments[nextIndex].trimmingCharacters(in: .whitespacesAndNewlines) + if next.isEmpty { + nextIndex = segments.index(after: nextIndex) + continue + } + if startsNewChallenge(next) { + break + } + parameters.append(next) + nextIndex = segments.index(after: nextIndex) + } + + return parameters.joined(separator: ",") + } + + return nil + } + + private func isBearerChallengeStart(_ segment: String) -> Bool { + guard segment.count >= OAuthTokenType.bearer.count else { return false } + let schemeEnd = segment.index(segment.startIndex, offsetBy: OAuthTokenType.bearer.count) + let scheme = segment[.. String { + guard segment.count > OAuthTokenType.bearer.count else { return "" } + let index = segment.index(segment.startIndex, offsetBy: OAuthTokenType.bearer.count) + return String(segment[index...]).trimmingCharacters(in: .whitespacesAndNewlines) + } + + private func startsNewChallenge(_ segment: String) -> Bool { + if isBearerChallengeStart(segment) { + return true + } + + if let equalsIndex = segment.firstIndex(of: "=") { + let parameterName = String(segment[.. Bool { + guard !value.isEmpty else { return false } + return value.rangeOfCharacter(from: tokenCharacters.inverted) == nil + } + + private func splitParameters(_ value: String) -> [String] { + var components: [String] = [] + var current = "" + var inQuotes = false + var escaping = false + + for character in value { + if escaping { + current.append(character) + escaping = false + continue + } + if character == "\\" { + current.append(character) + escaping = true + continue + } + if character == "\"" { + inQuotes.toggle() + current.append(character) + continue + } + if character == "," && !inQuotes { + components.append(current.trimmingCharacters(in: .whitespacesAndNewlines)) + current = "" + continue + } + current.append(character) + } + + if !current.isEmpty { + components.append(current.trimmingCharacters(in: .whitespacesAndNewlines)) + } + + return components.filter { !$0.isEmpty } + } +} + diff --git a/Sources/MCP/Base/Authorization/PKCE.swift b/Sources/MCP/Base/Authorization/PKCE.swift new file mode 100644 index 0000000..4515454 --- /dev/null +++ b/Sources/MCP/Base/Authorization/PKCE.swift @@ -0,0 +1,62 @@ +import Foundation + +#if canImport(CryptoKit) + import CryptoKit +#endif + +/// Pure PKCE (RFC 7636) helpers required by the authorization_code flow. +public enum PKCE { + + /// Generates a cryptographically random PKCE code verifier. + /// + /// - Parameter length: Number of characters in the verifier. Defaults to 64. + /// RFC 7636 requires 43–128 characters. + /// - Returns: A URL-safe random string suitable for use as a PKCE code verifier. + public static func makeVerifier(length: Int = 64) -> String { + let charset = Array("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-._~") + // 66 characters; 256 % 66 == 52, so reject bytes > 252 to eliminate modulo bias. + let limit = UInt8(255 - (255 % charset.count)) // 252 + var rng = SystemRandomNumberGenerator() + var result = "" + result.reserveCapacity(length) + while result.count < length { + let byte = UInt8.random(in: 0...255, using: &rng) + if byte <= limit { + result.append(charset[Int(byte % UInt8(charset.count))]) + } + } + return result + } + + /// Derives the PKCE S256 code challenge from a verifier. + /// + /// - Parameter verifier: A code verifier produced by ``makeVerifier(length:)``. + /// - Returns: The base64url-encoded SHA-256 hash of the verifier. + /// - Throws: ``OAuthAuthorizationError/pkceS256Unavailable`` on platforms without CryptoKit. + public static func makeChallenge(from verifier: String) throws -> String { + #if canImport(CryptoKit) + let hash = SHA256.hash(data: Data(verifier.utf8)) + return Data(hash).base64URLEncodedString() + #else + throw OAuthAuthorizationError.pkceS256Unavailable + #endif + } + + /// Verifies that the authorization server metadata advertises S256 PKCE support. + /// + /// - Parameter metadata: Authorization server metadata to inspect. + /// - Throws: ``OAuthAuthorizationError/pkceCodeChallengeMethodsMissing`` if + /// `code_challenge_methods_supported` is absent or empty, or + /// ``OAuthAuthorizationError/pkceS256NotSupported(advertisedMethods:)`` if S256 is not listed. + static func checkSupport(in metadata: OAuthAuthorizationServerMetadata) throws { + guard let methods = metadata.codeChallengeMethodsSupported, !methods.isEmpty else { + throw OAuthAuthorizationError.pkceCodeChallengeMethodsMissing + } + let supportsS256 = methods.contains { + $0.caseInsensitiveCompare(OAuthCodeChallengeMethod.s256) == .orderedSame + } + guard supportsS256 else { + throw OAuthAuthorizationError.pkceS256NotSupported(advertisedMethods: methods) + } + } +} diff --git a/Sources/MCP/Base/Authorization/TokenStorage.swift b/Sources/MCP/Base/Authorization/TokenStorage.swift new file mode 100644 index 0000000..07a5e3c --- /dev/null +++ b/Sources/MCP/Base/Authorization/TokenStorage.swift @@ -0,0 +1,37 @@ +import Foundation + +// MARK: - Token Storage Protocol + +/// Abstraction for persisting OAuth access tokens. +/// +/// Implement this protocol to provide custom token storage (e.g., Keychain-backed). +/// The default ``InMemoryTokenStorage`` stores tokens in memory only. +public protocol TokenStorage: AnyObject, Sendable { + func save(_ token: OAuthAccessToken) + func load() -> OAuthAccessToken? + func clear() +} + +// MARK: - In-Memory Implementation + +/// Default ``TokenStorage`` that stores the access token in memory only. +/// +/// The token is lost when the process exits. For persistent storage +/// (e.g., system Keychain), implement ``TokenStorage`` directly. +public final class InMemoryTokenStorage: TokenStorage, @unchecked Sendable { + private var token: OAuthAccessToken? + + public init() {} + + public func save(_ token: OAuthAccessToken) { + self.token = token + } + + public func load() -> OAuthAccessToken? { + token + } + + public func clear() { + token = nil + } +} diff --git a/Sources/MCP/Base/Error.swift b/Sources/MCP/Base/Error.swift index 0e9a705..b7196f0 100644 --- a/Sources/MCP/Base/Error.swift +++ b/Sources/MCP/Base/Error.swift @@ -1,7 +1,32 @@ import Foundation +#if canImport(System) + import System +#else + @preconcurrency import SystemPackage +#endif + +/// Information about a required URL elicitation +public struct URLElicitationInfo: Codable, Hashable, Sendable { + /// Elicitation mode (must be "url") + public var mode: String + /// Unique identifier for this elicitation + public var elicitationId: String + /// URL for the user to visit + public var url: String + /// Message describing the elicitation + public var message: String + + public init(mode: String = "url", elicitationId: String, url: String, message: String) { + self.mode = mode + self.elicitationId = elicitationId + self.url = url + self.message = message + } +} + /// A model context protocol error. -public enum Error: Sendable { +public enum MCPError: Swift.Error, Sendable { // Standard JSON-RPC 2.0 errors (-32700 to -32603) case parseError(String?) // -32700 case invalidRequest(String?) // -32600 @@ -12,6 +37,9 @@ public enum Error: Sendable { // Server errors (-32000 to -32099) case serverError(code: Int, message: String) + // MCP specific errors + case urlElicitationRequired(message: String, elicitations: [URLElicitationInfo]) // -32042 + // Transport specific errors case connectionClosed case transportError(Swift.Error) @@ -25,15 +53,30 @@ public enum Error: Sendable { case .invalidParams: return -32602 case .internalError: return -32603 case .serverError(let code, _): return code + case .urlElicitationRequired: return -32042 case .connectionClosed: return -32000 case .transportError: return -32001 } } + + /// Check if an error represents a "resource temporarily unavailable" condition + public static func isResourceTemporarilyUnavailable(_ error: Swift.Error) -> Bool { + #if canImport(System) + if let errno = error as? System.Errno, errno == .resourceTemporarilyUnavailable { + return true + } + #else + if let errno = error as? SystemPackage.Errno, errno == .resourceTemporarilyUnavailable { + return true + } + #endif + return false + } } // MARK: LocalizedError -extension Error: LocalizedError { +extension MCPError: LocalizedError { public var errorDescription: String? { switch self { case .parseError(let detail): @@ -48,6 +91,8 @@ extension Error: LocalizedError { return "Internal error" + (detail.map { ": \($0)" } ?? "") case .serverError(_, let message): return "Server error: \(message)" + case .urlElicitationRequired(let message, _): + return "URL elicitation required: \(message)" case .connectionClosed: return "Connection closed" case .transportError(let error): @@ -69,6 +114,8 @@ extension Error: LocalizedError { return "Internal JSON-RPC error" case .serverError: return "Server-defined error occurred" + case .urlElicitationRequired: + return "The server requires user authentication or input via external URL" case .connectionClosed: return "The connection to the server was closed" case .transportError(let error): @@ -86,6 +133,11 @@ extension Error: 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 .urlElicitationRequired(_, let elicitations): + if let first = elicitations.first { + return "Visit \(first.url) to complete the required authentication or input" + } + return "Complete the required URL-based elicitation" case .connectionClosed: return "Try reconnecting to the server" default: @@ -96,7 +148,7 @@ extension Error: LocalizedError { // MARK: CustomDebugStringConvertible -extension Error: CustomDebugStringConvertible { +extension MCPError: CustomDebugStringConvertible { public var debugDescription: String { switch self { case .transportError(let error): @@ -111,7 +163,7 @@ extension Error: CustomDebugStringConvertible { // MARK: Codable -extension Error: Codable { +extension MCPError: Codable { private enum CodingKeys: String, CodingKey { case code, message, data } @@ -119,7 +171,6 @@ extension Error: 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) // Encode additional data if available switch self { @@ -128,15 +179,34 @@ extension Error: Codable { .methodNotFound(let detail), .invalidParams(let detail), .internalError(let detail): + try container.encode(errorDescription ?? "Unknown error", forKey: .message) if let detail = detail { try container.encode(["detail": detail], forKey: .data) } case .serverError(_, _): // No additional data for server errors + try container.encode(errorDescription ?? "Unknown error", forKey: .message) break + case .urlElicitationRequired(let message, let elicitations): + // Encode the raw message so decode can round-trip without prefix doubling + try container.encode(message, forKey: .message) + // Encode elicitations array as structured data + let elicitationsData = elicitations.map { info -> [String: Value] in + return [ + "mode": .string(info.mode), + "elicitationId": .string(info.elicitationId), + "url": .string(info.url), + "message": .string(info.message) + ] + } + try container.encode( + ["elicitations": Value.array(elicitationsData.map { .object($0) })], + forKey: .data + ) case .connectionClosed: - break + try container.encode(errorDescription ?? "Unknown error", forKey: .message) case .transportError(let error): + try container.encode(errorDescription ?? "Unknown error", forKey: .message) try container.encode( ["error": error.localizedDescription], forKey: .data @@ -150,25 +220,57 @@ extension Error: Codable { 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 + } + switch code { case -32700: - self = .parseError(data?["detail"] as? String ?? message) + self = .parseError(unwrapDetail(message)) case -32600: - self = .invalidRequest(data?["detail"] as? String ?? message) + self = .invalidRequest(unwrapDetail(message)) case -32601: - self = .methodNotFound(data?["detail"] as? String ?? message) + self = .methodNotFound(unwrapDetail(message)) case -32602: - self = .invalidParams(data?["detail"] as? String ?? message) + self = .invalidParams(unwrapDetail(message)) case -32603: - self = .internalError(data?["detail"] as? String ?? message) + self = .internalError(unwrapDetail(nil)) + case -32042: + // Extract elicitations array from data + var elicitations: [URLElicitationInfo] = [] + if case .array(let items) = data?["elicitations"] { + for item in items { + if case .object(let dict) = item, + case .string(let mode) = dict["mode"], + case .string(let elicitationId) = dict["elicitationId"], + case .string(let url) = dict["url"], + case .string(let msg) = dict["message"] { + elicitations.append(URLElicitationInfo( + mode: mode, + elicitationId: elicitationId, + url: url, + message: msg)) + } + } + } + self = .urlElicitationRequired(message: message, elicitations: elicitations) case -32000: self = .connectionClosed case -32001: + // 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 self = .transportError( NSError( domain: "org.jsonrpc.error", code: code, - userInfo: [NSLocalizedDescriptionKey: message] + userInfo: [NSLocalizedDescriptionKey: underlyingErrorString] ) ) default: @@ -179,15 +281,29 @@ extension Error: Codable { // MARK: Equatable -extension Error: Equatable { - public static func == (lhs: Error, rhs: Error) -> Bool { - lhs.code == rhs.code +extension MCPError: Equatable { + public static func == (lhs: MCPError, rhs: MCPError) -> Bool { + switch (lhs, rhs) { + case (.parseError(let a), .parseError(let b)): return a == b + case (.invalidRequest(let a), .invalidRequest(let b)): return a == b + case (.methodNotFound(let a), .methodNotFound(let b)): return a == b + case (.invalidParams(let a), .invalidParams(let b)): return a == b + case (.internalError(let a), .internalError(let b)): return a == b + case (.serverError(let c1, let m1), .serverError(let c2, let m2)): + return c1 == c2 && m1 == m2 + case (.urlElicitationRequired(let m1, let e1), .urlElicitationRequired(let m2, let e2)): + return m1 == m2 && e1 == e2 + case (.connectionClosed, .connectionClosed): return true + case (.transportError(let a), .transportError(let b)): + return a.localizedDescription == b.localizedDescription + default: return false + } } } // MARK: Hashable -extension Error: Hashable { +extension MCPError: Hashable { public func hash(into hasher: inout Hasher) { hasher.combine(code) switch self { @@ -203,6 +319,9 @@ extension Error: Hashable { hasher.combine(detail) case .serverError(_, let message): hasher.combine(message) + case .urlElicitationRequired(let message, let elicitations): + hasher.combine(message) + hasher.combine(elicitations) case .connectionClosed: break case .transportError(let error): diff --git a/Sources/MCP/Base/Icon.swift b/Sources/MCP/Base/Icon.swift new file mode 100644 index 0000000..9587199 --- /dev/null +++ b/Sources/MCP/Base/Icon.swift @@ -0,0 +1,67 @@ +import Foundation + +/// An optionally-sized icon that can be displayed in a user interface. +/// +/// Icons can be used to provide visual representation of tools, resources, prompts, +/// and other MCP entities. +/// +/// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/schema#icon +public struct Icon: Hashable, Codable, Sendable { + /// The theme this icon is designed for. + public enum Theme: String, Hashable, Codable, Sendable { + /// Icon is designed for light backgrounds + case light + /// Icon is designed for dark backgrounds + case dark + } + + /// A standard URI pointing to an icon resource. + /// + /// May be an HTTP/HTTPS URL or a `data:` URI with Base64-encoded image data. + /// + /// - Note: Consumers SHOULD take steps to ensure URLs serving icons are from + /// the same domain as the client/server or a trusted domain. + /// - Note: Consumers SHOULD take appropriate precautions when consuming SVGs + /// as they can contain executable JavaScript. + public let src: String + + /// Optional MIME type override if the source MIME type is missing or generic. + /// + /// For example: `"image/png"`, `"image/jpeg"`, or `"image/svg+xml"`. + 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. + /// + /// - `light`: Icon is designed to be used with a light background + /// - `dark`: Icon is designed to be used with a dark background + /// + /// If not provided, the client should assume the icon can be used with any theme. + public let theme: Theme? + + /// Creates a new icon. + /// + /// - Parameters: + /// - src: A standard URI pointing to an icon resource (HTTP/HTTPS URL or data: URI) + /// - mimeType: Optional MIME type override + /// - sizes: Optional array of size strings (e.g., ["48x48", "96x96"]) + /// - theme: Optional theme specifier + 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 7d3e711..08ea967 100644 --- a/Sources/MCP/Base/Lifecycle.swift +++ b/Sources/MCP/Base/Lifecycle.swift @@ -46,6 +46,43 @@ public enum Initialize: Method { public let capabilities: Server.Capabilities public let serverInfo: Server.Info public let instructions: String? + public var _meta: Metadata? + + public init( + protocolVersion: String, + capabilities: Server.Capabilities, + serverInfo: Server.Info, + instructions: String? = nil, + _meta: Metadata? = nil + ) { + self.protocolVersion = protocolVersion + self.capabilities = capabilities + self.serverInfo = serverInfo + self.instructions = instructions + self._meta = _meta + } + + private enum CodingKeys: String, CodingKey, CaseIterable { + case protocolVersion, capabilities, serverInfo, instructions, _meta + } + + 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) + } + + 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(Metadata.self, forKey: ._meta) + } } } diff --git a/Sources/MCP/Base/Messages.swift b/Sources/MCP/Base/Messages.swift index fc8e985..e53bb72 100644 --- a/Sources/MCP/Base/Messages.swift +++ b/Sources/MCP/Base/Messages.swift @@ -1,4 +1,5 @@ -import Foundation +import class Foundation.JSONDecoder +import class Foundation.JSONEncoder private let jsonrpc = "2.0" @@ -10,10 +11,16 @@ public struct Empty: NotRequired, Hashable, Codable, Sendable { public init() {} } +extension Value: NotRequired { + public init() { + self = .null + } +} + // MARK: - /// A method that can be used to send requests and receive responses. -public protocol Method { +public protocol Method: Sendable { /// The parameters of the method. associatedtype Parameters: Codable, Hashable, Sendable = Empty /// The result of the method. @@ -23,14 +30,16 @@ public protocol Method { } /// Type-erased method for request/response handling -struct AnyMethod: Method { +struct AnyMethod: Method, Sendable { static var name: String { "" } typealias Parameters = Value typealias Result = Value } extension Method where Parameters == Empty { - public static func request(id: ID = .random) -> Request { + public static func request( + id: ID = .random + ) -> Request { Request(id: id, method: name, params: Empty()) } } @@ -43,17 +52,26 @@ extension Method where 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: ID = .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: ID, + result: Self.Result + ) -> Response { Response(id: id, result: result) } /// Create a response with the given error. - public static func response(id: ID, error: Error) -> Response { + public static func response( + id: ID, + error: MCPError + ) -> Response { Response(id: id, error: error) } } @@ -69,13 +87,17 @@ public struct Request: Hashable, Identifiable, Codable, Sendable { /// The request parameters. public let params: M.Parameters - init(id: ID = .random, method: String, params: M.Parameters) { + init( + id: ID = .random, + method: String, + params: M.Parameters + ) { self.id = id self.method = method self.params = params } - private enum CodingKeys: String, CodingKey { + private enum CodingKeys: String, CodingKey, CaseIterable { case jsonrpc, id, method, params } @@ -100,11 +122,31 @@ extension Request { method = try container.decode(String.self, forKey: .method) 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) + } else if let value = try? container.decode(M.Parameters.self, forKey: .params) { + // If params exists and can be decoded, use it + params = value + } else if !container.contains(.params) + || (try? container.decodeNil(forKey: .params)) == true + { + // 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 + } else { + throw DecodingError.dataCorrupted( + DecodingError.Context( + codingPath: container.codingPath, + debugDescription: "Missing required params field")) + } } else { - params = try container.decode(M.Parameters.self, forKey: .params) + throw DecodingError.dataCorrupted( + DecodingError.Context( + codingPath: container.codingPath, + debugDescription: "Invalid params field")) } } } @@ -112,9 +154,19 @@ extension Request { /// A type-erased request for request/response handling typealias AnyRequest = Request +extension AnyRequest { + init(_ request: Request) throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(request) + self = try decoder.decode(AnyRequest.self, from: data) + } +} + /// A box for request handlers that can be type-erased class RequestHandlerBox: @unchecked Sendable { - func callAsFunction(_ request: Request) async throws -> Response { + func callAsFunction(_ request: AnyRequest) async throws -> AnyResponse { fatalError("Must override") } } @@ -128,21 +180,22 @@ final class TypedRequestHandler: RequestHandlerBox, @unchecked Sendab super.init() } - override func callAsFunction(_ request: Request) async throws -> Response - { + override func callAsFunction(_ request: AnyRequest) async throws -> AnyResponse { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + // Create a concrete request from the type-erased one - let data = try JSONEncoder().encode(request.params) - let params = try JSONDecoder().decode(M.Parameters.self, from: data) - let typedRequest = Request(id: request.id, method: M.name, params: params) + let data = try encoder.encode(request) + let request = try decoder.decode(Request.self, from: data) // Handle with concrete type - let response = try await _handle(typedRequest) + let response = try await _handle(request) // Convert result to AnyMethod response switch response.result { case .success(let result): - let resultData = try JSONEncoder().encode(result) - let resultValue = try JSONDecoder().decode(Value.self, from: resultData) + let resultData = try encoder.encode(result) + let resultValue = try decoder.decode(Value.self, from: resultData) return Response(id: response.id, result: resultValue) case .failure(let error): return Response(id: response.id, error: error) @@ -157,19 +210,31 @@ public struct Response: Hashable, Identifiable, Codable, Sendable { /// The response ID. public let id: ID /// The response result. - public let result: Swift.Result + public let result: Swift.Result - public init(id: ID, result: M.Result) { + public init( + id: ID, + result: Swift.Result + ) { self.id = id - self.result = .success(result) + self.result = result } - public init(id: ID, error: Error) { - self.id = id - self.result = .failure(error) + public init( + id: ID, + result: M.Result + ) { + self.init(id: id, result: .success(result)) } - private enum CodingKeys: String, CodingKey { + public init( + id: ID, + error: MCPError + ) { + self.init(id: id, result: .failure(error)) + } + + private enum CodingKeys: String, CodingKey, CaseIterable { case jsonrpc, id, result, error } @@ -177,6 +242,7 @@ public struct Response: Hashable, Identifiable, Codable, Sendable { var container = encoder.container(keyedBy: CodingKeys.self) try container.encode(jsonrpc, forKey: .jsonrpc) try container.encode(id, forKey: .id) + switch result { case .success(let result): try container.encode(result, forKey: .result) @@ -195,7 +261,7 @@ public struct Response: Hashable, Identifiable, Codable, Sendable { id = try container.decode(ID.self, forKey: .id) if let result = try? container.decode(M.Result.self, forKey: .result) { self.result = .success(result) - } else if let error = try? container.decode(Error.self, forKey: .error) { + } else if let error = try? container.decode(MCPError.self, forKey: .error) { self.result = .failure(error) } else { throw DecodingError.dataCorrupted( @@ -209,10 +275,29 @@ public struct Response: Hashable, Identifiable, Codable, Sendable { /// A type-erased response for request/response handling typealias AnyResponse = Response +extension AnyResponse { + init(_ response: Response) throws { + switch response.result { + case .success(let result): + let data = try JSONEncoder().encode(result) + let resultValue = try JSONDecoder().decode(Value.self, from: data) + self = Response( + id: response.id, + result: .success(resultValue) + ) + case .failure(let error): + self = Response( + id: response.id, + result: .failure(error) + ) + } + } +} + // MARK: - /// A notification message. -public protocol Notification { +public protocol Notification: Hashable, Codable, Sendable { /// The parameters of the notification. associatedtype Parameters: Hashable, Codable, Sendable = Empty /// The name of the notification. @@ -220,9 +305,19 @@ public protocol Notification { } /// A type-erased notification for message handling -struct AnyNotification: Notification { +struct AnyNotification: Notification, Sendable { static var name: String { "" } - typealias Parameters = Empty + typealias Parameters = Value +} + +extension AnyNotification { + init(_ notification: some Notification) throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(notification) + self = try decoder.decode(AnyNotification.self, from: data) + } } /// A message that can be used to send notifications. @@ -259,10 +354,11 @@ public struct Message: Hashable, Codable, Sendable { } method = try container.decode(String.self, forKey: .method) - // Handle params field being optional - if N.Parameters.self == Empty.self { - // For Empty parameters, use Empty() if params is missing or null - params = Empty() as! N.Parameters + 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) } else if let value = try? container.decode(N.Parameters.self, forKey: .params) { // If params exists and can be decoded, use it params = value diff --git a/Sources/MCP/Base/PendingRequest.swift b/Sources/MCP/Base/PendingRequest.swift new file mode 100644 index 0000000..b6e0a83 --- /dev/null +++ b/Sources/MCP/Base/PendingRequest.swift @@ -0,0 +1,43 @@ +import class Foundation.JSONDecoder +import class Foundation.JSONEncoder + +/// An error indicating a type mismatch when decoding a pending request response. +struct TypeMismatchError: Swift.Error {} + +/// A pending request with a continuation for the result. +struct PendingRequest { + let continuation: CheckedContinuation +} + +/// A type-erased pending request. +struct AnyPendingRequest: Sendable { + private let _resume: @Sendable (Result) -> Void + + init(_ request: PendingRequest) { + _resume = { result in + switch result { + case .success(let value): + if let typedValue = value as? T { + request.continuation.resume(returning: typedValue) + } 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) + } else { + request.continuation.resume(throwing: TypeMismatchError()) + } + case .failure(let error): + request.continuation.resume(throwing: error) + } + } + } + + func resume(returning value: Any) { + _resume(.success(value)) + } + + func resume(throwing error: Swift.Error) { + _resume(.failure(error)) + } +} diff --git a/Sources/MCP/Base/Transport.swift b/Sources/MCP/Base/Transport.swift new file mode 100644 index 0000000..4e4350f --- /dev/null +++ b/Sources/MCP/Base/Transport.swift @@ -0,0 +1,20 @@ +import Logging + +import struct Foundation.Data + +/// Protocol defining the transport layer for MCP communication +public protocol Transport: Actor { + var logger: Logger { get } + + /// Establishes connection with the transport + func connect() async throws + + /// Disconnects from the transport + func disconnect() async + + /// Sends data + func send(_ data: Data) async throws + + /// Receives data in an async sequence + func receive() -> AsyncThrowingStream +} diff --git a/Sources/MCP/Base/Transports.swift b/Sources/MCP/Base/Transports.swift deleted file mode 100644 index b00827a..0000000 --- a/Sources/MCP/Base/Transports.swift +++ /dev/null @@ -1,412 +0,0 @@ -import Darwin -import Logging -import SystemPackage - -import struct Foundation.Data -import struct Foundation.Date -import class Foundation.JSONDecoder -import class Foundation.JSONEncoder - -/// Protocol defining the transport layer for MCP communication -public protocol Transport: Actor { - var logger: Logger { get } - - /// Establishes connection with the transport - func connect() async throws - - /// Disconnects from the transport - func disconnect() async - - /// Sends a message string - func send(_ message: String) async throws - - /// Receives message strings as an async sequence - func receive() -> AsyncThrowingStream -} - -/// Standard input/output transport implementation -public actor StdioTransport: Transport { - private let input: FileDescriptor - private let output: FileDescriptor - public nonisolated let logger: Logger - - private var isConnected = false - private let messageStream: AsyncStream - private let messageContinuation: AsyncStream.Continuation - - public init( - input: FileDescriptor = FileDescriptor.standardInput, - output: FileDescriptor = FileDescriptor.standardOutput, - logger: Logger? = nil - ) { - self.input = input - self.output = output - self.logger = - logger - ?? Logger( - label: "mcp.transport.stdio", - factory: { _ in SwiftLogNoOpLogHandler() }) - - // Create message stream - var continuation: AsyncStream.Continuation! - self.messageStream = AsyncStream { continuation = $0 } - self.messageContinuation = continuation - } - - public func connect() async throws { - guard !isConnected else { return } - - // Set non-blocking mode - try setNonBlocking(fileDescriptor: input) - try setNonBlocking(fileDescriptor: output) - - isConnected = true - logger.info("Transport connected successfully") - - // Start reading loop in background - Task { - await readLoop() - } - } - - private func setNonBlocking(fileDescriptor: FileDescriptor) throws { - let flags = fcntl(fileDescriptor.rawValue, F_GETFL) - guard flags >= 0 else { - throw Error.transportError(Errno.badFileDescriptor) - } - let result = fcntl(fileDescriptor.rawValue, F_SETFL, flags | O_NONBLOCK) - guard result >= 0 else { - throw Error.transportError(Errno.badFileDescriptor) - } - } - - private func readLoop() async { - let bufferSize = 4096 - var buffer = [UInt8](repeating: 0, count: bufferSize) - var pendingData = Data() - - while isConnected && !Task.isCancelled { - do { - let bytesRead = try buffer.withUnsafeMutableBufferPointer { pointer in - try input.read(into: UnsafeMutableRawBufferPointer(pointer)) - } - - if bytesRead == 0 { - logger.notice("EOF received") - break - } - - pendingData.append(Data(buffer[.. 0 { - remaining = remaining.dropFirst(written) - } - } catch let error as Errno where error == .resourceTemporarilyUnavailable { - try await Task.sleep(nanoseconds: 10_000_000) // 10ms backoff - continue - } catch { - throw Error.transportError(error) - } - } - } - - public func receive() -> AsyncThrowingStream { - return AsyncThrowingStream { continuation in - Task { - for await message in messageStream { - continuation.yield(message) - } - continuation.finish() - } - } - } -} - -#if canImport(Network) - import Network - - /// Network connection based transport implementation - public actor NetworkTransport: Transport { - private let connection: NWConnection - public nonisolated let logger: Logger - - private var isConnected = false - private let messageStream: AsyncThrowingStream - private let messageContinuation: AsyncThrowingStream.Continuation - - // Track connection state for continuations - private var connectionContinuationResumed = false - - public init(connection: NWConnection, logger: Logger? = nil) { - self.connection = connection - self.logger = - logger - ?? Logger( - label: "mcp.transport.network", - factory: { _ in SwiftLogNoOpLogHandler() } - ) - - // Create message stream - var continuation: AsyncThrowingStream.Continuation! - self.messageStream = AsyncThrowingStream { continuation = $0 } - self.messageContinuation = continuation - } - - /// Connects to the network transport - public func connect() async throws { - guard !isConnected else { return } - - // Reset continuation state - connectionContinuationResumed = false - - // Wait for connection to be ready - try await withCheckedThrowingContinuation { - [weak self] (continuation: CheckedContinuation) in - guard let self = self else { - continuation.resume(throwing: MCP.Error.internalError("Transport deallocated")) - return - } - - 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) - default: - // Wait for ready or failed state - break - } - } - } - - // Start the connection if it's not already started - if connection.state != .ready { - connection.start(queue: .main) - } else { - Task { @MainActor in - await self.handleConnectionReady(continuation: continuation) - } - } - } - } - - private func handleConnectionReady(continuation: CheckedContinuation) - async - { - if !connectionContinuationResumed { - connectionContinuationResumed = true - isConnected = true - logger.info("Network transport connected successfully") - continuation.resume() - // Start the receive loop after connection is established - Task { await self.receiveLoop() } - } - } - - private func handleConnectionFailed( - error: Swift.Error, continuation: CheckedContinuation - ) async { - if !connectionContinuationResumed { - connectionContinuationResumed = true - logger.error("Connection failed: \(error)") - continuation.resume(throwing: error) - } - } - - private func handleConnectionCancelled(continuation: CheckedContinuation) - async - { - if !connectionContinuationResumed { - connectionContinuationResumed = true - logger.warning("Connection cancelled") - continuation.resume(throwing: MCP.Error.internalError("Connection cancelled")) - } - } - - public func disconnect() async { - guard isConnected else { return } - isConnected = false - connection.cancel() - messageContinuation.finish() - logger.info("Network transport disconnected") - } - - public func send(_ message: String) async throws { - guard isConnected else { - throw MCP.Error.internalError("Transport not connected") - } - - guard let data = (message + "\n").data(using: .utf8) else { - throw MCP.Error.internalError("Failed to encode message") - } - - // 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 { - continuation.resume(throwing: MCP.Error.internalError("Transport deallocated")) - return - } - - connection.send( - content: data, - 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)") - continuation.resume( - throwing: MCP.Error.internalError("Send error: \(error)")) - } else { - continuation.resume() - } - } - } - }) - } - } - - public func receive() -> AsyncThrowingStream { - return AsyncThrowingStream { continuation in - Task { - do { - for try await message in messageStream { - continuation.yield(message) - } - continuation.finish() - } catch { - continuation.finish(throwing: error) - } - } - } - } - - private func receiveLoop() async { - var buffer = Data() - - while isConnected && !Task.isCancelled { - do { - let newData = try await receiveData() - buffer.append(newData) - - // Process complete messages - while let newlineIndex = buffer.firstIndex(of: UInt8(ascii: "\n")) { - let messageData = buffer[.. Data { - var receiveContinuationResumed = false - - return try await withCheckedThrowingContinuation { - [weak self] (continuation: CheckedContinuation) in - guard let self = self else { - continuation.resume(throwing: MCP.Error.internalError("Transport deallocated")) - return - } - - connection.receive(minimumIncompleteLength: 1, maximumLength: 65536) { - content, _, _, error in - Task { @MainActor in - if !receiveContinuationResumed { - receiveContinuationResumed = true - if let error = error { - continuation.resume(throwing: MCP.Error.transportError(error)) - } else if let content = content { - continuation.resume(returning: content) - } else { - continuation.resume( - throwing: MCP.Error.internalError("No data received")) - } - } - } - } - } - } - } -#endif diff --git a/Sources/MCP/Base/Transports/HTTPClientTransport.swift b/Sources/MCP/Base/Transports/HTTPClientTransport.swift new file mode 100644 index 0000000..060cd13 --- /dev/null +++ b/Sources/MCP/Base/Transports/HTTPClientTransport.swift @@ -0,0 +1,686 @@ +import Foundation +import Logging + +#if !os(Linux) + import EventSource +#endif + +#if canImport(FoundationNetworking) + import FoundationNetworking +#endif + +// MARK: - Timeout Helpers + +/// Error thrown when an operation times out +/// An implementation of the MCP Streamable HTTP transport protocol for clients. +/// +/// This transport implements the [Streamable HTTP transport](https://spec.modelcontextprotocol.io/specification/2025-11-25/basic/transports#streamable-http) +/// specification from the Model Context Protocol (version 2025-11-25). +/// +/// It supports: +/// - Sending JSON-RPC messages via HTTP POST requests +/// - Receiving responses via both direct JSON responses and SSE streams +/// - Session management using the `MCP-Session-Id` header +/// - Protocol version negotiation via `MCP-Protocol-Version` header +/// - Automatic reconnection for dropped SSE streams with resumability support +/// - Platform-specific optimizations for different operating systems +/// +/// The transport supports two modes: +/// - 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. +/// +/// ## Example Usage +/// +/// ```swift +/// import MCP +/// +/// // Create a streaming HTTP transport with bearer token authentication +/// let transport = HTTPClientTransport( +/// endpoint: URL(string: "https://api.example.com/mcp")!, +/// requestModifier: { request in +/// var modifiedRequest = request +/// modifiedRequest.addValue("Bearer your-token-here", forHTTPHeaderField: "Authorization") +/// return modifiedRequest +/// } +/// ) +/// +/// // Initialize the client with streaming transport +/// let client = Client(name: "MyApp", version: "1.0.0") +/// try await client.connect(transport: transport) +/// +/// // The transport will automatically handle SSE events +/// // and deliver them through the client's notification handlers +/// ``` +public actor HTTPClientTransport: Transport { + /// The server endpoint URL to connect to + public let endpoint: URL + private let session: URLSession + + /// The session ID assigned by the server, used for maintaining state across requests + public private(set) var sessionID: String? + + /// The negotiated protocol version to send in MCP-Protocol-Version header + public var protocolVersion: String? + + private let streaming: Bool + private var streamingTask: Task? + + /// Logger instance for transport-related events + public nonisolated let logger: Logger + + /// Maximum time to wait for a session ID before proceeding with SSE connection + public let sseInitializationTimeout: TimeInterval + + /// Closure to modify requests before they are sent + private let requestModifier: (URLRequest) -> URLRequest + + /// Optional OAuth 2.1 authorizer. + private let authorizer: (any HTTPClientAuthorizer)? + + private var isConnected = false + private let messageStream: AsyncThrowingStream + private let messageContinuation: AsyncThrowingStream.Continuation + + private var initialSessionIDSignalTask: Task? + private var initialSessionIDContinuation: CheckedContinuation? + + /// The last event ID received from the server for SSE stream resumability + private var lastEventID: String? + + /// The retry interval (in milliseconds) from the server's SSE `retry:` field + private var retryInterval: Int = 3000 // Default 3000ms per SSE spec + + /// The underlying URLSession task for the active GET SSE stream. + /// Used to trigger reconnection when a POST SSE stream closes without delivering data. + private var activeGETSessionTask: URLSessionDataTask? + + /// Creates a new HTTP transport client with the specified endpoint + /// + /// - Parameters: + /// - endpoint: The server URL to connect to + /// - 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) + /// - protocolVersion: The MCP protocol version to use (default: "2025-11-25") + /// - authorizer: Optional ``HTTPClientAuthorizer`` for automatic Bearer token acquisition and retries. + /// - requestModifier: Optional closure to customize requests before they are sent (default: no modification) + /// - logger: Optional logger instance for transport events + public init( + endpoint: URL, + configuration: URLSessionConfiguration = .default, + streaming: Bool = true, + sseInitializationTimeout: TimeInterval = 10, + protocolVersion: String = Version.latest, + authorizer: (any HTTPClientAuthorizer)? = nil, + requestModifier: @escaping (URLRequest) -> URLRequest = { $0 }, + logger: Logger? = nil + ) { + self.init( + endpoint: endpoint, + session: URLSession(configuration: configuration), + streaming: streaming, + sseInitializationTimeout: sseInitializationTimeout, + protocolVersion: protocolVersion, + authorizer: authorizer, + requestModifier: requestModifier, + logger: logger + ) + } + + internal init( + endpoint: URL, + session: URLSession, + streaming: Bool = false, + sseInitializationTimeout: TimeInterval = 10, + protocolVersion: String = Version.latest, + authorizer: (any HTTPClientAuthorizer)? = nil, + requestModifier: @escaping (URLRequest) -> URLRequest = { $0 }, + logger: Logger? = nil + ) { + self.endpoint = endpoint + self.session = session + self.streaming = streaming + self.sseInitializationTimeout = sseInitializationTimeout + self.protocolVersion = protocolVersion + self.requestModifier = requestModifier + self.authorizer = authorizer + + // Create message stream + var continuation: AsyncThrowingStream.Continuation! + self.messageStream = AsyncThrowingStream { continuation = $0 } + self.messageContinuation = continuation + + self.logger = + logger + ?? Logger( + label: "mcp.transport.http.client", + factory: { _ in SwiftLogNoOpLogHandler() } + ) + } + + // Setup the initial session ID signal + private func setupInitialSessionIDSignal() { + self.initialSessionIDSignalTask = Task { + await withCheckedContinuation { continuation in + self.initialSessionIDContinuation = 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 + logger.debug("✓ Initial session ID signal triggered for SSE task") + } else { + logger.debug("✗ No continuation to trigger - signal already consumed or SSE task not waiting") + } + } + + /// Establishes connection with the transport + public func connect() async throws { + guard !isConnected else { return } + isConnected = true + + setupInitialSessionIDSignal() + + if streaming { + streamingTask = Task { await startListeningForServerEvents() } + } + + logger.debug("HTTP transport connected") + } + + /// Disconnects from the transport + public func disconnect() async { + guard isConnected else { return } + isConnected = false + + streamingTask?.cancel() + streamingTask = nil + + session.invalidateAndCancel() + messageContinuation.finish() + + initialSessionIDSignalTask?.cancel() + initialSessionIDSignalTask = nil + initialSessionIDContinuation?.resume() + initialSessionIDContinuation = nil + + logger.debug("HTTP clienttransport disconnected") + } + + /// Updates the protocol version used for `MCP-Protocol-Version` headers on subsequent requests. + public func updateNegotiatedProtocolVersion(_ version: String) { + self.protocolVersion = version + } + + /// Sends data through an HTTP POST request + public func send(_ data: Data) async throws { + guard isConnected else { + throw MCPError.internalError("Transport not connected") + } + + if let authorizer { + do { + try authorizer.validateEndpointSecurity(for: endpoint) + } catch { + throw MCPError.internalError( + "Authorization flow failed: \(error.localizedDescription)" + ) + } + } + + if let authorizer { + try? await authorizer.prepareAuthorization(for: endpoint, session: session) + } + + var attempts = 0 + let operationKey = jsonRPCOperationKey(from: data) + + while true { + var request = URLRequest(url: endpoint) + request.httpMethod = "POST" + request.addValue( + "\(ContentType.json), \(ContentType.sse)", + forHTTPHeaderField: HTTPHeaderName.accept + ) + request.addValue(ContentType.json, forHTTPHeaderField: HTTPHeaderName.contentType) + request.httpBody = data + + if let protocolVersion = protocolVersion { + request.addValue(protocolVersion, forHTTPHeaderField: HTTPHeaderName.protocolVersion) + } + + if let sessionID = sessionID { + request.addValue(sessionID, forHTTPHeaderField: HTTPHeaderName.sessionID) + } + + if let authValue = authorizer?.authorizationHeader(for: endpoint) { + request.setValue(authValue, forHTTPHeaderField: HTTPHeaderName.authorization) + } + + request = requestModifier(request) + + do { + #if os(Linux) + let (responseData, response) = try await session.data(for: request) + try await processResponse(response: response, data: responseData) + #else + let (responseStream, response) = try await session.bytes(for: request) + try await processResponse(response: response, stream: responseStream) + #endif + return + } catch let authError as HTTPAuthenticationChallengeError { + guard let authorizer else { + throw mapAuthenticationChallengeError(authError) + } + + let handled: Bool + do { + handled = try await authorizer.handleChallenge( + statusCode: authError.statusCode, + headers: authError.headers, + endpoint: endpoint, + operationKey: operationKey, + session: session + ) + } catch { + throw MCPError.internalError( + "Authorization flow failed: \(error.localizedDescription)") + } + + attempts += 1 + + if handled, attempts < authorizer.maxAuthorizationAttempts { + continue + } + + throw mapAuthenticationChallengeError(authError) + } + } + } + + #if os(Linux) + private func processResponse(response: URLResponse, data: Data) async throws { + guard let httpResponse = response as? HTTPURLResponse else { + throw MCPError.internalError("Invalid HTTP response") + } + + let contentType = httpResponse.value(forHTTPHeaderField: HTTPHeaderName.contentType) ?? "" + + if let newSessionID = httpResponse.value(forHTTPHeaderField: HTTPHeaderName.sessionID) { + let wasSessionIDNil = (self.sessionID == nil) + self.sessionID = newSessionID + if wasSessionIDNil { + triggerInitialSessionIDSignal() + } + logger.debug("Session ID received", metadata: ["sessionID": "\(newSessionID)"]) + } + + try processHTTPResponse(httpResponse, contentType: contentType) + guard case 200..<300 = httpResponse.statusCode else { return } + + if contentType.contains(ContentType.sse) { + logger.warning("SSE responses aren't fully supported on Linux") + messageContinuation.yield(data) + } else if contentType.contains(ContentType.json) { + logger.trace("Received JSON response", metadata: ["size": "\(data.count)"]) + messageContinuation.yield(data) + } else { + logger.warning("Unexpected content type: \(contentType)") + } + } + #else + private func processResponse(response: URLResponse, stream: URLSession.AsyncBytes) + async throws + { + guard let httpResponse = response as? HTTPURLResponse else { + throw MCPError.internalError("Invalid HTTP response") + } + + let contentType = httpResponse.value(forHTTPHeaderField: HTTPHeaderName.contentType) ?? "" + + if let newSessionID = httpResponse.value(forHTTPHeaderField: HTTPHeaderName.sessionID) { + let wasSessionIDNil = (self.sessionID == nil) + self.sessionID = newSessionID + if wasSessionIDNil { + triggerInitialSessionIDSignal() + } + logger.debug("Session ID received", metadata: ["sessionID": "\(newSessionID)"]) + } + + try processHTTPResponse(httpResponse, contentType: contentType) + guard case 200..<300 = httpResponse.statusCode else { return } + + if contentType.contains(ContentType.sse) { + logger.trace("Received SSE response, processing in streaming task") + let hadData = try await self.processSSE(stream) + + if !hadData { + logger.debug("POST SSE stream closed without data, triggering GET reconnection") + self.activeGETSessionTask?.cancel() + } + } else if contentType.contains(ContentType.json) { + var buffer = Data() + for try await byte in stream { + buffer.append(byte) + } + logger.trace("Received JSON response", metadata: ["size": "\(buffer.count)"]) + messageContinuation.yield(buffer) + } else { + logger.warning("Unexpected content type: \(contentType)") + } + } + #endif + + private func processHTTPResponse(_ response: HTTPURLResponse, contentType: String) throws { + switch response.statusCode { + case 200..<300: + return + + case 400: + throw MCPError.internalError("Bad request") + + case 401: + throw HTTPAuthenticationChallengeError( + statusCode: response.statusCode, + headers: responseHeaders(from: response) + ) + + case 403: + throw HTTPAuthenticationChallengeError( + statusCode: response.statusCode, + headers: responseHeaders(from: response) + ) + + case 404: + if sessionID != nil { + logger.warning("Session has expired") + sessionID = nil + throw MCPError.internalError("Session expired") + } + throw MCPError.internalError("Endpoint not found") + + case 405: + if streaming { + self.streamingTask?.cancel() + throw MCPError.internalError("Server does not support streaming") + } + throw MCPError.internalError("Method not allowed") + + case 408: + throw MCPError.internalError("Request timeout") + + case 429: + throw MCPError.internalError("Too many requests") + + case 500..<600: + throw MCPError.internalError("Server error: \(response.statusCode)") + + default: + throw MCPError.internalError( + "Unexpected HTTP response: \(response.statusCode) (\(contentType))") + } + } + + private func mapAuthenticationChallengeError(_ error: HTTPAuthenticationChallengeError) -> MCPError { + switch error.statusCode { + case 401: + return MCPError.internalError("Authentication required") + case 403: + return MCPError.internalError("Access forbidden") + default: + return MCPError.internalError("HTTP authorization error: \(error.statusCode)") + } + } + + private func responseHeaders(from response: HTTPURLResponse) -> [String: String] { + var headers: [String: String] = [:] + for (key, value) in response.allHeaderFields { + guard let key = key as? String, let value = value as? String else { continue } + headers[key] = value + } + return headers + } + + private func jsonRPCOperationKey(from data: Data) -> String? { + guard + let jsonObject = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let method = jsonObject["method"] as? String + else { + return nil + } + + let normalized = method.trimmingCharacters(in: .whitespacesAndNewlines) + return normalized.isEmpty ? nil : normalized + } + + /// Receives data in an async sequence + public func receive() -> AsyncThrowingStream { + return messageStream + } + + // MARK: - SSE + + private func startListeningForServerEvents() async { + #if os(Linux) + if streaming { + logger.warning( + "SSE streaming was requested but is not fully supported on Linux. SSE connection will not be attempted." + ) + } + #else + guard isConnected else { return } + + if self.sessionID == nil, let signalTask = self.initialSessionIDSignalTask { + logger.debug("⏳ Waiting for session ID to be set (timeout: \(self.sseInitializationTimeout)s)...") + + let startTime = Date() + let timeout = self.sseInitializationTimeout + do { + try await withThrowingTaskGroup { group in + group.addTask { + try await Task.sleep(for: .seconds(timeout)) + } + + group.addTask { + await signalTask.value + } + + if let firstResult = try await group.next() { + group.cancelAll() + return firstResult + } + } + } catch { + logger.warning("⏱️ Timeout waiting for session ID (\(timeout)s). SSE stream will proceed anyway.") + } + + if self.sessionID != nil { + let elapsed = Date().timeIntervalSince(startTime) + logger.debug("✓ Session ID received after \(Int(elapsed * 1000))ms, proceeding with SSE connection") + } + } else { + logger.debug("✓ Session ID already available, proceeding with SSE connection immediately") + } + + var isFirstAttempt = true + var attemptCount = 0 + + logger.debug("🔄 Starting SSE retry loop", metadata: [ + "isConnected": "\(isConnected)", + "isCancelled": "\(Task.isCancelled)" + ]) + + while isConnected && !Task.isCancelled { + attemptCount += 1 + logger.debug("🔄 SSE retry loop iteration", metadata: [ + "attempt": "\(attemptCount)", + "isFirstAttempt": "\(isFirstAttempt)" + ]) + + do { + if !isFirstAttempt { + let delayMs = self.retryInterval + logger.debug("⏳ Waiting before SSE reconnection", metadata: ["retryMs": "\(delayMs)"]) + try await Task.sleep(for: .milliseconds(delayMs)) + logger.debug("✓ Wait complete, reconnecting now") + } + isFirstAttempt = false + + logger.debug("📡 Calling connectToEventStream (attempt #\(attemptCount))") + + try await self.connectToEventStream() + + logger.info("🔌 SSE stream closed gracefully, will reconnect", metadata: [ + "attempt": "\(attemptCount)", + "willRetryAfter": "\(self.retryInterval)ms" + ]) + } catch { + if !Task.isCancelled { + logger.error("❌ SSE connection error (attempt #\(attemptCount)): \(error)") + } else { + logger.debug("⏹️ SSE task cancelled") + } + } + + logger.debug("🔄 End of retry loop iteration", metadata: [ + "isConnected": "\(isConnected)", + "isCancelled": "\(Task.isCancelled)", + "willContinue": "\(isConnected && !Task.isCancelled)" + ]) + } + + logger.debug("⏹️ SSE retry loop exited", metadata: [ + "isConnected": "\(isConnected)", + "isCancelled": "\(Task.isCancelled)", + "totalAttempts": "\(attemptCount)" + ]) + #endif + } + + #if !os(Linux) + private func connectToEventStream() async throws { + guard isConnected else { + logger.debug("⚠️ Skipping connectToEventStream - transport not connected") + return + } + + if let authorizer { + do { + try authorizer.validateEndpointSecurity(for: endpoint) + } catch { + throw MCPError.internalError( + "Authorization flow failed: \(error.localizedDescription)" + ) + } + } + + if let authorizer { + try? await authorizer.prepareAuthorization(for: endpoint, session: session) + } + + logger.debug("🔌 Preparing SSE connection request") + + var request = URLRequest(url: endpoint) + request.httpMethod = "GET" + request.addValue(ContentType.sse, forHTTPHeaderField: HTTPHeaderName.accept) + request.addValue("no-cache", forHTTPHeaderField: HTTPHeaderName.cacheControl) + + if let protocolVersion = protocolVersion { + request.addValue(protocolVersion, forHTTPHeaderField: HTTPHeaderName.protocolVersion) + } + + if let sessionID = sessionID { + request.addValue(sessionID, forHTTPHeaderField: HTTPHeaderName.sessionID) + } + + if let lastEventID = lastEventID { + request.addValue(lastEventID, forHTTPHeaderField: HTTPHeaderName.lastEventID) + logger.info("→ Resuming SSE stream with Last-Event-ID", metadata: ["lastEventID": "\(lastEventID)"]) + } else { + logger.info("→ Connecting to SSE stream (no last event ID to resume from)") + } + + if let authValue = authorizer?.authorizationHeader(for: endpoint) { + request.setValue(authValue, forHTTPHeaderField: HTTPHeaderName.authorization) + } + + request = requestModifier(request) + + logger.debug("Starting SSE connection") + + let (stream, response) = try await session.bytes(for: request) + self.activeGETSessionTask = stream.task + + guard let httpResponse = response as? HTTPURLResponse else { + throw MCPError.internalError("Invalid HTTP response") + } + + guard httpResponse.statusCode == 200 else { + if httpResponse.statusCode == 405 { + self.streamingTask?.cancel() + } + throw MCPError.internalError("HTTP error: \(httpResponse.statusCode)") + } + + if let newSessionID = httpResponse.value(forHTTPHeaderField: HTTPHeaderName.sessionID) { + let wasSessionIDNil = (self.sessionID == nil) + self.sessionID = newSessionID + if wasSessionIDNil { + triggerInitialSessionIDSignal() + } + logger.debug("Session ID received", metadata: ["sessionID": "\(newSessionID)"]) + } + + defer { self.activeGETSessionTask = nil } + try await self.processSSE(stream) + } + + @discardableResult + private func processSSE(_ stream: URLSession.AsyncBytes) async throws -> Bool { + logger.debug("📥 Starting SSE event processing") + var eventCount = 0 + var hadDataEvent = false + + for try await event in stream.events { + eventCount += 1 + + if Task.isCancelled { + logger.debug("⏹️ SSE processing cancelled", metadata: ["eventsProcessed": "\(eventCount)"]) + break + } + + logger.trace( + "SSE event received", + metadata: [ + "type": "\(event.event ?? "message")", + "id": "\(event.id ?? "none")", + ] + ) + + if let eventID = event.id, !eventID.isEmpty { + self.lastEventID = eventID + logger.debug("Stored event ID for resumability", metadata: ["eventID": "\(eventID)"]) + } + + if let retry = event.retry { + self.retryInterval = retry + logger.debug("SSE retry interval updated", metadata: ["retryMs": "\(retry)"]) + } + + if !event.data.isEmpty, let data = event.data.data(using: .utf8) { + hadDataEvent = true + messageContinuation.yield(data) + } + } + + logger.debug("✓ SSE event stream completed", metadata: ["eventsProcessed": "\(eventCount)", "hadData": "\(hadDataEvent)"]) + return hadDataEvent + } + #endif +} diff --git a/Sources/MCP/Base/Transports/HTTPServer/HTTPRequestValidation.swift b/Sources/MCP/Base/Transports/HTTPServer/HTTPRequestValidation.swift new file mode 100644 index 0000000..51c7aae --- /dev/null +++ b/Sources/MCP/Base/Transports/HTTPServer/HTTPRequestValidation.swift @@ -0,0 +1,724 @@ +import Foundation + +// MARK: - Validation Protocol + +/// Validates an incoming HTTP request before the transport processes it. +/// +/// Validators are composed into a pipeline and executed in order. The first validator +/// that returns a non-nil response short-circuits the pipeline and that error response +/// is returned to the client. +/// +/// Conform to this protocol to add custom validation (e.g., authentication): +/// ```swift +/// struct BearerTokenValidator: HTTPRequestValidator { +/// func validate(_ request: HTTPRequest, context: HTTPValidationContext) -> HTTPResponse? { +/// guard let auth = request.header("Authorization"), +/// auth.hasPrefix("Bearer ") else { +/// return .error(statusCode: 401, .invalidRequest("Missing bearer token")) +/// } +/// return nil +/// } +/// } +/// ``` +public protocol HTTPRequestValidator: Sendable { + /// Validates the request. Returns an error response if invalid, or `nil` if valid. + func validate(_ request: HTTPRequest, context: HTTPValidationContext) -> HTTPResponse? +} + +// MARK: - Validation Context + +/// Context provided to validators for making validation decisions. +public struct HTTPValidationContext: Sendable { + /// The HTTP method of the request (GET, POST, DELETE). + public let httpMethod: String + + /// The current session ID, if any (nil in stateless mode or before initialization). + public let sessionID: String? + + /// Whether the request body contains an `initialize` JSON-RPC request. + public let isInitializationRequest: Bool + + /// The set of protocol versions this server supports. + public let supportedProtocolVersions: Set + + public init( + httpMethod: String, + sessionID: String? = nil, + isInitializationRequest: Bool = false, + supportedProtocolVersions: Set = Version.supported + ) { + self.httpMethod = httpMethod + self.sessionID = sessionID + self.isInitializationRequest = isInitializationRequest + self.supportedProtocolVersions = supportedProtocolVersions + } +} + +// MARK: - Accept Header Validator + +/// Validates the `Accept` header based on the HTTP method and transport response mode. +/// +/// - Stateful (SSE) mode: POST requests must accept both `application/json` and `text/event-stream` +/// - Stateless (JSON) mode: POST requests only need to accept `application/json` +/// - GET requests always require `text/event-stream` +public struct AcceptHeaderValidator: HTTPRequestValidator { + /// The response mode determines which content types are required. + public enum Mode: Sendable { + /// POST requires both `application/json` and `text/event-stream`. + case sseRequired + /// POST only requires `application/json`. + case jsonOnly + } + + public let mode: Mode + + public init(mode: Mode) { + self.mode = mode + } + + public func validate(_ request: HTTPRequest, context: HTTPValidationContext) -> HTTPResponse? { + let accept = request.header(HTTPHeaderName.accept) ?? "" + let acceptTypes = accept.split(separator: ",").map { + $0.trimmingCharacters(in: .whitespaces) + } + + let hasJSON = acceptTypes.contains { $0.hasPrefix(ContentType.json) } + let hasSSE = acceptTypes.contains { $0.hasPrefix(ContentType.sse) } + + switch context.httpMethod { + case "POST": + switch mode { + case .sseRequired: + guard hasJSON, hasSSE else { + return .error( + statusCode: 406, + .invalidRequest( + "Not Acceptable: Client must accept both application/json and text/event-stream" + ), + sessionID: context.sessionID + ) + } + case .jsonOnly: + guard hasJSON else { + return .error( + statusCode: 406, + .invalidRequest( + "Not Acceptable: Client must accept application/json" + ), + sessionID: context.sessionID + ) + } + } + case "GET": + guard hasSSE else { + return .error( + statusCode: 406, + .invalidRequest( + "Not Acceptable: Client must accept text/event-stream" + ), + sessionID: context.sessionID + ) + } + default: + break + } + + return nil + } +} + +// MARK: - Content-Type Validator + +/// Validates that POST requests have `Content-Type: application/json`. +public struct ContentTypeValidator: HTTPRequestValidator { + public init() {} + + public func validate(_ request: HTTPRequest, context: HTTPValidationContext) -> HTTPResponse? { + guard context.httpMethod == "POST" else { return nil } + + let contentType = request.header(HTTPHeaderName.contentType) ?? "" + let mainType = contentType.split(separator: ";").first? + .trimmingCharacters(in: .whitespaces) ?? "" + + guard mainType == ContentType.json else { + return .error( + statusCode: 415, + .invalidRequest( + "Unsupported Media Type: Content-Type must be application/json" + ), + sessionID: context.sessionID + ) + } + + return nil + } +} + +// MARK: - Protocol Version Validator + +/// Validates the `MCP-Protocol-Version` header against supported versions. +/// +/// Per spec: +/// - If the header is absent, the server assumes the default negotiated version +/// - If the header is present but unsupported, the server returns 400 Bad Request +/// - Initialization requests are exempt (protocol version comes from the request body) +public struct ProtocolVersionValidator: HTTPRequestValidator { + public init() {} + + public func validate(_ request: HTTPRequest, context: HTTPValidationContext) -> HTTPResponse? { + // Skip for initialization requests (version is in the body, not the header) + guard !context.isInitializationRequest else { return nil } + + // Skip for non-POST methods (GET/DELETE don't carry protocol version) + // Actually, per spec, all subsequent requests should include it + guard let version = request.header(HTTPHeaderName.protocolVersion) else { + // Per spec: if not received, assume default version + return nil + } + + guard context.supportedProtocolVersions.contains(version) else { + let supported = context.supportedProtocolVersions.sorted().joined(separator: ", ") + return .error( + statusCode: 400, + .invalidRequest( + "Bad Request: Unsupported protocol version: \(version). Supported: \(supported)" + ), + sessionID: context.sessionID + ) + } + + return nil + } +} + +// MARK: - Session Validator + +/// Validates the `Mcp-Session-Id` header for stateful transports. +/// +/// - Initialization requests are exempt (no session exists yet) +/// - Non-initialization requests must include the session ID header +/// - The session ID must match the active session +public struct SessionValidator: HTTPRequestValidator { + public init() {} + + public func validate(_ request: HTTPRequest, context: HTTPValidationContext) -> HTTPResponse? { + // Skip validation for initialization requests + guard !context.isInitializationRequest else { return nil } + + // Non-initialization requests require an established session. + guard let expectedSessionID = context.sessionID else { + return .error( + statusCode: 400, + .invalidRequest("Bad Request: Session not initialized"), + sessionID: nil + ) + } + + let requestSessionID = request.header(HTTPHeaderName.sessionID) + + guard let requestSessionID else { + return .error( + statusCode: 400, + .invalidRequest("Bad Request: Missing \(HTTPHeaderName.sessionID) header"), + sessionID: expectedSessionID + ) + } + + guard requestSessionID == expectedSessionID else { + return .error( + statusCode: 404, + .invalidRequest("Not Found: Invalid or expired session ID"), + sessionID: expectedSessionID + ) + } + + return nil + } +} + +// MARK: - Origin Validator + +/// DNS rebinding protection: validates `Origin` and `Host` headers. +/// +/// Per spec, servers MUST validate the Origin header to prevent DNS rebinding attacks. +/// This is particularly important for servers running on localhost. +/// +/// Use `.localhost()` for local development servers. +/// Use `.disabled` to skip validation (e.g., cloud deployments). +/// Use `init(allowedHosts:allowedOrigins:)` for custom configurations. +public struct OriginValidator: HTTPRequestValidator { + public let allowedHosts: [String] + public let allowedOrigins: [String] + private let enabled: Bool + + public init(allowedHosts: [String], allowedOrigins: [String]) { + self.allowedHosts = allowedHosts + self.allowedOrigins = allowedOrigins + self.enabled = true + } + + private init(disabled: Void) { + self.allowedHosts = [] + self.allowedOrigins = [] + self.enabled = false + } + + /// Protection for localhost-bound servers. + /// Allows requests from `localhost`, `127.0.0.1`, and `[::1]` with the specified port. + public static func localhost(port: Int? = nil) -> OriginValidator { + let portPattern = port.map { String($0) } ?? "*" + return OriginValidator( + allowedHosts: [ + "127.0.0.1:\(portPattern)", + "localhost:\(portPattern)", + "[::1]:\(portPattern)", + ], + allowedOrigins: [ + "http://127.0.0.1:\(portPattern)", + "http://localhost:\(portPattern)", + "http://[::1]:\(portPattern)", + ] + ) + } + + /// Disables DNS rebinding protection. + /// Use for cloud deployments where DNS rebinding is not a threat. + public static var disabled: OriginValidator { + OriginValidator(disabled: ()) + } + + public func validate(_ request: HTTPRequest, context: HTTPValidationContext) -> HTTPResponse? { + guard enabled else { return nil } + + // Validate Host header + if let host = request.header(HTTPHeaderName.host) { + let hostAllowed = allowedHosts.contains { pattern in + matchesPattern(host, pattern: pattern) + } + if !hostAllowed { + return .error( + statusCode: 421, + .invalidRequest("Misdirected Request: Host header not allowed"), + sessionID: context.sessionID + ) + } + } + + // Validate Origin header (only if present — non-browser clients won't send it) + if let origin = request.header(HTTPHeaderName.origin) { + let originAllowed = allowedOrigins.contains { pattern in + matchesPattern(origin, pattern: pattern) + } + if !originAllowed { + return .error( + statusCode: 403, + .invalidRequest("Forbidden: Origin not allowed"), + sessionID: context.sessionID + ) + } + } + + return nil + } + + /// Matches a value against a pattern that may contain a port wildcard `:*`. + /// + /// Examples: + /// - `"localhost:*"` matches `"localhost:8080"`, `"localhost:3000"` + /// - `"http://localhost:*"` matches `"http://localhost:8080"` + /// - `"localhost:8080"` matches only `"localhost:8080"` exactly + private func matchesPattern(_ value: String, pattern: String) -> Bool { + guard pattern.hasSuffix(":*") else { + return value == pattern + } + + let prefix = String(pattern.dropLast(2)) + guard value.hasPrefix(prefix + ":") else { return false } + + let portPart = value.dropFirst(prefix.count + 1) + return !portPart.isEmpty && portPart.allSatisfy(\.isNumber) + } +} + +// MARK: - Protected Resource Metadata Validator + +/// Serves the RFC 9728 Protected Resource Metadata document for discovery. +/// +/// Per the MCP authorization specification, servers **MUST** serve Protected Resource +/// Metadata at `/.well-known/oauth-protected-resource` so that clients can discover +/// authorization server endpoints automatically. +/// +/// Place this validator **before** ``BearerTokenValidator`` in the pipeline so that +/// unauthenticated metadata discovery requests succeed. +/// +/// ```swift +/// let prmValidator = ProtectedResourceMetadataValidator( +/// metadata: OAuthProtectedResourceServerMetadata( +/// resource: "https://api.example.com", +/// authorizationServers: [URL(string: "https://auth.example.com")!] +/// ) +/// ) +/// let pipeline = StandardValidationPipeline(validators: [ +/// prmValidator, +/// bearerTokenValidator, +/// // ... +/// ]) +/// ``` +public struct ProtectedResourceMetadataValidator: HTTPRequestValidator { + private let encodedMetadata: Data + + public init(metadata: OAuthProtectedResourceServerMetadata) { + self.encodedMetadata = (try? JSONEncoder().encode(metadata)) ?? Data() + } + + public func validate(_ request: HTTPRequest, context: HTTPValidationContext) -> HTTPResponse? { + guard context.httpMethod == "GET", + let path = request.path, + path == OAuthWellKnownPath.protectedResource + || path.hasPrefix("\(OAuthWellKnownPath.protectedResource)/") + else { + return nil + } + return .data(encodedMetadata, headers: [HTTPHeaderName.contentType: ContentType.json]) + } +} + +// MARK: - OAuth Bearer Validator + +/// Result produced by ``BearerTokenValidator`` when validating an access token. +public enum BearerTokenValidationResult: Sendable, Equatable { + /// Access token is valid for this request, with its extracted claims. + /// + /// Supply a ``BearerTokenInfo`` with `audience` and `expiresAt` populated so that + /// ``BearerTokenValidator`` can enforce expiry and audience checks automatically. + /// Pass `BearerTokenInfo()` (all `nil`) to delegate all enforcement to the caller. + case valid(BearerTokenInfo) + + /// Access token is missing required privileges, and new scopes are required. + case insufficientScope(requiredScopes: Set, errorDescription: String? = nil) + + /// Access token is invalid or expired. + case invalidToken(errorDescription: String? = nil) + + /// Authorization request is malformed. + case malformedRequest(errorDescription: String? = nil) +} + +/// Validates OAuth 2.1 Bearer authorization for protected MCP HTTP endpoints. +/// +/// This validator implements resource-server error semantics aligned with the MCP auth spec: +/// - `401` with `WWW-Authenticate: Bearer ...` for missing/invalid tokens +/// - `403` with `error="insufficient_scope"` for insufficient permissions +/// - `400` for malformed authorization requests +/// +/// Include this validator early in your pipeline, before `SessionValidator`, so unauthenticated +/// initialization requests can return a challenge. +/// +/// ## Audience Validation (MUST) +/// +/// Per the MCP authorization specification, **the resource server MUST validate the audience +/// (`aud` claim) of the access token** to ensure it matches the resource server's own identifier. +/// Failure to validate the audience allows token substitution attacks where a token intended +/// for a different resource is replayed against your server. +/// +/// Your ``TokenValidator`` closure **MUST** verify the audience. Example: +/// +/// ```swift +/// let validator = BearerTokenValidator( +/// resourceMetadataURL: metadataURL, +/// tokenValidator: { token, request, context in +/// guard let claims = verifyAndDecode(token) else { +/// return .invalidToken(errorDescription: "Token verification failed") +/// } +/// // MUST: Verify token audience matches this resource server +/// guard claims.audience.contains("https://api.example.com") else { +/// return .invalidToken(errorDescription: "Token audience mismatch") +/// } +/// return .valid +/// } +/// ) +/// ``` +public struct BearerTokenValidator: HTTPRequestValidator { + /// Validates a bearer token and returns token info for audience and expiry enforcement. + public typealias TokenValidator = @Sendable ( + _ token: String, + _ request: HTTPRequest, + _ context: HTTPValidationContext + ) -> BearerTokenValidationResult + + /// Closure that returns the scopes to advertise in `WWW-Authenticate` challenge headers. + /// + /// Return `nil` to omit the `scope` parameter from the challenge. + public typealias ChallengeScopeProvider = @Sendable ( + _ request: HTTPRequest, + _ context: HTTPValidationContext + ) -> Set? + + /// Closure that decides whether a request requires Bearer authentication. + /// + /// Return `false` to allow a request through unauthenticated (e.g., public health-check endpoints). + /// Defaults to requiring authentication on all requests. + public typealias RequirementPredicate = @Sendable ( + _ request: HTTPRequest, + _ context: HTTPValidationContext + ) -> Bool + + public let resourceMetadataURL: URL + public let resourceIdentifier: URL + private let tokenValidator: TokenValidator + private let challengeScopeProvider: ChallengeScopeProvider? + private let requiresAuthentication: RequirementPredicate + private let metadataDiscovery: any OAuthMetadataDiscovering + + /// Creates a `BearerTokenValidator`. + /// + /// - Parameters: + /// - resourceMetadataURL: Included in `WWW-Authenticate` challenge headers as the + /// `resource_metadata` parameter, pointing to the RFC 9728 Protected Resource Metadata document. + /// - resourceIdentifier: The canonical URI of this resource server. Used to validate the + /// `aud` claim in tokens that supply audience information via ``BearerTokenInfo``. + /// - tokenValidator: Validates the Bearer token and returns ``BearerTokenInfo`` with + /// claims for SDK-side expiry and audience enforcement. + /// - challengeScopeProvider: Optional closure supplying scopes to include in challenge headers. + /// - requiresAuthentication: Predicate controlling which requests require a Bearer token. + /// Defaults to requiring authentication on all requests. + /// - metadataDiscovery: Used for audience URL matching. Defaults to ``DefaultOAuthMetadataDiscovery``. + public init( + resourceMetadataURL: URL, + resourceIdentifier: URL, + tokenValidator: @escaping TokenValidator, + challengeScopeProvider: ChallengeScopeProvider? = nil, + requiresAuthentication: @escaping RequirementPredicate = { _, _ in true }, + metadataDiscovery: any OAuthMetadataDiscovering = DefaultOAuthMetadataDiscovery() + ) { + self.resourceMetadataURL = resourceMetadataURL + self.resourceIdentifier = resourceIdentifier + self.tokenValidator = tokenValidator + self.challengeScopeProvider = challengeScopeProvider + self.requiresAuthentication = requiresAuthentication + self.metadataDiscovery = metadataDiscovery + } + + public func validate(_ request: HTTPRequest, context: HTTPValidationContext) -> HTTPResponse? { + guard requiresAuthentication(request, context) else { return nil } + + guard let authorizationHeader = request.header(HTTPHeaderName.authorization) else { + return unauthorizedResponse( + challengeScope: challengeScopeProvider?(request, context), + error: nil, + errorDescription: nil, + sessionID: context.sessionID + ) + } + + let parsedToken: String + switch parseBearerToken(from: authorizationHeader) { + case .success(let token): + parsedToken = token + case .failure(let error): + return .error( + statusCode: 400, + .invalidRequest("Bad Request: \(error.message)"), + sessionID: context.sessionID + ) + } + + switch tokenValidator(parsedToken, request, context) { + case .valid(let info): + // Expiry check + if let exp = info.expiresAt, exp <= Date() { + return unauthorizedResponse( + challengeScope: challengeScopeProvider?(request, context), + error: "invalid_token", + errorDescription: "Token has expired", + sessionID: context.sessionID + ) + } + // Audience check — skipped for opaque tokens (audience == nil) + if let audience = info.audience { + let matches = audience.contains { audString in + guard let audURL = URL(string: audString) else { return false } + return metadataDiscovery.protectedResourceMatches( + resource: audURL, endpoint: resourceIdentifier) + } + if !matches { + return unauthorizedResponse( + challengeScope: challengeScopeProvider?(request, context), + error: "invalid_token", + errorDescription: "Token audience mismatch", + sessionID: context.sessionID + ) + } + } + return nil + + case .invalidToken(let errorDescription): + return unauthorizedResponse( + challengeScope: challengeScopeProvider?(request, context), + error: "invalid_token", + errorDescription: errorDescription, + sessionID: context.sessionID + ) + + case .insufficientScope(let requiredScopes, let errorDescription): + return forbiddenInsufficientScopeResponse( + requiredScopes: requiredScopes, + errorDescription: errorDescription, + sessionID: context.sessionID + ) + + case .malformedRequest(let errorDescription): + let message = errorDescription ?? "Malformed authorization request" + return .error( + statusCode: 400, + .invalidRequest("Bad Request: \(message)"), + sessionID: context.sessionID + ) + } + } + + private func unauthorizedResponse( + challengeScope: Set?, + error: String?, + errorDescription: String?, + sessionID: String? + ) -> HTTPResponse { + let challenge = makeBearerChallenge( + resourceMetadataURL: resourceMetadataURL, + scope: challengeScope, + error: error, + errorDescription: errorDescription + ) + return .error( + statusCode: 401, + .invalidRequest("Unauthorized"), + sessionID: sessionID, + extraHeaders: [HTTPHeaderName.wwwAuthenticate: challenge] + ) + } + + private func forbiddenInsufficientScopeResponse( + requiredScopes: Set, + errorDescription: String?, + sessionID: String? + ) -> HTTPResponse { + let challenge = makeBearerChallenge( + resourceMetadataURL: resourceMetadataURL, + scope: requiredScopes, + error: "insufficient_scope", + errorDescription: errorDescription + ) + return .error( + statusCode: 403, + .invalidRequest("Forbidden: Insufficient scope"), + sessionID: sessionID, + extraHeaders: [HTTPHeaderName.wwwAuthenticate: challenge] + ) + } + + private func makeBearerChallenge( + resourceMetadataURL: URL, + scope: Set?, + error: String?, + errorDescription: String? + ) -> String { + var parameters: [String] = [] + parameters.append("resource_metadata=\"\(escapeAuthParameter(resourceMetadataURL.absoluteString))\"") + + if let scope, !scope.isEmpty { + let serializedScope = scope.sorted().joined(separator: " ") + parameters.append("scope=\"\(escapeAuthParameter(serializedScope))\"") + } + + if let error { + parameters.append("error=\"\(escapeAuthParameter(error))\"") + } + + if let errorDescription, !errorDescription.isEmpty { + parameters.append("error_description=\"\(escapeAuthParameter(errorDescription))\"") + } + + return "\(OAuthTokenType.bearer) " + parameters.joined(separator: ", ") + } + + private func escapeAuthParameter(_ value: String) -> String { + value + .replacingOccurrences(of: "\\", with: "\\\\") + .replacingOccurrences(of: "\"", with: "\\\"") + } + + private struct BearerTokenParseError: Swift.Error { + let message: String + } + + private func parseBearerToken( + from authorizationHeader: String + ) -> Result { + let trimmed = authorizationHeader.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { + return .failure(.init(message: "Authorization header is empty")) + } + + let parts = trimmed.split( + maxSplits: 1, + whereSeparator: { $0.isWhitespace } + ) + + guard parts.count == 2 else { + return .failure( + .init(message: "Authorization header must be in the form: Bearer ") + ) + } + + guard String(parts[0]).caseInsensitiveCompare(OAuthTokenType.bearer) == .orderedSame else { + return .failure(.init(message: "Authorization scheme must be Bearer")) + } + + let token = String(parts[1]).trimmingCharacters(in: .whitespacesAndNewlines) + guard !token.isEmpty else { + return .failure(.init(message: "Bearer token is empty")) + } + + if token.contains(where: \.isWhitespace) { + return .failure(.init(message: "Bearer token must not contain whitespace")) + } + + return .success(token) + } +} + +// MARK: - Validation Pipeline Protocol + +/// Runs a validation pipeline against an HTTP request. +/// +/// Implementations execute a sequence of validators and return the first error, +/// or `nil` if all validations pass. +public protocol HTTPRequestValidationPipeline: Sendable { + /// Validates the request using the configured pipeline. + /// Returns an error response if validation fails, or `nil` if the request is valid. + func validate(_ request: HTTPRequest, context: HTTPValidationContext) -> HTTPResponse? +} + +// MARK: - Standard Validation Pipeline + +/// Standard implementation of `HTTPRequestValidationPipeline` that runs validators in sequence. +/// +/// The first validator that returns a non-nil error response short-circuits the pipeline. +public struct StandardValidationPipeline: HTTPRequestValidationPipeline { + private let validators: [any HTTPRequestValidator] + + /// Creates a pipeline with the given validators. + /// Validators are executed in the order provided. + public init(validators: [any HTTPRequestValidator]) { + self.validators = validators + } + + public func validate(_ request: HTTPRequest, context: HTTPValidationContext) -> HTTPResponse? { + for validator in validators { + if let errorResponse = validator.validate(request, context: context) { + return errorResponse + } + } + return nil + } +} diff --git a/Sources/MCP/Base/Transports/HTTPServer/HTTPServerTypes.swift b/Sources/MCP/Base/Transports/HTTPServer/HTTPServerTypes.swift new file mode 100644 index 0000000..bd56191 --- /dev/null +++ b/Sources/MCP/Base/Transports/HTTPServer/HTTPServerTypes.swift @@ -0,0 +1,285 @@ +import Foundation + +// MARK: - Session ID Generator + +/// Generates unique session identifiers for stateful HTTP server transports. +/// +/// Conform to this protocol to provide custom session ID generation logic. +/// Session IDs **MUST** contain only visible ASCII characters (0x21–0x7E) +/// per the MCP specification. +/// +/// A default implementation using UUID is provided via ``UUIDSessionIDGenerator``. +public protocol SessionIDGenerator: Sendable { + /// Generates a new unique session identifier. + func generateSessionID() -> String +} + +/// Default session ID generator that produces UUID strings. +/// +/// UUID strings consist of hexadecimal characters and hyphens, +/// which are all within the valid ASCII range (0x21–0x7E). +public struct UUIDSessionIDGenerator: SessionIDGenerator { + public init() {} + + public func generateSessionID() -> String { + UUID().uuidString + } +} + +// MARK: - HTTP Request + +/// A framework-agnostic HTTP request representation. +/// +/// This type decouples the transport from any specific HTTP framework. +/// The HTTP framework adapter converts its native request type into this before passing to the transport. +public struct HTTPRequest: Sendable { + /// The HTTP method (e.g., "GET", "POST", "DELETE"). + public let method: String + + /// HTTP headers as key-value pairs. + public let headers: [String: String] + + /// The request body data, if any. + public let body: Data? + + /// The request path (e.g., "/mcp", "/.well-known/oauth-protected-resource"). + /// Used by validators that need to match on specific paths. + public let path: String? + + public init(method: String, headers: [String: String] = [:], body: Data? = nil, path: String? = nil) { + self.method = method + self.headers = headers + self.body = body + self.path = path + } + + /// Case-insensitive header lookup. + public func header(_ name: String) -> String? { + let lowercased = name.lowercased() + return headers.first { $0.key.lowercased() == lowercased }?.value + } +} + +// MARK: - HTTP Response + +/// A framework-agnostic HTTP response. +/// +/// The HTTP framework adapter converts this into its native response type. +/// +/// Use computed properties (`statusCode`, `headers`, `bodyData`) for generic access, +/// or switch on the enum for case-specific handling (e.g., streaming): +/// +/// ```swift +/// let response = await transport.handleRequest(request) +/// switch response { +/// case .stream(let sseStream, _): +/// // Pipe the async stream to the HTTP response body +/// default: +/// // Use response.bodyData for the body +/// } +/// ``` +public enum HTTPResponse: Sendable { + /// 202 Accepted, no body. Used for notifications and client responses. + case accepted(headers: [String: String] = [:]) + + /// 200 OK, no body. Used for DELETE confirmation. + case ok(headers: [String: String] = [:]) + + /// 200 OK with data body (typically JSON). + case data(Data, headers: [String: String] = [:]) + + /// 200 OK with SSE streaming body. + case stream(AsyncThrowingStream, headers: [String: String] = [:]) + + /// Error response with a JSON-RPC error body. + /// The status code, headers, and body are derived automatically. + case error(statusCode: Int, MCPError, sessionID: String? = nil, extraHeaders: [String: String] = [:]) + + // MARK: - Computed Properties + + public var statusCode: Int { + switch self { + case .accepted: 202 + case .ok, .data, .stream: 200 + case .error(let code, _, _, _): code + } + } + + public var headers: [String: String] { + switch self { + case .accepted(let headers), .ok(let headers), .data(_, let headers), .stream(_, let headers): + return headers + case .error(_, _, let sessionID, let extraHeaders): + var headers: [String: String] = [HTTPHeaderName.contentType: ContentType.json] + if let sessionID { headers[HTTPHeaderName.sessionID] = sessionID } + headers.merge(extraHeaders) { _, new in new } + return headers + } + } + + /// The response body as data. `nil` for `.accepted`, `.ok`, and `.stream`. + public var bodyData: Data? { + switch self { + case .accepted, .ok, .stream: + return nil + case .data(let data, _): + return data + case .error(_, let error, _, _): + let errorBody: [String: Any] = [ + "jsonrpc": "2.0", + "error": [ + "code": error.code, + "message": error.errorDescription ?? "Unknown error", + ], + "id": NSNull(), + ] + return try? JSONSerialization.data(withJSONObject: errorBody) + } + } +} + +// MARK: - HTTP Header Names + +/// Standard header names used by the MCP Streamable HTTP transport. +public enum HTTPHeaderName { + public static let sessionID = "MCP-Session-Id" + public static let protocolVersion = "MCP-Protocol-Version" + public static let lastEventID = "Last-Event-ID" + public static let accept = "Accept" + public static let contentType = "Content-Type" + public static let authorization = "Authorization" + public static let wwwAuthenticate = "WWW-Authenticate" + public static let origin = "Origin" + public static let host = "Host" + public static let cacheControl = "Cache-Control" + public static let connection = "Connection" + public static let allow = "Allow" +} + +// MARK: - Content Types + +enum ContentType { + static let json = "application/json" + static let sse = "text/event-stream" +} + +// MARK: - SSE Event + +/// A Server-Sent Event (SSE) data structure. +/// +/// Formats according to the SSE specification: +/// https://html.spec.whatwg.org/multipage/server-sent-events.html +struct SSEEvent: Sendable { + var id: String? + var event: String? + var data: String + var retry: Int? + + /// Formats the event as SSE wire data. + func formatted() -> Data { + var result = "" + if let id { + result += "id: \(id)\n" + } + if let event { + result += "event: \(event)\n" + } + if let retry { + result += "retry: \(retry)\n" + } + result += "data: \(data)\n\n" + return Data(result.utf8) + } + + /// Creates a priming event with an empty data field. + /// Per spec, this is sent immediately to prime the client for reconnection. + static func priming(id: String, retry: Int? = nil) -> SSEEvent { + SSEEvent(id: id, event: nil, data: "", retry: retry) + } + + /// Creates a message event wrapping JSON-RPC data. + static func message(data: Data, id: String? = nil) -> SSEEvent { + SSEEvent( + id: id, + event: "message", + data: String(decoding: data, as: UTF8.self) + ) + } +} + +// MARK: - HTTP Context Providing + +/// A transport that can surface the originating HTTP request for an in-flight JSON-RPC +/// request, keyed by its JSON-RPC id. +/// +/// The server consults this during dispatch so that registered method handlers can +/// observe the HTTP request that triggered them (headers, auth, path, body) via +/// ``Server/currentHandlerContext``, without changing the `withMethodHandler` signature. +/// +/// Only JSON-RPC *requests* (with an id) are addressable — notifications have no id and +/// are not correlated. +public protocol HTTPContextProviding: Sendable { + /// Returns the HTTP request associated with the given JSON-RPC id, if the + /// transport still has it on hand. Implementations must return `nil` once the + /// corresponding response has been delivered or the session has terminated. + func httpRequestContext(for id: ID) async -> HTTPRequest? +} + +// MARK: - JSON-RPC Message Classification + +/// Classifies a raw JSON-RPC message for routing purposes. +/// +/// Used by transports to determine where to route outgoing messages: +/// - Responses are routed to the originating request's stream +/// - Notifications and server requests are routed to the standalone GET stream +package enum JSONRPCMessageKind { + case request(id: String, method: String) + case notification(method: String) + case response(id: String) + + /// Attempts to classify raw JSON-RPC data. + /// Returns `nil` if the data cannot be parsed or classified. + package init?(data: Data) { + guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { + return nil + } + + let id = Self.extractID(from: json) + + if let method = json["method"] as? String { + if let id { + self = .request(id: id, method: method) + } else { + self = .notification(method: method) + } + } else if json["result"] != nil || json["error"] != nil { + guard let id else { return nil } + self = .response(id: id) + } else { + return nil + } + } + + /// Whether this message is a JSON-RPC response (success or error). + var isResponse: Bool { + if case .response = self { return true } + return false + } + + /// Whether this message is an `initialize` request. + package var isInitializeRequest: Bool { + if case .request(_, let method) = self { + return method == "initialize" + } + return false + } + + private static func extractID(from json: [String: Any]) -> String? { + if let stringID = json["id"] as? String { + return stringID + } else if let intID = json["id"] as? Int { + return String(intID) + } + return nil + } +} diff --git a/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerTransport.swift b/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerTransport.swift new file mode 100644 index 0000000..c7a4d48 --- /dev/null +++ b/Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerTransport.swift @@ -0,0 +1,581 @@ +import Foundation +import Logging + +/// A stateful HTTP server transport that manages sessions and uses SSE for streaming responses. +/// +/// This transport implements the MCP Streamable HTTP specification with full session management: +/// - Assigns a session ID during initialization (via `Mcp-Session-Id` header) +/// - POST requests receive SSE-streamed responses +/// - GET requests establish a standalone SSE stream for server-initiated messages +/// - DELETE requests terminate the session +/// - Built-in event store for resumability (reconnection with `Last-Event-ID`) +/// +/// ## Usage +/// +/// ```swift +/// let transport = StatefulHTTPServerTransport() // Uses UUID by default +/// +/// // Start the MCP server with this transport +/// try await server.start(transport: transport) +/// +/// // In your HTTP framework handler: +/// let response = await transport.handleRequest(httpRequest) +/// // Convert response to your framework's response type and return it +/// ``` +/// +/// ## Framework Integration +/// +/// This transport is framework-agnostic. You provide incoming requests as `HTTPRequest` +/// and receive `HTTPResponse` values to convert to your framework's native types. +/// For SSE responses, the `.stream` case provides an `AsyncThrowingStream` +/// to pipe to the client. +public actor StatefulHTTPServerTransport: Transport, HTTPContextProviding { + public nonisolated let logger: Logger + + // MARK: - Dependencies + + private let sessionIDGenerator: any SessionIDGenerator + private let validationPipeline: any HTTPRequestValidationPipeline + private let retryInterval: Int? + + // MARK: - State + + private var sessionID: String? + private var terminated = false + private var started = false + + // MARK: - Incoming message stream (client → server) + + private let incomingStream: AsyncThrowingStream + private let incomingContinuation: AsyncThrowingStream.Continuation + + // MARK: - SSE streams for POST request responses + + /// Maps request ID → SSE stream continuation for active POST request streams. + private var requestSSEContinuations: [String: AsyncThrowingStream.Continuation] = [:] + + /// Maps request ID → originating HTTP request, surfaced to handlers via + /// ``Server/currentHTTPContext``. Cleared when the response is routed, + /// the stream is closed, or the session terminates. + private var httpRequestContexts: [String: HTTPRequest] = [:] + + // MARK: - Standalone GET SSE stream + + /// The standalone SSE stream continuation for server-initiated messages. + /// Only one GET stream is allowed per session. + private var standaloneSSEContinuation: AsyncThrowingStream.Continuation? + + /// Internal identifier for the standalone GET stream in the event store. + private let standaloneStreamID = "_GET_stream" + + // MARK: - Event Store (Resumability) + + private struct StoredEvent { + let streamID: String + let eventID: String + let message: Data? + } + + private var storedEvents: [StoredEvent] = [] + private var eventCounter: Int = 0 + + // MARK: - Init + + /// Creates a new stateful HTTP server transport. + /// + /// - Parameters: + /// - sessionIDGenerator: Generator for session IDs. The IDs MUST contain + /// only visible ASCII characters (0x21-0x7E) per the MCP specification. + /// Defaults to ``UUIDSessionIDGenerator``. + /// - validationPipeline: Custom validation pipeline. If `nil`, uses sensible defaults: + /// origin validation (localhost), Accept header (SSE required), Content-Type, + /// protocol version, and session validation. + /// - retryInterval: Retry interval in milliseconds for SSE priming events. + /// Controls how long clients wait before attempting to reconnect. + /// - logger: Optional logger. If `nil`, a no-op logger is used. + public init( + sessionIDGenerator: any SessionIDGenerator = UUIDSessionIDGenerator(), + validationPipeline: (any HTTPRequestValidationPipeline)? = nil, + retryInterval: Int? = nil, + logger: Logger? = nil + ) { + self.sessionIDGenerator = sessionIDGenerator + self.validationPipeline = validationPipeline ?? StandardValidationPipeline(validators: [ + OriginValidator.localhost(), + AcceptHeaderValidator(mode: .sseRequired), + ContentTypeValidator(), + ProtocolVersionValidator(), + SessionValidator(), + ]) + self.retryInterval = retryInterval + self.logger = logger ?? Logger( + label: "mcp.transport.http.server.stateful", + factory: { _ in SwiftLogNoOpLogHandler() } + ) + + let (stream, continuation) = AsyncThrowingStream.makeStream() + self.incomingStream = stream + self.incomingContinuation = continuation + } + + // MARK: - Transport Conformance + + public func connect() async throws { + guard !started else { + throw MCPError.internalError("Transport already started") + } + started = true + logger.debug("Stateful HTTP server transport started") + } + + public func disconnect() async { + terminate() + } + + /// Routes outgoing server messages to the appropriate client connection. + /// + /// - Responses are routed to the SSE stream matching the response's JSON-RPC ID. + /// - Notifications and server-initiated requests are routed to the standalone GET stream. + public func send(_ data: Data) async throws { + guard !terminated else { + throw MCPError.connectionClosed + } + + guard let kind = JSONRPCMessageKind(data: data) else { + logger.warning("Could not classify outgoing message for routing") + return + } + + switch kind { + case .response(let id): + routeResponse(data, requestID: id) + case .notification, .request: + routeServerInitiatedMessage(data) + } + } + + public func receive() -> AsyncThrowingStream { + incomingStream + } + + // MARK: - HTTP Request Handler + + /// Handles an incoming HTTP request from the framework adapter. + /// + /// Routes by HTTP method: + /// - **POST**: JSON-RPC messages (requests, notifications) + /// - **GET**: Establish standalone SSE stream for server-initiated messages + /// - **DELETE**: Terminate the session + /// - Others: 405 Method Not Allowed + public func handleRequest(_ request: HTTPRequest) async -> HTTPResponse { + if terminated { + return .error( + statusCode: 404, + .invalidRequest("Not Found: Session has been terminated"), + sessionID: sessionID + ) + } + + switch request.method.uppercased() { + case "POST": + return handlePost(request) + case "GET": + return handleGet(request) + case "DELETE": + return handleDelete(request) + default: + return .error( + statusCode: 405, + .invalidRequest("Method Not Allowed"), + sessionID: sessionID, + extraHeaders: [HTTPHeaderName.allow: "GET, POST, DELETE"] + ) + } + } + + // MARK: - POST Handler + + private func handlePost(_ request: HTTPRequest) -> HTTPResponse { + // Parse body first so we can determine if it's an initialization request + guard let body = request.body, !body.isEmpty else { + return .error( + statusCode: 400, + .parseError("Empty request body"), + sessionID: sessionID + ) + } + + guard let messageKind = JSONRPCMessageKind(data: body) else { + return .error( + statusCode: 400, + .parseError("Invalid JSON-RPC message"), + sessionID: sessionID + ) + } + + // Build validation context + let context = HTTPValidationContext( + httpMethod: "POST", + sessionID: sessionID, + isInitializationRequest: messageKind.isInitializeRequest, + supportedProtocolVersions: Version.supported + ) + + // Run validation pipeline + if let errorResponse = validationPipeline.validate(request, context: context) { + return errorResponse + } + + // Handle initialization request specially + if messageKind.isInitializeRequest { + return handleInitializationRequest(body, request: request) + } + + // Handle by message type + switch messageKind { + case .notification, .response: + // Yield to server and return 202 Accepted + incomingContinuation.yield(body) + return .accepted(headers: sessionHeaders()) + + case .request(let id, _): + return handleJSONRPCRequest(body, requestID: id, request: request) + } + } + + private func handleInitializationRequest(_ body: Data, request: HTTPRequest) -> HTTPResponse { + // Reject re-initialization at the transport level + if sessionID != nil { + return .error( + statusCode: 400, + .invalidRequest("Bad Request: Session already initialized"), + sessionID: sessionID + ) + } + + // Generate session ID + let newSessionID = sessionIDGenerator.generateSessionID() + + // Validate session ID contains only visible ASCII (0x21-0x7E) + guard isValidSessionID(newSessionID) else { + logger.error("Generated session ID contains invalid characters") + return .error( + statusCode: 500, + .internalError("Internal error: Invalid session ID generated") + ) + } + + self.sessionID = newSessionID + logger.info("Session initialized", metadata: ["sessionID": "\(newSessionID)"]) + + // Extract request ID for routing the response + guard case .request(let requestID, _) = JSONRPCMessageKind(data: body) else { + return .error( + statusCode: 400, + .parseError("Invalid initialize request"), + sessionID: newSessionID + ) + } + + // For the initialize request, use SSE streaming like any other request + return handleJSONRPCRequest(body, requestID: requestID, request: request) + } + + private func handleJSONRPCRequest(_ body: Data, requestID: String, request: HTTPRequest) -> HTTPResponse { + // Create SSE stream for this request + let (sseStream, sseContinuation) = AsyncThrowingStream.makeStream() + requestSSEContinuations[requestID] = sseContinuation + httpRequestContexts[requestID] = request + + // Extract protocol version for priming event decision + let protocolVersion = extractProtocolVersion(from: body, request: request) + + // Send priming event for resumability + sendPrimingEvent( + streamID: requestID, + continuation: sseContinuation, + protocolVersion: protocolVersion + ) + + // Yield the incoming message to the server + incomingContinuation.yield(body) + + // Build response headers + var headers = sessionHeaders() + headers[HTTPHeaderName.contentType] = ContentType.sse + headers[HTTPHeaderName.cacheControl] = "no-cache, no-transform" + headers[HTTPHeaderName.connection] = "keep-alive" + + return .stream(sseStream, headers: headers) + } + + // MARK: - GET Handler + + private func handleGet(_ request: HTTPRequest) -> HTTPResponse { + // Build validation context (GET is never an initialization request) + let context = HTTPValidationContext( + httpMethod: "GET", + sessionID: sessionID, + isInitializationRequest: false, + supportedProtocolVersions: Version.supported + ) + + // Run validation pipeline + if let errorResponse = validationPipeline.validate(request, context: context) { + return errorResponse + } + + // Handle resumability: check for Last-Event-ID header + if let lastEventID = request.header(HTTPHeaderName.lastEventID) { + return handleResumeRequest(lastEventID: lastEventID, request: request) + } + + // Only one standalone GET stream per session + guard standaloneSSEContinuation == nil else { + return .error( + statusCode: 409, + .invalidRequest("Conflict: Only one SSE stream is allowed per session"), + sessionID: sessionID + ) + } + + // Create standalone SSE stream + let (sseStream, sseContinuation) = AsyncThrowingStream.makeStream() + standaloneSSEContinuation = sseContinuation + + // Extract protocol version for priming event + let protocolVersion = request.header(HTTPHeaderName.protocolVersion) ?? Version.latest + + // Send priming event + sendPrimingEvent( + streamID: standaloneStreamID, + continuation: sseContinuation, + protocolVersion: protocolVersion + ) + + // Build response headers + var headers = sessionHeaders() + headers[HTTPHeaderName.contentType] = ContentType.sse + headers[HTTPHeaderName.cacheControl] = "no-cache, no-transform" + headers[HTTPHeaderName.connection] = "keep-alive" + + return .stream(sseStream, headers: headers) + } + + // MARK: - DELETE Handler + + private func handleDelete(_ request: HTTPRequest) -> HTTPResponse { + // Validate session + let context = HTTPValidationContext( + httpMethod: "DELETE", + sessionID: sessionID, + isInitializationRequest: false, + supportedProtocolVersions: Version.supported + ) + + if let errorResponse = validationPipeline.validate(request, context: context) { + return errorResponse + } + + terminate() + + return .ok(headers: sessionHeaders()) + } + + // MARK: - Message Routing + + /// Routes a message to a specific request's SSE stream without closing it. + /// Used for server-initiated messages during request handling. + private func routeToRequestStream(_ data: Data, requestID: String) { + let eventID = storeEvent(streamID: requestID, message: data) + + guard let continuation = requestSSEContinuations[requestID] else { + logger.debug( + "No active stream for request, message stored for replay", + metadata: ["requestID": "\(requestID)"] + ) + return + } + + // Format as SSE and yield (but don't close the stream) + let sseEvent = SSEEvent.message(data: data, id: eventID) + continuation.yield(sseEvent.formatted()) + } + + /// Routes a response to a specific request's SSE stream and closes it. + /// Used for final responses to client requests. + private func routeResponse(_ data: Data, requestID: String) { + routeToRequestStream(data, requestID: requestID) + + // Response means the request is complete — close the stream + if let continuation = requestSSEContinuations[requestID] { + continuation.finish() + requestSSEContinuations.removeValue(forKey: requestID) + } + httpRequestContexts.removeValue(forKey: requestID) + } + + private func routeServerInitiatedMessage(_ data: Data) { + let eventID = storeEvent(streamID: standaloneStreamID, message: data) + + guard let continuation = standaloneSSEContinuation else { + logger.debug("No standalone GET stream connected, message stored for replay") + return + } + + let sseEvent = SSEEvent.message(data: data, id: eventID) + continuation.yield(sseEvent.formatted()) + } + + // MARK: - Resumability + + private func handleResumeRequest(lastEventID: String, request: HTTPRequest) -> HTTPResponse { + guard let replay = replayEventsAfter(lastEventID) else { + return .error( + statusCode: 400, + .invalidRequest("Invalid Last-Event-ID"), + sessionID: sessionID + ) + } + + let (sseStream, sseContinuation) = AsyncThrowingStream.makeStream() + + // Replay stored events + for (eventID, message) in replay.events { + let sseEvent = SSEEvent.message(data: message, id: eventID) + sseContinuation.yield(sseEvent.formatted()) + } + + // Re-register the stream for future messages + if replay.streamID == standaloneStreamID { + standaloneSSEContinuation = sseContinuation + } else { + requestSSEContinuations[replay.streamID] = sseContinuation + } + + // Send a new priming event so the client can resume again if disconnected + let protocolVersion = request.header(HTTPHeaderName.protocolVersion) ?? Version.latest + sendPrimingEvent( + streamID: replay.streamID, + continuation: sseContinuation, + protocolVersion: protocolVersion + ) + + var headers = sessionHeaders() + headers[HTTPHeaderName.contentType] = ContentType.sse + headers[HTTPHeaderName.cacheControl] = "no-cache, no-transform" + headers[HTTPHeaderName.connection] = "keep-alive" + + return .stream(sseStream, headers: headers) + } + + // MARK: - Internal Event Store + + private func storeEvent(streamID: String, message: Data?) -> String { + eventCounter += 1 + let eventID = "\(streamID)_\(eventCounter)" + storedEvents.append(StoredEvent(streamID: streamID, eventID: eventID, message: message)) + return eventID + } + + private func replayEventsAfter(_ lastEventID: String) -> (streamID: String, events: [(eventID: String, message: Data)])? { + guard let index = storedEvents.firstIndex(where: { $0.eventID == lastEventID }) else { + return nil + } + let streamID = storedEvents[index].streamID + let eventsToReplay = storedEvents[(index + 1)...] + .filter { $0.streamID == streamID && $0.message != nil } + .map { (eventID: $0.eventID, message: $0.message!) } + return (streamID, eventsToReplay) + } + + // MARK: - SSE Helpers + + private func sendPrimingEvent( + streamID: String, + continuation: AsyncThrowingStream.Continuation, + protocolVersion: String + ) { + // Priming events with empty data are only safe for clients >= 2025-03-26 + guard protocolVersion >= "2025-03-26" else { return } + + let primingEventID = storeEvent(streamID: streamID, message: nil) + let primingEvent = SSEEvent.priming(id: primingEventID, retry: retryInterval) + continuation.yield(primingEvent.formatted()) + } + + private func extractProtocolVersion(from body: Data, request: HTTPRequest) -> String { + // For initialize requests, extract from the request body params + if let json = try? JSONSerialization.jsonObject(with: body) as? [String: Any], + let method = json["method"] as? String, method == "initialize", + let params = json["params"] as? [String: Any], + let version = params["protocolVersion"] as? String + { + return version + } + // For other requests, use the header + return request.header(HTTPHeaderName.protocolVersion) ?? Version.latest + } + + // MARK: - Session Helpers + + private func sessionHeaders() -> [String: String] { + var headers: [String: String] = [:] + if let sessionID { + headers[HTTPHeaderName.sessionID] = sessionID + } + return headers + } + + private func isValidSessionID(_ id: String) -> Bool { + guard !id.isEmpty else { return false } + return id.utf8.allSatisfy { $0 >= 0x21 && $0 <= 0x7E } + } + + // MARK: - Stream Control + + /// Closes the SSE stream for a specific request without sending a response. + /// + /// Used to trigger client reconnection mid-call (SEP-1699). The response, + /// when eventually sent, will be stored and replayed to the reconnected stream. + package func closeSSEStream(forRequestID requestID: String) { + guard let continuation = requestSSEContinuations[requestID] else { return } + continuation.finish() + requestSSEContinuations.removeValue(forKey: requestID) + httpRequestContexts.removeValue(forKey: requestID) + } + + // MARK: - HTTPContextProviding + + public func httpRequestContext(for id: ID) -> HTTPRequest? { + httpRequestContexts[id.description] + } + + // MARK: - Termination + + /// Terminates the session, closing all active streams. + /// After termination, all requests receive 404 Not Found. + private func terminate() { + guard !terminated else { return } + terminated = true + + logger.info("Terminating session", metadata: ["sessionID": "\(sessionID ?? "none")"]) + + // Close all request SSE streams + for (_, continuation) in requestSSEContinuations { + continuation.finish() + } + requestSSEContinuations.removeAll() + httpRequestContexts.removeAll() + + // Close standalone GET stream + standaloneSSEContinuation?.finish() + standaloneSSEContinuation = nil + + // Clear stored events + storedEvents.removeAll() + + // Close incoming stream + incomingContinuation.finish() + } +} diff --git a/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift b/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift new file mode 100644 index 0000000..abe9576 --- /dev/null +++ b/Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift @@ -0,0 +1,271 @@ +import Foundation +import Logging + +/// A stateless HTTP server transport that returns single JSON responses. +/// +/// This transport implements a minimal subset of the MCP Streamable HTTP specification: +/// - No session management (no `Mcp-Session-Id` header) +/// - POST requests receive direct JSON responses (no SSE streaming) +/// - GET and DELETE requests return 405 Method Not Allowed +/// +/// ## Usage +/// +/// ```swift +/// let transport = StatelessHTTPServerTransport() +/// +/// // Start the MCP server with this transport +/// try await server.start(transport: transport) +/// +/// // In your HTTP framework handler: +/// let response = await transport.handleRequest(httpRequest) +/// // Convert response to your framework's response type and return it +/// ``` +/// +/// ## When to Use +/// +/// Use this transport when: +/// - You don't need server-initiated messages (no GET SSE stream) +/// - You want simple request-response semantics +/// - Session management is handled externally or not needed +/// +/// For full streaming and session support, use ``StatefulHTTPServerTransport`` instead. +public actor StatelessHTTPServerTransport: Transport, HTTPContextProviding { + public nonisolated let logger: Logger + + // MARK: - Dependencies + + private let validationPipeline: any HTTPRequestValidationPipeline + + // MARK: - State + + private var terminated = false + private var started = false + + // MARK: - Incoming message stream (client → server) + + private let incomingStream: AsyncThrowingStream + private let incomingContinuation: AsyncThrowingStream.Continuation + + // MARK: - Response waiters + + /// Maps request ID → continuation waiting for the server's response. + /// When the server calls `send()` with a response, the matching continuation is resumed. + private var responseWaiters: [String: CheckedContinuation] = [:] + + /// Maps request ID → originating HTTP request, surfaced to handlers via + /// ``Server/currentHTTPContext``. Entries live only while a JSON-RPC request + /// is in flight. + private var httpRequestContexts: [String: HTTPRequest] = [:] + + // MARK: - Init + + /// Creates a new stateless HTTP server transport. + /// + /// - Parameters: + /// - validationPipeline: Custom validation pipeline. If `nil`, uses sensible defaults: + /// origin validation (localhost), Accept header (JSON only), Content-Type, + /// and protocol version validation. + /// - logger: Optional logger. If `nil`, a no-op logger is used. + public init( + validationPipeline: (any HTTPRequestValidationPipeline)? = nil, + logger: Logger? = nil + ) { + self.validationPipeline = validationPipeline ?? StandardValidationPipeline(validators: [ + OriginValidator.localhost(), + AcceptHeaderValidator(mode: .jsonOnly), + ContentTypeValidator(), + ProtocolVersionValidator(), + ]) + self.logger = logger ?? Logger( + label: "mcp.transport.http.server.stateless", + factory: { _ in SwiftLogNoOpLogHandler() } + ) + + let (stream, continuation) = AsyncThrowingStream.makeStream() + self.incomingStream = stream + self.incomingContinuation = continuation + } + + // MARK: - Transport Conformance + + public func connect() async throws { + guard !started else { + throw MCPError.internalError("Transport already started") + } + started = true + logger.debug("Stateless HTTP server transport started") + } + + public func disconnect() async { + await terminate() + } + + /// Routes outgoing server messages to the appropriate waiting HTTP handler. + /// + /// - Responses are matched by JSON-RPC ID and delivered to the waiting `handleRequest` call. + /// - Notifications and server-initiated requests are logged and dropped + /// (no streaming channel available in stateless mode). + public func send(_ data: Data) async throws { + guard !terminated else { + throw MCPError.connectionClosed + } + + guard let kind = JSONRPCMessageKind(data: data) else { + logger.warning("Could not classify outgoing message for routing") + return + } + + switch kind { + case .response(let id): + guard let continuation = responseWaiters.removeValue(forKey: id) else { + logger.debug( + "No waiter for response, may have timed out", + metadata: ["requestID": "\(id)"] + ) + return + } + continuation.resume(returning: data) + + case .notification(let method): + logger.debug( + "Server-initiated notification dropped in stateless mode (no GET SSE stream)", + metadata: ["method": "\(method)"] + ) + + case .request(_, let method): + logger.debug( + "Server-initiated request dropped in stateless mode (no GET SSE stream)", + metadata: ["method": "\(method)"] + ) + } + } + + public func receive() -> AsyncThrowingStream { + incomingStream + } + + // MARK: - HTTP Request Handler + + /// Handles an incoming HTTP request from the framework adapter. + /// + /// Only POST is supported: + /// - **POST**: JSON-RPC messages (requests, notifications) + /// - **GET**: 405 Method Not Allowed + /// - **DELETE**: 405 Method Not Allowed + /// - Others: 405 Method Not Allowed + public func handleRequest(_ request: HTTPRequest) async -> HTTPResponse { + if terminated { + return .error( + statusCode: 404, + .invalidRequest("Not Found: Transport has been terminated") + ) + } + + switch request.method.uppercased() { + case "POST": + return await handlePost(request) + default: + return .error( + statusCode: 405, + .invalidRequest("Method Not Allowed"), + extraHeaders: [HTTPHeaderName.allow: "POST"] + ) + } + } + + // MARK: - POST Handler + + private func handlePost(_ request: HTTPRequest) async -> HTTPResponse { + // Parse body first to determine message type + guard let body = request.body, !body.isEmpty else { + return .error( + statusCode: 400, + .parseError("Empty request body") + ) + } + + guard let messageKind = JSONRPCMessageKind(data: body) else { + return .error( + statusCode: 400, + .parseError("Invalid JSON-RPC message") + ) + } + + // Build validation context + let context = HTTPValidationContext( + httpMethod: "POST", + sessionID: nil, + isInitializationRequest: messageKind.isInitializeRequest, + supportedProtocolVersions: Version.supported + ) + + // Run validation pipeline + if let errorResponse = validationPipeline.validate(request, context: context) { + return errorResponse + } + + // Handle by message type + switch messageKind { + case .notification, .response: + // Yield to server and return 202 Accepted + incomingContinuation.yield(body) + return .accepted() + + case .request(let id, _): + return await handleJSONRPCRequest(body, requestID: id, request: request) + } + } + + private func handleJSONRPCRequest( + _ body: Data, + requestID: String, + request: HTTPRequest + ) async -> HTTPResponse { + httpRequestContexts[requestID] = request + // Yield the incoming message to the server + incomingContinuation.yield(body) + + // Wait for the server to process and send a response + let responseData: Data + do { + responseData = try await withCheckedThrowingContinuation { continuation in + responseWaiters[requestID] = continuation + } + } catch { + httpRequestContexts.removeValue(forKey: requestID) + return .error( + statusCode: 500, + .internalError("Error processing request: \(error.localizedDescription)") + ) + } + + httpRequestContexts.removeValue(forKey: requestID) + return .data(responseData, headers: [HTTPHeaderName.contentType: ContentType.json]) + } + + // MARK: - HTTPContextProviding + + public func httpRequestContext(for id: ID) -> HTTPRequest? { + httpRequestContexts[id.description] + } + + // MARK: - Termination + + private func terminate() async { + guard !terminated else { return } + terminated = true + + logger.debug("Stateless HTTP server transport terminated") + + // Cancel all waiting continuations + for (id, continuation) in responseWaiters { + continuation.resume(throwing: MCPError.connectionClosed) + logger.debug("Cancelled waiter for request", metadata: ["requestID": "\(id)"]) + } + responseWaiters.removeAll() + httpRequestContexts.removeAll() + + // Close incoming stream + incomingContinuation.finish() + } +} diff --git a/Sources/MCP/Base/Transports/InMemoryTransport.swift b/Sources/MCP/Base/Transports/InMemoryTransport.swift new file mode 100644 index 0000000..9234dfc --- /dev/null +++ b/Sources/MCP/Base/Transports/InMemoryTransport.swift @@ -0,0 +1,197 @@ +import Foundation +import Logging + +/// An in-memory transport implementation for direct communication within the same process. +/// +/// - Example: +/// ```swift +/// // Create a connected pair of transports +/// let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() +/// +/// // Use with client and server +/// let client = Client(name: "MyApp", version: "1.0.0") +/// let server = Server(name: "MyServer", version: "1.0.0") +/// +/// try await client.connect(transport: clientTransport) +/// try await server.connect(transport: serverTransport) +/// ``` +public actor InMemoryTransport: Transport { + /// Logger instance for transport-related events + public nonisolated let logger: Logger + + private var isConnected = false + private var pairedTransport: InMemoryTransport? + + // Message queues + private var incomingMessages: [Data] = [] + private var messageContinuation: AsyncThrowingStream.Continuation? + + /// Creates a new in-memory transport + /// + /// - Parameter logger: Optional logger instance for transport events + public init(logger: Logger? = nil) { + self.logger = + logger + ?? Logger( + label: "mcp.transport.in-memory", + factory: { _ in SwiftLogNoOpLogHandler() } + ) + } + + /// Creates a connected pair of in-memory transports + /// + /// This is the recommended way to create transports for client-server communication + /// within the same process. The returned transports are already paired and ready + /// to be connected. + /// + /// - Parameter logger: Optional logger instance shared by both transports + /// - Returns: A tuple of (clientTransport, serverTransport) ready for use + public static func createConnectedPair( + logger: Logger? = nil + ) async -> (client: InMemoryTransport, server: InMemoryTransport) { + let clientLogger: Logger + let serverLogger: Logger + + if let providedLogger = logger { + // If a logger is provided, use it directly for both transports + clientLogger = providedLogger + serverLogger = providedLogger + } else { + // Create default loggers with appropriate labels + clientLogger = Logger( + label: "mcp.transport.in-memory.client", + factory: { _ in SwiftLogNoOpLogHandler() } + ) + serverLogger = Logger( + label: "mcp.transport.in-memory.server", + factory: { _ in SwiftLogNoOpLogHandler() } + ) + } + + let clientTransport = InMemoryTransport(logger: clientLogger) + let serverTransport = InMemoryTransport(logger: serverLogger) + + // Perform pairing + await clientTransport.pair(with: serverTransport) + await serverTransport.pair(with: clientTransport) + + return (clientTransport, serverTransport) + } + + /// Pairs this transport with another for bidirectional communication + /// + /// - Parameter other: The transport to pair with + /// - Important: This method should typically not be called directly. + /// Use `createConnectedPair()` instead. + private func pair(with other: InMemoryTransport) { + self.pairedTransport = other + } + + /// Establishes connection with the transport + /// + /// For in-memory transports, this validates that the transport is properly + /// paired and sets up the message stream. + /// + /// - Throws: MCPError.internalError if the transport is not paired + public func connect() async throws { + guard !isConnected else { + logger.debug("Transport already connected") + return + } + + guard pairedTransport != nil else { + throw MCPError.internalError( + "Transport not paired. Use createConnectedPair() to create paired transports.") + } + + isConnected = true + logger.info("Transport connected successfully") + } + + /// Disconnects from the transport + /// + /// This closes the message stream and marks the transport as disconnected. + public func disconnect() async { + guard isConnected else { return } + + isConnected = false + messageContinuation?.finish() + messageContinuation = nil + + // Notify paired transport of disconnection + if let paired = pairedTransport { + await paired.handlePeerDisconnection() + } + + logger.info("Transport disconnected") + } + + /// Handles disconnection from the paired transport + private func handlePeerDisconnection() { + if isConnected { + messageContinuation?.finish(throwing: MCPError.connectionClosed) + messageContinuation = nil + isConnected = false + logger.info("Peer transport disconnected") + } + } + + /// Sends a message to the paired transport + /// + /// Messages are delivered directly to the paired transport's receive queue + /// without any additional encoding or framing. + /// + /// - Parameter data: The message data to send + /// - Throws: MCPError.internalError if not connected or no paired transport + public func send(_ data: Data) async throws { + guard isConnected else { + throw MCPError.internalError("Transport not connected") + } + + guard let paired = pairedTransport else { + throw MCPError.internalError("No paired transport") + } + + logger.debug("Sending message", metadata: ["size": "\(data.count)"]) + + // Deliver message to paired transport + await paired.deliverMessage(data) + } + + /// Delivers a message from the paired transport + private func deliverMessage(_ data: Data) { + guard isConnected else { + logger.warning("Received message while disconnected") + return + } + + logger.debug("Message received", metadata: ["size": "\(data.count)"]) + + if let continuation = messageContinuation { + continuation.yield(data) + } else { + // Queue message if stream not yet created + incomingMessages.append(data) + } + } + + /// Receives messages from the paired transport + /// + /// - Returns: An AsyncThrowingStream of Data 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) + } + self.incomingMessages.removeAll() + + // Check if already disconnected + if !self.isConnected { + continuation.finish() + } + } + } +} diff --git a/Sources/MCP/Base/Transports/NetworkTransport.swift b/Sources/MCP/Base/Transports/NetworkTransport.swift new file mode 100644 index 0000000..7c4d00f --- /dev/null +++ b/Sources/MCP/Base/Transports/NetworkTransport.swift @@ -0,0 +1,801 @@ +import Foundation +import Logging + +#if canImport(Network) + import Network + + /// Protocol that abstracts the Network.NWConnection functionality needed for NetworkTransport + @preconcurrency protocol NetworkConnectionProtocol { + var state: NWConnection.State { get } + var stateUpdateHandler: ((@Sendable (NWConnection.State) -> Void))? { get set } + + func start(queue: DispatchQueue) + func cancel() + func send( + content: Data?, contentContext: NWConnection.ContentContext, isComplete: Bool, + completion: NWConnection.SendCompletion) + func receive( + minimumIncompleteLength: Int, maximumLength: Int, + completion: @escaping @Sendable ( + Data?, NWConnection.ContentContext?, Bool, NWError? + ) -> Void) + } + + /// Extension to conform NWConnection to internal NetworkConnectionProtocol + extension NWConnection: NetworkConnectionProtocol {} + + /// An implementation of a custom MCP transport using Apple's Network framework. + /// + /// This transport allows MCP clients and servers to communicate over TCP/UDP connections + /// using Apple's Network framework. + /// + /// - Important: This transport is available exclusively on Apple platforms + /// (macOS, iOS, watchOS, tvOS, visionOS) as it depends on the Network framework. + /// + /// ## Example Usage + /// + /// ```swift + /// import MCP + /// import Network + /// + /// // Create a TCP connection to a server + /// let connection = NWConnection( + /// host: NWEndpoint.Host("localhost"), + /// port: NWEndpoint.Port(8080)!, + /// using: .tcp + /// ) + /// + /// // Initialize the transport with the connection + /// let transport = NetworkTransport(connection: connection) + /// + /// // For large messages (e.g., images), configure unlimited buffer size + /// let largeBufferTransport = NetworkTransport( + /// connection: connection, + /// bufferConfig: .unlimited + /// ) + /// + /// // Use the transport with an MCP client + /// let client = Client(name: "MyApp", version: "1.0.0") + /// try await client.connect(transport: transport) + /// ``` + public actor NetworkTransport: Transport { + /// Represents a heartbeat message for connection health monitoring. + public struct Heartbeat: RawRepresentable, Hashable, Sendable { + /// Magic bytes used to identify a heartbeat message. + private static let magicBytes: [UInt8] = [0xF0, 0x9F, 0x92, 0x93] + + /// The timestamp of when the heartbeat was created. + public let timestamp: Date + + /// Creates a new heartbeat with the current timestamp. + public init() { + self.timestamp = Date() + } + + /// Creates a heartbeat with a specific timestamp. + /// + /// - Parameter timestamp: The timestamp for the heartbeat. + public init(timestamp: Date) { + self.timestamp = timestamp + } + + // MARK: - RawRepresentable + + public typealias RawValue = [UInt8] + + /// Creates a heartbeat from its raw representation. + /// + /// - Parameter rawValue: The raw bytes of the heartbeat message. + /// - Returns: A heartbeat if the raw value is valid, nil otherwise. + public init?(rawValue: [UInt8]) { + // Check if the data has the correct format (magic bytes + timestamp) + guard rawValue.count >= 12, + rawValue.prefix(4).elementsEqual(Self.magicBytes) + else { + return nil + } + + // Extract the timestamp + let timestampData = Data(rawValue[4..<12]) + let timestamp = timestampData.withUnsafeBytes { + $0.load(as: UInt64.self) + } + + self.timestamp = Date( + timeIntervalSinceReferenceDate: TimeInterval(timestamp) / 1000.0) + } + + /// Converts the heartbeat to its raw representation. + public var rawValue: [UInt8] { + var result = Data(Self.magicBytes) + + // Add timestamp (milliseconds since reference date) + let timestamp = UInt64(self.timestamp.timeIntervalSinceReferenceDate * 1000) + withUnsafeBytes(of: timestamp) { buffer in + result.append(contentsOf: buffer) + } + + return Array(result) + } + + /// Converts the heartbeat to Data. + public var data: Data { + return Data(self.rawValue) + } + + /// Checks if the given data represents a heartbeat message. + /// + /// - Parameter data: The data to check. + /// - Returns: true if the data is a heartbeat message, false otherwise. + public static func isHeartbeat(_ data: Data) -> Bool { + guard data.count >= 4 else { + return false + } + + return data.prefix(4).elementsEqual(Self.magicBytes) + } + + /// Attempts to parse a heartbeat from the given data. + /// + /// - Parameter data: The data to parse. + /// - Returns: A heartbeat if the data is valid, nil otherwise. + public static func from(data: Data) -> Heartbeat? { + guard data.count >= 12 else { + return nil + } + + return Heartbeat(rawValue: Array(data)) + } + } + + /// Configuration for heartbeat behavior. + public struct HeartbeatConfiguration: Hashable, Sendable { + /// Whether heartbeats are enabled. + public let enabled: Bool + /// Interval between heartbeats in seconds. + public let interval: TimeInterval + + /// Creates a new heartbeat configuration. + /// + /// - Parameters: + /// - enabled: Whether heartbeats are enabled (default: true) + /// - interval: Interval in seconds between heartbeats (default: 15.0) + public init(enabled: Bool = true, interval: TimeInterval = 15.0) { + self.enabled = enabled + self.interval = interval + } + + /// Default heartbeat configuration. + public static let `default` = HeartbeatConfiguration() + + /// Configuration with heartbeats disabled. + public static let disabled = HeartbeatConfiguration(enabled: false) + } + + /// Configuration for connection retry behavior. + public struct ReconnectionConfiguration: Hashable, Sendable { + /// Whether the transport should attempt to reconnect on failure. + public let enabled: Bool + /// Maximum number of reconnection attempts. + public let maxAttempts: Int + /// Multiplier for exponential backoff on reconnect. + public let backoffMultiplier: Double + + /// Creates a new reconnection configuration. + /// + /// - Parameters: + /// - enabled: Whether reconnection should be attempted on failure (default: true) + /// - maxAttempts: Maximum number of reconnection attempts (default: 5) + /// - backoffMultiplier: Multiplier for exponential backoff on reconnect (default: 1.5) + public init( + enabled: Bool = true, + maxAttempts: Int = 5, + backoffMultiplier: Double = 1.5 + ) { + self.enabled = enabled + self.maxAttempts = maxAttempts + self.backoffMultiplier = backoffMultiplier + } + + /// Default reconnection configuration. + public static let `default` = ReconnectionConfiguration() + + /// Configuration with reconnection disabled. + public static let disabled = ReconnectionConfiguration(enabled: false) + + /// Calculates the backoff delay for a given attempt number. + /// + /// - Parameter attempt: The current attempt number (1-based) + /// - Returns: The delay in seconds before the next attempt + public func backoffDelay(for attempt: Int) -> TimeInterval { + let baseDelay = 0.5 // 500ms + return baseDelay * pow(backoffMultiplier, Double(attempt - 1)) + } + } + + /// Configuration for buffer behavior. + public struct BufferConfiguration: Hashable, Sendable { + /// Maximum buffer size for receiving data chunks. + /// Set to nil for unlimited (uses system default). + public let maxReceiveBufferSize: Int? + + /// Creates a new buffer configuration. + /// + /// - Parameter maxReceiveBufferSize: Maximum buffer size in bytes (default: 10MB, nil for unlimited) + public init(maxReceiveBufferSize: Int? = 10 * 1024 * 1024) { + self.maxReceiveBufferSize = maxReceiveBufferSize + } + + /// Default buffer configuration with 10MB limit. + public static let `default` = BufferConfiguration() + + /// Configuration with no buffer size limit. + public static let unlimited = BufferConfiguration(maxReceiveBufferSize: nil) + } + + // State tracking + private var isConnected = false + private var isStopping = false + private var reconnectAttempt = 0 + private var heartbeatTask: Task? + private var lastHeartbeatTime: Date? + 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 + + /// Logger instance for transport-related events + public nonisolated let logger: Logger + + // Configuration + private let heartbeatConfig: HeartbeatConfiguration + private let reconnectionConfig: ReconnectionConfiguration + private let bufferConfig: BufferConfiguration + + /// Creates a new NetworkTransport with the specified NWConnection + /// + /// - Parameters: + /// - connection: The NWConnection to use for communication + /// - logger: Optional logger instance for transport events + /// - reconnectionConfig: Configuration for reconnection behavior (default: .default) + /// - heartbeatConfig: Configuration for heartbeat behavior (default: .default) + /// - bufferConfig: Configuration for buffer behavior (default: .default) + public init( + connection: NWConnection, + logger: Logger? = nil, + heartbeatConfig: HeartbeatConfiguration = .default, + reconnectionConfig: ReconnectionConfiguration = .default, + bufferConfig: BufferConfiguration = .default + ) { + self.init( + connection, + logger: logger, + heartbeatConfig: heartbeatConfig, + reconnectionConfig: reconnectionConfig, + bufferConfig: bufferConfig + ) + } + + init( + _ connection: NetworkConnectionProtocol, + logger: Logger? = nil, + heartbeatConfig: HeartbeatConfiguration = .default, + reconnectionConfig: ReconnectionConfiguration = .default, + bufferConfig: BufferConfiguration = .default + ) { + self.connection = connection + self.logger = + logger + ?? Logger( + label: "mcp.transport.network", + factory: { _ in SwiftLogNoOpLogHandler() } + ) + self.reconnectionConfig = reconnectionConfig + self.heartbeatConfig = heartbeatConfig + self.bufferConfig = bufferConfig + + // Create message stream + var continuation: AsyncThrowingStream.Continuation! + self.messageStream = AsyncThrowingStream { continuation = $0 } + self.messageContinuation = continuation + } + + /// Establishes connection with the transport + /// + /// This initiates the NWConnection and waits for it to become ready. + /// Once the connection is established, it starts the message receiving loop. + /// + /// - Throws: Error if the connection fails to establish + public func connect() async throws { + guard !isConnected else { return } + + // Reset state for fresh connection + isStopping = false + reconnectAttempt = 0 + + // Retry loop with exponential backoff + while !isConnected && reconnectAttempt <= reconnectionConfig.maxAttempts { + do { + try await attemptConnection() + return // Success + } catch { + guard !isStopping, + reconnectionConfig.enabled, + reconnectAttempt < reconnectionConfig.maxAttempts + else { + throw error + } + + // Schedule retry with backoff + reconnectAttempt += 1 + let delay = reconnectionConfig.backoffDelay(for: reconnectAttempt) + logger.debug( + "Attempting reconnection (\(reconnectAttempt)/\(reconnectionConfig.maxAttempts))..." + ) + + try await Task.sleep(for: .seconds(delay)) + connection.cancel() + } + } + + throw MCPError.internalError("Failed to connect after \(reconnectAttempt) attempts") + } + + /// Attempts a single connection + /// + /// Creates a stream for connection state changes and waits for a terminal state. + /// + /// - Throws: Error if the connection fails + private func attemptConnection() async throws { + // Create stream for state changes with proper cleanup + let stateStream = AsyncStream { continuation in + continuation.onTermination = { [weak self] _ in + self?.connection.stateUpdateHandler = nil + } + + connection.stateUpdateHandler = { state in + continuation.yield(state) + switch state { + case .ready, .failed, .cancelled: + continuation.finish() + default: + break + } + } + } + + connection.start(queue: .main) + + // Process states until terminal state + for await state in stateStream { + switch state { + case .ready: + await handleConnectionReady() + return // Success - exits loop and stream + + case .failed(let error): + logger.error("Connection failed: \(error)") + throw error // Exits loop and stream + + case .cancelled: + logger.warning("Connection 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") + } + } + } + + /// Handles when the connection reaches the ready state + private func handleConnectionReady() async { + isConnected = true + reconnectAttempt = 0 + logger.debug("Network transport connected successfully") + + // Start the receive loop after connection is established + Task { await self.receiveLoop() } + + // Start heartbeat task if enabled + if heartbeatConfig.enabled { + startHeartbeat() + } + } + + /// Starts a task to periodically send heartbeats to check connection health + private func startHeartbeat() { + // Cancel any existing heartbeat task + heartbeatTask?.cancel() + + // Start a new heartbeat task + heartbeatTask = Task { [weak self] in + guard let self = self else { return } + + // Initial delay before starting heartbeats + try? await Task.sleep(for: .seconds(1)) + + while !Task.isCancelled { + do { + // Check actor-isolated properties first + let isStopping = await self.isStopping + let isConnected = await self.isConnected + + guard !isStopping && isConnected else { break } + + try await self.sendHeartbeat() + try await Task.sleep(for: .seconds(self.heartbeatConfig.interval)) + } catch { + // If heartbeat fails, log and retry after a shorter interval + self.logger.warning("Heartbeat failed: \(error)") + try? await Task.sleep(for: .seconds(2)) + } + } + } + } + + /// Sends a heartbeat message to verify connection health + private func sendHeartbeat() async throws { + guard isConnected && !isStopping else { return } + + // 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 { + continuation.resume(throwing: MCPError.internalError("Transport deallocated")) + return + } + + connection.send( + content: Heartbeat().data, + contentContext: .defaultMessage, + isComplete: true, + completion: .contentProcessed { [weak self] error in + if let error = error { + continuation.resume(throwing: error) + } else { + Task { [weak self] in + await self?.setLastHeartbeatTime(Date()) + } + continuation.resume() + } + }) + } + + logger.trace("Heartbeat sent") + } + + + /// Disconnects from the transport + /// + /// This cancels the NWConnection, finalizes the message stream, + /// and releases associated resources. + public func disconnect() async { + guard isConnected else { return } + + // Mark as stopping to prevent reconnection attempts during disconnect + isStopping = true + isConnected = false + + // Cancel heartbeat task if it exists + heartbeatTask?.cancel() + heartbeatTask = nil + + connection.cancel() + messageContinuation.finish() + logger.debug("Network transport disconnected") + } + + /// Sends data through the network connection + /// + /// This sends a JSON-RPC message through the NWConnection, adding a newline + /// delimiter to mark the end of the message. + /// + /// - Parameter message: The JSON-RPC message to send + /// - Throws: MCPError for transport failures or connection issues + public func send(_ message: Data) async throws { + guard isConnected else { + throw MCPError.internalError("Transport not connected") + } + + // Add newline as delimiter + var messageWithNewline = message + messageWithNewline.append(UInt8(ascii: "\n")) + + // Use a main-isolated variable to track continuation state + let sendContinuationResumed = MainFlag() + + try await withCheckedThrowingContinuation { + [weak self] (continuation: CheckedContinuation) in + guard let self = self else { + continuation.resume(throwing: MCPError.internalError("Transport deallocated")) + return + } + + connection.send( + content: messageWithNewline, + contentContext: .defaultMessage, + isComplete: true, + completion: .contentProcessed { [weak self] error in + guard let self = self else { return } + + Task { @MainActor in + if !sendContinuationResumed.flag { + sendContinuationResumed.flag = 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() + } + } + } + } + } + + continuation.resume( + throwing: MCPError.internalError("Send error: \(error)")) + } else { + continuation.resume() + } + } + } + }) + } + } + + /// Receives data in an async sequence + /// + /// This returns an AsyncThrowingStream that emits Data objects representing + /// each JSON-RPC message received from the network connection. + /// + /// - Returns: An AsyncThrowingStream of Data objects + public func receive() -> AsyncThrowingStream { + return messageStream + } + + /// Continuous loop to receive and process incoming messages + /// + /// This method runs continuously while the connection is active, + /// receiving data and yielding complete messages to the message stream. + /// Messages are delimited by newline characters. + private func receiveLoop() async { + var buffer = Data() + var consecutiveEmptyReads = 0 + let maxConsecutiveEmptyReads = 5 + + while isConnected && !Task.isCancelled && !isStopping { + do { + let newData = try await receiveData() + + // Check for EOF or empty data + if newData.isEmpty { + consecutiveEmptyReads += 1 + + if consecutiveEmptyReads >= maxConsecutiveEmptyReads { + logger.warning( + "Multiple consecutive empty reads (\(consecutiveEmptyReads)), possible connection issue" + ) + + // Check connection state + if connection.state != .ready { + logger.warning("Connection no longer ready, exiting receive loop") + break + } + } + + // Brief pause before retry + try await Task.sleep(for: .milliseconds(100)) + continue + } + + // Check if this is a heartbeat message + if Heartbeat.isHeartbeat(newData) { + logger.trace("Received heartbeat from peer") + + // Extract timestamp if available + if let heartbeat = Heartbeat.from(data: newData) { + logger.trace("Heartbeat timestamp: \(heartbeat.timestamp)") + } + + // Reset the counter since we got valid data + consecutiveEmptyReads = 0 + continue // Skip regular message processing for heartbeats + } + + // Reset counter on successful data read + consecutiveEmptyReads = 0 + buffer.append(newData) + + // Process complete messages + while let newlineIndex = buffer.firstIndex(of: UInt8(ascii: "\n")) { + let messageData = buffer[.. Data { + // Use a main-isolated variable to track continuation state + let receiveContinuationResumed = MainFlag() + + return try await withCheckedThrowingContinuation { + [weak self] (continuation: CheckedContinuation) in + guard let self = self else { + continuation.resume(throwing: MCPError.internalError("Transport deallocated")) + return + } + + let maxLength = bufferConfig.maxReceiveBufferSize ?? Int.max + connection.receive(minimumIncompleteLength: 1, maximumLength: maxLength) { + content, _, isComplete, error in + Task { @MainActor in + if !receiveContinuationResumed.flag { + receiveContinuationResumed.flag = 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()) + } + } + } + } + } + } + + private func setLastHeartbeatTime(_ time: Date) { + self.lastHeartbeatTime = time + } + + private func setIsConnected(_ connected: Bool) { + self.isConnected = connected + } + } + @MainActor final class MainFlag { var flag = false } + + extension NWError { + /// Whether this error indicates a connection has been lost or reset. + fileprivate var isConnectionLost: Bool { + let nsError = self as NSError + return nsError.code == 57 // Socket is not connected (EHOSTUNREACH or ENOTCONN) + || nsError.code == 54 // Connection reset by peer (ECONNRESET) + } + } +#endif diff --git a/Sources/MCP/Base/Transports/StdioTransport.swift b/Sources/MCP/Base/Transports/StdioTransport.swift new file mode 100644 index 0000000..45522fa --- /dev/null +++ b/Sources/MCP/Base/Transports/StdioTransport.swift @@ -0,0 +1,235 @@ +import Logging + +import struct Foundation.Data + +#if canImport(System) + import System +#else + @preconcurrency import SystemPackage +#endif + +// Import for specific low-level operations not yet in Swift System +#if canImport(Darwin) + import Darwin.POSIX +#elseif canImport(Glibc) + import Glibc +#elseif canImport(Musl) + import Musl +#endif + +#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-06-18/basic/transports#stdio) + /// specification from the Model Context Protocol. + /// + /// The stdio transport works by: + /// - Reading JSON-RPC messages from standard input + /// - Writing JSON-RPC messages to standard output + /// - Using newline characters as message delimiters + /// - Supporting non-blocking I/O operations + /// + /// This transport is the recommended option for most MCP applications due to its + /// simplicity and broad platform support. + /// + /// - Important: This transport is available on Apple platforms and Linux distributions with glibc + /// (Ubuntu, Debian, Fedora, CentOS, RHEL). + /// + /// ## Example Usage + /// + /// ```swift + /// import MCP + /// + /// // Initialize the client + /// let client = Client(name: "MyApp", version: "1.0.0") + /// + /// // Create a transport and connect + /// let transport = StdioTransport() + /// try await client.connect(transport: transport) + /// ``` + public actor StdioTransport: Transport { + private let input: FileDescriptor + private let output: FileDescriptor + /// Logger instance for transport-related events + public nonisolated let logger: Logger + + private var isConnected = false + private let messageStream: AsyncThrowingStream + private let messageContinuation: AsyncThrowingStream.Continuation + + /// Creates a new stdio transport with the specified file descriptors + /// + /// - Parameters: + /// - input: File descriptor for reading (defaults to standard input) + /// - output: File descriptor for writing (defaults to standard output) + /// - logger: Optional logger instance for transport events + public init( + input: FileDescriptor = FileDescriptor.standardInput, + output: FileDescriptor = FileDescriptor.standardOutput, + logger: Logger? = nil + ) { + self.input = input + self.output = output + self.logger = + logger + ?? Logger( + label: "mcp.transport.stdio", + factory: { _ in SwiftLogNoOpLogHandler() }) + + // Create message stream + var continuation: AsyncThrowingStream.Continuation! + self.messageStream = AsyncThrowingStream { continuation = $0 } + self.messageContinuation = continuation + } + + /// Establishes connection with the transport + /// + /// This method configures the file descriptors for non-blocking I/O + /// and starts the background message reading loop. + /// + /// - Throws: Error if the file descriptors cannot be configured + public func connect() async throws { + guard !isConnected else { return } + + // Set non-blocking mode + try setNonBlocking(fileDescriptor: input) + try setNonBlocking(fileDescriptor: output) + + isConnected = true + logger.debug("Transport connected successfully") + + // Start reading loop in background + Task { + await readLoop() + } + } + + /// Configures a file descriptor for non-blocking I/O + /// + /// - Parameter fileDescriptor: The file descriptor to configure + /// - Throws: Error if the operation fails + private func setNonBlocking(fileDescriptor: FileDescriptor) throws { + #if canImport(Darwin) || canImport(Glibc) || canImport(Musl) + // Get current flags + let flags = fcntl(fileDescriptor.rawValue, F_GETFL) + guard flags >= 0 else { + throw MCPError.transportError(Errno(rawValue: CInt(errno))) + } + + // Set non-blocking flag + let result = fcntl(fileDescriptor.rawValue, F_SETFL, flags | O_NONBLOCK) + guard result >= 0 else { + throw MCPError.transportError(Errno(rawValue: CInt(errno))) + } + #else + // For platforms where non-blocking operations aren't supported + throw MCPError.internalError( + "Setting non-blocking mode not supported on this platform") + #endif + } + + /// Continuous loop that reads and processes incoming messages + /// + /// This method runs in the background while the transport is connected, + /// parsing complete messages delimited by newlines and yielding them + /// to the message stream. + private func readLoop() async { + let bufferSize = 4096 + var buffer = [UInt8](repeating: 0, count: bufferSize) + var pendingData = Data() + + while isConnected && !Task.isCancelled { + do { + let bytesRead = try buffer.withUnsafeMutableBufferPointer { pointer in + try input.read(into: UnsafeMutableRawBufferPointer(pointer)) + } + + if bytesRead == 0 { + logger.notice("EOF received") + break + } + + pendingData.append(Data(buffer[.. 0 { + remaining = remaining.dropFirst(written) + } + } catch let error where MCPError.isResourceTemporarilyUnavailable(error) { + try await Task.sleep(for: .milliseconds(10)) + continue + } catch { + throw MCPError.transportError(error) + } + } + } + + /// Receives messages from the transport. + /// + /// Messages may be individual JSON-RPC requests, notifications, responses, + /// 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 { + return messageStream + } + } +#endif diff --git a/Sources/MCP/Base/UnitInterval.swift b/Sources/MCP/Base/UnitInterval.swift new file mode 100644 index 0000000..07f94e7 --- /dev/null +++ b/Sources/MCP/Base/UnitInterval.swift @@ -0,0 +1,126 @@ +/// A value constrained to the range 0.0 to 1.0, inclusive. +/// +/// `UnitInterval` represents a normalized value that is guaranteed to be within +/// the unit interval [0, 1]. This type is commonly used for representing +/// priorities in sampling request model preferences. +/// +/// The type provides safe initialization that returns `nil` for values outside +/// the valid range, ensuring that all instances contain valid unit interval values. +/// +/// - Example: +/// ```swift +/// let zero: UnitInterval = 0 // 0.0 +/// let half = UnitInterval(0.5)! // 0.5 +/// let one: UnitInterval = 1.0 // 1.0 +/// let invalid = UnitInterval(1.5) // nil +/// ``` +public struct UnitInterval: Hashable, Sendable { + private let value: Double + + /// Creates a unit interval value from a `Double`. + /// + /// - Parameter value: A double value that must be in the range 0.0...1.0 + /// - Returns: A `UnitInterval` instance if the value is valid, `nil` otherwise + /// + /// - Example: + /// ```swift + /// let valid = UnitInterval(0.75) // Optional(0.75) + /// let invalid = UnitInterval(-0.1) // nil + /// let boundary = UnitInterval(1.0) // Optional(1.0) + /// ``` + public init?(_ value: Double) { + guard (0...1).contains(value) else { return nil } + self.value = value + } + + /// The underlying double value. + /// + /// This property provides access to the raw double value that is guaranteed + /// to be within the range [0, 1]. + /// + /// - Returns: The double value between 0.0 and 1.0, inclusive + public var doubleValue: Double { value } +} + +// MARK: - Comparable + +extension UnitInterval: Comparable { + public static func < (lhs: UnitInterval, rhs: UnitInterval) -> Bool { + lhs.value < rhs.value + } +} + +// MARK: - CustomStringConvertible + +extension UnitInterval: CustomStringConvertible { + public var description: String { "\(value)" } +} + +// MARK: - Codable + +extension UnitInterval: Codable { + public init(from decoder: Decoder) throws { + let container = try decoder.singleValueContainer() + let doubleValue = try container.decode(Double.self) + guard let interval = UnitInterval(doubleValue) else { + throw DecodingError.dataCorrupted( + DecodingError.Context( + codingPath: decoder.codingPath, + debugDescription: "Value \(doubleValue) is not in range 0...1") + ) + } + self = interval + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.singleValueContainer() + try container.encode(value) + } +} + +// MARK: - ExpressibleByFloatLiteral + +extension UnitInterval: ExpressibleByFloatLiteral { + /// Creates a unit interval from a floating-point literal. + /// + /// This initializer allows you to create `UnitInterval` instances using + /// floating-point literals. The literal value must be in the range [0, 1] + /// or a runtime error will occur. + /// + /// - Parameter value: A floating-point literal between 0.0 and 1.0 + /// + /// - Warning: This initializer will crash if the literal is outside the valid range. + /// Use the failable initializer `init(_:)` for runtime validation. + /// + /// - Example: + /// ```swift + /// let quarter: UnitInterval = 0.25 + /// let half: UnitInterval = 0.5 + /// ``` + public init(floatLiteral value: Double) { + self.init(value)! + } +} + +// MARK: - ExpressibleByIntegerLiteral + +extension UnitInterval: ExpressibleByIntegerLiteral { + /// Creates a unit interval from an integer literal. + /// + /// This initializer allows you to create `UnitInterval` instances using + /// integer literals. Only the values 0 and 1 are valid. + /// + /// - Parameter value: An integer literal, either 0 or 1 + /// + /// - Warning: This initializer will crash if the literal is outside the valid range. + /// Use the failable initializer `init(_:)` for runtime validation. + /// + /// - Example: + /// ```swift + /// let zero: UnitInterval = 0 + /// let one: UnitInterval = 1 + /// ``` + public init(integerLiteral value: Int) { + self.init(Double(value))! + } +} diff --git a/Sources/MCP/Base/Utilities/Cancellation.swift b/Sources/MCP/Base/Utilities/Cancellation.swift new file mode 100644 index 0000000..a8635ce --- /dev/null +++ b/Sources/MCP/Base/Utilities/Cancellation.swift @@ -0,0 +1,29 @@ +import Foundation + +/// The Model Context Protocol supports cancellation of requests through +/// notification messages. Either side can send a cancellation notification +/// to indicate that a previously-issued request should be cancelled. +/// +/// Cancellation is advisory: the receiver should make a best-effort attempt +/// to cancel the operation, but may not always be able to do so. +/// +/// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/cancellation +public struct CancelledNotification: Notification { + public static let name: String = "notifications/cancelled" + + public struct Parameters: Hashable, Codable, Sendable { + /// The ID of the request to cancel. + /// + /// This MUST correspond to the ID of a request previously issued + /// in the same direction. Optional per the spec (`requestId?`). + public let requestId: ID? + + /// An optional human-readable reason for the cancellation. + public let reason: String? + + public init(requestId: ID? = nil, reason: String? = nil) { + self.requestId = requestId + self.reason = reason + } + } +} diff --git a/Sources/MCP/Base/Utilities/Progress.swift b/Sources/MCP/Base/Utilities/Progress.swift new file mode 100644 index 0000000..df2ab66 --- /dev/null +++ b/Sources/MCP/Base/Utilities/Progress.swift @@ -0,0 +1,153 @@ +import Foundation + +/// Progress notifications are used to report progress on long-running operations. +/// +/// The sender (either client or server) that issued the original request may +/// include a progress token in the request's `_meta` field. If the receiver +/// supports progress reporting, it can send progress notifications +/// containing that token to indicate how the operation is proceeding. +/// +/// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/basic/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. + /// + /// This is used to associate the progress notification with its + /// originating request. + public let progressToken: ProgressToken + + /// The current progress value. + /// + /// This should increase monotonically as the operation proceeds. + /// It represents the amount of work completed so far. + public let progress: Double + + /// The total expected progress value, if known. + /// + /// When provided, `progress / total` can be used to calculate + /// a percentage completion. + public let total: Double? + + /// An optional human-readable message describing the current progress. + public let message: String? + + public init( + progressToken: ProgressToken, + progress: Double, + total: Double? = nil, + message: String? = nil + ) { + self.progressToken = progressToken + self.progress = progress + self.total = total + self.message = message + } + } +} + +/// A token used to associate progress notifications with their originating request. +/// +/// Progress tokens can be either strings or integers. +/// Progress tokens MUST be unique across all active requests. +public enum ProgressToken: Hashable, Codable, Sendable { + case string(String) + case integer(Int) + + public init(from decoder: Decoder) throws { + let container = try decoder.singleValueContainer() + if let stringValue = try? container.decode(String.self) { + self = .string(stringValue) + } else if let intValue = try? container.decode(Int.self) { + self = .integer(intValue) + } else { + throw DecodingError.dataCorruptedError( + in: container, + debugDescription: "Progress token must be a string or integer" + ) + } + } + + 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) + } + } + + /// Creates a unique progress token using UUID. + public static func unique() -> ProgressToken { + .string(UUID().uuidString) + } +} + +/// Metadata that can be included in request parameters. +/// +/// This structure represents the `_meta` field in MCP request parameters, +/// which can contain a progress token for receiving progress notifications. +/// +/// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/progress +public struct Metadata: Hashable, Codable, Sendable { + public subscript(_ key: String) -> Value? { + fields[key] + } + + /// The underlying fields dictionary. + public var fields: [String: Value] + + /// The progress token for receiving progress notifications. + /// + /// If specified, the caller is requesting out-of-band progress notifications + /// for this request. The value of this parameter is an opaque token that will + /// be attached to any subsequent notifications. The receiver is not obligated + /// to provide these notifications. + public var progressToken: ProgressToken? { + get { + guard let value = fields["progressToken"] else { return nil } + if let stringValue = value.stringValue { + return .string(stringValue) + } else if let intValue = value.intValue { + return .integer(intValue) + } + return nil + } + set { + if let token = newValue { + switch token { + case .string(let s): + fields["progressToken"] = .string(s) + case .integer(let i): + fields["progressToken"] = .int(i) + } + } else { + fields.removeValue(forKey: "progressToken") + } + } + } + + /// Creates request metadata. + /// + /// - Parameters: + /// - progressToken: Optional progress token for receiving progress notifications. + /// If specified, the caller is requesting out-of-band progress notifications for this request. + /// - additionalFields: Optional dictionary of additional metadata fields. + /// These fields will be included in the `_meta` object alongside the progress token. + public init(progressToken: ProgressToken? = nil, additionalFields: [String: Value] = [:]) { + self.fields = additionalFields + self.progressToken = progressToken + } + + public init(from decoder: Decoder) throws { + let container = try decoder.singleValueContainer() + self.fields = try container.decode([String: Value].self) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.singleValueContainer() + try container.encode(fields) + } +} diff --git a/Sources/MCP/Base/Utilities/RequestContext.swift b/Sources/MCP/Base/Utilities/RequestContext.swift new file mode 100644 index 0000000..8caf1c4 --- /dev/null +++ b/Sources/MCP/Base/Utilities/RequestContext.swift @@ -0,0 +1,34 @@ +import Foundation + +/// A context object that wraps a pending request, providing both the request ID +/// and a Task handle for the asynchronous operation. +/// +/// This allows you to track and cancel in-flight requests by sending +/// a CancelledNotification to the receiver (client or server). +/// +/// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/cancellation +public struct RequestContext: Sendable { + /// The unique identifier for this request. + public let requestID: ID + + /// The Task representing the asynchronous work for this request. + private let requestTask: Task + + /// Convenience property to await the result. + /// + /// Example: + /// ```swift + /// let context = try await client.send(request) + /// let result = try await context.value + /// ``` + public var value: Output { + get async throws { + try await requestTask.value + } + } + + public init(requestID: ID, requestTask: Task) { + self.requestID = requestID + self.requestTask = requestTask + } +} diff --git a/Sources/MCP/Base/Versioning.swift b/Sources/MCP/Base/Versioning.swift index 45c8de0..0bed5c4 100644 --- a/Sources/MCP/Base/Versioning.swift +++ b/Sources/MCP/Base/Versioning.swift @@ -4,8 +4,27 @@ import Foundation /// following the format YYYY-MM-DD, to indicate /// the last date backwards incompatible changes were made. /// -/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2024-11-05/basic/versioning/ +/// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/ public enum Version { - /// The current protocol version. - public static let latest = "2024-11-05" + /// All protocol versions supported by this implementation, ordered from newest to oldest. + public static let supported: Set = [ + "2025-11-25", + "2025-06-18", + "2025-03-26", + "2024-11-05", + ] + + /// The latest protocol version supported by this implementation. + public static let latest = supported.max()! + + /// Negotiates the protocol version based on the client's request and server's capabilities. + /// - Parameter clientRequestedVersion: The protocol version requested by the client. + /// - Returns: The negotiated protocol version. If the client's requested version is supported, + /// that version is returned. Otherwise, the server's latest supported version is returned. + static func negotiate(clientRequestedVersion: String) -> String { + if supported.contains(clientRequestedVersion) { + return clientRequestedVersion + } + return latest + } } diff --git a/Sources/MCP/Client/Client.swift b/Sources/MCP/Client/Client.swift index c9aa96e..3eea124 100644 --- a/Sources/MCP/Client/Client.swift +++ b/Sources/MCP/Client/Client.swift @@ -1,5 +1,4 @@ import Logging -import SystemPackage import struct Foundation.Data import struct Foundation.Date @@ -8,16 +7,58 @@ import class Foundation.JSONEncoder /// Model Context Protocol client public actor Client { + /// The client configuration + public struct Configuration: Hashable, Codable, Sendable { + /// The default configuration. + public static let `default` = Configuration(strict: false) + + /// The strict configuration. + public static let strict = Configuration(strict: true) + + /// When strict mode is enabled, the client: + /// - Requires server capabilities to be initialized before making requests + /// - Rejects all requests that require capabilities before initialization + /// + /// While the MCP specification requires servers to respond to initialize requests + /// with their capabilities, some implementations may not follow this. + /// Disabling strict mode allows the client to be more lenient with non-compliant + /// servers, though this may lead to undefined behavior. + public var strict: Bool + + public init(strict: Bool = false) { + self.strict = strict + } + } + /// Implementation information public struct Info: Hashable, Codable, Sendable { /// The client name public var name: String + /// A human-readable title for display purposes + public var title: String? /// The client version public var version: String + /// Optional description of the client + public var description: String? + /// Optional website URL for the client + public var websiteUrl: String? + /// Optional set of sized icons for display in a user interface + public var icons: [Icon]? - public init(name: String, version: String) { + public init( + name: String, + version: String, + title: String? = nil, + description: String? = nil, + websiteUrl: String? = nil, + icons: [Icon]? = nil + ) { self.name = name + self.title = title self.version = version + self.description = description + self.websiteUrl = websiteUrl + self.icons = icons } } @@ -34,12 +75,55 @@ public actor Client { } /// The sampling capabilities - public struct Sampling: Hashable, Codable, Sendable { - public init() {} + public struct Sampling: Hashable, Sendable { + /// Tools sub-capability for sampling + public struct Tools: Hashable, Codable, Sendable { + public init() {} + } + + /// Context sub-capability for sampling + public struct Context: Hashable, Codable, Sendable { + public init() {} + } + + /// Whether tools are supported in sampling + public var tools: Tools? + /// Whether context is supported in sampling + public var context: Context? + + public init(tools: Tools? = nil, context: Context? = nil) { + self.tools = tools + self.context = context + } + } + + /// The elicitation capabilities + public struct Elicitation: Hashable, Sendable { + /// Form-based elicitation sub-capability + public struct Form: Hashable, Codable, Sendable { + public init() {} + } + + /// URL-based elicitation sub-capability + public struct URL: Hashable, Codable, Sendable { + public init() {} + } + + /// Whether form-based elicitation is supported + public var form: Form? + /// Whether URL-based elicitation is supported + public var url: URL? + + public init(form: Form? = Form(), url: URL? = nil) { + self.form = form + self.url = url + } } /// Whether the client supports sampling public var sampling: Sampling? + /// Whether the client supports elicitation + public var elicitation: Elicitation? /// Experimental features supported by the client public var experimental: [String: String]? /// Whether the client supports roots @@ -47,10 +131,12 @@ public actor Client { public init( sampling: Sampling? = nil, + elicitation: Elicitation? = nil, experimental: [String: String]? = nil, roots: Capabilities.Roots? = nil ) { self.sampling = sampling + self.elicitation = elicitation self.experimental = experimental self.roots = roots } @@ -69,11 +155,15 @@ public actor Client { private let clientInfo: Client.Info /// The client name public nonisolated var name: String { clientInfo.name } + /// A human-readable client title + public nonisolated var title: String? { clientInfo.title } /// The client version public nonisolated var version: String { clientInfo.version } /// The client capabilities public var capabilities: Client.Capabilities + /// The client configuration + public var configuration: Configuration /// The server capabilities private var serverCapabilities: Server.Capabilities? @@ -84,66 +174,41 @@ public actor Client { /// A dictionary of type-erased notification handlers, keyed by method name private var notificationHandlers: [String: [NotificationHandlerBox]] = [:] + /// Method handlers for server-to-client requests + private var methodHandlers: [String: RequestHandlerBox] = [:] /// The task for the message handling loop private var task: 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 - private struct AnyPendingRequest { - private let _resume: (Result) -> Void - - init(_ request: PendingRequest) { - _resume = { result in - switch result { - case .success(let value): - if let typedValue = value as? T { - request.continuation.resume(returning: typedValue) - } 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) - } else { - request.continuation.resume(throwing: TypeMismatchError()) - } - case .failure(let error): - request.continuation.resume(throwing: error) - } - } - } - func resume(returning value: Any) { - _resume(.success(value)) - } - - func resume(throwing error: Swift.Error) { - _resume(.failure(error)) - } - } - /// A dictionary of type-erased pending requests, keyed by request ID private var pendingRequests: [ID: AnyPendingRequest] = [:] + // Add reusable JSON encoder/decoder + private let encoder = JSONEncoder() + private let decoder = JSONDecoder() public init( name: String, - version: String + version: String, + title: String? = nil, + description: String? = nil, + websiteUrl: String? = nil, + icons: [Icon]? = nil, + capabilities: Capabilities = Capabilities(), + configuration: Configuration = .default ) { - self.clientInfo = Client.Info(name: name, version: version) - self.capabilities = Capabilities() + self.clientInfo = Client.Info( + name: name, version: version, title: title, + description: description, websiteUrl: websiteUrl, icons: icons) + self.capabilities = capabilities + self.configuration = configuration } /// Connect to the server using the given transport - public func connect(transport: any Transport) async throws { + @discardableResult + public func connect(transport: any Transport) async throws -> Initialize.Result { self.connection = transport try await self.connection?.connect() - await logger?.info( + await logger?.debug( "Client connected", metadata: ["name": "\(name)", "version": "\(version)"]) // Start message handling loop @@ -155,20 +220,17 @@ public actor Client { do { let stream = await connection.receive() - for try await string in stream { + for try await data in stream { if Task.isCancelled { break } // Check inside loop too - // Decode and handle incoming message - guard let data = string.data(using: .utf8) else { - throw Error.parseError("Invalid UTF-8 data") - } - - // Attempt to decode string data as AnyResponse or AnyMessage - let decoder = JSONDecoder() - if let response = try? decoder.decode(AnyResponse.self, from: data), - let request = pendingRequests[response.id] - { - await handleResponse(response, for: request) + // 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) { + await handleIncomingRequest(request) } else if let message = try? decoder.decode(AnyMessage.self, from: data) { await handleMessage(message) } else { @@ -177,11 +239,13 @@ public actor Client { metadata["message"] = .string(string) } await logger?.warning( - "Unexpected message received by client", metadata: metadata) + "Unexpected message received by client (not single/batch response or notification)", + metadata: metadata + ) } } - } catch let error as Errno where error == .resourceTemporarilyUnavailable { - try? await Task.sleep(nanoseconds: 10_000_000) // 10ms + } catch let error where MCPError.isResourceTemporarilyUnavailable(error) { + try? await Task.sleep(for: .milliseconds(10)) continue } catch { await logger?.error( @@ -189,23 +253,75 @@ public actor Client { break } } while true + await self.logger?.debug("Client message handling loop task is terminating.") } + + // Register cancellation notification handler + await self.onNotification(CancelledNotification.self) { [weak self] message in + guard let self = self else { return } + + let requestId = message.params.requestId + let reason = message.params.reason + + await self.logger?.debug( + "Received cancellation notification", + metadata: [ + "requestId": requestId.map { "\($0)" } ?? "none", + "reason": reason.map { "\($0)" } ?? "none", + ] + ) + + // Remove the pending request and resume with cancellation error + if let requestId = requestId, + let pendingRequest = await self.removePendingRequest(id: requestId) + { + pendingRequest.resume(throwing: CancellationError()) + } + } + + // Automatically initialize after connecting + return try await _initialize() } /// Disconnect the client and cancel all pending requests public func disconnect() async { - // Cancel all pending requests - for (id, request) in pendingRequests { - request.resume(throwing: Error.internalError("Client disconnected")) - pendingRequests.removeValue(forKey: id) + await logger?.debug("Initiating client disconnect...") + + // Part 1: Inside actor - Grab state and clear internal references + let taskToCancel = self.task + let connectionToDisconnect = self.connection + let pendingRequestsToCancel = self.pendingRequests + + self.task = nil + self.connection = nil + self.pendingRequests = [:] // Use empty dictionary literal + + // Part 2: Outside actor - Resume continuations, disconnect transport, await task + + // Resume continuations first + for (_, request) in pendingRequestsToCancel { + request.resume(throwing: MCPError.internalError("Client disconnected")) + } + await logger?.debug("Pending requests cancelled.") + + // Cancel the task + taskToCancel?.cancel() + await logger?.debug("Message loop task cancellation requested.") + + // Disconnect the transport *before* awaiting the task + // This should ensure the transport stream is finished, unblocking the loop. + if let conn = connectionToDisconnect { + await conn.disconnect() + await logger?.debug("Transport disconnected.") + } else { + await logger?.debug("No active transport connection to disconnect.") } - task?.cancel() - task = nil - if let connection = connection { - await connection.disconnect() - } - connection = nil + // Await the task completion *after* transport disconnect + _ = await taskToCancel?.value + await logger?.debug("Client message loop task finished.") + + await logger?.debug("Client disconnect complete.") } // MARK: - Registration @@ -221,54 +337,310 @@ public actor Client { return self } + /// Register a handler for server-to-client requests + @discardableResult + public func withMethodHandler( + _ type: M.Type, + handler: @escaping @Sendable (M.Parameters) async throws -> M.Result + ) -> Self { + methodHandlers[M.name] = TypedRequestHandler { (request: Request) -> Response in + let result = try await handler(request.params) + return Response(id: request.id, result: result) + } + 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 response back to the server for a server-to-client request + private func send(_ response: Response) async throws { + guard let connection = connection else { + throw MCPError.internalError("Client connection not initialized") + } + let responseData = try encoder.encode(response) + try await connection.send(responseData) + } + // MARK: - Requests - /// Send a request and receive its response - public func send(_ request: Request) async throws -> M.Result { + /// Send a request and return a RequestContext that wraps both the request ID and the Task. + /// + /// This allows you to track and cancel the request by sending a CancelledNotification + /// to the server using the requestID. + /// + /// Example: + /// ```swift + /// let context = try await client.send(request) + /// // Later, to cancel: + /// try await client.cancelRequest(context.requestID, reason: "User cancelled") + /// // Await the result: + /// let result = try await context.value + /// ``` + /// + /// - Parameter request: The request to send + /// - Returns: A RequestContext containing the request ID and Task + /// - Throws: MCPError if the client is not connected + /// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/cancellation + public func send(_ request: Request) throws -> RequestContext { guard let connection = connection else { - throw Error.internalError("Client connection not initialized") + throw MCPError.internalError("Client connection not initialized") } - let requestData = try JSONEncoder().encode(request) - guard let requestString = String(data: requestData, encoding: .utf8) else { - throw Error.internalError("Failed to encode request") - } + let requestData = try encoder.encode(request) - return try await withCheckedThrowingContinuation { continuation in - // Store the pending request first - Task { - self.addPendingRequest( - id: request.id, - continuation: continuation, - type: M.Result.self - ) + let requestTask = Task { + 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 + ) - // Send the request - do { - try await connection.send(requestString) - } catch { - continuation.resume(throwing: error) - self.removePendingRequest(id: request.id) + // Send the request data + do { + try await connection.send(requestData) + } catch { + // If send fails, try to remove the pending request. + if self.removePendingRequest(id: request.id) != nil { + continuation.resume(throwing: error) + } + } } } } + + return RequestContext(requestID: request.id, requestTask: requestTask) + } + + /// Cancel a request by sending a CancelledNotification to the server. + /// + /// According to the MCP specification, cancellation is advisory: + /// - The server SHOULD stop processing and free resources + /// - The server MAY ignore the cancellation if the request is unknown, already completed, + /// or cannot be cancelled + /// - The client SHOULD ignore any response that arrives after cancellation + /// + /// This method removes the pending request and resumes it with CancellationError, + /// ensuring that any response arriving after cancellation is ignored. + /// + /// - Parameters: + /// - requestID: The ID of the request to cancel + /// - reason: An optional human-readable reason for the cancellation + /// - Throws: MCPError if the notification cannot be sent + /// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/cancellation + public func cancelRequest(_ requestID: ID, reason: String? = nil) async throws { + // Remove the pending request and resume with cancellation error + // This ensures any response that arrives after cancellation is ignored + if let pendingRequest = removePendingRequest(id: requestID) { + pendingRequest.resume(throwing: CancellationError()) + } + + // Send cancellation notification to server + let notification = CancelledNotification.message( + .init(requestId: requestID, reason: reason) + ) + try await notify(notification) + } + + /// Send a request and receive its response immediately. + /// + /// Internal convenience method for cases where cancellation tracking is not needed. + /// + /// - Parameter request: The request to send + /// - Returns: The result of the request + /// - Throws: MCPError if the client is not connected + func sendAndAwait(_ request: Request) async throws -> M.Result { + let context = try send(request) + return try await context.value } private func addPendingRequest( id: ID, continuation: CheckedContinuation, - type: T.Type + type: T.Type // Keep type for AnyPendingRequest internal logic ) { - pendingRequests[id] = AnyPendingRequest(PendingRequest(continuation: continuation)) + pendingRequests[id] = AnyPendingRequest( + PendingRequest(continuation: continuation) + ) } - private func removePendingRequest(id: ID) { - pendingRequests.removeValue(forKey: id) + private func removePendingRequest(id: ID) -> AnyPendingRequest? { + return pendingRequests.removeValue(forKey: id) + } + + // MARK: - Batching + + /// A batch of requests. + /// + /// Objects of this type are passed as an argument to the closure + /// of the ``Client/withBatch(body:)`` 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)) + + // 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 + ) + } + } + } + } + } + + /// 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 @Sendable (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, @@ -276,118 +648,344 @@ public actor Client { clientInfo: clientInfo )) - let result = try await send(request) + let result = try await sendAndAwait(request) self.serverCapabilities = result.capabilities self.serverVersion = result.protocolVersion self.instructions = result.instructions + // For HTTP transport, ensure subsequent MCP-Protocol-Version headers + // reflect the negotiated lifecycle version. + if let httpTransport = connection as? HTTPClientTransport { + await httpTransport.updateNegotiatedProtocolVersion(result.protocolVersion) + } + + try await notify(InitializedNotification.message()) + return result } public func ping() async throws { let request = Ping.request() - _ = try await send(request) + _ = try await sendAndAwait(request) } // 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 checkCapability(\.prompts, "Prompts") + try validateServerCapability(\.prompts, "Prompts") let request = GetPrompt.request(.init(name: name, arguments: arguments)) - let result = try await send(request) + let result = try await sendAndAwait(request) return (description: result.description, messages: result.messages) } public func listPrompts(cursor: String? = nil) async throws -> (prompts: [Prompt], nextCursor: String?) { - _ = try checkCapability(\.prompts, "Prompts") + 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) + let result = try await sendAndAwait(request) return (prompts: result.prompts, nextCursor: result.nextCursor) } // MARK: - Resources public func readResource(uri: String) async throws -> [Resource.Content] { - _ = try checkCapability(\.resources, "Resources") + try validateServerCapability(\.resources, "Resources") let request = ReadResource.request(.init(uri: uri)) - let result = try await send(request) + let result = try await sendAndAwait(request) return result.contents } public func listResources(cursor: String? = nil) async throws -> ( resources: [Resource], nextCursor: String? ) { - _ = try checkCapability(\.resources, "Resources") + 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) + let result = try await sendAndAwait(request) return (resources: result.resources, nextCursor: result.nextCursor) } public func subscribeToResource(uri: String) async throws { - _ = try checkCapability(\.resources?.subscribe, "Resource subscription") + try validateServerCapability(\.resources?.subscribe, "Resource subscription") let request = ResourceSubscribe.request(.init(uri: uri)) - _ = try await send(request) + _ = try await sendAndAwait(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 sendAndAwait(request) + return (templates: result.templates, nextCursor: result.nextCursor) } // MARK: - Tools - public func listTools(cursor: String? = nil) async throws -> [Tool] { - _ = try checkCapability(\.tools, "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 result.tools + let result = try await sendAndAwait(request) + return (tools: result.tools, nextCursor: result.nextCursor) } - public func callTool(name: String, arguments: [String: Value]? = nil) async throws -> ( - content: [Tool.Content], isError: Bool? - ) { - _ = try checkCapability(\.tools, "Tools") - let request = CallTool.request(.init(name: name, arguments: arguments)) - let result = try await send(request) + /// Call a tool on the server. + /// + /// - Parameters: + /// - name: The name of the tool to call. + /// - arguments: Arguments to use for the tool call. + /// - meta: Optional request metadata including progress token. If `progressToken` is specified, + /// the caller is requesting out-of-band progress notifications for this request. + /// Use `onNotification(ProgressNotification.self)` to receive progress updates. + /// - Returns: A tuple containing the tool's content response and an optional error flag. + /// - Note: For advanced use cases requiring cancellation support, use `send()` directly to get a `RequestContext`. + /// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/server/tools/#calling-tools + public func callTool( + name: String, + arguments: [String: Value]? = nil, + meta: Metadata? = nil + ) async throws -> (content: [Tool.Content], isError: Bool?) { + try validateServerCapability(\.tools, "Tools") + let request = CallTool.request(.init(name: name, arguments: arguments, meta: meta)) + let result = try await sendAndAwait(request) return (content: result.content, isError: result.isError) } + /// Call a tool on the server. + /// + /// - Parameters: + /// - name: The name of the tool to call. + /// - arguments: Arguments to use for the tool call. + /// - meta: Optional request metadata including progress token. If `progressToken` is specified, + /// the caller is requesting out-of-band progress notifications for this request. + /// Use `onNotification(ProgressNotification.self)` to receive progress updates. + /// - Returns: A tuple containing the tool's content response and an optional error flag. + /// - Note: For advanced use cases requiring cancellation support, use `send()` directly to get a `RequestContext`. + /// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/server/tools/#calling-tools + public func callTool( + name: String, + arguments: [String: Value]? = nil, + meta: Metadata? = nil + ) throws -> RequestContext { + try validateServerCapability(\.tools, "Tools") + let request = CallTool.request(.init(name: name, arguments: arguments, meta: meta)) + return try send(request) + } + + // MARK: - Sampling + + /// Register a handler for sampling requests from servers + /// + /// Sampling allows servers to request LLM completions through the client, + /// enabling sophisticated agentic behaviors while maintaining human-in-the-loop control. + /// + /// The sampling flow follows these steps: + /// 1. Server sends a `sampling/createMessage` request to the client + /// 2. Client reviews the request and can modify it (via this handler) + /// 3. Client samples from an LLM (via this handler) + /// 4. Client reviews the completion (via this handler) + /// 5. Client returns the result to the server + /// + /// - Parameter handler: A closure that processes sampling requests and returns completions + /// - Returns: Self for method chaining + /// - SeeAlso: https://modelcontextprotocol.io/docs/concepts/sampling#how-sampling-works + @discardableResult + public func withSamplingHandler( + _ handler: + @escaping @Sendable (CreateSamplingMessage.Parameters) async throws -> + CreateSamplingMessage.Result + ) -> Self { + return withMethodHandler(CreateSamplingMessage.self, handler: handler) + } + + // MARK: - Elicitation + + /// Register a handler for elicitation requests from servers + /// + /// The elicitation flow lets servers collect structured input from users during + /// ongoing interactions. Clients remain in control by mediating the prompt, + /// collecting the response, and returning the chosen action to the server. + /// + /// - Parameter handler: A closure that processes elicitation requests and returns user actions + /// - Returns: Self for method chaining + /// - SeeAlso: https://modelcontextprotocol.io/specification/2025-06-18/client/elicitation + @discardableResult + public func withElicitationHandler( + _ handler: + @escaping @Sendable (CreateElicitation.Parameters) async throws -> + CreateElicitation.Result + ) -> Self { + return withMethodHandler(CreateElicitation.self, handler: handler) + } + + // MARK: - Roots + + /// Register a handler for roots/list requests from servers + /// + /// Roots define filesystem boundaries that servers can operate within. + /// Unlike other MCP features, roots use bidirectional communication where + /// servers send requests TO clients to discover available roots. + /// + /// - Parameter handler: A closure that returns the list of available roots + /// - Returns: Self for method chaining + /// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/client/roots + @discardableResult + public func withRootsHandler( + _ handler: @escaping @Sendable () async throws -> [Root] + ) -> Self { + return withMethodHandler(ListRoots.self) { _ in + let roots = try await handler() + return ListRoots.Result(roots: roots) + } + } + + /// Notify the server that the list of roots has changed + /// + /// Clients should send this notification when roots are added, removed, + /// or modified to inform connected servers of the change. + /// + /// - Throws: MCPError if the client is not connected + /// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/client/roots + public func notifyRootsChanged() async throws { + let notification = RootsListChangedNotification.message() + try await notify(notification) + } + + // MARK: - Logging + + /// Set the minimum logging level for server log messages. + /// + /// Servers that declare the `logging` capability will send log messages via + /// `notifications/message` notifications. Use this method to control which + /// severity levels the server should send. + /// + /// - Parameter level: The minimum log level to receive + /// - Throws: MCPError if the client is not connected or if the server doesn't support logging + /// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/server/utilities/logging/ + public func setLoggingLevel(_ level: LogLevel) async throws { + try validateServerCapability(\.logging, "Logging") + let request = SetLoggingLevel.request(.init(level: level)) + _ = try await sendAndAwait(request) + } + + // MARK: - Completions + + /// Request completion suggestions for a prompt argument. + /// + /// Servers that declare the `completions` capability can provide autocompletion + /// suggestions for prompt arguments as users type. + /// + /// - Parameters: + /// - promptName: The name of the prompt + /// - argumentName: The name of the argument being completed + /// - argumentValue: The current (partial) value of the argument + /// - context: Optional context with already-resolved arguments + /// - Returns: A completion result containing suggested values + /// - Throws: MCPError if the client is not connected or if the server doesn't support completions + /// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/server/utilities/completion/ + public func complete( + promptName: String, + argumentName: String, + argumentValue: String, + context: [String: String]? = nil + ) async throws -> Complete.Result.Completion { + try validateServerCapability(\.completions, "Completions") + let request = Complete.request( + .init( + ref: .prompt(.init(name: promptName)), + argument: .init(name: argumentName, value: argumentValue), + context: context.map { .init(arguments: $0) } + ) + ) + let result = try await sendAndAwait(request) + return result.completion + } + + /// Request completion suggestions for a resource template argument. + /// + /// Servers that declare the `completions` capability can provide autocompletion + /// suggestions for resource template arguments as users type. + /// + /// - Parameters: + /// - resourceURI: The URI of the resource template + /// - argumentName: The name of the argument being completed + /// - argumentValue: The current (partial) value of the argument + /// - context: Optional context with already-resolved arguments + /// - Returns: A completion result containing suggested values + /// - Throws: MCPError if the client is not connected or if the server doesn't support completions + /// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/server/utilities/completion/ + public func complete( + resourceURI: String, + argumentName: String, + argumentValue: String, + context: [String: String]? = nil + ) async throws -> Complete.Result.Completion { + try validateServerCapability(\.completions, "Completions") + let request = Complete.request( + .init( + ref: .resource(.init(uri: resourceURI)), + argument: .init(name: argumentName, value: argumentValue), + context: context.map { .init(arguments: $0) } + ) + ) + let result = try await sendAndAwait(request) + return result.completion + } + // MARK: - - private func handleResponse(_ response: Response, for request: AnyPendingRequest) - async - { - await logger?.debug( + private func handleResponse(_ response: Response) async { + await logger?.trace( "Processing response", metadata: ["id": "\(response.id)"]) - switch response.result { - case .success(let value): - request.resume(returning: value) - case .failure(let error): - request.resume(throwing: error) + // 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)"] + ) } - - removePendingRequest(id: response.id) } private func handleMessage(_ message: Message) async { - await logger?.debug( + await logger?.trace( "Processing notification", metadata: ["method": "\(message.method)"]) @@ -409,17 +1007,139 @@ public actor Client { } } + private func handleIncomingRequest(_ request: Request) async { + await logger?.trace( + "Processing incoming request from server", + metadata: ["method": "\(request.method)", "id": "\(request.id)"]) + + guard let handler = methodHandlers[request.method] else { + await logger?.warning( + "No handler registered for method", + metadata: ["method": "\(request.method)"]) + let error = MCPError.methodNotFound("Unknown method: \(request.method)") + let response = AnyMethod.response(id: request.id, error: error) + do { + try await send(response) + } catch { + await logger?.error( + "Failed to send error response", + metadata: ["error": "\(error)"]) + } + return + } + + do { + let response = try await handler(request) + try await send(response) + } catch { + let mcpError = error as? MCPError ?? MCPError.internalError(error.localizedDescription) + let response = AnyMethod.response(id: request.id, error: mcpError) + do { + try await send(response) + } catch { + await logger?.error( + "Failed to send error response", + metadata: ["error": "\(error)"]) + } + } + } + // MARK: - - private func checkCapability(_ keyPath: KeyPath, _ name: String) - throws -> T + /// 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 { - guard let capabilities = serverCapabilities else { - throw Error.methodNotFound("Server capabilities not initialized") + 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") + } } - guard let value = capabilities[keyPath: keyPath] else { - throw Error.methodNotFound("\(name) is not supported by the server") + } + + // 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)"] + ) + } } - return value + } +} + +// MARK: - Codable + +extension Client.Capabilities.Sampling: Codable { + private enum CodingKeys: String, CodingKey { + case tools, context + } + + public init(from decoder: Decoder) throws { + // Handle both empty object {} and object with sub-capabilities + if let container = try? decoder.container(keyedBy: CodingKeys.self) { + self.tools = try container.decodeIfPresent(Tools.self, forKey: .tools) + self.context = try container.decodeIfPresent(Context.self, forKey: .context) + } else { + // Empty object - no capabilities + self.tools = nil + self.context = nil + } + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encodeIfPresent(tools, forKey: .tools) + try container.encodeIfPresent(context, forKey: .context) + } +} + +extension Client.Capabilities.Elicitation: Codable { + private enum CodingKeys: String, CodingKey { + case form, url + } + + public init(from decoder: Decoder) throws { + // Handle both empty object {} and object with sub-capabilities + if let container = try? decoder.container(keyedBy: CodingKeys.self) { + self.form = try container.decodeIfPresent(Form.self, forKey: .form) + self.url = try container.decodeIfPresent(URL.self, forKey: .url) + // If both are nil, default to form for backward compatibility + if self.form == nil && self.url == nil { + self.form = Form() + } + } else { + // Empty object - default to form-only for backward compatibility + self.form = Form() + self.url = nil + } + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encodeIfPresent(form, forKey: .form) + try container.encodeIfPresent(url, forKey: .url) } } diff --git a/Sources/MCP/Client/Elicitation.swift b/Sources/MCP/Client/Elicitation.swift new file mode 100644 index 0000000..bf08519 --- /dev/null +++ b/Sources/MCP/Client/Elicitation.swift @@ -0,0 +1,233 @@ +import Foundation + +/// Types supporting the MCP elicitation flow. +/// +/// Servers use elicitation to collect structured input from users via the client. +/// The schema subset mirrors the 2025-11-25 revision of the specification. +public enum Elicitation { + /// Schema describing the expected response content. + public struct RequestSchema: Hashable, Codable, Sendable { + /// Supported top-level types. Currently limited to objects. + public enum SchemaType: String, Hashable, Codable, Sendable { + case object + } + + /// Schema title presented to users. + public var title: String? + /// Schema description providing additional guidance. + public var description: String? + /// Raw JSON Schema fragments describing the requested fields. + public var properties: [String: Value] + /// List of required field keys. + public var required: [String]? + /// Top-level schema type. Defaults to `object`. + public var type: SchemaType + + public init( + title: String? = nil, + description: String? = nil, + properties: [String: Value] = [:], + required: [String]? = nil, + type: SchemaType = .object + ) { + self.title = title + self.description = description + self.properties = properties + self.required = required + self.type = type + } + + private enum CodingKeys: String, CodingKey { + case title, description, properties, required, type + } + } + + /// Elicitation mode indicating how user input is collected + public enum Mode: String, Hashable, Codable, Sendable { + /// Form-based elicitation (client displays UI) + case form + /// URL-based elicitation (client opens external URL) + case url + } +} + +/// To request information from a user, servers send an `elicitation/create` request. +/// - SeeAlso: https://modelcontextprotocol.io/specification/2025-06-18/client/elicitation +public enum CreateElicitation: Method { + public static let name = "elicitation/create" + + public enum Parameters: Hashable, Sendable { + /// Form-based elicitation parameters + case form(FormParameters) + /// URL-based elicitation parameters + case url(URLParameters) + + /// Parameters for form-based elicitation + public struct FormParameters: Hashable, Codable, Sendable { + /// Message displayed to the user describing the request + public var message: String + /// Elicitation mode (optional for backward compatibility, defaults to form) + public var mode: Elicitation.Mode? + /// Schema describing the expected response content (required per spec) + public var requestedSchema: Elicitation.RequestSchema + /// Optional metadata + public var _meta: Metadata? + + public init( + message: String, + mode: Elicitation.Mode? = nil, + requestedSchema: Elicitation.RequestSchema = .init(), + _meta: Metadata? = nil + ) { + self.message = message + self.mode = mode + self.requestedSchema = requestedSchema + self._meta = _meta + } + } + + /// Parameters for URL-based elicitation + public struct URLParameters: Hashable, Codable, Sendable { + /// Message displayed to the user describing the request + public var message: String + /// Elicitation mode (always "url") + public var mode: Elicitation.Mode + /// URL for the user to visit + public var url: String + /// Unique identifier for this elicitation + public var elicitationId: String + /// Optional metadata + public var _meta: Metadata? + + public init( + message: String, + url: String, + elicitationId: String, + _meta: Metadata? = nil + ) { + self.message = message + self.mode = .url + self.url = url + self.elicitationId = elicitationId + self._meta = _meta + } + } + } + + public struct Result: Hashable, Codable, Sendable { + /// Indicates how the user responded to the request. + public enum Action: String, Hashable, Codable, Sendable { + case accept + case decline + case cancel + } + + /// Selected action. + public var action: Action + /// Submitted content when action is `.accept`. + public var content: [String: Value]? + /// Optional metadata about this result + public var _meta: Metadata? + + public init( + action: Action, + content: [String: Value]? = nil, + _meta: Metadata? = nil + ) { + self.action = action + self.content = content + self._meta = _meta + } + + private enum CodingKeys: String, CodingKey, CaseIterable { + case action, content, _meta + } + + 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) + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + action = try container.decode(Action.self, forKey: .action) + content = try container.decodeIfPresent([String: Value].self, forKey: .content) + _meta = try container.decodeIfPresent(Metadata.self, forKey: ._meta) + } + } +} + +// MARK: - Codable + +extension CreateElicitation.Parameters: Codable { + private enum CodingKeys: String, CodingKey { + case mode, message, requestedSchema, url, elicitationId + case _meta + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + + // Read mode field (may be missing for backward compatibility) + let mode = try container.decodeIfPresent(Elicitation.Mode.self, forKey: .mode) + + // Discriminate based on mode + if mode == .url { + // URL mode + let message = try container.decode(String.self, forKey: .message) + let url = try container.decode(String.self, forKey: .url) + let elicitationId = try container.decode(String.self, forKey: .elicitationId) + let _meta = try container.decodeIfPresent(Metadata.self, forKey: ._meta) + self = .url(URLParameters( + message: message, + url: url, + elicitationId: elicitationId, + _meta: _meta)) + } else { + // Form mode (default for backward compatibility) + let message = try container.decode(String.self, forKey: .message) + let requestedSchema = try container.decode( + Elicitation.RequestSchema.self, forKey: .requestedSchema) + let _meta = try container.decodeIfPresent(Metadata.self, forKey: ._meta) + self = .form(FormParameters( + message: message, + mode: mode, + requestedSchema: requestedSchema, + _meta: _meta)) + } + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + + switch self { + case .form(let params): + try container.encode(params.message, forKey: .message) + try container.encodeIfPresent(params.mode, forKey: .mode) + try container.encodeIfPresent(params.requestedSchema, forKey: .requestedSchema) + try container.encodeIfPresent(params._meta, forKey: ._meta) + case .url(let params): + try container.encode(params.message, forKey: .message) + try container.encode(params.mode, forKey: .mode) + try container.encode(params.url, forKey: .url) + try container.encode(params.elicitationId, forKey: .elicitationId) + try container.encodeIfPresent(params._meta, forKey: ._meta) + } + } +} + +/// Notification sent when a URL-based elicitation is complete +public struct ElicitationCompleteNotification: Notification { + public static let name = "notifications/elicitation/complete" + + public struct Parameters: Hashable, Codable, Sendable { + /// The elicitation ID that was completed + public var elicitationId: String + + public init(elicitationId: String) { + self.elicitationId = elicitationId + } + } +} diff --git a/Sources/MCP/Client/Roots.swift b/Sources/MCP/Client/Roots.swift new file mode 100644 index 0000000..215fe3a --- /dev/null +++ b/Sources/MCP/Client/Roots.swift @@ -0,0 +1,100 @@ +import Foundation + +/// The Model Context Protocol (MCP) provides a mechanism for clients to expose +/// filesystem boundaries to servers through roots. Roots allow servers to understand +/// the scope of filesystem access they can request, enabling safe and controlled +/// file operations. +/// +/// Unlike Resources/Tools/Prompts, Roots use bidirectional communication: +/// - Servers send `roots/list` requests TO clients +/// - Clients respond with available roots +/// - Clients send `notifications/roots/list_changed` when roots change +/// +/// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/client/roots +public struct Root: Hashable, Codable, Sendable { + /// The root URI (must use file:// scheme) + public let uri: String + /// Optional human-readable name for the root + public let name: String? + /// Optional metadata + public var _meta: Metadata? + + public init( + uri: String, + name: String? = nil, + _meta: Metadata? = nil + ) { + self.uri = uri + self.name = name + self._meta = _meta + } + + private enum CodingKeys: String, CodingKey { + case uri + case name + case _meta + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + uri = try container.decode(String.self, forKey: .uri) + name = try container.decodeIfPresent(String.self, forKey: .name) + _meta = try container.decodeIfPresent(Metadata.self, forKey: ._meta) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode(uri, forKey: .uri) + try container.encodeIfPresent(name, forKey: .name) + try container.encodeIfPresent(_meta, forKey: ._meta) + } +} + +// MARK: - + +/// To discover available roots, servers send a `roots/list` request to the client. +/// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/client/roots +public enum ListRoots: Method { + public static let name: String = "roots/list" + + public typealias Parameters = Empty + + public struct Result: Hashable, Codable, Sendable { + public let roots: [Root] + /// Optional metadata about this result + public var _meta: Metadata? + + public init( + roots: [Root], + _meta: Metadata? = nil + ) { + self.roots = roots + self._meta = _meta + } + + private enum CodingKeys: String, CodingKey, CaseIterable { + case roots, _meta + } + + 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) + } + + 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(Metadata.self, forKey: ._meta) + } + } +} + +/// When the list of roots changes, clients that declared the `roots` capability +/// SHOULD send this notification to inform servers. +/// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/client/roots +public struct RootsListChangedNotification: Notification { + public static let name: String = "notifications/roots/list_changed" + + public typealias Parameters = Empty +} diff --git a/Sources/MCP/Client/Sampling.swift b/Sources/MCP/Client/Sampling.swift new file mode 100644 index 0000000..5694f7e --- /dev/null +++ b/Sources/MCP/Client/Sampling.swift @@ -0,0 +1,600 @@ +import Foundation + +/// The Model Context Protocol (MCP) allows servers to request LLM completions +/// through the client, enabling sophisticated agentic behaviors while maintaining +/// security and privacy. +/// +/// - SeeAlso: https://modelcontextprotocol.io/docs/concepts/sampling#how-sampling-works +public enum Sampling { + /// A message in the conversation history. + public struct Message: Hashable, Sendable { + /// The message role + public enum Role: String, Hashable, Codable, Sendable { + /// A user message + case user + /// An assistant message + case assistant + } + + /// The message role + public let role: Role + /// The message content + public let content: Content + /// Optional metadata + public var _meta: Metadata? + + /// 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, _meta: Metadata? = nil) { + self.role = role + self.content = content + self._meta = _meta + } + + /// Private initializer for convenience methods to avoid deprecation warnings + private init(_role role: Role, _content content: Content, _meta: Metadata? = nil) { + self.role = role + self.content = content + self._meta = _meta + } + + /// Creates a user message with the specified content + public static func user(_ content: Content, _meta: Metadata? = nil) -> Message { + return Message(_role: .user, _content: content, _meta: _meta) + } + + /// Creates an assistant message with the specified content + public static func assistant(_ content: Content, _meta: Metadata? = nil) -> Message { + return Message(_role: .assistant, _content: content, _meta: _meta) + } + + /// Content types for sampling messages + public enum Content: Hashable, Sendable { + /// Single content block + case single(ContentBlock) + /// Multiple content blocks + case multiple([ContentBlock]) + + /// Individual content blocks in messages + public enum ContentBlock: Hashable, Sendable { + /// Text content + case text(String) + /// Image content + case image(data: String, mimeType: String) + /// Audio content + case audio(data: String, mimeType: String) + /// Tool use content + case toolUse(Sampling.ToolUseContent) + /// Tool result content + case toolResult(Sampling.ToolResultContent) + } + + /// Returns true if this is a single content block + public var isSingle: Bool { + if case .single = self { return true } + return false + } + + /// Returns content as an array of blocks + public var asArray: [ContentBlock] { + switch self { + case .single(let block): + return [block] + case .multiple(let blocks): + return blocks + } + } + + /// Creates content from a text string (convenience) + public static func text(_ text: String) -> Content { + .single(.text(text)) + } + + /// Creates content from an image (convenience) + public static func image(data: String, mimeType: String) -> Content { + .single(.image(data: data, mimeType: mimeType)) + } + + /// Creates content from audio (convenience) + public static func audio(data: String, mimeType: String) -> Content { + .single(.audio(data: data, mimeType: mimeType)) + } + } + } + + /// Model preferences for sampling requests + public struct ModelPreferences: Hashable, Codable, Sendable { + /// Model hints for selection + public struct Hint: Hashable, Codable, Sendable { + /// Suggested model name/family + public let name: String? + + public init(name: String? = nil) { + self.name = name + } + } + + /// Array of model name suggestions that clients can use to select an appropriate model + public let hints: [Hint]? + /// Importance of minimizing costs (0-1 normalized) + public let costPriority: UnitInterval? + /// Importance of low latency response (0-1 normalized) + public let speedPriority: UnitInterval? + /// Importance of advanced model capabilities (0-1 normalized) + public let intelligencePriority: UnitInterval? + + public init( + hints: [Hint]? = nil, + costPriority: UnitInterval? = nil, + speedPriority: UnitInterval? = nil, + intelligencePriority: UnitInterval? = nil + ) { + self.hints = hints + self.costPriority = costPriority + self.speedPriority = speedPriority + self.intelligencePriority = intelligencePriority + } + } + + /// Context inclusion options for sampling requests + public enum ContextInclusion: String, Hashable, Codable, Sendable { + /// No additional context + case none + /// Include context from the requesting server + case thisServer + /// Include context from all connected MCP servers + case allServers + } + + /// Stop reason for sampling completion. + /// + /// The spec defines this as an open string — any provider-specific value is valid. + /// The well-known values are exposed as static constants. + public struct StopReason: RawRepresentable, Hashable, Codable, Sendable, + ExpressibleByStringLiteral + { + public let rawValue: String + public init(rawValue: String) { self.rawValue = rawValue } + public init(stringLiteral value: String) { self.rawValue = value } + + /// 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 wants to use a tool + public static let toolUse = StopReason(rawValue: "toolUse") + } + + /// Content representing a tool use request from the model + public struct ToolUseContent: Hashable, Codable, Sendable { + /// Unique identifier for this tool use + public let id: String + /// Name of the tool being invoked + public let name: String + /// Input parameters for the tool + public let input: [String: Value] + /// Optional metadata + public var _meta: Metadata? + + public init(id: String, name: String, input: [String: Value], _meta: Metadata? = nil) { + self.id = id + self.name = name + self.input = input + self._meta = _meta + } + } + + /// Content representing the result of a tool execution + public struct ToolResultContent: Hashable, Codable, Sendable { + /// ID of the tool use this result corresponds to + public let toolUseId: String + /// Content blocks from tool execution + public let content: [ContentBlock] + /// Structured data from tool execution + public let structuredContent: [String: Value]? + /// Whether the tool execution resulted in an error + public let isError: Bool? + /// Optional metadata + public var _meta: Metadata? + + /// Individual content blocks in tool results + public enum ContentBlock: Hashable, Sendable { + /// Text content + case text(String) + /// Image content + case image(data: String, mimeType: String) + /// Audio content + case audio(data: String, mimeType: String) + /// Embedded resource content + case resource(resource: Resource.Content, annotations: Resource.Annotations?, _meta: Metadata?) + /// Resource link + case resourceLink(uri: String, name: String, title: String?, description: String?, mimeType: String?, annotations: Resource.Annotations?) + } + + public init( + toolUseId: String, + content: [ContentBlock], + structuredContent: [String: Value]? = nil, + isError: Bool? = nil, + _meta: Metadata? = nil + ) { + self.toolUseId = toolUseId + self.content = content + self.structuredContent = structuredContent + self.isError = isError + self._meta = _meta + } + } +} + +// MARK: - Codable + +extension Sampling.Message: Codable { + 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 = try container.decode(Content.self, forKey: .content) + _meta = try container.decodeIfPresent(Metadata.self, forKey: ._meta) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode(role, forKey: .role) + try container.encode(content, forKey: .content) + try container.encodeIfPresent(_meta, forKey: ._meta) + } +} + +extension Sampling.Message.Content.ContentBlock: Codable { + private enum CodingKeys: String, CodingKey { + case type, text, data, mimeType + case id, name, input, _meta + case toolUseId, content, structuredContent, isError + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + let type = try container.decode(String.self, forKey: .type) + + switch type { + case "text": + let text = try container.decode(String.self, forKey: .text) + self = .text(text) + case "image": + let data = try container.decode(String.self, forKey: .data) + let mimeType = try container.decode(String.self, forKey: .mimeType) + self = .image(data: data, mimeType: mimeType) + 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) + case "tool_use": + let id = try container.decode(String.self, forKey: .id) + let name = try container.decode(String.self, forKey: .name) + let input = try container.decode([String: Value].self, forKey: .input) + let _meta = try container.decodeIfPresent(Metadata.self, forKey: ._meta) + self = .toolUse(Sampling.ToolUseContent(id: id, name: name, input: input, _meta: _meta)) + case "tool_result": + let toolUseId = try container.decode(String.self, forKey: .toolUseId) + let content = try container.decode([Sampling.ToolResultContent.ContentBlock].self, forKey: .content) + let structuredContent = try container.decodeIfPresent([String: Value].self, forKey: .structuredContent) + let isError = try container.decodeIfPresent(Bool.self, forKey: .isError) + let _meta = try container.decodeIfPresent(Metadata.self, forKey: ._meta) + self = .toolResult(Sampling.ToolResultContent( + toolUseId: toolUseId, + content: content, + structuredContent: structuredContent, + isError: isError, + _meta: _meta)) + default: + throw DecodingError.dataCorruptedError( + forKey: .type, in: container, + debugDescription: "Unknown sampling message content block type") + } + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + + switch self { + case .text(let text): + try container.encode("text", forKey: .type) + try container.encode(text, forKey: .text) + case .image(let data, let mimeType): + try container.encode("image", forKey: .type) + try container.encode(data, forKey: .data) + try container.encode(mimeType, forKey: .mimeType) + case .audio(let data, let mimeType): + try container.encode("audio", forKey: .type) + try container.encode(data, forKey: .data) + try container.encode(mimeType, forKey: .mimeType) + case .toolUse(let toolUse): + try container.encode("tool_use", forKey: .type) + try container.encode(toolUse.id, forKey: .id) + try container.encode(toolUse.name, forKey: .name) + try container.encode(toolUse.input, forKey: .input) + try container.encodeIfPresent(toolUse._meta, forKey: ._meta) + case .toolResult(let toolResult): + try container.encode("tool_result", forKey: .type) + try container.encode(toolResult.toolUseId, forKey: .toolUseId) + try container.encode(toolResult.content, forKey: .content) + try container.encodeIfPresent(toolResult.structuredContent, forKey: .structuredContent) + try container.encodeIfPresent(toolResult.isError, forKey: .isError) + try container.encodeIfPresent(toolResult._meta, forKey: ._meta) + } + } +} + +extension Sampling.Message.Content: Codable { + public init(from decoder: Decoder) throws { + // Try to decode as an array first + if let blocks = try? [ContentBlock](from: decoder) { + // If it's a single-element array, unwrap it to single + if blocks.count == 1, let block = blocks.first { + self = .single(block) + } else { + self = .multiple(blocks) + } + } else { + // Try to decode as a single block + let block = try ContentBlock(from: decoder) + self = .single(block) + } + } + + public func encode(to encoder: Encoder) throws { + switch self { + case .single(let block): + // Encode single block directly (not as array) + try block.encode(to: encoder) + case .multiple(let blocks): + // Encode as array + try blocks.encode(to: encoder) + } + } +} + +// MARK: - ExpressibleByStringLiteral + +extension Sampling.Message.Content: ExpressibleByStringLiteral { + public init(stringLiteral value: String) { + self = .single(.text(value)) + } +} + +// MARK: - ExpressibleByStringInterpolation + +extension Sampling.Message.Content: ExpressibleByStringInterpolation { + public init(stringInterpolation: DefaultStringInterpolation) { + self = .single(.text(String(stringInterpolation: stringInterpolation))) + } +} + +extension Sampling.ToolResultContent.ContentBlock: Codable { + private enum CodingKeys: String, CodingKey { + case type, text, data, mimeType, resource, annotations, _meta + case uri, name, title, description + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + let type = try container.decode(String.self, forKey: .type) + + switch type { + case "text": + let text = try container.decode(String.self, forKey: .text) + self = .text(text) + case "image": + let data = try container.decode(String.self, forKey: .data) + let mimeType = try container.decode(String.self, forKey: .mimeType) + self = .image(data: data, mimeType: mimeType) + 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) + case "resource": + let resourceContent = try container.decode(Resource.Content.self, forKey: .resource) + let annotations = try container.decodeIfPresent(Resource.Annotations.self, forKey: .annotations) + let _meta = try container.decodeIfPresent(Metadata.self, forKey: ._meta) + self = .resource(resource: resourceContent, annotations: annotations, _meta: _meta) + case "resource_link": + let uri = try container.decode(String.self, forKey: .uri) + let name = try container.decode(String.self, forKey: .name) + let title = try container.decodeIfPresent(String.self, forKey: .title) + let description = try container.decodeIfPresent(String.self, forKey: .description) + let mimeType = try container.decodeIfPresent(String.self, forKey: .mimeType) + let annotations = try container.decodeIfPresent(Resource.Annotations.self, forKey: .annotations) + self = .resourceLink( + uri: uri, name: name, title: title, description: description, + mimeType: mimeType, annotations: annotations) + default: + throw DecodingError.dataCorruptedError( + forKey: .type, in: container, + debugDescription: "Unknown tool result content block type") + } + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + + switch self { + case .text(let text): + try container.encode("text", forKey: .type) + try container.encode(text, forKey: .text) + case .image(let data, let mimeType): + try container.encode("image", forKey: .type) + try container.encode(data, forKey: .data) + try container.encode(mimeType, forKey: .mimeType) + case .audio(let data, let mimeType): + try container.encode("audio", forKey: .type) + try container.encode(data, forKey: .data) + try container.encode(mimeType, forKey: .mimeType) + case .resource(let resourceContent, let annotations, let _meta): + try container.encode("resource", forKey: .type) + try container.encode(resourceContent, forKey: .resource) + try container.encodeIfPresent(annotations, forKey: .annotations) + try container.encodeIfPresent(_meta, forKey: ._meta) + case .resourceLink(let uri, let name, let title, let description, let mimeType, let annotations): + try container.encode("resource_link", forKey: .type) + try container.encode(uri, forKey: .uri) + try container.encode(name, forKey: .name) + try container.encodeIfPresent(title, forKey: .title) + try container.encodeIfPresent(description, forKey: .description) + try container.encodeIfPresent(mimeType, forKey: .mimeType) + try container.encodeIfPresent(annotations, forKey: .annotations) + } + } +} + +// MARK: - + +/// To request sampling from a client, servers send a `sampling/createMessage` request. +/// - SeeAlso: https://modelcontextprotocol.io/docs/concepts/sampling#how-sampling-works +public enum CreateSamplingMessage: Method { + public static let name = "sampling/createMessage" + + /// Tool choice configuration for sampling + public struct ToolChoice: Hashable, Codable, Sendable { + /// Tool choice mode + public enum Mode: String, Hashable, Codable, Sendable { + /// Automatically decide whether to use tools + case auto + /// Require using at least one tool + case required + /// Do not use any tools + case none + } + + /// The tool choice mode. If omitted, defaults to `.auto`. + public let mode: Mode? + + public init(mode: Mode? = .auto) { + self.mode = mode + } + + private enum CodingKeys: String, CodingKey { + case mode + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + mode = try container.decodeIfPresent(Mode.self, forKey: .mode) ?? .auto + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + if let mode, mode != .auto { + try container.encode(mode, forKey: .mode) + } + } + } + + 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]? + /// Optional request metadata + public var _meta: Metadata? + /// Tools available for the model to use + public let tools: [Tool]? + /// Tool choice configuration + public let toolChoice: ToolChoice? + /// Provider-specific metadata to pass to the LLM + public let metadata: [String: Value]? + + public init( + messages: [Sampling.Message], + modelPreferences: Sampling.ModelPreferences? = nil, + systemPrompt: String? = nil, + includeContext: Sampling.ContextInclusion? = nil, + temperature: Double? = nil, + maxTokens: Int, + stopSequences: [String]? = nil, + _meta: Metadata? = nil, + tools: [Tool]? = nil, + toolChoice: ToolChoice? = nil, + metadata: [String: Value]? = nil + ) { + self.messages = messages + self.modelPreferences = modelPreferences + self.systemPrompt = systemPrompt + self.includeContext = includeContext + self.temperature = temperature + self.maxTokens = maxTokens + self.stopSequences = stopSequences + self._meta = _meta + self.tools = tools + self.toolChoice = toolChoice + self.metadata = metadata + } + } + + public struct Result: Hashable, Codable, Sendable { + /// Name of the model used + public let model: String + /// Why sampling stopped + public let stopReason: Sampling.StopReason? + /// The role of the completion + public let role: Sampling.Message.Role + /// The completion content + public let content: Sampling.Message.Content + /// Optional metadata about this result + public var _meta: Metadata? + + public init( + model: String, + stopReason: Sampling.StopReason? = nil, + role: Sampling.Message.Role, + content: Sampling.Message.Content, + _meta: Metadata? = nil + ) { + self.model = model + self.stopReason = stopReason + self.role = role + self.content = content + self._meta = _meta + } + + private enum CodingKeys: String, CodingKey, CaseIterable { + case model, stopReason, role, content, _meta + } + + 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) + } + + 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( + Sampling.StopReason.self, forKey: .stopReason) + role = try container.decode(Sampling.Message.Role.self, forKey: .role) + content = try container.decode(Sampling.Message.Content.self, forKey: .content) + _meta = try container.decodeIfPresent(Metadata.self, forKey: ._meta) + } + } +} diff --git a/Sources/MCP/Server/Completion.swift b/Sources/MCP/Server/Completion.swift new file mode 100644 index 0000000..cd8e73a --- /dev/null +++ b/Sources/MCP/Server/Completion.swift @@ -0,0 +1,183 @@ +import Foundation + +/// The Model Context Protocol (MCP) provides a standardized way for servers to offer +/// autocompletion suggestions for the arguments of prompts and resource templates. +/// +/// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/server/utilities/completion/ + +// MARK: - Reference Types + +/// A reference to a prompt by name. +/// +/// This is a typealias for `Prompt.Reference` — the two types are equivalent. +public typealias PromptReference = Prompt.Reference + +/// A reference to a resource by URI +public struct ResourceReference: Hashable, Codable, Sendable { + /// The resource URI + public let uri: String + + public init(uri: String) { + self.uri = uri + } + + private enum CodingKeys: String, CodingKey { + case type, uri + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode("ref/resource", forKey: .type) + try container.encode(uri, forKey: .uri) + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + let type = try container.decode(String.self, forKey: .type) + guard type == "ref/resource" else { + throw DecodingError.dataCorruptedError( + forKey: .type, + in: container, + debugDescription: "Expected ref/resource type" + ) + } + uri = try container.decode(String.self, forKey: .uri) + } +} + +/// A reference type for completion requests (either prompt or resource) +public enum CompletionReference: Hashable, Codable, Sendable { + /// References a prompt by name + case prompt(PromptReference) + /// References a resource URI + case resource(ResourceReference) + + 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": + self = .prompt(try PromptReference(from: decoder)) + case "ref/resource": + self = .resource(try ResourceReference(from: decoder)) + 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 + +/// To get completion suggestions, clients send a `completion/complete` request. +/// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/server/utilities/completion/ +public enum Complete: Method { + public static let name = "completion/complete" + + public struct Parameters: Hashable, Codable, Sendable { + /// The reference to what is being completed + public let ref: CompletionReference + /// The argument being completed + public let argument: Argument + /// Optional context with already-resolved arguments + public let context: Context? + + public init( + ref: CompletionReference, + argument: Argument, + context: Context? = nil + ) { + self.ref = ref + self.argument = argument + self.context = context + } + + /// The argument being completed + public struct Argument: Hashable, Codable, Sendable { + /// The argument name + public let name: String + /// The current value (partial or complete) + public let value: String + + public init(name: String, value: String) { + self.name = name + self.value = value + } + } + + /// Context containing already-resolved arguments + public struct Context: Hashable, Codable, Sendable { + /// A mapping of already-resolved argument names to their values + public let arguments: [String: String] + + public init(arguments: [String: String]) { + self.arguments = arguments + } + } + } + + public struct Result: Hashable, Codable, Sendable { + /// The completion result + public let completion: Completion + /// Optional metadata about this result + public var _meta: Metadata? + + public init(completion: Completion, _meta: Metadata? = nil) { + self.completion = completion + self._meta = _meta + } + + private enum CodingKeys: String, CodingKey, CaseIterable { + case completion, _meta + } + + 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) + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + completion = try container.decode(Completion.self, forKey: .completion) + _meta = try container.decodeIfPresent(Metadata.self, forKey: ._meta) + } + + /// Completion result containing suggested values + public struct Completion: Hashable, Codable, Sendable { + /// Array of completion values (max 100 items) + public let values: [String] + /// Optional total number of available matches + public let total: Int? + /// Whether additional results exist + public let hasMore: Bool? + + public init( + values: [String], + total: Int? = nil, + hasMore: Bool? = nil + ) { + self.values = values + self.total = total + self.hasMore = hasMore + } + } + } +} diff --git a/Sources/MCP/Server/Logging.swift b/Sources/MCP/Server/Logging.swift new file mode 100644 index 0000000..cb1b8be --- /dev/null +++ b/Sources/MCP/Server/Logging.swift @@ -0,0 +1,72 @@ +import Foundation + +/// The Model Context Protocol (MCP) provides a standardized way for servers to send +/// structured log messages to clients. Clients can control logging verbosity by setting +/// minimum log levels, with servers sending notifications containing severity levels, +/// optional logger names, and arbitrary JSON-serializable data. +/// +/// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/server/utilities/logging/ +public enum LogLevel: String, Hashable, Codable, Sendable, CaseIterable { + /// Detailed debugging information + case debug + /// General informational messages + case info + /// Normal but significant events + case notice + /// Warning conditions + case warning + /// Error conditions + case error + /// Critical conditions + case critical + /// Action must be taken immediately + case alert + /// System is unusable + case emergency +} + +// MARK: - Set Log Level + +/// To configure the minimum log level, clients MAY send a `logging/setLevel` request. +/// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/server/utilities/logging/ +public enum SetLoggingLevel: Method { + public static let name = "logging/setLevel" + + public struct Parameters: Hashable, Codable, Sendable { + /// The minimum log level to set + public let level: LogLevel + + public init(level: LogLevel) { + self.level = level + } + } + + public typealias Result = Empty +} + +// MARK: - Log Message Notification + +/// Servers send log messages using `notifications/message` notifications. +/// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/server/utilities/logging/ +public struct LogMessageNotification: Notification { + public static let name = "notifications/message" + + public struct Parameters: Hashable, Codable, Sendable { + /// The severity level of the log message + public let level: LogLevel + /// Optional logger name to identify the source + public let logger: String? + /// Arbitrary JSON-serializable data for the log message + public let data: Value + + public init( + level: LogLevel, + logger: String? = nil, + data: Value + ) { + self.level = level + self.logger = logger + self.data = data + } + } +} diff --git a/Sources/MCP/Server/Prompts.swift b/Sources/MCP/Server/Prompts.swift index 2ebd3ae..acfdcad 100644 --- a/Sources/MCP/Server/Prompts.swift +++ b/Sources/MCP/Server/Prompts.swift @@ -7,32 +7,80 @@ 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://modelcontextprotocol.io/specification/2025-11-25/server/prompts/ public struct Prompt: Hashable, Codable, Sendable { /// The prompt name public let name: String + /// A human-readable prompt title + public let title: String? /// The prompt description public let description: String? /// The prompt arguments public let arguments: [Argument]? + /// Optional set of sized icons that the client can display in a user interface + public var icons: [Icon]? + /// Optional metadata about this prompt + public var _meta: Metadata? - public init(name: String, description: String? = nil, arguments: [Argument]? = nil) { + public init( + name: String, + title: String? = nil, + description: String? = nil, + arguments: [Argument]? = nil, + icons: [Icon]? = nil, + meta: Metadata? = nil + ) { self.name = name + self.title = title self.description = description self.arguments = arguments + self.icons = icons + self._meta = meta + } + + private enum CodingKeys: String, CodingKey { + case name, title, description, arguments, icons, _meta + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode(name, forKey: .name) + try container.encodeIfPresent(title, forKey: .title) + try container.encodeIfPresent(description, forKey: .description) + try container.encodeIfPresent(arguments, forKey: .arguments) + try container.encodeIfPresent(icons, forKey: .icons) + try container.encodeIfPresent(_meta, forKey: ._meta) + } + + 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) + arguments = try container.decodeIfPresent([Argument].self, forKey: .arguments) + icons = try container.decodeIfPresent([Icon].self, forKey: .icons) + _meta = try container.decodeIfPresent(Metadata.self, forKey: ._meta) } /// An argument for a prompt public struct Argument: Hashable, Codable, Sendable { /// The argument name public let name: String + /// A human-readable argument title + 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 } @@ -53,69 +101,43 @@ public struct Prompt: Hashable, Codable, Sendable { /// The message content public let content: Content + /// 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) { self.role = role self.content = content } + /// Private initializer for convenience methods to avoid deprecation warnings + private init(_role role: Role, _content content: Content) { + self.role = role + self.content = content + } + + /// Creates a user message with the specified content + public static func user(_ content: Content) -> Message { + return 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) + } + /// Content types for messages - public enum Content: Hashable, Codable, Sendable { + public enum Content: Hashable, Sendable { /// Text content case text(text: String) /// Image content case image(data: String, mimeType: String) - /// Embedded resource content - case resource(uri: String, mimeType: String, text: String?, blob: String?) - - private enum CodingKeys: String, CodingKey { - case type, text, data, mimeType, uri, blob - } - - public func encode(to encoder: Encoder) throws { - var container = encoder.container(keyedBy: CodingKeys.self) - - switch self { - case .text(let text): - try container.encode("text", forKey: .type) - try container.encode(text, forKey: .text) - case .image(let data, let mimeType): - try container.encode("image", forKey: .type) - try container.encode(data, forKey: .data) - try container.encode(mimeType, forKey: .mimeType) - case .resource(let uri, let mimeType, let text, let blob): - 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) - } - } - - public init(from decoder: Decoder) throws { - let container = try decoder.container(keyedBy: CodingKeys.self) - let type = try container.decode(String.self, forKey: .type) - - switch type { - case "text": - let text = try container.decode(String.self, forKey: .text) - self = .text(text: text) - case "image": - let data = try container.decode(String.self, forKey: .data) - let mimeType = try container.decode(String.self, forKey: .mimeType) - self = .image(data: data, mimeType: mimeType) - 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) - default: - throw DecodingError.dataCorruptedError( - forKey: .type, - in: container, - debugDescription: "Unknown content type") - } - } + /// Audio content + case audio(data: String, mimeType: String) + /// Embedded resource content (EmbeddedResource from spec) + case resource(resource: Resource.Content, annotations: Resource.Annotations? = nil, _meta: Metadata? = nil) + /// Resource link + case resourceLink(uri: String, name: String, title: String? = nil, description: String? = nil, mimeType: String? = nil, annotations: Resource.Annotations? = nil) } } @@ -123,39 +145,137 @@ public struct Prompt: Hashable, Codable, Sendable { public struct Reference: Hashable, Codable, Sendable { /// The prompt reference name public let name: String + /// A human-readable prompt title + 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) } } } +// MARK: - Codable + +extension Prompt.Message.Content: Codable { + private enum CodingKeys: String, CodingKey { + case type, text, data, mimeType, resource, annotations, _meta + case uri, name, title, description + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + + switch self { + case .text(let text): + try container.encode("text", forKey: .type) + try container.encode(text, forKey: .text) + case .image(let data, let mimeType): + try container.encode("image", forKey: .type) + try container.encode(data, forKey: .data) + try container.encode(mimeType, forKey: .mimeType) + case .audio(let data, let mimeType): + try container.encode("audio", forKey: .type) + try container.encode(data, forKey: .data) + try container.encode(mimeType, forKey: .mimeType) + case .resource(let resourceContent, let annotations, let _meta): + try container.encode("resource", forKey: .type) + try container.encode(resourceContent, forKey: .resource) + try container.encodeIfPresent(annotations, forKey: .annotations) + try container.encodeIfPresent(_meta, forKey: ._meta) + case .resourceLink(let uri, let name, let title, let description, let mimeType, let annotations): + try container.encode("resource_link", forKey: .type) + try container.encode(uri, forKey: .uri) + try container.encode(name, forKey: .name) + try container.encodeIfPresent(title, forKey: .title) + try container.encodeIfPresent(description, forKey: .description) + try container.encodeIfPresent(mimeType, forKey: .mimeType) + try container.encodeIfPresent(annotations, forKey: .annotations) + } + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + let type = try container.decode(String.self, forKey: .type) + + switch type { + case "text": + let text = try container.decode(String.self, forKey: .text) + self = .text(text: text) + case "image": + let data = try container.decode(String.self, forKey: .data) + let mimeType = try container.decode(String.self, forKey: .mimeType) + self = .image(data: data, mimeType: mimeType) + 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) + case "resource": + let resourceContent = try container.decode(Resource.Content.self, forKey: .resource) + let annotations = try container.decodeIfPresent(Resource.Annotations.self, forKey: .annotations) + let _meta = try container.decodeIfPresent(Metadata.self, forKey: ._meta) + self = .resource(resource: resourceContent, annotations: annotations, _meta: _meta) + case "resource_link": + let uri = try container.decode(String.self, forKey: .uri) + let name = try container.decode(String.self, forKey: .name) + let title = try container.decodeIfPresent(String.self, forKey: .title) + let description = try container.decodeIfPresent(String.self, forKey: .description) + let mimeType = try container.decodeIfPresent(String.self, forKey: .mimeType) + let annotations = try container.decodeIfPresent(Resource.Annotations.self, forKey: .annotations) + self = .resourceLink(uri: uri, name: name, title: title, description: description, mimeType: mimeType, annotations: annotations) + default: + throw DecodingError.dataCorruptedError( + forKey: .type, + in: container, + debugDescription: "Unknown content type") + } + } +} + +// MARK: - ExpressibleByStringLiteral + +extension Prompt.Message.Content: ExpressibleByStringLiteral { + public init(stringLiteral value: String) { + self = .text(text: value) + } +} + +// MARK: - ExpressibleByStringInterpolation + +extension Prompt.Message.Content: ExpressibleByStringInterpolation { + public init(stringInterpolation: DefaultStringInterpolation) { + self = .text(text: String(stringInterpolation: stringInterpolation)) + } +} + // MARK: - /// To retrieve available prompts, clients send a `prompts/list` request. -/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2024-11-05/server/prompts/#listing-prompts +/// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/server/prompts/#listing-prompts public enum ListPrompts: Method { public static let name: String = "prompts/list" public struct Parameters: NotRequired, Hashable, Codable, Sendable { public let cursor: String? - + public init() { self.cursor = nil } @@ -168,25 +288,49 @@ public enum ListPrompts: Method { public struct Result: Hashable, Codable, Sendable { public let prompts: [Prompt] public let nextCursor: String? + public var _meta: Metadata? - public init(prompts: [Prompt], nextCursor: String? = nil) { + public init( + prompts: [Prompt], + nextCursor: String? = nil, + _meta: Metadata? = nil + ) { self.prompts = prompts self.nextCursor = nextCursor + self._meta = _meta + } + + private enum CodingKeys: String, CodingKey, CaseIterable { + case prompts, nextCursor, _meta + } + + 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) + } + + 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(Metadata.self, forKey: ._meta) } } } /// To retrieve a specific prompt, clients send a `prompts/get` request. /// Arguments may be auto-completed through the completion API. -/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2024-11-05/server/prompts/#getting-a-prompt +/// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/server/prompts/#getting-a-prompt public enum GetPrompt: Method { public static let name: String = "prompts/get" public struct Parameters: Hashable, Codable, Sendable { public let name: String - public let arguments: [String: Value]? + public let arguments: [String: String]? - public init(name: String, arguments: [String: Value]? = nil) { + public init(name: String, arguments: [String: String]? = nil) { self.name = name self.arguments = arguments } @@ -195,16 +339,41 @@ public enum GetPrompt: Method { public struct Result: Hashable, Codable, Sendable { public let description: String? public let messages: [Prompt.Message] + /// Optional metadata about this result + public var _meta: Metadata? - public init(description: String?, messages: [Prompt.Message]) { + public init( + description: String? = nil, + messages: [Prompt.Message], + _meta: Metadata? = nil + ) { self.description = description self.messages = messages + self._meta = _meta + } + + private enum CodingKeys: String, CodingKey, CaseIterable { + case description, messages, _meta + } + + 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) + } + + 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(Metadata.self, forKey: ._meta) } } } /// When the list of available prompts changes, servers that declared the listChanged capability SHOULD send a notification. -/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2024-11-05/server/prompts/#list-changed-notification +/// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/server/prompts/#list-changed-notification public struct PromptListChangedNotification: Notification { public static let name: String = "notifications/prompts/list_changed" } diff --git a/Sources/MCP/Server/Resources.swift b/Sources/MCP/Server/Resources.swift index 311ff09..fb512a0 100644 --- a/Sources/MCP/Server/Resources.swift +++ b/Sources/MCP/Server/Resources.swift @@ -6,31 +6,85 @@ 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://modelcontextprotocol.io/specification/2025-11-25/server/resources/ public struct Resource: Hashable, Codable, Sendable { /// The resource name public var name: String + /// A human-readable resource title + 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]? + /// Raw content size in bytes + public var size: Int? + /// Annotations for this resource + public var annotations: Resource.Annotations? + /// Optional set of sized icons that the client can display in a user interface + public var icons: [Icon]? + /// Metadata fields for the resource (see spec for _meta usage) + public var _meta: Metadata? public init( name: String, uri: String, + title: String? = nil, description: String? = nil, mimeType: String? = nil, - metadata: [String: String]? = nil + size: Int? = nil, + annotations: Resource.Annotations? = nil, + icons: [Icon]? = nil, + _meta: Metadata? = 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.icons = icons + self._meta = _meta + } + + private enum CodingKeys: String, CodingKey { + case name + case uri + case title + 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) + name = try container.decode(String.self, forKey: .name) + uri = try container.decode(String.self, forKey: .uri) + title = try container.decodeIfPresent(String.self, forKey: .title) + 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(Resource.Annotations.self, forKey: .annotations) + icons = try container.decodeIfPresent([Icon].self, forKey: .icons) + _meta = try container.decodeIfPresent(Metadata.self, forKey: ._meta) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode(name, forKey: .name) + try container.encode(uri, forKey: .uri) + try container.encodeIfPresent(title, forKey: .title) + 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) } /// Content of a resource. @@ -43,27 +97,82 @@ public struct Resource: Hashable, Codable, Sendable { public let text: String? /// The resource binary content public let blob: String? + /// Metadata fields (see spec for _meta usage) + public var _meta: Metadata? - public static func text(_ content: String, uri: String, mimeType: String? = nil) -> Self { - .init(uri: uri, mimeType: mimeType, text: content) + public static func text( + _ content: String, + uri: String, + mimeType: String? = nil, + _meta: Metadata? = nil + ) -> Self { + .init(uri: uri, mimeType: mimeType, text: content, _meta: _meta) } - public static func binary(_ data: Data, uri: String, mimeType: String? = nil) -> Self { - .init(uri: uri, mimeType: mimeType, blob: data.base64EncodedString()) + public static func binary( + _ data: Data, + uri: String, + mimeType: String? = nil, + _meta: Metadata? = nil + ) -> Self { + .init( + uri: uri, + mimeType: mimeType, + blob: data.base64EncodedString(), + _meta: _meta + ) } - private init(uri: String, mimeType: String? = nil, text: String? = nil) { + private init( + uri: String, + mimeType: String? = nil, + text: String? = nil, + _meta: Metadata? = nil + ) { self.uri = uri self.mimeType = mimeType self.text = text self.blob = nil + self._meta = _meta } - private init(uri: String, mimeType: String? = nil, blob: String) { + private init( + uri: String, + mimeType: String? = nil, + blob: String, + _meta: Metadata? = nil + ) { self.uri = uri self.mimeType = mimeType self.text = nil self.blob = blob + self._meta = _meta + } + + private enum CodingKeys: String, CodingKey { + case uri + case mimeType + case text + case blob + case _meta + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + uri = try container.decode(String.self, forKey: .uri) + mimeType = try container.decodeIfPresent(String.self, forKey: .mimeType) + text = try container.decodeIfPresent(String.self, forKey: .text) + blob = try container.decodeIfPresent(String.self, forKey: .blob) + _meta = try container.decodeIfPresent(Metadata.self, forKey: ._meta) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encode(uri, forKey: .uri) + try container.encodeIfPresent(mimeType, forKey: .mimeType) + try container.encodeIfPresent(text, forKey: .text) + try container.encodeIfPresent(blob, forKey: .blob) + try container.encodeIfPresent(_meta, forKey: ._meta) } } @@ -73,21 +182,65 @@ public struct Resource: Hashable, Codable, Sendable { public var uriTemplate: String /// The template name public var name: String + /// A human-readable template title + public var title: String? /// The template description public var description: String? /// The resource MIME type public var mimeType: String? + /// Annotations for this template + public var annotations: Resource.Annotations? + /// Optional set of sized icons that the client can display in a user interface + public var icons: [Icon]? + /// Metadata fields for the template (see spec for _meta usage) + public var _meta: Metadata? public init( uriTemplate: String, name: String, + title: String? = nil, description: String? = nil, - mimeType: String? = nil + mimeType: String? = nil, + annotations: Resource.Annotations? = nil, + icons: [Icon]? = nil, + _meta: Metadata? = nil ) { self.uriTemplate = uriTemplate self.name = name + self.title = title self.description = description self.mimeType = mimeType + self.annotations = annotations + self.icons = icons + self._meta = _meta + } + } + + // A resource annotation. + public struct Annotations: Hashable, Codable, Sendable { + /// The intended audience for this resource. + public enum Audience: String, Hashable, Codable, Sendable { + /// Content intended for end users. + case user = "user" + /// Content intended for AI assistants. + case assistant = "assistant" + } + + /// An array indicating the intended audience(s) for this resource. For example, `[.user, .assistant]` indicates content useful for both. + public let audience: [Audience]? + /// A number from 0.0 to 1.0 indicating the importance of this resource. A value of 1 means "most important" (effectively required), while 0 means "least important". + public let priority: Double? + /// An ISO 8601 formatted timestamp indicating when the resource was last modified (e.g., "2025-01-12T15:00:58Z"). + public let lastModified: String? + + public init( + audience: [Audience]? = nil, + priority: Double? = nil, + lastModified: String? = nil + ) { + self.audience = audience + self.priority = priority + self.lastModified = lastModified } } } @@ -95,7 +248,7 @@ public struct Resource: Hashable, Codable, Sendable { // MARK: - /// To discover available resources, clients send a `resources/list` request. -/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2024-11-05/server/resources/#listing-resources +/// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/server/resources/#listing-resources public enum ListResources: Method { public static let name: String = "resources/list" @@ -105,7 +258,7 @@ public enum ListResources: Method { public init() { self.cursor = nil } - + public init(cursor: String) { self.cursor = cursor } @@ -114,16 +267,40 @@ public enum ListResources: Method { public struct Result: Hashable, Codable, Sendable { public let resources: [Resource] public let nextCursor: String? + public var _meta: Metadata? - public init(resources: [Resource], nextCursor: String? = nil) { + public init( + resources: [Resource], + nextCursor: String? = nil, + _meta: Metadata? = nil + ) { self.resources = resources self.nextCursor = nextCursor + self._meta = _meta + } + + private enum CodingKeys: String, CodingKey, CaseIterable { + case resources, nextCursor, _meta + } + + 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) + } + + 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(Metadata.self, forKey: ._meta) } } } /// To retrieve resource contents, clients send a `resources/read` request: -/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2024-11-05/server/resources/#reading-resources +/// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/server/resources/#reading-resources public enum ReadResource: Method { public static let name: String = "resources/read" @@ -137,35 +314,92 @@ public enum ReadResource: Method { public struct Result: Hashable, Codable, Sendable { public let contents: [Resource.Content] + /// Optional metadata about this result + public var _meta: Metadata? - public init(contents: [Resource.Content]) { + public init( + contents: [Resource.Content], + _meta: Metadata? = nil + ) { self.contents = contents + self._meta = _meta + } + + private enum CodingKeys: String, CodingKey, CaseIterable { + case contents, _meta + } + + 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) + } + + 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(Metadata.self, forKey: ._meta) } } } /// To discover available resource templates, clients send a `resources/templates/list` request. -/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2024-11-05/server/resources/#resource-templates +/// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/server/resources/#resource-templates public enum ListResourceTemplates: Method { public static let name: String = "resources/templates/list" - public typealias Parameters = Empty + public struct Parameters: NotRequired, Hashable, Codable, Sendable { + public let cursor: String? + + public init() { + self.cursor = nil + } + + public init(cursor: String) { + self.cursor = cursor + } + } public struct Result: Hashable, Codable, Sendable { public let templates: [Resource.Template] + public let nextCursor: String? + /// Optional metadata about this result + public var _meta: Metadata? - public init(templates: [Resource.Template]) { + public init( + templates: [Resource.Template], + nextCursor: String? = nil, + _meta: Metadata? = nil + ) { self.templates = templates + self.nextCursor = nextCursor + self._meta = _meta } - private enum CodingKeys: String, CodingKey { + private enum CodingKeys: String, CodingKey, CaseIterable { case templates = "resourceTemplates" + case nextCursor + case _meta + } + + 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) + } + + 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(Metadata.self, forKey: ._meta) } } } /// When the list of available resources changes, servers that declared the listChanged capability SHOULD send a notification. -/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2024-11-05/server/resources/#list-changed-notification +/// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/server/resources/#list-changed-notification public struct ResourceListChangedNotification: Notification { public static let name: String = "notifications/resources/list_changed" @@ -173,7 +407,7 @@ public struct ResourceListChangedNotification: Notification { } /// 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://modelcontextprotocol.io/specification/2025-11-25/server/resources/#subscriptions public enum ResourceSubscribe: Method { public static let name: String = "resources/subscribe" @@ -184,8 +418,20 @@ public enum ResourceSubscribe: Method { public typealias Result = Empty } +/// Sent from the client to request cancellation of resources/updated notifications from the server. This should follow a previous resources/subscribe request. +/// - SeeAlso: https://modelcontextprotocol.io/specification/2025-06-18/schema#unsubscriberequest +public enum ResourceUnsubscribe: Method { + public static let name: String = "resources/unsubscribe" + + public struct Parameters: Hashable, Codable, Sendable { + public let uri: String + } + + public typealias Result = Empty +} + /// When a resource changes, servers that declared the updated capability SHOULD send a notification to subscribed clients. -/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2024-11-05/server/resources/#subscriptions +/// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/server/resources/#subscriptions public struct ResourceUpdatedNotification: Notification { public static let name: String = "notifications/resources/updated" diff --git a/Sources/MCP/Server/Server.swift b/Sources/MCP/Server/Server.swift index 168b062..060b51f 100644 --- a/Sources/MCP/Server/Server.swift +++ b/Sources/MCP/Server/Server.swift @@ -1,5 +1,4 @@ import Logging -import SystemPackage import struct Foundation.Data import struct Foundation.Date @@ -31,12 +30,31 @@ public actor Server { public struct Info: Hashable, Codable, Sendable { /// The server name public let name: String + /// A human-readable server title for display + public let title: String? /// The server version public let version: String + /// Optional description of the server + public let description: String? + /// Optional website URL for the server + public let websiteUrl: String? + /// Optional set of sized icons for display in a user interface + public let icons: [Icon]? - public init(name: String, version: String) { + public init( + name: String, + version: String, + title: String? = nil, + description: String? = nil, + websiteUrl: String? = nil, + icons: [Icon]? = nil + ) { self.name = name + self.title = title self.version = version + self.description = description + self.websiteUrl = websiteUrl + self.icons = icons } } @@ -83,6 +101,13 @@ public actor Server { public init() {} } + /// Completions capabilities + public struct Completions: Hashable, Codable, Sendable { + public init() {} + } + + /// Completions capabilities + public var completions: Completions? /// Logging capabilities public var logging: Logging? /// Prompts capabilities @@ -93,11 +118,13 @@ public actor Server { public var tools: Tools? public init( + completions: Completions? = nil, logging: Logging? = nil, prompts: Prompts? = nil, resources: Resources? = nil, tools: Tools? = nil ) { + self.completions = completions self.logging = logging self.prompts = prompts self.resources = resources @@ -118,8 +145,17 @@ public actor Server { /// The server name public nonisolated var name: String { serverInfo.name } + /// A human-readable server title + public nonisolated var title: String? { serverInfo.title } /// The server version public nonisolated var version: String { serverInfo.version } + /// Instructions describing how to use the server and its features + /// + /// This can be used by clients to improve the LLM's understanding of + /// available tools, resources, etc. + /// It can be thought of like a "hint" to the model. + /// For example, this information MAY be added to the system prompt. + public nonisolated let instructions: String? /// The server capabilities public var capabilities: Capabilities /// The server configuration @@ -129,6 +165,11 @@ public actor Server { private var methodHandlers: [String: RequestHandlerBox] = [:] /// Notification handlers private var notificationHandlers: [String: [NotificationHandlerBox]] = [:] + /// Pending request tasks (for cancellation support) + private var pendingRequestTasks: [ID: Task, Error>] = [:] + + /// Pending requests sent to the client, awaiting responses + private var pendingRequests: [ID: AnyPendingRequest] = [:] /// Whether the server is initialized private var isInitialized = false @@ -146,40 +187,54 @@ public actor Server { public init( name: String, version: String, + title: 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) self.capabilities = capabilities self.configuration = configuration + self.instructions = instructions } /// Start the server - public func start(transport: any Transport) async throws { + /// - Parameters: + /// - transport: The transport to use for the server + /// - initializeHook: An optional hook that runs when the client sends an initialize request + public func start( + transport: any Transport, + initializeHook: (@Sendable (Client.Info, Client.Capabilities) async throws -> Void)? = nil + ) async throws { self.connection = transport - registerDefaultHandlers() + registerDefaultHandlers(initializeHook: initializeHook) + registerCancellationHandler() try await transport.connect() - await logger?.info( - "Server started", metadata: ["name": "\(name)", "version": "\(version)"]) + await logger?.debug( + "Server started", metadata: ["name": "\(name)", "version": "\(version)"] + ) // Start message handling loop task = Task { do { let stream = await transport.receive() - for try await string in stream { + for try await data in stream { if Task.isCancelled { break } // Check cancellation inside loop var requestID: ID? do { - guard let data = string.data(using: .utf8) else { - throw Error.parseError("Invalid UTF-8 data") - } - - // Attempt to decode string data as AnyRequest or AnyMessage + // Attempt to decode as batch first, then as individual response, request, or notification let decoder = JSONDecoder() - if let request = try? decoder.decode(AnyRequest.self, from: data) { - try await handleRequest(request) + if let batch = try? decoder.decode(Server.Batch.self, from: data) { + try await handleBatch(batch) + } 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 request in a separate task to avoid blocking the receive loop + Task { + _ = try? await self.handleRequest(request, sendResponse: true) + } } else if let message = try? decoder.decode(AnyMessage.self, from: data) { try await handleMessage(message) } else { @@ -194,18 +249,19 @@ public actor Server { requestID = .number(intValue) } } - throw Error.parseError("Invalid message format") + throw MCPError.parseError("Invalid message format") } - } catch let error as Errno where error == .resourceTemporarilyUnavailable { + } catch let error where MCPError.isResourceTemporarilyUnavailable(error) { // Resource temporarily unavailable, retry after a short delay - try? await Task.sleep(nanoseconds: 10_000_000) // 10ms + try? await Task.sleep(for: .milliseconds(10)) continue } catch { await logger?.error( "Error processing message", metadata: ["error": "\(error)"]) let response = AnyMethod.response( id: requestID ?? .random, - error: error as? Error ?? Error.internalError(error.localizedDescription) + error: error as? MCPError + ?? MCPError.internalError(error.localizedDescription) ) try? await send(response) } @@ -214,6 +270,7 @@ public actor Server { await logger?.error( "Fatal error in message handling loop", metadata: ["error": "\(error)"]) } + await logger?.debug("Server finished", metadata: [:]) } } @@ -221,12 +278,62 @@ public actor Server { public func stop() async { task?.cancel() task = nil + + // Clear pending requests with errors + let pendingRequestsToCancel = self.pendingRequests + self.pendingRequests = [:] + for (_, request) in pendingRequestsToCancel { + request.resume(throwing: MCPError.internalError("Server disconnected")) + } + if let connection = connection { await connection.disconnect() } connection = nil } + public func waitUntilCompleted() async { + await task?.value + } + + // MARK: - Request Context + + /// Per-dispatch context the SDK attaches to the currently executing handler. + /// + /// Exposed via the ``Server/currentHandlerContext`` task-local so method + /// handlers can observe the originating HTTP request (headers, auth, path, + /// body) without changing the `withMethodHandler` signature. + public struct HandlerContext: Sendable { + /// The JSON-RPC request id of the in-flight request. SDK-internal use + /// (e.g. transports closing an SSE stream mid-call per SEP-1699). + package let id: ID + + /// The originating HTTP request, if the active transport conforms to + /// ``HTTPContextProviding``. `nil` for transports that don't carry HTTP + /// context (stdio, in-memory) or for handlers reached off the dispatch + /// path. + public let httpContext: HTTPRequest? + + package init(id: ID, httpContext: HTTPRequest?) { + self.id = id + self.httpContext = httpContext + } + } + + /// The handler context for the currently executing method handler. + /// + /// Set via `@TaskLocal` before dispatching each request, so it propagates + /// automatically into the handler task. `nil` outside of a handler. + /// + /// When spawning `Task.detached { … }` from inside a handler, capture the + /// value up front — detached tasks do not inherit task-locals: + /// + /// ```swift + /// let ctx = Server.currentHandlerContext + /// Task.detached { await doWork(with: ctx?.httpContext) } + /// ``` + @TaskLocal public static var currentHandlerContext: HandlerContext? = nil + // MARK: - Registration /// Register a method handler @@ -255,32 +362,392 @@ public actor Server { // MARK: - Sending - /// Send a response to a client + /// Send a response to a request public func send(_ response: Response) async throws { guard let connection = connection else { - throw Error.internalError("Server connection not initialized") - } - let responseData = try JSONEncoder().encode(response) - if let responseStr = String(data: responseData, encoding: .utf8) { - try await connection.send(responseStr) + 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 Error.internalError("Server connection not initialized") + throw MCPError.internalError("Server connection not initialized") } - let notificationData = try JSONEncoder().encode(notification) - if let notificationStr = String(data: notificationData, encoding: .utf8) { - try await connection.send(notificationStr) + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + + let notificationData = try encoder.encode(notification) + try await connection.send(notificationData) + } + + /// Send a request to the client and return a Task for the response + private func send(_ request: Request) throws -> Task { + 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) + + let requestTask = Task { + try await withCheckedThrowingContinuation { continuation in + Task { + // Add pending response before sending + self.addPendingResponse( + id: request.id, + continuation: continuation, + type: M.Result.self + ) + + // Send the request + do { + try await connection.send(requestData) + } catch { + // If send fails, remove pending response and resume with error + if self.removePendingResponse(id: request.id) != nil { + continuation.resume(throwing: error) + } + } + } + } + } + + return requestTask + } + + /// Send a request and await its response + private func sendAndAwait(_ request: Request) async throws -> M.Result { + let task = try send(request) + return try await task.value + } + + private func addPendingResponse( + id: ID, + continuation: CheckedContinuation, + type: T.Type + ) { + pendingRequests[id] = AnyPendingRequest( + PendingRequest(continuation: continuation) + ) + } + + private func removePendingResponse(id: ID) -> AnyPendingRequest? { + return pendingRequests.removeValue(forKey: id) + } + + // MARK: - Sampling + + /// Request sampling from the connected client + /// + /// Sampling allows servers to request LLM completions through the client, + /// enabling sophisticated agentic behaviors while maintaining human-in-the-loop control. + /// + /// The sampling flow follows these steps: + /// 1. Server sends a `sampling/createMessage` request to the client + /// 2. Client reviews the request and can modify it + /// 3. Client samples from an LLM + /// 4. Client reviews the completion + /// 5. Client returns the result to the server + /// + /// - Parameters: + /// - messages: The conversation history to send to the LLM + /// - modelPreferences: Model selection preferences + /// - systemPrompt: Optional system prompt + /// - includeContext: What MCP context to include + /// - temperature: Controls randomness (0.0 to 1.0) + /// - maxTokens: Maximum tokens to generate + /// - stopSequences: Array of sequences that stop generation + /// - _meta: Optional request metadata + /// - 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, + _meta: Metadata? = nil + ) async throws -> CreateSamplingMessage.Result { + guard connection != nil else { + throw MCPError.internalError("Server connection not initialized") + } + + try validateClientCapability(\.sampling, "Sampling") + + let request = CreateSamplingMessage.request( + .init( + messages: messages, + modelPreferences: modelPreferences, + systemPrompt: systemPrompt, + includeContext: includeContext, + temperature: temperature, + maxTokens: maxTokens, + stopSequences: stopSequences, + _meta: _meta + ) + ) + + let result = try await sendAndAwait(request) + return result + } + + // MARK: - Elicitation + + /// Request user input from the client using form-based elicitation + /// + /// Elicitation allows servers to request user input during operations. + /// This is useful for collecting user feedback, confirmations, or data + /// that the server needs but doesn't have. + /// + /// The flow: + /// 1. Server requests elicitation with a message and optional schema + /// 2. Client displays the request to the user + /// 3. User provides input or declines + /// 4. Client returns the result to the server + /// + /// - Parameters: + /// - message: The message to display to the user + /// - mode: The elicitation mode (form or url) + /// - requestedSchema: Optional JSON schema describing the expected response + /// - _meta: Optional request metadata + /// - Returns: The elicitation result containing the action and optional content + /// - Throws: MCPError if the request fails + /// - SeeAlso: https://modelcontextprotocol.io/docs/concepts/elicitation + public func requestElicitation( + message: String, + requestedSchema: Elicitation.RequestSchema, + mode: Elicitation.Mode? = nil, + _meta: Metadata? = nil + ) async throws -> CreateElicitation.Result { + guard connection != nil else { + throw MCPError.internalError("Server connection not initialized") + } + + try validateClientCapability(\.elicitation, "Elicitation") + + let request = CreateElicitation.request( + .form( + .init( + message: message, + mode: mode, + requestedSchema: requestedSchema, + _meta: _meta + ) + ) + ) + + let result = try await sendAndAwait(request) + return result + } + + /// Request user input from the client using URL-based elicitation + /// + /// URL-based elicitation directs the user to an external URL for authentication + /// or data collection. This is useful for OAuth flows or other web-based input. + /// + /// - Parameters: + /// - message: The message to display to the user + /// - url: The URL to direct the user to + /// - elicitationId: Unique identifier for this elicitation + /// - _meta: Optional request metadata + /// - Returns: The elicitation result containing the action and optional content + /// - Throws: MCPError if the request fails + /// - SeeAlso: https://modelcontextprotocol.io/docs/concepts/elicitation + public func requestElicitation( + message: String, + url: String, + elicitationId: String, + _meta: Metadata? = nil + ) async throws -> CreateElicitation.Result { + guard connection != nil else { + throw MCPError.internalError("Server connection not initialized") + } + + try validateClientCapability(\.elicitation, "Elicitation") + + let request = CreateElicitation.request( + .url( + .init( + message: message, + url: url, + elicitationId: elicitationId, + _meta: _meta + ) + ) + ) + + let result = try await sendAndAwait(request) + return result + } + + // MARK: - Logging + + /// Send a log message notification to connected clients. + /// + /// Servers that declare the `logging` capability can send structured log messages + /// to clients. The client controls which severity levels it wants to receive via + /// the `logging/setLevel` request. + /// + /// - Parameters: + /// - level: The severity level of the log message + /// - logger: Optional logger name to identify the source + /// - data: Arbitrary JSON-serializable data for the log message + /// - Throws: MCPError if the server is not connected + /// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/server/utilities/logging/ + public func log( + level: LogLevel, + logger: String? = nil, + data: Value + ) async throws { + let notification = LogMessageNotification.message( + .init(level: level, logger: logger, data: data) + ) + try await notify(notification) + } + + /// Send a log message notification with codable data. + /// + /// Convenience method that encodes data to JSON before sending. + /// + /// - Parameters: + /// - level: The severity level of the log message + /// - logger: Optional logger name to identify the source + /// - data: Any codable data for the log message + /// - Throws: MCPError if the server is not connected or encoding fails + /// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/server/utilities/logging/ + public func log( + level: LogLevel, + logger: String? = nil, + data: T + ) async throws { + let value = try Value(data) + try await log(level: level, logger: logger, data: value) + } + + // MARK: - Roots + + /// Request the list of roots from the connected client + /// + /// Roots define filesystem boundaries that servers can operate within. + /// The client must have declared the `roots` capability and registered + /// a roots handler for this to work. + /// + /// - Returns: Array of Root objects representing accessible directories/files + /// - Throws: MCPError if the client doesn't support roots or request fails + /// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/client/roots + public func listRoots() async throws -> [Root] { + guard connection != nil else { + throw MCPError.internalError("Server connection not initialized") + } + + try validateClientCapability(\.roots, "Roots") + + let request = ListRoots.request() + let result = try await sendAndAwait(request) + return result.roots + } + + /// 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 } } - // MARK: - + /// Process a batch of requests and/or notifications + private func handleBatch(_ batch: Batch) async throws { + await logger?.trace("Processing batch request", metadata: ["size": "\(batch.items.count)"]) - private func handleRequest(_ request: Request) async throws { - await logger?.debug( + 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) + 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 + if !responses.isEmpty { + 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 + + /// 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? + { + // 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)", @@ -300,26 +767,87 @@ public actor Server { // Find handler for method name guard let handler = methodHandlers[request.method] else { - let error = Error.methodNotFound("Unknown method: \(request.method)") + let error = MCPError.methodNotFound("Unknown method: \(request.method)") let response = AnyMethod.response(id: request.id, error: error) - try await send(response) - throw error + + if sendResponse { + try await send(response) + return nil + } + + return response + } + + // Ask the transport for the originating HTTP request so handlers can observe + // headers/auth via `Server.currentHandlerContext?.httpContext`. Transports + // that don't carry HTTP context (stdio, in-memory) don't conform. + let httpContext = await (connection as? any HTTPContextProviding)? + .httpRequestContext(for: request.id) + let handlerContext = HandlerContext(id: request.id, httpContext: httpContext) + + // Create a task to handle the request with cancellation support. + // Set currentHandlerContext as a task local so handlers see it. + var handlerTask: Task, Error>! + Server.$currentHandlerContext.withValue(handlerContext) { + handlerTask = Task, Error> { + do { + // Check if task was cancelled before starting + try Task.checkCancellation() + + // Handle request and get response + let response = try await handler(request) + return response + } catch is CancellationError { + // Request was cancelled, don't send a response per MCP spec + await logger?.debug( + "Request cancelled", + metadata: ["id": "\(request.id)", "method": "\(request.method)"] + ) + throw CancellationError() + } catch { + let mcpError = + error as? MCPError ?? MCPError.internalError(error.localizedDescription) + return AnyMethod.response(id: request.id, error: mcpError) + } + } + } + + // Store the handler task for potential cancellation + pendingRequestTasks[request.id] = handlerTask + + // Ensure cleanup happens regardless of success or failure + defer { + pendingRequestTasks.removeValue(forKey: request.id) } do { - // Handle request and get response - let response = try await handler(request) - try await send(response) + let response = try await handlerTask.value + + if sendResponse { + try await send(response) + return nil + } + + return response + } catch is CancellationError { + // Request was cancelled, don't send a response per MCP spec + return nil } catch { - let mcpError = error as? Error ?? Error.internalError(error.localizedDescription) + // This should not happen as errors are caught in the task + let mcpError = error as? MCPError ?? MCPError.internalError(error.localizedDescription) let response = AnyMethod.response(id: request.id, error: mcpError) - try await send(response) - throw error + + if sendResponse { + try await send(response) + return nil + } + + return response } } private func handleMessage(_ message: Message) async throws { - await logger?.debug( + await logger?.trace( "Processing notification", metadata: ["method": "\(message.method)"]) @@ -348,49 +876,82 @@ public actor Server { } } - private func checkInitialized() throws { - guard isInitialized else { - throw Error.invalidRequest("Server is not initialized") + private func handleResponse(_ response: Response) async { + if let pendingRequest = self.removePendingResponse(id: 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)"] + ) } } - private func registerDefaultHandlers() { + private func checkInitialized() throws { + guard isInitialized else { + throw MCPError.invalidRequest("Server is not initialized") + } + } + + /// Validate the client capabilities. + /// Throws an error if the server is configured to be strict and the capability is not supported. + private func validateClientCapability( + _ keyPath: KeyPath, + _ name: String + ) + throws + { + if configuration.strict { + guard let capabilities = clientCapabilities else { + throw MCPError.methodNotFound("Client capabilities not initialized") + } + guard capabilities[keyPath: keyPath] != nil else { + throw MCPError.methodNotFound("\(name) is not supported by the client") + } + } + } + + private func registerDefaultHandlers( + initializeHook: (@Sendable (Client.Info, Client.Capabilities) async throws -> Void)? + ) { // Initialize withMethodHandler(Initialize.self) { [weak self] params in guard let self = self else { - throw Error.internalError("Server was deallocated") + throw MCPError.internalError("Server was deallocated") } guard await !self.isInitialized else { - throw Error.invalidRequest("Server is already initialized") + throw MCPError.invalidRequest("Server is already initialized") } - // Validate protocol version - guard Version.latest == params.protocolVersion else { - throw Error.invalidRequest( - "Unsupported protocol version: \(params.protocolVersion)") + // Call initialization hook if registered + if let hook = initializeHook { + try await hook(params.clientInfo, params.capabilities) } + // Perform version negotiation + let clientRequestedVersion = params.protocolVersion + let negotiatedProtocolVersion = Version.negotiate( + clientRequestedVersion: clientRequestedVersion) + + // Set initial state with the negotiated protocol version await self.setInitialState( clientInfo: params.clientInfo, clientCapabilities: params.capabilities, - protocolVersion: params.protocolVersion + protocolVersion: negotiatedProtocolVersion ) - let result = Initialize.Result( - protocolVersion: Version.latest, + return Initialize.Result( + protocolVersion: negotiatedProtocolVersion, capabilities: await self.capabilities, serverInfo: self.serverInfo, - instructions: nil + instructions: self.instructions ) - - // Send initialized notification after a short delay - Task { - try? await Task.sleep(nanoseconds: 100_000_000) // 100ms - try? await self.notify(InitializedNotification.message()) - } - - return result } // Ping @@ -407,4 +968,121 @@ public actor Server { self.protocolVersion = protocolVersion self.isInitialized = true } + + /// Cancel and remove a pending request task + private func removePendingRequest(id: ID) -> Task, Error>? { + pendingRequestTasks.removeValue(forKey: id) + } + + private func registerCancellationHandler() { + onNotification(CancelledNotification.self) { [weak self] message in + guard let self = self else { return } + + let requestId = message.params.requestId + let reason = message.params.reason + + await self.logger?.debug( + "Received cancellation notification", + metadata: [ + "requestId": requestId.map { "\($0)" } ?? "none", + "reason": reason.map { "\($0)" } ?? "none", + ] + ) + + guard let requestId = requestId else { + await self.logger?.warning( + "Received cancellation notification with no requestId (violates spec MUST)", + metadata: ["reason": reason.map { "\($0)" } ?? "none"] + ) + return + } + + // Cancel the pending request task if it exists and remove from tracking + if let task = await self.removePendingRequest(id: requestId) { + task.cancel() + await self.logger?.debug( + "Cancelled request", + metadata: ["requestId": "\(requestId)"] + ) + } else { + // Request may have already completed or is unknown + // Per MCP spec, we should ignore this gracefully + await self.logger?.trace( + "Cancellation notification for unknown or completed request", + metadata: ["requestId": "\(requestId)"] + ) + } + } + } + + /// Cancel a request by sending a CancelledNotification to the client. + /// + /// This is used when the server needs to cancel an in-progress request it made to the client + /// (e.g., a sampling request). + /// + /// According to the MCP specification, cancellation is advisory: + /// - The client SHOULD stop processing and free resources + /// - The client MAY ignore the cancellation if the request is unknown, already completed, + /// or cannot be cancelled + /// - The server SHOULD ignore any response that arrives after cancellation + /// + /// - Parameters: + /// - requestID: The ID of the request to cancel + /// - reason: An optional human-readable reason for the cancellation + /// - Throws: MCPError if the notification cannot be sent + /// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/cancellation + public func cancelRequest(_ requestID: ID, reason: String? = nil) async throws { + // Send cancellation notification to client + let notification = CancelledNotification.message( + .init(requestId: requestID, reason: reason) + ) + try await notify(notification) + } +} + +extension Server.Batch: Codable { + init(from decoder: Decoder) throws { + let container = try decoder.singleValueContainer() + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + var items: [Item] = [] + for item in try container.decode([Value].self) { + let data = try encoder.encode(item) + try items.append(decoder.decode(Item.self, from: data)) + } + + self.items = items + } + + func encode(to encoder: Encoder) throws { + var container = encoder.singleValueContainer() + try container.encode(items) + } +} + +extension Server.Batch.Item: Codable { + private enum CodingKeys: String, CodingKey { + case id + } + + init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + // Check if it's a request (has id) or notification (no id) + if container.contains(.id) { + self = .request(try Request(from: decoder)) + } else { + self = .notification(try Message(from: decoder)) + } + } + + func encode(to encoder: Encoder) throws { + switch self { + case .request(let request): + try request.encode(to: encoder) + case .notification(let notification): + try notification.encode(to: encoder) + } + } } diff --git a/Sources/MCP/Server/Tools.swift b/Sources/MCP/Server/Tools.swift index 7475557..478ba60 100644 --- a/Sources/MCP/Server/Tools.swift +++ b/Sources/MCP/Server/Tools.swift @@ -7,39 +7,178 @@ 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://modelcontextprotocol.io/specification/2025-11-25/server/tools/ public struct Tool: Hashable, Codable, Sendable { /// The tool name public let name: String + /// The human-readable name of the tool for display purposes. + public let title: String? /// The tool description - public let description: String + public let description: String? /// The tool input schema - public let inputSchema: Value? + public let inputSchema: Value + /// Optional set of sized icons that the client can display in a user interface + public var icons: [Icon]? + /// The tool output schema, defining expected output structure + public let outputSchema: Value? + /// Metadata fields for the tool (see spec for _meta usage) + public var _meta: Metadata? - public init(name: String, description: String, inputSchema: Value? = nil) { + /// Annotations that provide display-facing and operational information for a Tool. + /// + /// - Note: All properties in `ToolAnnotations` are **hints**. + /// They are not guaranteed to provide a faithful description of + /// tool behavior (including descriptive properties like `title`). + /// + /// Clients should never make tool use decisions based on `ToolAnnotations` + /// received from untrusted servers. + public struct Annotations: Hashable, Codable, Sendable, ExpressibleByNilLiteral { + /// A human-readable title for the tool + public var title: String? + + /// If true, the tool may perform destructive updates to its environment. + /// If false, the tool performs only additive updates. + /// (This property is meaningful only when `readOnlyHint == false`) + /// + /// When unspecified, the implicit default is `true`. + public var destructiveHint: Bool? + + /// If true, calling the tool repeatedly with the same arguments + /// will have no additional effect on its environment. + /// (This property is meaningful only when `readOnlyHint == false`) + /// + /// When unspecified, the implicit default is `false`. + public var idempotentHint: Bool? + + /// If true, this tool may interact with an "open world" of external + /// entities. If false, the tool's domain of interaction is closed. + /// For example, the world of a web search tool is open, whereas that + /// of a memory tool is not. + /// + /// When unspecified, the implicit default is `true`. + public var openWorldHint: Bool? + + /// If true, the tool does not modify its environment. + /// + /// When unspecified, the implicit default is `false`. + public var readOnlyHint: Bool? + + /// Returns true if all properties are nil + public var isEmpty: Bool { + title == nil && readOnlyHint == nil && destructiveHint == nil && idempotentHint == nil + && openWorldHint == nil + } + + public init( + title: String? = nil, + readOnlyHint: Bool? = nil, + destructiveHint: Bool? = nil, + idempotentHint: Bool? = nil, + openWorldHint: Bool? = nil + ) { + self.title = title + self.readOnlyHint = readOnlyHint + self.destructiveHint = destructiveHint + self.idempotentHint = idempotentHint + self.openWorldHint = openWorldHint + } + + /// Initialize an empty annotations object + public init(nilLiteral: ()) {} + } + + /// Annotations that provide display-facing and operational information + public var annotations: Annotations + + /// Initialize a tool with a name, description, input schema, annotations, and optional icons + public init( + name: String, + title: String? = nil, + description: String?, + inputSchema: Value, + annotations: Annotations = nil, + outputSchema: Value? = nil, + icons: [Icon]? = nil, + _meta: Metadata? = nil + ) { self.name = name + self.title = title self.description = description self.inputSchema = inputSchema + self.outputSchema = outputSchema + self.annotations = annotations + self._meta = _meta + self.icons = icons } /// Content types that can be returned by a tool public enum Content: Hashable, Codable, Sendable { /// Text content - case text(String) + case text(text: String, annotations: Resource.Annotations?, _meta: Metadata?) /// Image content - case image(data: String, mimeType: String, metadata: [String: String]?) - /// Embedded resource content - case resource(uri: String, mimeType: String, text: String?) + case image(data: String, mimeType: String, annotations: Resource.Annotations?, _meta: Metadata?) + /// Audio content + case audio(data: String, mimeType: String, annotations: Resource.Annotations?, _meta: Metadata?) + /// Embedded resource content (EmbeddedResource from spec) + case resource(resource: Resource.Content, annotations: Resource.Annotations? = nil, _meta: Metadata? = nil) + /// Resource link + case resourceLink( + uri: String, name: String, title: String? = nil, description: String? = nil, + mimeType: String? = nil, + annotations: Resource.Annotations? = nil + ) + + /// Deprecated compatibility factory for older call sites that used `.text("...")` and `.text("...", metadata: ...)`. + @available(*, deprecated, message: "Use .text(text:annotations:_meta:) instead.") + public static func text(_ text: String, metadata: Metadata? = nil) -> Self { + .text(text: text, annotations: nil, _meta: metadata) + } + + /// Deprecated compatibility factory for older call sites that used `.text(text: ..., metadata: ...)`. + @available(*, deprecated, message: "Use .text(text:annotations:_meta:) instead.") + public static func text(text: String, metadata: Metadata? = nil) -> Self { + .text(text: text, annotations: nil, _meta: metadata) + } + + /// Deprecated compatibility factory for older call sites that used `.image(..., metadata: ...)`. + @available(*, deprecated, message: "Use .image(data:mimeType:annotations:_meta:) instead.") + public static func image(_ data: String, _ mimeType: String, metadata: Metadata? = nil) -> Self { + .image(data: data, mimeType: mimeType, annotations: nil, _meta: metadata) + } + + /// Deprecated compatibility factory for older call sites that used `.image(data:mimeType:metadata:)`. + @available(*, deprecated, message: "Use .image(data:mimeType:annotations:_meta:) instead.") + public static func image(data: String, mimeType: String, metadata: Metadata? = nil) -> Self { + .image(data: data, mimeType: mimeType, annotations: nil, _meta: metadata) + } + + /// Deprecated compatibility factory for older call sites that used `.audio(..., metadata: ...)`. + @available(*, deprecated, message: "Use .audio(data:mimeType:annotations:_meta:) instead.") + public static func audio(_ data: String, _ mimeType: String, metadata: Metadata? = nil) -> Self { + .audio(data: data, mimeType: mimeType, annotations: nil, _meta: metadata) + } + + /// Deprecated compatibility factory for older call sites that used `.audio(data:mimeType:metadata:)`. + @available(*, deprecated, message: "Use .audio(data:mimeType:annotations:_meta:) instead.") + public static func audio(data: String, mimeType: String, metadata: Metadata? = nil) -> Self { + .audio(data: data, mimeType: mimeType, annotations: nil, _meta: metadata) + } private enum CodingKeys: String, CodingKey { case type case text case image case resource + case resource_link + case audio case uri + case name + case title + case description + case annotations case mimeType case data - case metadata + case _meta } public init(from decoder: Decoder) throws { @@ -49,18 +188,37 @@ 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(Resource.Annotations.self, forKey: .annotations) + let _meta = try container.decodeIfPresent(Metadata.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) - let metadata = try container.decodeIfPresent( - [String: String].self, forKey: .metadata) - self = .image(data: data, mimeType: mimeType, metadata: metadata) - case "resource": - let uri = try container.decode(String.self, forKey: .uri) + let annotations = try container.decodeIfPresent(Resource.Annotations.self, forKey: .annotations) + let _meta = try container.decodeIfPresent(Metadata.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) - let text = try container.decodeIfPresent(String.self, forKey: .text) - self = .resource(uri: uri, mimeType: mimeType, text: text) + let annotations = try container.decodeIfPresent(Resource.Annotations.self, forKey: .annotations) + let _meta = try container.decodeIfPresent(Metadata.self, forKey: ._meta) + self = .audio(data: data, mimeType: mimeType, annotations: annotations, _meta: _meta) + case "resource": + let resourceContent = try container.decode(Resource.Content.self, forKey: .resource) + let annotations = try container.decodeIfPresent(Resource.Annotations.self, forKey: .annotations) + let _meta = try container.decodeIfPresent(Metadata.self, forKey: ._meta) + self = .resource(resource: resourceContent, annotations: annotations, _meta: _meta) + case "resource_link": + let uri = try container.decode(String.self, forKey: .uri) + let name = try container.decode(String.self, forKey: .name) + let title = try container.decodeIfPresent(String.self, forKey: .title) + let description = try container.decodeIfPresent(String.self, forKey: .description) + let mimeType = try container.decodeIfPresent(String.self, forKey: .mimeType) + let annotations = try container.decodeIfPresent( + Resource.Annotations.self, forKey: .annotations) + self = .resourceLink( + uri: uri, name: name, title: title, description: description, + mimeType: mimeType, annotations: annotations) default: throw DecodingError.dataCorruptedError( forKey: .type, in: container, debugDescription: "Unknown tool content type") @@ -71,56 +229,90 @@ 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 .resource(let uri, let mimeType, let text): - try container.encode("resource", forKey: .type) - try container.encode(uri, forKey: .uri) + 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) - try container.encodeIfPresent(text, forKey: .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(resourceContent, forKey: .resource) + try container.encodeIfPresent(annotations, forKey: .annotations) + try container.encodeIfPresent(_meta, forKey: ._meta) + case .resourceLink( + let uri, let name, let title, let description, let mimeType, let annotations): + try container.encode("resource_link", forKey: .type) + try container.encode(uri, forKey: .uri) + try container.encode(name, forKey: .name) + try container.encodeIfPresent(title, forKey: .title) + try container.encodeIfPresent(description, forKey: .description) + try container.encodeIfPresent(mimeType, forKey: .mimeType) + try container.encodeIfPresent(annotations, forKey: .annotations) } } } private enum CodingKeys: String, CodingKey { case name + case title case description case inputSchema + case outputSchema + case annotations + case icons + case _meta } public init(from decoder: Decoder) throws { let container = try decoder.container(keyedBy: CodingKeys.self) name = try container.decode(String.self, forKey: .name) - description = try container.decode(String.self, forKey: .description) - inputSchema = try container.decodeIfPresent(Value.self, forKey: .inputSchema) + 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) + annotations = + try container.decodeIfPresent(Tool.Annotations.self, forKey: .annotations) ?? .init() + icons = try container.decodeIfPresent([Icon].self, forKey: .icons) + _meta = try container.decodeIfPresent(Metadata.self, forKey: ._meta) } public func encode(to encoder: Encoder) throws { var container = encoder.container(keyedBy: CodingKeys.self) try container.encode(name, forKey: .name) + try container.encodeIfPresent(title, forKey: .title) try container.encode(description, forKey: .description) - if let schema = inputSchema { - try container.encode(schema, forKey: .inputSchema) + try container.encode(inputSchema, forKey: .inputSchema) + try container.encodeIfPresent(outputSchema, forKey: .outputSchema) + if !annotations.isEmpty { + try container.encode(annotations, forKey: .annotations) } + try container.encodeIfPresent(icons, forKey: .icons) + try container.encodeIfPresent(_meta, forKey: ._meta) } } // MARK: - /// To discover available tools, clients send a `tools/list` request. -/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2024-11-05/server/tools/#listing-tools +/// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/server/tools/#listing-tools public enum ListTools: Method { public static let name = "tools/list" public struct Parameters: NotRequired, Hashable, Codable, Sendable { public let cursor: String? - + public init() { self.cursor = nil } @@ -133,42 +325,142 @@ public enum ListTools: Method { public struct Result: Hashable, Codable, Sendable { public let tools: [Tool] public let nextCursor: String? + public var _meta: Metadata? - public init(tools: [Tool], nextCursor: String? = nil) { + public init( + tools: [Tool], + nextCursor: String? = nil, + _meta: Metadata? = nil + ) { self.tools = tools self.nextCursor = nextCursor + self._meta = _meta + } + + private enum CodingKeys: String, CodingKey, CaseIterable { + case tools, nextCursor, _meta + } + + 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) + } + + 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(Metadata.self, forKey: ._meta) } } } /// To call a tool, clients send a `tools/call` request. -/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2024-11-05/server/tools/#calling-tools +/// - SeeAlso: https://modelcontextprotocol.io/specification/2025-11-25/server/tools/#calling-tools public enum CallTool: Method { public static let name = "tools/call" public struct Parameters: Hashable, Codable, Sendable { + /// Optional request metadata including progress token. + /// + /// If `progressToken` is specified, the caller is requesting out-of-band + /// progress notifications for this request. + public let _meta: Metadata? + + /// The name of the tool to call. public let name: String + + /// Arguments to use for the tool call. public let arguments: [String: Value]? - public init(name: String, arguments: [String: Value]? = nil) { + public init(name: String, arguments: [String: Value]? = nil, meta: Metadata? = nil) { + self._meta = meta self.name = name self.arguments = arguments } + + private enum CodingKeys: String, CodingKey { + case _meta + case name + case arguments + } + + public init(from decoder: Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + _meta = try container.decodeIfPresent(Metadata.self, forKey: ._meta) + name = try container.decode(String.self, forKey: .name) + arguments = try container.decodeIfPresent([String: Value].self, forKey: .arguments) + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + try container.encodeIfPresent(_meta, forKey: ._meta) + try container.encode(name, forKey: .name) + try container.encodeIfPresent(arguments, forKey: .arguments) + } } public struct Result: Hashable, Codable, Sendable { public let content: [Tool.Content] + public let structuredContent: Value? public let isError: Bool? + /// Optional metadata about this result + public var _meta: Metadata? - public init(content: [Tool.Content], isError: Bool? = nil) { + public init( + content: [Tool.Content] = [], + structuredContent: Value? = nil, + isError: Bool? = nil, + _meta: Metadata? = nil + ) { self.content = content + self.structuredContent = structuredContent self.isError = isError + self._meta = _meta + } + + public init( + content: [Tool.Content] = [], + structuredContent: Output, + isError: Bool? = nil, + _meta: Metadata? = nil + ) throws { + let encoded = try Value(structuredContent) + self.init( + content: content, + structuredContent: Optional.some(encoded), + isError: isError, + _meta: _meta + ) + } + + private enum CodingKeys: String, CodingKey, CaseIterable { + case content, structuredContent, isError, _meta + } + + 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) + } + + 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(Metadata.self, forKey: ._meta) } } } /// When the list of available tools changes, servers that declared the listChanged capability SHOULD send a notification: -/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2024-11-05/server/tools/#list-changed-notification +/// - SeeAlso: https://spec.modelcontextprotocol.io/specification/2025-06-18/server/tools/#list-changed-notification public struct ToolListChangedNotification: Notification { public static let name: String = "notifications/tools/list_changed" } diff --git a/Sources/MCPConformance/Client/main.swift b/Sources/MCPConformance/Client/main.swift new file mode 100644 index 0000000..60b4bae --- /dev/null +++ b/Sources/MCPConformance/Client/main.swift @@ -0,0 +1,737 @@ +/** + * Everything client - a single conformance test client that handles all scenarios. + * + * Usage: mcp-everything-client + * + * The scenario name is read from the MCP_CONFORMANCE_SCENARIO environment variable, + * which is set by the conformance test runner. + * + * This client routes to the appropriate behavior based on the scenario name, + * consolidating all the individual test clients into one. + */ + +import Foundation +import Logging +import MCP + +#if canImport(FoundationNetworking) + import FoundationNetworking +#endif + +// MARK: - Scenario Handlers + +typealias ScenarioHandler = ([String]) async throws -> Void + +// MARK: - Authorization Scenarios + +private func loadConformanceContext() -> [String: String] { + let env = ProcessInfo.processInfo.environment + + if let raw = env["MCP_CONFORMANCE_CONTEXT"], + let data = raw.data(using: .utf8), + let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] + { + var parsed: [String: String] = [:] + for (key, value) in json { + if let value = value as? String { + parsed[key] = value + } + } + return parsed + } + + var parsed: [String: String] = [:] + if let clientID = env["MCP_CONFORMANCE_CLIENT_ID"] { + parsed["client_id"] = clientID + } + if let clientSecret = env["MCP_CONFORMANCE_CLIENT_SECRET"] { + parsed["client_secret"] = clientSecret + } + return parsed +} + +private func percentEncodeFormValue(_ value: String) -> String { + let allowed = CharacterSet(charactersIn: "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-._~") + return value.addingPercentEncoding(withAllowedCharacters: allowed) ?? value +} + +private func formURLEncodedBody(_ parameters: [String: String]) -> Data { + let encoded = parameters + .sorted { $0.key < $1.key } + .map { key, value in + "\(percentEncodeFormValue(key))=\(percentEncodeFormValue(value))" + } + .joined(separator: "&") + return Data(encoded.utf8) +} + +private func clientAssertionAudience(from tokenEndpoint: URL) -> String { + guard var components = URLComponents(url: tokenEndpoint, resolvingAgainstBaseURL: false) else { + return tokenEndpoint.absoluteString + } + + components.query = nil + components.fragment = nil + + var path = components.path + if path.hasSuffix("/token") { + path = String(path.dropLast("/token".count)) + } else { + let parts = path.split(separator: "/") + if !parts.isEmpty { + let parent = parts.dropLast() + path = parent.isEmpty ? "" : "/" + parent.joined(separator: "/") + } + } + if path == "/" { + path = "" + } + components.path = path + + return components.url?.absoluteString ?? tokenEndpoint.absoluteString +} + +private func parsePrivateKeyJWTSigningAlgorithm( + _ signingAlgorithm: String +) throws -> OAuthConfiguration.PrivateKeyJWTSigningAlgorithm { + switch signingAlgorithm.uppercased() { + case OAuthConfiguration.PrivateKeyJWTSigningAlgorithm.ES256.rawValue: + return .ES256 + default: + throw ConformanceError.invalidArguments( + "Unsupported signing algorithm: \(signingAlgorithm)" + ) + } +} + +private func makeOAuthConfiguration( + for scenario: String, + context: [String: String] +) -> OAuthConfiguration { + let clientID = context["client_id"] ?? "test-client" + let clientSecret = context["client_secret"] ?? "test-secret" + + var configuration: OAuthConfiguration + switch scenario { + case "auth/pre-registration": + configuration = .init( + grantType: .authorizationCode, + authentication: .clientSecretBasic( + clientID: clientID, + clientSecret: clientSecret + ) + ) + + case "auth/token-endpoint-auth-basic": + configuration = .init( + grantType: .authorizationCode, + authentication: .clientSecretBasic( + clientID: clientID, + clientSecret: clientSecret + ) + ) + + case "auth/token-endpoint-auth-post": + configuration = .init( + grantType: .authorizationCode, + authentication: .clientSecretPost( + clientID: clientID, + clientSecret: clientSecret + ) + ) + + case "auth/client-credentials-basic": + configuration = .init( + authentication: .clientSecretBasic( + clientID: clientID, + clientSecret: clientSecret + ) + ) + + case "auth/client-credentials-jwt": + let privateKeyPEM = context["private_key_pem"] ?? "" + let signingAlgorithm = context["signing_algorithm"] ?? "ES256" + configuration = .init( + authentication: .privateKeyJWT( + clientID: clientID, + assertionFactory: { tokenEndpoint, clientID in + try OAuthConfiguration.makePrivateKeyJWTAssertion( + clientID: clientID, + tokenEndpoint: tokenEndpoint, + privateKeyPEM: privateKeyPEM, + signingAlgorithm: try parsePrivateKeyJWTSigningAlgorithm(signingAlgorithm), + audience: clientAssertionAudience(from: tokenEndpoint) + ) + } + ) + ) + + case "auth/basic-cimd": + configuration = .init( + grantType: .authorizationCode, + authentication: .none( + clientID: context["client_id"] + ?? "https://conformance-test.local/client-metadata.json") + ) + + case "auth/cross-app-access-complete-flow": + configuration = .init( + authentication: .clientSecretBasic( + clientID: clientID, + clientSecret: clientSecret + ), + accessTokenProvider: makeCrossAppAccessTokenProvider(context: context) + ) + + case let s where s.hasPrefix("auth/client-credentials"): + configuration = .init( + authentication: .none(clientID: clientID) + ) + + default: + configuration = .init( + grantType: .authorizationCode, + authentication: .none(clientID: clientID) + ) + } + + // Conformance harness currently uses loopback http AS endpoints. + configuration.allowLoopbackHTTPAuthorizationServerEndpoints = true + return configuration +} + +private struct ConformanceTokenResponse: Decodable { + let accessToken: String + + enum CodingKeys: String, CodingKey { + case accessToken = "access_token" + } +} + +private func requestOAuthToken( + url: URL, + parameters: [String: String], + authorizationHeader: String?, + session: URLSession +) async throws -> String { + var request = URLRequest(url: url) + request.httpMethod = "POST" + request.setValue("application/x-www-form-urlencoded", forHTTPHeaderField: "Content-Type") + request.setValue("application/json", forHTTPHeaderField: "Accept") + if let authorizationHeader { + request.setValue(authorizationHeader, forHTTPHeaderField: "Authorization") + } + request.httpBody = formURLEncodedBody(parameters) + + let (data, response) = try await session.data(for: request) + guard let httpResponse = response as? HTTPURLResponse else { + throw ConformanceError.invalidArguments("Token endpoint returned an invalid response") + } + guard (200..<300).contains(httpResponse.statusCode) else { + let body = String(data: data, encoding: .utf8) ?? "" + throw ConformanceError.invalidArguments( + "Token endpoint error (\(httpResponse.statusCode)): \(body)" + ) + } + + let token = try JSONDecoder().decode(ConformanceTokenResponse.self, from: data) + guard !token.accessToken.isEmpty else { + throw ConformanceError.invalidArguments("Token endpoint returned an empty access token") + } + return token.accessToken +} + +private func makeCrossAppAccessTokenProvider( + context: [String: String] +) -> OAuthConfiguration.AccessTokenProvider { + return { discovery, session in + guard let clientID = context["client_id"], + let clientSecret = context["client_secret"], + let idpClientID = context["idp_client_id"], + let idpIDToken = context["idp_id_token"], + let idpTokenEndpointValue = context["idp_token_endpoint"], + let idpTokenEndpoint = URL(string: idpTokenEndpointValue) + else { + throw ConformanceError.invalidArguments( + "Cross-app scenario requires client_id, client_secret, idp_client_id, idp_id_token, and idp_token_endpoint" + ) + } + + guard let authorizationServer = discovery.authorizationServer else { + throw ConformanceError.invalidArguments( + "SDK did not provide authorization server discovery context" + ) + } + guard let tokenEndpoint = discovery.tokenEndpoint else { + throw ConformanceError.invalidArguments( + "SDK did not provide token endpoint discovery context" + ) + } + let resource = discovery.resource.absoluteString + + let idJag = try await requestOAuthToken( + url: idpTokenEndpoint, + parameters: [ + "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", + "subject_token": idpIDToken, + "subject_token_type": "urn:ietf:params:oauth:token-type:id_token", + "requested_token_type": "urn:ietf:params:oauth:token-type:id-jag", + "audience": authorizationServer.absoluteString, + "resource": resource, + "client_id": idpClientID, + ], + authorizationHeader: nil, + session: session + ) + + var accessTokenParameters: [String: String] = [ + "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer", + "assertion": idJag, + "resource": resource, + ] + if let requestedScopes = discovery.requestedScopes, !requestedScopes.isEmpty { + accessTokenParameters["scope"] = requestedScopes.sorted().joined(separator: " ") + } + + let basicCredentials = Data("\(clientID):\(clientSecret)".utf8).base64EncodedString() + let accessToken = try await requestOAuthToken( + url: tokenEndpoint, + parameters: accessTokenParameters, + authorizationHeader: "Basic \(basicCredentials)", + session: session + ) + + return accessToken + } +} + +func runAuthorizationScenario(scenario: String, args: [String]) async throws { + var logger = Logger( + label: "mcp.conformance.client.auth", + factory: { StreamLogHandler.standardError(label: $0) } + ) + logger.logLevel = .debug + + logger.debug("Starting auth scenario", metadata: ["scenario": "\(scenario)"]) + + guard let serverURLString = args.last, + let serverURL = URL(string: serverURLString) + else { + throw ConformanceError.invalidArguments("Valid server URL is required") + } + + let context = loadConformanceContext() + let oauthConfig = makeOAuthConfiguration(for: scenario, context: context) + + let transport = HTTPClientTransport( + endpoint: serverURL, + streaming: true, + authorizer: OAuthAuthorizer(configuration: oauthConfig), + logger: logger + ) + + let client = Client(name: "test-client", version: "1.0.0") + + // Scenarios that expect the connection to fail with a specific error. + if scenario == "auth/resource-mismatch" { + do { + _ = try await client.connect(transport: transport) + throw ConformanceError.invalidArguments( + "Expected authorization to fail with resource mismatch, but connection succeeded" + ) + } catch let error as MCPError { + guard case .internalError(let detail) = error, + detail?.contains("resource mismatch") == true + else { + throw ConformanceError.invalidArguments( + "Connection failed, but not due to resource mismatch: \(error.localizedDescription)" + ) + } + logger.debug("Client correctly rejected mismatched PRM resource") + } + return + } + + _ = try await client.connect(transport: transport) + + // Exercise both initialization and regular request paths. + let (tools, _) = try await client.listTools() + logger.debug("Auth scenario listed tools", metadata: ["count": "\(tools.count)"]) + + // Trigger an additional request for scenarios that involve runtime scope behavior. + if scenario.contains("scope"), let firstTool = tools.first { + _ = try? await client.callTool(name: firstTool.name, arguments: [:]) + } + + await client.disconnect() + logger.debug("Auth scenario completed", metadata: ["scenario": "\(scenario)"]) +} + +// MARK: - Basic Scenarios + +/// Basic client that connects, initializes, and lists tools +func runInitializeScenario(_ args: [String]) async throws { + var logger = Logger( + label: "mcp.conformance.client.initialize", + factory: { StreamLogHandler.standardError(label: $0) } + ) + logger.logLevel = .debug + + logger.debug("Starting initialize scenario") + + // Get server URL from args + guard let serverURLString = args.last, + let serverURL = URL(string: serverURLString) else { + throw ConformanceError.invalidArguments("Valid server URL is required") + } + + // Create HTTP transport + let transport = HTTPClientTransport( + endpoint: serverURL, + logger: logger + ) + + // Create client + let client = Client(name: "test-client", version: "1.0.0") + + // Connect + let initResult = try await client.connect(transport: transport) + logger.debug("Successfully connected to MCP server", metadata: [ + "serverName": "\(initResult.serverInfo.name)", + "serverVersion": "\(initResult.serverInfo.version)" + ]) + + // List tools + let (tools, _) = try await client.listTools() + logger.debug("Successfully listed tools", metadata: [ + "toolCount": "\(tools.count)" + ]) + + // Disconnect + await client.disconnect() + + logger.debug("Initialize scenario completed successfully") +} + +/// Client that calls the add_numbers tool +func runToolsCallScenario(_ args: [String]) async throws { + var logger = Logger( + label: "mcp.conformance.client.tools_call", + factory: { StreamLogHandler.standardError(label: $0) } + ) + logger.logLevel = .debug + + logger.debug("Starting tools_call scenario") + + // Get server URL from args + guard let serverURLString = args.last, + let serverURL = URL(string: serverURLString) else { + throw ConformanceError.invalidArguments("Valid server URL is required") + } + + // Create HTTP transport + let transport = HTTPClientTransport( + endpoint: serverURL, + logger: logger + ) + + // Create client + let client = Client(name: "test-client", version: "1.0.0") + + // Connect + try await client.connect(transport: transport) + logger.debug("Successfully connected to MCP server") + + // List tools + let (tools, _) = try await client.listTools() + logger.debug("Successfully listed tools", metadata: [ + "toolCount": "\(tools.count)" + ]) + + // Call the add_numbers tool + if tools.contains(where: { $0.name == "add_numbers" }) { + let result = try await client.callTool( + name: "add_numbers", + arguments: ["a": 5, "b": 3] + ) + logger.debug("Tool call result", metadata: [ + "isError": "\(result.isError ?? false)", + "contentCount": "\(result.content.count)" + ]) + } else { + logger.warning("add_numbers tool not found") + } + + // Disconnect + await client.disconnect() + + logger.debug("Tools call scenario completed successfully") +} + +// MARK: - SSE Scenarios + +/// Handler for SSE-related scenarios (retry, reconnection, etc.) +func runSSEScenario(_ args: [String]) async throws { + var logger = Logger( + label: "mcp.conformance.client.sse", + factory: { StreamLogHandler.standardError(label: $0) } + ) + logger.logLevel = .debug + + logger.debug("Starting SSE scenario") + + // Get server URL from args + guard let serverURLString = args.last, + let serverURL = URL(string: serverURLString) else { + throw ConformanceError.invalidArguments("Valid server URL is required") + } + + // Create HTTP transport with streaming enabled + let transport = HTTPClientTransport( + endpoint: serverURL, + streaming: true, + logger: logger + ) + + // Create client + let client = Client(name: "test-client", version: "1.0.0") + + // Connect - this will start the SSE stream in the background + let initResult = try await client.connect(transport: transport) + logger.debug("Successfully connected to MCP server", metadata: [ + "serverName": "\(initResult.serverInfo.name)", + "serverVersion": "\(initResult.serverInfo.version)" + ]) + + // Give the GET SSE stream time to establish + try await Task.sleep(for: .milliseconds(500)) + + // Call the test_reconnection tool to trigger SSE stream closure and retry test. + // The server will close the POST SSE stream without the response, + // then deliver it on the GET SSE stream after we reconnect. + logger.debug("Calling test_reconnection tool...") + let result = try await client.callTool(name: "test_reconnection", arguments: [:]) + logger.debug("Tool call result received", metadata: [ + "isError": "\(result.isError ?? false)", + "contentCount": "\(result.content.count)" + ]) + + // Keep the connection open briefly for the test to collect timing data + try await Task.sleep(for: .seconds(2)) + + // Disconnect + await client.disconnect() + + logger.debug("SSE scenario completed") +} + +/// Client that handles elicitation-sep1034-client-defaults scenario +/// Tests that client properly applies default values for omitted fields +func runElicitationSEP1034ClientDefaults(_ args: [String]) async throws { + var logger = Logger( + label: "mcp.conformance.client.elicitation_client_defaults", + factory: { StreamLogHandler.standardError(label: $0) } + ) + logger.logLevel = .debug + + logger.debug("Starting elicitation-sep1034-client-defaults scenario") + + // Get server URL from args + guard let serverURLString = args.last, + let serverURL = URL(string: serverURLString) else { + throw ConformanceError.invalidArguments("Valid server URL is required") + } + + // Create HTTP transport with streaming enabled for bidirectional communication + let transport = HTTPClientTransport( + endpoint: serverURL, + streaming: true, + logger: logger + ) + + // Create client with elicitation capabilities + let client = Client( + name: "test-client", + version: "1.0.0", + capabilities: Client.Capabilities( + elicitation: Client.Capabilities.Elicitation(form: .init(), url: .init()) + ) + ) + + // Set up elicitation handler that accepts defaults BEFORE connecting + await client.withElicitationHandler { [logger] params in + let message: String + switch params { + case .form(let formParams): + message = formParams.message + case .url(let urlParams): + message = urlParams.message + } + + logger.debug("Elicitation handler invoked", metadata: [ + "message": "\(message)" + ]) + + // Accept with default values applied + // The schema has optional fields with defaults: + // name: "John Doe", age: 30, score: 95.5, status: "active", verified: true + return CreateElicitation.Result( + action: .accept, + content: [ + "name": "John Doe", + "age": 30, + "score": 95.5, + "status": "active", + "verified": true + ] + ) + } + + // Connect + try await client.connect(transport: transport) + logger.debug("Successfully connected to MCP server") + + // List tools + let (tools, _) = try await client.listTools() + logger.debug("Successfully listed tools", metadata: [ + "toolCount": "\(tools.count)" + ]) + + // Call the test_client_elicitation_defaults tool + if tools.contains(where: { $0.name == "test_client_elicitation_defaults" }) { + let result = try await client.callTool( + name: "test_client_elicitation_defaults", + arguments: [:] + ) + logger.debug("Tool call result", metadata: [ + "isError": "\(result.isError ?? false)", + "contentCount": "\(result.content.count)" + ]) + } else { + logger.warning("test_client_elicitation_defaults tool not found") + } + + // Disconnect + await client.disconnect() + + logger.debug("Elicitation client defaults scenario completed successfully") +} + +// MARK: - Default Handler for Unimplemented Scenarios + +/// Default handler that performs basic connection test for unimplemented scenarios +func runDefaultScenario(_ args: [String]) async throws { + var logger = Logger( + label: "mcp.conformance.client.default", + factory: { StreamLogHandler.standardError(label: $0) } + ) + logger.logLevel = .debug + + logger.debug("Running default scenario handler") + + // Get server URL from args + guard let serverURLString = args.last, + let serverURL = URL(string: serverURLString) else { + throw ConformanceError.invalidArguments("Valid server URL is required") + } + + // Create HTTP transport + let transport = HTTPClientTransport( + endpoint: serverURL, + logger: logger + ) + + // Create client + let client = Client(name: "test-client", version: "1.0.0") + + // Connect + let initResult = try await client.connect(transport: transport) + logger.debug("Successfully connected to MCP server", metadata: [ + "serverName": "\(initResult.serverInfo.name)", + "serverVersion": "\(initResult.serverInfo.version)" + ]) + + // Disconnect + await client.disconnect() + + logger.debug("Default scenario completed successfully") +} + +// MARK: - Scenario Registry + +nonisolated(unsafe) let scenarioHandlers: [String: ScenarioHandler] = [ + "initialize": runInitializeScenario, + "tools_call": runToolsCallScenario, + "sse-retry": runSSEScenario, + "elicitation-sep1034-client-defaults": runElicitationSEP1034ClientDefaults, +] + +// MARK: - Error Types + +enum ConformanceError: Error, CustomStringConvertible { + case missingScenario + case invalidArguments(String) + + var description: String { + switch self { + case .missingScenario: + return "MCP_CONFORMANCE_SCENARIO environment variable not set" + case .invalidArguments(let message): + return "Invalid arguments: \(message)" + } + } +} + +struct ConformanceClient { + static func run() async { + do { + // Get scenario from environment + guard let scenario = ProcessInfo.processInfo.environment["MCP_CONFORMANCE_SCENARIO"] else { + var stderr = StandardError() + print("Error: MCP_CONFORMANCE_SCENARIO environment variable not set", to: &stderr) + Foundation.exit(1) + } + + // Get server URL from arguments (last argument) + let args = Array(CommandLine.arguments.dropFirst()) + guard !args.isEmpty else { + var stderr = StandardError() + print("Usage: mcp-everything-client ", to: &stderr) + print("Error: Server URL is required", to: &stderr) + Foundation.exit(1) + } + + // Get handler for scenario + let handler: ScenarioHandler + if let explicitHandler = scenarioHandlers[scenario] { + handler = explicitHandler + } else if scenario.hasPrefix("auth/") { + handler = { args in + try await runAuthorizationScenario(scenario: scenario, args: args) + } + } else { + handler = runDefaultScenario + var stderr = StandardError() + print("⚠️ Scenario '\(scenario)' not fully implemented - using default handler", to: &stderr) + } + + // Run the scenario + try await handler(args) + Foundation.exit(0) + } catch { + var stderr = StandardError() + print("Error: \(error)", to: &stderr) + Foundation.exit(1) + } + } +} + +// MARK: - Helpers + +struct StandardError: TextOutputStream { + mutating func write(_ string: String) { + FileHandle.standardError.write(Data(string.utf8)) + } +} + +await ConformanceClient.run() diff --git a/Sources/MCPConformance/Server/HTTPApp.swift b/Sources/MCPConformance/Server/HTTPApp.swift new file mode 100644 index 0000000..6edba56 --- /dev/null +++ b/Sources/MCPConformance/Server/HTTPApp.swift @@ -0,0 +1,432 @@ +import Foundation +import Logging +import MCP +@preconcurrency import NIOCore +@preconcurrency import NIOPosix +@preconcurrency import NIOHTTP1 + +#if canImport(FoundationNetworking) + import FoundationNetworking +#endif + +actor HTTPApp { + /// Configuration for the HTTP application. + struct Configuration: Sendable { + /// The host address to bind to. + var host: String + + /// The port to bind to. + var port: Int + + /// The MCP endpoint path. + var endpoint: String + + /// Session timeout in seconds. + var sessionTimeout: TimeInterval + + /// SSE retry interval in milliseconds for priming events. + var retryInterval: Int? + + init( + host: String = "127.0.0.1", + port: Int = 3000, + endpoint: String = "/mcp", + sessionTimeout: TimeInterval = 3600, + retryInterval: Int? = nil + ) { + self.host = host + self.port = port + self.endpoint = endpoint + self.sessionTimeout = sessionTimeout + self.retryInterval = retryInterval + } + } + + /// Factory function to create MCP Server instances for each session. + typealias ServerFactory = @Sendable (String, StatefulHTTPServerTransport) async throws -> Server + + private let configuration: Configuration + private let serverFactory: ServerFactory + private let validationPipeline: (any HTTPRequestValidationPipeline)? + private var channel: Channel? + private var sessions: [String: SessionContext] = [:] + + nonisolated let logger: Logger + + struct SessionContext { + let server: Server + let transport: StatefulHTTPServerTransport + let createdAt: Date + var lastAccessedAt: Date + } + + // MARK: - Init + + /// Creates a new HTTP application. + /// + /// - Parameters: + /// - configuration: Application configuration. + /// - validationPipeline: Custom validation pipeline passed to each transport. + /// If `nil`, transports use their sensible defaults. + /// - serverFactory: Factory function to create Server instances for each session. + /// - logger: Optional logger instance. + init( + configuration: Configuration = Configuration(), + validationPipeline: (any HTTPRequestValidationPipeline)? = nil, + serverFactory: @escaping ServerFactory, + logger: Logger? = nil + ) { + self.configuration = configuration + self.serverFactory = serverFactory + self.validationPipeline = validationPipeline + self.logger = logger ?? Logger( + label: "mcp.http.app", + factory: { _ in SwiftLogNoOpLogHandler() } + ) + } + + /// Convenience initializer with individual parameters. + init( + host: String = "127.0.0.1", + port: Int = 3000, + endpoint: String = "/mcp", + serverFactory: @escaping ServerFactory, + logger: Logger? = nil + ) { + self.init( + configuration: Configuration(host: host, port: port, endpoint: endpoint), + serverFactory: serverFactory, + logger: logger + ) + } + + // MARK: - Lifecycle + + /// Starts the HTTP application. + /// + /// This starts the NIO HTTP server and begins accepting connections. + /// The call blocks until the server is shut down via ``stop()``. + func start() async throws { + let group = MultiThreadedEventLoopGroup(numberOfThreads: System.coreCount) + + let bootstrap = ServerBootstrap(group: group) + .serverChannelOption(ChannelOptions.backlog, value: 256) + .serverChannelOption(ChannelOptions.socketOption(.so_reuseaddr), value: 1) + .childChannelInitializer { channel in + channel.pipeline.configureHTTPServerPipeline().flatMap { + channel.pipeline.addHandler(HTTPHandler(app: self)) + } + } + .childChannelOption(ChannelOptions.socketOption(.so_reuseaddr), value: 1) + .childChannelOption(ChannelOptions.maxMessagesPerRead, value: 1) + + logger.info( + "Starting MCP HTTP application", + metadata: [ + "host": "\(configuration.host)", + "port": "\(configuration.port)", + "endpoint": "\(configuration.endpoint)", + ] + ) + + let channel = try await bootstrap.bind(host: configuration.host, port: configuration.port).get() + self.channel = channel + + Task { await sessionCleanupLoop() } + + try await channel.closeFuture.get() + } + + /// Stops the HTTP application gracefully, closing all sessions. + func stop() async { + await closeAllSessions() + try? await channel?.close() + channel = nil + logger.info("MCP HTTP application stopped") + } + + // MARK: - Request Routing + + var endpoint: String { configuration.endpoint } + + /// Routes an incoming HTTP request to the appropriate session transport. + /// + /// - Requests with a valid `Mcp-Session-Id` are forwarded to the matching transport. + /// - POST requests with an `initialize` body create a new session. + /// - All other requests without a session return an error. + func handleHTTPRequest(_ request: HTTPRequest) async -> HTTPResponse { + let sessionID = request.header(HTTPHeaderName.sessionID) + + // Route to existing session + if let sessionID, var session = sessions[sessionID] { + session.lastAccessedAt = Date() + sessions[sessionID] = session + + let response = await session.transport.handleRequest(request) + + // Clean up on successful DELETE + if request.method.uppercased() == "DELETE" && response.statusCode == 200 { + sessions.removeValue(forKey: sessionID) + } + + return response + } + + // No session — check for initialize request + if request.method.uppercased() == "POST", + let body = request.body, + let kind = JSONRPCMessageKind(data: body), + kind.isInitializeRequest + { + return await createSessionAndHandle(request) + } + + // No session and not initialize + if sessionID != nil { + return .error(statusCode: 404, .invalidRequest("Not Found: Session not found or expired")) + } + return .error( + statusCode: 400, + .invalidRequest("Bad Request: Missing \(HTTPHeaderName.sessionID) header") + ) + } + + // MARK: - Session Management + + private struct FixedSessionIDGenerator: SessionIDGenerator { + let sessionID: String + func generateSessionID() -> String { sessionID } + } + + private func createSessionAndHandle(_ request: HTTPRequest) async -> HTTPResponse { + let sessionID = UUID().uuidString + + let transport = StatefulHTTPServerTransport( + sessionIDGenerator: FixedSessionIDGenerator(sessionID: sessionID), + validationPipeline: validationPipeline, + retryInterval: configuration.retryInterval, + logger: logger + ) + + do { + let server = try await serverFactory(sessionID, transport) + try await server.start(transport: transport) + + sessions[sessionID] = SessionContext( + server: server, + transport: transport, + createdAt: Date(), + lastAccessedAt: Date() + ) + + let response = await transport.handleRequest(request) + + // If transport returned an error, clean up + if case .error = response { + sessions.removeValue(forKey: sessionID) + await transport.disconnect() + } + + return response + } catch { + await transport.disconnect() + return .error( + statusCode: 500, + .internalError("Failed to create session: \(error.localizedDescription)") + ) + } + } + + private func closeSession(_ sessionID: String) async { + guard let session = sessions.removeValue(forKey: sessionID) else { return } + await session.transport.disconnect() + logger.info("Closed session", metadata: ["sessionID": "\(sessionID)"]) + } + + private func closeAllSessions() async { + for sessionID in sessions.keys { + await closeSession(sessionID) + } + } + + private func sessionCleanupLoop() async { + while true { + try? await Task.sleep(for: .seconds(60)) + + let now = Date() + let expired = sessions.filter { _, context in + now.timeIntervalSince(context.lastAccessedAt) > configuration.sessionTimeout + } + + for (sessionID, _) in expired { + logger.info("Session expired", metadata: ["sessionID": "\(sessionID)"]) + await closeSession(sessionID) + } + } + } +} + +// MARK: - NIO HTTP Handler + +/// Thin NIO adapter that converts between NIO HTTP types and the framework-agnostic +/// `HTTPRequest`/`HTTPResponse` types, delegating all logic to the `HTTPApp`. +private final class HTTPHandler: ChannelInboundHandler, @unchecked Sendable { + typealias InboundIn = HTTPServerRequestPart + typealias OutboundOut = HTTPServerResponsePart + + private let app: HTTPApp + + private struct RequestState { + var head: HTTPRequestHead + var bodyBuffer: ByteBuffer + } + + private var requestState: RequestState? + + init(app: HTTPApp) { + self.app = app + } + + func channelRead(context: ChannelHandlerContext, data: NIOAny) { + let part = unwrapInboundIn(data) + + switch part { + case .head(let head): + requestState = RequestState( + head: head, + bodyBuffer: context.channel.allocator.buffer(capacity: 0) + ) + case .body(var buffer): + requestState?.bodyBuffer.writeBuffer(&buffer) + case .end: + guard let state = requestState else { return } + requestState = nil + + nonisolated(unsafe) let ctx = context + Task { @MainActor in + await self.handleRequest(state: state, context: ctx) + } + } + } + + // MARK: - Request Processing + + private func handleRequest(state: RequestState, context: ChannelHandlerContext) async { + let head = state.head + let path = head.uri.split(separator: "?").first.map(String.init) ?? head.uri + let endpoint = await app.endpoint + + guard path == endpoint else { + await writeResponse( + .error(statusCode: 404, .invalidRequest("Not Found")), + version: head.version, + context: context + ) + return + } + + let httpRequest = makeHTTPRequest(from: state) + let response = await app.handleHTTPRequest(httpRequest) + await writeResponse(response, version: head.version, context: context) + } + + // MARK: - NIO ↔ HTTPRequest/HTTPResponse Conversion + + private func makeHTTPRequest(from state: RequestState) -> HTTPRequest { + // Combine multiple header values per RFC 7230 + var headers: [String: String] = [:] + for (name, value) in state.head.headers { + if let existing = headers[name] { + headers[name] = existing + ", " + value + } else { + headers[name] = value + } + } + + let body: Data? + if state.bodyBuffer.readableBytes > 0, + let bytes = state.bodyBuffer.getBytes(at: 0, length: state.bodyBuffer.readableBytes) + { + body = Data(bytes) + } else { + body = nil + } + + let path = String(state.head.uri.split(separator: "?").first ?? Substring(state.head.uri)) + + return HTTPRequest( + method: state.head.method.rawValue, + headers: headers, + body: body, + path: path + ) + } + + private func writeResponse( + _ response: HTTPResponse, + version: HTTPVersion, + context: ChannelHandlerContext + ) async { + nonisolated(unsafe) let ctx = context + let eventLoop = ctx.eventLoop + + // Write response head + let statusCode = response.statusCode + let headers = response.headers + + switch response { + case .stream(let stream, _): + eventLoop.execute { + var head = HTTPResponseHead( + version: version, + status: HTTPResponseStatus(statusCode: statusCode) + ) + for (name, value) in headers { + head.headers.add(name: name, value: value) + } + ctx.write(self.wrapOutboundOut(.head(head)), promise: nil) + ctx.flush() + } + + // Await the SSE stream directly — no Task needed since we're already in one + do { + for try await chunk in stream { + eventLoop.execute { + var buffer = ctx.channel.allocator.buffer(capacity: chunk.count) + buffer.writeBytes(chunk) + ctx.writeAndFlush( + self.wrapOutboundOut(.body(.byteBuffer(buffer))), promise: nil) + } + } + } catch { + // Stream ended with error — close connection + } + + eventLoop.execute { + ctx.writeAndFlush(self.wrapOutboundOut(.end(nil)), promise: nil) + } + + default: + let bodyData = response.bodyData + eventLoop.execute { + var head = HTTPResponseHead( + version: version, + status: HTTPResponseStatus(statusCode: statusCode) + ) + for (name, value) in headers { + head.headers.add(name: name, value: value) + } + + ctx.write(self.wrapOutboundOut(.head(head)), promise: nil) + + if let body = bodyData { + var buffer = ctx.channel.allocator.buffer(capacity: body.count) + buffer.writeBytes(body) + ctx.write(self.wrapOutboundOut(.body(.byteBuffer(buffer))), promise: nil) + } + + ctx.writeAndFlush(self.wrapOutboundOut(.end(nil)), promise: nil) + } + } + } +} diff --git a/Sources/MCPConformance/Server/main.swift b/Sources/MCPConformance/Server/main.swift new file mode 100644 index 0000000..6a7ccea --- /dev/null +++ b/Sources/MCPConformance/Server/main.swift @@ -0,0 +1,519 @@ +/** + * MCP HTTP Server Wrapper + * + * HTTP server that wraps the MCP conformance server for testing with the + * official conformance framework. + * + * Usage: mcp-http-server [--port PORT] + */ + +import Foundation +import Logging +import MCP + +#if canImport(FoundationNetworking) + import FoundationNetworking +#endif + +// MARK: - Test Data + +private let testImageBase64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg==" +private let testAudioBase64 = "UklGRiYAAABXQVZFZm10IBAAAAABAAEAQB8AAAB9AAACABAAZGF0YQIAAAA=" + +// MARK: - Server State + +actor ServerState { + var resourceSubscriptions: Set = [] + var watchedResourceContent = "Watched resource content" + + func subscribe(to uri: String) { + resourceSubscriptions.insert(uri) + } + + func unsubscribe(from uri: String) { + resourceSubscriptions.remove(uri) + } + + func isSubscribed(to uri: String) -> Bool { + resourceSubscriptions.contains(uri) + } + + func updateWatchedResource(_ newContent: String) { + watchedResourceContent = newContent + } +} + +// MARK: - Server Setup + +func createConformanceServer(state: ServerState, transport: StatefulHTTPServerTransport) async -> Server { + let server = Server( + name: "mcp-conformance-test-server", + version: "1.0.0", + capabilities: Server.Capabilities( + completions: .init(), + logging: .init(), + prompts: .init(listChanged: true), + resources: .init(subscribe: true, listChanged: true), + tools: .init(listChanged: true) + ) + ) + + // Tools + await server.withMethodHandler(ListTools.self) { _ in + .init(tools: [ + Tool(name: "test_simple_text", description: "Tests simple text content response", inputSchema: .object(["type": "object", "properties": [:]])), + Tool(name: "test_image_content", description: "Tests image content response", inputSchema: .object(["type": "object", "properties": [:]])), + Tool(name: "test_audio_content", description: "Tests audio content response", inputSchema: .object(["type": "object", "properties": [:]])), + Tool(name: "test_embedded_resource", description: "Tests embedded resource content response", inputSchema: .object(["type": "object", "properties": [:]])), + Tool(name: "test_multiple_content_types", description: "Tests response with multiple content types", inputSchema: .object(["type": "object", "properties": [:]])), + Tool(name: "test_error_handling", description: "Tests error response handling", inputSchema: .object(["type": "object", "properties": [:]])), + Tool(name: "test_logging", description: "Tests logging capabilities", inputSchema: .object(["type": "object", "properties": [:]])), + Tool(name: "test_progress", description: "Tests progress notifications", inputSchema: .object(["type": "object", "properties": ["duration_ms": ["type": "number", "description": "Duration in milliseconds to report progress"]]])), + Tool(name: "add_numbers", description: "Adds two numbers together", inputSchema: .object(["type": "object", "properties": ["a": ["type": "number", "description": "First number"], "b": ["type": "number", "description": "Second number"]]])), + Tool(name: "test_tool_with_progress", description: "Tool reports progress notifications", inputSchema: .object(["type": "object", "properties": [:]])), + Tool(name: "test_tool_with_logging", description: "Tool sends log messages during execution", inputSchema: .object(["type": "object", "properties": [:]])), + Tool(name: "test_reconnection", description: "Tests SSE reconnection and resumption with Last-Event-ID", inputSchema: .object(["type": "object", "properties": [:]])), + Tool(name: "test_sampling", description: "Tests LLM sampling capabilities", inputSchema: .object(["type": "object", "properties": ["prompt": ["type": "string", "description": "Text to send to the LLM"]], "required": ["prompt"]])), + Tool(name: "test_elicitation", description: "Tests user input elicitation", inputSchema: .object(["type": "object", "properties": ["message": ["type": "string", "description": "Text displayed to user"]], "required": ["message"]])), + Tool(name: "test_elicitation_sep1034_defaults", description: "Tests elicitation with default values (SEP-1034)", inputSchema: .object(["type": "object", "properties": [:]])), + Tool(name: "test_elicitation_sep1330_enums", description: "Tests elicitation with enum variants (SEP-1330)", inputSchema: .object(["type": "object", "properties": [:]])), + Tool(name: "test_client_elicitation_defaults", description: "Tests that client applies defaults for omitted elicitation fields", inputSchema: .object(["type": "object", "properties": [:]])), + Tool(name: "json_schema_2020_12_tool", description: "Tool with JSON Schema 2020-12 features", inputSchema: .object([ + "$schema": .string("https://json-schema.org/draft/2020-12/schema"), + "type": .string("object"), + "$defs": .object([ + "address": .object([ + "type": .string("object"), + "properties": .object([ + "street": .object(["type": .string("string")]), + "city": .object(["type": .string("string")]) + ]) + ]) + ]), + "properties": .object([ + "name": .object(["type": .string("string")]), + "address": .object(["$ref": .string("#/$defs/address")]) + ]), + "additionalProperties": .bool(false) + ])) + ]) + } + + await server.withMethodHandler(CallTool.self) { [weak server, transport] params in + switch params.name { + case "test_simple_text": + return .init(content: [.text(text: "This is a simple text response for testing.", annotations: nil, _meta: nil)], isError: false) + case "test_image_content": + return .init(content: [.image(data: testImageBase64, mimeType: "image/png", annotations: nil, _meta: nil)], isError: false) + case "test_audio_content": + return .init(content: [.audio(data: testAudioBase64, mimeType: "audio/wav", annotations: nil, _meta: nil)], isError: false) + case "test_embedded_resource": + return .init(content: [.resource(resource: .text("This is an embedded resource content.", uri: "test://embedded-resource", mimeType: "text/plain"))], isError: false) + case "test_multiple_content_types": + return .init(content: [ + .text(text: "Multiple content types test:", annotations: nil, _meta: nil), + .image(data: testImageBase64, mimeType: "image/png", annotations: nil, _meta: nil), + .resource(resource: .text("{\"test\":\"data\",\"value\":123}", uri: "test://mixed-content-resource", mimeType: "application/json"))], isError: false) + case "test_error_handling": + return .init(content: [.text(text: "An error occurred during tool execution", annotations: nil, _meta: nil)], isError: true) + case "test_logging": + return .init(content: [.text(text: "Logging test completed", annotations: nil, _meta: nil)], isError: false) + case "test_progress": + let duration = params.arguments?["duration_ms"]?.intValue ?? 1000 + try? await Task.sleep(for: .milliseconds(duration)) + return .init(content: [.text(text: "Progress test completed", annotations: nil, _meta: nil)], isError: false) + case "add_numbers": + guard let a = params.arguments?["a"]?.intValue, let b = params.arguments?["b"]?.intValue else { + return .init(content: [.text(text: "Invalid arguments: expected numbers a and b", annotations: nil, _meta: nil)], isError: true) + } + return .init(content: [.text(text: "\(a + b)", annotations: nil, _meta: nil)], isError: false) + case "test_tool_with_progress": + if let token = params._meta?.progressToken { + let notification1 = ProgressNotification.message( + .init(progressToken: token, progress: 0, total: 100) + ) + try await server?.notify(notification1) + try await Task.sleep(for: .microseconds(50)) + + let notification2 = ProgressNotification.message( + .init(progressToken: token, progress: 50, total: 100) + ) + try await server?.notify(notification2) + try await Task.sleep(for: .microseconds(50)) + + let notification3 = ProgressNotification.message( + .init(progressToken: token, progress: 100, total: 100) + ) + try await server?.notify(notification3) + } + + return .init(content: [.text(text: "This is a simple text response for testing.", annotations: nil, _meta: nil)], isError: false) + case "json_schema_2020_12_tool": + return .init(content: [.text(text: "This is a simple text response for testing.", annotations: nil, _meta: nil)], isError: false) + case "test_tool_with_logging": + // Send first log message + let log1 = LogMessageNotification.message( + .init(level: .info, data: .string("Tool execution started")) + ) + try await server?.notify(log1) + + // Wait 50ms + try await Task.sleep(for: .milliseconds(50)) + + // Send second log message + let log2 = LogMessageNotification.message( + .init(level: .info, data: .string("Tool processing data")) + ) + try await server?.notify(log2) + + // Wait another 50ms + try await Task.sleep(for: .milliseconds(50)) + + // Send third log message + let log3 = LogMessageNotification.message( + .init(level: .info, data: .string("Tool execution completed")) + ) + try await server?.notify(log3) + + return .init(content: [.text(text: "Logging test completed", annotations: nil, _meta: nil)], isError: false) + case "test_reconnection": + // SEP-1699: Close the SSE stream mid-call to trigger client reconnection. + // The client should reconnect via GET with Last-Event-ID and receive the + // response on the new stream. + if let requestID = Server.currentHandlerContext?.id { + await transport.closeSSEStream(forRequestID: requestID.description) + } + // Wait briefly for the client to reconnect before sending the response. + try await Task.sleep(for: .milliseconds(100)) + return .init(content: [.text(text: "Reconnection test completed successfully. If you received this, the client properly reconnected after stream closure.", annotations: nil, _meta: nil)], isError: false) + case "test_sampling": + // Test LLM sampling - request sampling/createMessage from client + guard let prompt = params.arguments?["prompt"]?.stringValue else { + return .init(content: [.text(text: "Missing required argument: prompt", annotations: nil, _meta: nil)], isError: true) + } + + let samplingResult = try await server?.requestSampling( + messages: [.user(.text(prompt))], + maxTokens: 100 + ) + + let responseText = samplingResult?.content.asArray + .compactMap { block -> String? in + if case .text(let text) = block { + return text + } + return nil + } + .joined(separator: "\n") ?? "No response" + + return .init(content: [.text(text: responseText, annotations: nil, _meta: nil)], isError: false) + case "test_elicitation": + // Test elicitation - request user input for username and email + guard let message = params.arguments?["message"]?.stringValue else { + return .init(content: [.text(text: "Missing required argument: message", annotations: nil, _meta: nil)], isError: true) + } + + let elicitationResult = try await server?.requestElicitation( + message: message, + requestedSchema: Elicitation.RequestSchema( + properties: [ + "username": .object(["type": .string("string")]), + "email": .object(["type": .string("string")]) + ], + required: ["username", "email"] + ) + ) + + return .init( + content: [.text(text: "Elicitation completed: action=\(elicitationResult?.action.rawValue ?? "unknown"), content=\(elicitationResult?.content ?? [:])", annotations: nil, _meta: nil)], + isError: false + ) + case "test_elicitation_sep1034_defaults": + // Test elicitation with default values (SEP-1034) + let elicitationResult = try await server?.requestElicitation( + message: "Please provide the following information:", + requestedSchema: Elicitation.RequestSchema( + properties: [ + "name": .object([ + "type": .string("string"), + "default": .string("John Doe") + ]), + "age": .object([ + "type": .string("integer"), + "default": .int(30) + ]), + "score": .object([ + "type": .string("number"), + "default": .double(95.5) + ]), + "status": .object([ + "type": .string("string"), + "enum": .array([.string("active"), .string("inactive"), .string("pending")]), + "default": .string("active") + ]), + "verified": .object([ + "type": .string("boolean"), + "default": .bool(true) + ]) + ] + ) + ) + + return .init( + content: [.text(text: "Elicitation completed: action=\(elicitationResult?.action.rawValue ?? "unknown"), content=\(elicitationResult?.content ?? [:])", annotations: nil, _meta: nil)], + isError: false + ) + case "test_elicitation_sep1330_enums": + // Test elicitation with enum variants (SEP-1330) + let elicitationResult = try await server?.requestElicitation( + message: "Select options for enum testing:", + requestedSchema: Elicitation.RequestSchema( + properties: [ + // 1. Untitled single-select + "untitledSingle": .object([ + "type": .string("string"), + "enum": .array([.string("option1"), .string("option2"), .string("option3")]) + ]), + // 2. Titled single-select + "titledSingle": .object([ + "type": .string("string"), + "oneOf": .array([ + .object(["const": .string("opt1"), "title": .string("Option One")]), + .object(["const": .string("opt2"), "title": .string("Option Two")]), + .object(["const": .string("opt3"), "title": .string("Option Three")]) + ]) + ]), + // 3. Legacy titled (deprecated enumNames) + "legacyEnum": .object([ + "type": .string("string"), + "enum": .array([.string("legacy1"), .string("legacy2"), .string("legacy3")]), + "enumNames": .array([.string("Legacy One"), .string("Legacy Two"), .string("Legacy Three")]) + ]), + // 4. Untitled multi-select + "untitledMulti": .object([ + "type": .string("array"), + "items": .object([ + "type": .string("string"), + "enum": .array([.string("multi1"), .string("multi2"), .string("multi3")]) + ]) + ]), + // 5. Titled multi-select + "titledMulti": .object([ + "type": .string("array"), + "items": .object([ + "anyOf": .array([ + .object(["const": .string("titled1"), "title": .string("Titled One")]), + .object(["const": .string("titled2"), "title": .string("Titled Two")]), + .object(["const": .string("titled3"), "title": .string("Titled Three")]) + ]) + ]) + ]) + ] + ) + ) + + return .init( + content: [.text(text: "Elicitation completed: action=\(elicitationResult?.action.rawValue ?? "unknown"), content=\(elicitationResult?.content ?? [:])", annotations: nil, _meta: nil)], + isError: false + ) + case "test_client_elicitation_defaults": + // Tool for client-side elicitation defaults test + let elicitationResult = try await server?.requestElicitation( + message: "Please provide your information (defaults available):", + requestedSchema: Elicitation.RequestSchema( + properties: [ + "name": .object([ + "type": .string("string"), + "default": .string("John Doe") + ]), + "age": .object([ + "type": .string("integer"), + "default": .int(30) + ]), + "score": .object([ + "type": .string("number"), + "default": .double(95.5) + ]), + "status": .object([ + "type": .string("string"), + "enum": .array([.string("active"), .string("inactive"), .string("pending")]), + "default": .string("active") + ]), + "verified": .object([ + "type": .string("boolean"), + "default": .bool(true) + ]) + ] + ) + ) + + // Verify the client applied defaults correctly + guard let content = elicitationResult?.content, + let name = content["name"]?.stringValue, + let age = content["age"]?.intValue, + let score = content["score"]?.doubleValue, + let status = content["status"]?.stringValue, + let verified = content["verified"]?.boolValue else { + return .init(content: [.text(text: "Client did not provide all required fields with defaults", annotations: nil, _meta: nil)], isError: true) + } + + guard name == "John Doe", age == 30, score == 95.5, status == "active", verified == true else { + return .init(content: [.text(text: "Client defaults do not match expected values", annotations: nil, _meta: nil)], isError: true) + } + + return .init( + content: [.text(text: "Client correctly applied all default values", annotations: nil, _meta: nil)], + isError: false + ) + default: + return .init(content: [.text(text: "Unknown tool: \(params.name)", annotations: nil, _meta: nil)], isError: true) + } + } + + // Resources + await server.withMethodHandler(ListResources.self) { _ in + .init(resources: [ + Resource(name: "Static Text Resource", uri: "test://static-text", description: "A simple static text resource", mimeType: "text/plain"), + Resource(name: "Static Binary Resource", uri: "test://static-binary", description: "A simple static binary resource", mimeType: "application/octet-stream"), + Resource(name: "Watched Resource", uri: "test://watched", description: "A resource that can be subscribed to for updates", mimeType: "text/plain"), + Resource(name: "Template Resource", uri: "test://template/{id}", description: "A resource template with URI parameters", mimeType: "text/plain"), + ]) + } + + await server.withMethodHandler(ReadResource.self) { params in + switch params.uri { + case "test://static-text": + return .init(contents: [.text("This is static text content for testing.", uri: params.uri, mimeType: "text/plain")]) + case "test://static-binary": + guard let imageData = Data(base64Encoded: testImageBase64) else { + return .init(contents: [.text("Failed to decode binary data", uri: params.uri)]) + } + return .init(contents: [.binary(imageData, uri: params.uri, mimeType: "application/octet-stream")]) + case "test://watched": + let content = await state.watchedResourceContent + return .init(contents: [.text(content, uri: params.uri)]) + default: + if params.uri.hasPrefix("test://template/") { + let id = String(params.uri.dropFirst("test://template/".count)) + return .init(contents: [.text("Template resource with id: \(id)", uri: params.uri)]) + } + return .init(contents: [.text("Resource not found: \(params.uri)", uri: params.uri)]) + } + } + + await server.withMethodHandler(ResourceSubscribe.self) { params in + await state.subscribe(to: params.uri) + return Empty() + } + + await server.withMethodHandler(ResourceUnsubscribe.self) { params in + await state.unsubscribe(from: params.uri) + return Empty() + } + + // Prompts + await server.withMethodHandler(ListPrompts.self) { _ in + .init(prompts: [ + Prompt(name: "test_simple_prompt", description: "A simple prompt without arguments"), + Prompt(name: "test_prompt_with_arguments", description: "A prompt that accepts arguments", arguments: [Prompt.Argument(name: "arg1", description: "First test argument", required: true), Prompt.Argument(name: "arg2", description: "Second test argument", required: true)]), + Prompt(name: "test_prompt_with_embedded_resource", description: "A prompt that includes embedded resources", arguments: [Prompt.Argument(name: "resourceUri", description: "URI of the resource to embed", required: true)]), + Prompt(name: "test_prompt_with_image", description: "A prompt with image content"), + ]) + } + + await server.withMethodHandler(GetPrompt.self) { params in + switch params.name { + case "test_simple_prompt": + return .init(description: "Simple prompt response", messages: [.user(.text(text: "This is a simple prompt for testing."))]) + case "test_prompt_with_arguments": + let arg1 = params.arguments?["arg1"] ?? "default1" + let arg2 = params.arguments?["arg2"] ?? "default2" + return .init(description: "Prompt with arguments", messages: [.user(.text(text: "Prompt with arguments: arg1='\(arg1)', arg2='\(arg2)'"))]) + case "test_prompt_with_embedded_resource": + let resourceUri = params.arguments?["resourceUri"] ?? "test://default" + return .init(description: "Prompt with embedded resource", messages: [ + .user(.resource(resource: .text("Embedded resource content for testing.", uri: resourceUri, mimeType: "text/plain"))), + .user(.text(text: "Please process the embedded resource above.")) + ]) + case "test_prompt_with_image": + return .init(description: "Prompt with image", messages: [ + .user(.image(data: testImageBase64, mimeType: "image/png")), + .user(.text(text: "Please analyze the image above.")) + ]) + default: + throw MCPError.invalidRequest("Unknown prompt: \(params.name)") + } + } + + await server.withMethodHandler(SetLoggingLevel.self) { _ in + // Accept any logging level (debug, info, notice, warning, error, critical, alert, emergency) + // For conformance testing, we just accept it without doing anything + return Empty() + } + + await server.withMethodHandler(Complete.self) { _ in + return .init(completion: .init(values: [])) + } + + return server +} + +// MARK: - HTTP Server + +// HTTPApp handles all HTTP server functionality + +// MARK: - Main + +struct MCPHTTPServer { + static func run() async throws { + let args = CommandLine.arguments + var port = 3001 + + for (index, arg) in args.enumerated() { + if arg == "--port" && index + 1 < args.count { + if let p = Int(args[index + 1]) { + port = p + } + } + } + + var loggerConfig = Logger(label: "mcp.http.server", factory: { StreamLogHandler.standardError(label: $0) }) + loggerConfig.logLevel = .trace + let logger = loggerConfig + + let state = ServerState() + + logger.info("Starting MCP HTTP Server...", metadata: ["port": "\(port)"]) + + // Create HTTPApp with server factory + let app = HTTPApp( + configuration: .init( + host: "127.0.0.1", + port: port, + endpoint: "/mcp", + retryInterval: 1000 + ), + validationPipeline: StandardValidationPipeline(validators: [ + OriginValidator.localhost(port: port), + AcceptHeaderValidator(mode: .sseRequired), + ContentTypeValidator(), + ProtocolVersionValidator(), + SessionValidator(), + ]), + serverFactory: { sessionID, transport in + logger.debug("Creating server for session", metadata: ["sessionID": "\(sessionID)"]) + return await createConformanceServer(state: state, transport: transport) + }, + logger: logger + ) + + try await app.start() + } +} + +do { + try await MCPHTTPServer.run() +} catch { + print(error) + exit(1) +} diff --git a/Tests/MCPTests/Base64DataTests.swift b/Tests/MCPTests/Base64DataTests.swift new file mode 100644 index 0000000..9602593 --- /dev/null +++ b/Tests/MCPTests/Base64DataTests.swift @@ -0,0 +1,142 @@ +import Foundation +import Testing + +@testable import MCP + +@Suite("Base64 Data Tests") +struct Base64DataTests { + + @Test("Check valid data URLs") + func testIsDataURL() { + // Basic valid data URL + #expect(Data.isDataURL(string: "data:,A%20brief%20note")) + + // Valid image data URL from RFC example + #expect( + Data.isDataURL(string: "data:image/gif;base64,R0lGODdhMAAwAPAAAAAAAP///ywAAAAAMAAw")) + + // Valid data URL with charset parameter + #expect(Data.isDataURL(string: "data:text/plain;charset=iso-8859-7,%be%fg%be")) + + // Valid custom application type from RFC example + #expect( + Data.isDataURL( + string: "data:application/vnd-xxx-query,select_vcount,fcol_from_fieldtable/local")) + + // Valid with implicit mediatype (defaults to text/plain) + #expect(Data.isDataURL(string: "data:,hello%20world")) + + // Valid with just charset parameter + #expect(Data.isDataURL(string: "data:;charset=utf-8,hello")) + } + + @Test("Check invalid data URLs") + func testInvalidDataURLs() { + // Not a data URL + #expect(!Data.isDataURL(string: "https://example.com")) + + // Missing comma separator + #expect(!Data.isDataURL(string: "data:text/plain")) + + // Malformed scheme + #expect(!Data.isDataURL(string: "dat:,hello")) + + // Empty string + #expect(!Data.isDataURL(string: "")) + } + + @Test("Parse data URL with plain text") + func testParseDataURLPlainText() { + // Example from RFC 2397 + let result = Data.parseDataURL("data:,A%20brief%20note") + #expect(result != nil) + if let (mimeType, data) = result { + #expect(mimeType == "text/plain") + #expect(String(data: data, encoding: .utf8) == "A brief note") + } + } + + @Test("Parse data URL with charset") + func testParseDataURLWithCharset() { + // Modified example from RFC 2397 (using valid percent encoding) + let result = Data.parseDataURL( + "data:text/plain;charset=iso-8859-7,%CF%84%CE%B5%CF%83%CF%84") + #expect(result != nil) + if let (mimeType, _) = result { + #expect(mimeType == "text/plain;charset=iso-8859-7") + } + } + + @Test("Parse data URL with base64 encoding") + func testParseDataURLBase64() { + // Simple base64 encoded "Hello" + let result = Data.parseDataURL("data:text/plain;base64,SGVsbG8=") + #expect(result != nil) + if let (mimeType, data) = result { + #expect(mimeType == "text/plain") + #expect(String(data: data, encoding: .utf8) == "Hello") + } + } + + @Test("Parse data URL with implicit text/plain") + func testParseDataURLImplicitTextPlain() { + let result = Data.parseDataURL("data:,hello%20world") + #expect(result != nil) + if let (mimeType, data) = result { + #expect(mimeType == "text/plain") + #expect(String(data: data, encoding: .utf8) == "hello world") + } + } + + @Test("Parse invalid data URLs") + func testParseInvalidDataURLs() { + // Not a data URL + #expect(Data.parseDataURL("https://example.com") == nil) + + // Malformed base64 + #expect(Data.parseDataURL("data:text/plain;base64,SGVsbG8=!") == nil) + + // Empty base64 data (valid according to RFC 2397) + let result = Data.parseDataURL("data:text/plain;base64,") + #expect(result != nil) + if let (mimeType, data) = result { + #expect(mimeType == "text/plain") + #expect(data.isEmpty) + } + } + + @Test("Test data URL encoding") + func testDataURLEncoding() { + let testData = "Hello, world!".data(using: .utf8)! + + // Default mime type (text/plain) + let url1 = testData.dataURLEncoded() + #expect(url1.hasPrefix("data:text/plain;base64,")) + #expect(Data.isDataURL(string: url1)) + + // Custom mime type + let url2 = testData.dataURLEncoded(mimeType: "application/octet-stream") + #expect(url2.hasPrefix("data:application/octet-stream;base64,")) + #expect(Data.isDataURL(string: url2)) + + // Roundtrip test + if let (mimeType, data) = Data.parseDataURL(url2) { + #expect(mimeType == "application/octet-stream") + #expect(data == testData) + } else { + #expect(Bool(false), "Failed to parse encoded data URL") + } + } + + @Test("Test complex example from RFC 2397") + func testComplexExample() { + // Partial example of image from RFC (shortened for test brevity) + let dataURL = "data:image/gif;base64,R0lGODdhMAAwAPAAAAAAAP///ywAAAAAMAAw" + + let result = Data.parseDataURL(dataURL) + #expect(result != nil) + if let (mimeType, _) = result { + #expect(mimeType == "image/gif") + } + } +} diff --git a/Tests/MCPTests/CancellationTests.swift b/Tests/MCPTests/CancellationTests.swift new file mode 100644 index 0000000..9d4412e --- /dev/null +++ b/Tests/MCPTests/CancellationTests.swift @@ -0,0 +1,249 @@ +import Foundation +import Testing + +@testable import MCP + +@Suite("Cancellation Tests") +struct CancellationTests { + @Test("Client sends CancelledNotification") + func testClientSendsCancellation() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let client = Client(name: "TestClient", version: "1.0") + let server = Server( + name: "TestServer", + version: "1.0", + capabilities: .init() + ) + + // Start server + try await server.start(transport: serverTransport) + + // Connect client + _ = try await client.connect(transport: clientTransport) + try await Task.sleep(for: .milliseconds(50)) + + // Send a ping request + let pingRequest = Ping.request() + let context = try await client.send(pingRequest) + + try await Task.sleep(for: .milliseconds(10)) + + // Cancel the request + try await client.cancelRequest(context.requestID, reason: "Test cancellation") + + try await Task.sleep(for: .milliseconds(50)) + + // Verify cancellation was sent (server should have received it) + // The test passes if no errors occur and the request is cancelled + + await client.disconnect() + await server.stop() + } + + @Test("Client receives and processes CancelledNotification") + func testClientReceivesCancellation() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let client = Client(name: "TestClient", version: "1.0") + let server = Server( + name: "TestServer", + version: "1.0", + capabilities: .init() + ) + + // Start server and connect client + try await server.start(transport: serverTransport) + + // Register a slow ping handler (must be after start to override default) + await server.withMethodHandler(Ping.self) { _ in + try await Task.sleep(for: .seconds(5)) + return Empty() + } + _ = try await client.connect(transport: clientTransport) + try await Task.sleep(for: .milliseconds(50)) + + // Send a request using send + let pingRequest = Ping.request() + let context = try await client.send(pingRequest) + + // Server cancels the request while it's being awaited + try await Task.sleep(for: .milliseconds(50)) + try await server.cancelRequest(pingRequest.id, reason: "Server cancelled") + + // Try to get result - should throw CancellationError + do { + _ = try await context.value + Issue.record("Expected CancellationError but request succeeded") + } catch is CancellationError { + // Expected + } catch { + Issue.record("Expected CancellationError but got: \(error)") + } + + await client.disconnect() + await server.stop() + } + + @Test("RequestContext structure") + func testRequestContextStructure() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let client = Client(name: "TestClient", version: "1.0") + let server = Server( + name: "TestServer", + version: "1.0", + capabilities: .init() + ) + + // Start server and connect client + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + try await Task.sleep(for: .milliseconds(50)) + + // Create a request with send + let pingRequest = Ping.request() + let context: RequestContext = try await client.send(pingRequest) + + // Verify the context has the correct requestID + #expect(context.requestID == pingRequest.id) + + // Await the result through the context + let result = try await context.value + #expect(result == Empty()) + + await client.disconnect() + await server.stop() + } + + @Test("callTool with RequestContext overload") + func testCallToolWithRequestContext() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let client = Client(name: "TestClient", version: "1.0") + let server = Server( + name: "TestServer", + version: "1.0", + capabilities: .init(tools: .init()) + ) + + // Register a tool handler + await server.withMethodHandler(ListTools.self) { _ in + .init(tools: [Tool(name: "testTool", description: "A test tool", inputSchema: .object([:]))]) + } + + await server.withMethodHandler(CallTool.self) { params in + return .init(content: [.text(text: "Result for \(params.name)", annotations: nil, _meta: nil)], isError: false) + } + + // Start server and connect client + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + try await Task.sleep(for: .milliseconds(50)) + + // Use the callTool overload that returns RequestContext (non-async version) + let context: RequestContext = try await client.callTool(name: "testTool", arguments: ["test": "value"]) + + // Verify we got a context + #expect(context.requestID != ID(stringLiteral: "")) + + // Get the result + let result = try await context.value + #expect(result.content == [.text(text: "Result for testTool", annotations: nil, _meta: nil)]) + #expect(result.isError == false) + + await client.disconnect() + await server.stop() + } + + @Test("Cancel callTool using RequestContext overload") + func testCancelCallToolWithRequestContext() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let client = Client(name: "TestClient", version: "1.0") + let server = Server( + name: "TestServer", + version: "1.0", + capabilities: .init(tools: .init()) + ) + + // Register a tool handler that takes time + await server.withMethodHandler(ListTools.self) { _ in + .init(tools: [Tool(name: "slowTool", description: "A slow tool", inputSchema: .object([:]))]) + } + + await server.withMethodHandler(CallTool.self) { params in + try await Task.sleep(for: .seconds(5)) + return .init(content: [.text(text: "Should not reach here", annotations: nil, _meta: nil)], isError: false) + } + + // Start server and connect client + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + try await Task.sleep(for: .milliseconds(50)) + + // Use the callTool overload that returns RequestContext (non-async version) + let context: RequestContext = try await client.callTool(name: "slowTool", arguments: [:]) + + // Cancel after a short delay + try await Task.sleep(for: .milliseconds(50)) + try await client.cancelRequest(context.requestID, reason: "Test timeout") + + // Try to get result - should throw CancellationError + do { + _ = try await context.value + Issue.record("Expected CancellationError but request succeeded") + } catch is CancellationError { + // Expected + } catch { + Issue.record("Expected CancellationError but got: \(error)") + } + + await client.disconnect() + await server.stop() + } + + @Test("CancelledNotification prevents response") + func testCancellationPreventsResponse() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let client = Client(name: "TestClient", version: "1.0") + let server = Server( + name: "TestServer", + version: "1.0", + capabilities: .init() + ) + + // Start server and connect client + try await server.start(transport: serverTransport) + + // Register a slow handler (must be after start to override default) + await server.withMethodHandler(Ping.self) { _ in + try await Task.sleep(for: .seconds(10)) + return Empty() + } + _ = try await client.connect(transport: clientTransport) + try await Task.sleep(for: .milliseconds(50)) + + // Send a ping request + let pingRequest = Ping.request() + let context = try await client.send(pingRequest) + + // Cancel the request while it's being awaited + try await Task.sleep(for: .milliseconds(50)) + try await client.cancelRequest(context.requestID, reason: "Test cancellation") + + // Try to get result - should throw CancellationError (proving no response was sent) + do { + _ = try await context.value + Issue.record("Expected CancellationError but request succeeded") + } catch is CancellationError { + // Expected - this proves the server didn't send a response + } catch { + Issue.record("Expected CancellationError but got: \(error)") + } + + await client.disconnect() + await server.stop() + } +} diff --git a/Tests/MCPTests/ClientTests.swift b/Tests/MCPTests/ClientTests.swift index a2cba6b..8a2882a 100644 --- a/Tests/MCPTests/ClientTests.swift +++ b/Tests/MCPTests/ClientTests.swift @@ -1,8 +1,9 @@ +import Foundation import Testing @testable import MCP -@Suite("Client Tests") +@Suite("Client Tests", .timeLimit(.minutes(1))) struct ClientTests { @Test("Client connect and disconnect") func testClientConnectAndDisconnect() async throws { @@ -10,43 +11,33 @@ struct ClientTests { let client = Client(name: "TestClient", version: "1.0") #expect(await transport.isConnected == false) - try await client.connect(transport: transport) - #expect(await transport.isConnected == true) - await client.disconnect() - #expect(await transport.isConnected == false) - } - @Test( - "Initialize request", - .timeLimit(.minutes(1)) - ) - func testClientInitialize() async throws { - let transport = MockTransport() - let client = Client(name: "TestClient", version: "1.0") - - try await client.connect(transport: transport) - // Small delay to ensure message loop is started - try await Task.sleep(nanoseconds: 10_000_000) // 10ms - - // Create a task for initialize that we'll cancel + // Set up a task to handle the initialize response let initTask = Task { - try await client.initialize() + 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) + } } - // Give it a moment to send the request - try await Task.sleep(nanoseconds: 10_000_000) // 10ms - - #expect(await transport.sentMessages.count == 1) - #expect(await transport.sentMessages[0].contains(Initialize.name)) - #expect(await transport.sentMessages[0].contains(client.name)) - #expect(await transport.sentMessages[0].contains(client.version)) - - // Cancel the initialize task - initTask.cancel() - - // Disconnect client to clean up message loop and give time for continuation cleanup + let result = try await client.connect(transport: transport) + #expect(await transport.isConnected == true) + #expect(result.protocolVersion == Version.latest) await client.disconnect() - try await Task.sleep(nanoseconds: 50_000_000) // 50ms + #expect(await transport.isConnected == false) + initTask.cancel() } @Test( @@ -57,27 +48,53 @@ struct ClientTests { let transport = MockTransport() let client = Client(name: "TestClient", version: "1.0") - try await client.connect(transport: transport) - // Small delay to ensure message loop is started - try await Task.sleep(nanoseconds: 10_000_000) // 10ms + // Queue a response for the initialize request + try await Task.sleep(for: .milliseconds(10)) // Wait for request to be sent - // Create a task for the ping that we'll cancel - let pingTask = Task { - try await client.ping() + if let lastMessage = await transport.sentMessages.last, + let data = lastMessage.data(using: .utf8), + let request = try? JSONDecoder().decode(Request.self, from: data) + { + // Create a valid initialize response + 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) + + // Now complete the connect call which will automatically initialize + let result = try await client.connect(transport: transport) + #expect(result.protocolVersion == Version.latest) + #expect(result.serverInfo.name == "TestServer") + #expect(result.serverInfo.version == "1.0") + + // Small delay to ensure message loop is started + try await Task.sleep(for: .milliseconds(10)) + + // Create a task for the ping + let pingTask = Task { + try await client.ping() + } + + // Give it a moment to send the request + try await Task.sleep(for: .milliseconds(10)) + + #expect(await transport.sentMessages.count == 2) // Initialize + Ping + #expect(await transport.sentMessages.last?.contains(Ping.name) == true) + + // Cancel the ping task + pingTask.cancel() } - // Give it a moment to send the request - try await Task.sleep(nanoseconds: 10_000_000) // 10ms - - #expect(await transport.sentMessages.count == 1) - #expect(await transport.sentMessages[0].contains(Ping.name)) - - // Cancel the ping task - pingTask.cancel() - // Disconnect client to clean up message loop and give time for continuation cleanup await client.disconnect() - try await Task.sleep(nanoseconds: 50_000_000) // 50ms + try await Task.sleep(for: .milliseconds(50)) } @Test("Connection failure handling") @@ -89,8 +106,8 @@ struct ClientTests { do { try await client.connect(transport: transport) #expect(Bool(false), "Expected connection to fail") - } catch let error as Error { - if case Error.transportError = error { + } catch let error as MCPError { + if case MCPError.transportError = error { #expect(Bool(true)) } else { #expect(Bool(false), "Expected transport error") @@ -103,22 +120,656 @@ struct ClientTests { @Test("Send failure handling") func testClientSendFailure() async throws { let transport = MockTransport() - await transport.setFailSend(true) 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) + } + } + + // Connect first without failure try await client.connect(transport: transport) + try await Task.sleep(for: .milliseconds(10)) + initTask.cancel() + + // Now set the transport to fail sends + await transport.setFailSend(true) do { try await client.ping() #expect(Bool(false), "Expected ping to fail") - } catch let error as Error { - if case Error.transportError = error { + } catch let error as MCPError { + if case MCPError.transportError = error { #expect(Bool(true)) } else { - #expect(Bool(false), "Expected transport error") + #expect(Bool(false), "Expected transport error, got \(error)") } } catch { #expect(Bool(false), "Expected MCP.Error") } + + await client.disconnect() + } + + @Test("Strict configuration - capabilities check") + func testStrictConfiguration() async throws { + let transport = MockTransport() + let config = Client.Configuration.strict + let client = Client(name: "TestClient", version: "1.0", configuration: config) + + // 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) + + // Create a task for listPrompts + let promptsTask = Task { + do { + _ = try await client.listPrompts() + #expect(Bool(false), "Expected listPrompts to fail in strict mode") + } catch let error as MCPError { + if case MCPError.methodNotFound = error { + #expect(Bool(true)) + } else { + #expect(Bool(false), "Expected methodNotFound error, got \(error)") + } + } catch { + #expect(Bool(false), "Expected MCP.Error") + } + } + + // Give it a short time to execute the task + try await Task.sleep(for: .milliseconds(50)) + + // Cancel the task if it's still running + promptsTask.cancel() + initTask.cancel() + + // Disconnect client + await client.disconnect() + try await Task.sleep(for: .milliseconds(50)) + } + + @Test("Non-strict configuration - capabilities check") + func testNonStrictConfiguration() async throws { + let transport = MockTransport() + let config = Client.Configuration.default + let client = Client(name: "TestClient", version: "1.0", configuration: config) + + // 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) + + // Make sure init task is complete + initTask.cancel() + + // Wait a bit for any setup to complete + try await Task.sleep(for: .milliseconds(10)) + + // Send the listPrompts request and immediately provide an error response + let promptsTask = Task { + do { + // Start the request + try await Task.sleep(for: .seconds(1)) + + // Get the last sent message and extract the request ID + if let lastMessage = await transport.sentMessages.last, + let data = lastMessage.data(using: .utf8), + let decodedRequest = try? JSONDecoder().decode( + Request.self, from: data) + { + + // Create an error response with the same ID + let errorResponse = Response( + id: decodedRequest.id, + error: MCPError.methodNotFound("Test: Prompts capability not available") + ) + try await transport.queue(response: errorResponse) + + // Try the request now that we have a response queued + do { + _ = try await client.listPrompts() + #expect(Bool(false), "Expected listPrompts to fail in non-strict mode") + } catch let error as MCPError { + if case MCPError.methodNotFound = error { + #expect(Bool(true)) + } else { + #expect(Bool(false), "Expected methodNotFound error, got \(error)") + } + } catch { + #expect(Bool(false), "Expected MCP.Error") + } + } + } catch { + // Ignore task cancellation + if !(error is CancellationError) { + throw error + } + } + } + + // Wait for the task to complete or timeout + let timeoutTask = Task { + try await Task.sleep(for: .milliseconds(500)) + promptsTask.cancel() + } + + // Wait for the task to complete + _ = await promptsTask.result + + // Cancel the timeout task + timeoutTask.cancel() + + // Disconnect client + await client.disconnect() + } + + @Test("Batch request - success") + func testBatchRequestSuccess() async throws { + 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)) // Allow connection tasks + initTask.cancel() + + let request1 = Ping.request() + let request2 = Ping.request() + nonisolated(unsafe) var resultTask1: Task? + nonisolated(unsafe) var resultTask2: Task? + + try await client.withBatch { batch in + resultTask1 = try await batch.addRequest(request1) + resultTask2 = try await batch.addRequest(request2) + } + + // Check if batch message was sent (after initialize and initialized notification) + let sentMessages = await transport.sentMessages + #expect(sentMessages.count == 3) // Initialize request + Initialized notification + Batch + + guard let batchData = sentMessages.last?.data(using: .utf8) else { + #expect(Bool(false), "Failed to get batch data") + return + } + + // Verify the sent batch contains the two requests + let decoder = JSONDecoder() + let sentRequests = try decoder.decode([AnyRequest].self, from: batchData) + #expect(sentRequests.count == 2) + #expect(sentRequests.first?.id == request1.id) + #expect(sentRequests.first?.method == Ping.name) + #expect(sentRequests.last?.id == request2.id) + #expect(sentRequests.last?.method == Ping.name) + + // Prepare batch response + let response1 = Response(id: request1.id, result: .init()) + let response2 = Response(id: request2.id, result: .init()) + let anyResponse1 = try AnyResponse(response1) + let anyResponse2 = try AnyResponse(response2) + + // Queue the batch response + try await transport.queue(batch: [anyResponse1, anyResponse2]) + + // Wait for results and verify + guard let task1 = resultTask1, let task2 = resultTask2 else { + #expect(Bool(false), "Result tasks not created") + return + } + + _ = try await task1.value // Should succeed + _ = try await task2.value // Should succeed + + #expect(Bool(true)) // Reaching here means success + + await client.disconnect() + } + + @Test("Batch request - mixed success/error") + func testBatchRequestMixed() async throws { + 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() + + let request1 = Ping.request() // Success + let request2 = Ping.request() // Error + + nonisolated(unsafe) var resultTasks: [Task] = [] + + try await client.withBatch { batch in + resultTasks.append(try await batch.addRequest(request1)) + resultTasks.append(try await batch.addRequest(request2)) + } + + // Check if batch message was sent (after initialize and initialized notification) + #expect(await transport.sentMessages.count == 3) // Initialize request + Initialized notification + Batch + + // Prepare batch response (success for 1, error for 2) + let response1 = Response(id: request1.id, result: .init()) + let error = MCPError.internalError("Simulated batch error") + let response2 = Response(id: request2.id, error: error) + let anyResponse1 = try AnyResponse(response1) + let anyResponse2 = try AnyResponse(response2) + + // Queue the batch response + try await transport.queue(batch: [anyResponse1, anyResponse2]) + + // Wait for results and verify + #expect(resultTasks.count == 2) + guard resultTasks.count == 2 else { + #expect(Bool(false), "Expected 2 result tasks") + return + } + + let task1 = resultTasks[0] + let task2 = resultTasks[1] + + _ = try await task1.value // Task 1 should succeed + + do { + _ = try await task2.value // Task 2 should fail + #expect(Bool(false), "Task 2 should have thrown an error") + } catch let mcpError as MCPError { + if case .internalError(let message) = mcpError { + #expect(message == "Simulated batch error") + } else { + #expect(Bool(false), "Expected internalError, got \(mcpError)") + } + } catch { + #expect(Bool(false), "Expected MCPError, got \(error)") + } + + await client.disconnect() + } + + @Test("Batch request - empty") + func testBatchRequestEmpty() async throws { + 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() + + // Call withBatch but don't add any requests + try await client.withBatch { _ in + // No requests added + } + + // Check that only initialize message and initialized notification were sent + #expect(await transport.sentMessages.count == 2) // Initialize request + Initialized notification + + await client.disconnect() + } + + @Test("Notify method sends notifications") + func testClientNotify() async throws { + 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() + + // Create a test notification + let notification = InitializedNotification.message() + try await client.notify(notification) + + // Verify notification was sent (in addition to initialize and initialized notification) + #expect(await transport.sentMessages.count == 3) // Initialize request + Initialized notification + Custom notification + + if let sentMessage = await transport.sentMessages.last, + let data = sentMessage.data(using: .utf8) + { + + // Decode as Message + let decoder = JSONDecoder() + do { + let decodedNotification = try decoder.decode( + Message.self, from: data) + #expect(decodedNotification.method == InitializedNotification.name) + } catch { + #expect(Bool(false), "Failed to decode notification: \(error)") + } + } else { + #expect(Bool(false), "No message was sent") + } + + await client.disconnect() + } + + @Test("Initialize sends initialized notification") + func testClientInitializeNotification() async throws { + let transport = MockTransport() + let client = Client(name: "TestClient", version: "1.0") + + // Create a task for initialize + let initTask = Task { + // Queue a response for the initialize request + try await Task.sleep(for: .milliseconds(10)) // Wait for request to be sent + + if let lastMessage = await transport.sentMessages.last, + let data = lastMessage.data(using: .utf8), + let request = try? JSONDecoder().decode(Request.self, from: data) + { + + // Create a valid initialize response + 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) + + // Now complete the initialize call + try await client.connect(transport: transport) + try await Task.sleep(for: .milliseconds(10)) + + // Verify that two messages were sent: initialize request and initialized notification + #expect(await transport.sentMessages.count == 2) + + // Check that the second message is the initialized notification + let notifications = await transport.sentMessages + if notifications.count >= 2 { + let notificationJson = notifications[1] + if let notificationData = notificationJson.data(using: .utf8) { + do { + let decoder = JSONDecoder() + let decodedNotification = try decoder.decode( + Message.self, from: notificationData) + #expect(decodedNotification.method == InitializedNotification.name) + } catch { + #expect(Bool(false), "Failed to decode notification: \(error)") + } + } else { + #expect(Bool(false), "Could not convert notification to data") + } + } else { + #expect( + Bool(false), "Expected both initialize request and initialized notification" + ) + } + } + } + + // Wait with timeout + let timeoutTask = Task { + try await Task.sleep(for: .seconds(1)) + initTask.cancel() + } + + // Wait for the task to complete + do { + _ = try await initTask.value + } catch is CancellationError { + #expect(Bool(false), "Test timed out") + } catch { + #expect(Bool(false), "Unexpected error: \(error)") + } + + timeoutTask.cancel() + + await client.disconnect() + } + + @Test("Race condition between send error and response") + func testSendErrorResponseRace() async throws { + 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() + + // Set up the transport to fail sends from the start + await transport.setFailSend(true) + + // Create a ping request to get the ID + let request = Ping.request() + + // Create a response for the request and queue it immediately + let response = Response(id: request.id, result: .init()) + let anyResponse = try AnyResponse(response) + try await transport.queue(response: anyResponse) + + // Now attempt to send the request - this should fail due to send error + // but the response handler might also try to process the queued response + do { + _ = try await client.ping() + #expect(Bool(false), "Expected send to fail") + } catch let error as MCPError { + if case .transportError = error { + #expect(Bool(true)) + } else { + #expect(Bool(false), "Expected transport error, got \(error)") + } + } catch { + #expect(Bool(false), "Expected MCPError, got \(error)") + } + + // Verify no continuation misuse occurred + // (If it did, the test would have crashed) + + await client.disconnect() + } + + @Test("Race condition between response and send error") + func testResponseSendErrorRace() async throws { + 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() + + // Create a ping request to get the ID + let request = Ping.request() + + // Create a response for the request and queue it immediately + let response = Response(id: request.id, result: .init()) + let anyResponse = try AnyResponse(response) + try await transport.queue(response: anyResponse) + + // Set up the transport to fail sends + await transport.setFailSend(true) + + // Now attempt to send the request + // The response might be processed before the send error occurs + do { + _ = try await client.ping() + // In this case, the response handler won the race and the request succeeded + #expect(Bool(true), "Response handler won the race - request succeeded") + } catch let error as MCPError { + if case .transportError = error { + // In this case, the send error handler won the race + #expect(Bool(true), "Send error handler won the race - request failed") + } else { + #expect(Bool(false), "Expected transport error, got \(error)") + } + } catch { + #expect(Bool(false), "Expected MCPError, got \(error)") + } + + // Verify no continuation misuse occurred + // (If it did, the test would have crashed) + + await client.disconnect() } } diff --git a/Tests/MCPTests/CompletionTests.swift b/Tests/MCPTests/CompletionTests.swift new file mode 100644 index 0000000..38b19e4 --- /dev/null +++ b/Tests/MCPTests/CompletionTests.swift @@ -0,0 +1,598 @@ +import Foundation +import Testing + +@testable import MCP + +@Suite("Completion Tests") +struct CompletionTests { + // MARK: - Reference Types Tests + + @Test("PromptReference initialization and encoding") + func testPromptReferenceEncodingDecoding() throws { + let ref = PromptReference(name: "code_review") + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + + let data = try encoder.encode(ref) + let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] + + #expect(json?["type"] as? String == "ref/prompt") + #expect(json?["name"] as? String == "code_review") + + // Test decoding + let decoder = JSONDecoder() + let decoded = try decoder.decode(PromptReference.self, from: data) + #expect(decoded.name == "code_review") + } + + @Test("ResourceReference initialization and encoding") + func testResourceReferenceEncodingDecoding() throws { + let ref = ResourceReference(uri: "file:///path/to/resource") + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + + let data = try encoder.encode(ref) + let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] + + #expect(json?["type"] as? String == "ref/resource") + #expect(json?["uri"] as? String == "file:///path/to/resource") + + // Test decoding + let decoder = JSONDecoder() + let decoded = try decoder.decode(ResourceReference.self, from: data) + #expect(decoded.uri == "file:///path/to/resource") + } + + @Test("CompletionReference prompt case encoding") + func testCompletionReferencePromptEncoding() throws { + let ref = CompletionReference.prompt(PromptReference(name: "test")) + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + + let data = try encoder.encode(ref) + let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] + + #expect(json?["type"] as? String == "ref/prompt") + #expect(json?["name"] as? String == "test") + } + + @Test("CompletionReference resource case encoding") + func testCompletionReferenceResourceEncoding() throws { + let ref = CompletionReference.resource(ResourceReference(uri: "file:///test")) + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + + let data = try encoder.encode(ref) + let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] + + #expect(json?["type"] as? String == "ref/resource") + #expect(json?["uri"] as? String == "file:///test") + } + + @Test("CompletionReference decoding prompt type") + func testCompletionReferenceDecodingPrompt() throws { + let json = """ + { + "type": "ref/prompt", + "name": "code_review" + } + """ + + let decoder = JSONDecoder() + let ref = try decoder.decode(CompletionReference.self, from: json.data(using: .utf8)!) + + if case .prompt(let promptRef) = ref { + #expect(promptRef.name == "code_review") + } else { + Issue.record("Expected prompt reference") + } + } + + @Test("CompletionReference decoding resource type") + func testCompletionReferenceDecodingResource() throws { + let json = """ + { + "type": "ref/resource", + "uri": "file:///path" + } + """ + + let decoder = JSONDecoder() + let ref = try decoder.decode(CompletionReference.self, from: json.data(using: .utf8)!) + + if case .resource(let resourceRef) = ref { + #expect(resourceRef.uri == "file:///path") + } else { + Issue.record("Expected resource reference") + } + } + + // MARK: - Complete Request Tests + + @Test("Complete request initialization") + func testCompleteRequestInitialization() throws { + let ref = CompletionReference.prompt(PromptReference(name: "code_review")) + let argument = Complete.Parameters.Argument(name: "language", value: "py") + let request = Complete.request(.init(ref: ref, argument: argument)) + + #expect(request.method == "completion/complete") + #expect(request.params.argument.name == "language") + #expect(request.params.argument.value == "py") + } + + @Test("Complete request with context") + func testCompleteRequestWithContext() throws { + let ref = CompletionReference.prompt(PromptReference(name: "code_review")) + let argument = Complete.Parameters.Argument(name: "framework", value: "fla") + let context = Complete.Parameters.Context(arguments: ["language": "python"]) + + let request = Complete.request(.init(ref: ref, argument: argument, context: context)) + + #expect(request.params.context != nil) + #expect(request.params.context?.arguments["language"] == "python") + } + + @Test("Complete request encoding") + func testCompleteRequestEncoding() throws { + let ref = CompletionReference.prompt(PromptReference(name: "code_review")) + let argument = Complete.Parameters.Argument(name: "language", value: "py") + let request = Complete.request(.init(ref: ref, argument: argument)) + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + + let data = try encoder.encode(request) + let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] + + #expect(json?["jsonrpc"] as? String == "2.0") + #expect(json?["method"] as? String == "completion/complete") + + guard let params = json?["params"] as? [String: Any] else { + Issue.record("Failed to get params") + return + } + guard let refDict = params["ref"] as? [String: Any] else { + Issue.record("Failed to get ref") + return + } + #expect(refDict["type"] as? String == "ref/prompt") + #expect(refDict["name"] as? String == "code_review") + + guard let arg = params["argument"] as? [String: Any] else { + Issue.record("Failed to get argument") + return + } + #expect(arg["name"] as? String == "language") + #expect(arg["value"] as? String == "py") + } + + @Test("Complete request decoding") + func testCompleteRequestDecoding() throws { + let json = """ + { + "jsonrpc": "2.0", + "id": "test-id", + "method": "completion/complete", + "params": { + "ref": { + "type": "ref/prompt", + "name": "code_review" + }, + "argument": { + "name": "language", + "value": "py" + } + } + } + """ + + let decoder = JSONDecoder() + let request = try decoder.decode(Request.self, from: json.data(using: .utf8)!) + + #expect(request.method == "completion/complete") + #expect(request.params.argument.name == "language") + #expect(request.params.argument.value == "py") + + if case .prompt(let promptRef) = request.params.ref { + #expect(promptRef.name == "code_review") + } else { + Issue.record("Expected prompt reference") + } + } + + // MARK: - Complete Result Tests + + @Test("Complete result initialization") + func testCompleteResultInitialization() throws { + let completion = Complete.Result.Completion( + values: ["python", "pytorch", "pyside"], + total: 10, + hasMore: true + ) + let result = Complete.Result(completion: completion) + + #expect(result.completion.values.count == 3) + #expect(result.completion.values[0] == "python") + #expect(result.completion.total == 10) + #expect(result.completion.hasMore == true) + } + + @Test("Complete result encoding") + func testCompleteResultEncoding() throws { + let completion = Complete.Result.Completion( + values: ["python", "pytorch"], + total: 2, + hasMore: false + ) + let result = Complete.Result(completion: completion) + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + + let data = try encoder.encode(result) + let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] + + let completionDict = json?["completion"] as? [String: Any] + let values = completionDict?["values"] as? [String] + #expect(values == ["python", "pytorch"]) + #expect(completionDict?["total"] as? Int == 2) + #expect(completionDict?["hasMore"] as? Bool == false) + } + + @Test("Complete result decoding") + func testCompleteResultDecoding() throws { + let json = """ + { + "completion": { + "values": ["python", "pytorch", "pyside"], + "total": 10, + "hasMore": true + } + } + """ + + let decoder = JSONDecoder() + let result = try decoder.decode(Complete.Result.self, from: json.data(using: .utf8)!) + + #expect(result.completion.values.count == 3) + #expect(result.completion.values == ["python", "pytorch", "pyside"]) + #expect(result.completion.total == 10) + #expect(result.completion.hasMore == true) + } + + @Test("Complete result with optional fields") + func testCompleteResultWithOptionalFields() throws { + let completion = Complete.Result.Completion( + values: ["value1"], + total: nil, + hasMore: nil + ) + + #expect(completion.values == ["value1"]) + #expect(completion.total == nil) + #expect(completion.hasMore == nil) + } + + // MARK: - Client Integration Tests + + @Test("Client complete for prompt argument") + func testClientCompleteForPrompt() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let client = Client(name: "TestClient", version: "1.0") + let server = Server( + name: "TestServer", + version: "1.0", + capabilities: .init(completions: .init()) + ) + + // Register handler for complete on server + await server.withMethodHandler(Complete.self) { params in + #expect(params.argument.name == "language") + #expect(params.argument.value == "py") + + if case .prompt(let promptRef) = params.ref { + #expect(promptRef.name == "code_review") + } else { + Issue.record("Expected prompt reference") + } + + return .init( + completion: .init( + values: ["python", "pytorch", "pyside"], + total: 10, + hasMore: true + ) + ) + } + + try await server.start(transport: serverTransport) + let initResult = try await client.connect(transport: clientTransport) + + // Verify completions capability is advertised + #expect(initResult.capabilities.completions != nil) + + // Request completions + let completion = try await client.complete( + promptName: "code_review", + argumentName: "language", + argumentValue: "py" + ) + + #expect(completion.values == ["python", "pytorch", "pyside"]) + #expect(completion.total == 10) + #expect(completion.hasMore == true) + + await client.disconnect() + await server.stop() + } + + @Test("Client complete for resource argument") + func testClientCompleteForResource() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let client = Client(name: "TestClient", version: "1.0") + let server = Server( + name: "TestServer", + version: "1.0", + capabilities: .init(completions: .init()) + ) + + // Register handler for complete on server + await server.withMethodHandler(Complete.self) { params in + #expect(params.argument.name == "path") + #expect(params.argument.value == "/usr/") + + if case .resource(let resourceRef) = params.ref { + #expect(resourceRef.uri == "file:///{path}") + } else { + Issue.record("Expected resource reference") + } + + return .init( + completion: .init( + values: ["/usr/bin", "/usr/lib", "/usr/local"], + total: 3, + hasMore: false + ) + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Request completions for resource + let completion = try await client.complete( + resourceURI: "file:///{path}", + argumentName: "path", + argumentValue: "/usr/" + ) + + #expect(completion.values == ["/usr/bin", "/usr/lib", "/usr/local"]) + #expect(completion.total == 3) + #expect(completion.hasMore == false) + + await client.disconnect() + await server.stop() + } + + @Test("Client complete with context") + func testClientCompleteWithContext() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let client = Client(name: "TestClient", version: "1.0") + let server = Server( + name: "TestServer", + version: "1.0", + capabilities: .init(completions: .init()) + ) + + // Register handler for complete on server + await server.withMethodHandler(Complete.self) { params in + #expect(params.argument.name == "framework") + #expect(params.argument.value == "fla") + #expect(params.context != nil) + #expect(params.context?.arguments["language"] == "python") + + return .init( + completion: .init( + values: ["flask"], + total: 1, + hasMore: false + ) + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Request completions with context + let completion = try await client.complete( + promptName: "code_review", + argumentName: "framework", + argumentValue: "fla", + context: ["language": "python"] + ) + + #expect(completion.values == ["flask"]) + #expect(completion.total == 1) + #expect(completion.hasMore == false) + + await client.disconnect() + await server.stop() + } + + @Test("Client complete fails without completions capability") + func testClientCompleteFailsWithoutCapability() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let client = Client(name: "TestClient", version: "1.0", configuration: .strict) + let server = Server( + name: "TestServer", + version: "1.0", + capabilities: .init() // No completions capability + ) + + try await server.start(transport: serverTransport) + let initResult = try await client.connect(transport: clientTransport) + + // Verify completions capability is NOT advertised + #expect(initResult.capabilities.completions == nil) + + // Attempt to request completions should fail in strict mode + await #expect(throws: MCPError.self) { + try await client.complete( + promptName: "test", + argumentName: "arg", + argumentValue: "val" + ) + } + + await client.disconnect() + await server.stop() + } + + @Test("Empty completion values") + func testEmptyCompletionValues() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let client = Client(name: "TestClient", version: "1.0") + let server = Server( + name: "TestServer", + version: "1.0", + capabilities: .init(completions: .init()) + ) + + // Register handler that returns empty results + await server.withMethodHandler(Complete.self) { _ in + return .init( + completion: .init( + values: [], + total: 0, + hasMore: false + ) + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let completion = try await client.complete( + promptName: "test", + argumentName: "arg", + argumentValue: "xyz" + ) + + #expect(completion.values.isEmpty) + #expect(completion.total == 0) + #expect(completion.hasMore == false) + + await client.disconnect() + await server.stop() + } + + @Test("Maximum completion values (100 items)") + func testMaximumCompletionValues() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let client = Client(name: "TestClient", version: "1.0") + let server = Server( + name: "TestServer", + version: "1.0", + capabilities: .init(completions: .init()) + ) + + // Register handler that returns 100 items + await server.withMethodHandler(Complete.self) { _ in + let values = (1...100).map { "value\($0)" } + return .init( + completion: .init( + values: values, + total: 200, + hasMore: true + ) + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + let completion = try await client.complete( + promptName: "test", + argumentName: "arg", + argumentValue: "" + ) + + #expect(completion.values.count == 100) + #expect(completion.values.first == "value1") + #expect(completion.values.last == "value100") + #expect(completion.total == 200) + #expect(completion.hasMore == true) + + await client.disconnect() + await server.stop() + } + + @Test("Fuzzy matching completion scenario") + func testFuzzyMatchingScenario() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let client = Client(name: "TestClient", version: "1.0") + let server = Server( + name: "TestServer", + version: "1.0", + capabilities: .init(completions: .init()) + ) + + // Register handler that implements fuzzy matching + await server.withMethodHandler(Complete.self) { params in + let input = params.argument.value.lowercased() + let allLanguages = ["python", "perl", "php", "pascal", "prolog", "javascript", "java"] + + // Simple prefix matching + let matches = allLanguages.filter { $0.lowercased().hasPrefix(input) } + + return .init( + completion: .init( + values: matches, + total: matches.count, + hasMore: false + ) + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Test with "p" prefix + let completion1 = try await client.complete( + promptName: "language_selector", + argumentName: "language", + argumentValue: "p" + ) + #expect(completion1.values.count == 5) // python, perl, php, pascal, prolog + + // Test with "py" prefix + let completion2 = try await client.complete( + promptName: "language_selector", + argumentName: "language", + argumentValue: "py" + ) + #expect(completion2.values == ["python"]) + + // Test with "ja" prefix + let completion3 = try await client.complete( + promptName: "language_selector", + argumentName: "language", + argumentValue: "ja" + ) + #expect(completion3.values == ["javascript", "java"]) + + await client.disconnect() + await server.stop() + } +} diff --git a/Tests/MCPTests/ElicitationTests.swift b/Tests/MCPTests/ElicitationTests.swift new file mode 100644 index 0000000..1706440 --- /dev/null +++ b/Tests/MCPTests/ElicitationTests.swift @@ -0,0 +1,631 @@ +import Testing + +import class Foundation.JSONDecoder +import class Foundation.JSONEncoder + +@testable import MCP + +@Suite("Elicitation Tests") +struct ElicitationTests { + @Test("Request schema encoding and decoding") + func testSchemaCoding() throws { + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + let decoder = JSONDecoder() + + let schema = Elicitation.RequestSchema( + title: "Contact Information", + description: "Used to follow up after onboarding", + properties: [ + "name": [ + "type": "string", + "title": "Full Name", + "description": "Enter your legal name", + "minLength": 2, + "maxLength": 120, + ], + "email": [ + "type": "string", + "title": "Email Address", + "description": "Where we can reach you", + "format": "email", + ], + "age": [ + "type": "integer", + "minimum": 18, + "maximum": 110, + ], + "marketingOptIn": [ + "type": "boolean", + "title": "Marketing opt-in", + "default": false, + ], + ], + required: ["name", "email"] + ) + + let data = try encoder.encode(schema) + let decoded = try decoder.decode(Elicitation.RequestSchema.self, from: data) + + #expect(decoded == schema) + } + + @Test("Enumeration support") + func testEnumerationSupport() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let property: Value = [ + "type": "string", + "title": "Department", + "enum": ["engineering", "design", "product"], + "enumNames": ["Engineering", "Design", "Product"], + ] + + let data = try encoder.encode(property) + let decoded = try decoder.decode(Value.self, from: data) + + let object = decoded.objectValue + let enumeration = object?["enum"]?.arrayValue?.compactMap { $0.stringValue } + let enumNames = object?["enumNames"]?.arrayValue?.compactMap { $0.stringValue } + + #expect(enumeration == ["engineering", "design", "product"]) + #expect(enumNames == ["Engineering", "Design", "Product"]) + } + + @Test("CreateElicitation.Parameters coding") + func testParametersCoding() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let schema = Elicitation.RequestSchema( + properties: [ + "username": [ + "type": "string", + "minLength": 2, + "maxLength": 39, + ] + ], + required: ["username"] + ) + + let parameters = CreateElicitation.Parameters.form( + .init( + message: "Please share your GitHub username", + requestedSchema: schema, + _meta: Metadata(additionalFields: ["flow": "onboarding"]) + ) + ) + + let data = try encoder.encode(parameters) + let decoded = try decoder.decode(CreateElicitation.Parameters.self, from: data) + + #expect(decoded == parameters) + } + + @Test("CreateElicitation.Result coding") + func testResultCoding() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let result = CreateElicitation.Result( + action: .accept, + content: ["username": "octocat", "age": 30] + ) + + let data = try encoder.encode(result) + let decoded = try decoder.decode(CreateElicitation.Result.self, from: data) + + #expect(decoded == result) + } + + @Test("Client capabilities include elicitation") + func testClientCapabilitiesIncludeElicitation() throws { + let capabilities = Client.Capabilities( + elicitation: .init() + ) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(capabilities) + let decoded = try decoder.decode(Client.Capabilities.self, from: data) + + #expect(decoded == capabilities) + } + + @Test("Client elicitation handler registration") + func testClientElicitationHandlerRegistration() async throws { + let client = Client(name: "TestClient", version: "1.0") + + let handlerClient = await client.withElicitationHandler { parameters in + if case .form(let formParams) = parameters { + #expect(formParams.message == "Collect input") + } + return CreateElicitation.Result(action: .decline) + } + + #expect(handlerClient === client) + } +} + +@Suite("Elicitation 2025-11-25 Spec Tests") +struct Elicitation2025_11_25Tests { + @Test("URL mode parameters encoding and decoding") + func testURLModeParameters() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let params = CreateElicitation.Parameters.url( + .init( + message: "Please authenticate", + url: "https://example.com/auth", + elicitationId: "elicit-123" + ) + ) + + let data = try encoder.encode(params) + let decoded = try decoder.decode(CreateElicitation.Parameters.self, from: data) + + #expect(decoded == params) + } + + @Test("Form mode backward compatibility") + func testFormModeBackwardCompatibility() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let params = CreateElicitation.Parameters.form( + .init(message: "Enter your name", requestedSchema: .init()) + ) + + let data = try encoder.encode(params) + let decoded = try decoder.decode(CreateElicitation.Parameters.self, from: data) + + #expect(decoded == params) + } + + @Test("ElicitationCompleteNotification") + func testElicitationCompleteNotification() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let notification = ElicitationCompleteNotification.Parameters( + elicitationId: "elicit-456" + ) + + let data = try encoder.encode(notification) + let decoded = try decoder.decode( + ElicitationCompleteNotification.Parameters.self, from: data + ) + + #expect(decoded == notification) + } + + @Test("Client elicitation capabilities with sub-capabilities") + func testElicitationSubCapabilities() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let capabilities = Client.Capabilities( + elicitation: .init(form: .init(), url: .init()) + ) + + let data = try encoder.encode(capabilities) + let decoded = try decoder.decode(Client.Capabilities.self, from: data) + + #expect(decoded == capabilities) + } + + @Test("URLElicitationRequiredError encoding and decoding") + func testURLElicitationRequiredError() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let elicitationInfo = URLElicitationInfo( + elicitationId: "elicit-789", + url: "https://example.com/verify", + message: "Please verify your identity" + ) + + let error = MCPError.urlElicitationRequired( + message: "Authentication required", + elicitations: [elicitationInfo] + ) + + let data = try encoder.encode(error) + let decoded = try decoder.decode(MCPError.self, from: data) + + #expect(decoded == error) + } +} + +@Suite("Elicitation Integration Tests") +struct ElicitationIntegrationTests { + + @Test("Form-based elicitation flow") + func testFormElicitationFlow() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "FormTestServer", + version: "1.0" + ) + + let client = Client( + name: "FormTestClient", + version: "1.0", + capabilities: .init(elicitation: .init()) + ) + + // Register handler on client that validates parameters + await client.withElicitationHandler { parameters in + if case .form(let formParams) = parameters { + #expect(formParams.message == "Please enter your details") + #expect(formParams.requestedSchema.properties["email"] != nil) + #expect(formParams._meta?["flow"]?.stringValue == "onboarding") + + // Return accepted response + return CreateElicitation.Result( + action: .accept, + content: [ + "email": "user@example.com", + "name": "Test User" + ] + ) + } else { + Issue.record("Expected form parameters") + return CreateElicitation.Result(action: .decline) + } + } + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Server requests elicitation with form parameters + let schema = Elicitation.RequestSchema( + title: "User Details", + properties: [ + "email": ["type": "string", "format": "email"], + "name": ["type": "string", "minLength": 2] + ], + required: ["email"] + ) + + let result = try await server.requestElicitation( + message: "Please enter your details", + requestedSchema: schema, + _meta: Metadata(additionalFields: ["flow": "onboarding"]) + ) + + // Verify the response + #expect(result.action == .accept) + #expect(result.content?["email"]?.stringValue == "user@example.com") + #expect(result.content?["name"]?.stringValue == "Test User") + + await server.stop() + await client.disconnect() + } + + @Test("URL-based elicitation flow") + func testURLElicitationFlow() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "URLTestServer", + version: "1.0" + ) + + let client = Client( + name: "URLTestClient", + version: "1.0", + capabilities: .init(elicitation: .init(url: .init())) + ) + + // Register handler on client that validates URL parameters + await client.withElicitationHandler { parameters in + if case .url(let urlParams) = parameters { + #expect(urlParams.message == "Please authenticate") + #expect(urlParams.url == "https://example.com/auth") + #expect(urlParams.elicitationId.isEmpty == false) + #expect(urlParams._meta?["provider"]?.stringValue == "oauth") + + // Return accepted response + return CreateElicitation.Result( + action: .accept, + content: ["token": "auth-token-123"] + ) + } else { + Issue.record("Expected URL parameters") + return CreateElicitation.Result(action: .decline) + } + } + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Server requests elicitation with URL parameters + let result = try await server.requestElicitation( + message: "Please authenticate", + url: "https://example.com/auth", + elicitationId: "elicit-test-123", + _meta: Metadata(additionalFields: ["provider": "oauth"]) + ) + + // Verify the response + #expect(result.action == .accept) + #expect(result.content?["token"]?.stringValue == "auth-token-123") + + await server.stop() + await client.disconnect() + } + + @Test("Declined elicitation") + func testDeclinedElicitation() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "DeclineTestServer", + version: "1.0" + ) + + let client = Client( + name: "DeclineTestClient", + version: "1.0", + capabilities: .init(elicitation: .init()) + ) + + // Register handler that declines + await client.withElicitationHandler { _ in + CreateElicitation.Result(action: .decline) + } + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + let result = try await server.requestElicitation( + message: "Optional question", + requestedSchema: .init() + ) + + #expect(result.action == .decline) + #expect(result.content == nil) + + await server.stop() + await client.disconnect() + } + + @Test("Elicitation without handler fails") + func testElicitationWithoutHandlerFails() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ErrorTestServer", + version: "1.0" + ) + + let client = Client( + name: "ErrorTestClient", + version: "1.0" + ) + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Should throw an error because client doesn't have elicitation capability + await #expect(throws: MCPError.self) { + _ = try await server.requestElicitation( + message: "Test message", + requestedSchema: .init() + ) + } + + await server.stop() + await client.disconnect() + } + + @Test("Strict mode succeeds when client declares elicitation capability") + func testElicitationStrictCapabilitiesSuccess() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "StrictTestServer", + version: "1.0", + configuration: .strict + ) + + let client = Client( + name: "StrictTestClient", + version: "1.0", + capabilities: .init(elicitation: .init()), + configuration: .strict + ) + + // Register elicitation handler + await client.withElicitationHandler { _ in + CreateElicitation.Result( + action: .accept, + content: ["response": "strict mode success"] + ) + } + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Should succeed because client declares elicitation capability + let result = try await server.requestElicitation( + message: "Test message", + requestedSchema: .init() + ) + + #expect(result.action == .accept) + #expect(result.content?["response"]?.stringValue == "strict mode success") + + await server.stop() + await client.disconnect() + } + + @Test("Strict mode fails when client doesn't declare elicitation capability") + func testElicitationStrictCapabilitiesError() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "StrictTestServer", + version: "1.0", + configuration: .strict + ) + + // Client WITHOUT elicitation capability in strict mode + let client = Client( + name: "StrictTestClient", + version: "1.0", + capabilities: .init(), + configuration: .strict + ) + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Should fail because client doesn't declare elicitation capability in strict mode + await #expect(throws: MCPError.self) { + _ = try await server.requestElicitation( + message: "Test message", + requestedSchema: .init() + ) + } + + await server.stop() + await client.disconnect() + } + + @Test("Non-strict mode succeeds even without client capability declaration") + func testElicitationNonStrictCapabilities() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "NonStrictTestServer", + version: "1.0", + configuration: .default // Non-strict mode + ) + + // Client WITHOUT elicitation capability in non-strict mode + let client = Client( + name: "NonStrictTestClient", + version: "1.0", + capabilities: .init(), + configuration: .default + ) + + // Register elicitation handler anyway + await client.withElicitationHandler { _ in + CreateElicitation.Result( + action: .accept, + content: ["response": "non-strict mode success"] + ) + } + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Should succeed because server is in non-strict mode + let result = try await server.requestElicitation( + message: "Test message", + requestedSchema: .init() + ) + + #expect(result.action == .accept) + #expect(result.content?["response"]?.stringValue == "non-strict mode success") + + await server.stop() + await client.disconnect() + } + + @Test("Complex schema validation") + func testComplexSchemaValidation() async throws { + let schema = Elicitation.RequestSchema( + title: "User Profile", + description: "Complete user profile information", + properties: [ + "username": [ + "type": "string", + "minLength": 3, + "maxLength": 20, + "pattern": "^[a-zA-Z0-9_]+$" + ], + "email": [ + "type": "string", + "format": "email" + ], + "age": [ + "type": "integer", + "minimum": 18, + "maximum": 120 + ], + "preferences": [ + "type": "object", + "properties": [ + "theme": ["type": "string", "enum": ["light", "dark"]], + "notifications": ["type": "boolean"] + ] + ] + ], + required: ["username", "email"] + ) + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .prettyPrinted] + let decoder = JSONDecoder() + + let data = try encoder.encode(schema) + let decoded = try decoder.decode(Elicitation.RequestSchema.self, from: data) + + #expect(decoded == schema) + } + + @Test("Multiple elicitation requests in sequence") + func testSequentialElicitations() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "SequentialTestServer", + version: "1.0" + ) + + let client = Client( + name: "SequentialTestClient", + version: "1.0", + capabilities: .init(elicitation: .init()) + ) + + // Register handler that echoes the message + await client.withElicitationHandler { parameters in + if case .form(let formParams) = parameters { + return CreateElicitation.Result( + action: .accept, + content: ["echo": Value(stringLiteral: formParams.message)] + ) + } else { + return CreateElicitation.Result(action: .decline) + } + } + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Make multiple sequential requests + let result1 = try await server.requestElicitation(message: "First question", requestedSchema: .init()) + #expect(result1.action == .accept) + #expect(result1.content?["echo"]?.stringValue == "First question") + + let result2 = try await server.requestElicitation(message: "Second question", requestedSchema: .init()) + #expect(result2.action == .accept) + #expect(result2.content?["echo"]?.stringValue == "Second question") + + let result3 = try await server.requestElicitation(message: "Third question", requestedSchema: .init()) + #expect(result3.action == .accept) + #expect(result3.content?["echo"]?.stringValue == "Third question") + + await server.stop() + await client.disconnect() + } +} diff --git a/Tests/MCPTests/HTTPClientTransportTests.swift b/Tests/MCPTests/HTTPClientTransportTests.swift new file mode 100644 index 0000000..bd8c38f --- /dev/null +++ b/Tests/MCPTests/HTTPClientTransportTests.swift @@ -0,0 +1,2398 @@ +@preconcurrency import Foundation +import Logging +import Testing + +@testable import MCP + +#if canImport(FoundationNetworking) + import FoundationNetworking +#endif + +#if swift(>=6.1) + + // MARK: - Test trait + + /// A test trait that automatically manages the mock URL protocol handler for HTTP client transport tests. + struct HTTPClientTransportTestSetupTrait: TestTrait, TestScoping { + func provideScope( + for test: Test, testCase: Test.Case?, + performing function: @Sendable () async throws -> Void + ) async throws { + // Clear handler before test + await MockURLProtocol.requestHandlerStorage.clearHandler() + + // Execute the test + try await function() + + // Clear handler after test + await MockURLProtocol.requestHandlerStorage.clearHandler() + } + } + + extension Trait where Self == HTTPClientTransportTestSetupTrait { + static var httpClientTransportSetup: Self { Self() } + } + + // MARK: - Mock Handler Registry Actor + + actor RequestHandlerStorage { + var requestHandler: + (@Sendable (URLRequest) async throws -> (HTTPURLResponse, Data))? + private var callCounts: [URL: Int] = [:] + + func setHandler( + _ handler: @Sendable @escaping (URLRequest) async throws -> (HTTPURLResponse, Data) + ) { + requestHandler = handler + } + + func clearHandler() { + requestHandler = nil + callCounts = [:] + } + + func executeHandler(for request: URLRequest) async throws -> (HTTPURLResponse, Data) { + guard let handler = requestHandler else { + throw NSError( + domain: "MockURLProtocolError", code: 0, + userInfo: [ + NSLocalizedDescriptionKey: "No request handler set" + ]) + } + if let url = request.url { + callCounts[url, default: 0] += 1 + } + return try await handler(request) + } + + func callCount(for url: URL) -> Int { + callCounts[url, default: 0] + } + } + + // MARK: - Helper Methods + + extension URLRequest { + fileprivate func readBody() -> Data? { + if let httpBodyData = self.httpBody { + return httpBodyData + } + + guard let bodyStream = self.httpBodyStream else { return nil } + bodyStream.open() + defer { bodyStream.close() } + + let bufferSize: Int = 4096 + let buffer = UnsafeMutablePointer.allocate(capacity: bufferSize) + defer { buffer.deallocate() } + + var data = Data() + while bodyStream.hasBytesAvailable { + let bytesRead = bodyStream.read(buffer, maxLength: bufferSize) + data.append(buffer, count: bytesRead) + } + return data + } + } + + // MARK: - Mock URL Protocol + + final class MockURLProtocol: URLProtocol, @unchecked Sendable { + static let requestHandlerStorage = RequestHandlerStorage() + + static func setHandler( + _ handler: @Sendable @escaping (URLRequest) async throws -> (HTTPURLResponse, Data) + ) async { + await requestHandlerStorage.setHandler { request in + try await handler(request) + } + } + + func executeHandler(for request: URLRequest) async throws -> (HTTPURLResponse, Data) { + return try await Self.requestHandlerStorage.executeHandler(for: request) + } + + override class func canInit(with request: URLRequest) -> Bool { + return true + } + + override class func canonicalRequest(for request: URLRequest) -> URLRequest { + return request + } + + override func startLoading() { + Task { + do { + let (response, data) = try await self.executeHandler(for: request) + client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed) + client?.urlProtocol(self, didLoad: data) + client?.urlProtocolDidFinishLoading(self) + } catch { + client?.urlProtocol(self, didFailWithError: error) + } + } + } + + override func stopLoading() {} + + static func verifyCallCounts( + _ expected: [URL: Int], + sourceLocation: SourceLocation = #_sourceLocation + ) async { + for (url, expectedCount) in expected { + let actual = await requestHandlerStorage.callCount(for: url) + #expect( + actual == expectedCount, + "Expected \(expectedCount) call(s) to \(url.lastPathComponent), got \(actual)", + sourceLocation: sourceLocation + ) + } + } + } + + // MARK: - + + @Suite("HTTP Client Transport Tests", .serialized) + struct HTTPClientTransportTests { + let testEndpoint = URL(string: "https://localhost:8080/test")! + + @Test("Connect and Disconnect", .httpClientTransportSetup) + func testConnectAndDisconnect() 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() + await transport.disconnect() + } + + @Test("Send and Receive JSON Response", .httpClientTransportSetup) + func testSendAndReceiveJSON() 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 messageData = #"{"jsonrpc":"2.0","method":"initialize","id":1}"#.data(using: .utf8)! + let responseData = #"{"jsonrpc":"2.0","result":{},"id":1}"#.data(using: .utf8)! + + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + #expect(request.url == testEndpoint) + #expect(request.httpMethod == "POST") + #expect(request.readBody() == messageData) + #expect(request.value(forHTTPHeaderField: "Content-Type") == "application/json") + #expect( + request.value(forHTTPHeaderField: "Accept") + == "application/json, text/event-stream" + ) + + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "application/json"])! + return (response, responseData) + } + + try await transport.send(messageData) + + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + let receivedData = try await iterator.next() + + #expect(receivedData == responseData) + } + + @Test("Send and Receive Session ID", .httpClientTransportSetup) + func testSendAndReceiveSessionID() 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 messageData = #"{"jsonrpc":"2.0","method":"initialize","id":1}"#.data(using: .utf8)! + let newSessionID = "session-12345" + + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + #expect(request.value(forHTTPHeaderField: "MCP-Session-Id") == nil) + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [ + "Content-Type": "application/json", + "MCP-Session-Id": newSessionID, + ])! + return (response, Data()) + } + + try await transport.send(messageData) + + let storedSessionID = await transport.sessionID + #expect(storedSessionID == newSessionID) + } + + @Test("Send With Existing Session ID", .httpClientTransportSetup) + func testSendWithExistingSessionID() 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 initialSessionID = "existing-session-abc" + let firstMessageData = #"{"jsonrpc":"2.0","method":"initialize","id":1}"#.data( + using: .utf8)! + let secondMessageData = #"{"jsonrpc":"2.0","method":"ping","id":2}"#.data( + using: .utf8)! + + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + #expect(request.readBody() == firstMessageData) + #expect(request.value(forHTTPHeaderField: "MCP-Session-Id") == nil) + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [ + "Content-Type": "application/json", + "MCP-Session-Id": initialSessionID, + ])! + return (response, Data()) + } + try await transport.send(firstMessageData) + #expect(await transport.sessionID == initialSessionID) + + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + #expect(request.readBody() == secondMessageData) + #expect(request.value(forHTTPHeaderField: "MCP-Session-Id") == initialSessionID) + + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "application/json"])! + return (response, Data()) + } + try await transport.send(secondMessageData) + + #expect(await transport.sessionID == initialSessionID) + } + + @Test("HTTP 404 Not Found Error", .httpClientTransportSetup) + func testHTTPNotFoundError() async throws { + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let messageData = #"{"jsonrpc":"2.0","method":"test","id":3}"#.data(using: .utf8)! + + // Set up the handler BEFORE creating the transport + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 404, httpVersion: "HTTP/1.1", headerFields: nil)! + return (response, Data("Not Found".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 404") + } catch let error as MCPError { + guard case .internalError(let message) = error else { + Issue.record("Expected MCPError.internalError, got \(error)") + throw error + } + #expect(message?.contains("Endpoint not found") ?? false) + } catch { + Issue.record("Expected MCPError, got \(error)") + throw error + } + } + + @Test("HTTP 500 Server Error", .httpClientTransportSetup) + func testHTTPServerError() async throws { + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let messageData = #"{"jsonrpc":"2.0","method":"test","id":4}"#.data(using: .utf8)! + + // Set up the handler BEFORE creating the transport + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 500, httpVersion: "HTTP/1.1", headerFields: nil)! + return (response, Data("Server Error".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 500") + } catch let error as MCPError { + guard case .internalError(let message) = error else { + Issue.record("Expected MCPError.internalError, got \(error)") + throw error + } + #expect(message?.contains("Server error: 500") ?? false) + } catch { + Issue.record("Expected MCPError, got \(error)") + throw error + } + } + + @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":40}"#.data(using: .utf8)! + + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (_: 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 Without OAuth", .httpClientTransportSetup) + func testHTTPUnauthorizedErrorWithoutOAuth() async throws { + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let messageData = #"{"jsonrpc":"2.0","method":"test","id":41}"#.data(using: .utf8)! + + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (_: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, + statusCode: 401, + httpVersion: "HTTP/1.1", + headerFields: [ + "WWW-Authenticate": "Bearer scope=\"files:read\"" + ])! + return (response, Data()) + } + + 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 without OAuth") + } 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 Without OAuth", .httpClientTransportSetup) + func testHTTPForbiddenErrorWithoutOAuth() async throws { + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let messageData = #"{"jsonrpc":"2.0","method":"test","id":42}"#.data(using: .utf8)! + + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint] (_: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, + statusCode: 403, + httpVersion: "HTTP/1.1", + headerFields: [ + "WWW-Authenticate": + "Bearer error=\"insufficient_scope\", scope=\"files:write\"" + ])! + return (response, Data()) + } + + 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 without OAuth") + } 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("OAuth scope step-up retries after 403 insufficient_scope", .httpClientTransportSetup) + func testOAuthStepUpRetryAfter403InsufficientScope() async throws { + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let testEndpoint = URL(string: "https://localhost:8080/step-up")! + let resourceMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-protected-resource/step-up")! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + let finalResponseData = #"{"jsonrpc":"2.0","result":{"ok":true},"id":43}"#.data( + using: .utf8)! + + actor CallTracker { + var tokenCalls = 0 + + func nextTokenCall() -> Int { + tokenCalls += 1 + return tokenCalls + } + } + let tracker = CallTracker() + + await MockURLProtocol.requestHandlerStorage.setHandler { + [tracker, testEndpoint, resourceMetadataURL, asMetadataURL, tokenEndpointURL, finalResponseData] request in + guard let url = request.url else { + throw NSError( + domain: "MockURLProtocolError", + code: 0, + userInfo: [NSLocalizedDescriptionKey: "Missing request URL"]) + } + + switch url { + case testEndpoint: + switch request.value(forHTTPHeaderField: "Authorization") { + case nil: + let response = HTTPURLResponse( + url: url, + statusCode: 401, + httpVersion: "HTTP/1.1", + headerFields: [ + "WWW-Authenticate": + "Bearer resource_metadata=\"\(resourceMetadataURL.absoluteString)\", scope=\"files:read\"" + ])! + return (response, Data()) + + case "Bearer access-token-read": + let response = HTTPURLResponse( + url: url, + statusCode: 403, + httpVersion: "HTTP/1.1", + headerFields: [ + "WWW-Authenticate": + "Bearer error=\"insufficient_scope\", scope=\"files:write\", resource_metadata=\"\(resourceMetadataURL.absoluteString)\", error_description=\"Additional file write permission required\"" + ])! + return (response, Data()) + + case "Bearer access-token-read-write": + let response = HTTPURLResponse( + url: url, + statusCode: 200, + httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "application/json"])! + return (response, finalResponseData) + + default: + throw NSError( + domain: "MockURLProtocolError", + code: 0, + userInfo: [ + NSLocalizedDescriptionKey: + "Unexpected Authorization value: \(request.value(forHTTPHeaderField: "Authorization") ?? "")" + ]) + } + + case resourceMetadataURL: + let metadata = + #"{ "authorization_servers": ["https://localhost:8080/auth"], "scopes_supported": ["files:read","files:write"] }"# + .data(using: .utf8)! + let response = HTTPURLResponse( + url: url, + statusCode: 200, + httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "application/json"])! + return (response, metadata) + + case asMetadataURL: + let metadata = #"{ "issuer": "https://localhost:8080/auth", "token_endpoint": "https://localhost:8080/oauth/token" }"# + .data(using: .utf8)! + let response = HTTPURLResponse( + url: url, + statusCode: 200, + httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "application/json"])! + return (response, metadata) + + case tokenEndpointURL: + let tokenCall = await tracker.nextTokenCall() + let body = String(data: request.readBody() ?? Data(), encoding: .utf8) ?? "" + #expect(body.contains("grant_type=client_credentials")) + #expect(body.contains("client_id=test-client")) + #expect(body.contains("resource=https%3A%2F%2Flocalhost%3A8080%2Fstep-up")) + + if tokenCall == 1 { + #expect(body.contains("scope=files%3Aread")) + let tokenResponse = + #"{ "access_token": "access-token-read", "token_type": "Bearer", "expires_in": 3600 }"# + .data(using: .utf8)! + let response = HTTPURLResponse( + url: url, + statusCode: 200, + httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "application/json"])! + return (response, tokenResponse) + } + + #expect(tokenCall == 2) + #expect(body.contains("scope=files%3Aread%20files%3Awrite")) + let tokenResponse = + #"{ "access_token": "access-token-read-write", "token_type": "Bearer", "expires_in": 3600 }"# + .data(using: .utf8)! + let response = HTTPURLResponse( + url: url, + statusCode: 200, + httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "application/json"])! + return (response, tokenResponse) + + default: + throw NSError( + domain: "MockURLProtocolError", + code: 0, + userInfo: [ + NSLocalizedDescriptionKey: + "Unexpected URL: \(url.absoluteString)" + ]) + } + } + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: false, + authorizer: OAuthAuthorizer(configuration: .init(authentication: .none(clientID: "test-client"))), + logger: nil + ) + + try await transport.connect() + let messageData = #"{"jsonrpc":"2.0","method":"ping","id":43}"#.data(using: .utf8)! + try await transport.send(messageData) + + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + let received = try await iterator.next() + #expect(received == finalResponseData) + #expect(await tracker.tokenCalls == 2) + + await transport.disconnect() + } + + @Test("OAuth scope upgrade tracking is scoped per operation", .httpClientTransportSetup) + func testOAuthScopeUpgradeTrackingPerOperation() async throws { + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let testEndpoint = URL(string: "https://localhost:8080/operation-tracking")! + let resourceMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-protected-resource/operation-tracking")! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + let finalResponseData = #"{"jsonrpc":"2.0","result":{"ok":true},"id":62}"#.data( + using: .utf8)! + + actor CallTracker { + var tokenCalls = 0 + var opAForbiddenCalls = 0 + var opBForbiddenCalls = 0 + + func nextTokenCall() -> Int { + tokenCalls += 1 + return tokenCalls + } + + func incrementOpAForbiddenCalls() { + opAForbiddenCalls += 1 + } + + func incrementOpBForbiddenCalls() { + opBForbiddenCalls += 1 + } + } + let tracker = CallTracker() + + await MockURLProtocol.requestHandlerStorage.setHandler { + [tracker, testEndpoint, resourceMetadataURL, asMetadataURL, tokenEndpointURL, finalResponseData] request in + guard let url = request.url else { + throw NSError( + domain: "MockURLProtocolError", + code: 0, + userInfo: [NSLocalizedDescriptionKey: "Missing request URL"]) + } + + switch url { + case testEndpoint: + let body = String(data: request.readBody() ?? Data(), encoding: .utf8) ?? "" + let isOperationA = body.contains(#""method":"tools/callA""#) + let isOperationB = body.contains(#""method":"tools/callB""#) + let authorization = request.value(forHTTPHeaderField: "Authorization") + + if authorization == nil { + let response = HTTPURLResponse( + url: url, + statusCode: 401, + httpVersion: "HTTP/1.1", + headerFields: [ + "WWW-Authenticate": + "Bearer resource_metadata=\"\(resourceMetadataURL.absoluteString)\", scope=\"files:read\"" + ])! + return (response, Data()) + } + + if isOperationA { + if authorization == "Bearer access-token-read" + || authorization == "Bearer access-token-read-write" + { + await tracker.incrementOpAForbiddenCalls() + let response = HTTPURLResponse( + url: url, + statusCode: 403, + httpVersion: "HTTP/1.1", + headerFields: [ + "WWW-Authenticate": + "Bearer error=\"insufficient_scope\", scope=\"files:write\", resource_metadata=\"\(resourceMetadataURL.absoluteString)\"" + ])! + return (response, Data()) + } + + throw NSError( + domain: "MockURLProtocolError", + code: 0, + userInfo: [ + NSLocalizedDescriptionKey: + "Unexpected Authorization for opA: \(authorization ?? "")" + ]) + } + + if isOperationB { + if authorization == "Bearer access-token-read-write" { + await tracker.incrementOpBForbiddenCalls() + let response = HTTPURLResponse( + url: url, + statusCode: 403, + httpVersion: "HTTP/1.1", + headerFields: [ + "WWW-Authenticate": + "Bearer error=\"insufficient_scope\", scope=\"files:write\", resource_metadata=\"\(resourceMetadataURL.absoluteString)\"" + ])! + return (response, Data()) + } + + if authorization == "Bearer access-token-opb" { + let response = HTTPURLResponse( + url: url, + statusCode: 200, + httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "application/json"])! + return (response, finalResponseData) + } + + throw NSError( + domain: "MockURLProtocolError", + code: 0, + userInfo: [ + NSLocalizedDescriptionKey: + "Unexpected Authorization for opB: \(authorization ?? "")" + ]) + } + + throw NSError( + domain: "MockURLProtocolError", + code: 0, + userInfo: [ + NSLocalizedDescriptionKey: + "Unexpected request body: \(body)" + ]) + + case resourceMetadataURL: + let metadata = + #"{ "authorization_servers": ["https://localhost:8080/auth"], "scopes_supported": ["files:read","files:write"] }"# + .data(using: .utf8)! + let response = HTTPURLResponse( + url: url, + statusCode: 200, + httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "application/json"])! + return (response, metadata) + + case asMetadataURL: + let metadata = #"{ "issuer": "https://localhost:8080/auth", "token_endpoint": "https://localhost:8080/oauth/token" }"# + .data(using: .utf8)! + let response = HTTPURLResponse( + url: url, + statusCode: 200, + httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "application/json"])! + return (response, metadata) + + case tokenEndpointURL: + let tokenCall = await tracker.nextTokenCall() + let body = String(data: request.readBody() ?? Data(), encoding: .utf8) ?? "" + #expect(body.contains("grant_type=client_credentials")) + #expect(body.contains("client_id=test-client")) + #expect( + body.contains( + "resource=https%3A%2F%2Flocalhost%3A8080%2Foperation-tracking") + ) + + switch tokenCall { + case 1: + #expect(body.contains("scope=files%3Aread")) + let tokenResponse = + #"{ "access_token": "access-token-read", "token_type": "Bearer", "expires_in": 3600 }"# + .data(using: .utf8)! + let response = HTTPURLResponse( + url: url, + statusCode: 200, + httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "application/json"])! + return (response, tokenResponse) + + case 2: + #expect(body.contains("scope=files%3Aread%20files%3Awrite")) + let tokenResponse = + #"{ "access_token": "access-token-read-write", "token_type": "Bearer", "expires_in": 3600 }"# + .data(using: .utf8)! + let response = HTTPURLResponse( + url: url, + statusCode: 200, + httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "application/json"])! + return (response, tokenResponse) + + case 3: + #expect(body.contains("scope=files%3Aread%20files%3Awrite")) + let tokenResponse = + #"{ "access_token": "access-token-opb", "token_type": "Bearer", "expires_in": 3600 }"# + .data(using: .utf8)! + let response = HTTPURLResponse( + url: url, + statusCode: 200, + httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "application/json"])! + return (response, tokenResponse) + + default: + throw NSError( + domain: "MockURLProtocolError", + code: 0, + userInfo: [ + NSLocalizedDescriptionKey: + "Unexpected token call count: \(tokenCall)" + ]) + } + + default: + throw NSError( + domain: "MockURLProtocolError", + code: 0, + userInfo: [ + NSLocalizedDescriptionKey: + "Unexpected URL: \(url.absoluteString)" + ]) + } + } + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: false, + authorizer: OAuthAuthorizer(configuration: .init( + authentication: .none(clientID: "test-client"), + retryPolicy: .init(maxAuthorizationAttempts: 8, maxScopeUpgradeAttempts: 1) + )), + logger: nil + ) + + try await transport.connect() + + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + + let operationAData = #"{"jsonrpc":"2.0","method":"tools/callA","id":61}"#.data( + using: .utf8)! + do { + try await transport.send(operationAData) + Issue.record("Expected operation A to fail after scope-upgrade retry limit") + } 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 + } + + let operationBData = #"{"jsonrpc":"2.0","method":"tools/callB","id":62}"#.data( + using: .utf8)! + try await transport.send(operationBData) + + let received = try await iterator.next() + #expect(received == finalResponseData) + + #expect(await tracker.tokenCalls == 3) + #expect(await tracker.opAForbiddenCalls == 2) + #expect(await tracker.opBForbiddenCalls == 1) + + await transport.disconnect() + } + + @Test("Session Expired Error (404 with Session ID)", .httpClientTransportSetup) + func testSessionExpiredError() async throws { + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let initialSessionID = "expired-session-xyz" + let firstMessageData = #"{"jsonrpc":"2.0","method":"initialize","id":1}"#.data( + using: .utf8)! + let secondMessageData = #"{"jsonrpc":"2.0","method":"ping","id":2}"#.data( + using: .utf8)! + + // Set up the first handler BEFORE creating the transport + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint, initialSessionID] (request: URLRequest) in + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [ + "Content-Type": "application/json", + "MCP-Session-Id": initialSessionID, + ])! + return (response, Data()) + } + + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: false, + logger: nil + ) + try await transport.connect() + + try await transport.send(firstMessageData) + #expect(await transport.sessionID == initialSessionID) + + // 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) + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 404, httpVersion: "HTTP/1.1", headerFields: nil)! + return (response, Data("Not Found".utf8)) + } + + do { + try await transport.send(secondMessageData) + Issue.record("Expected send to throw session expired error") + } catch let error as MCPError { + guard case .internalError(let message) = error else { + Issue.record("Expected MCPError.internalError, got \(error)") + throw error + } + #expect(message?.contains("Session expired") ?? false) + #expect(await transport.sessionID == nil) + } catch { + Issue.record("Expected MCPError, got \(error)") + throw error + } + } + + // Skip SSE tests on platforms that don't support streaming + #if !canImport(FoundationNetworking) + @Test("Receive Server-Sent Event (SSE)", .httpClientTransportSetup) + func testReceiveSSE() 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\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: [ + "Content-Type": "text/plain", + "Mcp-Session-Id": "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: "MCP-Session-Id") == "test-session-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 + } + + 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: [ + "Content-Type": "text/plain", + "Mcp-Session-Id": "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: "MCP-Session-Id") == "test-session-123") + + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "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: ["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") + } + + 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(text: "42", annotations: nil, _meta: nil)]) + 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: ["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)" + ]) + } + } + + // 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 let .text(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: ["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 + }, + 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() + } + + @Test("OAuth client credentials performs discovery and retries after 401", .httpClientTransportSetup) + func testOAuthClientCredentialsRetryAfter401() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage.configureOAuthClientCredentialsRetryAfter401() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: scenario.streaming, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + try await transport.send(scenario.messageData) + + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + let received = try await iterator.next() + #expect(received == scenario.expectedResponseData) + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test("OAuth scope selection falls back to scopes_supported when challenge scope is absent", .httpClientTransportSetup) + func testOAuthScopeSelectionFallsBackToScopesSupportedWhenChallengeScopeMissing() + async throws + { + let scenario = await MockURLProtocol.requestHandlerStorage.configureOAuthScopeSelectionFallsBackToScopesSupported() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: scenario.streaming, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + try await transport.send(scenario.messageData) + + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + let received = try await iterator.next() + #expect(received == scenario.expectedResponseData) + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test("OAuth omits scope parameter when challenge scope and scopes_supported are unavailable", .httpClientTransportSetup) + func testOAuthScopeSelectionOmitsScopeWhenNoHintsAvailable() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage.configureOAuthScopeOmittedWhenNoHints() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: scenario.streaming, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + try await transport.send(scenario.messageData) + + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + let received = try await iterator.next() + #expect(received == scenario.expectedResponseData) + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test("OAuth includes canonical resource in both authorization and token requests", .httpClientTransportSetup) + func testOAuthResourceParameterIncludedInAuthorizationAndTokenRequests() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage.configureOAuthResourceParameterInAuthorizationAndToken() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: scenario.streaming, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + try await transport.send(scenario.messageData) + + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + let received = try await iterator.next() + #expect(received == scenario.expectedResponseData) + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test( + "OAuth rejects authorization when AS metadata omits code_challenge_methods_supported", + .httpClientTransportSetup + ) + func testOAuthRejectsAuthorizationWithoutPKCEMetadata() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage.configureOAuthRejectsAuthorizationWithoutPKCEMetadata() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: false, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + + do { + try await transport.send(scenario.messageData) + Issue.record("Expected send to throw an error") + } catch let error as MCPError { + guard case .internalError(let detail) = error else { + Issue.record("Expected MCPError.internalError, got \(error)") + throw error + } + #expect(detail?.contains(scenario.expectedErrorSubstring!) == true) + for unexpected in scenario.unexpectedErrorSubstrings { + #expect(detail?.contains(unexpected) == false) + } + } catch { + Issue.record("Expected MCPError, got \(error)") + throw error + } + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test( + "OAuth rejects authorization when AS metadata lacks PKCE S256 support", + .httpClientTransportSetup + ) + func testOAuthRejectsAuthorizationWithoutS256PKCE() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage.configureOAuthRejectsAuthorizationWithoutS256PKCE() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: false, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + + do { + try await transport.send(scenario.messageData) + Issue.record("Expected send to throw an error") + } catch let error as MCPError { + guard case .internalError(let detail) = error else { + Issue.record("Expected MCPError.internalError, got \(error)") + throw error + } + #expect(detail?.contains(scenario.expectedErrorSubstring!) == true) + for unexpected in scenario.unexpectedErrorSubstrings { + #expect(detail?.contains(unexpected) == false) + } + } catch { + Issue.record("Expected MCPError, got \(error)") + throw error + } + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test( + "OAuth rejects authorization response redirect URI mismatch", + .httpClientTransportSetup + ) + func testOAuthRejectsAuthorizationResponseRedirectMismatch() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage.configureOAuthRejectsAuthorizationResponseRedirectMismatch() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: false, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + + do { + try await transport.send(scenario.messageData) + Issue.record("Expected send to throw an error") + } catch let error as MCPError { + guard case .internalError(let detail) = error else { + Issue.record("Expected MCPError.internalError, got \(error)") + throw error + } + #expect(detail?.contains(scenario.expectedErrorSubstring!) == true) + for unexpected in scenario.unexpectedErrorSubstrings { + #expect(detail?.contains(unexpected) == false) + } + } catch { + Issue.record("Expected MCPError, got \(error)") + throw error + } + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test( + "OAuth rejects authorization response state mismatch", + .httpClientTransportSetup + ) + func testOAuthRejectsAuthorizationResponseStateMismatch() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage.configureOAuthRejectsAuthorizationResponseStateMismatch() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: false, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + + do { + try await transport.send(scenario.messageData) + Issue.record("Expected send to throw an error") + } catch let error as MCPError { + guard case .internalError(let detail) = error else { + Issue.record("Expected MCPError.internalError, got \(error)") + throw error + } + #expect(detail?.contains(scenario.expectedErrorSubstring!) == true) + for unexpected in scenario.unexpectedErrorSubstrings { + #expect(detail?.contains(unexpected) == false) + } + } catch { + Issue.record("Expected MCPError, got \(error)") + throw error + } + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test("OAuth sends access token only via Authorization header", .httpClientTransportSetup) + func testOAuthDoesNotSendAccessTokenInBodyOrQuery() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage.configureOAuthAccessTokenOnlyViaAuthorizationHeader() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: scenario.streaming, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + try await transport.send(scenario.messageData) + + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + let received = try await iterator.next() + #expect(received == scenario.expectedResponseData) + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test("OAuth sends Bearer Authorization header on every request in a logical session", .httpClientTransportSetup) + func testOAuthUsesAuthorizationHeaderForEveryRequestInSession() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage.configureOAuthAuthorizationHeaderForEveryRequestInSession() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: false, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + + try await transport.send(scenario.messageData) + let receivedFirst = try await iterator.next() + #expect(receivedFirst == scenario.expectedResponseData) + + try await transport.send(scenario.secondMessageData!) + let receivedSecond = try await iterator.next() + #expect(receivedSecond == scenario.secondExpectedResponseData) + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test("OAuth streaming GET requests use Authorization header and not query token", .httpClientTransportSetup) + func testOAuthStreamingGETUsesAuthorizationHeaderOnly() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage.configureOAuthStreamingGETUsesAuthorizationHeaderOnly() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: scenario.streaming, + sseInitializationTimeout: scenario.sseInitializationTimeout!, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + try await transport.send(scenario.messageData) + + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + let received = try await iterator.next() + #expect(received == scenario.expectedResponseData) + + await transport.disconnect() + } + + @Test("OAuth rejects non-Bearer token_type for MCP resource requests", .httpClientTransportSetup) + func testOAuthRejectsNonBearerTokenType() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage.configureOAuthRejectsNonBearerTokenType() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: false, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + + do { + try await transport.send(scenario.messageData) + Issue.record("Expected send to throw an error") + } catch let error as MCPError { + guard case .internalError(let detail) = error else { + Issue.record("Expected MCPError.internalError, got \(error)") + throw error + } + #expect(detail?.contains(scenario.expectedErrorSubstring!) == true) + for unexpected in scenario.unexpectedErrorSubstrings { + #expect(detail?.contains(unexpected) == false) + } + } catch { + Issue.record("Expected MCPError, got \(error)") + throw error + } + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test("OAuth token endpoint failures redact raw response body", .httpClientTransportSetup) + func testOAuthTokenEndpointFailureRedactsResponseBody() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage.configureOAuthTokenEndpointFailureRedactsResponseBody() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: false, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + + do { + try await transport.send(scenario.messageData) + Issue.record("Expected send to throw an error") + } catch let error as MCPError { + guard case .internalError(let detail) = error else { + Issue.record("Expected MCPError.internalError, got \(error)") + throw error + } + #expect(detail?.contains(scenario.expectedErrorSubstring!) == true) + for unexpected in scenario.unexpectedErrorSubstrings { + #expect(detail?.contains(unexpected) == false) + } + } catch { + Issue.record("Expected MCPError, got \(error)") + throw error + } + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test("OAuth rejects non-HTTPS token endpoint from AS metadata", .httpClientTransportSetup) + func testOAuthRejectsNonHTTPSTokenEndpoint() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage.configureOAuthRejectsNonHTTPSTokenEndpoint() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: false, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + + do { + try await transport.send(scenario.messageData) + Issue.record("Expected send to throw an error") + } catch let error as MCPError { + guard case .internalError(let detail) = error else { + Issue.record("Expected MCPError.internalError, got \(error)") + throw error + } + #expect(detail?.contains(scenario.expectedErrorSubstring!) == true) + for unexpected in scenario.unexpectedErrorSubstrings { + #expect(detail?.contains(unexpected) == false) + } + } catch { + Issue.record("Expected MCPError, got \(error)") + throw error + } + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test( + "OAuth allows loopback http authorization server endpoints when explicitly enabled", + .httpClientTransportSetup + ) + func testOAuthAllowsLoopbackHTTPAuthorizationServerEndpointsWhenEnabled() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage.configureOAuthAllowsLoopbackHTTPAuthorizationServerEndpoints() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: false, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + try await transport.send(scenario.messageData) + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test("OAuth accessTokenProvider receives SDK discovery context", .httpClientTransportSetup) + func testOAuthAccessTokenProviderReceivesDiscoveryContext() async throws { + let (scenario, providerTracker) = await MockURLProtocol.requestHandlerStorage + .configureOAuthAccessTokenProviderReceivesDiscoveryContext() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: false, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + try await transport.send(scenario.messageData) + + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + let received = try await iterator.next() + #expect(received == scenario.expectedResponseData) + + let capturedContext = await providerTracker.capturedContext + #expect(capturedContext != nil) + #expect(capturedContext?.statusCode == 401) + #expect(capturedContext?.endpoint == scenario.testEndpoint) + #expect(capturedContext?.resource == scenario.testEndpoint) + #expect(capturedContext?.authorizationServer == URL(string: "https://localhost:8080/auth")) + #expect(capturedContext?.tokenEndpoint == URL(string: "https://localhost:8080/oauth/token")) + #expect(capturedContext?.challengedScope == "files:read files:write") + #expect(Set(capturedContext?.scopesSupported ?? []) == Set(["files:read", "files:write"])) + #expect(capturedContext?.requestedScopes == Set(["files:read", "files:write"])) + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test("OAuth discovery uses resource_metadata URL from WWW-Authenticate when present", .httpClientTransportSetup) + func testOAuthDiscoveryUsesHeaderResourceMetadataWhenPresent() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage + .configureOAuthDiscoveryUsesHeaderResourceMetadata() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: false, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + try await transport.send(scenario.messageData) + + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + let received = try await iterator.next() + #expect(received == scenario.expectedResponseData) + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test("OAuth discovery falls back to well-known metadata URLs in required order", .httpClientTransportSetup) + func testOAuthDiscoveryFallbackWellKnownOrder() async throws { + let (scenario, tracker) = await MockURLProtocol.requestHandlerStorage + .configureOAuthDiscoveryFallbackWellKnownOrder() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: false, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + try await transport.send(scenario.messageData) + + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + let received = try await iterator.next() + #expect(received == scenario.expectedResponseData) + + let metadataRequests = await tracker.requests + let fallbackPathMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-protected-resource/public/mcp")! + let fallbackRootMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-protected-resource")! + #expect(metadataRequests == [fallbackPathMetadataURL, fallbackRootMetadataURL]) + + await transport.disconnect() + } + + @Test("OAuth discovery fails when protected resource metadata is unavailable", .httpClientTransportSetup) + func testOAuthDiscoveryFailsWhenMetadataUnavailable() async throws { + let (scenario, tracker) = await MockURLProtocol.requestHandlerStorage + .configureOAuthDiscoveryFailsWhenMetadataUnavailable() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: false, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + + do { + try await transport.send(scenario.messageData) + Issue.record("Expected send to fail when PRM discovery fails") + } catch let error as MCPError { + guard case .internalError(let detail) = error else { + Issue.record("Expected MCPError.internalError, got \(error)") + throw error + } + #expect(detail?.contains(scenario.expectedErrorSubstring!) == true) + } catch { + Issue.record("Expected MCPError, got \(error)") + throw error + } + + let protectedResourceMetadataRequests = await tracker.requests + let fallbackPathMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-protected-resource/public/mcp")! + let fallbackRootMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-protected-resource")! + #expect(protectedResourceMetadataRequests == [fallbackPathMetadataURL, fallbackRootMetadataURL]) + + await transport.disconnect() + } + + @Test("OAuth authorization server metadata discovery tries path issuer URLs in RFC order", .httpClientTransportSetup) + func testOAuthAuthorizationServerMetadataDiscoveryOrderForPathIssuer() async throws { + let (scenario, tracker) = await MockURLProtocol.requestHandlerStorage + .configureOAuthASMetadataDiscoveryOrderForPathIssuer() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: false, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + try await transport.send(scenario.messageData) + + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + let received = try await iterator.next() + #expect(received == scenario.expectedResponseData) + + let requests = await tracker.requests + let asMetadataOAuthInsertedURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/tenant1")! + let asMetadataOIDCInsertedURL = URL( + string: "https://localhost:8080/.well-known/openid-configuration/tenant1")! + let asMetadataOIDCAppendedURL = URL( + string: "https://localhost:8080/tenant1/.well-known/openid-configuration")! + #expect(requests == [asMetadataOAuthInsertedURL, asMetadataOIDCInsertedURL, asMetadataOIDCAppendedURL]) + + await transport.disconnect() + } + + @Test("OAuth authorization server metadata discovery tries root issuer URLs in RFC order", .httpClientTransportSetup) + func testOAuthAuthorizationServerMetadataDiscoveryOrderForRootIssuer() async throws { + let (scenario, tracker) = await MockURLProtocol.requestHandlerStorage + .configureOAuthASMetadataDiscoveryOrderForRootIssuer() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: false, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + try await transport.send(scenario.messageData) + + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + let received = try await iterator.next() + #expect(received == scenario.expectedResponseData) + + let requests = await tracker.requests + let asMetadataOAuthURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server")! + let asMetadataOIDCURL = URL( + string: "https://localhost:8080/.well-known/openid-configuration")! + #expect(requests == [asMetadataOAuthURL, asMetadataOIDCURL]) + + await transport.disconnect() + } + + @Test("OAuth registration prefers CIMD when AS advertises support", .httpClientTransportSetup) + func testOAuthRegistrationPrefersCIMDWhenAdvertised() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage + .configureOAuthRegistrationPrefersCIMDWhenAdvertised() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: false, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + try await transport.send(scenario.messageData) + + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + let received = try await iterator.next() + #expect(received == scenario.expectedResponseData) + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test("OAuth pre-registration uses static client credentials without dynamic registration", .httpClientTransportSetup) + func testOAuthPreRegistrationUsesStaticCredentialsWithoutDynamicRegistration() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage + .configureOAuthPreRegistrationUsesStaticCredentials() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: false, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + try await transport.send(scenario.messageData) + + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + let received = try await iterator.next() + #expect(received == scenario.expectedResponseData) + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test("OAuth registration falls back to dynamic registration when CIMD is not advertised", .httpClientTransportSetup) + func testOAuthRegistrationFallsBackToDynamicRegistrationWhenCIMDNotAdvertised() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage + .configureOAuthRegistrationFallsBackToDynamicRegistrationCIMDNotAdvertised() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: false, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + try await transport.send(scenario.messageData) + + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + let received = try await iterator.next() + #expect(received == scenario.expectedResponseData) + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test("OAuth registration falls back to dynamic registration when CIMD capability is absent", .httpClientTransportSetup) + func testOAuthRegistrationFallsBackToDynamicRegistrationWhenCIMDCapabilityMissing() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage + .configureOAuthRegistrationFallsBackToDynamicRegistrationCIMDCapabilityMissing() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: false, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + try await transport.send(scenario.messageData) + + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + let received = try await iterator.next() + #expect(received == scenario.expectedResponseData) + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test("OAuth registration surfaces actionable error when no supported mechanism is available", .httpClientTransportSetup) + func testOAuthRegistrationMissingMechanismReturnsActionableError() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage + .configureOAuthRegistrationMissingMechanismReturnsActionableError() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: false, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + + do { + try await transport.send(scenario.messageData) + Issue.record("Expected send to throw an error") + } catch let error as MCPError { + guard case .internalError(let detail) = error else { + Issue.record("Expected MCPError.internalError, got \(error)") + throw error + } + #expect(detail?.contains(scenario.expectedErrorSubstring!) == true) + for unexpected in scenario.unexpectedErrorSubstrings { + #expect(detail?.contains(unexpected) == false) + } + } catch { + Issue.record("Expected MCPError, got \(error)") + throw error + } + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test("OAuth CIMD rejects non-HTTPS client_id URL when AS advertises support", .httpClientTransportSetup) + func testOAuthCIMDRejectsNonHTTPSClientIDURL() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage + .configureOAuthCIMDRejectsNonHTTPSClientIDURL() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: false, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + + do { + try await transport.send(scenario.messageData) + Issue.record("Expected send to throw an error") + } catch let error as MCPError { + guard case .internalError(let detail) = error else { + Issue.record("Expected MCPError.internalError, got \(error)") + throw error + } + #expect(detail?.contains(scenario.expectedErrorSubstring!) == true) + for unexpected in scenario.unexpectedErrorSubstrings { + #expect(detail?.contains(unexpected) == false) + } + } catch { + Issue.record("Expected MCPError, got \(error)") + throw error + } + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test("OAuth rejects insecure MCP endpoint URL", .httpClientTransportSetup) + func testOAuthRejectsInsecureMCPEndpointURL() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage + .configureOAuthRejectsInsecureMCPEndpointURL() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: false, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + + do { + try await transport.send(scenario.messageData) + Issue.record("Expected send to throw an error") + } catch let error as MCPError { + guard case .internalError(let detail) = error else { + Issue.record("Expected MCPError.internalError, got \(error)") + throw error + } + #expect(detail?.contains(scenario.expectedErrorSubstring!) == true) + for unexpected in scenario.unexpectedErrorSubstrings { + #expect(detail?.contains(unexpected) == false) + } + } catch { + Issue.record("Expected MCPError, got \(error)") + throw error + } + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test("OAuth rejects non-loopback http redirect URI", .httpClientTransportSetup) + func testOAuthRejectsNonLoopbackHTTPRedirectURI() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage + .configureOAuthRejectsNonLoopbackHTTPRedirectURI() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: false, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + + do { + try await transport.send(scenario.messageData) + Issue.record("Expected send to throw an error") + } catch let error as MCPError { + guard case .internalError(let detail) = error else { + Issue.record("Expected MCPError.internalError, got \(error)") + throw error + } + #expect(detail?.contains(scenario.expectedErrorSubstring!) == true) + for unexpected in scenario.unexpectedErrorSubstrings { + #expect(detail?.contains(unexpected) == false) + } + } catch { + Issue.record("Expected MCPError, got \(error)") + throw error + } + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test("OAuth PRM cache is invalidated when resource_metadata URL changes between challenges", .httpClientTransportSetup) + func testOAuthPRMCacheInvalidatedOnResourceMetadataURLChange() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage + .configureOAuthPRMCacheInvalidatedOnResourceMetadataURLChange() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: false, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + + try await transport.send(scenario.messageData) + let receivedFirst = try await iterator.next() + #expect(receivedFirst == scenario.expectedResponseData) + + try await transport.send(scenario.secondMessageData!) + let receivedSecond = try await iterator.next() + #expect(receivedSecond == scenario.secondExpectedResponseData) + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test( + "OAuth token request uses PRM resource field as resource indicator", + .httpClientTransportSetup) + func testOAuthResourceUsesPRMResourceField() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage + .configureOAuthResourceUsesPRMResourceField() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: scenario.streaming, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + try await transport.send(scenario.messageData) + + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + let received = try await iterator.next() + #expect(received == scenario.expectedResponseData) + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test( + "OAuth tries second authorization server when first returns no metadata", + .httpClientTransportSetup) + func testOAuthSecondAuthorizationServerTriedWhenFirstFails() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage + .configureOAuthSecondAuthorizationServerTriedWhenFirstFails() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: scenario.streaming, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + try await transport.send(scenario.messageData) + + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + let received = try await iterator.next() + #expect(received == scenario.expectedResponseData) + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test( + "OAuth re-registers client after client_secret_expires_at has passed", + .httpClientTransportSetup) + func testOAuthReRegistersAfterClientSecretExpiry() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage + .configureOAuthReRegistersAfterClientSecretExpiry() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: false, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + + try await transport.send(scenario.messageData) + let receivedFirst = try await iterator.next() + #expect(receivedFirst == scenario.expectedResponseData) + + try await transport.send(scenario.secondMessageData!) + let receivedSecond = try await iterator.next() + #expect(receivedSecond == scenario.secondExpectedResponseData) + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test( + "OAuth skips AS metadata with wrong issuer and uses next URL variant", + .httpClientTransportSetup) + func testOAuthIssuerMismatchTriesNextURLVariant() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage + .configureOAuthIssuerMismatchTriesNextURLVariant() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: scenario.streaming, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + try await transport.send(scenario.messageData) + + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + let received = try await iterator.next() + #expect(received == scenario.expectedResponseData) + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test( + "OAuth proactively refreshes token when within proactive refresh window", + .httpClientTransportSetup) + func testOAuthProactiveTokenRefreshWithinWindow() async throws { + let scenario = await MockURLProtocol.requestHandlerStorage + .configureOAuthProactiveTokenRefreshWithinWindow() + + let transport = HTTPClientTransport( + endpoint: scenario.testEndpoint, + configuration: MockResponses.ephemeralConfiguration(), + streaming: false, + authorizer: OAuthAuthorizer(configuration: scenario.oauthConfiguration), + logger: nil + ) + + try await transport.connect() + + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + + try await transport.send(scenario.messageData) + let receivedFirst = try await iterator.next() + #expect(receivedFirst == scenario.expectedResponseData) + + try await transport.send(scenario.secondMessageData!) + let receivedSecond = try await iterator.next() + #expect(receivedSecond == scenario.secondExpectedResponseData) + + await MockURLProtocol.verifyCallCounts(scenario.expectedCallCounts) + await transport.disconnect() + } + + @Test("Send With Protocol Version Header", .httpClientTransportSetup) + func testProtocolVersionHeader() async throws { + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockURLProtocol.self] + + let protocolVersion = "2025-11-25" + let transport = HTTPClientTransport( + endpoint: testEndpoint, + configuration: configuration, + streaming: false, + protocolVersion: protocolVersion, + logger: nil + ) + try await transport.connect() + + let messageData = #"{"jsonrpc":"2.0","method":"test","id":6}"#.data(using: .utf8)! + + await MockURLProtocol.requestHandlerStorage.setHandler { + [testEndpoint, protocolVersion] (request: URLRequest) in + // Verify the protocol version header is present + #expect( + request.value(forHTTPHeaderField: "MCP-Protocol-Version") + == protocolVersion) + + let response = HTTPURLResponse( + url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "application/json"])! + return (response, Data()) + } + + try await transport.send(messageData) + await transport.disconnect() + } + #endif // !canImport(FoundationNetworking) + } +#endif // swift(>=6.1) diff --git a/Tests/MCPTests/HTTPServerTransportTests.swift b/Tests/MCPTests/HTTPServerTransportTests.swift new file mode 100644 index 0000000..8e1c89b --- /dev/null +++ b/Tests/MCPTests/HTTPServerTransportTests.swift @@ -0,0 +1,1269 @@ +import Foundation +import Testing + +@testable import MCP + +// MARK: - Test Helpers + +private struct FixedSessionIDGenerator: SessionIDGenerator { + let sessionID: String + func generateSessionID() -> String { sessionID } +} + +private func makeInitializeBody(id: String = "1") -> Data { + let json: [String: Any] = [ + "jsonrpc": "2.0", + "id": id, + "method": "initialize", + "params": [ + "protocolVersion": "2025-11-25", + "capabilities": [:] as [String: Any], + "clientInfo": ["name": "test", "version": "1.0"], + ] as [String: Any], + ] + return try! JSONSerialization.data(withJSONObject: json) +} + +private func makeNotificationBody(method: String = "notifications/initialized") -> Data { + let json: [String: Any] = ["jsonrpc": "2.0", "method": method] + return try! JSONSerialization.data(withJSONObject: json) +} + +private func makeRequestBody(id: String = "2", method: String = "tools/list") -> Data { + let json: [String: Any] = [ + "jsonrpc": "2.0", + "id": id, + "method": method, + "params": [:] as [String: Any], + ] + return try! JSONSerialization.data(withJSONObject: json) +} + +private func makeResponseBody(id: String = "2") -> Data { + let json: [String: Any] = [ + "jsonrpc": "2.0", + "id": id, + "result": ["tools": []] as [String: Any], + ] + return try! JSONSerialization.data(withJSONObject: json) +} + +private func makeStatefulPOSTRequest( + body: Data, + sessionID: String? = nil, + authorization: String? = nil +) -> HTTPRequest { + var headers: [String: String] = [ + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + ] + if let sessionID { + headers["Mcp-Session-Id"] = sessionID + } + if let authorization { + headers[HTTPHeaderName.authorization] = authorization + } + return HTTPRequest(method: "POST", headers: headers, body: body) +} + +private func makeGETRequest(sessionID: String, lastEventID: String? = nil) -> HTTPRequest { + var headers: [String: String] = [ + "Accept": "text/event-stream", + "Mcp-Session-Id": sessionID, + ] + if let lastEventID { + headers["Last-Event-ID"] = lastEventID + } + return HTTPRequest(method: "GET", headers: headers) +} + +private func makeDELETERequest(sessionID: String) -> HTTPRequest { + HTTPRequest( + method: "DELETE", + headers: ["Mcp-Session-Id": sessionID] + ) +} + +private func makeStatelessPOSTRequest(body: Data) -> HTTPRequest { + HTTPRequest( + method: "POST", + headers: [ + "Content-Type": "application/json", + "Accept": "application/json", + ], + body: body + ) +} + +private func makeStatefulTransport( + sessionIDGenerator: any SessionIDGenerator = UUIDSessionIDGenerator() +) -> StatefulHTTPServerTransport { + StatefulHTTPServerTransport( + sessionIDGenerator: sessionIDGenerator, + validationPipeline: StandardValidationPipeline(validators: []) + ) +} + +private let authResourceMetadataURL = + URL(string: "https://mcp.example.com/.well-known/oauth-protected-resource/mcp")! +private let authResourceIdentifier = URL(string: "https://mcp.example.com/mcp")! + +private func makeAuthenticatedStatefulTransport( + challengeScopes: Set? = nil, + tokenValidator: @escaping BearerTokenValidator.TokenValidator +) -> StatefulHTTPServerTransport { + let validator = BearerTokenValidator( + resourceMetadataURL: authResourceMetadataURL, + resourceIdentifier: authResourceIdentifier, + tokenValidator: tokenValidator, + challengeScopeProvider: { _, _ in challengeScopes } + ) + return StatefulHTTPServerTransport( + validationPipeline: StandardValidationPipeline(validators: [ + validator, + SessionValidator(), + ]) + ) +} + +private func makeStatelessTransport() -> StatelessHTTPServerTransport { + StatelessHTTPServerTransport( + validationPipeline: StandardValidationPipeline(validators: []) + ) +} + +/// Drains an SSE stream, collecting raw SSE chunks. +private actor ChunkCollector { + var chunks: [Data] = [] + func append(_ data: Data) { chunks.append(data) } + func getChunks() -> [Data] { chunks } +} + +private func drainSSEStream( + _ response: HTTPResponse, + maxChunks: Int = 10, + timeout: Duration = .seconds(2) +) async -> [Data] { + guard case .stream(let stream, _) = response else { return [] } + let collector = ChunkCollector() + let task = Task { + for try await chunk in stream { + await collector.append(chunk) + if await collector.getChunks().count >= maxChunks { break } + } + } + // Wait for stream to finish or timeout + try? await Task.sleep(for: timeout) + task.cancel() + return await collector.getChunks() +} + +/// Initializes a stateful transport session and returns the session ID. +/// Spawns a background task to consume the receive stream and send the init response. +private func initializeSession( + transport: StatefulHTTPServerTransport, + sessionID: String? = nil +) async throws -> String { + try await transport.connect() + + let initBody = makeInitializeBody() + + // Background task: read the init request from receive() and send back a response + let respondTask = Task { + let stream = await transport.receive() + for try await data in stream { + // Check if this is the initialize request + if let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let method = json["method"] as? String, method == "initialize", + let id = json["id"] + { + let idString: String + if let s = id as? String { idString = s } + else if let n = id as? Int { idString = String(n) } + else { continue } + + let responseJSON: [String: Any] = [ + "jsonrpc": "2.0", + "id": idString, + "result": [ + "protocolVersion": "2025-11-25", + "serverInfo": ["name": "test", "version": "1.0"], + "capabilities": [:] as [String: Any], + ] as [String: Any], + ] + let responseData = try JSONSerialization.data(withJSONObject: responseJSON) + try await transport.send(responseData) + return + } + } + } + + let response = await transport.handleRequest( + makeStatefulPOSTRequest(body: initBody) + ) + + // Extract session ID + guard let sid = response.headers[HTTPHeaderName.sessionID] else { + throw MCPError.internalError("No session ID in init response") + } + + // Drain the SSE stream so the response task can complete + if case .stream(let stream, _) = response { + Task { for try await _ in stream {} } + } + + // Wait for the respond task + try? await respondTask.value + + return sid +} + +// MARK: - StatefulHTTPServerTransport Tests + +@Suite("StatefulHTTPServerTransport Tests") +struct StatefulHTTPServerTransportTests { + + // MARK: - Lifecycle + + @Test("Connect succeeds") + func testConnectSucceeds() async throws { + let transport = makeStatefulTransport() + try await transport.connect() + await transport.disconnect() + } + + @Test("Double connect throws") + func testDoubleConnectThrows() async throws { + let transport = makeStatefulTransport() + try await transport.connect() + do { + try await transport.connect() + Issue.record("Expected error on double connect") + } catch { + // Expected + } + await transport.disconnect() + } + + @Test("Send after disconnect throws connectionClosed") + func testSendAfterDisconnectThrows() async throws { + let transport = makeStatefulTransport() + try await transport.connect() + await transport.disconnect() + do { + try await transport.send(Data("test".utf8)) + Issue.record("Expected connectionClosed error") + } catch let error as MCPError { + #expect(error == .connectionClosed) + } + } + + // MARK: - POST Initialize + + @Test("Initialize creates session and returns SSE stream") + func testInitializeCreatesSession() async throws { + let transport = makeStatefulTransport( + sessionIDGenerator: FixedSessionIDGenerator(sessionID: "test-session-42") + ) + try await transport.connect() + + let response = await transport.handleRequest( + makeStatefulPOSTRequest(body: makeInitializeBody()) + ) + + #expect(response.statusCode == 200) + #expect(response.headers[HTTPHeaderName.sessionID] == "test-session-42") + + if case .stream = response { + // Expected + } else { + Issue.record("Expected .stream response, got \(response)") + } + + // Drain stream + if case .stream(let stream, _) = response { + Task { for try await _ in stream {} } + } + await transport.disconnect() + } + + @Test("Initialize with invalid session ID returns 500") + func testInitializeWithInvalidSessionIDReturns500() async throws { + // Control character \t is 0x09, outside valid range 0x21-0x7E + let transport = makeStatefulTransport( + sessionIDGenerator: FixedSessionIDGenerator(sessionID: "bad\tsession") + ) + try await transport.connect() + + let response = await transport.handleRequest( + makeStatefulPOSTRequest(body: makeInitializeBody()) + ) + + #expect(response.statusCode == 500) + } + + @Test("Custom SessionIDGenerator is used") + func testCustomSessionIDGenerator() async throws { + let transport = makeStatefulTransport( + sessionIDGenerator: FixedSessionIDGenerator(sessionID: "custom-id-abc") + ) + try await transport.connect() + + let response = await transport.handleRequest( + makeStatefulPOSTRequest(body: makeInitializeBody()) + ) + + #expect(response.headers[HTTPHeaderName.sessionID] == "custom-id-abc") + + if case .stream(let stream, _) = response { + Task { for try await _ in stream {} } + } + await transport.disconnect() + } + + @Test("Default UUIDSessionIDGenerator produces valid session ID") + func testDefaultGeneratorProducesUUID() async throws { + let transport = makeStatefulTransport() + try await transport.connect() + + let response = await transport.handleRequest( + makeStatefulPOSTRequest(body: makeInitializeBody()) + ) + + let sessionID = response.headers[HTTPHeaderName.sessionID] + #expect(sessionID != nil) + // UUID format: 8-4-4-4-12 hex chars + if let sid = sessionID { + #expect(sid.count == 36) + #expect(sid.contains("-")) + } + + if case .stream(let stream, _) = response { + Task { for try await _ in stream {} } + } + await transport.disconnect() + } + + // MARK: - POST Notification + + @Test("Notification returns 202 Accepted") + func testNotificationReturns202() async throws { + let transport = makeStatefulTransport() + let sessionID = try await initializeSession(transport: transport) + + let response = await transport.handleRequest( + makeStatefulPOSTRequest( + body: makeNotificationBody(), + sessionID: sessionID + ) + ) + + #expect(response.statusCode == 202) + await transport.disconnect() + } + + @Test("Notification yields to receive stream") + func testNotificationYieldsToReceive() async throws { + let transport = makeStatefulTransport() + let sessionID = try await initializeSession(transport: transport) + + let notificationBody = makeNotificationBody(method: "notifications/test") + + // Start receiving + let receiveTask = Task { + let stream = await transport.receive() + for try await data in stream { + // Skip init request if still in stream + if let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let method = json["method"] as? String, method == "notifications/test" + { + return data + } + } + return nil + } + + // Small delay to let receive() start + try await Task.sleep(for: .milliseconds(50)) + + _ = await transport.handleRequest( + makeStatefulPOSTRequest(body: notificationBody, sessionID: sessionID) + ) + + let received = try await receiveTask.value + #expect(received != nil) + + await transport.disconnect() + } + + // MARK: - POST Request/Response + + @Test("POST request returns SSE stream") + func testRequestReturnsSSEStream() async throws { + let transport = makeStatefulTransport() + let sessionID = try await initializeSession(transport: transport) + + let response = await transport.handleRequest( + makeStatefulPOSTRequest( + body: makeRequestBody(id: "req-1"), + sessionID: sessionID + ) + ) + + #expect(response.statusCode == 200) + if case .stream = response { + // Expected + } else { + Issue.record("Expected .stream response") + } + + if case .stream(let stream, _) = response { + Task { for try await _ in stream {} } + } + await transport.disconnect() + } + + @Test("Response is routed to matching request SSE stream") + func testResponseRoutedToRequestStream() async throws { + let transport = makeStatefulTransport() + let sessionID = try await initializeSession(transport: transport) + + let requestID = "route-test-1" + + // POST a request + let response = await transport.handleRequest( + makeStatefulPOSTRequest( + body: makeRequestBody(id: requestID, method: "tools/list"), + sessionID: sessionID + ) + ) + + guard case .stream(let stream, _) = response else { + Issue.record("Expected .stream response") + return + } + + // Collect SSE chunks in background + let collectTask = Task { + var chunks: [Data] = [] + for try await chunk in stream { + chunks.append(chunk) + } + return chunks + } + + // Give stream time to start + try await Task.sleep(for: .milliseconds(50)) + + // Consume the request from receive and send the response + let responseBody = makeResponseBody(id: requestID) + try await transport.send(responseBody) + + // Collect all SSE chunks + let chunks = try await collectTask.value + + // Should have at least one chunk containing the response data + let allText = chunks.map { String(decoding: $0, as: UTF8.self) }.joined() + #expect(allText.contains("data:")) + #expect(allText.contains(requestID)) + + await transport.disconnect() + } + + // MARK: - GET Stream + + @Test("GET returns standalone SSE stream") + func testGetReturnsSSEStream() async throws { + let transport = makeStatefulTransport() + let sessionID = try await initializeSession(transport: transport) + + let response = await transport.handleRequest( + makeGETRequest(sessionID: sessionID) + ) + + #expect(response.statusCode == 200) + if case .stream = response { + // Expected + } else { + Issue.record("Expected .stream response for GET") + } + + if case .stream(let stream, _) = response { + Task { for try await _ in stream {} } + } + await transport.disconnect() + } + + @Test("Server-initiated message routed to GET stream") + func testServerMessageRoutedToGetStream() async throws { + let transport = makeStatefulTransport() + let sessionID = try await initializeSession(transport: transport) + + // Open GET stream + let getResponse = await transport.handleRequest( + makeGETRequest(sessionID: sessionID) + ) + + guard case .stream(let stream, _) = getResponse else { + Issue.record("Expected .stream response for GET") + return + } + + // Collect chunks + let collectTask = Task { + var chunks: [Data] = [] + for try await chunk in stream { + chunks.append(chunk) + // priming + message + if chunks.count >= 2 { break } + } + return chunks + } + + try await Task.sleep(for: .milliseconds(50)) + + // Send a notification (server-initiated) + let notification: [String: Any] = [ + "jsonrpc": "2.0", + "method": "notifications/test", + "params": [:] as [String: Any], + ] + let notifData = try JSONSerialization.data(withJSONObject: notification) + try await transport.send(notifData) + + let chunks = try await collectTask.value + let allText = chunks.map { String(decoding: $0, as: UTF8.self) }.joined() + #expect(allText.contains("data:")) + // JSONSerialization may escape "/" as "\/" in some configurations + #expect(allText.contains("notifications/test") || allText.contains("notifications\\/test")) + + await transport.disconnect() + } + + @Test("Second GET returns 409 Conflict") + func testSecondGetReturns409() async throws { + let transport = makeStatefulTransport() + let sessionID = try await initializeSession(transport: transport) + + // First GET + let first = await transport.handleRequest(makeGETRequest(sessionID: sessionID)) + #expect(first.statusCode == 200) + + // Second GET + let second = await transport.handleRequest(makeGETRequest(sessionID: sessionID)) + #expect(second.statusCode == 409) + + if case .stream(let stream, _) = first { + Task { for try await _ in stream {} } + } + await transport.disconnect() + } + + // MARK: - DELETE + + @Test("DELETE terminates session") + func testDeleteTerminatesSession() async throws { + let transport = makeStatefulTransport() + let sessionID = try await initializeSession(transport: transport) + + let response = await transport.handleRequest( + makeDELETERequest(sessionID: sessionID) + ) + + #expect(response.statusCode == 200) + } + + @Test("Requests after DELETE return 404") + func testRequestsAfterDeleteReturn404() async throws { + let transport = makeStatefulTransport() + let sessionID = try await initializeSession(transport: transport) + + // DELETE + _ = await transport.handleRequest(makeDELETERequest(sessionID: sessionID)) + + // POST after delete + let response = await transport.handleRequest( + makeStatefulPOSTRequest( + body: makeRequestBody(), + sessionID: sessionID + ) + ) + + #expect(response.statusCode == 404) + } + + // MARK: - Terminated State + + @Test("All methods return 404 when terminated") + func testAllMethodsReturn404WhenTerminated() async throws { + let transport = makeStatefulTransport() + let sessionID = try await initializeSession(transport: transport) + await transport.disconnect() + + let post = await transport.handleRequest( + makeStatefulPOSTRequest(body: makeRequestBody(), sessionID: sessionID) + ) + #expect(post.statusCode == 404) + + let get = await transport.handleRequest(makeGETRequest(sessionID: sessionID)) + #expect(get.statusCode == 404) + + let delete = await transport.handleRequest(makeDELETERequest(sessionID: sessionID)) + #expect(delete.statusCode == 404) + } + + // MARK: - Error Cases + + @Test("Unsupported method returns 405") + func testUnsupportedMethodReturns405() async throws { + let transport = makeStatefulTransport() + let sessionID = try await initializeSession(transport: transport) + + let response = await transport.handleRequest( + HTTPRequest( + method: "PUT", + headers: ["Mcp-Session-Id": sessionID], + body: Data("test".utf8) + ) + ) + + #expect(response.statusCode == 405) + await transport.disconnect() + } + + @Test("Empty body returns 400") + func testEmptyBodyReturns400() async throws { + let transport = makeStatefulTransport() + let sessionID = try await initializeSession(transport: transport) + + let response = await transport.handleRequest( + makeStatefulPOSTRequest(body: Data(), sessionID: sessionID) + ) + + #expect(response.statusCode == 400) + await transport.disconnect() + } + + @Test("Invalid JSON body returns 400") + func testInvalidJSONReturns400() async throws { + let transport = makeStatefulTransport() + let sessionID = try await initializeSession(transport: transport) + + let response = await transport.handleRequest( + makeStatefulPOSTRequest(body: Data("not json".utf8), sessionID: sessionID) + ) + + #expect(response.statusCode == 400) + await transport.disconnect() + } + + // MARK: - OAuth Bearer Validation + + @Test("Bearer auth validator returns 401 with challenge when authorization is missing") + func testBearerAuthValidatorMissingAuthorizationReturns401() async throws { + let transport = makeAuthenticatedStatefulTransport( + challengeScopes: ["files:read"] + ) { _, _, _ in + .valid(BearerTokenInfo()) + } + try await transport.connect() + + let response = await transport.handleRequest( + makeStatefulPOSTRequest(body: makeInitializeBody()) + ) + + #expect(response.statusCode == 401) + let challenge = response.headers[HTTPHeaderName.wwwAuthenticate] + #expect(challenge?.contains("Bearer ") == true) + #expect( + challenge?.contains("resource_metadata=\"\(authResourceMetadataURL.absoluteString)\"") + == true + ) + #expect(challenge?.contains("scope=\"files:read\"") == true) + + await transport.disconnect() + } + + @Test("Bearer auth validator returns 400 when authorization header is malformed") + func testBearerAuthValidatorMalformedAuthorizationReturns400() async throws { + let transport = makeAuthenticatedStatefulTransport { _, _, _ in .valid(BearerTokenInfo()) } + try await transport.connect() + + let response = await transport.handleRequest( + makeStatefulPOSTRequest( + body: makeInitializeBody(), + authorization: "Basic dGVzdA==" + ) + ) + + #expect(response.statusCode == 400) + #expect(response.headers[HTTPHeaderName.wwwAuthenticate] == nil) + + await transport.disconnect() + } + + @Test("Bearer auth validator returns 401 invalid_token for rejected token") + func testBearerAuthValidatorInvalidTokenReturns401() async throws { + let transport = makeAuthenticatedStatefulTransport { _, _, _ in + .invalidToken(errorDescription: "Token expired") + } + try await transport.connect() + + let response = await transport.handleRequest( + makeStatefulPOSTRequest( + body: makeInitializeBody(), + authorization: "Bearer expired-token" + ) + ) + + #expect(response.statusCode == 401) + let challenge = response.headers[HTTPHeaderName.wwwAuthenticate] + #expect(challenge?.contains("error=\"invalid_token\"") == true) + #expect(challenge?.contains("error_description=\"Token expired\"") == true) + #expect( + challenge?.contains("resource_metadata=\"\(authResourceMetadataURL.absoluteString)\"") + == true + ) + + await transport.disconnect() + } + + @Test("Bearer auth validator returns 403 insufficient_scope with scope challenge") + func testBearerAuthValidatorInsufficientScopeReturns403() async throws { + let transport = makeAuthenticatedStatefulTransport { token, _, context in + if context.isInitializationRequest, token == "init-token" { + return .valid(BearerTokenInfo()) + } + if token == "read-only-token" { + return .insufficientScope( + requiredScopes: ["files:read", "files:write"], + errorDescription: "Additional file write permission required" + ) + } + return .invalidToken(errorDescription: "Unknown token") + } + try await transport.connect() + + let initResponse = await transport.handleRequest( + makeStatefulPOSTRequest( + body: makeInitializeBody(), + authorization: "Bearer init-token" + ) + ) + let sessionID = initResponse.headers[HTTPHeaderName.sessionID] + #expect(sessionID != nil) + + guard let sessionID else { + await transport.disconnect() + return + } + + let response = await transport.handleRequest( + makeStatefulPOSTRequest( + body: makeNotificationBody(), + sessionID: sessionID, + authorization: "Bearer read-only-token" + ) + ) + + #expect(response.statusCode == 403) + let challenge = response.headers[HTTPHeaderName.wwwAuthenticate] + #expect(challenge?.contains("error=\"insufficient_scope\"") == true) + #expect(challenge?.contains("scope=\"files:read files:write\"") == true) + #expect( + challenge?.contains("resource_metadata=\"\(authResourceMetadataURL.absoluteString)\"") + == true + ) + #expect( + challenge?.contains( + "error_description=\"Additional file write permission required\"" + ) == true + ) + + await transport.disconnect() + } + + // MARK: - Resumability + + @Test("GET with Last-Event-ID replays stored events") + func testGetWithLastEventIDReplaysEvents() async throws { + let transport = makeStatefulTransport() + let sessionID = try await initializeSession(transport: transport) + + // POST a request to create events in the store + let requestID = "resume-test" + let postResponse = await transport.handleRequest( + makeStatefulPOSTRequest( + body: makeRequestBody(id: requestID), + sessionID: sessionID + ) + ) + + guard case .stream(let postStream, _) = postResponse else { + Issue.record("Expected .stream") + return + } + + // Collect the priming event to get its ID + let eventIDHolder = ChunkCollector() + let collectTask = Task { + for try await chunk in postStream { + await eventIDHolder.append(chunk) + break // Just get the first chunk (priming) + } + } + + try await Task.sleep(for: .milliseconds(50)) + + // Send the response to create a stored event + try await transport.send(makeResponseBody(id: requestID)) + + try? await collectTask.value + + // Parse event ID from the collected priming event + let collectedChunks = await eventIDHolder.getChunks() + let primingEventID: String? = collectedChunks.first.flatMap { chunk in + let text = String(decoding: chunk, as: UTF8.self) + guard let range = text.range(of: "id: ") else { return nil } + let afterID = text[range.upperBound...] + guard let newline = afterID.firstIndex(of: "\n") else { return nil } + return String(afterID[...Continuation? + + private(set) var sentData: [Data] = [] + var sentMessages: [String] { + return sentData.compactMap { data in + guard let string = String(data: data, encoding: .utf8) else { + logger.error("Failed to decode sent data as UTF-8") + return nil + } + return string + } + } + + private var dataToReceive: [Data] = [] + private(set) var receivedMessages: [String] = [] + + private var dataStreamContinuation: AsyncThrowingStream.Continuation? + var shouldFailConnect = false var shouldFailSend = false @@ -17,85 +39,39 @@ actor MockTransport: Transport { self.logger = logger } - func connect() async throws { + public func connect() async throws { if shouldFailConnect { - throw Error.transportError(POSIXError(.ECONNREFUSED)) + throw MCPError.transportError(POSIXError(.ECONNREFUSED)) } isConnected = true } - func disconnect() async { + public func disconnect() async { isConnected = false - messageStreamContinuation?.finish() - messageStreamContinuation = nil + dataStreamContinuation?.finish() + dataStreamContinuation = nil } - func send(_ message: T) async throws { + public func send(_ message: Data) async throws { if shouldFailSend { - throw Error.transportError(POSIXError(.EIO)) + throw MCPError.transportError(POSIXError(.EIO)) } - let data = try JSONEncoder().encode(message) - let str = String(data: data, encoding: .utf8)! - sentMessages.append(str) + sentData.append(message) } - func receive() -> AsyncThrowingStream { - return AsyncThrowingStream { continuation in - messageStreamContinuation = continuation - // Send any queued messages - for message in messagesToReceive { + public func receive() -> AsyncThrowingStream { + return AsyncThrowingStream { continuation in + dataStreamContinuation = continuation + for message in dataToReceive { continuation.yield(message) + if let string = String(data: message, encoding: .utf8) { + receivedMessages.append(string) + } } - messagesToReceive.removeAll() + dataToReceive.removeAll() } } - func queueRequest(_ request: Request) throws { - let data = try JSONEncoder().encode(request) - let str = String(data: data, encoding: .utf8)! - if let continuation = messageStreamContinuation { - continuation.yield(str) - } else { - sentMessages.append(str) - } - } - - func queueResponse(_ response: Response) throws { - let data = try JSONEncoder().encode(response) - let str = String(data: data, encoding: .utf8)! - if let continuation = messageStreamContinuation { - continuation.yield(str) - } else { - messagesToReceive.append(str) - } - } - - func queueNotification(_ notification: Message) throws { - let data = try JSONEncoder().encode(notification) - let str = String(data: data, encoding: .utf8)! - if let continuation = messageStreamContinuation { - continuation.yield(str) - } else { - messagesToReceive.append(str) - } - } - - func getLastSentMessage() -> T? { - print("SENT:", sentMessages) - guard let lastMessage = sentMessages.last else { return nil } - do { - let data = lastMessage.data(using: .utf8)! - return try JSONDecoder().decode(T.self, from: data) - } catch { - return nil - } - } - - func clearMessages() { - sentMessages.removeAll() - messagesToReceive.removeAll() - } - func setFailConnect(_ shouldFail: Bool) { shouldFailConnect = shouldFail } @@ -103,4 +79,46 @@ actor MockTransport: Transport { func setFailSend(_ shouldFail: Bool) { shouldFailSend = shouldFail } + + 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 queue(response: Response) throws { + queue(data: try encoder.encode(response)) + } + + func queue(notification: Message) throws { + queue(data: try encoder.encode(notification)) + } + + func queue(batch requests: [AnyRequest]) throws { + queue(data: try encoder.encode(requests)) + } + + func queue(batch responses: [AnyResponse]) throws { + queue(data: try encoder.encode(responses)) + } + + func decodeLastSentMessage() -> T? { + guard let lastMessage = sentData.last else { return nil } + do { + return try decoder.decode(T.self, from: lastMessage) + } catch { + return nil + } + } + + func clearMessages() { + sentData.removeAll() + dataToReceive.removeAll() + } } diff --git a/Tests/MCPTests/InMemoryTransportTests.swift b/Tests/MCPTests/InMemoryTransportTests.swift new file mode 100644 index 0000000..d8228e7 --- /dev/null +++ b/Tests/MCPTests/InMemoryTransportTests.swift @@ -0,0 +1,409 @@ +import Foundation +import Logging +import Testing + +@testable import MCP + +@Suite("InMemory Transport Tests") +struct InMemoryTransportTests { + + @Test("Create connected pair") + func testCreateConnectedPair() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + // Verify both transports can be connected + try await clientTransport.connect() + try await serverTransport.connect() + + // Clean up + await clientTransport.disconnect() + await serverTransport.disconnect() + } + + @Test("Connect without pairing throws error") + func testConnectWithoutPairing() async throws { + let transport = InMemoryTransport() + + // Attempt to connect without pairing should throw + do { + try await transport.connect() + #expect(Bool(false), "Expected connect to throw an error") + } catch let error as MCPError { + if case .internalError(let message) = error { + #expect( + message + == "Transport not paired. Use createConnectedPair() to create paired transports." + ) + } else { + #expect(Bool(false), "Expected MCPError.internalError") + } + } + } + + @Test("Multiple connect calls are idempotent") + func testMultipleConnectCalls() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + // Connect multiple times should not throw + try await clientTransport.connect() + try await clientTransport.connect() // Should be safe + try await clientTransport.connect() // Should be safe + + // Clean up + await clientTransport.disconnect() + await serverTransport.disconnect() + } + + @Test("Send and receive messages") + func testSendAndReceiveMessages() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + try await clientTransport.connect() + try await serverTransport.connect() + + // Start receiving on server + let serverReceiveTask = Task { + var messages: [Data] = [] + for try await message in await serverTransport.receive() { + messages.append(message) + if messages.count >= 3 { + break + } + } + return messages + } + + // Send messages from client + let message1 = "Hello".data(using: .utf8)! + let message2 = "World".data(using: .utf8)! + let message3 = "!".data(using: .utf8)! + + try await clientTransport.send(message1) + try await clientTransport.send(message2) + try await clientTransport.send(message3) + + // Wait for messages to be received + let receivedMessages = try await serverReceiveTask.value + + #expect(receivedMessages.count == 3) + #expect(receivedMessages[0] == message1) + #expect(receivedMessages[1] == message2) + #expect(receivedMessages[2] == message3) + + // Clean up + await clientTransport.disconnect() + await serverTransport.disconnect() + } + + @Test("Bidirectional communication") + func testBidirectionalCommunication() async throws { + let (transport1, transport2) = await InMemoryTransport.createConnectedPair() + + try await transport1.connect() + try await transport2.connect() + + // 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) { + messages.append(message) + if messages.count >= 2 { + break + } + } + } + return messages + } + + let receive2Task = Task { + var messages: [String] = [] + for try await data in await transport2.receive() { + if let message = String(data: data, encoding: .utf8) { + messages.append(message) + if messages.count >= 2 { + break + } + } + } + return messages + } + + // Send messages in both directions + try await transport1.send("From transport 1 - message 1".data(using: .utf8)!) + try await transport2.send("From transport 2 - message 1".data(using: .utf8)!) + try await transport1.send("From transport 1 - message 2".data(using: .utf8)!) + try await transport2.send("From transport 2 - message 2".data(using: .utf8)!) + + // Verify both sides received messages + let messages1 = try await receive1Task.value + let messages2 = try await receive2Task.value + + #expect(messages1.count == 2) + #expect(messages1[0] == "From transport 2 - message 1") + #expect(messages1[1] == "From transport 2 - message 2") + + #expect(messages2.count == 2) + #expect(messages2[0] == "From transport 1 - message 1") + #expect(messages2[1] == "From transport 1 - message 2") + + // Clean up + await transport1.disconnect() + await transport2.disconnect() + } + + @Test("Send without connection throws error") + func testSendWithoutConnection() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + // Try to send without connecting + do { + try await clientTransport.send("test".data(using: .utf8)!) + #expect(Bool(false), "Expected send to throw an error") + } catch let error as MCPError { + if case .internalError(let message) = error { + #expect(message == "Transport not connected") + } else { + #expect(Bool(false), "Expected MCPError.internalError") + } + } + + // Clean up (connect server to avoid dangling connections) + try await serverTransport.connect() + await serverTransport.disconnect() + } + + @Test("Disconnect stops message stream") + func testDisconnectStopsMessageStream() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + try await clientTransport.connect() + try await serverTransport.connect() + + // Start receiving + let receiveTask = Task { + var messageCount = 0 + do { + for try await _ in await serverTransport.receive() { + messageCount += 1 + } + } catch { + // Expected when disconnected + } + return messageCount + } + + // Send a message + try await clientTransport.send("message".data(using: .utf8)!) + + // Give some time for message to be received + try await Task.sleep(for: .milliseconds(100)) + + // Disconnect + await serverTransport.disconnect() + + // The receive stream should complete + let messageCount = await receiveTask.value + #expect(messageCount >= 1) + } + + @Test("Peer disconnection handling") + func testPeerDisconnectionHandling() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + try await clientTransport.connect() + try await serverTransport.connect() + + // Start receiving on server + let receiveTask = Task { () -> Swift.Error? in + do { + for try await _ in await serverTransport.receive() { + // Keep receiving + } + return nil + } catch { + return error + } + } + + // Give a moment for the receive stream to be set up + try await Task.sleep(for: .milliseconds(10)) + + // Disconnect client (peer) + await clientTransport.disconnect() + + receiveTask.cancel() + } + + @Test("Multiple disconnects are safe") + func testMultipleDisconnects() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + try await clientTransport.connect() + try await serverTransport.connect() + + // Multiple disconnects should be safe + await clientTransport.disconnect() + await clientTransport.disconnect() + await clientTransport.disconnect() + + // Clean up server + await serverTransport.disconnect() + } + + @Test("Message queueing before stream creation") + func testMessageQueueingBeforeStreamCreation() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + try await clientTransport.connect() + try await serverTransport.connect() + + // Send messages before receive stream is created + try await clientTransport.send("message1".data(using: .utf8)!) + try await clientTransport.send("message2".data(using: .utf8)!) + try await clientTransport.send("message3".data(using: .utf8)!) + + // Now create receive stream + let messages = await serverTransport.receive() + var receivedMessages: [String] = [] + + for try await data in messages { + if let message = String(data: data, encoding: .utf8) { + receivedMessages.append(message) + if receivedMessages.count >= 3 { + break + } + } + } + + #expect(receivedMessages.count == 3) + #expect(receivedMessages[0] == "message1") + #expect(receivedMessages[1] == "message2") + #expect(receivedMessages[2] == "message3") + + // Clean up + await clientTransport.disconnect() + await serverTransport.disconnect() + } + + @Test("Receive after disconnect returns completed stream") + func testReceiveAfterDisconnect() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + try await clientTransport.connect() + try await serverTransport.connect() + + // Disconnect before receiving + await serverTransport.disconnect() + + // Create receive stream after disconnect + let messages = await serverTransport.receive() + var messageCount = 0 + + for try await _ in messages { + messageCount += 1 + } + + // Stream should complete immediately + #expect(messageCount == 0) + + // Clean up + await clientTransport.disconnect() + } + + @Test("Large message handling") + func testLargeMessageHandling() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + try await clientTransport.connect() + try await serverTransport.connect() + + // Create a large message (1MB) + let largeData = Data(repeating: 0xFF, count: 1024 * 1024) + + // Start receiving + let receiveTask = Task { + for try await data in await serverTransport.receive() { + return data + } + return Data() + } + + // Send large message + try await clientTransport.send(largeData) + + // Verify it was received correctly + let receivedData = try await receiveTask.value + #expect(receivedData == largeData) + + // Clean up + await clientTransport.disconnect() + await serverTransport.disconnect() + } + + @Test("Concurrent send operations") + func testConcurrentSendOperations() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + try await clientTransport.connect() + try await serverTransport.connect() + + // Start receiving + let receiveTask = Task { + var messages: [String] = [] + for try await data in await serverTransport.receive() { + if let message = String(data: data, encoding: .utf8) { + messages.append(message) + if messages.count >= 10 { + break + } + } + } + return messages + } + + // Send messages concurrently + await withTaskGroup(of: Void.self) { group in + for i in 0..<10 { + group.addTask { + try? await clientTransport.send("Message \(i)".data(using: .utf8)!) + } + } + } + + // Verify all messages were received + let receivedMessages = try await receiveTask.value + #expect(receivedMessages.count == 10) + + // Check that all messages are present (order may vary due to concurrency) + let expectedMessages = Set((0..<10).map { "Message \($0)" }) + let actualMessages = Set(receivedMessages) + #expect(actualMessages == expectedMessages) + + // Clean up + await clientTransport.disconnect() + await serverTransport.disconnect() + } + + @Test("Custom logger usage") + func testCustomLoggerUsage() async throws { + // Create a custom logger (in real tests, you might use a test logger that captures output) + let logger = Logger(label: "test.in-memory.transport") + + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair( + logger: logger) + + // Verify loggers are set correctly - when a custom logger is provided, it's used for both + #expect(clientTransport.logger.label == "test.in-memory.transport") + #expect(serverTransport.logger.label == "test.in-memory.transport") + + // Test basic operations with custom logger + try await clientTransport.connect() + try await serverTransport.connect() + + try await clientTransport.send("test".data(using: .utf8)!) + + await clientTransport.disconnect() + await serverTransport.disconnect() + } +} diff --git a/Tests/MCPTests/LoggingTests.swift b/Tests/MCPTests/LoggingTests.swift new file mode 100644 index 0000000..47e89b4 --- /dev/null +++ b/Tests/MCPTests/LoggingTests.swift @@ -0,0 +1,509 @@ +import Foundation +import Testing + +@testable import MCP + +@Suite("Logging Tests") +struct LoggingTests { + // MARK: - LogLevel Tests + + @Test("LogLevel case values") + func testLogLevelCases() throws { + #expect(LogLevel.debug.rawValue == "debug") + #expect(LogLevel.info.rawValue == "info") + #expect(LogLevel.notice.rawValue == "notice") + #expect(LogLevel.warning.rawValue == "warning") + #expect(LogLevel.error.rawValue == "error") + #expect(LogLevel.critical.rawValue == "critical") + #expect(LogLevel.alert.rawValue == "alert") + #expect(LogLevel.emergency.rawValue == "emergency") + } + + @Test("LogLevel encoding and decoding") + func testLogLevelEncodingDecoding() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + for level in LogLevel.allCases { + let data = try encoder.encode(level) + let decoded = try decoder.decode(LogLevel.self, from: data) + #expect(decoded == level) + + // Verify it encodes to the raw string value + let jsonString = String(data: data, encoding: .utf8) + #expect(jsonString == "\"\(level.rawValue)\"") + } + } + + // MARK: - SetLoggingLevel Tests + + @Test("SetLoggingLevel request initialization") + func testSetLoggingLevelRequest() throws { + let request = SetLoggingLevel.request(.init(level: .info)) + + #expect(request.method == "logging/setLevel") + #expect(request.params.level == LogLevel.info) + } + + @Test("SetLoggingLevel request encoding") + func testSetLoggingLevelRequestEncoding() throws { + let request = SetLoggingLevel.request(.init(level: .error)) + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + + let data = try encoder.encode(request) + let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] + + #expect(json?["jsonrpc"] as? String == "2.0") + #expect(json?["method"] as? String == "logging/setLevel") + + let params = json?["params"] as? [String: Any] + #expect(params?["level"] as? String == "error") + } + + @Test("SetLoggingLevel request decoding") + func testSetLoggingLevelRequestDecoding() throws { + let json = """ + { + "jsonrpc": "2.0", + "id": "test-id", + "method": "logging/setLevel", + "params": { + "level": "warning" + } + } + """ + + let decoder = JSONDecoder() + let request = try decoder.decode(Request.self, from: json.data(using: .utf8)!) + + #expect(request.method == "logging/setLevel") + #expect(request.params.level == .warning) + } + + @Test("SetLoggingLevel response") + func testSetLoggingLevelResponse() throws { + let response = SetLoggingLevel.response(id: .random) + + if case .success = response.result { + // Success case + } else { + Issue.record("Expected success result") + } + } + + // MARK: - LogMessageNotification Tests + + @Test("LogMessageNotification initialization") + func testLogMessageNotificationInitialization() throws { + let data = Value.object([ + "message": Value.string("Test log message"), + "code": Value.int(42) + ]) + + let params = LogMessageNotification.Parameters( + level: .info, + logger: "test-logger", + data: data + ) + + #expect(params.level == LogLevel.info) + #expect(params.logger == "test-logger") + #expect(params.data == data) + } + + @Test("LogMessageNotification with nil logger") + func testLogMessageNotificationWithNilLogger() throws { + let data = Value.object(["message": Value.string("Test")]) + + let params = LogMessageNotification.Parameters( + level: .debug, + logger: nil, + data: data + ) + + #expect(params.level == LogLevel.debug) + #expect(params.logger == nil) + #expect(params.data == data) + } + + @Test("LogMessageNotification encoding") + func testLogMessageNotificationEncoding() throws { + let data = Value.object([ + "error": Value.string("Connection failed"), + "details": .object([ + "host": Value.string("localhost"), + "port": Value.int(5432) + ]) + ]) + + let notification = LogMessageNotification.message( + .init(level: .error, logger: "database", data: data) + ) + + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys, .withoutEscapingSlashes] + + let encodedData = try encoder.encode(notification) + let json = try JSONSerialization.jsonObject(with: encodedData) as? [String: Any] + + guard let jsonValue = json else { + Issue.record("Failed to parse JSON") + return + } + + #expect(jsonValue["jsonrpc"] as? String == "2.0") + #expect(jsonValue["method"] as? String == "notifications/message") + + guard let params = jsonValue["params"] as? [String: Any] else { + Issue.record("Failed to get params") + return + } + #expect(params["level"] as? String == "error") + #expect(params["logger"] as? String == "database") + + guard let dataDict = params["data"] as? [String: Any] else { + Issue.record("Failed to get data") + return + } + #expect(dataDict["error"] as? String == "Connection failed") + } + + @Test("LogMessageNotification decoding") + func testLogMessageNotificationDecoding() throws { + let json = """ + { + "jsonrpc": "2.0", + "method": "notifications/message", + "params": { + "level": "info", + "logger": "app", + "data": { + "message": "Server started", + "port": 8080 + } + } + } + """ + + let decoder = JSONDecoder() + let notification = try decoder.decode(Message.self, from: json.data(using: .utf8)!) + + #expect(notification.method == "notifications/message") + #expect(notification.params.level == LogLevel.info) + #expect(notification.params.logger == "app") + + if case .object(let dataDict) = notification.params.data { + #expect(dataDict["message"] == Value.string("Server started")) + #expect(dataDict["port"] == Value.int(8080)) + } else { + Issue.record("Expected object data") + } + } + + // MARK: - Client Integration Tests + + @Test("Client setLoggingLevel sends correct request") + func testClientSetLoggingLevel() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let client = Client(name: "TestClient", version: "1.0") + let server = Server( + name: "TestServer", + version: "1.0", + capabilities: .init(logging: .init()) + ) + + actor TestState { + var receivedLevel: LogLevel? + func setLevel(_ level: LogLevel) { receivedLevel = level } + func getLevel() -> LogLevel? { receivedLevel } + } + + let state = TestState() + + // Register handler for setLoggingLevel on server + await server.withMethodHandler(SetLoggingLevel.self) { params in + await state.setLevel(params.level) + return Empty() + } + + try await server.start(transport: serverTransport) + let initResult = try await client.connect(transport: clientTransport) + + // Verify logging capability is advertised + #expect(initResult.capabilities.logging != nil) + + // Call setLoggingLevel + try await client.setLoggingLevel(.warning) + + // Give time for message processing + try await Task.sleep(for: .milliseconds(100)) + + // Verify the handler was called + #expect(await state.getLevel() == .warning) + + await client.disconnect() + await server.stop() + } + + @Test("Client setLoggingLevel fails without logging capability") + func testClientSetLoggingLevelFailsWithoutCapability() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let client = Client(name: "TestClient", version: "1.0", configuration: .strict) + let server = Server( + name: "TestServer", + version: "1.0", + capabilities: .init() // No logging capability + ) + + try await server.start(transport: serverTransport) + let initResult = try await client.connect(transport: clientTransport) + + // Verify logging capability is NOT advertised + #expect(initResult.capabilities.logging == nil) + + // Attempt to set logging level should fail in strict mode + await #expect(throws: MCPError.self) { + try await client.setLoggingLevel(.info) + } + + await client.disconnect() + await server.stop() + } + + // MARK: - Server Integration Tests + + @Test("Server log method sends notification") + func testServerLogMethod() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let client = Client(name: "TestClient", version: "1.0") + let server = Server( + name: "TestServer", + version: "1.0", + capabilities: .init(logging: .init()) + ) + + actor TestState { + var logMessages: [(level: LogLevel, logger: String?, data: Value)] = [] + func addLog(level: LogLevel, logger: String?, data: Value) { + logMessages.append((level, logger, data)) + } + func getLogs() -> [(level: LogLevel, logger: String?, data: Value)] { logMessages } + } + + let state = TestState() + + // Register handler for log notifications on client + await client.onNotification(LogMessageNotification.self) { message in + await state.addLog( + level: message.params.level, + logger: message.params.logger, + data: message.params.data + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Send a log message + let logData = Value.object([ + "message": Value.string("Test log"), + "count": Value.int(42) + ]) + + try await server.log(level: .info, logger: "test", data: logData) + + // Wait for message processing + try await Task.sleep(for: .milliseconds(100)) + + // Verify the notification was received + let logs = await state.getLogs() + #expect(logs.count == 1) + #expect(logs[0].level == LogLevel.info) + #expect(logs[0].logger == "test") + #expect(logs[0].data == logData) + + await client.disconnect() + await server.stop() + } + + @Test("Server log method with codable data") + func testServerLogMethodWithCodableData() async throws { + struct LogData: Codable, Hashable { + let message: String + let timestamp: String + let code: Int + } + + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let client = Client(name: "TestClient", version: "1.0") + let server = Server( + name: "TestServer", + version: "1.0", + capabilities: .init(logging: .init()) + ) + + actor TestState { + var logMessages: [(level: LogLevel, logger: String?, data: Value)] = [] + func addLog(level: LogLevel, logger: String?, data: Value) { + logMessages.append((level, logger, data)) + } + func getLogs() -> [(level: LogLevel, logger: String?, data: Value)] { logMessages } + } + + let state = TestState() + + // Register handler for log notifications on client + await client.onNotification(LogMessageNotification.self) { message in + await state.addLog( + level: message.params.level, + logger: message.params.logger, + data: message.params.data + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Send a log message with codable data + let logData = LogData( + message: "Error occurred", + timestamp: "2025-01-29T12:00:00Z", + code: 500 + ) + + try await server.log(level: .error, logger: "api", data: logData) + + // Wait for message processing + try await Task.sleep(for: .milliseconds(100)) + + // Verify the notification was received + let logs = await state.getLogs() + #expect(logs.count == 1) + #expect(logs[0].level == LogLevel.error) + #expect(logs[0].logger == "api") + + // Verify data content + if case .object(let dataDict) = logs[0].data { + #expect(dataDict["message"] == Value.string("Error occurred")) + #expect(dataDict["timestamp"] == Value.string("2025-01-29T12:00:00Z")) + #expect(dataDict["code"] == Value.int(500)) + } else { + Issue.record("Expected object data") + } + + await client.disconnect() + await server.stop() + } + + @Test("Server log without logger name") + func testServerLogWithoutLoggerName() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let client = Client(name: "TestClient", version: "1.0") + let server = Server( + name: "TestServer", + version: "1.0", + capabilities: .init(logging: .init()) + ) + + actor TestState { + var logMessages: [(level: LogLevel, logger: String?, data: Value)] = [] + func addLog(level: LogLevel, logger: String?, data: Value) { + logMessages.append((level, logger, data)) + } + func getLogs() -> [(level: LogLevel, logger: String?, data: Value)] { logMessages } + } + + let state = TestState() + + // Register handler for log notifications on client + await client.onNotification(LogMessageNotification.self) { message in + await state.addLog( + level: message.params.level, + logger: message.params.logger, + data: message.params.data + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Send a log message without logger name + let logData = Value.object(["message": Value.string("Generic log")]) + try await server.log(level: .debug, data: logData) + + // Wait for message processing + try await Task.sleep(for: .milliseconds(100)) + + // Verify the notification was received + let logs = await state.getLogs() + #expect(logs.count == 1) + #expect(logs[0].level == LogLevel.debug) + #expect(logs[0].logger == nil) + + await client.disconnect() + await server.stop() + } + + @Test("Multiple log levels sent correctly") + func testMultipleLogLevels() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let client = Client(name: "TestClient", version: "1.0") + let server = Server( + name: "TestServer", + version: "1.0", + capabilities: .init(logging: .init()) + ) + + actor TestState { + var logMessages: [(level: LogLevel, logger: String?, data: Value)] = [] + func addLog(level: LogLevel, logger: String?, data: Value) { + logMessages.append((level, logger, data)) + } + func getLogs() -> [(level: LogLevel, logger: String?, data: Value)] { logMessages } + } + + let state = TestState() + + // Register handler for log notifications on client + await client.onNotification(LogMessageNotification.self) { message in + await state.addLog( + level: message.params.level, + logger: message.params.logger, + data: message.params.data + ) + } + + try await server.start(transport: serverTransport) + _ = try await client.connect(transport: clientTransport) + + // Send log messages at different levels + try await server.log(level: .debug, data: Value.object(["msg": Value.string("Debug message")])) + try await server.log(level: .info, data: Value.object(["msg": Value.string("Info message")])) + try await server.log(level: .warning, data: Value.object(["msg": Value.string("Warning message")])) + try await server.log(level: .error, data: Value.object(["msg": Value.string("Error message")])) + try await server.log(level: .critical, data: Value.object(["msg": Value.string("Critical message")])) + + // Wait for message processing + try await Task.sleep(for: .milliseconds(200)) + + // Verify all notifications were received + let logs = await state.getLogs() + #expect(logs.count == 5) + #expect(logs[0].level == LogLevel.debug) + #expect(logs[1].level == LogLevel.info) + #expect(logs[2].level == LogLevel.warning) + #expect(logs[3].level == LogLevel.error) + #expect(logs[4].level == LogLevel.critical) + + await client.disconnect() + await server.stop() + } +} diff --git a/Tests/MCPTests/MetaFieldsTests.swift b/Tests/MCPTests/MetaFieldsTests.swift new file mode 100644 index 0000000..5b8fbbd --- /dev/null +++ b/Tests/MCPTests/MetaFieldsTests.swift @@ -0,0 +1,297 @@ +import Testing + +import class Foundation.JSONDecoder +import class Foundation.JSONEncoder +import class Foundation.JSONSerialization + +@testable import MCP + +@Suite("Meta Fields") +struct MetaFieldsTests { + private struct Payload: Codable, Hashable, Sendable { + let message: String + } + + private enum TestMethod: Method { + static let name = "test.general" + typealias Parameters = Payload + typealias Result = Payload + } + + @Test("Tool encoding and decoding with general fields") + func testToolGeneralFields() throws { + let meta = Metadata(additionalFields: [ + "vendor.example/outputTemplate": .string("ui://widget/kanban-board.html") + ]) + + let tool = Tool( + name: "kanban-board", + title: "Kanban Board", + description: "Display kanban widget", + inputSchema: try Value(["type": "object"]), + _meta: meta + ) + + let data = try JSONEncoder().encode(tool) + let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] + + let metaObject = json?["_meta"] as? [String: Any] + #expect( + metaObject?["vendor.example/outputTemplate"] as? String + == "ui://widget/kanban-board.html") + + let decoded = try JSONDecoder().decode(Tool.self, from: data) + #expect( + decoded._meta?["vendor.example/outputTemplate"] + == .string("ui://widget/kanban-board.html") + ) + } + + @Test("Meta keys allow nested prefixes") + func testMetaKeyNestedPrefixes() throws { + let meta = Metadata(additionalFields: [ + "vendor.example/toolInvocation/invoking": .bool(true) + ]) + + let tool = Tool( + name: "invoke", + description: "Invoke tool", + inputSchema: [:], + _meta: meta + ) + + let data = try JSONEncoder().encode(tool) + let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] + let metaObject = json?["_meta"] as? [String: Any] + #expect(metaObject?["vendor.example/toolInvocation/invoking"] as? Bool == true) + + let decoded = try JSONDecoder().decode(Tool.self, from: data) + #expect(decoded._meta?["vendor.example/toolInvocation/invoking"] == .bool(true)) + } + + @Test("Resource content encodes meta") + func testResourceContentGeneralFields() throws { + let meta = Metadata(additionalFields: [ + "vendor.example/widgetPrefersBorder": .bool(true) + ]) + + let content = Resource.Content.text( + "
Widget
", + uri: "ui://widget/kanban-board.html", + mimeType: "text/html", + _meta: meta + ) + + let data = try JSONEncoder().encode(content) + let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] + let metaObject = json?["_meta"] as? [String: Any] + + #expect(metaObject?["vendor.example/widgetPrefersBorder"] as? Bool == true) + + let decoded = try JSONDecoder().decode(Resource.Content.self, from: data) + #expect(decoded._meta?["vendor.example/widgetPrefersBorder"] == .bool(true)) + } + + @Test("Initialize.Result encoding with meta") + func testInitializeResultGeneralFields() throws { + let meta = Metadata(additionalFields: ["vendor.example/build": .string("v1.0.0")]) + + let result = Initialize.Result( + protocolVersion: "2024-11-05", + capabilities: Server.Capabilities(), + serverInfo: Server.Info(name: "test", version: "1.0"), + instructions: "Test server", + _meta: meta + ) + + let data = try JSONEncoder().encode(result) + let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] + + let metaObject = json?["_meta"] as? [String: Any] + #expect(metaObject?["vendor.example/build"] as? String == "v1.0.0") + + let decoded = try JSONDecoder().decode(Initialize.Result.self, from: data) + #expect(decoded._meta?["vendor.example/build"] == .string("v1.0.0")) + } + + @Test("ListTools.Result encoding with meta") + func testListToolsResultGeneralFields() throws { + let meta = Metadata(additionalFields: ["vendor.example/page": .int(1)]) + + let tool = Tool( + name: "test", + description: "A test tool", + inputSchema: try Value(["type": "object"]) + ) + + let result = ListTools.Result( + tools: [tool], + nextCursor: "page2", + _meta: meta + ) + + let data = try JSONEncoder().encode(result) + let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] + + let metaObject = json?["_meta"] as? [String: Any] + #expect(metaObject?["vendor.example/page"] as? Int == 1) + + let decoded = try JSONDecoder().decode(ListTools.Result.self, from: data) + #expect(decoded._meta?["vendor.example/page"] == .int(1)) + } + + @Test("CallTool.Result encoding with meta") + func testCallToolResultGeneralFields() throws { + let meta = Metadata(additionalFields: ["vendor.example/executionTime": .int(150)]) + + let result = CallTool.Result( + content: [.text(text: "Result data", annotations: nil, _meta: nil)], + isError: false, + _meta: meta + ) + + let data = try JSONEncoder().encode(result) + let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] + + let metaObject = json?["_meta"] as? [String: Any] + #expect(metaObject?["vendor.example/executionTime"] as? Int == 150) + + let decoded = try JSONDecoder().decode(CallTool.Result.self, from: data) + #expect(decoded._meta?["vendor.example/executionTime"] == .int(150)) + } + + @Test("ListResources.Result encoding with meta") + func testListResourcesResultGeneralFields() throws { + let meta = Metadata(additionalFields: ["vendor.example/cacheControl": .string("max-age=3600")]) + + let resource = Resource( + name: "test.txt", + uri: "file://test.txt", + description: "Test resource", + mimeType: "text/plain" + ) + + let result = ListResources.Result( + resources: [resource], + nextCursor: nil, + _meta: meta + ) + + let data = try JSONEncoder().encode(result) + let json = try JSONSerialization.jsonObject(with: data) as! [String: Any] + + let metaObject = json["_meta"] as! [String: Any] + #expect(metaObject["vendor.example/cacheControl"] as? String == "max-age=3600") + + let decoded = try JSONDecoder().decode(ListResources.Result.self, from: data) + #expect(decoded._meta?["vendor.example/cacheControl"] == Value.string("max-age=3600")) + } + + @Test("ReadResource.Result encoding with meta") + func testReadResourceResultGeneralFields() throws { + let meta = Metadata(additionalFields: ["vendor.example/encoding": .string("utf-8")]) + + let result = ReadResource.Result( + contents: [.text("file contents", uri: "file://test.txt")], + _meta: meta + ) + + let data = try JSONEncoder().encode(result) + let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] + + let metaObject = json?["_meta"] as? [String: Any] + #expect(metaObject?["vendor.example/encoding"] as? String == "utf-8") + + let decoded = try JSONDecoder().decode(ReadResource.Result.self, from: data) + #expect(decoded._meta?["vendor.example/encoding"] == .string("utf-8")) + } + + @Test("ListPrompts.Result encoding with meta") + func testListPromptsResultGeneralFields() throws { + let meta = Metadata(additionalFields: ["vendor.example/category": .string("system")]) + + let prompt = Prompt( + name: "greeting", + description: "A greeting prompt" + ) + + let result = ListPrompts.Result( + prompts: [prompt], + nextCursor: nil, + _meta: meta + ) + + let data = try JSONEncoder().encode(result) + let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] + + let metaObject = json?["_meta"] as? [String: Any] + #expect(metaObject?["vendor.example/category"] as? String == "system") + + let decoded = try JSONDecoder().decode(ListPrompts.Result.self, from: data) + #expect(decoded._meta?["vendor.example/category"] == .string("system")) + } + + @Test("GetPrompt.Result encoding with meta") + func testGetPromptResultGeneralFields() throws { + let meta = Metadata(additionalFields: ["vendor.example/version": .int(2)]) + + let message = Prompt.Message.user("Hello") + + let result = GetPrompt.Result( + description: "A test prompt", + messages: [message], + _meta: meta + ) + + let data = try JSONEncoder().encode(result) + let json = try JSONSerialization.jsonObject(with: data) as! [String: Any] + + let metaObject = json["_meta"] as! [String: Any] + #expect(metaObject["vendor.example/version"] as? Int == 2) + + let decoded = try JSONDecoder().decode(GetPrompt.Result.self, from: data) + #expect(decoded._meta?["vendor.example/version"] == Value.int(2)) + } + + @Test("CreateSamplingMessage.Result encoding with meta") + func testSamplingResultGeneralFields() throws { + let meta = Metadata(additionalFields: ["vendor.example/model-version": .string("gpt-4-0613")]) + + let result = CreateSamplingMessage.Result( + model: "gpt-4", + stopReason: .endTurn, + role: .assistant, + content: .text("Hello!"), + _meta: meta + ) + + let data = try JSONEncoder().encode(result) + let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] + + let metaObject = json?["_meta"] as? [String: Any] + #expect(metaObject?["vendor.example/model-version"] as? String == "gpt-4-0613") + + let decoded = try JSONDecoder().decode(CreateSamplingMessage.Result.self, from: data) + #expect(decoded._meta?["vendor.example/model-version"] == .string("gpt-4-0613")) + } + + @Test("CreateElicitation.Result encoding with meta") + func testElicitationResultGeneralFields() throws { + let meta = Metadata(additionalFields: ["vendor.example/timestamp": .int(1_640_000_000)]) + + let result = CreateElicitation.Result( + action: .accept, + content: ["response": .string("user input")], + _meta: meta + ) + + let data = try JSONEncoder().encode(result) + let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] + + let metaObject = json?["_meta"] as? [String: Any] + #expect(metaObject?["vendor.example/timestamp"] as? Int == 1_640_000_000) + + let decoded = try JSONDecoder().decode(CreateElicitation.Result.self, from: data) + #expect(decoded._meta?["vendor.example/timestamp"] == .int(1_640_000_000)) + } +} diff --git a/Tests/MCPTests/MockResponses.swift b/Tests/MCPTests/MockResponses.swift new file mode 100644 index 0000000..1c7728c --- /dev/null +++ b/Tests/MCPTests/MockResponses.swift @@ -0,0 +1,216 @@ +@preconcurrency import Foundation +import Testing + +@testable import MCP + +#if canImport(FoundationNetworking) + import FoundationNetworking +#endif + +#if swift(>=6.1) + + enum MockResponses { + typealias Route = @Sendable (URLRequest) async throws -> (HTTPURLResponse, Data) + + static func ephemeralConfiguration() -> URLSessionConfiguration { + let config = URLSessionConfiguration.ephemeral + config.protocolClasses = [MockURLProtocol.self] + return config + } + + static func mockError(_ message: String) -> NSError { + NSError( + domain: "MockURLProtocolError", code: 0, + userInfo: [NSLocalizedDescriptionKey: message] + ) + } + + static func jsonRPCResult(id: Int) -> Data { + #"{"jsonrpc":"2.0","result":{"ok":true},"id":\#(id)}"#.data(using: .utf8)! + } + + // MARK: - Route Builders + + static func jsonSuccess(body: Data) -> Route { + { request in + let response = HTTPURLResponse( + url: request.url!, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "application/json"])! + return (response, body) + } + } + + static func bearerChallenge( + statusCode: Int = 401, + resourceMetadataURL: URL? = nil, + scope: String? = nil, + error: String? = nil, + errorDescription: String? = nil + ) -> Route { + { request in + var params: [String] = [] + if let url = resourceMetadataURL { + params.append("resource_metadata=\"\(url.absoluteString)\"") + } + if let scope { params.append("scope=\"\(scope)\"") } + if let error { params.append("error=\"\(error)\"") } + if let errorDescription { params.append("error_description=\"\(errorDescription)\"") } + let headerValue = "Bearer \(params.joined(separator: ", "))" + let response = HTTPURLResponse( + url: request.url!, statusCode: statusCode, httpVersion: "HTTP/1.1", + headerFields: ["WWW-Authenticate": headerValue])! + return (response, Data()) + } + } + + static func resourceMetadata( + authorizationServers: [String], + scopesSupported: [String]? = nil, + resource: String? = nil + ) -> Route { + { request in + var dict: [String: Any] = ["authorization_servers": authorizationServers] + if let scopes = scopesSupported { dict["scopes_supported"] = scopes } + if let resource { dict["resource"] = resource } + let data = try JSONSerialization.data(withJSONObject: dict) + let response = HTTPURLResponse( + url: request.url!, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "application/json"])! + return (response, data) + } + } + + static func asMetadata( + issuer: String, + tokenEndpoint: String, + authorizationEndpoint: String? = nil, + registrationEndpoint: String? = nil, + codeChallengeMethodsSupported: [String]? = nil, + tokenEndpointAuthMethodsSupported: [String]? = nil, + clientIDMetadataDocumentSupported: Bool? = nil + ) -> Route { + { request in + var dict: [String: Any] = ["issuer": issuer, "token_endpoint": tokenEndpoint] + if let v = authorizationEndpoint { dict["authorization_endpoint"] = v } + if let v = registrationEndpoint { dict["registration_endpoint"] = v } + if let v = codeChallengeMethodsSupported { + dict["code_challenge_methods_supported"] = v + } + if let v = tokenEndpointAuthMethodsSupported { + dict["token_endpoint_auth_methods_supported"] = v + } + if let v = clientIDMetadataDocumentSupported { + dict["client_id_metadata_document_supported"] = v + } + let data = try JSONSerialization.data(withJSONObject: dict) + let response = HTTPURLResponse( + url: request.url!, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "application/json"])! + return (response, data) + } + } + + static func tokenSuccess( + accessToken: String, + expiresIn: Int = 3600, + scope: String? = nil, + refreshToken: String? = nil + ) -> Route { + { request in + var dict: [String: Any] = [ + "access_token": accessToken, "token_type": "Bearer", "expires_in": expiresIn, + ] + if let scope { dict["scope"] = scope } + if let refreshToken { dict["refresh_token"] = refreshToken } + let data = try JSONSerialization.data(withJSONObject: dict) + let response = HTTPURLResponse( + url: request.url!, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "application/json"])! + return (response, data) + } + } + + static func tokenResponse( + accessToken: String, + tokenType: String, + expiresIn: Int = 3600 + ) -> Route { + { request in + let dict: [String: Any] = [ + "access_token": accessToken, "token_type": tokenType, "expires_in": expiresIn, + ] + let data = try JSONSerialization.data(withJSONObject: dict) + let response = HTTPURLResponse( + url: request.url!, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "application/json"])! + return (response, data) + } + } + + static func tokenError( + statusCode: Int = 400, + error: String, + errorDescription: String? = nil, + extraFields: [String: String] = [:] + ) -> Route { + { request in + var dict: [String: Any] = ["error": error] + if let errorDescription { dict["error_description"] = errorDescription } + for (key, value) in extraFields { dict[key] = value } + let data = try JSONSerialization.data(withJSONObject: dict) + let response = HTTPURLResponse( + url: request.url!, statusCode: statusCode, httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "application/json"])! + return (response, data) + } + } + + static func httpError(statusCode: Int) -> Route { + { request in + let response = HTTPURLResponse( + url: request.url!, statusCode: statusCode, httpVersion: "HTTP/1.1", + headerFields: nil)! + return (response, Data()) + } + } + + static func registrationSuccess( + clientID: String, + clientSecret: String? = nil + ) -> Route { + { request in + var dict: [String: Any] = ["client_id": clientID] + if let clientSecret { dict["client_secret"] = clientSecret } + let data = try JSONSerialization.data(withJSONObject: dict) + let response = HTTPURLResponse( + url: request.url!, statusCode: 201, httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "application/json"])! + return (response, data) + } + } + + static func redirect(to location: String, statusCode: Int = 302) -> Route { + { request in + let response = HTTPURLResponse( + url: request.url!, statusCode: statusCode, httpVersion: "HTTP/1.1", + headerFields: ["Location": location])! + return (response, Data()) + } + } + + // MARK: - Routing Handler + + static func routingHandler( + routes: [URL: Route] + ) -> @Sendable (URLRequest) async throws -> (HTTPURLResponse, Data) { + { request in + guard let url = request.url else { throw mockError("Missing request URL") } + guard let handler = routes[url] else { + throw mockError("Unexpected URL: \(url.absoluteString)") + } + return try await handler(request) + } + } + } + +#endif // swift(>=6.1) diff --git a/Tests/MCPTests/NetworkTransportTests.swift b/Tests/MCPTests/NetworkTransportTests.swift new file mode 100644 index 0000000..9c02c90 --- /dev/null +++ b/Tests/MCPTests/NetworkTransportTests.swift @@ -0,0 +1,665 @@ +import Foundation +import Logging +import Testing + +@testable import MCP + +#if canImport(Network) + import Network + + /// A mock implementation of NetworkConnectionProtocol for testing + final class MockNetworkConnection: NetworkConnectionProtocol, @unchecked Sendable { + /// Current state of the connection + private var mockState: NWConnection.State = .setup + + /// Error to be returned on send/receive operations + private var mockError: Swift.Error? + + /// Data queue for testing + private var dataToReceive: [Data] = [] + private var sentData: [Data] = [] + + /// The state update handler + public var stateUpdateHandler: ((@Sendable (NWConnection.State) -> Void))? + + /// Current state + var state: NWConnection.State { + return mockState + } + + /// Initialize a mock connection + init() {} + + /// Start the connection + func start(queue: DispatchQueue) { + // Simulate successful connection by default + Task { @MainActor in + if mockError != nil { + self.updateState(mockState) + } else { + self.updateState(.ready) + } + } + } + + /// Send data through the connection + func send( + content: Data?, + contentContext: NWConnection.ContentContext, + isComplete: Bool, + completion: NWConnection.SendCompletion + ) { + if let content = content { + sentData.append(content) + } + + switch completion { + case .contentProcessed(let handler): + Task { @MainActor in + handler(self.mockError as? NWError) + } + default: + break + } + } + + /// Receive data from the connection + func receive( + minimumIncompleteLength: Int, + maximumLength: Int, + completion: @escaping @Sendable ( + Data?, NWConnection.ContentContext?, Bool, NWError? + ) -> Void + ) { + Task { @MainActor in + if self.mockState == .cancelled { + completion( + nil, nil, true, + NWError.posix(POSIXErrorCode.ECANCELED) + ) + return + } + + if let error = self.mockError { + completion(nil, nil, false, error as? NWError) + return + } + + if self.dataToReceive.isEmpty { + completion(Data(), nil, false, nil) + return + } + + let data = self.dataToReceive.removeFirst() + completion(data, nil, self.dataToReceive.isEmpty, nil) + } + } + + /// Cancel the connection + func cancel() { + updateState(.cancelled) + } + + // Test helpers + + /// Simulate the connection becoming ready + func simulateReady() { + updateState(.ready) + } + + /// Simulate the connection becoming preparing + func simulatePreparing() { + updateState(.preparing) + } + + /// Simulate a connection failure + func simulateFailure( + error: Swift.Error? = nil + ) { + mockError = error + if let nwError = error as? NWError { + updateState(.failed(nwError)) + } else { + updateState(.failed(NWError.posix(POSIXErrorCode(rawValue: 57)!))) + } + } + + /// Simulate connection cancellation + func simulateCancellation() { + updateState(.cancelled) + } + + /// Update the connection state and notify handler + private func updateState(_ newState: NWConnection.State) { + mockState = newState + Task { @MainActor in + self.stateUpdateHandler?(newState) + } + } + + /// Queue data to be received + func queueDataForReceiving(_ data: Data) { + dataToReceive.append(data) + } + + /// Queue a heartbeat message to be received + func queueHeartbeat() { + // Create a mock heartbeat message that matches the format + let magicBytes: [UInt8] = [0xF0, 0x9F, 0x92, 0x93] // Magic bytes for heartbeat + var data = Data(magicBytes) + let timestamp = UInt64(Date().timeIntervalSinceReferenceDate * 1000) + withUnsafeBytes(of: timestamp) { buffer in + data.append(contentsOf: buffer) + } + queueDataForReceiving(data) + } + + /// Queue text message to be received + func queueTextMessage(_ text: String) { + guard let data = text.data(using: .utf8) else { return } + queueDataForReceiving(data) + } + + /// Get all sent data + func getSentData() -> [Data] { + return sentData + } + + /// Clear sent data buffer + func clearSentData() { + sentData.removeAll() + } + } + + @Suite("Network Transport Tests", .serialized) + struct NetworkTransportTests { + @Test("Heartbeat Creation And Parsing") + func testHeartbeatCreationAndParsing() { + // Create a heartbeat + let heartbeat = NetworkTransport.Heartbeat() + + // Convert to data and back + let data = heartbeat.data + let parsed = NetworkTransport.Heartbeat.from(data: data) + + #expect(parsed != nil) + + // Time should be very close (within 1 second) + if let parsed = parsed { + let timeDifference = abs(parsed.timestamp.timeIntervalSince(heartbeat.timestamp)) + #expect(timeDifference < 1.0) + } + + // Test invalid data + let invalidData = Data([0x01, 0x02, 0x03]) + #expect(NetworkTransport.Heartbeat.from(data: invalidData) == nil) + #expect(NetworkTransport.Heartbeat.isHeartbeat(invalidData) == false) + #expect(NetworkTransport.Heartbeat.isHeartbeat(data) == true) + } + + @Test("Reconnection Configuration") + func testReconnectionConfiguration() { + // Create custom config + let config = NetworkTransport.ReconnectionConfiguration( + enabled: true, + maxAttempts: 3, + backoffMultiplier: 2.0 + ) + + #expect(config.enabled == true) + #expect(config.maxAttempts == 3) + #expect(config.backoffMultiplier == 2.0) + + // Test backoff delay calculation + let firstDelay = config.backoffDelay(for: 1) + let secondDelay = config.backoffDelay(for: 2) + let thirdDelay = config.backoffDelay(for: 3) + + // Check delays are approximately correct (within 0.001) + #expect(abs(firstDelay - 0.5) < 0.001) + #expect(abs(secondDelay - 1.0) < 0.001) + #expect(abs(thirdDelay - 2.0) < 0.001) + + // Test disabled config + let disabledConfig = NetworkTransport.ReconnectionConfiguration.disabled + #expect(disabledConfig.enabled == false) + } + + @Test("Heartbeat Configuration") + func testHeartbeatConfiguration() { + // Create custom config + let config = NetworkTransport.HeartbeatConfiguration( + enabled: true, + interval: 5.0 + ) + + #expect(config.enabled == true) + #expect(config.interval == 5.0) + + // Test default config + let defaultConfig = NetworkTransport.HeartbeatConfiguration.default + #expect(defaultConfig.enabled == true) + #expect(defaultConfig.interval == 15.0) + + // Test disabled config + let disabledConfig = NetworkTransport.HeartbeatConfiguration.disabled + #expect(disabledConfig.enabled == false) + } + + @Test("Connect Success") + func testNetworkTransportConnectSuccess() async throws { + let mockConnection = MockNetworkConnection() + let transport = NetworkTransport( + mockConnection, + heartbeatConfig: .disabled // Disable heartbeats for simplified testing + ) + + try await transport.connect() + + // Verify connection state + #expect(mockConnection.state == .ready) + + await transport.disconnect() + } + + @Test("Connect Failure") + func testNetworkTransportConnectFailure() async { + let mockConnection = MockNetworkConnection() + let transport = NetworkTransport( + mockConnection, + reconnectionConfig: .disabled // Disable reconnection for this test + ) + + // Simulate failure before connecting + mockConnection.simulateFailure(error: NWError.posix(POSIXErrorCode.ECONNRESET)) + + do { + try await transport.connect() + Issue.record("Expected connect to throw an error") + } catch let error as MCPError { + // Expected failure + #expect(error.localizedDescription.contains("Connection failed")) + } catch let error as NWError { + // Also accept NWError since it's the underlying error + #expect((error as NSError).code == POSIXErrorCode.ECONNRESET.rawValue) + } catch { + Issue.record("Unexpected error type: \(type(of: error))") + } + + await transport.disconnect() + } + + @Test("Send Message") + func testNetworkTransportSendMessage() async throws { + let mockConnection = MockNetworkConnection() + let transport = NetworkTransport( + mockConnection, + heartbeatConfig: .disabled + ) + + try await transport.connect() + + // Test sending a simple message + let message = #"{"key":"value"}"#.data(using: .utf8)! + try await transport.send(message) + + // Verify the message was sent with a newline delimiter + let sentData = mockConnection.getSentData() + #expect(sentData.count == 1) + + if sentData.count > 0 { + let expectedOutput = message + "\n".data(using: .utf8)! + #expect(sentData[0] == expectedOutput) + } + + await transport.disconnect() + } + + @Test("Receive Message") + func testNetworkTransportReceiveMessage() async throws { + let mockConnection = MockNetworkConnection() + let transport = NetworkTransport( + mockConnection, + heartbeatConfig: .disabled + ) + + // Queue a message to be received + let message = #"{"key":"value"}"# + let messageWithNewline = message + "\n" + mockConnection.queueTextMessage(messageWithNewline) + + try await transport.connect() + + // Start receiving messages + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + + // Get first message + let received = try await iterator.next() + #expect(received != nil) + + if let received = received { + #expect(received == message.data(using: .utf8)!) + } + + await transport.disconnect() + } + + @Test("Heartbeat Send and Receive") + func testNetworkTransportHeartbeat() async throws { + let mockConnection = MockNetworkConnection() + + // Create transport with rapid heartbeats + let heartbeatConfig = NetworkTransport.HeartbeatConfiguration( + enabled: true, + interval: 0.1 // Short interval for testing + ) + + let transport = NetworkTransport( + mockConnection, + heartbeatConfig: heartbeatConfig + ) + + try await transport.connect() + + // Wait for initial connection setup + try await Task.sleep(for: .milliseconds(100)) + + // Wait for initial heartbeat delay (1 second) plus a small buffer + try await Task.sleep(for: .seconds(1.2)) + + // Check if heartbeat was sent + let sentData = mockConnection.getSentData() + #expect(sentData.count >= 1, "No heartbeat was sent after \(sentData.count) attempts") + + if let firstSent = sentData.first { + #expect( + NetworkTransport.Heartbeat.isHeartbeat(firstSent), + "Sent data is not a heartbeat") + } + + // Queue a heartbeat to be received + mockConnection.queueHeartbeat() + + // Wait for heartbeat processing + try await Task.sleep(for: .milliseconds(100)) + + await transport.disconnect() + } + + @Test("Reconnection") + func testNetworkTransportReconnection() async throws { + let mockConnection = MockNetworkConnection() + + // Configure for quick reconnection + let reconnectionConfig = NetworkTransport.ReconnectionConfiguration( + enabled: true, + maxAttempts: 2, + backoffMultiplier: 1.0 + ) + + let transport = NetworkTransport( + mockConnection, + heartbeatConfig: .disabled, + reconnectionConfig: reconnectionConfig + ) + + try await transport.connect() + + // Simulate connection failure during operation + mockConnection.simulateFailure(error: NWError.posix(POSIXErrorCode.ECONNRESET)) + + // Wait a bit to ensure failure is processed + try await Task.sleep(for: .milliseconds(100)) + + // Try to send message after failure - should trigger reconnection process + let message = #"{"test":"reconnect"}"#.data(using: .utf8)! + + do { + try await transport.send(message) + Issue.record("Expected send to fail after connection lost") + } catch { + // Expected error + #expect(error is MCPError, "Expected MCPError but got \(type(of: error))") + } + + // Wait for potential reconnection attempt + try await Task.sleep(for: .milliseconds(600)) + + await transport.disconnect() + } + + @Test("Multiple Messages") + func testNetworkTransportMultipleMessages() async throws { + let mockConnection = MockNetworkConnection() + let transport = NetworkTransport( + mockConnection, + heartbeatConfig: .disabled + ) + + // Queue multiple messages + let messages = [ + #"{"id":1,"method":"test1"}"#, + #"{"id":2,"method":"test2"}"#, + #"{"id":3,"method":"test3"}"#, + ] + + for message in messages { + mockConnection.queueTextMessage(message + "\n") + } + + try await transport.connect() + + // Receive and verify all messages + let stream = await transport.receive() + var receiveCount = 0 + + for try await data in stream { + if let receivedStr = String(data: data, encoding: .utf8) { + #expect(messages.contains(receivedStr)) + receiveCount += 1 + + if receiveCount >= messages.count { + break + } + } + } + + await transport.disconnect() + } + + @Test("Disconnect During Receive") + func testNetworkTransportDisconnectDuringReceive() async throws { + let mockConnection = MockNetworkConnection() + let transport = NetworkTransport( + mockConnection, + heartbeatConfig: .disabled + ) + + try await transport.connect() + + // Start a task to receive messages + let receiveTask = Task { + var count = 0 + for try await _ in await transport.receive() { + count += 1 + if count > 10 { + // Prevent infinite loop in test + break + } + } + } + + // Let the receive loop start + try await Task.sleep(for: .milliseconds(100)) + + // Disconnect while receiving + await transport.disconnect() + + // Wait for the receive task to complete + _ = await receiveTask.result + } + + @Test("Partial Message Reception") + func testPartialMessageReception() async throws { + let mockConnection = MockNetworkConnection() + let transport = NetworkTransport( + mockConnection, + heartbeatConfig: .disabled + ) + + try await transport.connect() + + // Split a message into multiple parts + let message = #"{"key":"value"}"# + let parts = [ + message.prefix(5).data(using: .utf8)!, + message.dropFirst(5).data(using: .utf8)!, + "\n".data(using: .utf8)!, + ] + + // Queue the parts + for part in parts { + mockConnection.queueDataForReceiving(part) + } + + // Start receiving messages + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + + // Get the complete message + let received = try await iterator.next() + #expect(received != nil) + if let received = received { + #expect(received == message.data(using: .utf8)!) + } + + await transport.disconnect() + } + + @Test("Large Message Handling") + func testLargeMessageHandling() async throws { + let mockConnection = MockNetworkConnection() + let transport = NetworkTransport( + mockConnection, + heartbeatConfig: .disabled + ) + + try await transport.connect() + + // Create a large message (larger than typical receive buffer) + let largeMessage = String(repeating: "x", count: 100_000) + let messageWithNewline = largeMessage + "\n" + mockConnection.queueTextMessage(messageWithNewline) + + // Start receiving messages + let stream = await transport.receive() + var iterator = stream.makeAsyncIterator() + + // Get the message + let received = try await iterator.next() + #expect(received != nil) + if let received = received { + #expect(received.count == largeMessage.count) + } + + await transport.disconnect() + } + + @Test("Reconnection Backoff") + func testReconnectionBackoff() async throws { + let mockConnection = MockNetworkConnection() + let startTime = Date() + + // Configure for quick reconnection with known backoff + let reconnectionConfig = NetworkTransport.ReconnectionConfiguration( + enabled: true, + maxAttempts: 3, + backoffMultiplier: 2.0 + ) + + let transport = NetworkTransport( + mockConnection, + heartbeatConfig: .disabled, + reconnectionConfig: reconnectionConfig + ) + + try await transport.connect() + + // Simulate failure + mockConnection.simulateFailure(error: NWError.posix(POSIXErrorCode.ECONNRESET)) + + // Wait for reconnection attempts + try await Task.sleep(for: .seconds(4)) + + let endTime = Date() + let duration = endTime.timeIntervalSince(startTime) + + // Verify we had enough time for reconnection attempts + #expect(duration >= 3.5) // Should have time for 3 attempts with backoff + + await transport.disconnect() + } + + @Test("Heartbeat Failure Handling") + func testHeartbeatFailureHandling() async throws { + let mockConnection = MockNetworkConnection() + + // Create transport with rapid heartbeats + let heartbeatConfig = NetworkTransport.HeartbeatConfiguration( + enabled: true, + interval: 0.1 + ) + + let transport = NetworkTransport( + mockConnection, + heartbeatConfig: heartbeatConfig + ) + + try await transport.connect() + + // Wait for initial heartbeat + try await Task.sleep(for: .seconds(1.2)) + + // Simulate failure during heartbeat + mockConnection.simulateFailure(error: NWError.posix(POSIXErrorCode.ECONNRESET)) + + // Wait for potential recovery + try await Task.sleep(for: .milliseconds(500)) + + // Verify connection state + if case .failed = mockConnection.state { + // expected + } else { + Issue.record("Expected state to be failed") + } + + await transport.disconnect() + } + + @Test("Resource Cleanup") + func testResourceCleanup() async throws { + weak var weakConnection: MockNetworkConnection? + + do { + let mockConnection = MockNetworkConnection() + weakConnection = mockConnection + + // Create and use transport in a separate scope + let transport = NetworkTransport( + mockConnection, + heartbeatConfig: .disabled + ) + + try await transport.connect() + await transport.disconnect() + } + + // Wait for potential async cleanup + try await Task.sleep(for: .milliseconds(100)) + + // Verify connection is cleaned up + #expect(weakConnection == nil, "Connection was not properly cleaned up") + } + } +#endif diff --git a/Tests/MCPTests/NotificationTests.swift b/Tests/MCPTests/NotificationTests.swift index 39886e3..4bdd187 100644 --- a/Tests/MCPTests/NotificationTests.swift +++ b/Tests/MCPTests/NotificationTests.swift @@ -86,4 +86,87 @@ struct NotificationTests { #expect(decoded.method == InitializedNotification.name) } + + @Test("Resource updated notification with parameters") + func testResourceUpdatedNotification() throws { + let params = ResourceUpdatedNotification.Parameters(uri: "test://resource") + let notification = ResourceUpdatedNotification.message(params) + + #expect(notification.method == ResourceUpdatedNotification.name) + #expect(notification.params.uri == "test://resource") + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(notification) + + // Verify the exact JSON structure + let json = try JSONDecoder().decode([String: Value].self, from: data) + #expect(json["jsonrpc"] == "2.0") + #expect(json["method"] == "notifications/resources/updated") + #expect(json["params"] != nil) + #expect(json.count == 3, "Should contain jsonrpc, method, and params fields") + + // Verify we can decode it back + let decoded = try decoder.decode(Message.self, from: data) + #expect(decoded.method == ResourceUpdatedNotification.name) + #expect(decoded.params.uri == "test://resource") + } + + @Test("AnyNotification decoding - without params") + func testAnyNotificationDecodingWithoutParams() throws { + // Test decoding when params field is missing + let jsonString = """ + {"jsonrpc":"2.0","method":"notifications/initialized"} + """ + let data = jsonString.data(using: .utf8)! + + let decoder = JSONDecoder() + let decoded = try decoder.decode(AnyMessage.self, from: data) + + #expect(decoded.method == InitializedNotification.name) + } + + @Test("AnyNotification decoding - with null params") + func testAnyNotificationDecodingWithNullParams() throws { + // Test decoding when params field is null + let jsonString = """ + {"jsonrpc":"2.0","method":"notifications/initialized","params":null} + """ + let data = jsonString.data(using: .utf8)! + + let decoder = JSONDecoder() + let decoded = try decoder.decode(AnyMessage.self, from: data) + + #expect(decoded.method == InitializedNotification.name) + } + + @Test("AnyNotification decoding - with empty params") + func testAnyNotificationDecodingWithEmptyParams() throws { + // Test decoding when params field is empty + let jsonString = """ + {"jsonrpc":"2.0","method":"notifications/initialized","params":{}} + """ + let data = jsonString.data(using: .utf8)! + + let decoder = JSONDecoder() + let decoded = try decoder.decode(AnyMessage.self, from: data) + + #expect(decoded.method == InitializedNotification.name) + } + + @Test("AnyNotification decoding - with non-empty params") + func testAnyNotificationDecodingWithNonEmptyParams() throws { + // Test decoding when params field has values + let jsonString = """ + {"jsonrpc":"2.0","method":"notifications/resources/updated","params":{"uri":"test://resource"}} + """ + let data = jsonString.data(using: .utf8)! + + let decoder = JSONDecoder() + let decoded = try decoder.decode(AnyMessage.self, from: data) + + #expect(decoded.method == ResourceUpdatedNotification.name) + #expect(decoded.params.objectValue?["uri"]?.stringValue == "test://resource") + } } diff --git a/Tests/MCPTests/OAuthAuthorizationCodeFlowTests.swift b/Tests/MCPTests/OAuthAuthorizationCodeFlowTests.swift new file mode 100644 index 0000000..85c610e --- /dev/null +++ b/Tests/MCPTests/OAuthAuthorizationCodeFlowTests.swift @@ -0,0 +1,178 @@ +import Foundation +import Testing + +@testable import MCP + +#if canImport(FoundationNetworking) + import FoundationNetworking +#endif + +@Suite("OAuthAuthorizationCodeFlow") +struct OAuthAuthorizationCodeFlowTests { + + let flow = OAuthAuthorizationCodeFlow() + let authEndpoint = URL(string: "https://auth.example.com/authorize")! + let resource = URL(string: "https://api.example.com")! + let redirectURI = URL(string: "https://app.example.com/callback")! + let scopeSelector = DefaultOAuthScopeSelector() + + // MARK: - buildURL + + @Test("buildURL includes all required parameters") + func testBuildURLRequiredParams() throws { + let url = try flow.buildURL( + authorizationEndpoint: authEndpoint, + resource: resource, + redirectURI: redirectURI, + clientID: "my-client", + codeChallenge: "abc123", + scopes: nil, + state: "state-xyz", + scopeSerializer: scopeSelector + ) + let components = URLComponents(url: url, resolvingAgainstBaseURL: false)! + let items = Dictionary( + uniqueKeysWithValues: components.queryItems!.map { ($0.name, $0.value ?? "") }) + + #expect(items["response_type"] == "code") + #expect(items["client_id"] == "my-client") + #expect(items["redirect_uri"] == redirectURI.absoluteString) + #expect(items["state"] == "state-xyz") + #expect(items["resource"] == resource.absoluteString) + #expect(items["code_challenge"] == "abc123") + #expect(items["code_challenge_method"] == "S256") + } + + @Test("buildURL includes scope when provided") + func testBuildURLWithScope() throws { + let url = try flow.buildURL( + authorizationEndpoint: authEndpoint, + resource: resource, + redirectURI: redirectURI, + clientID: "my-client", + codeChallenge: "abc123", + scopes: Set(["read", "write"]), + state: "state-xyz", + scopeSerializer: scopeSelector + ) + let components = URLComponents(url: url, resolvingAgainstBaseURL: false)! + let items = Dictionary( + uniqueKeysWithValues: components.queryItems!.map { ($0.name, $0.value ?? "") }) + + let scope = items["scope"] ?? "" + #expect(scope.contains("read")) + #expect(scope.contains("write")) + } + + @Test("buildURL omits scope for nil scopes") + func testBuildURLOmitsScopeWhenNil() throws { + let url = try flow.buildURL( + authorizationEndpoint: authEndpoint, + resource: resource, + redirectURI: redirectURI, + clientID: "my-client", + codeChallenge: "abc123", + scopes: nil, + state: "state-xyz", + scopeSerializer: scopeSelector + ) + let components = URLComponents(url: url, resolvingAgainstBaseURL: false)! + let items = components.queryItems ?? [] + #expect(!items.contains(where: { $0.name == "scope" })) + } + + @Test("buildURL omits scope for empty scope set") + func testBuildURLOmitsScopeWhenEmpty() throws { + let url = try flow.buildURL( + authorizationEndpoint: authEndpoint, + resource: resource, + redirectURI: redirectURI, + clientID: "my-client", + codeChallenge: "abc123", + scopes: Set(), + state: "state-xyz", + scopeSerializer: scopeSelector + ) + let components = URLComponents(url: url, resolvingAgainstBaseURL: false)! + let items = components.queryItems ?? [] + #expect(!items.contains(where: { $0.name == "scope" })) + } + + // MARK: - extractCode + + @Test("extractCode returns code from valid redirect URL") + func testExtractCodeSuccess() throws { + let redirectURL = URL(string: + "https://app.example.com/callback?code=auth-code-123&state=my-state")! + let code = try flow.extractCode( + from: redirectURL, + expectedRedirectURI: redirectURI, + expectedState: "my-state" + ) + #expect(code == "auth-code-123") + } + + @Test("extractCode throws state mismatch") + func testExtractCodeThrowsStateMismatch() { + let redirectURL = URL(string: + "https://app.example.com/callback?code=auth-code-123&state=wrong-state")! + #expect(throws: OAuthAuthorizationError.self) { + try flow.extractCode( + from: redirectURL, + expectedRedirectURI: redirectURI, + expectedState: "my-state" + ) + } + } + + @Test("extractCode throws missing state") + func testExtractCodeThrowsMissingState() { + let redirectURL = URL(string: + "https://app.example.com/callback?code=auth-code-123")! + #expect(throws: OAuthAuthorizationError.self) { + try flow.extractCode( + from: redirectURL, + expectedRedirectURI: redirectURI, + expectedState: "my-state" + ) + } + } + + @Test("extractCode throws missing code") + func testExtractCodeThrowsMissingCode() { + let redirectURL = URL(string: + "https://app.example.com/callback?state=my-state")! + #expect(throws: OAuthAuthorizationError.self) { + try flow.extractCode( + from: redirectURL, + expectedRedirectURI: redirectURI, + expectedState: "my-state" + ) + } + } + + @Test("extractCode throws redirect URI mismatch") + func testExtractCodeThrowsRedirectMismatch() { + let redirectURL = URL(string: + "https://evil.example.com/callback?code=auth-code-123&state=my-state")! + #expect(throws: OAuthAuthorizationError.self) { + try flow.extractCode( + from: redirectURL, + expectedRedirectURI: redirectURI, + expectedState: "my-state" + ) + } + } + + @Test("extractCode normalizes host case in redirect URI comparison") + func testExtractCodeCaseInsensitiveHost() throws { + let redirectURL = URL(string: + "https://APP.EXAMPLE.COM/callback?code=auth-code-123&state=my-state")! + let code = try flow.extractCode( + from: redirectURL, + expectedRedirectURI: redirectURI, + expectedState: "my-state" + ) + #expect(code == "auth-code-123") + } +} diff --git a/Tests/MCPTests/OAuthAuthorizationTests.swift b/Tests/MCPTests/OAuthAuthorizationTests.swift new file mode 100644 index 0000000..b28026f --- /dev/null +++ b/Tests/MCPTests/OAuthAuthorizationTests.swift @@ -0,0 +1,735 @@ +import Foundation +import Testing + +#if canImport(CryptoKit) +import CryptoKit +#endif + +@testable import MCP + +@Suite("OAuth Authorization Helpers") +struct OAuthAuthorizationTests { + #if canImport(CryptoKit) + // Generated fresh each test run — no hardcoded key material in source. + private static let testPrivateKeyPEM: String = P256.Signing.PrivateKey().pemRepresentation + #else + private static let testPrivateKeyPEM: String = "" + #endif + + private var metadataDiscovery: DefaultOAuthMetadataDiscovery { DefaultOAuthMetadataDiscovery() } + + private static func decodeBase64URL(_ input: Substring) -> Data? { + var base64 = String(input) + .replacingOccurrences(of: "-", with: "+") + .replacingOccurrences(of: "_", with: "/") + let remainder = base64.count % 4 + if remainder != 0 { + base64 += String(repeating: "=", count: 4 - remainder) + } + return Data(base64Encoded: base64) + } + + @Test("Parse Bearer challenge with resource metadata and scope") + func testParseBearerChallenge() { + let headers = [ + "WWW-Authenticate": + "Bearer resource_metadata=\"https://mcp.example.com/.well-known/oauth-protected-resource\", scope=\"files:read files:write\", error=\"insufficient_scope\"" + ] + + let challenge = DefaultOAuthWWWAuthenticateParser().parseBearer(from: headers) + + #expect(challenge != nil) + #expect( + challenge?.resourceMetadataURL + == URL(string: "https://mcp.example.com/.well-known/oauth-protected-resource")) + #expect(challenge?.scope == "files:read files:write") + #expect(challenge?.error == "insufficient_scope") + } + + @Test("Parse Bearer challenge with optional error description") + func testParseBearerChallengeErrorDescription() { + let headers = [ + "WWW-Authenticate": + "Bearer error=\"insufficient_scope\", scope=\"files:read files:write\", resource_metadata=\"https://mcp.example.com/.well-known/oauth-protected-resource\", error_description=\"Additional file write permission required\"" + ] + + let challenge = DefaultOAuthWWWAuthenticateParser().parseBearer(from: headers) + + #expect(challenge != nil) + #expect(challenge?.error == "insufficient_scope") + #expect(challenge?.scope == "files:read files:write") + #expect( + challenge?.resourceMetadataURL + == URL(string: "https://mcp.example.com/.well-known/oauth-protected-resource")) + #expect(challenge?.errorDescription == "Additional file write permission required") + } + + @Test("Parse Bearer challenge when another auth scheme appears first") + func testParseBearerChallengeWhenBearerIsNotFirst() { + let headers = [ + "WWW-Authenticate": + "Basic realm=\"legacy\", Bearer resource_metadata=\"https://mcp.example.com/.well-known/oauth-protected-resource\", scope=\"files:read\"" + ] + + let challenge = DefaultOAuthWWWAuthenticateParser().parseBearer(from: headers) + + #expect(challenge != nil) + #expect( + challenge?.resourceMetadataURL + == URL(string: "https://mcp.example.com/.well-known/oauth-protected-resource")) + #expect(challenge?.scope == "files:read") + } + + @Test("Parse Bearer challenge stops before the next auth scheme") + func testParseBearerChallengeStopsBeforeNextScheme() { + let headers = [ + "WWW-Authenticate": + "Bearer resource_metadata=\"https://mcp.example.com/.well-known/oauth-protected-resource\", scope=\"files:read\", DPoP algs=\"ES256\"" + ] + + let challenge = DefaultOAuthWWWAuthenticateParser().parseBearer(from: headers) + + #expect(challenge != nil) + #expect( + challenge?.resourceMetadataURL + == URL(string: "https://mcp.example.com/.well-known/oauth-protected-resource")) + #expect(challenge?.scope == "files:read") + #expect(challenge?.parameters["algs"] == nil) + } + + @Test("Protected resource metadata discovery fallback URLs") + func testProtectedResourceMetadataURLs() { + let endpoint = URL(string: "https://example.com/public/mcp")! + let urls = metadataDiscovery.protectedResourceMetadataURLs(for: endpoint) + + #expect( + urls == [ + URL(string: "https://example.com/.well-known/oauth-protected-resource/public/mcp")!, + URL(string: "https://example.com/.well-known/oauth-protected-resource")!, + ]) + } + + @Test("Authorization server metadata discovery URLs for issuer with path") + func testAuthorizationServerMetadataURLsWithPath() { + let issuer = URL(string: "https://auth.example.com/tenant1")! + let urls = metadataDiscovery.authorizationServerMetadataURLs(for: issuer) + + #expect( + urls == [ + URL(string: "https://auth.example.com/.well-known/oauth-authorization-server/tenant1")!, + URL(string: "https://auth.example.com/.well-known/openid-configuration/tenant1")!, + URL(string: "https://auth.example.com/tenant1/.well-known/openid-configuration")!, + ]) + } + + @Test("Authorization server metadata discovery URLs for issuer without path") + func testAuthorizationServerMetadataURLsWithoutPath() { + let issuer = URL(string: "https://auth.example.com")! + let urls = metadataDiscovery.authorizationServerMetadataURLs(for: issuer) + + #expect( + urls == [ + URL(string: "https://auth.example.com/.well-known/oauth-authorization-server")!, + URL(string: "https://auth.example.com/.well-known/openid-configuration")!, + ]) + } + + @Test("Canonical resource URI normalization") + func testCanonicalResourceURINormalization() throws { + let endpoint = URL(string: "HTTPS://MCP.EXAMPLE.COM/?q=1")! + let canonical = try metadataDiscovery.canonicalResourceURI(from: endpoint) + #expect(canonical.absoluteString == "https://mcp.example.com") + } + + @Test("Canonical resource URI supports explicit port and root slash normalization") + func testCanonicalResourceURIWithExplicitPort() throws { + let endpoint = URL(string: "HTTPS://MCP.EXAMPLE.COM:8443/")! + let canonical = try metadataDiscovery.canonicalResourceURI(from: endpoint) + #expect(canonical.absoluteString == "https://mcp.example.com:8443") + } + + @Test("Canonical resource URI preserves specific server path") + func testCanonicalResourceURIPreservesPath() throws { + let endpoint = URL(string: "https://mcp.example.com/server/mcp")! + let canonical = try metadataDiscovery.canonicalResourceURI(from: endpoint) + #expect(canonical.absoluteString == "https://mcp.example.com/server/mcp") + } + + @Test("Canonical resource URI rejects missing scheme") + func testCanonicalResourceURIRejectsMissingScheme() { + #expect(throws: OAuthAuthorizationError.self) { + _ = try metadataDiscovery.canonicalResourceURI(from: URL(string: "mcp.example.com")!) + } + } + + @Test("Canonical resource URI rejects insecure non-loopback http scheme") + func testCanonicalResourceURIRejectsNonLoopbackHTTP() { + #expect(throws: OAuthAuthorizationError.self) { + _ = try metadataDiscovery.canonicalResourceURI( + from: URL(string: "http://mcp.example.com/resource")! + ) + } + } + + @Test("Canonical resource URI allows loopback http scheme") + func testCanonicalResourceURIAllowsLoopbackHTTP() throws { + let canonical = try metadataDiscovery.canonicalResourceURI( + from: URL(string: "http://localhost:8080/mcp")! + ) + #expect(canonical.absoluteString == "http://localhost:8080/mcp") + } + + @Test("Canonical resource URI rejects fragment") + func testCanonicalResourceURIRejectsFragment() { + #expect(throws: OAuthAuthorizationError.self) { + _ = try metadataDiscovery.canonicalResourceURI( + from: URL(string: "https://mcp.example.com#fragment")! + ) + } + } + + @Test("Protected resource matching allows same-origin parent resource") + func testProtectedResourceMatchingParentResource() { + let resource = URL(string: "https://mcp.example.com")! + let endpoint = URL(string: "https://mcp.example.com/mcp")! + #expect(metadataDiscovery.protectedResourceMatches(resource: resource, endpoint: endpoint)) + } + + @Test("Protected resource matching enforces path boundaries") + func testProtectedResourceMatchingPathBoundary() { + let resource = URL(string: "https://mcp.example.com/mcp")! + let validEndpoint = URL(string: "https://mcp.example.com/mcp/tools")! + let invalidEndpoint = URL(string: "https://mcp.example.com/mcp2")! + + #expect(metadataDiscovery.protectedResourceMatches(resource: resource, endpoint: validEndpoint)) + #expect(!metadataDiscovery.protectedResourceMatches(resource: resource, endpoint: invalidEndpoint)) + } + + @Test("Protected resource matching rejects origin mismatches") + func testProtectedResourceMatchingOriginMismatch() { + let resource = URL(string: "https://evil.example.com/mcp")! + let endpoint = URL(string: "https://mcp.example.com/mcp")! + #expect(!metadataDiscovery.protectedResourceMatches(resource: resource, endpoint: endpoint)) + } + + @Test("Scope selection prefers challenge scope") + func testScopeSelection() { + let selected = DefaultOAuthScopeSelector().selectScopes( + challengeScope: "files:read", + scopesSupported: ["files:read", "files:write"] + ) + #expect(selected == Set(["files:read"])) + + let fallback = DefaultOAuthScopeSelector().selectScopes( + challengeScope: nil, + scopesSupported: ["files:read", "files:write"] + ) + #expect(fallback == Set(["files:read", "files:write"])) + + let omitted = DefaultOAuthScopeSelector().selectScopes(challengeScope: nil, scopesSupported: nil) + #expect(omitted == nil) + } + + #if canImport(CryptoKit) + @Test("private_key_jwt helper builds signed JWT with expected claims") + func testPrivateKeyJWTAssertionHelper() throws { + let tokenEndpoint = URL(string: "https://auth.example.com/oauth/token")! + let assertion = try OAuthConfiguration.makePrivateKeyJWTAssertion( + clientID: "test-client", + tokenEndpoint: tokenEndpoint, + privateKeyPEM: Self.testPrivateKeyPEM, + audience: "https://auth.example.com", + issuedAt: Date(timeIntervalSince1970: 1_700_000_000), + expiresIn: 300 + ) + + let parts = assertion.split(separator: ".") + #expect(parts.count == 3) + + let headerData = Self.decodeBase64URL(parts[0]) + let payloadData = Self.decodeBase64URL(parts[1]) + #expect(headerData != nil) + #expect(payloadData != nil) + + let header = try JSONSerialization.jsonObject(with: headerData!) as? [String: Any] + let payload = try JSONSerialization.jsonObject(with: payloadData!) as? [String: Any] + + #expect(header?["alg"] as? String == "ES256") + #expect(header?["typ"] as? String == "JWT") + #expect(payload?["iss"] as? String == "test-client") + #expect(payload?["sub"] as? String == "test-client") + #expect(payload?["aud"] as? String == "https://auth.example.com") + #expect(payload?["iat"] as? Int == 1_700_000_000) + #expect(payload?["exp"] as? Int == 1_700_000_300) + #expect((payload?["jti"] as? String)?.isEmpty == false) + } + #endif + + @Test("private_key_jwt helper rejects non-positive lifetime") + func testPrivateKeyJWTAssertionRejectsNonPositiveLifetime() { + #expect(throws: OAuthConfiguration.PrivateKeyJWTAssertionError.self) { + _ = try OAuthConfiguration.makePrivateKeyJWTAssertion( + clientID: "test-client", + tokenEndpoint: URL(string: "https://auth.example.com/token")!, + privateKeyPEM: Self.testPrivateKeyPEM, + expiresIn: 0 + ) + } + } + + // MARK: - WWW-Authenticate Parser Edge Cases + + @Test("Parse Bearer returns nil when WWW-Authenticate header is absent") + func testParseBearerReturnsNilWhenHeaderAbsent() { + let challenge = DefaultOAuthWWWAuthenticateParser().parseBearer(from: ["Content-Type": "text/plain"]) + #expect(challenge == nil) + } + + @Test("Parse Bearer returns nil for empty header value") + func testParseBearerReturnsNilForEmptyValue() { + let challenge = DefaultOAuthWWWAuthenticateParser().parseBearer(from: ["WWW-Authenticate": ""]) + #expect(challenge == nil) + } + + @Test("Parse Bearer returns empty challenge for bare Bearer with no parameters") + func testParseBearerBareBearerScheme() { + let challenge = DefaultOAuthWWWAuthenticateParser().parseBearer(from: ["WWW-Authenticate": "Bearer"]) + #expect(challenge != nil) + #expect(challenge?.parameters.isEmpty == true) + } + + @Test("Parse Bearer header name lookup is case-insensitive") + func testParseBearerCaseInsensitiveHeaderName() { + let challenge = DefaultOAuthWWWAuthenticateParser().parseBearer( + from: ["www-authenticate": "Bearer scope=\"read\""] + ) + #expect(challenge != nil) + #expect(challenge?.scope == "read") + } + + @Test("Parse Bearer returns nil when only non-Bearer schemes present") + func testParseBearerReturnsNilForNonBearerScheme() { + let challenge = DefaultOAuthWWWAuthenticateParser().parseBearer( + from: ["WWW-Authenticate": "Basic realm=\"example\""] + ) + #expect(challenge == nil) + } + + // MARK: - Metadata Discovery Edge Cases + + @Test("Protected resource metadata returns empty for non-HTTPS endpoint") + func testProtectedResourceMetadataRejectsHTTP() { + let urls = metadataDiscovery.protectedResourceMetadataURLs( + for: URL(string: "http://remote.example.com/mcp")! + ) + #expect(urls.isEmpty) + } + + @Test("Protected resource metadata for root endpoint produces path-specific and root URLs") + func testProtectedResourceMetadataRootEndpoint() { + let urls = metadataDiscovery.protectedResourceMetadataURLs( + for: URL(string: "https://example.com")! + ) + #expect(urls.count == 2) + #expect(urls.allSatisfy { + $0 == URL(string: "https://example.com/.well-known/oauth-protected-resource")! + }) + } + + @Test("Protected resource metadata allows loopback HTTP") + func testProtectedResourceMetadataAllowsLoopbackHTTP() { + let urls = metadataDiscovery.protectedResourceMetadataURLs( + for: URL(string: "http://localhost:8080/mcp")! + ) + #expect(!urls.isEmpty) + } + + @Test("Authorization server metadata returns empty for non-HTTPS issuer") + func testAuthorizationServerMetadataRejectsHTTP() { + let urls = metadataDiscovery.authorizationServerMetadataURLs( + for: URL(string: "http://remote.example.com")! + ) + #expect(urls.isEmpty) + } + + @Test("Authorization server fallback issuer derives origin from endpoint") + func testAuthorizationServerFallbackIssuer() throws { + let issuer = try metadataDiscovery.authorizationServerFallbackIssuer( + from: URL(string: "https://mcp.example.com:8443/server/mcp")! + ) + #expect(issuer.absoluteString == "https://mcp.example.com:8443") + } + + @Test("Authorization server fallback issuer normalizes case and strips query") + func testAuthorizationServerFallbackIssuerNormalization() throws { + let issuer = try metadataDiscovery.authorizationServerFallbackIssuer( + from: URL(string: "HTTPS://AUTH.EXAMPLE.COM/path?q=1")! + ) + #expect(issuer.absoluteString == "https://auth.example.com") + } + + @Test("Authorization server fallback issuer rejects non-HTTPS non-loopback") + func testAuthorizationServerFallbackIssuerRejectsInsecure() { + #expect(throws: OAuthAuthorizationError.self) { + _ = try metadataDiscovery.authorizationServerFallbackIssuer( + from: URL(string: "http://remote.example.com/mcp")! + ) + } + } + + @Test("Protected resource matching with port mismatch") + func testProtectedResourceMatchingPortMismatch() { + let resource = URL(string: "https://mcp.example.com:8443")! + let endpoint = URL(string: "https://mcp.example.com:9443/mcp")! + #expect(!metadataDiscovery.protectedResourceMatches(resource: resource, endpoint: endpoint)) + } + + @Test("Protected resource matching with exact path") + func testProtectedResourceMatchingExactPath() { + let resource = URL(string: "https://mcp.example.com/mcp")! + let endpoint = URL(string: "https://mcp.example.com/mcp")! + #expect(metadataDiscovery.protectedResourceMatches(resource: resource, endpoint: endpoint)) + } + + // MARK: - Scope Selector Edge Cases + + @Test("parseScopeString splits on whitespace and ignores empty tokens") + func testParseScopeString() { + let scopes = DefaultOAuthScopeSelector().parseScopeString("files:read files:write\tprofile") + #expect(scopes == Set(["files:read", "files:write", "profile"])) + } + + @Test("parseScopeString returns empty set for blank string") + func testParseScopeStringEmpty() { + let scopes = DefaultOAuthScopeSelector().parseScopeString(" ") + #expect(scopes.isEmpty) + } + + @Test("serialize produces sorted space-separated string") + func testSerializeScopes() { + let result = DefaultOAuthScopeSelector().serialize(Set(["write", "read", "admin"])) + #expect(result == "admin read write") + } + + @Test("serialize returns nil for empty set") + func testSerializeScopesEmpty() { + let result = DefaultOAuthScopeSelector().serialize(Set()) + #expect(result == nil) + } + + @Test("selectScopes returns nil for empty challenge scope string") + func testSelectScopesEmptyChallengeScope() { + let result = DefaultOAuthScopeSelector().selectScopes( + challengeScope: " ", + scopesSupported: ["files:read"] + ) + #expect(result == nil) + } + + @Test("selectScopes returns nil for empty scopesSupported array") + func testSelectScopesEmptyScopesSupported() { + let result = DefaultOAuthScopeSelector().selectScopes( + challengeScope: nil, + scopesSupported: [] + ) + #expect(result == nil) + } + + // MARK: - Token Type Validation + + @Test("makePrivateKeyJWTAssertion — token type empty string is rejected") + func testEmptyTokenTypeIsRejected() throws { + // OAuthTokenResponse decodes token_type from JSON; simulate an empty value via + // direct struct construction using the internal initializer path. + let json = #"{"access_token":"tok","token_type":"","expires_in":3600}"# + let data = json.data(using: .utf8)! + let decoded = try JSONDecoder().decode(OAuthTokenResponse.self, from: data) + // The token_type is empty; the authorizer guard must reject this. + let tokenType = decoded.tokenType.trimmingCharacters(in: .whitespacesAndNewlines) + let isValid = !tokenType.isEmpty + && tokenType.caseInsensitiveCompare("Bearer") == .orderedSame + #expect(!isValid, "Empty token_type must not be accepted as Bearer") + } + + @Test("Token type whitespace-only is rejected") + func testWhitespaceTokenTypeIsRejected() throws { + let json = #"{"access_token":"tok","token_type":" ","expires_in":3600}"# + let data = json.data(using: .utf8)! + let decoded = try JSONDecoder().decode(OAuthTokenResponse.self, from: data) + let tokenType = decoded.tokenType.trimmingCharacters(in: .whitespacesAndNewlines) + let isValid = !tokenType.isEmpty + && tokenType.caseInsensitiveCompare("Bearer") == .orderedSame + #expect(!isValid, "Whitespace-only token_type must not be accepted as Bearer") + } + + // MARK: - InMemoryTokenStorage + + @Test("InMemoryTokenStorage save and load round-trip") + func testTokenStorageSaveAndLoad() { + let storage = InMemoryTokenStorage() + #expect(storage.load() == nil) + + let token = OAuthAccessToken( + value: "access-123", + tokenType: "Bearer", + expiresAt: nil, + scopes: ["read"], + authorizationServer: nil, + refreshToken: nil + ) + storage.save(token) + + let loaded = storage.load() + #expect(loaded?.value == "access-123") + #expect(loaded?.tokenType == "Bearer") + #expect(loaded?.scopes == Set(["read"])) + } + + @Test("InMemoryTokenStorage clear removes stored token") + func testTokenStorageClear() { + let storage = InMemoryTokenStorage() + let token = OAuthAccessToken( + value: "access-456", + tokenType: "Bearer", + expiresAt: nil, + scopes: [], + authorizationServer: nil, + refreshToken: nil + ) + storage.save(token) + #expect(storage.load() != nil) + + storage.clear() + #expect(storage.load() == nil) + } + + @Test("InMemoryTokenStorage save overwrites previous token") + func testTokenStorageSaveOverwrites() { + let storage = InMemoryTokenStorage() + let first = OAuthAccessToken( + value: "first", + tokenType: "Bearer", + expiresAt: nil, + scopes: [], + authorizationServer: nil, + refreshToken: nil + ) + let second = OAuthAccessToken( + value: "second", + tokenType: "Bearer", + expiresAt: nil, + scopes: [], + authorizationServer: nil, + refreshToken: nil + ) + storage.save(first) + storage.save(second) + + #expect(storage.load()?.value == "second") + } + + // MARK: - Proactive Token Refresh + + @Test("prepareAuthorization does nothing when proactiveRefreshWindowSeconds is zero") + func testPrepareAuthorizationSkipsWhenWindowIsZero() async throws { + let storage = InMemoryTokenStorage() + // Token expires in 50 s — within a large proactive window but not within the 30 s default skew. + storage.save(OAuthAccessToken( + value: "original-token", + tokenType: "Bearer", + expiresAt: Date().addingTimeInterval(50), + scopes: [], + authorizationServer: nil, + refreshToken: "some-refresh-token" + )) + let config = OAuthConfiguration( + authentication: .none(clientID: "client"), + proactiveRefreshWindowSeconds: 0 + ) + let authorizer = OAuthAuthorizer(configuration: config, tokenStorage: storage) + + try await authorizer.prepareAuthorization( + for: URL(string: "https://example.com/mcp")!, + session: .shared + ) + + // Token must be unchanged — proactive refresh is disabled. + #expect(storage.load()?.value == "original-token") + } + + @Test("prepareAuthorization does nothing without cached authorization server metadata") + func testPrepareAuthorizationSkipsWithoutCachedASMetadata() async throws { + let storage = InMemoryTokenStorage() + // Token expires in 50 s — within the 400 s proactive window. + storage.save(OAuthAccessToken( + value: "original-token", + tokenType: "Bearer", + expiresAt: Date().addingTimeInterval(50), + scopes: [], + authorizationServer: nil, + refreshToken: "some-refresh-token" + )) + let config = OAuthConfiguration( + authentication: .none(clientID: "client"), + proactiveRefreshWindowSeconds: 400 + ) + // No handleChallenge call → authorizationServerMetadata remains nil. + let authorizer = OAuthAuthorizer(configuration: config, tokenStorage: storage) + + try await authorizer.prepareAuthorization( + for: URL(string: "https://example.com/mcp")!, + session: .shared + ) + + // Token must be unchanged — refresh requires cached AS metadata. + #expect(storage.load()?.value == "original-token") + } + + // MARK: - WWW-Authenticate Bare Scheme Detection + + @Test("Bare scheme Digest terminates preceding Bearer parameter collection") + func testBareSchemeTerminatesBearer() { + // "Digest" has no params and no '=', so it must start a new challenge and stop + // the Bearer parameter collector before "scope" from the second Bearer leaks in. + let header = #"Bearer scope="a", Digest, Bearer scope="b""# + let challenge = DefaultOAuthWWWAuthenticateParser().parseBearer( + from: ["WWW-Authenticate": header]) + #expect(challenge != nil) + // The outer loop stops at the first Bearer, so scope must be "a", not "b". + #expect(challenge?.scope == "a") + } + + @Test("Non-Bearer leading scheme followed by Bearer is parsed correctly") + func testNonBearerLeadingSchemeFollowedByBearer() { + let header = #"SomeScheme, Bearer scope="test""# + let challenge = DefaultOAuthWWWAuthenticateParser().parseBearer( + from: ["WWW-Authenticate": header]) + #expect(challenge != nil) + #expect(challenge?.scope == "test") + } + + // MARK: - Private IP Blocking + + @Test("privateIPAddressBlocked error has informative description") + func testPrivateIPAddressBlockedErrorDescription() { + let url = "https://169.254.169.254/.well-known/oauth-protected-resource" + let error = OAuthAuthorizationError.privateIPAddressBlocked( + context: "Protected resource metadata URL", url: url) + let description = error.errorDescription ?? "" + #expect(description.contains("private or reserved IP")) + #expect(description.contains(url)) + } + + @Test("cimdNotSupported is thrown when CIMD URL provided but server does not support it") + func testCIMDNotSupportedError() { + // Verify the error case exists and its description is meaningful. + let error = OAuthAuthorizationError.cimdNotSupported( + clientID: "https://client.example.com/client-metadata.json") + let description = error.errorDescription ?? "" + #expect(description.contains("Client ID Metadata Document")) + #expect(description.contains("client.example.com")) + } + + // MARK: - BearerTokenValidator Audience and Expiry Tests + + private let testResourceMetadataURL = + URL(string: "https://api.example.com/.well-known/oauth-protected-resource")! + private let testResourceIdentifier = URL(string: "https://api.example.com")! + + private func makeBearerValidator( + tokenValidator: @escaping BearerTokenValidator.TokenValidator + ) -> BearerTokenValidator { + BearerTokenValidator( + resourceMetadataURL: testResourceMetadataURL, + resourceIdentifier: testResourceIdentifier, + tokenValidator: tokenValidator + ) + } + + private func makeRequest(authorization: String) -> HTTPRequest { + HTTPRequest( + method: "POST", + headers: [HTTPHeaderName.authorization: authorization], + path: "/mcp" + ) + } + + private func makeContext() -> HTTPValidationContext { + HTTPValidationContext(httpMethod: "POST", isInitializationRequest: false) + } + + @Test("BearerTokenValidator allows valid token with nil audience (opaque token)") + func testBearerValidatorNilAudienceAllows() { + let validator = makeBearerValidator { _, _, _ in + .valid(BearerTokenInfo(audience: nil)) + } + let result = validator.validate(makeRequest(authorization: "Bearer tok"), context: makeContext()) + #expect(result == nil) + } + + @Test("BearerTokenValidator allows token with matching audience") + func testBearerValidatorMatchingAudienceAllows() { + let validator = makeBearerValidator { _, _, _ in + .valid(BearerTokenInfo(audience: ["https://api.example.com"])) + } + let result = validator.validate(makeRequest(authorization: "Bearer tok"), context: makeContext()) + #expect(result == nil) + } + + @Test("BearerTokenValidator allows token when one aud entry matches") + func testBearerValidatorOneOfMultipleAudienceMatches() { + let validator = makeBearerValidator { _, _, _ in + .valid(BearerTokenInfo(audience: ["https://other.example.com", "https://api.example.com"])) + } + let result = validator.validate(makeRequest(authorization: "Bearer tok"), context: makeContext()) + #expect(result == nil) + } + + @Test("BearerTokenValidator returns 401 invalid_token when audience does not match") + func testBearerValidatorAudienceMismatchReturns401() { + let validator = makeBearerValidator { _, _, _ in + .valid(BearerTokenInfo(audience: ["https://other.example.com"])) + } + let result = validator.validate(makeRequest(authorization: "Bearer tok"), context: makeContext()) + #expect(result?.statusCode == 401) + let challenge = result?.headers[HTTPHeaderName.wwwAuthenticate] + #expect(challenge?.contains("error=\"invalid_token\"") == true) + #expect(challenge?.contains("Token audience mismatch") == true) + } + + @Test("BearerTokenValidator returns 401 invalid_token when token is expired") + func testBearerValidatorExpiredTokenReturns401() { + let pastDate = Date(timeIntervalSinceNow: -3600) + let validator = makeBearerValidator { _, _, _ in + .valid(BearerTokenInfo(expiresAt: pastDate)) + } + let result = validator.validate(makeRequest(authorization: "Bearer tok"), context: makeContext()) + #expect(result?.statusCode == 401) + let challenge = result?.headers[HTTPHeaderName.wwwAuthenticate] + #expect(challenge?.contains("error=\"invalid_token\"") == true) + #expect(challenge?.contains("Token has expired") == true) + } + + @Test("BearerTokenValidator allows token that has not yet expired") + func testBearerValidatorNonExpiredTokenAllows() { + let futureDate = Date(timeIntervalSinceNow: 3600) + let validator = makeBearerValidator { _, _, _ in + .valid(BearerTokenInfo(expiresAt: futureDate)) + } + let result = validator.validate(makeRequest(authorization: "Bearer tok"), context: makeContext()) + #expect(result == nil) + } + + @Test("BearerTokenValidator checks expiry before audience") + func testBearerValidatorExpiryCheckedBeforeAudience() { + let pastDate = Date(timeIntervalSinceNow: -1) + let validator = makeBearerValidator { _, _, _ in + // Expired token but matching audience — expiry should win + .valid(BearerTokenInfo(audience: ["https://api.example.com"], expiresAt: pastDate)) + } + let result = validator.validate(makeRequest(authorization: "Bearer tok"), context: makeContext()) + #expect(result?.statusCode == 401) + let challenge = result?.headers[HTTPHeaderName.wwwAuthenticate] + #expect(challenge?.contains("Token has expired") == true) + } +} diff --git a/Tests/MCPTests/OAuthAuthorizerTests.swift b/Tests/MCPTests/OAuthAuthorizerTests.swift new file mode 100644 index 0000000..0384307 --- /dev/null +++ b/Tests/MCPTests/OAuthAuthorizerTests.swift @@ -0,0 +1,429 @@ +@preconcurrency import Foundation +import Testing + +@testable import MCP + +#if canImport(FoundationNetworking) + import FoundationNetworking +#endif + +// MARK: - Mock Implementations + +final class MockURLValidator: OAuthURLValidating, @unchecked Sendable { + var validateHTTPSOrLoopbackCallCount = 0 + var validateAuthorizationServerCallCount = 0 + var validateRedirectURICallCount = 0 + var shouldThrow: Error? + + func validateHTTPSOrLoopback(_ url: URL, context: String) throws { + validateHTTPSOrLoopbackCallCount += 1 + if let error = shouldThrow { throw error } + } + + func validateAuthorizationServer(_ url: URL, context: String) throws { + validateAuthorizationServerCallCount += 1 + if let error = shouldThrow { throw error } + } + + func validateRedirectURI(_ url: URL) throws { + validateRedirectURICallCount += 1 + if let error = shouldThrow { throw error } + } + + func isPrivateIPHost(_ host: String) -> Bool { false } +} + +final class MockDiscoveryClient: OAuthDiscoveryFetching, @unchecked Sendable { + var fetchProtectedResourceMetadataCallCount = 0 + var fetchAuthorizationServerMetadataCallCount = 0 + let metadataDiscovery: any OAuthMetadataDiscovering = DefaultOAuthMetadataDiscovery() + + var protectedResourceMetadataResult: OAuthProtectedResourceMetadata + var authorizationServerMetadataResult: (server: URL, metadata: OAuthAuthorizationServerMetadata) + + init( + authorizationServer: URL = URL(string: "https://auth.example.com")!, + tokenEndpoint: URL = URL(string: "https://auth.example.com/token")! + ) { + self.protectedResourceMetadataResult = OAuthProtectedResourceMetadata( + resource: nil, + authorizationServers: [authorizationServer], + scopesSupported: nil + ) + self.authorizationServerMetadataResult = ( + server: authorizationServer, + metadata: OAuthAuthorizationServerMetadata( + issuer: authorizationServer, + authorizationEndpoint: URL(string: "https://auth.example.com/authorize"), + tokenEndpoint: tokenEndpoint, + registrationEndpoint: nil, + codeChallengeMethodsSupported: ["S256"], + tokenEndpointAuthMethodsSupported: nil, + clientIDMetadataDocumentSupported: nil + ) + ) + } + + func fetchProtectedResourceMetadata(candidates: [URL], fallbackIssuer: URL?, session: URLSession) async throws -> OAuthProtectedResourceMetadata { + fetchProtectedResourceMetadataCallCount += 1 + return protectedResourceMetadataResult + } + + func fetchAuthorizationServerMetadata(candidates: [URL], session: URLSession) async throws -> (server: URL, metadata: OAuthAuthorizationServerMetadata) { + fetchAuthorizationServerMetadataCallCount += 1 + return authorizationServerMetadataResult + } +} + +final class MockTokenClient: OAuthTokenRequesting, @unchecked Sendable { + var requestCallCount = 0 + var capturedParameters: [String: String]? + var tokenResponse = OAuthTokenResponse( + accessToken: "mock-access-token", + tokenType: "Bearer", + expiresIn: 3600, + scope: nil, + refreshToken: nil + ) + + func request( + parameters: inout [String: String], + endpoint: URL, + authentication: OAuthConfiguration.TokenEndpointAuthentication, + session: URLSession + ) async throws -> OAuthTokenResponse { + requestCallCount += 1 + capturedParameters = parameters + return tokenResponse + } +} + +final class MockClientRegistrar: OAuthClientRegistering, @unchecked Sendable { + var registerCallCount = 0 + var registrationResult: ( + response: OAuthClientRegistrationResponse, + updatedAuthentication: OAuthConfiguration.TokenEndpointAuthentication + )? + + func register( + configuration: OAuthConfiguration, + asMetadata: OAuthAuthorizationServerMetadata, + session: URLSession + ) async throws -> ( + response: OAuthClientRegistrationResponse, + updatedAuthentication: OAuthConfiguration.TokenEndpointAuthentication + )? { + registerCallCount += 1 + return registrationResult + } +} + +final class MockAuthCodeFlow: OAuthAuthorizationCodeFlowing, @unchecked Sendable { + var buildURLCallCount = 0 + var performCallCount = 0 + var authorizationCode = "mock-auth-code" + + func buildURL( + authorizationEndpoint: URL, + resource: URL, + redirectURI: URL, + clientID: String, + codeChallenge: String, + scopes: Set?, + state: String, + scopeSerializer: any OAuthScopeSelecting + ) throws -> URL { + buildURLCallCount += 1 + return URL(string: "https://auth.example.com/authorize?code=stub")! + } + + func perform( + authorizationURL: URL, + redirectURI: URL, + state: String, + delegate: (any OAuthAuthorizationDelegate)?, + session: URLSession + ) async throws -> String { + performCallCount += 1 + return authorizationCode + } +} + +// MARK: - OAuthAuthorizer Invocation Tests + +@Suite("OAuthAuthorizer dependency invocations") +struct OAuthAuthorizerTests { + + let endpoint = URL(string: "https://mcp.example.com/mcp")! + let headers401 = [ + "WWW-Authenticate": + "Bearer resource_metadata=\"https://mcp.example.com/.well-known/oauth-protected-resource\"" + ] + + func makeAuthorizer( + grantType: OAuthConfiguration.GrantType = .clientCredentials, + urlValidator: MockURLValidator = MockURLValidator(), + discoveryClient: MockDiscoveryClient = MockDiscoveryClient(), + tokenClient: MockTokenClient = MockTokenClient(), + registrar: MockClientRegistrar = MockClientRegistrar(), + authCodeFlow: MockAuthCodeFlow = MockAuthCodeFlow() + ) -> OAuthAuthorizer { + let config = OAuthConfiguration( + grantType: grantType, + authentication: .clientSecretBasic(clientID: "client", clientSecret: "secret") + ) + return OAuthAuthorizer( + configuration: config, + urlValidator: urlValidator, + discoveryClient: discoveryClient, + tokenEndpointClient: tokenClient, + clientRegistrar: registrar, + authCodeFlow: authCodeFlow + ) + } + + // MARK: - validateEndpointSecurity + + @Test("validateEndpointSecurity calls urlValidator") + func testValidateEndpointSecurityCallsURLValidator() throws { + let validator = MockURLValidator() + let authorizer = makeAuthorizer(urlValidator: validator) + + try authorizer.validateEndpointSecurity(for: endpoint) + + #expect(validator.validateHTTPSOrLoopbackCallCount == 1) + } + + @Test("validateEndpointSecurity propagates validation error") + func testValidateEndpointSecurityPropagatesError() { + let validator = MockURLValidator() + validator.shouldThrow = OAuthAuthorizationError.insecureOAuthEndpoint( + context: "test", url: "http://example.com") + let authorizer = makeAuthorizer(urlValidator: validator) + + #expect(throws: OAuthAuthorizationError.self) { + try authorizer.validateEndpointSecurity(for: endpoint) + } + } + + // MARK: - handleChallenge (401 — client_credentials) + + @Test("handleChallenge 401 calls discovery and token clients") + func testHandleChallenge401CallsDiscoveryAndTokenClient() async throws { + let discovery = MockDiscoveryClient() + let tokenClient = MockTokenClient() + + let authorizer = makeAuthorizer( + discoveryClient: discovery, + tokenClient: tokenClient + ) + + let handled = try await authorizer.handleChallenge( + statusCode: 401, + headers: headers401, + endpoint: endpoint, + operationKey: nil, + session: .shared + ) + + #expect(handled == true) + #expect(discovery.fetchProtectedResourceMetadataCallCount >= 1) + #expect(discovery.fetchAuthorizationServerMetadataCallCount >= 1) + #expect(tokenClient.requestCallCount == 1) + } + + @Test("handleChallenge 401 uses client_credentials grant type parameter") + func testHandleChallenge401ClientCredentialsGrantType() async throws { + let tokenClient = MockTokenClient() + let authorizer = makeAuthorizer(tokenClient: tokenClient) + + _ = try await authorizer.handleChallenge( + statusCode: 401, + headers: headers401, + endpoint: endpoint, + operationKey: nil, + session: .shared + ) + + #expect(tokenClient.capturedParameters?["grant_type"] == "client_credentials") + } + + @Test("handleChallenge 401 attaches resource parameter") + func testHandleChallenge401AttachesResourceParameter() async throws { + let tokenClient = MockTokenClient() + let authorizer = makeAuthorizer(tokenClient: tokenClient) + + _ = try await authorizer.handleChallenge( + statusCode: 401, + headers: headers401, + endpoint: endpoint, + operationKey: nil, + session: .shared + ) + + #expect(tokenClient.capturedParameters?["resource"] != nil) + } + + // MARK: - handleChallenge (authorization_code) + + #if canImport(CryptoKit) + @Test("handleChallenge 401 calls authCodeFlow for authorization_code grant") + func testHandleChallenge401AuthorizationCodeCallsFlow() async throws { + let authCodeFlow = MockAuthCodeFlow() + let tokenClient = MockTokenClient() + + let config = OAuthConfiguration( + grantType: .authorizationCode, + authentication: .none(clientID: "my-client"), + authorizationRedirectURI: URL(string: "https://app.example.com/callback")! + ) + let authorizer = OAuthAuthorizer( + configuration: config, + urlValidator: MockURLValidator(), + discoveryClient: MockDiscoveryClient(), + tokenEndpointClient: tokenClient, + clientRegistrar: MockClientRegistrar(), + authCodeFlow: authCodeFlow + ) + + _ = try await authorizer.handleChallenge( + statusCode: 401, + headers: headers401, + endpoint: endpoint, + operationKey: nil, + session: .shared + ) + + #expect(authCodeFlow.buildURLCallCount == 1) + #expect(authCodeFlow.performCallCount == 1) + #expect(tokenClient.capturedParameters?["grant_type"] == "authorization_code") + #expect(tokenClient.capturedParameters?["code"] == "mock-auth-code") + } + #endif + + // MARK: - handleChallenge (403) + + @Test("handleChallenge 403 returns false for non-insufficient_scope error") + func testHandleChallenge403NonInsufficientScope() async throws { + let authorizer = makeAuthorizer() + + let handled = try await authorizer.handleChallenge( + statusCode: 403, + headers: ["WWW-Authenticate": "Bearer error=\"access_denied\""], + endpoint: endpoint, + operationKey: nil, + session: .shared + ) + + #expect(handled == false) + } + + @Test("handleChallenge 403 insufficient_scope acquires token with upgraded scopes") + func testHandleChallenge403InsufficientScope() async throws { + let tokenClient = MockTokenClient() + let discovery = MockDiscoveryClient() + + let authorizer = makeAuthorizer( + discoveryClient: discovery, + tokenClient: tokenClient + ) + + let handled = try await authorizer.handleChallenge( + statusCode: 403, + headers: [ + "WWW-Authenticate": + "Bearer error=\"insufficient_scope\", scope=\"admin\"" + ], + endpoint: endpoint, + operationKey: nil, + session: .shared + ) + + #expect(handled == true) + #expect(tokenClient.requestCallCount == 1) + } + + // MARK: - Client registration + + @Test("handleChallenge calls client registrar when authentication is .none") + func testHandleChallengeCallsRegistrar() async throws { + let registrar = MockClientRegistrar() + let config = OAuthConfiguration( + authentication: .none(clientID: "plain-client")) + let authorizer = OAuthAuthorizer( + configuration: config, + urlValidator: MockURLValidator(), + discoveryClient: MockDiscoveryClient(), + tokenEndpointClient: MockTokenClient(), + clientRegistrar: registrar, + authCodeFlow: MockAuthCodeFlow() + ) + + _ = try await authorizer.handleChallenge( + statusCode: 401, + headers: headers401, + endpoint: endpoint, + operationKey: nil, + session: .shared + ) + + #expect(registrar.registerCallCount == 1) + } + + @Test("handleChallenge skips client registrar when credentials are already configured") + func testHandleChallengeSkipsRegistrarWithCredentials() async throws { + let registrar = MockClientRegistrar() + let authorizer = makeAuthorizer(registrar: registrar) + + _ = try await authorizer.handleChallenge( + statusCode: 401, + headers: headers401, + endpoint: endpoint, + operationKey: nil, + session: .shared + ) + + #expect(registrar.registerCallCount == 0) + } + + @Test("handleChallenge persists the DCR-assigned clientID on the saved token") + func testHandleChallengePersistsDCRClientIDOnToken() async throws { + let assignedClientID = "dcr-assigned-client-id" + let tokenStorage = InMemoryTokenStorage() + let registrar = MockClientRegistrar() + registrar.registrationResult = ( + response: OAuthClientRegistrationResponse( + clientID: assignedClientID, + clientSecret: nil, + tokenEndpointAuthMethod: nil, + clientSecretExpiresAt: nil + ), + updatedAuthentication: .none(clientID: assignedClientID) + ) + + let config = OAuthConfiguration( + authentication: .none(clientID: "") + ) + let authorizer = OAuthAuthorizer( + configuration: config, + tokenStorage: tokenStorage, + urlValidator: MockURLValidator(), + discoveryClient: MockDiscoveryClient(), + tokenEndpointClient: MockTokenClient(), + clientRegistrar: registrar, + authCodeFlow: MockAuthCodeFlow() + ) + + let handled = try await authorizer.handleChallenge( + statusCode: 401, + headers: headers401, + endpoint: endpoint, + operationKey: nil, + session: .shared + ) + + #expect(handled == true) + #expect(registrar.registerCallCount == 1) + #expect(tokenStorage.load()?.clientID == assignedClientID) + } +} diff --git a/Tests/MCPTests/OAuthClientRegistrarTests.swift b/Tests/MCPTests/OAuthClientRegistrarTests.swift new file mode 100644 index 0000000..7bdec5d --- /dev/null +++ b/Tests/MCPTests/OAuthClientRegistrarTests.swift @@ -0,0 +1,194 @@ +@preconcurrency import Foundation +import Testing + +@testable import MCP + +#if canImport(FoundationNetworking) + import FoundationNetworking +#endif + +#if swift(>=6.1) && !os(Linux) + + @Suite("OAuthClientRegistrar", .serialized) + struct OAuthClientRegistrarTests { + + let registrar = OAuthClientRegistrar(urlValidator: OAuthURLValidator()) + let registrationEndpoint = URL(string: "https://auth.example.com/register")! + + func makeASMetadata(registrationEndpoint: URL? = nil) -> OAuthAuthorizationServerMetadata { + OAuthAuthorizationServerMetadata( + issuer: URL(string: "https://auth.example.com"), + authorizationEndpoint: URL(string: "https://auth.example.com/authorize"), + tokenEndpoint: URL(string: "https://auth.example.com/token"), + registrationEndpoint: registrationEndpoint, + codeChallengeMethodsSupported: ["S256"], + tokenEndpointAuthMethodsSupported: nil, + clientIDMetadataDocumentSupported: nil + ) + } + + func makeConfig( + authentication: OAuthConfiguration.TokenEndpointAuthentication = .none(clientID: "") + ) -> OAuthConfiguration { + OAuthConfiguration(authentication: authentication) + } + + func successRegistrationBody(clientID: String = "registered-client") throws -> Data { + let dict: [String: Any] = ["client_id": clientID] + return try JSONSerialization.data(withJSONObject: dict) + } + + // MARK: - Skip Conditions + + @Test("Returns nil when authentication is not .none") + func testRegisterReturnsNilForNonNoneAuth() async throws { + let config = makeConfig( + authentication: .clientSecretBasic(clientID: "id", clientSecret: "secret")) + let (session, _) = makeIsolatedSession() + let result = try await registrar.register( + configuration: config, + asMetadata: makeASMetadata(), + session: session + ) + #expect(result == nil) + } + + @Test("Returns nil when no registration endpoint and no CIMD") + func testRegisterReturnsNilWithoutRegistrationEndpoint() async throws { + let config = makeConfig(authentication: .none(clientID: "plain-client-id")) + let (session, _) = makeIsolatedSession() + let result = try await registrar.register( + configuration: config, + asMetadata: makeASMetadata(registrationEndpoint: nil), + session: session + ) + #expect(result == nil) + } + + // MARK: - CIMD Errors + + @Test("Throws when clientID is HTTPS URL with path but server does not support CIMD") + func testRegisterThrowsCIMDNotSupported() async throws { + let config = makeConfig( + authentication: .none(clientID: "https://client.example.com/metadata.json")) + let asMetadata = OAuthAuthorizationServerMetadata( + issuer: nil, authorizationEndpoint: nil, tokenEndpoint: nil, + registrationEndpoint: nil, codeChallengeMethodsSupported: nil, + tokenEndpointAuthMethodsSupported: nil, clientIDMetadataDocumentSupported: false + ) + let (session, _) = makeIsolatedSession() + + await #expect(throws: OAuthAuthorizationError.self) { + try await registrar.register( + configuration: config, + asMetadata: asMetadata, + session: session + ) + } + } + + // MARK: - Successful Registration + + @Test("Returns registration response and updated authentication on success") + func testRegisterSucceeds() async throws { + let body = try successRegistrationBody(clientID: "registered-client") + let (session, key) = makeIsolatedSession() + await IsolatedMockURLProtocol.setHandler(key: key) { _ in + let response = HTTPURLResponse( + url: self.registrationEndpoint, statusCode: 201, + httpVersion: nil, headerFields: nil)! + return (response, body) + } + + let config = makeConfig(authentication: .none(clientID: "")) + let result = try await registrar.register( + configuration: config, + asMetadata: makeASMetadata(registrationEndpoint: registrationEndpoint), + session: session + ) + + let resultValue = try #require(result) + let expected = OAuthConfiguration.TokenEndpointAuthentication.none(clientID: "registered-client") + #expect(resultValue.updatedAuthentication == expected) + } + + @Test("Throws on 4xx registration response") + func testRegisterThrowsOn4xx() async throws { + let errorBody = try JSONSerialization.data( + withJSONObject: ["error": "invalid_client_metadata"]) + let (session, key) = makeIsolatedSession() + await IsolatedMockURLProtocol.setHandler(key: key) { _ in + let response = HTTPURLResponse( + url: self.registrationEndpoint, statusCode: 400, + httpVersion: nil, headerFields: nil)! + return (response, errorBody) + } + + let config = makeConfig(authentication: .none(clientID: "")) + await #expect(throws: OAuthAuthorizationError.self) { + try await registrar.register( + configuration: config, + asMetadata: makeASMetadata(registrationEndpoint: registrationEndpoint), + session: session + ) + } + } + + @Test("Throws tokenRequestFailed on non-2xx non-4xx response") + func testRegisterThrowsOn5xx() async throws { + let (session, key) = makeIsolatedSession() + await IsolatedMockURLProtocol.setHandler(key: key) { _ in + let response = HTTPURLResponse( + url: self.registrationEndpoint, statusCode: 503, + httpVersion: nil, headerFields: nil)! + return (response, Data()) + } + + let config = makeConfig(authentication: .none(clientID: "")) + let error = await #expect(throws: OAuthAuthorizationError.self) { + try await registrar.register( + configuration: config, + asMetadata: makeASMetadata(registrationEndpoint: registrationEndpoint), + session: session + ) + } + guard case .tokenRequestFailed(let statusCode, let oauthError) = error else { + Issue.record("Expected tokenRequestFailed, got \(String(describing: error))") + return + } + #expect(statusCode == 503) + #expect(oauthError == nil) + } + + // MARK: - updatedAuthentication helper + + @Test("updatedAuthentication updates client ID and secret for basic auth") + func testUpdatedAuthenticationBasic() { + let registration = OAuthClientRegistrationResponse( + clientID: "new-id", clientSecret: "new-secret", + tokenEndpointAuthMethod: nil, clientSecretExpiresAt: nil + ) + let result = OAuthClientRegistrar.updatedAuthentication( + from: registration, + current: .clientSecretBasic(clientID: "old-id", clientSecret: "old-secret") + ) + let expected = OAuthConfiguration.TokenEndpointAuthentication.clientSecretBasic(clientID: "new-id", clientSecret: "new-secret") + #expect(result == expected) + } + + @Test("updatedAuthentication falls back to existing secret when not returned") + func testUpdatedAuthenticationFallsBackToCurrentSecret() { + let registration = OAuthClientRegistrationResponse( + clientID: "new-id", clientSecret: nil, + tokenEndpointAuthMethod: nil, clientSecretExpiresAt: nil + ) + let result = OAuthClientRegistrar.updatedAuthentication( + from: registration, + current: .clientSecretBasic(clientID: "old-id", clientSecret: "kept-secret") + ) + let expected = OAuthConfiguration.TokenEndpointAuthentication.clientSecretBasic(clientID: "new-id", clientSecret: "kept-secret") + #expect(result == expected) + } + } + +#endif diff --git a/Tests/MCPTests/OAuthComponentTestHelpers.swift b/Tests/MCPTests/OAuthComponentTestHelpers.swift new file mode 100644 index 0000000..a89cd0f --- /dev/null +++ b/Tests/MCPTests/OAuthComponentTestHelpers.swift @@ -0,0 +1,90 @@ +@preconcurrency import Foundation + +#if canImport(FoundationNetworking) + import FoundationNetworking +#endif + +/// Key-isolated mock URLProtocol. +/// +/// Each test creates a session with a unique key and registers its handler under that key. +/// Multiple suites run concurrently without interference because each test's requests carry +/// a different key, so they never read the wrong handler. +#if swift(>=6.1) + + actor IsolatedMockStorage { + static let shared = IsolatedMockStorage() + private var handlers: + [String: @Sendable (URLRequest) async throws -> (HTTPURLResponse, Data)] = [:] + + func set( + key: String, + handler: @escaping @Sendable (URLRequest) async throws -> (HTTPURLResponse, Data) + ) { + handlers[key] = handler + } + + func execute(key: String, request: URLRequest) async throws -> (HTTPURLResponse, Data) { + guard let handler = handlers[key] else { + throw NSError( + domain: "IsolatedMockError", code: 0, + userInfo: [NSLocalizedDescriptionKey: "No handler for key '\(key)'"]) + } + return try await handler(request) + } + } + + final class IsolatedMockURLProtocol: URLProtocol, @unchecked Sendable { + static let sessionKeyHeader = "X-Mock-Key" + + static func makeSession(key: String) -> URLSession { + let config = URLSessionConfiguration.ephemeral + config.protocolClasses = [IsolatedMockURLProtocol.self] + config.httpAdditionalHeaders = [sessionKeyHeader: key] + return URLSession(configuration: config) + } + + static func setHandler( + key: String, + _ handler: @escaping @Sendable (URLRequest) async throws -> (HTTPURLResponse, Data) + ) async { + await IsolatedMockStorage.shared.set(key: key, handler: handler) + } + + func executeHandler(for request: URLRequest) async throws -> (HTTPURLResponse, Data) { + let key = request.value(forHTTPHeaderField: IsolatedMockURLProtocol.sessionKeyHeader) + ?? "unknown" + return try await IsolatedMockStorage.shared.execute(key: key, request: request) + } + + override class func canInit(with request: URLRequest) -> Bool { true } + override class func canonicalRequest(for request: URLRequest) -> URLRequest { request } + override func startLoading() { + Task { + do { + let (response, data) = try await self.executeHandler(for: request) + client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed) + client?.urlProtocol(self, didLoad: data) + client?.urlProtocolDidFinishLoading(self) + } catch { + client?.urlProtocol(self, didFailWithError: error) + } + } + } + override func stopLoading() {} + } + + // MARK: - Per-test session factory + + /// Creates an isolated URLSession + unique key for one test. + /// + /// Usage: + /// ```swift + /// let (session, key) = makeIsolatedSession() + /// await IsolatedMockURLProtocol.setHandler(key: key) { _ in (response, data) } + /// ``` + func makeIsolatedSession() -> (session: URLSession, key: String) { + let key = UUID().uuidString + return (IsolatedMockURLProtocol.makeSession(key: key), key) + } + +#endif diff --git a/Tests/MCPTests/OAuthDiscoveryClientTests.swift b/Tests/MCPTests/OAuthDiscoveryClientTests.swift new file mode 100644 index 0000000..b962bb6 --- /dev/null +++ b/Tests/MCPTests/OAuthDiscoveryClientTests.swift @@ -0,0 +1,239 @@ +@preconcurrency import Foundation +import Testing + +@testable import MCP + +#if canImport(FoundationNetworking) + import FoundationNetworking +#endif + +#if swift(>=6.1) && !os(Linux) + + @Suite("OAuthDiscoveryClient", .serialized) + struct OAuthDiscoveryClientTests { + + let urlValidator = OAuthURLValidator(allowLoopbackHTTPForAuthorizationServer: true) + let metadataDiscovery = DefaultOAuthMetadataDiscovery() + + func makeClient() -> OAuthDiscoveryClient { + OAuthDiscoveryClient(metadataDiscovery: metadataDiscovery, urlValidator: urlValidator) + } + + func makeProtectedResourceBody(authorizationServers: [String]) throws -> Data { + let dict: [String: Any] = ["authorization_servers": authorizationServers] + return try JSONSerialization.data(withJSONObject: dict) + } + + func makeASMetadataBody(issuer: String) throws -> Data { + let dict: [String: Any] = [ + "issuer": issuer, + "token_endpoint": "https://auth.example.com/token", + "code_challenge_methods_supported": ["S256"], + ] + return try JSONSerialization.data(withJSONObject: dict) + } + + // MARK: - fetchProtectedResourceMetadata + + @Test("Returns metadata from first successful candidate") + func testFetchProtectedResourceMetadataSuccess() async throws { + let body = try makeProtectedResourceBody( + authorizationServers: ["https://auth.example.com"]) + let (session, key) = makeIsolatedSession() + await IsolatedMockURLProtocol.setHandler(key: key) { _ in + let response = HTTPURLResponse( + url: URL(string: "https://example.com/.well-known/oauth-protected-resource")!, + statusCode: 200, httpVersion: nil, headerFields: nil)! + return (response, body) + } + + let metadata = try await makeClient().fetchProtectedResourceMetadata( + candidates: [URL(string: "https://example.com/.well-known/oauth-protected-resource")!], + fallbackIssuer: nil, + session: session + ) + let expected = OAuthProtectedResourceMetadata( + resource: nil, + authorizationServers: [URL(string: "https://auth.example.com")!], + scopesSupported: nil) + #expect(metadata == expected) + } + + @Test("Skips candidates that return non-2xx status") + func testFetchProtectedResourceMetadataSkipsNon2xx() async throws { + let body = try makeProtectedResourceBody( + authorizationServers: ["https://auth.example.com"]) + let (session, key) = makeIsolatedSession() + await IsolatedMockURLProtocol.setHandler(key: key) { request in + let statusCode = request.url?.lastPathComponent == "mcp" ? 404 : 200 + let response = HTTPURLResponse( + url: request.url!, statusCode: statusCode, + httpVersion: nil, headerFields: nil)! + return (response, statusCode == 200 ? body : Data()) + } + + let metadata = try await makeClient().fetchProtectedResourceMetadata( + candidates: [ + URL(string: "https://example.com/.well-known/oauth-protected-resource/mcp")!, + URL(string: "https://example.com/.well-known/oauth-protected-resource")!, + ], + fallbackIssuer: nil, + session: session + ) + let expected = OAuthProtectedResourceMetadata( + resource: nil, + authorizationServers: [URL(string: "https://auth.example.com")!], + scopesSupported: nil) + #expect(metadata == expected) + } + + @Test("Skips candidates with empty authorizationServers array") + func testFetchProtectedResourceMetadataSkipsEmptyAuthServers() async throws { + let emptyBody = try makeProtectedResourceBody(authorizationServers: []) + let validBody = try makeProtectedResourceBody( + authorizationServers: ["https://auth.example.com"]) + let (session, key) = makeIsolatedSession() + await IsolatedMockURLProtocol.setHandler(key: key) { request in + let body = request.url?.lastPathComponent == "mcp" ? emptyBody : validBody + let response = HTTPURLResponse( + url: request.url!, statusCode: 200, httpVersion: nil, headerFields: nil)! + return (response, body) + } + + let metadata = try await makeClient().fetchProtectedResourceMetadata( + candidates: [ + URL(string: "https://example.com/.well-known/oauth-protected-resource/mcp")!, + URL(string: "https://example.com/.well-known/oauth-protected-resource")!, + ], + fallbackIssuer: nil, + session: session + ) + let expected = OAuthProtectedResourceMetadata( + resource: nil, + authorizationServers: [URL(string: "https://auth.example.com")!], + scopesSupported: nil) + #expect(metadata == expected) + } + + @Test("Throws metadataDiscoveryFailed when all candidates fail and no fallback issuer") + func testFetchProtectedResourceMetadataThrowsWhenAllFail() async throws { + let (session, key) = makeIsolatedSession() + await IsolatedMockURLProtocol.setHandler(key: key) { request in + let response = HTTPURLResponse( + url: request.url!, statusCode: 404, httpVersion: nil, headerFields: nil)! + return (response, Data()) + } + + await #expect(throws: OAuthAuthorizationError.self) { + try await makeClient().fetchProtectedResourceMetadata( + candidates: [ + URL(string: "https://example.com/.well-known/oauth-protected-resource")! + ], + fallbackIssuer: nil, + session: session + ) + } + } + + @Test("Returns synthetic metadata with fallback issuer when all candidates fail") + func testFetchProtectedResourceMetadataUsesFallbackIssuer() async throws { + let fallback = URL(string: "https://example.com")! + let (session, key) = makeIsolatedSession() + await IsolatedMockURLProtocol.setHandler(key: key) { request in + let response = HTTPURLResponse( + url: request.url!, statusCode: 404, httpVersion: nil, headerFields: nil)! + return (response, Data()) + } + + let metadata = try await makeClient().fetchProtectedResourceMetadata( + candidates: [URL(string: "https://example.com/.well-known/oauth-protected-resource")!], + fallbackIssuer: fallback, + session: session + ) + let expected = OAuthProtectedResourceMetadata( + resource: nil, + authorizationServers: [fallback], + scopesSupported: nil) + #expect(metadata == expected) + } + + // MARK: - fetchAuthorizationServerMetadata + + @Test("Returns server and metadata when issuer matches") + func testFetchAuthorizationServerMetadataSuccess() async throws { + let issuer = "https://auth.example.com" + let body = try makeASMetadataBody(issuer: issuer) + let (session, key) = makeIsolatedSession() + await IsolatedMockURLProtocol.setHandler(key: key) { _ in + let response = HTTPURLResponse( + url: URL(string: "\(issuer)/.well-known/oauth-authorization-server")!, + statusCode: 200, httpVersion: nil, headerFields: nil)! + return (response, body) + } + + let (server, metadata) = try await makeClient().fetchAuthorizationServerMetadata( + candidates: [URL(string: issuer)!], + session: session + ) + let expectedServer = URL(string: issuer)! + let expectedMetadata = OAuthAuthorizationServerMetadata( + issuer: URL(string: issuer), + authorizationEndpoint: nil, + tokenEndpoint: URL(string: "https://auth.example.com/token"), + registrationEndpoint: nil, + codeChallengeMethodsSupported: ["S256"], + tokenEndpointAuthMethodsSupported: nil, + clientIDMetadataDocumentSupported: nil) + #expect(server == expectedServer) + #expect(metadata == expectedMetadata) + } + + @Test("Uses metadata issuer as server identity when it differs from candidate URL") + func testFetchAuthorizationServerMetadataUsesMetadataIssuer() async throws { + let metadataIssuer = "https://other.example.com" + let body = try makeASMetadataBody(issuer: metadataIssuer) + let (session, key) = makeIsolatedSession() + await IsolatedMockURLProtocol.setHandler(key: key) { _ in + let response = HTTPURLResponse( + url: URL(string: "https://auth.example.com")!, + statusCode: 200, httpVersion: nil, headerFields: nil)! + return (response, body) + } + + let (server, _) = try await makeClient().fetchAuthorizationServerMetadata( + candidates: [URL(string: "https://auth.example.com")!], + session: session + ) + #expect(server == URL(string: metadataIssuer)!) + } + + @Test("Skips private IP candidates without making HTTP calls") + func testFetchAuthorizationServerMetadataSkipsPrivateIP() async throws { + let (session, _) = makeIsolatedSession() + await #expect(throws: OAuthAuthorizationError.self) { + try await makeClient().fetchAuthorizationServerMetadata( + candidates: [URL(string: "https://10.0.0.1")!], + session: session + ) + } + } + + @Test("Throws when all candidates return non-2xx") + func testFetchAuthorizationServerMetadataThrowsWhenAllFail() async throws { + let (session, key) = makeIsolatedSession() + await IsolatedMockURLProtocol.setHandler(key: key) { request in + let response = HTTPURLResponse( + url: request.url!, statusCode: 500, httpVersion: nil, headerFields: nil)! + return (response, Data()) + } + + await #expect(throws: OAuthAuthorizationError.self) { + try await makeClient().fetchAuthorizationServerMetadata( + candidates: [URL(string: "https://auth.example.com")!], + session: session + ) + } + } + } + +#endif diff --git a/Tests/MCPTests/OAuthTestScenarios.swift b/Tests/MCPTests/OAuthTestScenarios.swift new file mode 100644 index 0000000..62d7457 --- /dev/null +++ b/Tests/MCPTests/OAuthTestScenarios.swift @@ -0,0 +1,2361 @@ +@preconcurrency import Foundation +import Testing + +@testable import MCP + +#if canImport(FoundationNetworking) + import FoundationNetworking +#endif + +#if swift(>=6.1) + + // MARK: - Scenario Context + + struct OAuthScenarioContext: Sendable { + let testEndpoint: URL + let oauthConfiguration: OAuthConfiguration + let messageData: Data + let expectedResponseData: Data? + let expectedCallCounts: [URL: Int] + let streaming: Bool + let sseInitializationTimeout: TimeInterval? + let expectedErrorSubstring: String? + let unexpectedErrorSubstrings: [String] + let secondMessageData: Data? + let secondExpectedResponseData: Data? + + init( + testEndpoint: URL, + oauthConfiguration: OAuthConfiguration, + messageData: Data, + expectedResponseData: Data? = nil, + expectedCallCounts: [URL: Int] = [:], + streaming: Bool = false, + sseInitializationTimeout: TimeInterval? = nil, + expectedErrorSubstring: String? = nil, + unexpectedErrorSubstrings: [String] = [], + secondMessageData: Data? = nil, + secondExpectedResponseData: Data? = nil + ) { + self.testEndpoint = testEndpoint + self.oauthConfiguration = oauthConfiguration + self.messageData = messageData + self.expectedResponseData = expectedResponseData + self.expectedCallCounts = expectedCallCounts + self.streaming = streaming + self.sseInitializationTimeout = sseInitializationTimeout + self.expectedErrorSubstring = expectedErrorSubstring + self.unexpectedErrorSubstrings = unexpectedErrorSubstrings + self.secondMessageData = secondMessageData + self.secondExpectedResponseData = secondExpectedResponseData + } + } + + // MARK: - Request Body Helper + + func readRequestBody(_ request: URLRequest) -> Data? { + if let data = request.httpBody { return data } + guard let stream = request.httpBodyStream else { return nil } + stream.open() + defer { stream.close() } + let bufferSize = 4096 + let buffer = UnsafeMutablePointer.allocate(capacity: bufferSize) + defer { buffer.deallocate() } + var data = Data() + while stream.hasBytesAvailable { + let bytesRead = stream.read(buffer, maxLength: bufferSize) + data.append(buffer, count: bytesRead) + } + return data + } + + // MARK: - Trackers + + actor ProviderTracker: Sendable { + var capturedContext: OAuthConfiguration.AccessTokenProviderContext? + func capture(_ context: OAuthConfiguration.AccessTokenProviderContext) { + capturedContext = context + } + } + + actor OrderTracker: Sendable { + var requests: [URL] = [] + func append(_ url: URL) { requests.append(url) } + func count() -> Int { requests.count } + } + + /// Dispenses tokens in order; repeats the last token when the list is exhausted. + actor TokenDispenser: Sendable { + private var tokens: [String] + init(_ tokens: String...) { self.tokens = Array(tokens) } + func next() -> String { + if tokens.count > 1 { return tokens.removeFirst() } + return tokens.first ?? "fallback-token" + } + } + + // MARK: - Scenario Definitions + + #if !canImport(FoundationNetworking) + + extension RequestHandlerStorage { + + // MARK: 1 - Client Credentials Retry After 401 + + func configureOAuthClientCredentialsRetryAfter401() -> OAuthScenarioContext { + let testEndpoint = URL(string: "https://localhost:8080/test")! + let resourceMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-protected-resource/test")! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + let finalResponseData = MockResponses.jsonRPCResult(id: 1) + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { request in + if request.value(forHTTPHeaderField: "Authorization") == nil { + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL, scope: "files:read" + )(request) + } + #expect( + request.value(forHTTPHeaderField: "Authorization") + == "Bearer access-token-123") + return try await MockResponses.jsonSuccess(body: finalResponseData)(request) + }, + resourceMetadataURL: MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth"], + scopesSupported: ["files:read", "files:write"] + ), + asMetadataURL: MockResponses.asMetadata( + issuer: "https://localhost:8080/auth", + tokenEndpoint: "https://localhost:8080/oauth/token" + ), + tokenEndpointURL: { request in + #expect(request.httpMethod == "POST") + let body = + String(data: readRequestBody(request) ?? Data(), encoding: .utf8) ?? "" + #expect(body.contains("grant_type=client_credentials")) + #expect( + body.contains("resource=https%3A%2F%2Flocalhost%3A8080%2Ftest")) + #expect(body.contains("scope=files%3Aread")) + #expect(body.contains("client_id=test-client")) + return try await MockResponses.tokenSuccess( + accessToken: "access-token-123")(request) + }, + ]) + + return OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init(authentication: .none(clientID: "test-client")), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":1}"#.data(using: .utf8)!, + expectedResponseData: finalResponseData, + expectedCallCounts: [ + testEndpoint: 2, resourceMetadataURL: 1, asMetadataURL: 1, + tokenEndpointURL: 1, + ] + ) + } + + // MARK: 2 - Scope Fallback to scopes_supported + + func configureOAuthScopeSelectionFallsBackToScopesSupported() -> OAuthScenarioContext { + let testEndpoint = URL(string: "https://localhost:8080/fallback-scope")! + let resourceMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/fallback-scope" + )! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + let finalResponseData = MockResponses.jsonRPCResult(id: 21) + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { request in + if request.value(forHTTPHeaderField: "Authorization") == nil { + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL + )(request) + } + #expect( + request.value(forHTTPHeaderField: "Authorization") + == "Bearer access-token-fallback") + return try await MockResponses.jsonSuccess(body: finalResponseData)(request) + }, + resourceMetadataURL: MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth"], + scopesSupported: ["files:write", "files:read"] + ), + asMetadataURL: MockResponses.asMetadata( + issuer: "https://localhost:8080/auth", + tokenEndpoint: "https://localhost:8080/oauth/token" + ), + tokenEndpointURL: { request in + let body = + String(data: readRequestBody(request) ?? Data(), encoding: .utf8) ?? "" + #expect(body.contains("grant_type=client_credentials")) + #expect( + body.contains( + "resource=https%3A%2F%2Flocalhost%3A8080%2Ffallback-scope")) + #expect(body.contains("scope=files%3Aread%20files%3Awrite")) + #expect(body.contains("client_id=test-client")) + return try await MockResponses.tokenSuccess( + accessToken: "access-token-fallback")(request) + }, + ]) + + return OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init(authentication: .none(clientID: "test-client")), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":21}"#.data(using: .utf8)!, + expectedResponseData: finalResponseData, + expectedCallCounts: [ + testEndpoint: 2, resourceMetadataURL: 1, asMetadataURL: 1, + tokenEndpointURL: 1, + ] + ) + } + + // MARK: 3 - Scope Omitted When No Hints + + func configureOAuthScopeOmittedWhenNoHints() -> OAuthScenarioContext { + let testEndpoint = URL(string: "https://localhost:8080/no-scope-hints")! + let resourceMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/no-scope-hints" + )! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + let finalResponseData = MockResponses.jsonRPCResult(id: 22) + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { request in + if request.value(forHTTPHeaderField: "Authorization") == nil { + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL + )(request) + } + #expect( + request.value(forHTTPHeaderField: "Authorization") + == "Bearer access-token-no-scope") + return try await MockResponses.jsonSuccess(body: finalResponseData)(request) + }, + resourceMetadataURL: MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth"] + ), + asMetadataURL: MockResponses.asMetadata( + issuer: "https://localhost:8080/auth", + tokenEndpoint: "https://localhost:8080/oauth/token" + ), + tokenEndpointURL: { request in + let body = + String(data: readRequestBody(request) ?? Data(), encoding: .utf8) ?? "" + #expect(body.contains("grant_type=client_credentials")) + #expect( + body.contains( + "resource=https%3A%2F%2Flocalhost%3A8080%2Fno-scope-hints")) + #expect(!body.contains("scope=")) + #expect(body.contains("client_id=test-client")) + return try await MockResponses.tokenSuccess( + accessToken: "access-token-no-scope")(request) + }, + ]) + + return OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init(authentication: .none(clientID: "test-client")), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":22}"#.data(using: .utf8)!, + expectedResponseData: finalResponseData, + expectedCallCounts: [ + testEndpoint: 2, resourceMetadataURL: 1, asMetadataURL: 1, + tokenEndpointURL: 1, + ] + ) + } + + // MARK: 4 - Resource Parameter in Authorization Code Flow + + func configureOAuthResourceParameterInAuthorizationAndToken() -> OAuthScenarioContext { + let testEndpoint = URL(string: "https://localhost:8080/public/mcp?foo=bar")! + let canonicalResource = "https://localhost:8080/public/mcp" + let resourceMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/public/mcp")! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + let authorizationEndpointURL = URL( + string: "https://localhost:8080/oauth/authorize")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + let finalResponseData = MockResponses.jsonRPCResult(id: 23) + + requestHandler = { + [testEndpoint, canonicalResource, resourceMetadataURL, asMetadataURL, authorizationEndpointURL, tokenEndpointURL, finalResponseData] + request in + guard let url = request.url else { + throw MockResponses.mockError("Missing request URL") + } + + if url.scheme == authorizationEndpointURL.scheme, + url.host == authorizationEndpointURL.host, + url.port == authorizationEndpointURL.port, + url.path == authorizationEndpointURL.path + { + let queryItems = + URLComponents(url: url, resolvingAgainstBaseURL: false)?.queryItems + ?? [] + #expect( + queryItems.contains(where: { + $0.name == "resource" && $0.value == canonicalResource + })) + #expect( + queryItems.contains(where: { + $0.name == "response_type" && $0.value == "code" + })) + #expect( + queryItems.contains(where: { + $0.name == "client_id" && $0.value == "test-client" + })) + #expect( + queryItems.contains(where: { + $0.name == "scope" && $0.value == "files:read" + })) + let state = queryItems.first(where: { $0.name == "state" })?.value + let redirectURI = + queryItems.first(where: { $0.name == "redirect_uri" })?.value + #expect(state != nil) + #expect(redirectURI != nil) + + var redirectComponents = URLComponents(string: redirectURI ?? "") + var redirectQueryItems = redirectComponents?.queryItems ?? [] + redirectQueryItems.append(.init(name: "code", value: "test")) + redirectQueryItems.append(.init(name: "state", value: state)) + redirectComponents?.queryItems = redirectQueryItems + let locationValue = + redirectComponents?.url?.absoluteString + ?? "http://127.0.0.1:3000/callback?code=test&state=\(state ?? "")" + let response = HTTPURLResponse( + url: url, statusCode: 302, httpVersion: "HTTP/1.1", + headerFields: ["Location": locationValue])! + return (response, Data()) + } + + switch url { + case testEndpoint: + if request.value(forHTTPHeaderField: "Authorization") == nil { + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL, scope: "files:read" + )(request) + } + #expect( + request.value(forHTTPHeaderField: "Authorization") + == "Bearer access-token-resource") + #expect(request.url?.query == "foo=bar") + return try await MockResponses.jsonSuccess(body: finalResponseData)( + request) + + case resourceMetadataURL: + return try await MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth"], + scopesSupported: ["files:read"] + )(request) + + case asMetadataURL: + return try await MockResponses.asMetadata( + issuer: "https://localhost:8080/auth", + tokenEndpoint: "https://localhost:8080/oauth/token", + authorizationEndpoint: "https://localhost:8080/oauth/authorize", + codeChallengeMethodsSupported: ["S256"] + )(request) + + case tokenEndpointURL: + let body = + String(data: readRequestBody(request) ?? Data(), encoding: .utf8) ?? "" + #expect(body.contains("grant_type=authorization_code")) + #expect( + body.contains( + "resource=https%3A%2F%2Flocalhost%3A8080%2Fpublic%2Fmcp")) + #expect(!body.contains("%3Ffoo%3Dbar")) + #expect(body.contains("scope=files%3Aread")) + #expect(body.contains("client_id=test-client")) + #expect(body.contains("code_verifier=")) + #expect(body.contains("code=test")) + return try await MockResponses.tokenSuccess( + accessToken: "access-token-resource")(request) + + default: + throw MockResponses.mockError("Unexpected URL: \(url.absoluteString)") + } + } + + return OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init( + grantType: .authorizationCode, + authentication: .none(clientID: "test-client") + ), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":23}"#.data(using: .utf8)!, + expectedResponseData: finalResponseData, + expectedCallCounts: [ + testEndpoint: 2, resourceMetadataURL: 1, asMetadataURL: 1, + tokenEndpointURL: 1, + ] + ) + } + + // MARK: 5 - Rejects Authorization Without PKCE Metadata + + func configureOAuthRejectsAuthorizationWithoutPKCEMetadata() -> OAuthScenarioContext { + let testEndpoint = URL(string: "https://localhost:8080/pkce-metadata-missing")! + let resourceMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/pkce-metadata-missing" + )! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + let authorizationEndpointURL = URL( + string: "https://localhost:8080/oauth/authorize")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + + requestHandler = { + [testEndpoint, resourceMetadataURL, asMetadataURL, authorizationEndpointURL, tokenEndpointURL] + request in + guard let url = request.url else { + throw MockResponses.mockError("Missing request URL") + } + + if url.scheme == authorizationEndpointURL.scheme, + url.host == authorizationEndpointURL.host, + url.port == authorizationEndpointURL.port, + url.path == authorizationEndpointURL.path + { + let response = HTTPURLResponse( + url: url, statusCode: 302, httpVersion: "HTTP/1.1", + headerFields: [ + "Location": "http://127.0.0.1:3000/callback?code=test" + ])! + return (response, Data()) + } + + switch url { + case testEndpoint: + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL, scope: "files:read" + )(request) + + case resourceMetadataURL: + return try await MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth"], + scopesSupported: ["files:read"] + )(request) + + case asMetadataURL: + return try await MockResponses.asMetadata( + issuer: "https://localhost:8080/auth", + tokenEndpoint: "https://localhost:8080/oauth/token", + authorizationEndpoint: "https://localhost:8080/oauth/authorize" + )(request) + + case tokenEndpointURL: + return try await MockResponses.tokenSuccess( + accessToken: "should-not-be-issued")(request) + + default: + throw MockResponses.mockError("Unexpected URL: \(url.absoluteString)") + } + } + + return OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init( + grantType: .authorizationCode, + authentication: .none(clientID: "test-client") + ), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":24}"#.data(using: .utf8)!, + expectedCallCounts: [tokenEndpointURL: 0], + expectedErrorSubstring: "code_challenge_methods_supported" + ) + } + + // MARK: 6 - Rejects Authorization Without S256 PKCE + + func configureOAuthRejectsAuthorizationWithoutS256PKCE() -> OAuthScenarioContext { + let testEndpoint = URL(string: "https://localhost:8080/pkce-s256-missing")! + let resourceMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/pkce-s256-missing" + )! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + let authorizationEndpointURL = URL( + string: "https://localhost:8080/oauth/authorize")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + + requestHandler = { + [testEndpoint, resourceMetadataURL, asMetadataURL, authorizationEndpointURL, tokenEndpointURL] + request in + guard let url = request.url else { + throw MockResponses.mockError("Missing request URL") + } + + if url.scheme == authorizationEndpointURL.scheme, + url.host == authorizationEndpointURL.host, + url.port == authorizationEndpointURL.port, + url.path == authorizationEndpointURL.path + { + let response = HTTPURLResponse( + url: url, statusCode: 302, httpVersion: "HTTP/1.1", + headerFields: [ + "Location": "http://127.0.0.1:3000/callback?code=test" + ])! + return (response, Data()) + } + + switch url { + case testEndpoint: + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL, scope: "files:read" + )(request) + + case resourceMetadataURL: + return try await MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth"], + scopesSupported: ["files:read"] + )(request) + + case asMetadataURL: + return try await MockResponses.asMetadata( + issuer: "https://localhost:8080/auth", + tokenEndpoint: "https://localhost:8080/oauth/token", + authorizationEndpoint: "https://localhost:8080/oauth/authorize", + codeChallengeMethodsSupported: ["plain"] + )(request) + + case tokenEndpointURL: + return try await MockResponses.tokenSuccess( + accessToken: "should-not-be-issued")(request) + + default: + throw MockResponses.mockError("Unexpected URL: \(url.absoluteString)") + } + } + + return OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init( + grantType: .authorizationCode, + authentication: .none(clientID: "test-client") + ), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":25}"#.data(using: .utf8)!, + expectedCallCounts: [tokenEndpointURL: 0], + expectedErrorSubstring: "must support PKCE S256" + ) + } + + // MARK: 7 - Rejects Authorization Response Redirect Mismatch + + func configureOAuthRejectsAuthorizationResponseRedirectMismatch() + -> OAuthScenarioContext + { + let testEndpoint = URL( + string: "https://localhost:8080/authorization-redirect-mismatch")! + let resourceMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/authorization-redirect-mismatch" + )! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + let authorizationEndpointURL = URL( + string: "https://localhost:8080/oauth/authorize")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + + requestHandler = { + [testEndpoint, resourceMetadataURL, asMetadataURL, authorizationEndpointURL, tokenEndpointURL] + request in + guard let url = request.url else { + throw MockResponses.mockError("Missing request URL") + } + + if url.scheme == authorizationEndpointURL.scheme, + url.host == authorizationEndpointURL.host, + url.port == authorizationEndpointURL.port, + url.path == authorizationEndpointURL.path + { + let queryItems = + URLComponents(url: url, resolvingAgainstBaseURL: false)?.queryItems + ?? [] + let state = + queryItems.first(where: { $0.name == "state" })?.value ?? "" + let locationValue = + "https://evil.example.com/callback?code=test&state=\(state)" + let response = HTTPURLResponse( + url: url, statusCode: 302, httpVersion: "HTTP/1.1", + headerFields: ["Location": locationValue])! + return (response, Data()) + } + + switch url { + case testEndpoint: + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL, scope: "files:read" + )(request) + + case resourceMetadataURL: + return try await MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth"], + scopesSupported: ["files:read"] + )(request) + + case asMetadataURL: + return try await MockResponses.asMetadata( + issuer: "https://localhost:8080/auth", + tokenEndpoint: "https://localhost:8080/oauth/token", + authorizationEndpoint: "https://localhost:8080/oauth/authorize", + codeChallengeMethodsSupported: ["S256"] + )(request) + + case tokenEndpointURL: + return try await MockResponses.tokenSuccess( + accessToken: "should-not-be-issued")(request) + + default: + throw MockResponses.mockError("Unexpected URL: \(url.absoluteString)") + } + } + + return OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init( + grantType: .authorizationCode, + authentication: .none(clientID: "test-client") + ), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":26}"#.data(using: .utf8)!, + expectedCallCounts: [tokenEndpointURL: 0], + expectedErrorSubstring: "redirect URI mismatch" + ) + } + + // MARK: 8 - Rejects Authorization Response State Mismatch + + func configureOAuthRejectsAuthorizationResponseStateMismatch() + -> OAuthScenarioContext + { + let testEndpoint = URL( + string: "https://localhost:8080/authorization-state-mismatch")! + let resourceMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/authorization-state-mismatch" + )! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + let authorizationEndpointURL = URL( + string: "https://localhost:8080/oauth/authorize")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + + requestHandler = { + [testEndpoint, resourceMetadataURL, asMetadataURL, authorizationEndpointURL, tokenEndpointURL] + request in + guard let url = request.url else { + throw MockResponses.mockError("Missing request URL") + } + + if url.scheme == authorizationEndpointURL.scheme, + url.host == authorizationEndpointURL.host, + url.port == authorizationEndpointURL.port, + url.path == authorizationEndpointURL.path + { + let queryItems = + URLComponents(url: url, resolvingAgainstBaseURL: false)?.queryItems + ?? [] + let redirectURI = + queryItems.first(where: { $0.name == "redirect_uri" })?.value + var redirectComponents = URLComponents(string: redirectURI ?? "") + var redirectQueryItems = redirectComponents?.queryItems ?? [] + redirectQueryItems.append(.init(name: "code", value: "test")) + redirectQueryItems.append( + .init(name: "state", value: "unexpected-state")) + redirectComponents?.queryItems = redirectQueryItems + let locationValue = + redirectComponents?.url?.absoluteString + ?? "http://127.0.0.1:3000/callback?code=test&state=unexpected-state" + let response = HTTPURLResponse( + url: url, statusCode: 302, httpVersion: "HTTP/1.1", + headerFields: ["Location": locationValue])! + return (response, Data()) + } + + switch url { + case testEndpoint: + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL, scope: "files:read" + )(request) + + case resourceMetadataURL: + return try await MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth"], + scopesSupported: ["files:read"] + )(request) + + case asMetadataURL: + return try await MockResponses.asMetadata( + issuer: "https://localhost:8080/auth", + tokenEndpoint: "https://localhost:8080/oauth/token", + authorizationEndpoint: "https://localhost:8080/oauth/authorize", + codeChallengeMethodsSupported: ["S256"] + )(request) + + case tokenEndpointURL: + return try await MockResponses.tokenSuccess( + accessToken: "should-not-be-issued")(request) + + default: + throw MockResponses.mockError("Unexpected URL: \(url.absoluteString)") + } + } + + return OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init( + grantType: .authorizationCode, + authentication: .none(clientID: "test-client") + ), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":27}"#.data(using: .utf8)!, + expectedCallCounts: [tokenEndpointURL: 0], + expectedErrorSubstring: "state mismatch" + ) + } + + // MARK: 9 - Access Token Only Via Authorization Header + + func configureOAuthAccessTokenOnlyViaAuthorizationHeader() -> OAuthScenarioContext { + let testEndpoint = URL(string: "https://localhost:8080/test?foo=bar")! + let resourceMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-protected-resource/test")! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + let finalResponseData = MockResponses.jsonRPCResult(id: 12) + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { request in + if request.value(forHTTPHeaderField: "Authorization") == nil { + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL, scope: "files:read" + )(request) + } + #expect( + request.value(forHTTPHeaderField: "Authorization") + == "Bearer access-token-xyz") + #expect(request.url?.query == "foo=bar") + #expect( + !(request.url?.absoluteString.contains("access_token=") ?? false)) + let requestBody = + String(data: readRequestBody(request) ?? Data(), encoding: .utf8) ?? "" + #expect(!requestBody.contains("access_token=")) + #expect( + request.value(forHTTPHeaderField: "Content-Type") + == "application/json") + return try await MockResponses.jsonSuccess(body: finalResponseData)( + request) + }, + resourceMetadataURL: MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth"], + scopesSupported: ["files:read"] + ), + asMetadataURL: MockResponses.asMetadata( + issuer: "https://localhost:8080/auth", + tokenEndpoint: "https://localhost:8080/oauth/token" + ), + tokenEndpointURL: MockResponses.tokenSuccess( + accessToken: "access-token-xyz"), + ]) + + return OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init(authentication: .none(clientID: "test-client")), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":12}"#.data(using: .utf8)!, + expectedResponseData: finalResponseData, + expectedCallCounts: [ + testEndpoint: 2, resourceMetadataURL: 1, asMetadataURL: 1, + tokenEndpointURL: 1, + ] + ) + } + + // MARK: 10 - Authorization Header For Every Request In Session + + func configureOAuthAuthorizationHeaderForEveryRequestInSession() + -> OAuthScenarioContext + { + let testEndpoint = URL(string: "https://localhost:8080/session-auth")! + let resourceMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/session-auth")! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + let firstResponseData = #"{"jsonrpc":"2.0","result":{"ok":true},"id":31}"#.data( + using: .utf8)! + let secondResponseData = #"{"jsonrpc":"2.0","result":{"ok":true},"id":32}"#.data( + using: .utf8)! + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { + [firstResponseData, secondResponseData] request in + if request.value(forHTTPHeaderField: "Authorization") == nil { + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL, scope: "files:read" + )(request) + } + #expect( + request.value(forHTTPHeaderField: "Authorization") + == "Bearer session-access-token") + #expect( + !(request.url?.absoluteString.contains("access_token=") ?? false)) + let requestBody = + String(data: readRequestBody(request) ?? Data(), encoding: .utf8) ?? "" + #expect(!requestBody.contains("access_token=")) + + let responseBody: Data + if requestBody.contains("\"id\":31") { + responseBody = firstResponseData + } else if requestBody.contains("\"id\":32") { + responseBody = secondResponseData + } else { + throw MockResponses.mockError( + "Unexpected JSON-RPC body: \(requestBody)") + } + + let response = HTTPURLResponse( + url: request.url!, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [ + "Content-Type": "application/json", + "MCP-Session-Id": "session-123", + ])! + return (response, responseBody) + }, + resourceMetadataURL: MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth"], + scopesSupported: ["files:read"] + ), + asMetadataURL: MockResponses.asMetadata( + issuer: "https://localhost:8080/auth", + tokenEndpoint: "https://localhost:8080/oauth/token" + ), + tokenEndpointURL: { request in + let body = + String(data: readRequestBody(request) ?? Data(), encoding: .utf8) ?? "" + #expect(body.contains("grant_type=client_credentials")) + #expect( + body.contains( + "resource=https%3A%2F%2Flocalhost%3A8080%2Fsession-auth")) + #expect(body.contains("scope=files%3Aread")) + return try await MockResponses.tokenSuccess( + accessToken: "session-access-token")(request) + }, + ]) + + return OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init(authentication: .none(clientID: "test-client")), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":31}"#.data(using: .utf8)!, + expectedResponseData: firstResponseData, + expectedCallCounts: [ + testEndpoint: 3, resourceMetadataURL: 1, asMetadataURL: 1, + tokenEndpointURL: 1, + ], + secondMessageData: #"{"jsonrpc":"2.0","method":"ping","id":32}"#.data( + using: .utf8)!, + secondExpectedResponseData: secondResponseData + ) + } + + // MARK: 11 - Streaming GET Uses Authorization Header Only + + func configureOAuthStreamingGETUsesAuthorizationHeaderOnly() -> OAuthScenarioContext { + let testEndpoint = URL(string: "https://localhost:8080/stream-auth?foo=bar")! + let resourceMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/stream-auth")! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + let sseEventData = "id: evt-1\ndata: {\"stream\":\"ok\"}\n\n".data(using: .utf8)! + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { [sseEventData] request in + if request.httpMethod == "GET" { + #expect( + request.value(forHTTPHeaderField: "Authorization") + == "Bearer stream-access-token") + #expect(request.url?.query == "foo=bar") + #expect( + !(request.url?.absoluteString.contains("access_token=") ?? false)) + #expect( + request.value(forHTTPHeaderField: "MCP-Session-Id") + == "stream-session-id") + let response = HTTPURLResponse( + url: request.url!, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "text/event-stream"])! + return (response, sseEventData) + } + + if request.value(forHTTPHeaderField: "Authorization") == nil { + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL, scope: "files:read" + )(request) + } + + #expect( + request.value(forHTTPHeaderField: "Authorization") + == "Bearer stream-access-token") + #expect( + !(request.url?.absoluteString.contains("access_token=") ?? false)) + let requestBody = + String(data: readRequestBody(request) ?? Data(), encoding: .utf8) ?? "" + #expect(!requestBody.contains("access_token=")) + + let response = HTTPURLResponse( + url: request.url!, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: [ + "Content-Type": "text/plain", + "MCP-Session-Id": "stream-session-id", + ])! + return (response, Data()) + }, + resourceMetadataURL: MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth"], + scopesSupported: ["files:read"] + ), + asMetadataURL: MockResponses.asMetadata( + issuer: "https://localhost:8080/auth", + tokenEndpoint: "https://localhost:8080/oauth/token" + ), + tokenEndpointURL: { request in + let body = + String(data: readRequestBody(request) ?? Data(), encoding: .utf8) ?? "" + #expect(body.contains("grant_type=client_credentials")) + #expect( + body.contains( + "resource=https%3A%2F%2Flocalhost%3A8080%2Fstream-auth")) + #expect(body.contains("scope=files%3Aread")) + return try await MockResponses.tokenSuccess( + accessToken: "stream-access-token")(request) + }, + ]) + + let expectedEventPayload = #"{"stream":"ok"}"#.data(using: .utf8)! + return OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init(authentication: .none(clientID: "test-client")), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":41}"#.data( + using: .utf8)!, + expectedResponseData: expectedEventPayload, + streaming: true, + sseInitializationTimeout: 1 + ) + } + + // MARK: 12 - Rejects Non-Bearer Token Type + + func configureOAuthRejectsNonBearerTokenType() -> OAuthScenarioContext { + let testEndpoint = URL( + string: "https://localhost:8080/non-bearer-token-type")! + let resourceMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/non-bearer-token-type" + )! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { request in + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL, scope: "files:read" + )(request) + }, + resourceMetadataURL: MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth"], + scopesSupported: ["files:read"] + ), + asMetadataURL: MockResponses.asMetadata( + issuer: "https://localhost:8080/auth", + tokenEndpoint: "https://localhost:8080/oauth/token" + ), + tokenEndpointURL: MockResponses.tokenResponse( + accessToken: "non-bearer-token", tokenType: "DPoP"), + ]) + + return OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init(authentication: .none(clientID: "test-client")), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":51}"#.data( + using: .utf8)!, + expectedErrorSubstring: "Token response is invalid" + ) + } + + // MARK: 13 - Token Endpoint Failure Redacts Response Body + + func configureOAuthTokenEndpointFailureRedactsResponseBody() -> OAuthScenarioContext { + let testEndpoint = URL( + string: "https://localhost:8080/token-error-redaction")! + let resourceMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/token-error-redaction" + )! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { request in + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL, scope: "files:read" + )(request) + }, + resourceMetadataURL: MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth"], + scopesSupported: ["files:read"] + ), + asMetadataURL: MockResponses.asMetadata( + issuer: "https://localhost:8080/auth", + tokenEndpoint: "https://localhost:8080/oauth/token" + ), + tokenEndpointURL: MockResponses.tokenError( + error: "invalid_client", + errorDescription: "leaked-secret-value", + extraFields: ["access_token": "should-not-leak"] + ), + ]) + + return OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init(authentication: .none(clientID: "test-client")), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":52}"#.data( + using: .utf8)!, + expectedErrorSubstring: "oauth_error: invalid_client", + unexpectedErrorSubstrings: ["leaked-secret-value", "should-not-leak"] + ) + } + + // MARK: 14 - Rejects Non-HTTPS Token Endpoint + + func configureOAuthRejectsNonHTTPSTokenEndpoint() -> OAuthScenarioContext { + let testEndpoint = URL( + string: "https://localhost:8080/non-https-token-endpoint")! + let resourceMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/non-https-token-endpoint" + )! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { request in + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL, scope: "files:read" + )(request) + }, + resourceMetadataURL: MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth"], + scopesSupported: ["files:read"] + ), + asMetadataURL: MockResponses.asMetadata( + issuer: "https://localhost:8080/auth", + tokenEndpoint: "http://localhost:8080/oauth/token" + ), + ]) + + var oauth = OAuthConfiguration( + authentication: .none(clientID: "test-client")) + oauth.allowLoopbackHTTPAuthorizationServerEndpoints = false + + return OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: oauth, + messageData: #"{"jsonrpc":"2.0","method":"ping","id":53}"#.data( + using: .utf8)!, + expectedErrorSubstring: "Token endpoint must use https" + ) + } + + // MARK: 15 - Allows Loopback HTTP Authorization Server Endpoints + + func configureOAuthAllowsLoopbackHTTPAuthorizationServerEndpoints() + -> OAuthScenarioContext + { + let testEndpoint = URL( + string: "https://localhost:8080/loopback-http-auth-server-enabled")! + let resourceMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/loopback-http-auth-server-enabled" + )! + let asMetadataURL = URL( + string: "http://localhost:8080/.well-known/oauth-authorization-server/auth")! + let tokenEndpointURL = URL(string: "http://localhost:8080/oauth/token")! + let finalResponseData = MockResponses.jsonRPCResult(id: 54) + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { request in + if request.value(forHTTPHeaderField: "Authorization") != nil { + return try await MockResponses.jsonSuccess(body: finalResponseData)( + request) + } + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL, scope: "files:read" + )(request) + }, + resourceMetadataURL: MockResponses.resourceMetadata( + authorizationServers: ["http://localhost:8080/auth"], + scopesSupported: ["files:read"] + ), + asMetadataURL: MockResponses.asMetadata( + issuer: "http://localhost:8080/auth", + tokenEndpoint: "http://localhost:8080/oauth/token" + ), + tokenEndpointURL: { request in + let body = + String(data: readRequestBody(request) ?? Data(), encoding: .utf8) ?? "" + #expect(body.contains("grant_type=client_credentials")) + return try await MockResponses.tokenSuccess( + accessToken: "loopback-http-token")(request) + }, + ]) + + var oauth = OAuthConfiguration( + authentication: .none(clientID: "test-client")) + oauth.allowLoopbackHTTPAuthorizationServerEndpoints = true + + return OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: oauth, + messageData: #"{"jsonrpc":"2.0","method":"ping","id":54}"#.data( + using: .utf8)!, + expectedResponseData: finalResponseData, + expectedCallCounts: [tokenEndpointURL: 1] + ) + } + + // MARK: 16 - Access Token Provider Receives Discovery Context + + func configureOAuthAccessTokenProviderReceivesDiscoveryContext() + -> (OAuthScenarioContext, ProviderTracker) + { + let testEndpoint = URL(string: "https://localhost:8080/test")! + let resourceMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-protected-resource/test")! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + let finalResponseData = MockResponses.jsonRPCResult(id: 2) + let providerTracker = ProviderTracker() + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { request in + if request.value(forHTTPHeaderField: "Authorization") == nil { + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL, + scope: "files:read files:write" + )(request) + } + #expect( + request.value(forHTTPHeaderField: "Authorization") + == "Bearer provider-access-token") + return try await MockResponses.jsonSuccess(body: finalResponseData)( + request) + }, + resourceMetadataURL: MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth"], + scopesSupported: ["files:read", "files:write"] + ), + asMetadataURL: MockResponses.asMetadata( + issuer: "https://localhost:8080/auth", + tokenEndpoint: "https://localhost:8080/oauth/token", + authorizationEndpoint: "https://localhost:8080/oauth/authorize", + codeChallengeMethodsSupported: ["S256"] + ), + tokenEndpointURL: { request in + let response = HTTPURLResponse( + url: request.url!, statusCode: 200, httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "application/json"])! + return (response, Data()) + }, + ]) + + let oauthConfiguration = OAuthConfiguration( + authentication: .none(clientID: "test-client"), + accessTokenProvider: { [providerTracker] providerContext, _ in + await providerTracker.capture(providerContext) + return "provider-access-token" + } + ) + + let context = OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: oauthConfiguration, + messageData: #"{"jsonrpc":"2.0","method":"ping","id":2}"#.data(using: .utf8)!, + expectedResponseData: finalResponseData, + expectedCallCounts: [ + testEndpoint: 2, resourceMetadataURL: 1, asMetadataURL: 1, + tokenEndpointURL: 0, + ] + ) + return (context, providerTracker) + } + + // MARK: 17 - Discovery Uses Header Resource Metadata + + func configureOAuthDiscoveryUsesHeaderResourceMetadata() -> OAuthScenarioContext { + let testEndpoint = URL(string: "https://localhost:8080/public/mcp")! + let headerMetadataURL = URL( + string: "https://localhost:8080/custom-metadata")! + let fallbackPathMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/public/mcp")! + let fallbackRootMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-protected-resource")! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + let finalResponseData = MockResponses.jsonRPCResult(id: 3) + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { request in + if request.value(forHTTPHeaderField: "Authorization") == nil { + return try await MockResponses.bearerChallenge( + resourceMetadataURL: headerMetadataURL, scope: "files:read" + )(request) + } + #expect( + request.value(forHTTPHeaderField: "Authorization") + == "Bearer access-token-123") + return try await MockResponses.jsonSuccess(body: finalResponseData)( + request) + }, + headerMetadataURL: MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth"], + scopesSupported: ["files:read"] + ), + fallbackPathMetadataURL: MockResponses.httpError(statusCode: 500), + fallbackRootMetadataURL: MockResponses.httpError(statusCode: 500), + asMetadataURL: MockResponses.asMetadata( + issuer: "https://localhost:8080/auth", + tokenEndpoint: "https://localhost:8080/oauth/token" + ), + tokenEndpointURL: MockResponses.tokenSuccess( + accessToken: "access-token-123"), + ]) + + return OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init(authentication: .none(clientID: "test-client")), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":3}"#.data(using: .utf8)!, + expectedResponseData: finalResponseData, + expectedCallCounts: [ + testEndpoint: 2, headerMetadataURL: 1, + fallbackPathMetadataURL: 0, fallbackRootMetadataURL: 0, + ] + ) + } + + // MARK: 18 - Discovery Fallback Well-Known Order + + func configureOAuthDiscoveryFallbackWellKnownOrder() + -> (OAuthScenarioContext, OrderTracker) + { + let testEndpoint = URL(string: "https://localhost:8080/public/mcp")! + let fallbackPathMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/public/mcp")! + let fallbackRootMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-protected-resource")! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + let finalResponseData = MockResponses.jsonRPCResult(id: 4) + let tracker = OrderTracker() + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { request in + if request.value(forHTTPHeaderField: "Authorization") == nil { + return try await MockResponses.bearerChallenge( + scope: "files:read")(request) + } + #expect( + request.value(forHTTPHeaderField: "Authorization") + == "Bearer access-token-456") + return try await MockResponses.jsonSuccess(body: finalResponseData)( + request) + }, + fallbackPathMetadataURL: { [tracker] request in + await tracker.append(request.url!) + return try await MockResponses.httpError(statusCode: 404)(request) + }, + fallbackRootMetadataURL: { [tracker] request in + await tracker.append(request.url!) + return try await MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth"], + scopesSupported: ["files:read"] + )(request) + }, + asMetadataURL: MockResponses.asMetadata( + issuer: "https://localhost:8080/auth", + tokenEndpoint: "https://localhost:8080/oauth/token" + ), + tokenEndpointURL: MockResponses.tokenSuccess( + accessToken: "access-token-456"), + ]) + + let context = OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init(authentication: .none(clientID: "test-client")), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":4}"#.data(using: .utf8)!, + expectedResponseData: finalResponseData + ) + return (context, tracker) + } + + // MARK: 19 - Discovery Fails When Metadata Unavailable + + func configureOAuthDiscoveryFailsWhenMetadataUnavailable() + -> (OAuthScenarioContext, OrderTracker) + { + let testEndpoint = URL(string: "https://localhost:8080/public/mcp")! + let fallbackPathMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/public/mcp")! + let fallbackRootMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-protected-resource")! + let tracker = OrderTracker() + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { request in + return try await MockResponses.bearerChallenge( + scope: "files:read")(request) + }, + fallbackPathMetadataURL: { [tracker] request in + await tracker.append(request.url!) + return try await MockResponses.httpError(statusCode: 404)(request) + }, + fallbackRootMetadataURL: { [tracker] request in + await tracker.append(request.url!) + return try await MockResponses.httpError(statusCode: 404)(request) + }, + ]) + + let context = OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init(authentication: .none(clientID: "test-client")), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":7}"#.data(using: .utf8)!, + expectedErrorSubstring: "metadata" + ) + return (context, tracker) + } + + // MARK: 20 - AS Metadata Discovery Order For Path Issuer + + func configureOAuthASMetadataDiscoveryOrderForPathIssuer() + -> (OAuthScenarioContext, OrderTracker) + { + let testEndpoint = URL(string: "https://localhost:8080/public/mcp")! + let resourceMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/public/mcp")! + let authorizationServer = URL(string: "https://localhost:8080/tenant1")! + let asMetadataOAuthInsertedURL = URL( + string: + "https://localhost:8080/.well-known/oauth-authorization-server/tenant1")! + let asMetadataOIDCInsertedURL = URL( + string: "https://localhost:8080/.well-known/openid-configuration/tenant1")! + let asMetadataOIDCAppendedURL = URL( + string: "https://localhost:8080/tenant1/.well-known/openid-configuration")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + let finalResponseData = MockResponses.jsonRPCResult(id: 5) + let tracker = OrderTracker() + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { request in + if request.value(forHTTPHeaderField: "Authorization") == nil { + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL, scope: "files:read" + )(request) + } + #expect( + request.value(forHTTPHeaderField: "Authorization") + == "Bearer path-issuer-token") + return try await MockResponses.jsonSuccess(body: finalResponseData)( + request) + }, + resourceMetadataURL: MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/tenant1"], + scopesSupported: ["files:read"] + ), + asMetadataOAuthInsertedURL: { [tracker] request in + await tracker.append(request.url!) + return try await MockResponses.httpError(statusCode: 404)(request) + }, + asMetadataOIDCInsertedURL: { [tracker] request in + await tracker.append(request.url!) + return try await MockResponses.httpError(statusCode: 404)(request) + }, + asMetadataOIDCAppendedURL: { [tracker] request in + await tracker.append(request.url!) + return try await MockResponses.asMetadata( + issuer: "https://localhost:8080/tenant1", + tokenEndpoint: "https://localhost:8080/oauth/token" + )(request) + }, + tokenEndpointURL: { request in + let body = + String(data: readRequestBody(request) ?? Data(), encoding: .utf8) ?? "" + #expect(body.contains("grant_type=client_credentials")) + #expect(body.contains("scope=files%3Aread")) + return try await MockResponses.tokenSuccess( + accessToken: "path-issuer-token")(request) + }, + authorizationServer: { _ in + throw MockResponses.mockError( + "Unexpected direct issuer request") + }, + ]) + + let context = OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init(authentication: .none(clientID: "test-client")), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":5}"#.data(using: .utf8)!, + expectedResponseData: finalResponseData + ) + return (context, tracker) + } + + // MARK: 21 - AS Metadata Discovery Order For Root Issuer + + func configureOAuthASMetadataDiscoveryOrderForRootIssuer() + -> (OAuthScenarioContext, OrderTracker) + { + let testEndpoint = URL(string: "https://localhost:8080/public/mcp")! + let resourceMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/public/mcp")! + let authorizationServer = URL(string: "https://localhost:8080")! + let asMetadataOAuthURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server")! + let asMetadataOIDCURL = URL( + string: "https://localhost:8080/.well-known/openid-configuration")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + let finalResponseData = MockResponses.jsonRPCResult(id: 6) + let tracker = OrderTracker() + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { request in + if request.value(forHTTPHeaderField: "Authorization") == nil { + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL, scope: "files:read" + )(request) + } + #expect( + request.value(forHTTPHeaderField: "Authorization") + == "Bearer root-issuer-token") + return try await MockResponses.jsonSuccess(body: finalResponseData)( + request) + }, + resourceMetadataURL: MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080"], + scopesSupported: ["files:read"] + ), + asMetadataOAuthURL: { [tracker] request in + await tracker.append(request.url!) + return try await MockResponses.httpError(statusCode: 404)(request) + }, + asMetadataOIDCURL: { [tracker] request in + await tracker.append(request.url!) + return try await MockResponses.asMetadata( + issuer: "https://localhost:8080", + tokenEndpoint: "https://localhost:8080/oauth/token" + )(request) + }, + tokenEndpointURL: MockResponses.tokenSuccess( + accessToken: "root-issuer-token"), + authorizationServer: { _ in + throw MockResponses.mockError( + "Unexpected direct issuer request") + }, + ]) + + let context = OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init(authentication: .none(clientID: "test-client")), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":6}"#.data(using: .utf8)!, + expectedResponseData: finalResponseData + ) + return (context, tracker) + } + + // MARK: 22 - Registration Prefers CIMD When Advertised + + func configureOAuthRegistrationPrefersCIMDWhenAdvertised() -> OAuthScenarioContext { + let testEndpoint = URL(string: "https://localhost:8080/public/mcp")! + let resourceMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/public/mcp")! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + let registrationEndpointURL = URL( + string: "https://localhost:8080/register")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + let clientMetadataDocumentID = "https://client.example.com/metadata.json" + let finalResponseData = MockResponses.jsonRPCResult(id: 8) + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { request in + if request.value(forHTTPHeaderField: "Authorization") == nil { + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL, scope: "files:read" + )(request) + } + #expect( + request.value(forHTTPHeaderField: "Authorization") + == "Bearer cimd-token") + return try await MockResponses.jsonSuccess(body: finalResponseData)( + request) + }, + resourceMetadataURL: MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth"], + scopesSupported: ["files:read"] + ), + asMetadataURL: MockResponses.asMetadata( + issuer: "https://localhost:8080/auth", + tokenEndpoint: "https://localhost:8080/oauth/token", + registrationEndpoint: "https://localhost:8080/register", + clientIDMetadataDocumentSupported: true + ), + registrationEndpointURL: MockResponses.httpError(statusCode: 500), + tokenEndpointURL: { request in + let body = + String(data: readRequestBody(request) ?? Data(), encoding: .utf8) ?? "" + #expect(body.contains("grant_type=client_credentials")) + #expect( + body.contains( + "client_id=https%3A%2F%2Fclient.example.com%2Fmetadata.json")) + return try await MockResponses.tokenSuccess( + accessToken: "cimd-token")(request) + }, + ]) + + return OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init( + authentication: .none(clientID: clientMetadataDocumentID)), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":8}"#.data(using: .utf8)!, + expectedResponseData: finalResponseData, + expectedCallCounts: [registrationEndpointURL: 0] + ) + } + + // MARK: 23 - Pre-Registration Uses Static Credentials + + func configureOAuthPreRegistrationUsesStaticCredentials() -> OAuthScenarioContext { + let testEndpoint = URL(string: "https://localhost:8080/public/mcp")! + let resourceMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/public/mcp")! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + let registrationEndpointURL = URL( + string: "https://localhost:8080/register")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + let finalResponseData = MockResponses.jsonRPCResult(id: 13) + + let expectedClientID = "pre-registered-client" + let expectedClientSecret = "pre-registered-secret" + let expectedBasic = Data( + "\(expectedClientID):\(expectedClientSecret)".utf8 + ).base64EncodedString() + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { request in + if request.value(forHTTPHeaderField: "Authorization") == nil { + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL, scope: "files:read" + )(request) + } + #expect( + request.value(forHTTPHeaderField: "Authorization") + == "Bearer pre-registered-token") + return try await MockResponses.jsonSuccess(body: finalResponseData)( + request) + }, + resourceMetadataURL: MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth"], + scopesSupported: ["files:read"] + ), + asMetadataURL: MockResponses.asMetadata( + issuer: "https://localhost:8080/auth", + tokenEndpoint: "https://localhost:8080/oauth/token", + registrationEndpoint: "https://localhost:8080/register" + ), + registrationEndpointURL: MockResponses.httpError(statusCode: 500), + tokenEndpointURL: { [expectedBasic] request in + #expect(request.httpMethod == "POST") + #expect( + request.value(forHTTPHeaderField: "Authorization") + == "Basic \(expectedBasic)") + let body = + String(data: readRequestBody(request) ?? Data(), encoding: .utf8) ?? "" + #expect(body.contains("grant_type=client_credentials")) + #expect( + body.contains( + "resource=https%3A%2F%2Flocalhost%3A8080%2Fpublic%2Fmcp")) + #expect(body.contains("scope=files%3Aread")) + #expect(!body.contains("client_id=")) + #expect(!body.contains("client_secret=")) + return try await MockResponses.tokenSuccess( + accessToken: "pre-registered-token")(request) + }, + ]) + + return OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init( + authentication: .clientSecretBasic( + clientID: expectedClientID, + clientSecret: expectedClientSecret + )), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":13}"#.data( + using: .utf8)!, + expectedResponseData: finalResponseData, + expectedCallCounts: [registrationEndpointURL: 0] + ) + } + + // MARK: 24 - Registration Falls Back to Dynamic Registration (CIMD Not Advertised) + + func configureOAuthRegistrationFallsBackToDynamicRegistrationCIMDNotAdvertised() + -> OAuthScenarioContext + { + let testEndpoint = URL(string: "https://localhost:8080/public/mcp")! + let resourceMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/public/mcp")! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + let registrationEndpointURL = URL( + string: "https://localhost:8080/register")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + let clientMetadataDocumentID = "https://client.example.com/metadata.json" + let finalResponseData = MockResponses.jsonRPCResult(id: 9) + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { request in + if request.value(forHTTPHeaderField: "Authorization") == nil { + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL, scope: "files:read" + )(request) + } + #expect( + request.value(forHTTPHeaderField: "Authorization") + == "Bearer dynamic-registration-token") + return try await MockResponses.jsonSuccess(body: finalResponseData)( + request) + }, + resourceMetadataURL: MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth"], + scopesSupported: ["files:read"] + ), + asMetadataURL: MockResponses.asMetadata( + issuer: "https://localhost:8080/auth", + tokenEndpoint: "https://localhost:8080/oauth/token", + registrationEndpoint: "https://localhost:8080/register", + clientIDMetadataDocumentSupported: false + ), + registrationEndpointURL: { request in + let body = + String(data: readRequestBody(request) ?? Data(), encoding: .utf8) ?? "" + #expect(body.contains("\"grant_types\":[\"client_credentials\"]")) + #expect(body.contains("\"token_endpoint_auth_method\":\"none\"")) + #expect(!body.contains("\"response_types\"")) + #expect(!body.contains("redirect_uris")) + #expect(!body.contains("authorization_code")) + return try await MockResponses.registrationSuccess( + clientID: "registered-client")(request) + }, + tokenEndpointURL: { request in + let body = + String(data: readRequestBody(request) ?? Data(), encoding: .utf8) ?? "" + #expect(body.contains("grant_type=client_credentials")) + #expect(body.contains("client_id=registered-client")) + #expect( + !body.contains( + "client_id=https%3A%2F%2Fclient.example.com%2Fmetadata.json")) + return try await MockResponses.tokenSuccess( + accessToken: "dynamic-registration-token")(request) + }, + ]) + + return OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init( + authentication: .none(clientID: clientMetadataDocumentID)), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":9}"#.data(using: .utf8)!, + expectedResponseData: finalResponseData, + expectedCallCounts: [registrationEndpointURL: 1] + ) + } + + // MARK: 25 - Registration Falls Back to Dynamic Registration (CIMD Capability Missing) + + func configureOAuthRegistrationFallsBackToDynamicRegistrationCIMDCapabilityMissing() + -> OAuthScenarioContext + { + let testEndpoint = URL(string: "https://localhost:8080/public/mcp")! + let resourceMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/public/mcp")! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + let registrationEndpointURL = URL( + string: "https://localhost:8080/register")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + let clientMetadataDocumentID = "https://client.example.com/metadata.json" + let finalResponseData = MockResponses.jsonRPCResult(id: 14) + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { request in + if request.value(forHTTPHeaderField: "Authorization") == nil { + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL, scope: "files:read" + )(request) + } + #expect( + request.value(forHTTPHeaderField: "Authorization") + == "Bearer dynamic-registration-token-missing-capability") + return try await MockResponses.jsonSuccess(body: finalResponseData)( + request) + }, + resourceMetadataURL: MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth"], + scopesSupported: ["files:read"] + ), + asMetadataURL: MockResponses.asMetadata( + issuer: "https://localhost:8080/auth", + tokenEndpoint: "https://localhost:8080/oauth/token", + registrationEndpoint: "https://localhost:8080/register" + ), + registrationEndpointURL: { request in + let body = + String(data: readRequestBody(request) ?? Data(), encoding: .utf8) ?? "" + #expect(body.contains("\"grant_types\":[\"client_credentials\"]")) + #expect(!body.contains("redirect_uris")) + return try await MockResponses.registrationSuccess( + clientID: "registered-client-2")(request) + }, + tokenEndpointURL: { request in + let body = + String(data: readRequestBody(request) ?? Data(), encoding: .utf8) ?? "" + #expect(body.contains("grant_type=client_credentials")) + #expect(body.contains("client_id=registered-client-2")) + return try await MockResponses.tokenSuccess( + accessToken: "dynamic-registration-token-missing-capability")( + request) + }, + ]) + + return OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init( + authentication: .none(clientID: clientMetadataDocumentID)), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":14}"#.data( + using: .utf8)!, + expectedResponseData: finalResponseData, + expectedCallCounts: [registrationEndpointURL: 1] + ) + } + + // MARK: 26 - Registration Missing Mechanism Returns Actionable Error + + func configureOAuthRegistrationMissingMechanismReturnsActionableError() + -> OAuthScenarioContext + { + let testEndpoint = URL(string: "https://localhost:8080/public/mcp")! + let resourceMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/public/mcp")! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + let clientMetadataDocumentID = "https://client.example.com/metadata.json" + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { request in + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL, scope: "files:read" + )(request) + }, + resourceMetadataURL: MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth"], + scopesSupported: ["files:read"] + ), + asMetadataURL: MockResponses.asMetadata( + issuer: "https://localhost:8080/auth", + tokenEndpoint: "https://localhost:8080/oauth/token", + clientIDMetadataDocumentSupported: false + ), + tokenEndpointURL: MockResponses.httpError(statusCode: 500), + ]) + + return OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init( + authentication: .none(clientID: clientMetadataDocumentID)), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":10}"#.data( + using: .utf8)!, + expectedErrorSubstring: + "Authorization server does not support Client ID Metadata Documents" + ) + } + + // MARK: 27 - CIMD Rejects Non-HTTPS Client ID URL + + func configureOAuthCIMDRejectsNonHTTPSClientIDURL() -> OAuthScenarioContext { + let testEndpoint = URL(string: "https://localhost:8080/public/mcp")! + let resourceMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/public/mcp")! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + let invalidClientID = "http://client.example.com/metadata.json" + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { request in + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL, scope: "files:read" + )(request) + }, + resourceMetadataURL: MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth"], + scopesSupported: ["files:read"] + ), + asMetadataURL: MockResponses.asMetadata( + issuer: "https://localhost:8080/auth", + tokenEndpoint: "https://localhost:8080/oauth/token", + clientIDMetadataDocumentSupported: true + ), + ]) + + return OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init( + authentication: .none(clientID: invalidClientID)), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":11}"#.data( + using: .utf8)!, + expectedErrorSubstring: + "Client ID metadata document URL must use https and include a path" + ) + } + + // MARK: 28 - Rejects Insecure MCP Endpoint URL + + func configureOAuthRejectsInsecureMCPEndpointURL() -> OAuthScenarioContext { + let insecureEndpoint = URL(string: "http://example.com/public/mcp")! + + return OAuthScenarioContext( + testEndpoint: insecureEndpoint, + oauthConfiguration: .init(authentication: .none(clientID: "test-client")), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":12}"#.data( + using: .utf8)!, + expectedErrorSubstring: + "MCP endpoint must use https or loopback http" + ) + } + + // MARK: 30 - PRM Cache Invalidated When resource_metadata URL Changes + + func configureOAuthPRMCacheInvalidatedOnResourceMetadataURLChange() + -> OAuthScenarioContext + { + let testEndpoint = URL(string: "https://localhost:8080/cache-invalidation")! + let resourceMetadataURL_A = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/cache-a")! + let resourceMetadataURL_B = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/cache-b")! + let asMetadataURL_A = URL( + string: + "https://localhost:8080/.well-known/oauth-authorization-server/auth-a")! + let asMetadataURL_B = URL( + string: + "https://localhost:8080/.well-known/oauth-authorization-server/auth-b")! + let tokenEndpointURL_A = URL( + string: "https://localhost:8080/oauth/token-a")! + let tokenEndpointURL_B = URL( + string: "https://localhost:8080/oauth/token-b")! + let firstResponseData = MockResponses.jsonRPCResult(id: 71) + let secondResponseData = MockResponses.jsonRPCResult(id: 72) + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { request in + let authHeader = request.value(forHTTPHeaderField: "Authorization") + let bodyStr = + String(data: readRequestBody(request) ?? Data(), encoding: .utf8) ?? "" + if authHeader == nil { + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL_A, scope: "files:read" + )(request) + } else if authHeader == "Bearer token-a" { + if bodyStr.contains("\"id\":71") { + return try await MockResponses.jsonSuccess( + body: firstResponseData)(request) + } else { + // Second message with old token → return 401 with new URL_B + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL_B, scope: "files:read" + )(request) + } + } else if authHeader == "Bearer token-b" { + return try await MockResponses.jsonSuccess( + body: secondResponseData)(request) + } + throw MockResponses.mockError( + "Unexpected Authorization header: \(String(describing: authHeader))") + }, + resourceMetadataURL_A: MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth-a"], + scopesSupported: ["files:read"] + ), + resourceMetadataURL_B: MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth-b"], + scopesSupported: ["files:read"] + ), + asMetadataURL_A: MockResponses.asMetadata( + issuer: "https://localhost:8080/auth-a", + tokenEndpoint: "https://localhost:8080/oauth/token-a" + ), + asMetadataURL_B: MockResponses.asMetadata( + issuer: "https://localhost:8080/auth-b", + tokenEndpoint: "https://localhost:8080/oauth/token-b" + ), + tokenEndpointURL_A: MockResponses.tokenSuccess(accessToken: "token-a"), + tokenEndpointURL_B: MockResponses.tokenSuccess(accessToken: "token-b"), + ]) + + return OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init(authentication: .none(clientID: "test-client")), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":71}"#.data(using: .utf8)!, + expectedResponseData: firstResponseData, + expectedCallCounts: [ + resourceMetadataURL_A: 1, + resourceMetadataURL_B: 1, + tokenEndpointURL_A: 1, + tokenEndpointURL_B: 1, + ], + secondMessageData: #"{"jsonrpc":"2.0","method":"ping","id":72}"#.data( + using: .utf8)!, + secondExpectedResponseData: secondResponseData + ) + } + + // MARK: 29 - Rejects Non-Loopback HTTP Redirect URI + + func configureOAuthRejectsNonLoopbackHTTPRedirectURI() -> OAuthScenarioContext { + let testEndpoint = URL(string: "https://localhost:8080/public/mcp")! + let resourceMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/public/mcp")! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { request in + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL, scope: "files:read" + )(request) + }, + resourceMetadataURL: MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth"], + scopesSupported: ["files:read"] + ), + asMetadataURL: MockResponses.asMetadata( + issuer: "https://localhost:8080/auth", + tokenEndpoint: "https://localhost:8080/oauth/token", + authorizationEndpoint: "https://localhost:8080/oauth/authorize", + codeChallengeMethodsSupported: ["S256"] + ), + ]) + + return OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init( + grantType: .authorizationCode, + authentication: .none(clientID: "test-client"), + authorizationRedirectURI: URL( + string: "http://evil.example.com/callback")! + ), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":13}"#.data( + using: .utf8)!, + expectedErrorSubstring: + "Redirect URI must use https or loopback http and must not include fragments" + ) + } + + // MARK: 31 - Resource Uses PRM Resource Field (Gap 1) + + func configureOAuthResourceUsesPRMResourceField() -> OAuthScenarioContext { + let testEndpoint = URL(string: "https://localhost:8080/mcp/tools")! + let resourceMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/mcp/tools")! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + let finalResponseData = MockResponses.jsonRPCResult(id: 91) + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { request in + if request.value(forHTTPHeaderField: "Authorization") == nil { + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL)(request) + } + return try await MockResponses.jsonSuccess(body: finalResponseData)(request) + }, + // PRM resource = "https://localhost:8080" (origin only, no path) + resourceMetadataURL: MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth"], + resource: "https://localhost:8080" + ), + asMetadataURL: MockResponses.asMetadata( + issuer: "https://localhost:8080/auth", + tokenEndpoint: "https://localhost:8080/oauth/token" + ), + tokenEndpointURL: { request in + let body = + String(data: readRequestBody(request) ?? Data(), encoding: .utf8) ?? "" + // Must use the PRM resource field, not the specific endpoint path + #expect(body.contains("resource=https%3A%2F%2Flocalhost%3A8080")) + #expect(!body.contains("resource=https%3A%2F%2Flocalhost%3A8080%2Fmcp")) + return try await MockResponses.tokenSuccess( + accessToken: "access-token-prm-resource")(request) + }, + ]) + + return OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init(authentication: .none(clientID: "test-client")), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":91}"#.data(using: .utf8)!, + expectedResponseData: finalResponseData, + expectedCallCounts: [ + testEndpoint: 2, resourceMetadataURL: 1, asMetadataURL: 1, + tokenEndpointURL: 1, + ] + ) + } + + // MARK: 32 - Second Authorization Server Tried When First Fails (Gap 2) + + func configureOAuthSecondAuthorizationServerTriedWhenFirstFails() -> OAuthScenarioContext + { + let testEndpoint = URL(string: "https://localhost:8080/test-as-fallback")! + let resourceMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/test-as-fallback" + )! + // AS1 discovery URLs — all return 404 + let as1MetadataURL1 = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/as1")! + let as1MetadataURL2 = URL( + string: "https://localhost:8080/.well-known/openid-configuration/as1")! + let as1MetadataURL3 = URL( + string: "https://localhost:8080/as1/.well-known/openid-configuration")! + // AS2 first discovery URL — returns valid metadata + let as2MetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/as2")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token-as2")! + let finalResponseData = MockResponses.jsonRPCResult(id: 92) + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { request in + if request.value(forHTTPHeaderField: "Authorization") == nil { + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL)(request) + } + return try await MockResponses.jsonSuccess(body: finalResponseData)(request) + }, + resourceMetadataURL: MockResponses.resourceMetadata( + authorizationServers: [ + "https://localhost:8080/as1", + "https://localhost:8080/as2", + ] + ), + // AS1: all well-known URLs return 404 + as1MetadataURL1: MockResponses.httpError(statusCode: 404), + as1MetadataURL2: MockResponses.httpError(statusCode: 404), + as1MetadataURL3: MockResponses.httpError(statusCode: 404), + // AS2: first well-known URL returns valid metadata + as2MetadataURL: MockResponses.asMetadata( + issuer: "https://localhost:8080/as2", + tokenEndpoint: "https://localhost:8080/oauth/token-as2" + ), + tokenEndpointURL: MockResponses.tokenSuccess( + accessToken: "access-token-from-as2"), + ]) + + return OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init(authentication: .none(clientID: "test-client")), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":92}"#.data(using: .utf8)!, + expectedResponseData: finalResponseData, + expectedCallCounts: [ + testEndpoint: 2, + resourceMetadataURL: 1, + as1MetadataURL1: 1, + as1MetadataURL2: 1, + as1MetadataURL3: 1, + as2MetadataURL: 1, + tokenEndpointURL: 1, + ] + ) + } + + // MARK: 33 - Re-Registration After Client Secret Expiry (Gap 3) + + func configureOAuthReRegistersAfterClientSecretExpiry() -> OAuthScenarioContext { + let testEndpoint = URL(string: "https://localhost:8080/test-secret-expiry")! + let resourceMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/test-secret-expiry" + )! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + let registrationEndpointURL = URL(string: "https://localhost:8080/register")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + let firstResponseData = MockResponses.jsonRPCResult(id: 93) + let secondResponseData = MockResponses.jsonRPCResult(id: 94) + let regTracker = OrderTracker() + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { request in + let auth = request.value(forHTTPHeaderField: "Authorization") + let body = + String(data: readRequestBody(request) ?? Data(), encoding: .utf8) ?? "" + if auth == nil { + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL)(request) + } + if auth == "Bearer token-v1" && body.contains("\"id\":93") { + return try await MockResponses.jsonSuccess(body: firstResponseData)( + request) + } + // Second request with old token → trigger re-auth + if auth == "Bearer token-v1" && body.contains("\"id\":94") { + return try await MockResponses.bearerChallenge()(request) + } + if auth == "Bearer token-v2" { + return try await MockResponses.jsonSuccess(body: secondResponseData)( + request) + } + throw MockResponses.mockError( + "Unexpected auth: \(String(describing: auth))") + }, + resourceMetadataURL: MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth"] + ), + asMetadataURL: MockResponses.asMetadata( + issuer: "https://localhost:8080/auth", + tokenEndpoint: "https://localhost:8080/oauth/token", + registrationEndpoint: "https://localhost:8080/register" + ), + registrationEndpointURL: { [regTracker] request in + await regTracker.append(request.url!) + let count = await regTracker.count() + // Return different client IDs per registration call + let clientID = count == 1 ? "client-v1" : "client-v2" + // client_secret_expires_at = 1 (Unix epoch — always in the past) + let dict: [String: Any] = [ + "client_id": clientID, + "client_secret_expires_at": 1, + ] + let data = try JSONSerialization.data(withJSONObject: dict) + let response = HTTPURLResponse( + url: request.url!, statusCode: 201, httpVersion: "HTTP/1.1", + headerFields: ["Content-Type": "application/json"])! + return (response, data) + }, + tokenEndpointURL: { request in + let body = + String(data: readRequestBody(request) ?? Data(), encoding: .utf8) ?? "" + if body.contains("client_id=client-v1") { + return try await MockResponses.tokenSuccess(accessToken: "token-v1")( + request) + } + if body.contains("client_id=client-v2") { + return try await MockResponses.tokenSuccess(accessToken: "token-v2")( + request) + } + throw MockResponses.mockError("Unexpected client_id in token request") + }, + ]) + + return OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init(authentication: .none(clientID: "anon")), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":93}"#.data(using: .utf8)!, + expectedResponseData: firstResponseData, + expectedCallCounts: [ + registrationEndpointURL: 2, + tokenEndpointURL: 2, + ], + secondMessageData: #"{"jsonrpc":"2.0","method":"ping","id":94}"#.data( + using: .utf8)!, + secondExpectedResponseData: secondResponseData + ) + } + + // MARK: 34 - Issuer Mismatch Causes Next Discovery URL Variant to Be Tried (Gap 4) + + func configureOAuthIssuerMismatchTriesNextURLVariant() -> OAuthScenarioContext { + let testEndpoint = URL(string: "https://localhost:8080/test-issuer-check")! + let resourceMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/test-issuer-check" + )! + // Discovery URL 1 returns wrong issuer → skipped + let wrongIssuerURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + // Discovery URL 2 returns correct issuer → used + let correctIssuerURL = URL( + string: "https://localhost:8080/.well-known/openid-configuration/auth")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + let finalResponseData = MockResponses.jsonRPCResult(id: 95) + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { request in + if request.value(forHTTPHeaderField: "Authorization") == nil { + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL)(request) + } + return try await MockResponses.jsonSuccess(body: finalResponseData)(request) + }, + resourceMetadataURL: MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth"] + ), + // URL 1: returns metadata with wrong issuer field → SDK skips it + wrongIssuerURL: MockResponses.asMetadata( + issuer: "https://evil.example.com", + tokenEndpoint: "https://evil.example.com/token" + ), + // URL 2: returns metadata with correct issuer → SDK accepts it + correctIssuerURL: MockResponses.asMetadata( + issuer: "https://localhost:8080/auth", + tokenEndpoint: "https://localhost:8080/oauth/token" + ), + tokenEndpointURL: MockResponses.tokenSuccess( + accessToken: "access-token-after-issuer-check"), + ]) + + return OAuthScenarioContext( + testEndpoint: testEndpoint, + oauthConfiguration: .init(authentication: .none(clientID: "test-client")), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":95}"#.data(using: .utf8)!, + expectedResponseData: finalResponseData, + expectedCallCounts: [ + testEndpoint: 2, + resourceMetadataURL: 1, + wrongIssuerURL: 1, + correctIssuerURL: 1, + tokenEndpointURL: 1, + ] + ) + } + + // MARK: 35 - Proactive Token Refresh Within Window (Gap 5) + + func configureOAuthProactiveTokenRefreshWithinWindow() -> OAuthScenarioContext { + let testEndpoint = URL(string: "https://localhost:8080/test-proactive-refresh")! + let resourceMetadataURL = URL( + string: + "https://localhost:8080/.well-known/oauth-protected-resource/test-proactive-refresh" + )! + let asMetadataURL = URL( + string: "https://localhost:8080/.well-known/oauth-authorization-server/auth")! + let tokenEndpointURL = URL(string: "https://localhost:8080/oauth/token")! + let firstResponseData = MockResponses.jsonRPCResult(id: 96) + let secondResponseData = MockResponses.jsonRPCResult(id: 97) + + requestHandler = MockResponses.routingHandler(routes: [ + testEndpoint: { request in + let auth = request.value(forHTTPHeaderField: "Authorization") + let body = + String(data: readRequestBody(request) ?? Data(), encoding: .utf8) ?? "" + if auth == nil { + return try await MockResponses.bearerChallenge( + resourceMetadataURL: resourceMetadataURL)(request) + } + if body.contains("\"id\":96") { + return try await MockResponses.jsonSuccess(body: firstResponseData)( + request) + } + // Second request: must be sent with the proactively refreshed token + if auth == "Bearer refreshed-token" { + return try await MockResponses.jsonSuccess(body: secondResponseData)( + request) + } + throw MockResponses.mockError( + "Expected refreshed token on second request, got: \(String(describing: auth))" + ) + }, + resourceMetadataURL: MockResponses.resourceMetadata( + authorizationServers: ["https://localhost:8080/auth"] + ), + asMetadataURL: MockResponses.asMetadata( + issuer: "https://localhost:8080/auth", + tokenEndpoint: "https://localhost:8080/oauth/token" + ), + tokenEndpointURL: { request in + let body = + String(data: readRequestBody(request) ?? Data(), encoding: .utf8) ?? "" + if body.contains("grant_type=client_credentials") { + // Initial token: expires in 300s (within 400s proactive window) + return try await MockResponses.tokenSuccess( + accessToken: "initial-token", + expiresIn: 300, + refreshToken: "rt-1" + )(request) + } + if body.contains("grant_type=refresh_token") { + #expect(body.contains("refresh_token=rt-1")) + return try await MockResponses.tokenSuccess( + accessToken: "refreshed-token", + expiresIn: 3600 + )(request) + } + throw MockResponses.mockError("Unexpected grant type in token request") + }, + ]) + + return OAuthScenarioContext( + testEndpoint: testEndpoint, + // proactiveRefreshWindowSeconds = 400 > expires_in = 300 + oauthConfiguration: .init( + authentication: .none(clientID: "test-client"), + proactiveRefreshWindowSeconds: 400 + ), + messageData: #"{"jsonrpc":"2.0","method":"ping","id":96}"#.data(using: .utf8)!, + expectedResponseData: firstResponseData, + expectedCallCounts: [ + testEndpoint: 3, // initial 401 + first send retry + second send + tokenEndpointURL: 2, // initial client_credentials + refresh_token + ], + secondMessageData: #"{"jsonrpc":"2.0","method":"ping","id":97}"#.data( + using: .utf8)!, + secondExpectedResponseData: secondResponseData + ) + } + + } + + #endif // !canImport(FoundationNetworking) + +#endif // swift(>=6.1) diff --git a/Tests/MCPTests/OAuthTokenEndpointClientTests.swift b/Tests/MCPTests/OAuthTokenEndpointClientTests.swift new file mode 100644 index 0000000..be9704c --- /dev/null +++ b/Tests/MCPTests/OAuthTokenEndpointClientTests.swift @@ -0,0 +1,194 @@ +@preconcurrency import Foundation +import Testing + +@testable import MCP + +#if canImport(FoundationNetworking) + import FoundationNetworking +#endif + +#if swift(>=6.1) && !os(Linux) + + @Suite("OAuthTokenEndpointClient", .serialized) + struct OAuthTokenEndpointClientTests { + + let client = OAuthTokenEndpointClient(urlValidator: OAuthURLValidator()) + let tokenEndpoint = URL(string: "https://auth.example.com/token")! + + func successBody( + accessToken: String = "access-token", + tokenType: String = "Bearer", + scope: String? = nil, + refreshToken: String? = nil + ) throws -> Data { + var dict: [String: Any] = [ + "access_token": accessToken, + "token_type": tokenType, + "expires_in": 3600, + ] + if let scope { dict["scope"] = scope } + if let refreshToken { dict["refresh_token"] = refreshToken } + return try JSONSerialization.data(withJSONObject: dict) + } + + // MARK: - Success + + @Test("Returns decoded token response on success") + func testRequestReturnsToken() async throws { + let body = try successBody() + let (session, key) = makeIsolatedSession() + await IsolatedMockURLProtocol.setHandler(key: key) { _ in + let response = HTTPURLResponse( + url: self.tokenEndpoint, statusCode: 200, + httpVersion: nil, headerFields: nil)! + return (response, body) + } + + var params = ["grant_type": "client_credentials"] + let result = try await client.request( + parameters: ¶ms, + endpoint: tokenEndpoint, + authentication: .none(clientID: "client-id"), + session: session + ) + let expected = OAuthTokenResponse( + accessToken: "access-token", tokenType: "Bearer", + expiresIn: 3600, scope: nil, refreshToken: nil) + #expect(result == expected) + } + + @Test("Parses optional scope and refresh token") + func testRequestParsesOptionalFields() async throws { + let body = try successBody(scope: "read write", refreshToken: "refresh-xyz") + let (session, key) = makeIsolatedSession() + await IsolatedMockURLProtocol.setHandler(key: key) { _ in + let response = HTTPURLResponse( + url: self.tokenEndpoint, statusCode: 200, + httpVersion: nil, headerFields: nil)! + return (response, body) + } + + var params = ["grant_type": "client_credentials"] + let result = try await client.request( + parameters: ¶ms, + endpoint: tokenEndpoint, + authentication: .none(clientID: "client-id"), + session: session + ) + let expected = OAuthTokenResponse( + accessToken: "access-token", tokenType: "Bearer", + expiresIn: 3600, scope: "read write", refreshToken: "refresh-xyz") + #expect(result == expected) + } + + // MARK: - Error Responses + + @Test("Throws for non-2xx status") + func testRequestThrowsForNon2xx() async throws { + let errorBody = try JSONSerialization.data(withJSONObject: ["error": "invalid_client"]) + let (session, key) = makeIsolatedSession() + await IsolatedMockURLProtocol.setHandler(key: key) { _ in + let response = HTTPURLResponse( + url: self.tokenEndpoint, statusCode: 401, + httpVersion: nil, headerFields: nil)! + return (response, errorBody) + } + + var params = ["grant_type": "client_credentials"] + await #expect(throws: OAuthAuthorizationError.self) { + try await client.request( + parameters: ¶ms, + endpoint: tokenEndpoint, + authentication: .none(clientID: "client-id"), + session: session + ) + } + } + + @Test("Throws for empty access_token") + func testRequestThrowsForEmptyAccessToken() async throws { + let body = try successBody(accessToken: "") + let (session, key) = makeIsolatedSession() + await IsolatedMockURLProtocol.setHandler(key: key) { _ in + let response = HTTPURLResponse( + url: self.tokenEndpoint, statusCode: 200, + httpVersion: nil, headerFields: nil)! + return (response, body) + } + + var params = ["grant_type": "client_credentials"] + await #expect(throws: OAuthAuthorizationError.self) { + try await client.request( + parameters: ¶ms, + endpoint: tokenEndpoint, + authentication: .none(clientID: "client-id"), + session: session + ) + } + } + + @Test("Throws for non-Bearer token_type") + func testRequestThrowsForNonBearerTokenType() async throws { + let body = try successBody(tokenType: "MAC") + let (session, key) = makeIsolatedSession() + await IsolatedMockURLProtocol.setHandler(key: key) { _ in + let response = HTTPURLResponse( + url: self.tokenEndpoint, statusCode: 200, + httpVersion: nil, headerFields: nil)! + return (response, body) + } + + var params = ["grant_type": "client_credentials"] + await #expect(throws: OAuthAuthorizationError.self) { + try await client.request( + parameters: ¶ms, + endpoint: tokenEndpoint, + authentication: .none(clientID: "client-id"), + session: session + ) + } + } + + // MARK: - Form Encoding + + @Test("Sends parameters as form-encoded POST body") + func testRequestSendsFormEncodedBody() async throws { + let body = try successBody() + let (session, key) = makeIsolatedSession() + actor RequestCapture { var value: URLRequest?; func set(_ r: URLRequest) { value = r } } + let capture = RequestCapture() + await IsolatedMockURLProtocol.setHandler(key: key) { request in + await capture.set(request) + let response = HTTPURLResponse( + url: self.tokenEndpoint, statusCode: 200, + httpVersion: nil, headerFields: nil)! + return (response, body) + } + + var params = ["grant_type": "client_credentials", "resource": "https://api.example.com"] + _ = try await client.request( + parameters: ¶ms, + endpoint: tokenEndpoint, + authentication: .none(clientID: "client-id"), + session: session + ) + + let capturedRequest = await capture.value + let bodyData: Data = { + if let data = capturedRequest?.httpBody { return data } + guard let stream = capturedRequest?.httpBodyStream else { return Data() } + stream.open(); defer { stream.close() } + var data = Data() + let buf = UnsafeMutablePointer.allocate(capacity: 4096) + defer { buf.deallocate() } + while stream.hasBytesAvailable { data.append(buf, count: stream.read(buf, maxLength: 4096)) } + return data + }() + let bodyString = String(data: bodyData, encoding: .utf8) ?? "" + #expect(bodyString.contains("grant_type=client_credentials")) + #expect(bodyString.contains("resource=")) + #expect(capturedRequest?.httpMethod == "POST") + } + } + +#endif diff --git a/Tests/MCPTests/OAuthURLValidatorTests.swift b/Tests/MCPTests/OAuthURLValidatorTests.swift new file mode 100644 index 0000000..0071c04 --- /dev/null +++ b/Tests/MCPTests/OAuthURLValidatorTests.swift @@ -0,0 +1,141 @@ +import Foundation +import Testing + +@testable import MCP + +#if canImport(FoundationNetworking) + import FoundationNetworking +#endif + +@Suite("OAuthURLValidator") +struct OAuthURLValidatorTests { + + // MARK: - validateHTTPSOrLoopback + + @Test("Accepts HTTPS URL") + func testValidateHTTPSOrLoopbackAcceptsHTTPS() throws { + try OAuthURLValidator().validateHTTPSOrLoopback( + URL(string: "https://example.com/mcp")!, context: "test") + } + + @Test("Accepts loopback HTTP") + func testValidateHTTPSOrLoopbackAcceptsLoopback() throws { + try OAuthURLValidator().validateHTTPSOrLoopback( + URL(string: "http://localhost:8080/mcp")!, context: "test") + try OAuthURLValidator().validateHTTPSOrLoopback( + URL(string: "http://127.0.0.1:9000/mcp")!, context: "test") + } + + @Test("Rejects remote HTTP") + func testValidateHTTPSOrLoopbackRejectsRemoteHTTP() { + #expect(throws: OAuthAuthorizationError.self) { + try OAuthURLValidator().validateHTTPSOrLoopback( + URL(string: "http://example.com/mcp")!, context: "test") + } + } + + @Test("Rejects URL with fragment") + func testValidateHTTPSOrLoopbackRejectsFragment() { + #expect(throws: OAuthAuthorizationError.self) { + try OAuthURLValidator().validateHTTPSOrLoopback( + URL(string: "https://example.com/mcp#frag")!, context: "test") + } + } + + // MARK: - validateAuthorizationServer + + @Test("Accepts HTTPS authorization server") + func testValidateAuthorizationServerAcceptsHTTPS() throws { + try OAuthURLValidator().validateAuthorizationServer( + URL(string: "https://auth.example.com")!, context: "test") + } + + @Test("Rejects HTTP authorization server by default") + func testValidateAuthorizationServerRejectsHTTP() { + #expect(throws: OAuthAuthorizationError.self) { + try OAuthURLValidator().validateAuthorizationServer( + URL(string: "http://auth.example.com")!, context: "test") + } + } + + @Test("Accepts loopback HTTP when flag is set") + func testValidateAuthorizationServerAcceptsLoopbackWhenAllowed() throws { + let v = OAuthURLValidator(allowLoopbackHTTPForAuthorizationServer: true) + try v.validateAuthorizationServer( + URL(string: "http://localhost:8080")!, context: "test") + } + + @Test("Rejects loopback HTTP when flag is not set") + func testValidateAuthorizationServerRejectsLoopbackWhenNotAllowed() { + #expect(throws: OAuthAuthorizationError.self) { + try OAuthURLValidator().validateAuthorizationServer( + URL(string: "http://localhost:8080")!, context: "test") + } + } + + // MARK: - validateRedirectURI + + @Test("Accepts HTTPS redirect URI") + func testValidateRedirectURIAcceptsHTTPS() throws { + try OAuthURLValidator().validateRedirectURI( + URL(string: "https://app.example.com/callback")!) + } + + @Test("Accepts loopback HTTP redirect URI") + func testValidateRedirectURIAcceptsLoopback() throws { + try OAuthURLValidator().validateRedirectURI( + URL(string: "http://localhost:8080/callback")!) + } + + @Test("Rejects remote HTTP redirect URI") + func testValidateRedirectURIRejectsRemoteHTTP() { + #expect(throws: OAuthAuthorizationError.self) { + try OAuthURLValidator().validateRedirectURI( + URL(string: "http://app.example.com/callback")!) + } + } + + @Test("Rejects redirect URI with fragment") + func testValidateRedirectURIRejectsFragment() { + #expect(throws: OAuthAuthorizationError.self) { + try OAuthURLValidator().validateRedirectURI( + URL(string: "https://app.example.com/callback#section")!) + } + } + + // MARK: - isPrivateIPHost + + @Test("Identifies private IPv4 ranges") + func testIsPrivateIPHostIPv4() { + let v = OAuthURLValidator() + #expect(v.isPrivateIPHost("10.0.0.1")) + #expect(v.isPrivateIPHost("172.16.0.1")) + #expect(v.isPrivateIPHost("172.31.255.255")) + #expect(v.isPrivateIPHost("192.168.1.1")) + #expect(v.isPrivateIPHost("169.254.169.254")) + #expect(v.isPrivateIPHost("100.64.0.1")) + } + + @Test("Does not block public IPv4 addresses") + func testIsPrivateIPHostPublicIPv4() { + let v = OAuthURLValidator() + #expect(!v.isPrivateIPHost("1.2.3.4")) + #expect(!v.isPrivateIPHost("8.8.8.8")) + #expect(!v.isPrivateIPHost("203.0.113.1")) + } + + @Test("Identifies private IPv6 ULA and link-local addresses") + func testIsPrivateIPHostIPv6() { + let v = OAuthURLValidator() + #expect(v.isPrivateIPHost("fc00::1")) + #expect(v.isPrivateIPHost("fd12:3456::1")) + #expect(v.isPrivateIPHost("fe80::1")) + } + + @Test("Does not block public hostnames") + func testIsPrivateIPHostPublicHostname() { + let v = OAuthURLValidator() + #expect(!v.isPrivateIPHost("example.com")) + #expect(!v.isPrivateIPHost("localhost")) + } +} diff --git a/Tests/MCPTests/ProgressTests.swift b/Tests/MCPTests/ProgressTests.swift new file mode 100644 index 0000000..fdbb79a --- /dev/null +++ b/Tests/MCPTests/ProgressTests.swift @@ -0,0 +1,390 @@ +import Foundation +import Testing + +@testable import MCP + +@Suite("Progress Tests") +struct ProgressTests { + // MARK: - ProgressToken Tests + + @Test("ProgressToken string encoding and decoding") + func testProgressTokenStringEncodingDecoding() throws { + let token = ProgressToken.string("test-token-123") + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(token) + let decoded = try decoder.decode(ProgressToken.self, from: data) + + #expect(decoded == token) + if case .string(let value) = decoded { + #expect(value == "test-token-123") + } else { + #expect(Bool(false), "Expected string token") + } + } + + @Test("ProgressToken integer encoding and decoding") + func testProgressTokenIntegerEncodingDecoding() throws { + let token = ProgressToken.integer(42) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(token) + let decoded = try decoder.decode(ProgressToken.self, from: data) + + #expect(decoded == token) + if case .integer(let value) = decoded { + #expect(value == 42) + } else { + #expect(Bool(false), "Expected integer token") + } + } + + @Test("ProgressToken unique generation") + func testProgressTokenUnique() { + let token1 = ProgressToken.unique() + let token2 = ProgressToken.unique() + + #expect(token1 != token2) + + if case .string(let value1) = token1, case .string(let value2) = token2 { + #expect(value1 != value2) + } else { + #expect(Bool(false), "Expected string tokens") + } + } + + @Test("ProgressToken hashable") + func testProgressTokenHashable() { + let token1 = ProgressToken.string("test") + let token2 = ProgressToken.string("test") + let token3 = ProgressToken.integer(1) + let token4 = ProgressToken.integer(1) + + #expect(token1 == token2) + #expect(token3 == token4) + #expect(token1 != token3) + + let set: Set = [token1, token2, token3, token4] + #expect(set.count == 2) + } + + // MARK: - RequestMeta Tests + + @Test("RequestMeta empty initialization") + func testRequestMetaEmptyInit() throws { + let meta = Metadata() + + #expect(meta.progressToken == nil) + #expect(meta.fields.isEmpty) + } + + @Test("RequestMeta with progress token") + func testRequestMetaWithProgressToken() throws { + let token = ProgressToken.string("my-token") + let meta = Metadata(progressToken: token) + + #expect(meta.progressToken == token) + #expect(meta.fields["progressToken"] == .string("my-token")) + } + + @Test("RequestMeta with integer progress token") + func testRequestMetaWithIntegerProgressToken() throws { + let token = ProgressToken.integer(42) + let meta = Metadata(progressToken: token) + + #expect(meta.progressToken == token) + #expect(meta.fields["progressToken"] == .int(42)) + } + + @Test("RequestMeta encoding with progress token") + func testRequestMetaEncodingWithProgressToken() throws { + let token = ProgressToken.string("my-token") + let meta = Metadata(progressToken: token) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(meta) + let decoded = try decoder.decode(Metadata.self, from: data) + + #expect(decoded.progressToken == token) + } + + @Test("RequestMeta encoding without progress token") + func testRequestMetaEncodingWithoutProgressToken() throws { + let meta = Metadata() + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(meta) + let decoded = try decoder.decode(Metadata.self, from: data) + + #expect(decoded.progressToken == nil) + } + + @Test("RequestMeta JSON representation with progress token") + func testRequestMetaJSONWithProgressToken() throws { + let token = ProgressToken.string("test-token") + let meta = Metadata(progressToken: token) + + let encoder = JSONEncoder() + let data = try encoder.encode(meta) + let jsonString = String(data: data, encoding: .utf8)! + + #expect(jsonString.contains("progressToken")) + #expect(jsonString.contains("test-token")) + } + + @Test("RequestMeta with additional fields") + func testRequestMetaWithAdditionalFields() throws { + let meta = Metadata( + progressToken: .string("token"), + additionalFields: ["customField": .string("customValue")] + ) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(meta) + let decoded = try decoder.decode(Metadata.self, from: data) + + #expect(decoded.progressToken == .string("token")) + #expect(decoded.fields["customField"] == .string("customValue")) + } + + @Test("RequestMeta with only additional fields") + func testRequestMetaWithOnlyAdditionalFields() throws { + let meta = Metadata(additionalFields: [ + "customKey": .int(123), + "anotherKey": .string("value") + ]) + + #expect(meta.progressToken == nil) + #expect(meta.fields["customKey"] == .int(123)) + #expect(meta.fields["anotherKey"] == .string("value")) + } + + + // MARK: - ProgressNotification Tests + + @Test("ProgressNotification name") + func testProgressNotificationName() { + #expect(ProgressNotification.name == "notifications/progress") + } + + @Test("ProgressNotification parameters encoding and decoding") + func testProgressNotificationParametersEncodingDecoding() throws { + let params = ProgressNotification.Parameters( + progressToken: .string("test-token"), + progress: 50, + total: 100, + message: "Processing..." + ) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(params) + let decoded = try decoder.decode(ProgressNotification.Parameters.self, from: data) + + #expect(decoded.progressToken == .string("test-token")) + #expect(decoded.progress == 50) + #expect(decoded.total == 100) + #expect(decoded.message == "Processing...") + } + + @Test("ProgressNotification parameters without optional fields") + func testProgressNotificationParametersWithoutOptionals() throws { + let params = ProgressNotification.Parameters( + progressToken: .integer(42), + progress: 75 + ) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(params) + let decoded = try decoder.decode(ProgressNotification.Parameters.self, from: data) + + #expect(decoded.progressToken == .integer(42)) + #expect(decoded.progress == 75) + #expect(decoded.total == nil) + #expect(decoded.message == nil) + } + + @Test("ProgressNotification message creation") + func testProgressNotificationMessage() throws { + let params = ProgressNotification.Parameters( + progressToken: .string("my-token"), + progress: 30, + total: 100, + message: "Step 3 of 10" + ) + + let message = ProgressNotification.message(params) + + #expect(message.method == "notifications/progress") + #expect(message.params.progressToken == .string("my-token")) + #expect(message.params.progress == 30) + #expect(message.params.total == 100) + #expect(message.params.message == "Step 3 of 10") + } + + // MARK: - CallTool with _meta Tests + + @Test("CallTool parameters without meta") + func testCallToolParametersWithoutMeta() throws { + let params = CallTool.Parameters(name: "test_tool", arguments: ["key": "value"]) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(params) + let decoded = try decoder.decode(CallTool.Parameters.self, from: data) + + #expect(decoded.name == "test_tool") + #expect(decoded.arguments?["key"] == .string("value")) + #expect(decoded._meta == nil) + + // Verify _meta is not included in JSON when nil + let jsonString = String(data: data, encoding: .utf8)! + #expect(!jsonString.contains("_meta")) + } + + @Test("CallTool parameters with progress token") + func testCallToolParametersWithProgressToken() throws { + let token = ProgressToken.string("call-tool-token") + let meta = Metadata(progressToken: token) + let params = CallTool.Parameters(name: "test_tool", arguments: ["key": "value"], meta: meta) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(params) + let decoded = try decoder.decode(CallTool.Parameters.self, from: data) + + #expect(decoded.name == "test_tool") + #expect(decoded.arguments?["key"] == .string("value")) + #expect(decoded._meta?.progressToken == token) + + // Verify _meta is included in JSON + let jsonString = String(data: data, encoding: .utf8)! + #expect(jsonString.contains("_meta")) + #expect(jsonString.contains("progressToken")) + #expect(jsonString.contains("call-tool-token")) + } + + @Test("CallTool request encoding with progress token") + func testCallToolRequestEncodingWithProgressToken() throws { + let token = ProgressToken.string("request-token") + let meta = Metadata(progressToken: token) + let request = CallTool.request(.init(name: "my_tool", arguments: ["arg": 42], meta: meta)) + + let encoder = JSONEncoder() + let data = try encoder.encode(request) + let jsonString = String(data: data, encoding: .utf8)! + + // Note: JSON encoding may escape forward slashes as \/ + #expect(jsonString.contains("tools") && jsonString.contains("call")) + #expect(jsonString.contains("my_tool")) + #expect(jsonString.contains("_meta")) + #expect(jsonString.contains("progressToken")) + #expect(jsonString.contains("request-token")) + } + + @Test("CallTool request decoding with _meta") + func testCallToolRequestDecodingWithMeta() throws { + let jsonString = """ + { + "jsonrpc": "2.0", + "id": "test-id", + "method": "tools/call", + "params": { + "_meta": { + "progressToken": "decoded-token" + }, + "name": "decoded_tool", + "arguments": {"x": 10} + } + } + """ + let data = jsonString.data(using: .utf8)! + + let decoder = JSONDecoder() + let request = try decoder.decode(Request.self, from: data) + + #expect(request.id == "test-id") + #expect(request.method == "tools/call") + #expect(request.params.name == "decoded_tool") + #expect(request.params.arguments?["x"] == .int(10)) + #expect(request.params._meta?.progressToken == .string("decoded-token")) + } + + @Test("CallTool request decoding without _meta") + func testCallToolRequestDecodingWithoutMeta() throws { + let jsonString = """ + { + "jsonrpc": "2.0", + "id": "test-id", + "method": "tools/call", + "params": { + "name": "tool_without_meta", + "arguments": {"y": 20} + } + } + """ + let data = jsonString.data(using: .utf8)! + + let decoder = JSONDecoder() + let request = try decoder.decode(Request.self, from: data) + + #expect(request.params.name == "tool_without_meta") + #expect(request.params.arguments?["y"] == .int(20)) + #expect(request.params._meta == nil) + } + + @Test("Client notification updates should fire when server sends them") + func testClientNotificationUpdates() async throws { + let pair = await InMemoryTransport.createConnectedPair() + let client = Client(name: "testClient", version: "1") + let token = ProgressToken.unique() + + var progresses: [Double] = [] + await client.onNotification(ProgressNotification.self) { message in + let receivedToken = message.params.progressToken + #expect(receivedToken == token) + await MainActor.run { + progresses.append(message.params.progress) + } + } + + let server = Server(name: "testServer", version: "1") + let expectedToolCallResult = CallTool.Result(content: [.text(text: "success", annotations: nil, _meta: nil)]) + await server.withMethodHandler(CallTool.self) { params in + if let token = params._meta?.progressToken { + for i in 1...5 { + let notification = ProgressNotification.message( + .init(progressToken: token, progress: Double(i * 20)) + ) + try await server.notify(notification) + } + } + + return .init(content: [.text(text: "success", annotations: nil, _meta: nil)]) + } + + try await server.start(transport: pair.server) + try await client.connect(transport: pair.client) + let context: RequestContext = try await client.callTool(name: "random", meta: Metadata(progressToken: token)) + let result = try await context.value + + #expect(progresses == [20, 40, 60, 80, 100]) + #expect(result.content == expectedToolCallResult.content) + #expect(result.isError == nil) + } +} diff --git a/Tests/MCPTests/PromptTests.swift b/Tests/MCPTests/PromptTests.swift index 083d463..3aad91f 100644 --- a/Tests/MCPTests/PromptTests.swift +++ b/Tests/MCPTests/PromptTests.swift @@ -9,30 +9,31 @@ struct PromptTests { func testPromptInitialization() throws { let argument = Prompt.Argument( name: "test_arg", + title: "Test Argument Title", description: "A test argument", required: true ) let prompt = Prompt( name: "test_prompt", + title: "Test Prompt Title", description: "A test prompt", arguments: [argument] ) #expect(prompt.name == "test_prompt") + #expect(prompt.title == "Test Prompt Title") #expect(prompt.description == "A test prompt") #expect(prompt.arguments?.count == 1) #expect(prompt.arguments?[0].name == "test_arg") + #expect(prompt.arguments?[0].title == "Test Argument Title") #expect(prompt.arguments?[0].description == "A test argument") #expect(prompt.arguments?[0].required == true) } @Test("Prompt Message encoding and decoding") func testPromptMessageEncodingDecoding() throws { - let textMessage = Prompt.Message( - role: .user, - content: .text(text: "Hello, world!") - ) + let textMessage: Prompt.Message = .user("Hello, world!") let encoder = JSONEncoder() let decoder = JSONDecoder() @@ -63,6 +64,18 @@ struct PromptTests { #expect(Bool(false), "Expected text content") } + // Test audio content + let audioContent = Prompt.Message.Content.audio( + 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 { + #expect(data == "base64audiodata") + #expect(mimeType == "audio/wav") + } else { + #expect(Bool(false), "Expected audio content") + } + // Test image content let imageContent = Prompt.Message.Content.image(data: "base64data", mimeType: "image/png") let imageData = try encoder.encode(imageContent) @@ -75,19 +88,20 @@ struct PromptTests { } // Test resource content - let resourceContent = Prompt.Message.Content.resource( + let textResourceContent = Resource.Content.text( + "Sample text", uri: "file://test.txt", - mimeType: "text/plain", - text: "Sample text", - blob: "blob_data" + mimeType: "text/plain" ) + let resourceContent = Prompt.Message.Content.resource(resource: textResourceContent, annotations: nil, _meta: nil) 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, let annotations, let _meta) = decodedResource { + #expect(resourceData.uri == "file://test.txt") + #expect(resourceData.mimeType == "text/plain") + #expect(resourceData.text == "Sample text") + #expect(annotations == nil) + #expect(_meta == nil) } else { #expect(Bool(false), "Expected resource content") } @@ -95,8 +109,7 @@ struct PromptTests { @Test("Prompt Reference validation") func testPromptReference() throws { - let reference = Prompt.Reference(name: "test_prompt") - #expect(reference.name == "test_prompt") + let reference = Prompt.Reference(name: "test_prompt", title: "Test Prompt Title") let encoder = JSONEncoder() let decoder = JSONDecoder() @@ -104,45 +117,61 @@ struct PromptTests { let data = try encoder.encode(reference) let decoded = try decoder.decode(Prompt.Reference.self, from: data) - #expect(decoded.name == "test_prompt") + #expect(decoded == reference) } @Test("GetPrompt parameters validation") func testGetPromptParameters() throws { - let arguments: [String: Value] = [ - "param1": .string("value1"), - "param2": .int(42), + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + 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)) + let data = try encoder.encode(params) + let decoded = try decoder.decode(GetPrompt.Parameters.self, from: data) + + #expect(decoded == params) } @Test("GetPrompt result validation") func testGetPromptResult() throws { - let messages = [ - Prompt.Message(role: .user, content: .text(text: "User message")), - Prompt.Message(role: .assistant, content: .text(text: "Assistant response")), + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let messages: [Prompt.Message] = [ + .user("User message"), + .assistant("Assistant response"), ] let result = GetPrompt.Result(description: "Test description", messages: messages) - #expect(result.description == "Test description") - #expect(result.messages.count == 2) - #expect(result.messages[0].role == .user) - #expect(result.messages[1].role == .assistant) + let data = try encoder.encode(result) + let decoded = try decoder.decode(GetPrompt.Result.self, from: data) + + #expect(decoded == result) } @Test("ListPrompts parameters validation") func testListPromptsParameters() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + let params = ListPrompts.Parameters(cursor: "next_page") - #expect(params.cursor == "next_page") + let data = try encoder.encode(params) + let decoded = try decoder.decode(ListPrompts.Parameters.self, from: data) + + #expect(decoded == params) let emptyParams = ListPrompts.Parameters() - #expect(emptyParams.cursor == nil) + let emptyData = try encoder.encode(emptyParams) + let decodedEmpty = try decoder.decode(ListPrompts.Parameters.self, from: emptyData) + + #expect(decodedEmpty == emptyParams) } - + @Test("ListPrompts request decoding with omitted params") func testListPromptsRequestDecodingWithOmittedParams() throws { // Test decoding when params field is omitted @@ -157,7 +186,7 @@ struct PromptTests { #expect(decoded.id == "test-id") #expect(decoded.method == ListPrompts.name) } - + @Test("ListPrompts request decoding with null params") func testListPromptsRequestDecodingWithNullParams() throws { // Test decoding when params field is null @@ -175,20 +204,689 @@ struct PromptTests { @Test("ListPrompts result validation") func testListPromptsResult() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + let prompts = [ Prompt(name: "prompt1", description: "First prompt"), Prompt(name: "prompt2", description: "Second prompt"), ] let result = ListPrompts.Result(prompts: prompts, nextCursor: "next_page") - #expect(result.prompts.count == 2) - #expect(result.prompts[0].name == "prompt1") - #expect(result.prompts[1].name == "prompt2") - #expect(result.nextCursor == "next_page") + let data = try encoder.encode(result) + let decoded = try decoder.decode(ListPrompts.Result.self, from: data) + + #expect(decoded == result) } @Test("PromptListChanged notification name validation") func testPromptListChangedNotification() throws { #expect(PromptListChangedNotification.name == "notifications/prompts/list_changed") } + + @Test("Prompt Message factory methods") + func testPromptMessageFactoryMethods() throws { + // Test user message factory method + let userMessage: Prompt.Message = .user("Hello, world!") + #expect(userMessage.role == .user) + if case .text(let text) = userMessage.content { + #expect(text == "Hello, world!") + } else { + #expect(Bool(false), "Expected text content") + } + + // Test assistant message factory method + let assistantMessage: Prompt.Message = .assistant("Hi there!") + #expect(assistantMessage.role == .assistant) + if case .text(let text) = assistantMessage.content { + #expect(text == "Hi there!") + } else { + #expect(Bool(false), "Expected text content") + } + + // 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 { + #expect(data == "base64data") + #expect(mimeType == "image/png") + } else { + #expect(Bool(false), "Expected image content") + } + + // Test with audio content + let audioMessage: Prompt.Message = .assistant( + .audio(data: "base64audio", mimeType: "audio/wav")) + #expect(audioMessage.role == .assistant) + if case .audio(let data, let mimeType) = audioMessage.content { + #expect(data == "base64audio") + #expect(mimeType == "audio/wav") + } else { + #expect(Bool(false), "Expected audio content") + } + + // Test with resource content + let resourceContent = Resource.Content.text( + "Sample text", + uri: "file://test.txt", + mimeType: "text/plain" + ) + let resourceMessage: Prompt.Message = .user(.resource(resource: resourceContent, annotations: nil, _meta: nil)) + #expect(resourceMessage.role == .user) + if case .resource(let resource, let annotations, let _meta) = resourceMessage.content { + #expect(resource.uri == "file://test.txt") + #expect(resource.mimeType == "text/plain") + #expect(resource.text == "Sample text") + #expect(annotations == nil) + #expect(_meta == nil) + } else { + #expect(Bool(false), "Expected resource content") + } + } + + @Test("Prompt Content ExpressibleByStringLiteral") + func testPromptContentExpressibleByStringLiteral() throws { + // Test string literal assignment + let content: Prompt.Message.Content = "Hello from string literal" + + if case .text(let text) = content { + #expect(text == "Hello from string literal") + } else { + #expect(Bool(false), "Expected text content") + } + + // Test in message creation + let message: Prompt.Message = .user("Direct string literal") + if case .text(let text) = message.content { + #expect(text == "Direct string literal") + } else { + #expect(Bool(false), "Expected text content") + } + + // Test in array context + let messages: [Prompt.Message] = [ + .user("First message"), + .assistant("Second message"), + .user("Third message"), + ] + + #expect(messages.count == 3) + #expect(messages[0].role == .user) + #expect(messages[1].role == .assistant) + #expect(messages[2].role == .user) + } + + @Test("Prompt Content ExpressibleByStringInterpolation") + func testPromptContentExpressibleByStringInterpolation() throws { + let userName = "Alice" + let position = "Software Engineer" + let company = "TechCorp" + + // Test string interpolation + let content: Prompt.Message.Content = + "Hello \(userName), welcome to your \(position) interview at \(company)" + + 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") + } + + // 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 { + #expect(text == "Hi Alice, I'm excited about the Software Engineer role at TechCorp") + } else { + #expect(Bool(false), "Expected text content") + } + + // Test complex interpolation + let skills = ["Swift", "Python", "JavaScript"] + let experience = 5 + let interviewMessage: Prompt.Message = .assistant( + "I see you have \(experience) years of experience with \(skills.joined(separator: ", ")). That's impressive!" + ) + + if case .text(let text) = interviewMessage.content { + #expect( + text + == "I see you have 5 years of experience with Swift, Python, JavaScript. That's impressive!" + ) + } else { + #expect(Bool(false), "Expected text content") + } + } + + @Test("Prompt Message factory methods with string interpolation") + func testPromptMessageFactoryMethodsWithStringInterpolation() throws { + let candidateName = "Bob" + let position = "Data Scientist" + let company = "DataCorp" + let experience = 3 + + // Test user message with interpolation + 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 { + #expect(text == "Hello, I'm Bob and I'm interviewing for the Data Scientist position") + } else { + #expect(Bool(false), "Expected text content") + } + + // Test assistant message with interpolation + let assistantMessage: Prompt.Message = .assistant( + "Welcome \(candidateName)! Tell me about your \(experience) years of experience in data science" + ) + #expect(assistantMessage.role == .assistant) + 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") + } + + // Test in conversation array + let conversation: [Prompt.Message] = [ + .user("Hi, I'm \(candidateName) applying for \(position) at \(company)"), + .assistant("Welcome \(candidateName)! How many years of experience do you have?"), + .user("I have \(experience) years of experience in the field"), + .assistant( + "Great! \(experience) years is solid experience for a \(position) role at \(company)" + ), + ] + + #expect(conversation.count == 4) + + // Verify interpolated 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") + } + } + + @Test("Prompt ergonomic API usage patterns") + func testPromptErgonomicAPIUsagePatterns() throws { + // Test various ergonomic usage patterns enabled by the new API + + // Pattern 1: Simple interview conversation + let interviewConversation: [Prompt.Message] = [ + .user("Tell me about yourself"), + .assistant("I'm a software engineer with 5 years of experience"), + .user("What's your biggest strength?"), + .assistant("I'm great at problem-solving and team collaboration"), + ] + #expect(interviewConversation.count == 4) + + // Pattern 2: Dynamic content with interpolation + let candidateName = "Sarah" + let role = "Product Manager" + let yearsExp = 7 + + let dynamicConversation: [Prompt.Message] = [ + .user("Welcome \(candidateName) to the \(role) interview"), + .assistant("Thank you! I'm excited about this \(role) opportunity"), + .user("I see you have \(yearsExp) years of experience. Tell me about your background"), + .assistant( + "In my \(yearsExp) years as a \(role), I've led multiple successful product launches" + ), + ] + #expect(dynamicConversation.count == 4) + + // Pattern 3: Mixed content types + let mixedContent: [Prompt.Message] = [ + .user("Please review this design mockup"), + .assistant(.image(data: "design_mockup_data", mimeType: "image/png")), + .user("What do you think of the user flow?"), + .assistant( + "The design looks clean and intuitive. I particularly like the navigation structure." + ), + ] + #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 + { + // All content types are correct + } else { + #expect(Bool(false), "Content types don't match expected pattern") + } + + // Pattern 4: Encoding/decoding still works + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(interviewConversation) + let decoded = try decoder.decode([Prompt.Message].self, from: data) + + #expect(decoded.count == 4) + #expect(decoded[0].role == .user) + #expect(decoded[1].role == .assistant) + #expect(decoded[2].role == .user) + #expect(decoded[3].role == .assistant) + } +} + +@Suite("Prompt Integration Tests") +struct PromptIntegrationTests { + + @Test("List prompts with empty results") + func testListPromptsEmpty() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "PromptTestServer", + version: "1.0", + capabilities: .init(prompts: .init()) + ) + + // Register list prompts handler + await server.withMethodHandler(ListPrompts.self) { _ in + ListPrompts.Result(prompts: []) + } + + let client = Client( + name: "PromptTestClient", + version: "1.0" + ) + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + let (prompts, nextCursor) = try await client.listPrompts() + + #expect(prompts.isEmpty) + #expect(nextCursor == nil) + + await server.stop() + await client.disconnect() + } + + @Test("List prompts with multiple results") + func testListPromptsWithResults() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "PromptTestServer", + version: "1.0", + capabilities: .init(prompts: .init()) + ) + + let expectedPrompts = [ + Prompt( + name: "greeting", + title: "Greeting Prompt", + description: "A friendly greeting prompt" + ), + Prompt( + name: "interview", + title: "Interview Prompt", + description: "An interview preparation prompt", + arguments: [ + Prompt.Argument( + name: "position", + title: "Job Position", + description: "The job position to interview for", + required: true + ) + ] + ), + ] + + // Register list prompts handler + await server.withMethodHandler(ListPrompts.self) { _ in + ListPrompts.Result(prompts: expectedPrompts, nextCursor: "page2") + } + + let client = Client( + name: "PromptTestClient", + version: "1.0" + ) + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + let (prompts, nextCursor) = try await client.listPrompts() + + #expect(prompts == expectedPrompts) + #expect(nextCursor == "page2") + + await server.stop() + await client.disconnect() + } + + @Test("Get prompt with messages") + func testGetPrompt() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "PromptTestServer", + version: "1.0", + capabilities: .init(prompts: .init()) + ) + + let expectedMessages: [Prompt.Message] = [ + .user("Hello, I'd like to schedule an interview for the Software Engineer position"), + .assistant("I'd be happy to help you prepare for the Software Engineer interview. Let's discuss your background."), + ] + + // Register get prompt handler + await server.withMethodHandler(GetPrompt.self) { params in + #expect(params.name == "interview") + #expect(params.arguments?["position"] == "Software Engineer") + + return GetPrompt.Result( + description: "Interview preparation prompt", + messages: expectedMessages + ) + } + + let client = Client( + name: "PromptTestClient", + version: "1.0" + ) + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + let (description, messages) = try await client.getPrompt( + name: "interview", + arguments: ["position": "Software Engineer"] + ) + + #expect(description == "Interview preparation prompt") + #expect(messages == expectedMessages) + + await server.stop() + await client.disconnect() + } + + @Test("Get prompt with mixed content types") + func testGetPromptMixedContent() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "PromptTestServer", + version: "1.0", + capabilities: .init(prompts: .init()) + ) + + let expectedMessages: [Prompt.Message] = [ + .user("Please review this design mockup:"), + .user(.image(data: "base64_image_data", mimeType: "image/png")), + .assistant("I'll analyze the design mockup for you."), + .assistant(.audio(data: "base64_audio_data", mimeType: "audio/mp3")), + ] + + // Register get prompt handler + await server.withMethodHandler(GetPrompt.self) { params in + #expect(params.name == "design_review") + + return GetPrompt.Result( + description: "Design review prompt with multimedia", + messages: expectedMessages + ) + } + + let client = Client( + name: "PromptTestClient", + version: "1.0" + ) + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + let (description, messages) = try await client.getPrompt(name: "design_review") + + #expect(description == "Design review prompt with multimedia") + #expect(messages == expectedMessages) + + await server.stop() + await client.disconnect() + } + + @Test("Prompt with resource content") + func testPromptWithResourceContent() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "PromptTestServer", + version: "1.0", + capabilities: .init(prompts: .init()) + ) + + let resourceContent = Resource.Content.text( + "Code review content", + uri: "file://code.swift", + mimeType: "text/plain" + ) + + let expectedMessages: [Prompt.Message] = [ + .user("Review this code:"), + .user(.resource(resource: resourceContent, annotations: nil, _meta: nil)), + .assistant("I'll review the code for you."), + ] + + // Register get prompt handler + await server.withMethodHandler(GetPrompt.self) { _ in + GetPrompt.Result( + description: "Code review prompt", + messages: expectedMessages + ) + } + + let client = Client( + name: "PromptTestClient", + version: "1.0" + ) + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + let (description, messages) = try await client.getPrompt(name: "code_review") + + #expect(description == "Code review prompt") + #expect(messages == expectedMessages) + + await server.stop() + await client.disconnect() + } + + @Test("List prompts with pagination") + func testListPromptsWithPagination() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "PromptTestServer", + version: "1.0", + capabilities: .init(prompts: .init()) + ) + + // Register list prompts handler with pagination + await server.withMethodHandler(ListPrompts.self) { params in + if let cursor = params.cursor { + #expect(cursor == "page2") + return ListPrompts.Result( + prompts: [ + Prompt(name: "prompt3", description: "Third prompt"), + Prompt(name: "prompt4", description: "Fourth prompt"), + ], + nextCursor: nil + ) + } else { + return ListPrompts.Result( + prompts: [ + Prompt(name: "prompt1", description: "First prompt"), + Prompt(name: "prompt2", description: "Second prompt"), + ], + nextCursor: "page2" + ) + } + } + + let client = Client( + name: "PromptTestClient", + version: "1.0" + ) + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // First page + let (page1Prompts, page1Cursor) = try await client.listPrompts() + #expect(page1Prompts.count == 2) + #expect(page1Prompts[0].name == "prompt1") + #expect(page1Prompts[1].name == "prompt2") + #expect(page1Cursor == "page2") + + // Second page + let (page2Prompts, page2Cursor) = try await client.listPrompts(cursor: "page2") + #expect(page2Prompts.count == 2) + #expect(page2Prompts[0].name == "prompt3") + #expect(page2Prompts[1].name == "prompt4") + #expect(page2Cursor == nil) + + await server.stop() + await client.disconnect() + } + + @Test("Prompt without capability fails") + func testPromptWithoutCapabilityFails() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + // Server WITHOUT prompts capability + let server = Server( + name: "PromptTestServer", + version: "1.0" + ) + + let client = Client( + name: "PromptTestClient", + version: "1.0" + ) + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Should throw an error because server doesn't have prompts capability + await #expect(throws: MCPError.self) { + _ = try await client.listPrompts() + } + + await server.stop() + await client.disconnect() + } + + @Test("Strict mode succeeds when server declares prompts capability") + func testPromptStrictCapabilitiesSuccess() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "StrictPromptTestServer", + version: "1.0", + capabilities: .init(prompts: .init()), + configuration: .strict + ) + + // Register list prompts handler + await server.withMethodHandler(ListPrompts.self) { _ in + ListPrompts.Result( + prompts: [ + Prompt(name: "test", description: "Test prompt") + ] + ) + } + + let client = Client( + name: "StrictPromptTestClient", + version: "1.0", + configuration: .strict + ) + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Should succeed because server declares prompts capability + let (prompts, _) = try await client.listPrompts() + + #expect(prompts.count == 1) + #expect(prompts[0].name == "test") + + await server.stop() + await client.disconnect() + } + + @Test("Strict mode fails when server doesn't declare prompts capability") + func testPromptStrictCapabilitiesError() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + // Server WITHOUT prompts capability in strict mode + let server = Server( + name: "StrictPromptTestServer", + version: "1.0", + capabilities: .init(), + configuration: .strict + ) + + let client = Client( + name: "StrictPromptTestClient", + version: "1.0", + configuration: .strict + ) + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Should fail because server doesn't declare prompts capability in strict mode + await #expect(throws: MCPError.self) { + _ = try await client.listPrompts() + } + + await server.stop() + await client.disconnect() + } + + @Test("Non-strict mode succeeds even without server capability declaration") + func testPromptNonStrictCapabilities() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + // Server WITHOUT prompts capability in non-strict mode + let server = Server( + name: "NonStrictPromptTestServer", + version: "1.0", + capabilities: .init(), + configuration: .default + ) + + // Register list prompts handler anyway + await server.withMethodHandler(ListPrompts.self) { _ in + ListPrompts.Result( + prompts: [ + Prompt(name: "non_strict_test", description: "Non-strict test prompt") + ] + ) + } + + let client = Client( + name: "NonStrictPromptTestClient", + version: "1.0", + configuration: .default + ) + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Should succeed because client is in non-strict mode + let (prompts, _) = try await client.listPrompts() + + #expect(prompts.count == 1) + #expect(prompts[0].name == "non_strict_test") + + await server.stop() + await client.disconnect() + } } diff --git a/Tests/MCPTests/RequestTests.swift b/Tests/MCPTests/RequestTests.swift index 7d4e8dd..86747e2 100644 --- a/Tests/MCPTests/RequestTests.swift +++ b/Tests/MCPTests/RequestTests.swift @@ -23,33 +23,33 @@ struct RequestTests { @Test("Request initialization with parameters") func testRequestInitialization() throws { - let id: ID = "test-id" - let params = TestMethod.Parameters(value: "test") - let request = Request(id: id, method: TestMethod.name, params: params) + let id: ID = 1 + let params = CallTool.Parameters(name: "test-tool") + let request = Request(id: id, method: CallTool.name, params: params) #expect(request.id == id) - #expect(request.method == TestMethod.name) - #expect(request.params.value == "test") + #expect(request.method == CallTool.name) + #expect(request.params.name == "test-tool") } @Test("Request encoding and decoding") func testRequestEncodingDecoding() throws { - let request = TestMethod.request(id: "test-id", TestMethod.Parameters(value: "test")) + let request = CallTool.request(id: 1, CallTool.Parameters(name: "test-tool")) let encoder = JSONEncoder() let decoder = JSONDecoder() let data = try encoder.encode(request) - let decoded = try decoder.decode(Request.self, from: data) + let decoded = try decoder.decode(Request.self, from: data) #expect(decoded.id == request.id) #expect(decoded.method == request.method) - #expect(decoded.params.value == request.params.value) + #expect(decoded.params.name == request.params.name) } @Test("Empty parameters request encoding") func testEmptyParametersRequestEncoding() throws { - let request = EmptyMethod.request(id: "test-id") + let request = EmptyMethod.request(id: 1) let encoder = JSONEncoder() let decoder = JSONDecoder() @@ -66,14 +66,14 @@ struct RequestTests { func testEmptyParametersRequestDecoding() throws { // Create a minimal JSON string let jsonString = """ - {"jsonrpc":"2.0","id":"test-id","method":"empty.method"} + {"jsonrpc":"2.0","id":1,"method":"empty.method"} """ let data = jsonString.data(using: .utf8)! let decoder = JSONDecoder() let decoded = try decoder.decode(Request.self, from: data) - #expect(decoded.id == "test-id") + #expect(decoded.id == 1) #expect(decoded.method == EmptyMethod.name) } @@ -81,14 +81,14 @@ struct RequestTests { func testNotRequiredParametersRequestDecodingWithParams() throws { // Test decoding when params field is present let jsonString = """ - {"jsonrpc":"2.0","id":"test-id","method":"ping","params":{}} + {"jsonrpc":"2.0","id":1,"method":"ping","params":{}} """ let data = jsonString.data(using: .utf8)! let decoder = JSONDecoder() let decoded = try decoder.decode(Request.self, from: data) - #expect(decoded.id == "test-id") + #expect(decoded.id == 1) #expect(decoded.method == Ping.name) } @@ -96,14 +96,14 @@ struct RequestTests { func testNotRequiredParametersRequestDecodingWithoutParams() throws { // Test decoding when params field is missing let jsonString = """ - {"jsonrpc":"2.0","id":"test-id","method":"ping"} + {"jsonrpc":"2.0","id":1,"method":"ping"} """ let data = jsonString.data(using: .utf8)! let decoder = JSONDecoder() let decoded = try decoder.decode(Request.self, from: data) - #expect(decoded.id == "test-id") + #expect(decoded.id == 1) #expect(decoded.method == Ping.name) } @@ -111,14 +111,198 @@ struct RequestTests { func testNotRequiredParametersRequestDecodingWithNullParams() throws { // Test decoding when params field is null let jsonString = """ - {"jsonrpc":"2.0","id":"test-id","method":"ping","params":null} + {"jsonrpc":"2.0","id":1,"method":"ping","params":null} """ let data = jsonString.data(using: .utf8)! let decoder = JSONDecoder() let decoded = try decoder.decode(Request.self, from: data) - #expect(decoded.id == "test-id") + #expect(decoded.id == 1) + #expect(decoded.method == Ping.name) + } + + @Test("Required parameters request decoding - missing params") + func testRequiredParametersRequestDecodingMissingParams() throws { + let jsonString = """ + {"jsonrpc":"2.0","id":1,"method":"tools/call"} + """ + let data = jsonString.data(using: .utf8)! + + let decoder = JSONDecoder() + #expect(throws: DecodingError.self) { + _ = try decoder.decode(Request.self, from: data) + } + } + + @Test("Required parameters request decoding - null params") + func testRequiredParametersRequestDecodingNullParams() throws { + let jsonString = """ + {"jsonrpc":"2.0","id":1,"method":"tools/call","params":null} + """ + let data = jsonString.data(using: .utf8)! + + let decoder = JSONDecoder() + #expect(throws: DecodingError.self) { + _ = try decoder.decode(Request.self, from: data) + } + } + + @Test("Empty parameters request decoding - with null params") + func testEmptyParametersRequestDecodingNullParams() throws { + let jsonString = """ + {"jsonrpc":"2.0","id":1,"method":"empty.method","params":null} + """ + let data = jsonString.data(using: .utf8)! + + let decoder = JSONDecoder() + let decoded = try decoder.decode(Request.self, from: data) + + #expect(decoded.id == 1) + #expect(decoded.method == EmptyMethod.name) + } + + @Test("Empty parameters request decoding - with empty object params") + func testEmptyParametersRequestDecodingEmptyParams() throws { + let jsonString = """ + {"jsonrpc":"2.0","id":1,"method":"empty.method","params":{}} + """ + let data = jsonString.data(using: .utf8)! + + let decoder = JSONDecoder() + let decoded = try decoder.decode(Request.self, from: data) + + #expect(decoded.id == 1) + #expect(decoded.method == EmptyMethod.name) + } + + @Test("Initialize request decoding - requires params") + func testInitializeRequestDecodingRequiresParams() throws { + // Test missing params field + let missingParams = """ + {"jsonrpc":"2.0","id":"test-id","method":"initialize"} + """ + let decoder = JSONDecoder() + #expect(throws: DecodingError.self) { + _ = try decoder.decode( + Request.self, from: missingParams.data(using: .utf8)!) + } + + // Test null params + let nullParams = """ + {"jsonrpc":"2.0","id":"test-id","method":"initialize","params":null} + """ + #expect(throws: DecodingError.self) { + _ = try decoder.decode(Request.self, from: nullParams.data(using: .utf8)!) + } + + // Verify that empty object params works (since fields have defaults) + let emptyParams = """ + {"jsonrpc":"2.0","id":"test-id","method":"initialize","params":{}} + """ + let decoded = try decoder.decode( + Request.self, from: emptyParams.data(using: .utf8)!) + #expect(decoded.params.protocolVersion == Version.latest) + #expect(decoded.params.clientInfo.name == "unknown") + } + + @Test("Invalid parameters request decoding") + func testInvalidParametersRequestDecoding() throws { + let jsonString = """ + {"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"invalid":"value"}} + """ + let data = jsonString.data(using: .utf8)! + + let decoder = JSONDecoder() + #expect(throws: DecodingError.self) { + _ = try decoder.decode(Request.self, from: data) + } + } + + @Test("NotRequired parameters request decoding") + func testNotRequiredParametersRequestDecoding() throws { + // Test with missing params + let missingParams = """ + {"jsonrpc":"2.0","id":1,"method":"tools/list"} + """ + let decoder = JSONDecoder() + let decodedMissing = try decoder.decode( + Request.self, + from: missingParams.data(using: .utf8)!) + #expect(decodedMissing.id == 1) + #expect(decodedMissing.method == ListTools.name) + #expect(decodedMissing.params.cursor == nil) + + // Test with null params + let nullParams = """ + {"jsonrpc":"2.0","id":1,"method":"tools/list","params":null} + """ + let decodedNull = try decoder.decode( + Request.self, + from: nullParams.data(using: .utf8)!) + #expect(decodedNull.params.cursor == nil) + + // Test with empty object params + let emptyParams = """ + {"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}} + """ + let decodedEmpty = try decoder.decode( + Request.self, + from: emptyParams.data(using: .utf8)!) + #expect(decodedEmpty.params.cursor == nil) + + // Test with provided cursor + let withCursor = """ + {"jsonrpc":"2.0","id":1,"method":"tools/list","params":{"cursor":"next-page"}} + """ + let decodedWithCursor = try decoder.decode( + Request.self, + from: withCursor.data(using: .utf8)!) + #expect(decodedWithCursor.params.cursor == "next-page") + } + + @Test("AnyRequest parameters request decoding - without params") + func testAnyRequestParametersRequestDecodingWithoutParams() throws { + // Test decoding when params field is missing + let jsonString = """ + {"jsonrpc":"2.0","id":1,"method":"ping"} + """ + let data = jsonString.data(using: .utf8)! + + let decoder = JSONDecoder() + let decoded = try decoder.decode(AnyRequest.self, from: data) + + #expect(decoded.id == 1) + #expect(decoded.method == Ping.name) + } + + @Test("AnyRequest parameters request decoding - with null params") + func testAnyRequestParametersRequestDecodingWithNullParams() throws { + // Test decoding when params field is null + let jsonString = """ + {"jsonrpc":"2.0","id":1,"method":"ping","params":null} + """ + let data = jsonString.data(using: .utf8)! + + let decoder = JSONDecoder() + let decoded = try decoder.decode(Request.self, from: data) + + #expect(decoded.id == 1) + #expect(decoded.method == Ping.name) + } + + @Test("AnyRequest parameters request decoding - with empty params") + func testAnyRequestParametersRequestDecodingWithEmptyParams() throws { + // Test decoding when params field is null + let jsonString = """ + {"jsonrpc":"2.0","id":1,"method":"ping","params":{}} + """ + let data = jsonString.data(using: .utf8)! + + let decoder = JSONDecoder() + let decoded = try decoder.decode(Request.self, from: data) + + #expect(decoded.id == 1) #expect(decoded.method == Ping.name) } } diff --git a/Tests/MCPTests/ResourceTests.swift b/Tests/MCPTests/ResourceTests.swift index 5403632..4c08d66 100644 --- a/Tests/MCPTests/ResourceTests.swift +++ b/Tests/MCPTests/ResourceTests.swift @@ -10,16 +10,18 @@ struct ResourceTests { let resource = Resource( name: "test_resource", uri: "file://test.txt", + title: "Test Resource Title", description: "A test resource", mimeType: "text/plain", - metadata: ["key": "value"] + size: 1024 ) #expect(resource.name == "test_resource") #expect(resource.uri == "file://test.txt") + #expect(resource.title == "Test Resource Title") #expect(resource.description == "A test resource") #expect(resource.mimeType == "text/plain") - #expect(resource.metadata?["key"] == "value") + #expect(resource.size == 1024) } @Test("Resource encoding and decoding") @@ -27,9 +29,10 @@ struct ResourceTests { let resource = Resource( name: "test_resource", uri: "file://test.txt", + title: "Test Resource Title", description: "Test resource description", mimeType: "text/plain", - metadata: ["key1": "value1", "key2": "value2"] + size: 512 ) let encoder = JSONEncoder() @@ -40,9 +43,10 @@ struct ResourceTests { #expect(decoded.name == resource.name) #expect(decoded.uri == resource.uri) + #expect(decoded.title == resource.title) #expect(decoded.description == resource.description) #expect(decoded.mimeType == resource.mimeType) - #expect(decoded.metadata == resource.metadata) + #expect(decoded.size == resource.size) } @Test("Resource.Content text initialization and encoding") @@ -86,7 +90,7 @@ struct ResourceTests { let emptyParams = ListResources.Parameters() #expect(emptyParams.cursor == nil) } - + @Test("ListResources request decoding with omitted params") func testListResourcesRequestDecodingWithOmittedParams() throws { // Test decoding when params field is omitted @@ -101,7 +105,7 @@ struct ResourceTests { #expect(decoded.id == "test-id") #expect(decoded.method == ListResources.name) } - + @Test("ListResources request decoding with null params") func testListResourcesRequestDecodingWithNullParams() throws { // Test decoding when params field is null diff --git a/Tests/MCPTests/ResponseTests.swift b/Tests/MCPTests/ResponseTests.swift index f09a916..cb91852 100644 --- a/Tests/MCPTests/ResponseTests.swift +++ b/Tests/MCPTests/ResponseTests.swift @@ -43,7 +43,7 @@ struct ResponseTests { @Test("Error response initialization and encoding") func testErrorResponse() throws { let id: ID = "test-id" - let error = MCP.Error.parseError(nil) + let error = MCPError.parseError(nil) let response = Response(id: id, error: error) let encoder = JSONEncoder() @@ -65,7 +65,7 @@ struct ResponseTests { @Test("Error response with detail") func testErrorResponseWithDetail() throws { let id: ID = "test-id" - let error = MCP.Error.parseError("Invalid syntax") + let error = MCPError.parseError("Invalid syntax") let response = Response(id: id, error: error) let encoder = JSONEncoder() @@ -78,7 +78,7 @@ struct ResponseTests { #expect(decodedError.code == -32700) #expect( decodedError.localizedDescription - == "Parse error: Invalid JSON: Parse error: Invalid JSON: Invalid syntax") + == "Parse error: Invalid JSON: Invalid syntax") } else { #expect(Bool(false), "Expected error result") } diff --git a/Tests/MCPTests/RootsTests.swift b/Tests/MCPTests/RootsTests.swift new file mode 100644 index 0000000..86a6e5e --- /dev/null +++ b/Tests/MCPTests/RootsTests.swift @@ -0,0 +1,494 @@ +import Foundation +import Testing + +@testable import MCP + +@Suite("Roots Tests") +struct RootsTests { + @Test("Root initialization with file:// URI") + func testRootInitialization() throws { + let root = Root(uri: "file:///workspace", name: "Workspace") + + #expect(root.uri == "file:///workspace") + #expect(root.name == "Workspace") + } + + @Test("Root initialization without name") + func testRootInitializationWithoutName() throws { + let root = Root(uri: "file:///home/user/docs") + + #expect(root.uri == "file:///home/user/docs") + #expect(root.name == nil) + } + + @Test("Root encoding and decoding") + func testRootEncodingDecoding() throws { + let root = Root(uri: "file:///workspace", name: "Workspace") + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(root) + let decoded = try decoder.decode(Root.self, from: data) + + #expect(decoded.uri == root.uri) + #expect(decoded.name == root.name) + } + + @Test("Root encoding and decoding without name") + func testRootEncodingDecodingWithoutName() throws { + let root = Root(uri: "file:///home/user/docs") + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(root) + let decoded = try decoder.decode(Root.self, from: data) + + #expect(decoded.uri == root.uri) + #expect(decoded.name == nil) + } + + @Test("Root JSON encoding format") + func testRootJSONFormat() throws { + let root = Root(uri: "file:///workspace", name: "Workspace") + + let encoder = JSONEncoder() + let data = try encoder.encode(root) + let jsonString = String(data: data, encoding: .utf8) + + #expect(jsonString?.contains("\"uri\"") == true) + #expect(jsonString?.contains("\"name\"") == true) + // The URI might be escaped differently, so just check it's present + #expect(jsonString?.contains("workspace") == true) + #expect(jsonString?.contains("Workspace") == true) + } + + @Test("ListRoots method name") + func testListRootsMethodName() throws { + #expect(ListRoots.name == "roots/list") + } + + @Test("ListRoots request creation") + func testListRootsRequestCreation() throws { + let request = ListRoots.request() + + #expect(request.method == "roots/list") + #expect(request.params == Empty()) + } + + @Test("ListRoots request decoding with omitted params") + func testListRootsRequestDecodingWithOmittedParams() throws { + // Test decoding when params field is omitted + let jsonString = """ + {"jsonrpc":"2.0","id":"test-id","method":"roots/list"} + """ + let data = jsonString.data(using: .utf8)! + + let decoder = JSONDecoder() + let decoded = try decoder.decode(Request.self, from: data) + + #expect(decoded.id == "test-id") + #expect(decoded.method == ListRoots.name) + } + + @Test("ListRoots request decoding with null params") + func testListRootsRequestDecodingWithNullParams() throws { + // Test decoding when params field is null + let jsonString = """ + {"jsonrpc":"2.0","id":"test-id","method":"roots/list","params":null} + """ + let data = jsonString.data(using: .utf8)! + + let decoder = JSONDecoder() + let decoded = try decoder.decode(Request.self, from: data) + + #expect(decoded.id == "test-id") + #expect(decoded.method == ListRoots.name) + } + + @Test("ListRoots result validation") + func testListRootsResult() throws { + let roots = [ + Root(uri: "file:///workspace", name: "Workspace"), + Root(uri: "file:///home/user/docs", name: "Documents"), + ] + + let result = ListRoots.Result(roots: roots) + #expect(result.roots.count == 2) + #expect(result.roots[0].uri == "file:///workspace") + #expect(result.roots[0].name == "Workspace") + #expect(result.roots[1].uri == "file:///home/user/docs") + #expect(result.roots[1].name == "Documents") + } + + @Test("ListRoots result encoding and decoding") + func testListRootsResultEncodingDecoding() throws { + let roots = [ + Root(uri: "file:///workspace", name: "Workspace"), + Root(uri: "file:///home/user/docs"), + ] + let result = ListRoots.Result(roots: roots) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(result) + let decoded = try decoder.decode(ListRoots.Result.self, from: data) + + #expect(decoded.roots.count == result.roots.count) + #expect(decoded.roots[0].uri == result.roots[0].uri) + #expect(decoded.roots[0].name == result.roots[0].name) + #expect(decoded.roots[1].uri == result.roots[1].uri) + #expect(decoded.roots[1].name == result.roots[1].name) + } + + @Test("ListRoots response format") + func testListRootsResponseFormat() throws { + let response = ListRoots.response( + id: "test-id", + result: ListRoots.Result( + roots: [ + Root(uri: "file:///workspace", name: "Workspace") + ] + ) + ) + + let encoder = JSONEncoder() + let data = try encoder.encode(response) + let jsonString = String(data: data, encoding: .utf8) + + #expect(jsonString?.contains("\"jsonrpc\"") == true) + #expect(jsonString?.contains("\"2.0\"") == true) + #expect(jsonString?.contains("\"id\"") == true) + #expect(jsonString?.contains("\"result\"") == true) + #expect(jsonString?.contains("\"roots\"") == true) + } + + @Test("RootsListChangedNotification name") + func testRootsListChangedNotificationName() throws { + #expect(RootsListChangedNotification.name == "notifications/roots/list_changed") + } + + @Test("RootsListChangedNotification message creation") + func testRootsListChangedNotificationMessage() throws { + let notification = RootsListChangedNotification.message() + + #expect(notification.method == "notifications/roots/list_changed") + #expect(notification.params == Empty()) + } + + @Test("RootsListChangedNotification encoding") + func testRootsListChangedNotificationEncoding() throws { + let notification = RootsListChangedNotification.message() + + let encoder = JSONEncoder() + let data = try encoder.encode(notification) + let jsonString = String(data: data, encoding: .utf8) + + #expect(jsonString?.contains("\"jsonrpc\"") == true) + #expect(jsonString?.contains("\"2.0\"") == true) + #expect(jsonString?.contains("\"method\"") == true) + // The method name might be escaped differently, so check key parts + #expect(jsonString?.contains("roots") == true) + #expect(jsonString?.contains("list_changed") == true) + } + + @Test("Client.Capabilities.Roots initialization") + func testClientCapabilitiesRoots() throws { + let roots = Client.Capabilities.Roots(listChanged: true) + + #expect(roots.listChanged == true) + } + + @Test("Client.Capabilities.Roots optional listChanged") + func testClientCapabilitiesRootsOptional() throws { + let roots = Client.Capabilities.Roots() + + #expect(roots.listChanged == nil) + } + + @Test("Client.Capabilities with roots") + func testClientCapabilitiesWithRoots() throws { + let capabilities = Client.Capabilities( + roots: Client.Capabilities.Roots(listChanged: true) + ) + + #expect(capabilities.roots?.listChanged == true) + } + + @Test("Root hashable conformance") + func testRootHashable() throws { + let root1 = Root(uri: "file:///workspace", name: "Workspace") + let root2 = Root(uri: "file:///workspace", name: "Workspace") + let root3 = Root(uri: "file:///other", name: "Other") + + #expect(root1 == root2) + #expect(root1 != root3) + + // Test in a Set + let set: Set = [root1, root2, root3] + #expect(set.count == 2) // root1 and root2 should be the same + } + + @Test("Client Capabilities struct with roots encoding") + func testClientCapabilitiesWithRootsEncoding() throws { + let capabilities = Client.Capabilities( + roots: Client.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?.listChanged == true) + } + + // MARK: - Integration Tests + + @Test("Server listRoots with client that has roots capability") + func testServerListRootsWithCapableClient() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "test-server", + version: "1.0.0", + configuration: .strict + ) + + let testRoots = [ + Root(uri: "file:///workspace", name: "Workspace"), + Root(uri: "file:///home/user/docs", name: "Documents"), + ] + + let client = Client( + name: "test-client", + version: "1.0.0", + capabilities: Client.Capabilities( + roots: Client.Capabilities.Roots(listChanged: true) + ) + ) + + await client.withRootsHandler { + return testRoots + } + + // Start server and client + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Wait for initialization to complete + try await Task.sleep(for: .milliseconds(100)) + + // Server requests roots from client - should succeed + let roots = try await server.listRoots() + + #expect(roots == testRoots) + + // Cleanup + await server.stop() + await client.disconnect() + } + + @Test("Server listRoots fails when client lacks roots capability (strict mode)") + func testServerListRootsFailsWithoutCapability() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "test-server", + version: "1.0.0", + configuration: .strict + ) + + // Client has NO roots capability (default empty capabilities) + let client = Client( + name: "test-client", + version: "1.0.0" + ) + + // Start server and client + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Wait for initialization to complete + try await Task.sleep(for: .milliseconds(100)) + + // Server tries to request roots - should fail + await #expect(throws: MCPError.self) { + try await server.listRoots() + } + + // Cleanup + await server.stop() + await client.disconnect() + } + + @Test("Server listRoots succeeds when client lacks capability (non-strict mode)") + func testServerListRootsNonStrictMode() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "test-server", + version: "1.0.0", + configuration: .default // Non-strict mode + ) + + let testRoots = [Root(uri: "file:///workspace")] + + // Client has NO roots capability but registers handler anyway + let client = Client( + name: "test-client", + version: "1.0.0" + ) + + await client.withRootsHandler { + return testRoots + } + + // Start server and client + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Wait for initialization to complete + try await Task.sleep(for: .milliseconds(100)) + + // Server requests roots - should succeed in non-strict mode + let roots = try await server.listRoots() + #expect(roots == testRoots) + + // Cleanup + await server.stop() + await client.disconnect() + } + + @Test("Client sends roots list changed notification") + func testClientSendsRootsListChangedNotification() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "test-server", + version: "1.0.0" + ) + + let client = Client( + name: "test-client", + version: "1.0.0", + capabilities: Client.Capabilities( + roots: Client.Capabilities.Roots(listChanged: true) + ) + ) + + // Register notification handler on server + nonisolated(unsafe) var didReceive = false + await server.onNotification(RootsListChangedNotification.self) { _ in + didReceive = true + } + + // Start server and client + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Wait for initialization + try await Task.sleep(for: .milliseconds(100)) + + // Client sends notification + try await client.notifyRootsChanged() + + // Wait for notification to be processed + try await Task.sleep(for: .milliseconds(100)) + + // Verify notification was received + #expect(didReceive) + + // Cleanup + await server.stop() + await client.disconnect() + } + + @Test("Server listRoots fails when client has no handler registered") + func testServerListRootsFailsWithoutHandler() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "test-server", + version: "1.0.0", + configuration: .default + ) + + let client = Client( + name: "test-client", + version: "1.0.0", + capabilities: Client.Capabilities( + roots: Client.Capabilities.Roots(listChanged: true) + ) + ) + + // NO handler registered on client + + // Start server and client + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Wait for initialization + try await Task.sleep(for: .milliseconds(100)) + + // Server requests roots - should fail with method not found + await #expect(throws: MCPError.self) { + try await server.listRoots() + } + + // Cleanup + await server.stop() + await client.disconnect() + } + + @Test("Multiple roots requests work correctly") + func testMultipleRootsRequests() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "test-server", + version: "1.0.0" + ) + + nonisolated(unsafe) var count = 0 + + let client = Client( + name: "test-client", + version: "1.0.0", + capabilities: Client.Capabilities( + roots: Client.Capabilities.Roots(listChanged: true) + ) + ) + + await client.withRootsHandler { + count += 1 + return [Root(uri: "file:///workspace\(count)")] + } + + // Start server and client + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + try await Task.sleep(for: .milliseconds(100)) + + // Make multiple requests + let roots1 = try await server.listRoots() + #expect(roots1.count == 1) + #expect(roots1[0].uri == "file:///workspace1") + + let roots2 = try await server.listRoots() + #expect(roots2.count == 1) + #expect(roots2[0].uri == "file:///workspace2") + + let roots3 = try await server.listRoots() + #expect(roots3.count == 1) + #expect(roots3[0].uri == "file:///workspace3") + + // Cleanup + await server.stop() + await client.disconnect() + } +} diff --git a/Tests/MCPTests/RoundtripTests.swift b/Tests/MCPTests/RoundtripTests.swift index 0de1210..b91825d 100644 --- a/Tests/MCPTests/RoundtripTests.swift +++ b/Tests/MCPTests/RoundtripTests.swift @@ -1,10 +1,15 @@ import Foundation import Logging -import SystemPackage import Testing @testable import MCP +#if canImport(System) + import System +#else + @preconcurrency import SystemPackage +#endif + @Suite("Roundtrip Tests") struct RoundtripTests { @Test( @@ -52,17 +57,17 @@ struct RoundtripTests { } await server.withMethodHandler(CallTool.self) { request in guard request.name == "add" else { - return CallTool.Result(content: [.text("Invalid tool name")], isError: true) + return CallTool.Result(content: [.text(text: "Invalid tool name", annotations: nil, _meta: nil)], isError: true) } guard let a = request.arguments?["a"]?.intValue, let b = request.arguments?["b"]?.intValue else { return CallTool.Result( - content: [.text("Did not receive valid arguments")], isError: true) + content: [.text(text: "Did not receive valid arguments", annotations: nil, _meta: nil)], isError: true) } - return CallTool.Result(content: [.text("\(a + b)")]) + return CallTool.Result(content: [.text(text: "\(a + b)", annotations: nil, _meta: nil)]) } // Add resource handlers to server @@ -94,10 +99,9 @@ struct RoundtripTests { let client = Client(name: "TestClient", version: "1.0") try await server.start(transport: serverTransport) - try await client.connect(transport: clientTransport) let initTask = Task { - let result = try await client.initialize() + let result = try await client.connect(transport: clientTransport) #expect(result.serverInfo.name == "TestServer") #expect(result.serverInfo.version == "1.0.0") @@ -122,7 +126,7 @@ struct RoundtripTests { let pingTask = Task { try await client.ping() // Ping doesn't return anything, so just getting here without throwing is success - #expect(true) // Test passed if we reach this point + #expect(Bool(true)) } try await withThrowingTaskGroup(of: Void.self) { group in @@ -139,15 +143,15 @@ struct RoundtripTests { } let listToolsTask = Task { - let result = try await client.listTools() - #expect(result.count == 1) - #expect(result[0].name == "add") + let (tools, _) = try await client.listTools() + #expect(tools.count == 1) + #expect(tools[0].name == "add") } let callToolTask = Task { let result = try await client.callTool(name: "add", arguments: ["a": 1, "b": 2]) #expect(result.isError == nil) - #expect(result.content == [.text("3")]) + #expect(result.content == [.text(text: "3", annotations: nil, _meta: nil)]) } try await withThrowingTaskGroup(of: Void.self) { group in diff --git a/Tests/MCPTests/SamplingTests.swift b/Tests/MCPTests/SamplingTests.swift new file mode 100644 index 0000000..0ecafbc --- /dev/null +++ b/Tests/MCPTests/SamplingTests.swift @@ -0,0 +1,1093 @@ +import Logging +import Testing + +import class Foundation.JSONDecoder +import class Foundation.JSONEncoder + +@testable import MCP + +@Suite("Sampling Tests") +struct SamplingTests { + @Test("Sampling.Message encoding and decoding") + func testSamplingMessageCoding() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + // Test text content + let textMessage: Sampling.Message = .user("Hello, world!") + + let textData = try encoder.encode(textMessage) + let decodedTextMessage = try decoder.decode(Sampling.Message.self, from: textData) + + #expect(decodedTextMessage.role == .user) + if case .single(.text(let text)) = decodedTextMessage.content { + #expect(text == "Hello, world!") + } else { + #expect(Bool(false), "Expected text content") + } + + // Test image content + let imageMessage: Sampling.Message = .assistant( + .image(data: "base64imagedata", mimeType: "image/png")) + + let imageData = try encoder.encode(imageMessage) + let decodedImageMessage = try decoder.decode(Sampling.Message.self, from: imageData) + + #expect(decodedImageMessage.role == .assistant) + if case .single(.image(let data, let mimeType)) = decodedImageMessage.content { + #expect(data == "base64imagedata") + #expect(mimeType == "image/png") + } else { + #expect(Bool(false), "Expected image content") + } + } + + @Test("ModelPreferences encoding and decoding") + func testModelPreferencesCoding() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let preferences = Sampling.ModelPreferences( + hints: [ + Sampling.ModelPreferences.Hint(name: "claude-4"), + Sampling.ModelPreferences.Hint(name: "gpt-4.1"), + ], + costPriority: 0.8, + speedPriority: 0.3, + intelligencePriority: 0.9 + ) + + let data = try encoder.encode(preferences) + let decoded = try decoder.decode(Sampling.ModelPreferences.self, from: data) + + #expect(decoded.hints?.count == 2) + #expect(decoded.hints?[0].name == "claude-4") + #expect(decoded.hints?[1].name == "gpt-4.1") + #expect(decoded.costPriority?.doubleValue == 0.8) + #expect(decoded.speedPriority?.doubleValue == 0.3) + #expect(decoded.intelligencePriority?.doubleValue == 0.9) + } + + @Test("ContextInclusion encoding and decoding") + func testContextInclusionCoding() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let contexts: [Sampling.ContextInclusion] = [.none, .thisServer, .allServers] + + for context in contexts { + let data = try encoder.encode(context) + let decoded = try decoder.decode(Sampling.ContextInclusion.self, from: data) + #expect(decoded == context) + } + } + + @Test("StopReason encoding and decoding") + func testStopReasonCoding() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let reasons: [Sampling.StopReason] = [.endTurn, .stopSequence, .maxTokens] + + for reason in reasons { + let data = try encoder.encode(reason) + let decoded = try decoder.decode(Sampling.StopReason.self, from: data) + #expect(decoded == reason) + } + } + + @Test("CreateMessage request parameters") + func testCreateMessageParameters() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let messages: [Sampling.Message] = [ + .user("What is the weather like?"), + .assistant("I need to check the weather for you."), + ] + + let modelPreferences = Sampling.ModelPreferences( + hints: [Sampling.ModelPreferences.Hint(name: "claude-4-sonnet")], + costPriority: 0.5, + speedPriority: 0.7, + intelligencePriority: 0.9 + ) + + let parameters = CreateSamplingMessage.Parameters( + messages: messages, + modelPreferences: modelPreferences, + systemPrompt: "You are a helpful weather assistant.", + includeContext: .thisServer, + temperature: 0.7, + maxTokens: 150, + stopSequences: ["END", "STOP"], + _meta: Metadata(additionalFields: ["provider": "test"]) + ) + + let data = try encoder.encode(parameters) + let decoded = try decoder.decode(CreateSamplingMessage.Parameters.self, from: data) + + #expect(decoded.messages.count == 2) + #expect(decoded.messages[0].role == .user) + #expect(decoded.systemPrompt == "You are a helpful weather assistant.") + #expect(decoded.includeContext == .thisServer) + #expect(decoded.temperature == 0.7) + #expect(decoded.maxTokens == 150) + #expect(decoded.stopSequences?.count == 2) + #expect(decoded.stopSequences?[0] == "END") + #expect(decoded.stopSequences?[1] == "STOP") + #expect(decoded._meta?["provider"]?.stringValue == "test") + } + + @Test("CreateMessage result") + func testCreateMessageResult() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let result = CreateSamplingMessage.Result( + model: "claude-4-sonnet", + stopReason: .endTurn, + role: .assistant, + content: .text("The weather is sunny and 75°F.") + ) + + let data = try encoder.encode(result) + let decoded = try decoder.decode(CreateSamplingMessage.Result.self, from: data) + + #expect(decoded.model == "claude-4-sonnet") + #expect(decoded.stopReason == .endTurn) + #expect(decoded.role == .assistant) + + if case .single(.text(let text)) = decoded.content { + #expect(text == "The weather is sunny and 75°F.") + } else { + #expect(Bool(false), "Expected text content") + } + } + + @Test("CreateMessage request creation") + func testCreateMessageRequest() throws { + let messages: [Sampling.Message] = [ + .user("Hello") + ] + + let request = CreateSamplingMessage.request( + .init( + messages: messages, + maxTokens: 100 + ) + ) + + #expect(request.method == "sampling/createMessage") + #expect(request.params.messages.count == 1) + #expect(request.params.maxTokens == 100) + } + + @Test("Client capabilities include sampling") + func testClientCapabilitiesIncludeSampling() throws { + let capabilities = Client.Capabilities( + sampling: .init() + ) + + #expect(capabilities.sampling != nil) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(capabilities) + let decoded = try decoder.decode(Client.Capabilities.self, from: data) + + #expect(decoded.sampling != nil) + } + + @Test("Client sampling handler registration") + func testClientSamplingHandlerRegistration() async throws { + let client = Client(name: "TestClient", version: "1.0") + + // Test that sampling handler can be registered + let handlerClient = await client.withSamplingHandler { parameters in + // Mock handler that returns a simple response + return CreateSamplingMessage.Result( + model: "test-model", + stopReason: .endTurn, + role: .assistant, + content: .text("Test response") + ) + } + + // Should return self for method chaining + #expect(handlerClient === client) + } + + @Test("Sampling message content JSON format") + func testSamplingMessageContentJSONFormat() throws { + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + + // Test text content JSON format + let textContent: Sampling.Message.Content = .text("Hello") + let textData = try encoder.encode(textContent) + let textJSON = String(data: textData, encoding: .utf8)! + + #expect(textJSON.contains("\"type\":\"text\"")) + #expect(textJSON.contains("\"text\":\"Hello\"")) + + // Test image content JSON format + let imageContent: Sampling.Message.Content = .image( + data: "base64data", mimeType: "image/png") + let imageData = try encoder.encode(imageContent) + let imageJSON = String(data: imageData, encoding: .utf8)! + + #expect(imageJSON.contains("\"type\":\"image\"")) + #expect(imageJSON.contains("\"data\":\"base64data\"")) + #expect(imageJSON.contains("\"mimeType\":\"image\\/png\"")) + } + + @Test("UnitInterval in Sampling.ModelPreferences") + func testUnitIntervalInModelPreferences() throws { + // Test that UnitInterval validation works in Sampling.ModelPreferences + let validPreferences = Sampling.ModelPreferences( + costPriority: 0.5, + speedPriority: 1.0, + intelligencePriority: 0.0 + ) + + #expect(validPreferences.costPriority?.doubleValue == 0.5) + #expect(validPreferences.speedPriority?.doubleValue == 1.0) + #expect(validPreferences.intelligencePriority?.doubleValue == 0.0) + + // Test JSON encoding/decoding preserves UnitInterval constraints + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(validPreferences) + let decoded = try decoder.decode(Sampling.ModelPreferences.self, from: data) + + #expect(decoded.costPriority?.doubleValue == 0.5) + #expect(decoded.speedPriority?.doubleValue == 1.0) + #expect(decoded.intelligencePriority?.doubleValue == 0.0) + } + + @Test("Message factory methods") + func testMessageFactoryMethods() throws { + // Test user message factory method + let userMessage: Sampling.Message = .user("Hello, world!") + #expect(userMessage.role == .user) + if case .single(.text(let text)) = userMessage.content { + #expect(text == "Hello, world!") + } else { + #expect(Bool(false), "Expected text content") + } + + // Test assistant message factory method + let assistantMessage: Sampling.Message = .assistant("Hi there!") + #expect(assistantMessage.role == .assistant) + if case .single(.text(let text)) = assistantMessage.content { + #expect(text == "Hi there!") + } else { + #expect(Bool(false), "Expected text content") + } + + // Test with image content + let imageMessage: Sampling.Message = .user( + .image(data: "base64data", mimeType: "image/png")) + #expect(imageMessage.role == .user) + if case .single(.image(let data, let mimeType)) = imageMessage.content { + #expect(data == "base64data") + #expect(mimeType == "image/png") + } else { + #expect(Bool(false), "Expected image content") + } + } + + @Test("Content ExpressibleByStringLiteral") + func testContentExpressibleByStringLiteral() throws { + // Test string literal assignment + let content: Sampling.Message.Content = "Hello from string literal" + + if case .single(.text(let text)) = content { + #expect(text == "Hello from string literal") + } else { + #expect(Bool(false), "Expected text content") + } + + // Test in message creation + let message: Sampling.Message = .user("Direct string literal") + if case .single(.text(let text)) = message.content { + #expect(text == "Direct string literal") + } else { + #expect(Bool(false), "Expected text content") + } + + // Test in array context + let messages: [Sampling.Message] = [ + .user("First message"), + .assistant("Second message"), + .user("Third message"), + ] + + #expect(messages.count == 3) + #expect(messages[0].role == .user) + #expect(messages[1].role == .assistant) + #expect(messages[2].role == .user) + } + + @Test("Content ExpressibleByStringInterpolation") + func testContentExpressibleByStringInterpolation() throws { + let userName = "Alice" + let temperature = 72 + let location = "San Francisco" + + // Test string interpolation + let content: Sampling.Message.Content = + "Hello \(userName), the temperature in \(location) is \(temperature)°F" + + if case .single(.text(let text)) = content { + #expect(text == "Hello Alice, the temperature in San Francisco is 72°F") + } else { + #expect(Bool(false), "Expected text content") + } + + // Test in message creation with interpolation + let message = Sampling.Message.user( + "Welcome \(userName)! Today's weather in \(location) is \(temperature)°F") + if case .single(.text(let text)) = message.content { + #expect(text == "Welcome Alice! Today's weather in San Francisco is 72°F") + } else { + #expect(Bool(false), "Expected text content") + } + + // Test complex interpolation + let items = ["apples", "bananas", "oranges"] + let count = items.count + let listMessage: Sampling.Message = .assistant( + "You have \(count) items: \(items.joined(separator: ", "))") + + if case .single(.text(let text)) = listMessage.content { + #expect(text == "You have 3 items: apples, bananas, oranges") + } else { + #expect(Bool(false), "Expected text content") + } + } + + @Test("Message factory methods with string interpolation") + func testMessageFactoryMethodsWithStringInterpolation() throws { + let customerName = "Bob" + let orderNumber = "ORD-12345" + let issueType = "delivery delay" + + // Test user message with interpolation + let userMessage: Sampling.Message = .user( + "Hi, I'm \(customerName) and I have an issue with order \(orderNumber)") + #expect(userMessage.role == .user) + if case .single(.text(let text)) = userMessage.content { + #expect(text == "Hi, I'm Bob and I have an issue with order ORD-12345") + } else { + #expect(Bool(false), "Expected text content") + } + + // Test assistant message with interpolation + let assistantMessage: Sampling.Message = .assistant( + "Hello \(customerName), I can help you with your \(issueType) issue for order \(orderNumber)" + ) + #expect(assistantMessage.role == .assistant) + if case .single(.text(let text)) = assistantMessage.content { + #expect( + text + == "Hello Bob, I can help you with your delivery delay issue for order ORD-12345" + ) + } else { + #expect(Bool(false), "Expected text content") + } + + // Test in conversation array + let conversation: [Sampling.Message] = [ + .user("Hello, I'm \(customerName)"), + .assistant("Hi \(customerName), how can I help you today?"), + .user("I have an issue with order \(orderNumber) - it's a \(issueType)"), + .assistant( + "I understand you're experiencing a \(issueType) with order \(orderNumber). Let me look into that for you." + ), + ] + + #expect(conversation.count == 4) + + // Verify interpolated content + if case .single(.text(let text)) = conversation[2].content { + #expect(text == "I have an issue with order ORD-12345 - it's a delivery delay") + } else { + #expect(Bool(false), "Expected text content") + } + } + + @Test("Ergonomic API usage patterns") + func testErgonomicAPIUsagePatterns() throws { + // Test various ergonomic usage patterns enabled by the new API + + // Pattern 1: Simple conversation + let simpleConversation: [Sampling.Message] = [ + .user("What's the weather like?"), + .assistant("I'd be happy to help you check the weather!"), + .user("Thanks!"), + ] + #expect(simpleConversation.count == 3) + + // Pattern 2: Dynamic content with interpolation + let productName = "Smart Thermostat" + let price = 199.99 + let discount = 20 + + let salesConversation: [Sampling.Message] = [ + .user("Tell me about the \(productName)"), + .assistant("The \(productName) is priced at $\(String(format: "%.2f", price))"), + .user("Do you have any discounts?"), + .assistant( + "Yes! We currently have a \(discount)% discount, bringing the price to $\(String(format: "%.2f", price * (1.0 - Double(discount)/100.0)))" + ), + ] + #expect(salesConversation.count == 4) + + // Pattern 3: Mixed content types + let mixedContent: [Sampling.Message] = [ + .user("Can you analyze this image?"), + .assistant(.image(data: "analysis_chart_data", mimeType: "image/png")), + .user("What does it show?"), + .assistant("The chart shows a clear upward trend in sales."), + ] + #expect(mixedContent.count == 4) + + // Verify content types + if case .single(.text) = mixedContent[0].content, + case .single(.image) = mixedContent[1].content, + case .single(.text) = mixedContent[2].content, + case .single(.text) = mixedContent[3].content + { + // All content types are correct + } else { + #expect(Bool(false), "Content types don't match expected pattern") + } + + // Pattern 4: Encoding/decoding still works + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(simpleConversation) + let decoded = try decoder.decode([Sampling.Message].self, from: data) + + #expect(decoded.count == 3) + #expect(decoded[0].role == .user) + #expect(decoded[1].role == .assistant) + #expect(decoded[2].role == .user) + } +} + +@Suite("Sampling Integration Tests") +struct SamplingIntegrationTests { + + @Test + func testSamplingHandlerRegistration() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "SamplingHandlerTestServer", + version: "1.0" + ) + + let client = Client( + name: "SamplingHandlerTestClient", + version: "1.0" + ) + + nonisolated(unsafe) var handlerCalled = false + + // Register sampling handler + await client.withSamplingHandler { parameters in + handlerCalled = true + #expect(parameters.messages.count == 1) + + // 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.") + ) + } + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Test that the handler actually gets called + let messages: [Sampling.Message] = [.user("Test")] + let result = try await server.requestSampling(messages: messages, maxTokens: 100) + + #expect(handlerCalled, "Sampling handler should have been called") + #expect(result.model == "test-model-v1") + #expect(result.stopReason == .endTurn) + + await server.stop() + await client.disconnect() + } + + @Test + func testServerSamplingRequestAPI() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "SamplingRequestTestServer", + version: "1.0" + ) + + let client = Client( + name: "SamplingTestClient", + version: "1.0" + ) + + // Register sampling handler on client to respond to server's request + let responseContentString = "Based on the analysis, sales show strong growth through Q3 with Q4 stabilization." + await client.withSamplingHandler { parameters in + // Verify the request parameters were passed correctly + #expect(parameters.messages.count == 3) + #expect(parameters.systemPrompt == "You are a business analyst expert.") + #expect(parameters.includeContext == .thisServer) + #expect(parameters.temperature == 0.7) + #expect(parameters.maxTokens == 500) + #expect(parameters.stopSequences?.count == 2) + #expect(parameters._meta?["requestId"]?.stringValue == "test-123") + + // Return mock LLM response + return CreateSamplingMessage.Result( + model: "claude-4-sonnet", + stopReason: .endTurn, + role: .assistant, + content: .text(responseContentString) + ) + } + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // 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 works end-to-end + let result = 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---"], + _meta: Metadata(additionalFields: [ + "requestId": "test-123", + "priority": "high", + "department": "analytics", + ]) + ) + + // Verify the response + #expect(result.model == "claude-4-sonnet") + #expect(result.stopReason == .endTurn) + #expect(result.role == .assistant) + if case .single(.text(let text)) = result.content { + #expect(text == responseContentString) + } else { + Issue.record("Expected text content") + } + + await server.stop() + await client.disconnect() + } + + @Test + func testSamplingMessageTypes() async throws { + // Test comprehensive message content types + let textMessage: Sampling.Message = .user("What do you see in this data?") + + let imageMessage: Sampling.Message = .user( + .image( + data: + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==", + mimeType: "image/png" + )) + + // Test encoding/decoding of different message types + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + // Test text message + let textData = try encoder.encode(textMessage) + let decodedTextMessage = try decoder.decode(Sampling.Message.self, from: textData) + #expect(decodedTextMessage == textMessage) + + // Test image message + let imageData = try encoder.encode(imageMessage) + let decodedImageMessage = try decoder.decode(Sampling.Message.self, from: imageData) + #expect(decodedImageMessage == imageMessage) + } + + @Test + func testSamplingResultTypes() async throws { + // Test different result content types and stop reasons + let textResult = CreateSamplingMessage.Result( + model: "claude-4-sonnet", + stopReason: .endTurn, + role: .assistant, + content: .text( + "Based on the sales data analysis, I can see a strong upward trend through Q3, with a slight decline in Q4. This suggests seasonal factors or market saturation." + ) + ) + + let imageResult = CreateSamplingMessage.Result( + model: "dall-e-3", + stopReason: .maxTokens, + role: .assistant, + content: .image( + data: "generated_chart_base64_data_here", + mimeType: "image/png" + ) + ) + + let stopSequenceResult = CreateSamplingMessage.Result( + model: "gpt-4.1", + stopReason: .stopSequence, + role: .assistant, + content: .text("Analysis complete.\nEND_ANALYSIS") + ) + + // Test encoding/decoding of different result types + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + // Test text result + let textData = try encoder.encode(textResult) + let decodedTextResult = try decoder.decode( + CreateSamplingMessage.Result.self, from: textData + ) + #expect(decodedTextResult == textResult) + + // Test image result + let imageData = try encoder.encode(imageResult) + let decodedImageResult = try decoder.decode( + CreateSamplingMessage.Result.self, from: imageData + ) + #expect(decodedImageResult == imageResult) + + // Test stop sequence result + let stopData = try encoder.encode(stopSequenceResult) + let decodedStopResult = try decoder.decode( + CreateSamplingMessage.Result.self, from: stopData + ) + #expect(decodedStopResult == stopSequenceResult) + } + + @Test + func testSamplingErrorHandling() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "ErrorTestServer", + version: "1.0" + ) + + // Client WITHOUT sampling capability + let client = Client( + name: "ErrorTestClient", + version: "1.0" + ) + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Test sampling request - should fail because client doesn't support sampling + let messages: [Sampling.Message] = [ + .user("Test message") + ] + + await #expect(throws: MCPError.self) { + _ = try await server.requestSampling( + messages: messages, + maxTokens: 100 + ) + } + + await server.stop() + await client.disconnect() + } + + @Test("Strict mode succeeds when client declares sampling capability") + func testSamplingStrictCapabilitiesSuccess() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "StrictTestServer", + version: "1.0", + configuration: .strict + ) + + let client = Client( + name: "StrictTestClient", + version: "1.0", + capabilities: .init(sampling: .init()), + configuration: .strict + ) + + // Register sampling handler + await client.withSamplingHandler { _ in + CreateSamplingMessage.Result( + model: "test-model", + stopReason: .endTurn, + role: .assistant, + content: .text("Strict mode success") + ) + } + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Should succeed because client declares sampling capability + let result = try await server.requestSampling( + messages: [.user("Test message")], + maxTokens: 100 + ) + + #expect(result.model == "test-model") + #expect(result.role == .assistant) + if case .single(.text(let text)) = result.content { + #expect(text == "Strict mode success") + } else { + Issue.record("Expected text content") + } + + await server.stop() + await client.disconnect() + } + + @Test("Strict mode fails when client doesn't declare sampling capability") + func testSamplingStrictCapabilitiesError() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "StrictTestServer", + version: "1.0", + configuration: .strict + ) + + // Client WITHOUT sampling capability in strict mode + let client = Client( + name: "StrictTestClient", + version: "1.0", + capabilities: .init(), + configuration: .strict + ) + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Should fail because client doesn't declare sampling capability in strict mode + await #expect(throws: MCPError.self) { + _ = try await server.requestSampling( + messages: [.user("Test message")], + maxTokens: 100 + ) + } + + await server.stop() + await client.disconnect() + } + + @Test("Non-strict mode succeeds even without client capability declaration") + func testSamplingNonStrictCapabilities() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + + let server = Server( + name: "NonStrictTestServer", + version: "1.0", + configuration: .default // Non-strict mode + ) + + // Client WITHOUT sampling capability in non-strict mode + let client = Client( + name: "NonStrictTestClient", + version: "1.0", + capabilities: .init(), + configuration: .default + ) + + // Register sampling handler anyway + await client.withSamplingHandler { _ in + CreateSamplingMessage.Result( + model: "test-model", + stopReason: .endTurn, + role: .assistant, + content: .text("Non-strict mode success") + ) + } + + try await server.start(transport: serverTransport) + try await client.connect(transport: clientTransport) + + // Should succeed because server is in non-strict mode + let result = try await server.requestSampling( + messages: [.user("Test message")], + maxTokens: 100 + ) + + #expect(result.model == "test-model") + if case .single(.text(let text)) = result.content { + #expect(text == "Non-strict mode success") + } else { + Issue.record("Expected text content") + } + + await server.stop() + await client.disconnect() + } + + @Test + func testSamplingParameterValidation() async throws { + // Test parameter validation and edge cases + let validMessages: [Sampling.Message] = [ + .user("Valid message") + ] + + // Test with valid parameters + let validParams = CreateSamplingMessage.Parameters( + messages: validMessages, + maxTokens: 100 + ) + #expect(validParams.messages.count == 1) + #expect(validParams.maxTokens == 100) + + // Test with comprehensive parameters + let comprehensiveParams = CreateSamplingMessage.Parameters( + messages: validMessages, + modelPreferences: Sampling.ModelPreferences( + hints: [Sampling.ModelPreferences.Hint(name: "claude-4")], + costPriority: 0.5, + speedPriority: 0.8, + intelligencePriority: 0.9 + ), + systemPrompt: "You are a helpful assistant.", + includeContext: .allServers, + temperature: 0.7, + maxTokens: 500, + stopSequences: ["STOP", "END"], + _meta: Metadata(additionalFields: [ + "sessionId": "test-session-123", + "userId": "user-456", + ]) + ) + + #expect(comprehensiveParams.messages.count == 1) + #expect(comprehensiveParams.modelPreferences?.hints?.count == 1) + #expect(comprehensiveParams.systemPrompt == "You are a helpful assistant.") + #expect(comprehensiveParams.includeContext == .allServers) + #expect(comprehensiveParams.temperature == 0.7) + #expect(comprehensiveParams.maxTokens == 500) + #expect(comprehensiveParams.stopSequences?.count == 2) + #expect(comprehensiveParams._meta?.fields.count == 2) + + // Test encoding/decoding of comprehensive parameters + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(comprehensiveParams) + let decoded = try decoder.decode(CreateSamplingMessage.Parameters.self, from: data) + + #expect(decoded == comprehensiveParams) + } + + @Test + func testSamplingWorkflowScenarios() async throws { + // Test realistic sampling workflow scenarios + + // Scenario 1: Data Analysis Request + let dataAnalysisMessages: [Sampling.Message] = [ + .user("Please analyze the following customer feedback data:"), + .user( + """ + Feedback Summary: + - 85% positive sentiment + - Top complaints: shipping delays (12%), product quality (8%) + - Top praise: customer service (45%), product features (40%) + - NPS Score: 72 + """), + ] + + let dataAnalysisParams = CreateSamplingMessage.Parameters( + messages: dataAnalysisMessages, + modelPreferences: Sampling.ModelPreferences( + hints: [Sampling.ModelPreferences.Hint(name: "claude-4-sonnet")], + speedPriority: 0.3, + intelligencePriority: 0.9 + ), + systemPrompt: "You are an expert business analyst. Provide actionable insights.", + includeContext: .thisServer, + temperature: 0.3, // Lower temperature for analytical tasks + maxTokens: 400, + stopSequences: ["---END---"], + _meta: Metadata(additionalFields: ["analysisType": "customer-feedback"]) + ) + + // Scenario 2: Creative Content Generation + let creativeMessages: [Sampling.Message] = [ + .user( + "Write a compelling product description for a new smart home device.") + ] + + let creativeParams = CreateSamplingMessage.Parameters( + messages: creativeMessages, + modelPreferences: Sampling.ModelPreferences( + hints: [Sampling.ModelPreferences.Hint(name: "gpt-4.1")], + costPriority: 0.4, + speedPriority: 0.6, + intelligencePriority: 0.8 + ), + systemPrompt: "You are a creative marketing copywriter.", + temperature: 0.8, // Higher temperature for creativity + maxTokens: 200, + _meta: Metadata(additionalFields: ["contentType": "marketing-copy"]) + ) + + // Test parameter encoding for both scenarios + let encoder = JSONEncoder() + + let analysisData = try encoder.encode(dataAnalysisParams) + let creativeData = try encoder.encode(creativeParams) + + // Test decoding + let decoder = JSONDecoder() + let decodedAnalysis = try decoder.decode( + CreateSamplingMessage.Parameters.self, from: analysisData) + let decodedCreative = try decoder.decode( + CreateSamplingMessage.Parameters.self, from: creativeData) + + #expect(decodedAnalysis == dataAnalysisParams) + #expect(decodedCreative == creativeParams) + } +} + +@Suite("Sampling 2025-11-25 Spec Tests") +struct Sampling2025_11_25Tests { + @Test("Audio content encoding and decoding") + func testAudioContent() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let audioMessage: Sampling.Message = .user( + .audio(data: "base64audiodata", mimeType: "audio/mp3")) + + let data = try encoder.encode(audioMessage) + let decoded = try decoder.decode(Sampling.Message.self, from: data) + + #expect(decoded.role == .user) + if case .single(.audio(let audioData, let mimeType)) = decoded.content { + #expect(audioData == "base64audiodata") + #expect(mimeType == "audio/mp3") + } else { + #expect(Bool(false), "Expected audio content") + } + } + + @Test("StopReason.toolUse encoding and decoding") + func testToolUseStopReason() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let result = CreateSamplingMessage.Result( + model: "claude-4", + stopReason: .toolUse, + role: .assistant, + content: .single(.text("I need to use a tool")) + ) + + let data = try encoder.encode(result) + let decoded = try decoder.decode(CreateSamplingMessage.Result.self, from: data) + + #expect(decoded == result) + } + + @Test("Multiple content blocks") + func testMultipleContentBlocks() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let blocks: [Sampling.Message.Content.ContentBlock] = [ + .text("Here's an image:"), + .image(data: "imagedata", mimeType: "image/png"), + .text("And some audio:"), + .audio(data: "audiodata", mimeType: "audio/mp3") + ] + + let content = Sampling.Message.Content.multiple(blocks) + + let message = Sampling.Message.assistant(content) + let data = try encoder.encode(message) + let decoded = try decoder.decode(Sampling.Message.self, from: data) + + #expect(decoded == message) + } + + @Test("Tools parameter encoding and decoding") + func testToolsParameter() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let tool = Tool( + name: "get_weather", + description: "Get current weather", + inputSchema: [ + "type": "object", + "properties": [ + "location": ["type": "string"] + ] + ] + ) + + let params = CreateSamplingMessage.Parameters( + messages: [.user("What's the weather?")], + maxTokens: 100, + tools: [tool], + toolChoice: .init(mode: .auto) + ) + + let data = try encoder.encode(params) + let decoded = try decoder.decode(CreateSamplingMessage.Parameters.self, from: data) + + #expect(decoded == params) + } + + @Test("Client sampling capabilities with sub-capabilities") + func testSamplingSubCapabilities() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let capabilities = Client.Capabilities( + sampling: .init(tools: .init(), context: .init()) + ) + + #expect(capabilities.sampling?.tools != nil) + #expect(capabilities.sampling?.context != nil) + + let data = try encoder.encode(capabilities) + let decoded = try decoder.decode(Client.Capabilities.self, from: data) + + #expect(decoded == capabilities) + } +} diff --git a/Tests/MCPTests/ServerTests.swift b/Tests/MCPTests/ServerTests.swift index e733c01..a7b6166 100644 --- a/Tests/MCPTests/ServerTests.swift +++ b/Tests/MCPTests/ServerTests.swift @@ -22,8 +22,8 @@ struct ServerTests { let transport = MockTransport() // Queue an initialize request - try await transport.queueRequest( - Initialize.request( + try await transport.queue( + request: Initialize.request( .init( protocolVersion: Version.latest, capabilities: .init(), @@ -39,13 +39,181 @@ struct ServerTests { try await server.start(transport: transport) // Wait for message processing and response - try await Task.sleep(nanoseconds: 10_000_000) // 10ms + try await Task.sleep(for: .milliseconds(200)) #expect(await transport.sentMessages.count == 1) - #expect(await transport.sentMessages[0].contains(Initialize.name)) + + let messages = await transport.sentMessages + if let response = messages.first { + #expect(response.contains("serverInfo")) + } // Clean up await server.stop() await transport.disconnect() } + + @Test("Initialize hook - successful") + func testInitializeHookSuccess() async throws { + let transport = MockTransport() + + actor TestState { + var hookCalled = false + func setHookCalled() { hookCalled = true } + func wasHookCalled() -> Bool { hookCalled } + } + + let state = TestState() + let server = Server(name: "TestServer", version: "1.0") + + // Start with the hook directly + try await server.start(transport: transport) { clientInfo, capabilities in + #expect(clientInfo.name == "TestClient") + #expect(clientInfo.version == "1.0") + await state.setHookCalled() + } + + // Wait for server to initialize + try await Task.sleep(for: .milliseconds(10)) + + // Queue an initialize request + try await transport.queue( + request: Initialize.request( + .init( + protocolVersion: Version.latest, + capabilities: .init(), + clientInfo: .init(name: "TestClient", version: "1.0") + ) + )) + + // Wait for message processing and hook execution + try await Task.sleep(for: .milliseconds(500)) + + #expect(await state.wasHookCalled() == true) + #expect(await transport.sentMessages.count >= 1) + + let messages = await transport.sentMessages + if let response = messages.first { + #expect(response.contains("serverInfo")) + } + + await server.stop() + await transport.disconnect() + } + + @Test("Initialize hook - rejection") + func testInitializeHookRejection() async throws { + let transport = MockTransport() + + let server = Server(name: "TestServer", version: "1.0") + + try await server.start(transport: transport) { clientInfo, _ in + if clientInfo.name == "BlockedClient" { + throw MCPError.invalidRequest("Client not allowed") + } + } + + // Wait for server to initialize + try await Task.sleep(for: .milliseconds(10)) + + // Queue an initialize request from blocked client + try await transport.queue( + request: Initialize.request( + .init( + protocolVersion: Version.latest, + capabilities: .init(), + clientInfo: .init(name: "BlockedClient", version: "1.0") + ) + )) + + // Wait for message processing + try await Task.sleep(for: .milliseconds(200)) + + #expect(await transport.sentMessages.count >= 1) + + let messages = await transport.sentMessages + if let response = messages.first { + #expect(response.contains("error")) + #expect(response.contains("Client not allowed")) + } + + await server.stop() + await transport.disconnect() + } + + @Test("JSON-RPC batch processing") + func testJSONRPCBatchProcessing() async throws { + let (clientTransport, serverTransport) = await InMemoryTransport.createConnectedPair() + let server = Server(name: "TestServer", version: "1.0") + + // Connect transports + try await clientTransport.connect() + try await serverTransport.connect() + + // Start receiving messages on client side + let receiveTask = Task { + var responses: [String] = [] + for try await data in await clientTransport.receive() { + if let response = String(data: data, encoding: .utf8) { + responses.append(response) + } + // Stop after receiving 2 responses (initialize + batch) + if responses.count == 2 { + break + } + } + return responses + } + + // Start the server + try await server.start(transport: serverTransport) + + // Initialize the server first + let initRequest = Initialize.request( + .init( + protocolVersion: Version.latest, + capabilities: .init(), + clientInfo: .init(name: "TestClient", version: "1.0") + ) + ) + let initData = try JSONEncoder().encode(AnyRequest(initRequest)) + try await clientTransport.send(initData) + + // Wait for initialization + try await Task.sleep(for: .milliseconds(100)) + + // Create a batch with multiple requests + let batchJSON = """ + [ + {"jsonrpc":"2.0","id":1,"method":"ping","params":{}}, + {"jsonrpc":"2.0","id":2,"method":"ping","params":{}} + ] + """ + let batchData = batchJSON.data(using: .utf8)! + try await clientTransport.send(batchData) + + // Wait for batch processing + try await Task.sleep(for: .milliseconds(200)) + + // Get responses + let responses = try await receiveTask.value + #expect(responses.count == 2) + + // Verify the batch response (second response) + if responses.count >= 2 { + let batchResponse = responses[1] + + // Should be an array + #expect(batchResponse.hasPrefix("[")) + #expect(batchResponse.hasSuffix("]")) + + // Should contain both request IDs + #expect(batchResponse.contains("\"id\":1")) + #expect(batchResponse.contains("\"id\":2")) + } + + await server.stop() + await clientTransport.disconnect() + await serverTransport.disconnect() + } } diff --git a/Tests/MCPTests/TransportTests.swift b/Tests/MCPTests/StdioTransportTests.swift similarity index 88% rename from Tests/MCPTests/TransportTests.swift rename to Tests/MCPTests/StdioTransportTests.swift index c56f646..e395fbe 100644 --- a/Tests/MCPTests/TransportTests.swift +++ b/Tests/MCPTests/StdioTransportTests.swift @@ -1,9 +1,14 @@ import Foundation -import SystemPackage import Testing @testable import MCP +#if canImport(System) + import System +#else + @preconcurrency import SystemPackage +#endif + @Suite("Stdio Transport Tests") struct StdioTransportTests { @Test("Connection") @@ -24,7 +29,7 @@ struct StdioTransportTests { // Test sending a simple message let message = #"{"key":"value"}"# - try await transport.send(message) + try await transport.send(message.data(using: .utf8)!) // Read and verify the output var buffer = [UInt8](repeating: 0, count: 1024) @@ -52,12 +57,12 @@ struct StdioTransportTests { try writer.close() // Start receiving messages - let stream: AsyncThrowingStream = await transport.receive() + let stream: AsyncThrowingStream = await transport.receive() var iterator = stream.makeAsyncIterator() // Get first message let received = try await iterator.next() - #expect(received == #"{"key":"value"}"#) + #expect(received == #"{"key":"value"}"#.data(using: .utf8)!) await transport.disconnect() } @@ -74,7 +79,7 @@ struct StdioTransportTests { try writer.writeAll(invalidJSON.data(using: .utf8)!) try writer.close() - let stream: AsyncThrowingStream = await transport.receive() + let stream: AsyncThrowingStream = await transport.receive() var iterator = stream.makeAsyncIterator() _ = try await iterator.next() @@ -95,7 +100,7 @@ struct StdioTransportTests { try await transport.connect() #expect(Bool(false), "Expected connect to throw an error") } catch { - #expect(error is MCP.Error) + #expect(error is MCPError) } await transport.disconnect() @@ -114,7 +119,7 @@ struct StdioTransportTests { try await transport.connect() #expect(Bool(false), "Expected connect to throw an error") } catch { - #expect(error is MCP.Error) + #expect(error is MCPError) } await transport.disconnect() diff --git a/Tests/MCPTests/ToolTests.swift b/Tests/MCPTests/ToolTests.swift index a9a03fc..d91fb49 100644 --- a/Tests/MCPTests/ToolTests.swift +++ b/Tests/MCPTests/ToolTests.swift @@ -11,13 +11,173 @@ struct ToolTests { name: "test_tool", description: "A test tool", inputSchema: .object([ - "param1": .string("Test parameter") + "properties": .object([ + "param1": .string("Test parameter") + ]) ]) ) #expect(tool.name == "test_tool") #expect(tool.description == "A test tool") #expect(tool.inputSchema != nil) + #expect(tool.title == nil) + #expect(tool.outputSchema == nil) + } + + @Test("Tool Annotations initialization and properties") + func testToolAnnotationsInitialization() throws { + // Empty annotations + let emptyAnnotations = Tool.Annotations() + #expect(emptyAnnotations.isEmpty) + #expect(emptyAnnotations.title == nil) + #expect(emptyAnnotations.readOnlyHint == nil) + #expect(emptyAnnotations.destructiveHint == nil) + #expect(emptyAnnotations.idempotentHint == nil) + #expect(emptyAnnotations.openWorldHint == nil) + + // Full annotations + let fullAnnotations = Tool.Annotations( + title: "Test Tool", + readOnlyHint: true, + destructiveHint: false, + idempotentHint: true, + openWorldHint: false + ) + + #expect(!fullAnnotations.isEmpty) + #expect(fullAnnotations.title == "Test Tool") + #expect(fullAnnotations.readOnlyHint == true) + #expect(fullAnnotations.destructiveHint == false) + #expect(fullAnnotations.idempotentHint == true) + #expect(fullAnnotations.openWorldHint == false) + + // Partial annotations - should not be empty + let partialAnnotations = Tool.Annotations(title: "Partial Test") + #expect(!partialAnnotations.isEmpty) + #expect(partialAnnotations.title == "Partial Test") + + // Initialize with nil literal + let nilAnnotations: Tool.Annotations = nil + #expect(nilAnnotations.isEmpty) + } + + @Test("Tool Annotations encoding and decoding") + func testToolAnnotationsEncodingDecoding() throws { + let annotations = Tool.Annotations( + title: "Test Tool", + readOnlyHint: true, + destructiveHint: false, + idempotentHint: true, + openWorldHint: false + ) + + #expect(!annotations.isEmpty) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(annotations) + let decoded = try decoder.decode(Tool.Annotations.self, from: data) + + #expect(decoded.title == annotations.title) + #expect(decoded.readOnlyHint == annotations.readOnlyHint) + #expect(decoded.destructiveHint == annotations.destructiveHint) + #expect(decoded.idempotentHint == annotations.idempotentHint) + #expect(decoded.openWorldHint == annotations.openWorldHint) + + // Test that empty annotations are encoded as expected + let emptyAnnotations = Tool.Annotations() + let emptyData = try encoder.encode(emptyAnnotations) + let decodedEmpty = try decoder.decode(Tool.Annotations.self, from: emptyData) + + #expect(decodedEmpty.isEmpty) + } + + @Test("Tool with annotations encoding and decoding") + func testToolWithAnnotationsEncodingDecoding() throws { + let annotations = Tool.Annotations( + title: "Calculator", + destructiveHint: false + ) + + let tool = Tool( + name: "calculate", + description: "Performs calculations", + inputSchema: .object([ + "properties": .object([ + "expression": .string("Mathematical expression to evaluate") + ]) + ]), + annotations: annotations + ) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(tool) + let decoded = try decoder.decode(Tool.self, from: data) + + #expect(decoded.name == tool.name) + #expect(decoded.description == tool.description) + #expect(decoded.annotations.title == annotations.title) + #expect(decoded.annotations.destructiveHint == annotations.destructiveHint) + + // Verify that the annotations field is properly included in the JSON + let jsonString = String(data: data, encoding: .utf8)! + #expect(jsonString.contains("\"annotations\"")) + #expect(jsonString.contains("\"title\":\"Calculator\"")) + } + + @Test("Tool with empty annotations") + func testToolWithEmptyAnnotations() throws { + var tool = Tool( + name: "test_tool", + description: "Test tool description", + inputSchema: [:] + ) + + do { + #expect(tool.annotations.isEmpty) + + let encoder = JSONEncoder() + let data = try encoder.encode(tool) + + // Verify that empty annotations are not included in the JSON + let jsonString = String(data: data, encoding: .utf8)! + #expect(!jsonString.contains("\"annotations\"")) + } + + do { + tool.annotations.title = "Test" + + #expect(!tool.annotations.isEmpty) + + let encoder = JSONEncoder() + let data = try encoder.encode(tool) + + // Verify that empty annotations are not included in the JSON + let jsonString = String(data: data, encoding: .utf8)! + #expect(jsonString.contains("\"annotations\"")) + } + } + + @Test("Tool with nil literal annotations") + func testToolWithNilLiteralAnnotations() throws { + let tool = Tool( + name: "test_tool", + description: "Test tool description", + inputSchema: [:], + annotations: nil + ) + + #expect(tool.annotations.isEmpty) + + let encoder = JSONEncoder() + let data = try encoder.encode(tool) + + // Verify that nil literal annotations are not included in the JSON + let jsonString = String(data: data, encoding: .utf8)! + #expect(!jsonString.contains("\"annotations\"")) } @Test("Tool encoding and decoding") @@ -26,8 +186,10 @@ struct ToolTests { name: "test_tool", description: "Test tool description", inputSchema: .object([ - "param1": .string("String parameter"), - "param2": .int(42), + "properties": .object([ + "param1": .string("String parameter"), + "param2": .int(42), + ]) ]) ) @@ -42,16 +204,50 @@ struct ToolTests { #expect(decoded.inputSchema == tool.inputSchema) } + @Test("Tool encoding and decoding with title and output schema") + func testToolEncodingDecodingWithTitleAndOutputSchema() throws { + let tool = Tool( + name: "test_tool", + title: "Readable Test Tool", + description: "Test tool description", + inputSchema: .object([ + "type": .string("object"), + "properties": .object([ + "param1": .string("String parameter") + ]), + ]), + outputSchema: .object([ + "type": .string("object"), + "properties": .object([ + "result": .string("String result") + ]), + ]) + ) + + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(tool) + let decoded = try decoder.decode(Tool.self, from: data) + + #expect(decoded.title == tool.title) + #expect(decoded.outputSchema == tool.outputSchema) + + let jsonString = String(decoding: data, as: UTF8.self) + #expect(jsonString.contains("\"title\":\"Readable Test Tool\"")) + #expect(jsonString.contains("\"outputSchema\"")) + } + @Test("Text content encoding and decoding") func testToolContentTextEncoding() throws { - let content = Tool.Content.text("Hello, world!") + let content = Tool.Content.text(text: "Hello, world!", annotations: nil, _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 .text(let text) = decoded { + if case .text(let text, _, _) = decoded { #expect(text == "Hello, world!") } else { #expect(Bool(false), "Expected text content") @@ -63,7 +259,8 @@ struct ToolTests { let content = Tool.Content.image( data: "base64data", mimeType: "image/png", - metadata: ["width": "100", "height": "100"] + annotations: nil, + _meta: nil ) let encoder = JSONEncoder() let decoder = JSONDecoder() @@ -71,11 +268,9 @@ struct ToolTests { 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") } @@ -83,10 +278,67 @@ struct ToolTests { @Test("Resource content encoding and decoding") func testToolContentResourceEncoding() throws { - let content = Tool.Content.resource( + let resourceContent = Resource.Content.text( + "Sample text", uri: "file://test.txt", - mimeType: "text/plain", - text: "Sample text" + mimeType: "text/plain" + ) + let content = Tool.Content.resource(resource: resourceContent, annotations: nil, _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 resource, let annotations, let _meta) = decoded { + #expect(resource.uri == "file://test.txt") + #expect(resource.mimeType == "text/plain") + #expect(resource.text == "Sample text") + #expect(annotations == nil) + #expect(_meta == nil) + } else { + #expect(Bool(false), "Expected resource content") + } + } + + @Test("Resource link content includes title") + func testToolContentResourceLinkEncoding() throws { + let content = Tool.Content.resourceLink( + uri: "file://resource.txt", + name: "resource_name", + title: "Resource Title", + description: "Resource description", + mimeType: "text/plain" + ) + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let data = try encoder.encode(content) + let jsonString = String(decoding: data, as: UTF8.self) + #expect(jsonString.contains("\"title\":\"Resource Title\"")) + + let decoded = try decoder.decode(Tool.Content.self, from: data) + if case .resourceLink( + let uri, let name, let title, let description, let mimeType, let annotations + ) = decoded { + #expect(uri == "file://resource.txt") + #expect(name == "resource_name") + #expect(title == "Resource Title") + #expect(description == "Resource description") + #expect(mimeType == "text/plain") + #expect(annotations == nil) + } else { + #expect(Bool(false), "Expected resourceLink content") + } + } + + @Test("Audio content encoding and decoding") + func testToolContentAudioEncoding() throws { + let content = Tool.Content.audio( + data: "base64audiodata", + mimeType: "audio/wav", + annotations: nil, + _meta: nil ) let encoder = JSONEncoder() let decoder = JSONDecoder() @@ -94,15 +346,79 @@ 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 .audio(let data, let mimeType, _, _) = decoded { + #expect(data == "base64audiodata") + #expect(mimeType == "audio/wav") } else { - #expect(Bool(false), "Expected resource content") + #expect(Bool(false), "Expected audio content") } } + @Test("Text content encoding includes annotations and meta fields") + func testToolContentTextEncodingIncludesAnnotationsAndMeta() throws { + let annotations = Resource.Annotations( + audience: [.assistant], + priority: 0.75, + lastModified: "2026-01-01T00:00:00Z" + ) + let meta = Metadata(additionalFields: [ + "vendor.example/source": .string("tool-text") + ]) + let expected = Tool.Content.text( + text: "Hello with metadata", + annotations: annotations, + _meta: meta + ) + + let data = try JSONEncoder().encode(expected) + let decoded = try JSONDecoder().decode(Tool.Content.self, from: data) + #expect(decoded == expected) + } + + @Test("Image content encoding includes annotations and meta fields") + func testToolContentImageEncodingIncludesAnnotationsAndMeta() throws { + let annotations = Resource.Annotations( + audience: [.user], + priority: 0.9, + lastModified: "2026-01-02T00:00:00Z" + ) + let meta = Metadata(additionalFields: [ + "vendor.example/source": .string("tool-image") + ]) + let expected = Tool.Content.image( + data: "base64-image", + mimeType: "image/jpeg", + annotations: annotations, + _meta: meta + ) + + let data = try JSONEncoder().encode(expected) + let decoded = try JSONDecoder().decode(Tool.Content.self, from: data) + #expect(decoded == expected) + } + + @Test("Audio content encoding includes annotations and meta fields") + func testToolContentAudioEncodingIncludesAnnotationsAndMeta() throws { + let annotations = Resource.Annotations( + audience: [.assistant], + priority: 0.4, + lastModified: "2026-01-03T00:00:00Z" + ) + let meta = Metadata(additionalFields: [ + "vendor.example/source": .string("tool-audio") + ]) + let expected = Tool.Content.audio( + data: "base64-audio", + mimeType: "audio/mp3", + annotations: annotations, + _meta: meta + ) + + let data = try JSONEncoder().encode(expected) + let decoded = try JSONDecoder().decode(Tool.Content.self, from: data) + #expect(decoded == expected) + } + @Test("ListTools parameters validation") func testListToolsParameters() throws { let params = ListTools.Parameters(cursor: "next_page") @@ -111,7 +427,7 @@ struct ToolTests { let emptyParams = ListTools.Parameters() #expect(emptyParams.cursor == nil) } - + @Test("ListTools request decoding with omitted params") func testListToolsRequestDecodingWithOmittedParams() throws { // Test decoding when params field is omitted @@ -126,7 +442,7 @@ struct ToolTests { #expect(decoded.id == "test-id") #expect(decoded.method == ListTools.name) } - + @Test("ListTools request decoding with null params") func testListToolsRequestDecodingWithNullParams() throws { // Test decoding when params field is null @@ -145,8 +461,8 @@ struct ToolTests { @Test("ListTools result validation") func testListToolsResult() throws { let tools = [ - Tool(name: "tool1", description: "First tool", inputSchema: nil), - Tool(name: "tool2", description: "Second tool", inputSchema: nil), + Tool(name: "tool1", description: "First tool", inputSchema: [:]), + Tool(name: "tool2", description: "Second tool", inputSchema: [:]), ] let result = ListTools.Result(tools: tools, nextCursor: "next_page") @@ -172,15 +488,15 @@ struct ToolTests { @Test("CallTool success result validation") func testCallToolResult() throws { let content = [ - Tool.Content.text("Result 1"), - Tool.Content.text("Result 2"), + Tool.Content.text(text: "Result 1", annotations: nil, _meta: nil), + Tool.Content.text(text: "Result 2", annotations: nil, _meta: nil), ] let result = CallTool.Result(content: content) #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") @@ -189,12 +505,12 @@ struct ToolTests { @Test("CallTool error result validation") func testCallToolErrorResult() throws { - let errorContent = [Tool.Content.text("Error message")] + let errorContent = [Tool.Content.text(text: "Error message", annotations: nil, _meta: nil)] let errorResult = CallTool.Result(content: errorContent, isError: true) #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") @@ -205,4 +521,58 @@ struct ToolTests { func testToolListChangedNotification() throws { #expect(ToolListChangedNotification.name == "notifications/tools/list_changed") } + + @Test("ListTools handler invocation without params") + func testListToolsHandlerWithoutParams() async throws { + let jsonString = """ + {"jsonrpc":"2.0","id":1,"method":"tools/list"} + """ + let jsonData = jsonString.data(using: .utf8)! + + let anyRequest = try JSONDecoder().decode(AnyRequest.self, from: jsonData) + + let handler = TypedRequestHandler { request in + #expect(request.method == ListTools.name) + #expect(request.id == 1) + #expect(request.params.cursor == nil) + + let testTool = Tool( + name: "test_tool", + description: "Test tool for verification", + inputSchema: [:] + ) + return ListTools.response(id: request.id, result: ListTools.Result(tools: [testTool])) + } + + let response = try await handler(anyRequest) + + if case .success(let value) = response.result { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + let data = try encoder.encode(value) + let result = try decoder.decode(ListTools.Result.self, from: data) + + #expect(result.tools.count == 1) + #expect(result.tools[0].name == "test_tool") + } else { + #expect(Bool(false), "Expected success result") + } + } + + @Test("Tool with missing description") + func testToolWithMissingDescription() throws { + let jsonString = """ + { + "name": "test_tool", + "inputSchema": {} + } + """ + 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 == [:]) + } } diff --git a/Tests/MCPTests/UnitIntervalTests.swift b/Tests/MCPTests/UnitIntervalTests.swift new file mode 100644 index 0000000..0269251 --- /dev/null +++ b/Tests/MCPTests/UnitIntervalTests.swift @@ -0,0 +1,219 @@ +import Testing + +import class Foundation.JSONDecoder +import class Foundation.JSONEncoder + +@testable import MCP + +@Suite("UnitInterval Tests") +struct UnitIntervalTests { + @Test("Valid literal initialization") + func testValidLiteralInitialization() throws { + let zero: UnitInterval = 0.0 + #expect(zero.doubleValue == 0.0) + + let half: UnitInterval = 0.5 + #expect(half.doubleValue == 0.5) + + let one: UnitInterval = 1.0 + #expect(one.doubleValue == 1.0) + + let quarter: UnitInterval = 0.25 + #expect(quarter.doubleValue == 0.25) + } + + @Test("Valid failable initialization with runtime values") + func testValidFailableInitialization() throws { + // Test with runtime computed values to force use of failable initializer + let values = [0.0, 0.5, 1.0, 0.25] + + for value in values { + let computed = value * 1.0 // Force runtime computation + let interval = UnitInterval(computed) + #expect(interval != nil) + #expect(interval!.doubleValue == value) + } + } + + @Test("Invalid failable initialization") + func testInvalidFailableInitialization() throws { + // Test with runtime computed values to force use of failable initializer + let invalidValues = [-0.1, 1.1, 100.0, -100.0] + + for value in invalidValues { + let computed = value * 1.0 // Force runtime computation + let interval = UnitInterval(computed) + #expect(interval == nil) + } + } + + @Test("Boundary and edge case values") + func testBoundaryAndEdgeCaseValues() throws { + // Test exact boundary values + let exactZero = 0.0 * 1.0 + let zero = UnitInterval(exactZero) + #expect(zero != nil) + + let exactOne = 1.0 * 1.0 + let one = UnitInterval(exactOne) + #expect(one != nil) + + // Test machine precision boundaries + let justAboveZero = Double.ulpOfOne * 1.0 + let aboveZero = UnitInterval(justAboveZero) + #expect(aboveZero != nil) + + let justBelowOne = (1.0 - Double.ulpOfOne) * 1.0 + let belowOne = UnitInterval(justBelowOne) + #expect(belowOne != nil) + + // Test very small positive value + let tinyValue = 1e-10 * 1.0 + let tiny = UnitInterval(tinyValue) + #expect(tiny != nil) + #expect(tiny!.doubleValue == 1e-10) + + // Test value very close to 1 + let almostOneValue = 0.9999999999 * 1.0 + let almostOne = UnitInterval(almostOneValue) + #expect(almostOne != nil) + #expect(almostOne!.doubleValue == 0.9999999999) + } + + @Test("Float literal initialization") + func testFloatLiteralInitialization() throws { + let zero: UnitInterval = 0.0 + #expect(zero.doubleValue == 0.0) + + let half: UnitInterval = 0.5 + #expect(half.doubleValue == 0.5) + + let one: UnitInterval = 1.0 + #expect(one.doubleValue == 1.0) + + let quarter: UnitInterval = 0.25 + #expect(quarter.doubleValue == 0.25) + } + + @Test("Integer literal initialization") + func testIntegerLiteralInitialization() throws { + let zero: UnitInterval = 0 + #expect(zero.doubleValue == 0.0) + + let one: UnitInterval = 1 + #expect(one.doubleValue == 1.0) + } + + @Test("Comparable conformance") + func testComparable() throws { + let zero: UnitInterval = 0.0 + let quarter: UnitInterval = 0.25 + let half: UnitInterval = 0.5 + let one: UnitInterval = 1.0 + + #expect(zero < quarter) + #expect(quarter < half) + #expect(half < one) + #expect(zero < one) + + #expect(!(quarter < zero)) + #expect(!(half < quarter)) + #expect(!(one < half)) + + #expect(zero <= quarter) + #expect(quarter <= half) + #expect(half <= one) + #expect(zero <= zero) + + #expect(quarter > zero) + #expect(half > quarter) + #expect(one > half) + + #expect(quarter >= zero) + #expect(half >= quarter) + #expect(one >= half) + #expect(one >= one) + } + + @Test("Equality and hashing") + func testEqualityAndHashing() throws { + let half1: UnitInterval = 0.5 + let half2: UnitInterval = 0.5 + let quarter: UnitInterval = 0.25 + + #expect(half1 == half2) + #expect(half1 != quarter) + #expect(half1.hashValue == half2.hashValue) + } + + @Test("String description") + func testStringDescription() throws { + let zero: UnitInterval = 0.0 + #expect(zero.description == "0.0") + + let half: UnitInterval = 0.5 + #expect(half.description == "0.5") + + let one: UnitInterval = 1.0 + #expect(one.description == "1.0") + + let quarter: UnitInterval = 0.25 + #expect(quarter.description == "0.25") + } + + @Test("JSON encoding and decoding") + func testJSONCodable() throws { + let encoder = JSONEncoder() + let decoder = JSONDecoder() + + let original: UnitInterval = 0.75 + + let data = try encoder.encode(original) + let decoded = try decoder.decode(UnitInterval.self, from: data) + + #expect(decoded == original) + #expect(decoded.doubleValue == 0.75) + } + + @Test("JSON decoding with invalid values") + func testJSONDecodingInvalidValues() throws { + let decoder = JSONDecoder() + + // Test negative value + let negativeJSON = "-0.5".data(using: .utf8)! + #expect(throws: DecodingError.self) { + try decoder.decode(UnitInterval.self, from: negativeJSON) + } + + // Test value greater than 1 + let tooLargeJSON = "1.5".data(using: .utf8)! + #expect(throws: DecodingError.self) { + try decoder.decode(UnitInterval.self, from: tooLargeJSON) + } + } + + @Test("JSON encoding produces expected format") + func testJSONEncodingFormat() throws { + let encoder = JSONEncoder() + + let half: UnitInterval = 0.5 + let data = try encoder.encode(half) + let jsonString = String(data: data, encoding: .utf8) + + #expect(jsonString == "0.5") + } + + @Test("Double value property") + func testDoubleValueProperty() throws { + let values = [0.0, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0] + + for value in values { + let computed = value * 1.0 // Force runtime computation + if let interval = UnitInterval(computed) { + #expect(interval.doubleValue == value) + } else { + #expect(Bool(false), "UnitInterval(\(value)) should succeed") + } + } + } +} diff --git a/Tests/MCPTests/VersioningTests.swift b/Tests/MCPTests/VersioningTests.swift new file mode 100644 index 0000000..2071f90 --- /dev/null +++ b/Tests/MCPTests/VersioningTests.swift @@ -0,0 +1,62 @@ +import Testing + +@testable import MCP + +@Suite("Version Negotiation Tests") +struct VersioningTests { + @Test("Client requests latest supported version") + func testClientRequestsLatestSupportedVersion() { + let clientVersion = Version.latest + let negotiatedVersion = Version.negotiate(clientRequestedVersion: clientVersion) + #expect(negotiatedVersion == Version.latest) + } + + @Test("Client requests older supported version") + func testClientRequestsOlderSupportedVersion() { + let clientVersion = "2024-11-05" + let negotiatedVersion = Version.negotiate(clientRequestedVersion: clientVersion) + #expect(negotiatedVersion == "2024-11-05") + } + + @Test("Client requests unsupported version") + func testClientRequestsUnsupportedVersion() { + let clientVersion = "2023-01-01" // An unsupported version + let negotiatedVersion = Version.negotiate(clientRequestedVersion: clientVersion) + #expect(negotiatedVersion == Version.latest) + } + + @Test("Client requests empty version string") + func testClientRequestsEmptyVersionString() { + let clientVersion = "" + let negotiatedVersion = Version.negotiate(clientRequestedVersion: clientVersion) + #expect(negotiatedVersion == Version.latest) + } + + @Test("Client requests garbage version string") + func testClientRequestsGarbageVersionString() { + let clientVersion = "not-a-version" + let negotiatedVersion = Version.negotiate(clientRequestedVersion: clientVersion) + #expect(negotiatedVersion == Version.latest) + } + + @Test("Server's supported versions correctly defined") + func testServerSupportedVersions() { + #expect(Version.supported.contains("2025-11-25")) + #expect(Version.supported.contains("2025-06-18")) + #expect(Version.supported.contains("2025-03-26")) + #expect(Version.supported.contains("2024-11-05")) + #expect(Version.supported.count == 4) + } + + @Test("Server's latest version is correct") + func testServerLatestVersion() { + #expect(Version.latest == "2025-11-25") + } + + @Test("Client requests new 2025-11-25 version") + func testClientRequests2025_11_25Version() { + let clientVersion = "2025-11-25" + let negotiatedVersion = Version.negotiate(clientRequestedVersion: clientVersion) + #expect(negotiatedVersion == "2025-11-25") + } +} diff --git a/conformance-baseline.yml b/conformance-baseline.yml new file mode 100644 index 0000000..7127a6b --- /dev/null +++ b/conformance-baseline.yml @@ -0,0 +1 @@ +client: [] diff --git a/scripts/run-conformance.sh b/scripts/run-conformance.sh new file mode 100755 index 0000000..6812461 --- /dev/null +++ b/scripts/run-conformance.sh @@ -0,0 +1,113 @@ +#!/bin/bash +set -euo pipefail + +# Color output helpers +RED='\033[0;31m' +GREEN='\033[0;32m' +YELLOW='\033[1;33m' +NC='\033[0m' + +log_info() { echo -e "${GREEN}[INFO]${NC} $*"; } +log_warn() { echo -e "${YELLOW}[WARN]${NC} $*"; } +log_error() { echo -e "${RED}[ERROR]${NC} $*"; } + +# Configuration +CONFORMANCE_PKG="@modelcontextprotocol/conformance" +CLIENT_EXEC="mcp-everything-client" +SERVER_EXEC="mcp-everything-server" +BASELINE_FILE="${BASELINE_FILE:-conformance-baseline.yml}" +MODE="${MODE:-both}" + +# Parse arguments +while [[ $# -gt 0 ]]; do + case $1 in + --mode) + MODE="$2" + shift 2 + ;; + --baseline) + BASELINE_FILE="$2" + shift 2 + ;; + *) + log_error "Unknown option: $1" + exit 1 + ;; + esac +done + +# Validate mode +if [[ ! "$MODE" =~ ^(client|server|both)$ ]]; then + log_error "Invalid mode: $MODE. Must be one of: client, server, both" + exit 1 +fi + +# Build Swift executables +log_info "Building Swift executables..." +swift build --product "$CLIENT_EXEC" || { + log_error "Failed to build client" + exit 1 +} +swift build --product "$SERVER_EXEC" || { + log_error "Failed to build server" + exit 1 +} + +CLIENT_PATH="$(swift build --show-bin-path)/$CLIENT_EXEC" +SERVER_PATH="$(swift build --show-bin-path)/$SERVER_EXEC" + +log_info "Client executable: $CLIENT_PATH" +log_info "Server executable: $SERVER_PATH" + +# Check for baseline file +BASELINE_ARG="" +if [[ -f "$BASELINE_FILE" ]]; then + log_info "Using baseline file: $BASELINE_FILE" + BASELINE_ARG="--expected-failures $BASELINE_FILE" +else + log_warn "No baseline file found at $BASELINE_FILE" +fi + +# Run client tests +if [[ "$MODE" == "client" || "$MODE" == "both" ]]; then + log_info "Running client conformance tests..." + npx "$CONFORMANCE_PKG" client \ + --command "$CLIENT_PATH" \ + --suite core \ + $BASELINE_ARG || { + log_error "Client conformance tests failed" + exit 1 + } + log_info "Client tests completed" +fi + +# Run server tests +if [[ "$MODE" == "server" || "$MODE" == "both" ]]; then + log_info "Starting server for conformance testing..." + + # Start server in background + "$SERVER_PATH" & + SERVER_PID=$! + + # Wait for server to be ready + log_info "Waiting for server to start (PID: $SERVER_PID)..." + sleep 3 + + # Run server tests + log_info "Running server conformance tests..." + npx "$CONFORMANCE_PKG" server \ + --url http://localhost:3001/mcp \ + --suite all \ + $BASELINE_ARG || { + log_error "Server conformance tests failed" + kill $SERVER_PID 2>/dev/null || true + exit 1 + } + + # Cleanup + log_info "Stopping server..." + kill $SERVER_PID 2>/dev/null || true + log_info "Server tests completed" +fi + +log_info "All conformance tests completed successfully"