Add requestModifier property to HTTPClientTransport (#152)

* Add requestModifier property to HTTPClientTransport

Co-authored-by: Tim Wang <tim.wang@ing.com>

* Add test coverage

---------

Co-authored-by: Tim Wang <tim.wang@ing.com>

Sources/MCP/Base/Transports/HTTPClientTransport.swift

@@ -32,9 +32,14 @@ import Logging
 /// ```swift
 /// import MCP
 ///
-/// // Create a streaming HTTP transport
+/// // Create a streaming HTTP transport with bearer token authentication
 /// let transport = HTTPClientTransport(
-///     endpoint: URL(string: "http://localhost:8080")!,
+///     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
@@ -60,6 +65,9 @@ public actor HTTPClientTransport: Transport {
     /// 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
+
     private var isConnected = false
     private let messageStream: AsyncThrowingStream<Data, Swift.Error>
     private let messageContinuation: AsyncThrowingStream<Data, Swift.Error>.Continuation
@@ -74,12 +82,14 @@ public actor HTTPClientTransport: Transport {
     ///   - configuration: URLSession configuration to use for HTTP requests
     ///   - streaming: Whether to enable SSE streaming mode (default: true)
     ///   - sseInitializationTimeout: Maximum time to wait for session ID before proceeding with SSE (default: 10 seconds)
+    ///   - 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,
+        requestModifier: @escaping (URLRequest) -> URLRequest = { $0 },
         logger: Logger? = nil
     ) {
         self.init(
@@ -87,6 +97,7 @@ public actor HTTPClientTransport: Transport {
             session: URLSession(configuration: configuration),
             streaming: streaming,
             sseInitializationTimeout: sseInitializationTimeout,
+            requestModifier: requestModifier,
             logger: logger
         )
     }
@@ -96,12 +107,14 @@ public actor HTTPClientTransport: Transport {
         session: URLSession,
         streaming: Bool = false,
         sseInitializationTimeout: TimeInterval = 10,
+        requestModifier: @escaping (URLRequest) -> URLRequest = { $0 },
         logger: Logger? = nil
     ) {
         self.endpoint = endpoint
         self.session = session
         self.streaming = streaming
         self.sseInitializationTimeout = sseInitializationTimeout
+        self.requestModifier = requestModifier
 
         // Create message stream
         var continuation: AsyncThrowingStream<Data, Swift.Error>.Continuation!
@@ -211,6 +224,9 @@ public actor HTTPClientTransport: Transport {
             request.addValue(sessionID, forHTTPHeaderField: "Mcp-Session-Id")
         }
 
+        // Apply request modifier
+        request = requestModifier(request)
+
         #if os(Linux)
             // Linux implementation using data(for:) instead of bytes(for:)
             let (responseData, response) = try await session.data(for: request)
@@ -480,6 +496,9 @@ public actor HTTPClientTransport: Transport {
                 request.addValue(sessionID, forHTTPHeaderField: "Mcp-Session-Id")
             }
 
+            // Apply request modifier
+            request = requestModifier(request)
+
             logger.debug("Starting SSE connection")
 
             // Create URLSession task for SSE

Tests/MCPTests/HTTPClientTransportTests.swift

@@ -279,6 +279,16 @@ import Testing
             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,
@@ -287,15 +297,6 @@ import Testing
             )
             try await transport.connect()
 
-            let messageData = #"{"jsonrpc":"2.0","method":"test","id":3}"#.data(using: .utf8)!
-
-            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))
-            }
-
             do {
                 try await transport.send(messageData)
                 Issue.record("Expected send to throw an error for 404")
@@ -316,6 +317,16 @@ import Testing
             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,
@@ -324,15 +335,6 @@ import Testing
             )
             try await transport.connect()
 
-            let messageData = #"{"jsonrpc":"2.0","method":"test","id":4}"#.data(using: .utf8)!
-
-            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))
-            }
-
             do {
                 try await transport.send(messageData)
                 Issue.record("Expected send to throw an error for 500")
@@ -353,22 +355,15 @@ import Testing
             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 = "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] (request: URLRequest) in
+                [testEndpoint, initialSessionID] (request: URLRequest) in
                 let response = HTTPURLResponse(
                     url: testEndpoint, statusCode: 200, httpVersion: "HTTP/1.1",
                     headerFields: [
@@ -377,11 +372,21 @@ import Testing
                     ])!
                 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] (request: URLRequest) in
+                [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)!
@@ -528,7 +533,6 @@ import Testing
 
                 await transport.disconnect()
             }
-        #endif  // !canImport(FoundationNetworking)
 
             @Test(
                 "Client with HTTP Transport complete flow", .httpClientTransportSetup,
@@ -576,7 +580,8 @@ import Testing
                                 == true)
                     case "POST":
                         #expect(
-                        request.allHTTPHeaderFields?["Accept"]?.contains("application/json") == true
+                            request.allHTTPHeaderFields?["Accept"]?.contains("application/json")
+                                == true
                         )
                     default:
                         Issue.record(
@@ -588,13 +593,15 @@ import Testing
                     let bodyData = request.readBody()
 
                     guard let bodyData = bodyData,
-                    let json = try JSONSerialization.jsonObject(with: bodyData) as? [String: Any],
+                        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)"
+                                NSLocalizedDescriptionKey:
+                                    "Invalid JSON-RPC message \(#file):\(#line)"
                             ])
                     }
 
@@ -617,7 +624,8 @@ import Testing
                         return (httpResponse, responseData)
                     } else if method == "tools/call" {
                         // Verify initialize was called first
-                    if let lastRequest = await tracker.getLastRequest(), lastRequest != .initialize
+                        if let lastRequest = await tracker.getLastRequest(),
+                            lastRequest != .initialize
                         {
                             #expect(Bool(false), "Initialize should be called before callTool")
                         }
@@ -673,5 +681,49 @@ import Testing
                 // 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()
+            }
+        #endif  // !canImport(FoundationNetworking)
     }
 #endif  // swift(>=6.1)